flow_adapter_ai/providers/

claude.rs

use async_trait::async_trait;
use futures_util::StreamExt;
use serde::Deserialize;
use serde_json::json;
use std::time::Instant;

use crate::error::{redact_error_body, CloudAiError};
use crate::registry::{CloudAiProvider, ModelCapabilityInfo};
use crate::request::{CloudAiRequest, CloudAiResponse, ToolDispatcher, ToolSpec};
use crate::stream::{LlmStreamEvent, LlmStreamSink};

const NAME: &str = "claude";
const ENV_VAR: &str = "ANTHROPIC_API_KEY";
const ENDPOINT: &str = "https://api.anthropic.com/v1/messages";
const ANTHROPIC_VERSION: &str = "2023-06-01";

const DEFAULT_MODELS: &[&str] = &["claude-opus-4-7", "claude-sonnet-4-6", "claude-haiku-4-5"];

pub struct ClaudeProvider {
    client: reqwest::Client,
}

impl ClaudeProvider {
    pub fn new() -> Self {
        Self {
            client: reqwest::Client::new(),
        }
    }
}

impl Default for ClaudeProvider {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl CloudAiProvider for ClaudeProvider {
    fn name(&self) -> &str {
        NAME
    }

    fn env_var(&self) -> &str {
        ENV_VAR
    }

    fn default_models(&self) -> &[&str] {
        DEFAULT_MODELS
    }

    fn model_capabilities(&self) -> Vec<ModelCapabilityInfo> {
        vec![
            ModelCapabilityInfo::new(
                "claude-opus-4-7",
                &["reasoning", "vision", "tool_use", "code"],
            ),
            ModelCapabilityInfo::new(
                "claude-sonnet-4-6",
                &["reasoning", "vision", "tool_use", "code"],
            ),
            ModelCapabilityInfo::new("claude-haiku-4-5", &["vision", "tool_use", "code"]),
        ]
    }

    async fn invoke(&self, req: &CloudAiRequest) -> Result<CloudAiResponse, CloudAiError> {
        let body = build_body(req);

        let started = Instant::now();
        let resp = self
            .client
            .post(ENDPOINT)
            .header("x-api-key", &req.api_key)
            .header("anthropic-version", ANTHROPIC_VERSION)
            .header("content-type", "application/json")
            .json(&body)
            .send()
            .await
            .map_err(|e| CloudAiError::Http {
                provider: NAME.into(),
                reason: e.to_string(),
            })?;

        let status = resp.status();
        let raw_body = resp.text().await.map_err(|e| CloudAiError::Http {
            provider: NAME.into(),
            reason: e.to_string(),
        })?;
        let latency_ms = started.elapsed().as_millis() as u64;

        if !status.is_success() {
            return Err(CloudAiError::Status {
                provider: NAME.into(),
                status: status.as_u16(),
                body: redact_error_body(&raw_body),
            });
        }

        parse_response(&raw_body, latency_ms)
    }

    async fn invoke_stream(
        &self,
        req: &CloudAiRequest,
        sink: &dyn LlmStreamSink,
    ) -> Result<CloudAiResponse, CloudAiError> {
        let call_id = req.call_id.clone().unwrap_or_default();
        let mut body = build_body(req);
        body["stream"] = json!(true);
        let started = Instant::now();
        let resp = self
            .client
            .post(ENDPOINT)
            .header("x-api-key", &req.api_key)
            .header("anthropic-version", ANTHROPIC_VERSION)
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .json(&body)
            .send()
            .await
            .map_err(|e| CloudAiError::Http {
                provider: NAME.into(),
                reason: e.to_string(),
            })?;

        let status = resp.status();
        if !status.is_success() {
            let raw = resp.text().await.unwrap_or_default();
            let err = CloudAiError::Status {
                provider: NAME.into(),
                status: status.as_u16(),
                body: redact_error_body(&raw),
            };
            sink.emit(LlmStreamEvent::error(&call_id, NAME, err.to_string()))
                .await;
            return Err(err);
        }

        let mut stream = resp.bytes_stream();
        let mut buffer = String::new();
        let mut full_text = String::new();
        let mut model_name = String::new();
        let mut finish_reason = String::from("unknown");
        let mut input_tokens: u32 = 0;
        let mut output_tokens: u32 = 0;
        let mut errored: Option<String> = None;

        while let Some(chunk) = stream.next().await {
            let bytes = match chunk {
                Ok(b) => b,
                Err(e) => {
                    errored = Some(e.to_string());
                    break;
                }
            };
            buffer.push_str(&String::from_utf8_lossy(&bytes));
            // Anthropic frames are separated by `\n\n` like generic SSE.
            while let Some(idx) = buffer.find("\n\n") {
                let frame = buffer[..idx].to_string();
                buffer.drain(..idx + 2);
                for raw_line in frame.split('\n') {
                    let line = raw_line.trim_start();
                    let Some(payload) = line.strip_prefix("data:") else {
                        continue;
                    };
                    let payload = payload.trim();
                    if payload.is_empty() {
                        continue;
                    }
                    let value: serde_json::Value = match serde_json::from_str(payload) {
                        Ok(v) => v,
                        Err(_) => continue,
                    };
                    let kind = value.get("type").and_then(|t| t.as_str()).unwrap_or("");
                    match kind {
                        "message_start" => {
                            if let Some(msg) = value.get("message") {
                                if let Some(m) = msg.get("model").and_then(|m| m.as_str()) {
                                    model_name = m.to_string();
                                }
                                if let Some(u) = msg.get("usage") {
                                    input_tokens = u
                                        .get("input_tokens")
                                        .and_then(|v| v.as_u64())
                                        .map(|v| v as u32)
                                        .unwrap_or(input_tokens);
                                }
                            }
                        }
                        "content_block_delta" => {
                            if let Some(delta) = value
                                .get("delta")
                                .and_then(|d| d.get("text"))
                                .and_then(|t| t.as_str())
                            {
                                if !delta.is_empty() {
                                    full_text.push_str(delta);
                                    sink.emit(LlmStreamEvent {
                                        call_id: call_id.clone(),
                                        provider: NAME.to_string(),
                                        model: model_name.clone(),
                                        delta: delta.to_string(),
                                        done: false,
                                        error: None,
                                    })
                                    .await;
                                }
                            }
                        }
                        "message_delta" => {
                            if let Some(reason) = value
                                .get("delta")
                                .and_then(|d| d.get("stop_reason"))
                                .and_then(|r| r.as_str())
                            {
                                finish_reason = reason.to_string();
                            }
                            if let Some(u) = value.get("usage") {
                                output_tokens = u
                                    .get("output_tokens")
                                    .and_then(|v| v.as_u64())
                                    .map(|v| v as u32)
                                    .unwrap_or(output_tokens);
                            }
                        }
                        "message_stop" => {
                            sink.emit(LlmStreamEvent {
                                call_id: call_id.clone(),
                                provider: NAME.to_string(),
                                model: model_name.clone(),
                                delta: String::new(),
                                done: true,
                                error: None,
                            })
                            .await;
                            let latency_ms = started.elapsed().as_millis() as u64;
                            return Ok(CloudAiResponse {
                                provider: NAME.to_string(),
                                model: if model_name.is_empty() {
                                    "unknown".to_string()
                                } else {
                                    model_name
                                },
                                text: full_text,
                                finish_reason,
                                input_tokens,
                                output_tokens,
                                latency_ms,
                            });
                        }
                        _ => {}
                    }
                }
            }
        }

        if let Some(err) = errored {
            sink.emit(LlmStreamEvent::error(&call_id, NAME, err.clone()))
                .await;
            return Err(CloudAiError::Http {
                provider: NAME.into(),
                reason: err,
            });
        }
        // Stream ended without an explicit `message_stop`. Treat it as
        // best-effort: emit the terminal event and return what we have.
        sink.emit(LlmStreamEvent {
            call_id: call_id.clone(),
            provider: NAME.to_string(),
            model: model_name.clone(),
            delta: String::new(),
            done: true,
            error: None,
        })
        .await;
        if full_text.is_empty() {
            return Err(CloudAiError::Shape {
                provider: NAME.into(),
                detail: "stream closed before any text was emitted".into(),
            });
        }
        let latency_ms = started.elapsed().as_millis() as u64;
        Ok(CloudAiResponse {
            provider: NAME.to_string(),
            model: if model_name.is_empty() {
                "unknown".to_string()
            } else {
                model_name
            },
            text: full_text,
            finish_reason,
            input_tokens,
            output_tokens,
            latency_ms,
        })
    }

    async fn invoke_tools(
        &self,
        req: &CloudAiRequest,
        tools: &[ToolSpec],
        dispatcher: &dyn ToolDispatcher,
        max_iters: usize,
    ) -> Result<CloudAiResponse, CloudAiError> {
        let started = Instant::now();
        let mut messages = vec![json!({ "role": "user", "content": user_content(req) })];
        let tools_body: Vec<_> = tools
            .iter()
            .map(|t| {
                json!({ "name": t.name, "description": t.description, "input_schema": t.parameters })
            })
            .collect();
        let mut input_tokens: u32 = 0;
        let mut output_tokens: u32 = 0;
        let mut model_out = req.model.clone();

        for _ in 0..max_iters.max(1) {
            let mut body = json!({
                "model": req.model,
                "max_tokens": req.max_tokens.unwrap_or(1024),
                "messages": messages,
                "tools": tools_body,
            });
            if let Some(t) = req.temperature {
                body["temperature"] = json!(t);
            }
            if let Some(s) = req.system.as_deref().filter(|s| !s.is_empty()) {
                body["system"] = json!(s);
            }
            apply_thinking(&mut body, req);

            let resp = self
                .client
                .post(ENDPOINT)
                .header("x-api-key", &req.api_key)
                .header("anthropic-version", ANTHROPIC_VERSION)
                .header("content-type", "application/json")
                .json(&body)
                .send()
                .await
                .map_err(|e| CloudAiError::Http {
                    provider: NAME.into(),
                    reason: e.to_string(),
                })?;
            let status = resp.status();
            let raw = resp.text().await.unwrap_or_default();
            if !status.is_success() {
                return Err(CloudAiError::Status {
                    provider: NAME.into(),
                    status: status.as_u16(),
                    body: redact_error_body(&raw),
                });
            }
            let v: serde_json::Value =
                serde_json::from_str(&raw).map_err(|e| CloudAiError::Parse {
                    provider: NAME.into(),
                    reason: e.to_string(),
                })?;
            if let Some(u) = v.get("usage") {
                input_tokens = input_tokens
                    .saturating_add(u.get("input_tokens").and_then(|x| x.as_u64()).unwrap_or(0) as u32);
                output_tokens = output_tokens.saturating_add(
                    u.get("output_tokens").and_then(|x| x.as_u64()).unwrap_or(0) as u32,
                );
            }
            if let Some(m) = v.get("model").and_then(|m| m.as_str()) {
                model_out = m.to_string();
            }
            let stop_reason = v
                .get("stop_reason")
                .and_then(|r| r.as_str())
                .unwrap_or("end_turn")
                .to_string();
            let content = v
                .get("content")
                .and_then(|c| c.as_array())
                .cloned()
                .unwrap_or_default();

            let mut text = String::new();
            let mut tool_uses: Vec<(String, String, serde_json::Value)> = Vec::new();
            for block in &content {
                match block.get("type").and_then(|t| t.as_str()) {
                    Some("text") => {
                        if let Some(t) = block.get("text").and_then(|t| t.as_str()) {
                            text.push_str(t);
                        }
                    }
                    Some("tool_use") => {
                        tool_uses.push((
                            block.get("id").and_then(|i| i.as_str()).unwrap_or("").to_string(),
                            block.get("name").and_then(|n| n.as_str()).unwrap_or("").to_string(),
                            block.get("input").cloned().unwrap_or_else(|| json!({})),
                        ));
                    }
                    _ => {}
                }
            }

            if tool_uses.is_empty() {
                if text.is_empty() {
                    return Err(CloudAiError::Shape {
                        provider: NAME.into(),
                        detail: "no text content blocks in response".into(),
                    });
                }
                return Ok(CloudAiResponse {
                    provider: NAME.into(),
                    model: model_out,
                    text,
                    finish_reason: stop_reason,
                    input_tokens,
                    output_tokens,
                    latency_ms: started.elapsed().as_millis() as u64,
                });
            }

            // Echo the assistant turn, then return each tool result as a user
            // turn of `tool_result` blocks (errors fed back so the model can
            // recover).
            messages.push(json!({ "role": "assistant", "content": content }));
            let mut results = Vec::new();
            for (id, name, input) in &tool_uses {
                let out = match dispatcher.call(name, input).await {
                    Ok(v) => v.to_string(),
                    Err(e) => json!({ "error": e }).to_string(),
                };
                results.push(json!({ "type": "tool_result", "tool_use_id": id, "content": out }));
            }
            messages.push(json!({ "role": "user", "content": results }));
        }

        Err(CloudAiError::Shape {
            provider: NAME.into(),
            detail: format!("tool loop did not converge within {max_iters} iterations"),
        })
    }
}

/// User-message content: a plain string, or `text` + `image` blocks when the
/// request carries images (Claude multimodal).
fn user_content(req: &CloudAiRequest) -> serde_json::Value {
    if req.images.is_empty() {
        return json!(req.prompt);
    }
    let mut blocks = vec![json!({ "type": "text", "text": req.prompt })];
    for img in &req.images {
        blocks.push(image_block(img));
    }
    json!(blocks)
}

/// A Claude `image` content block: a `data:` URL becomes a base64 source, an
/// `http(s)` URL a url source.
fn image_block(img: &str) -> serde_json::Value {
    if let Some(rest) = img.strip_prefix("data:") {
        if let Some((media_type, data)) = rest.split_once(";base64,") {
            return json!({
                "type": "image",
                "source": { "type": "base64", "media_type": media_type, "data": data },
            });
        }
    }
    json!({ "type": "image", "source": { "type": "url", "url": img } })
}

/// Apply the optional extended-thinking config. Claude requires
/// `budget_tokens < max_tokens` and disallows a custom temperature while
/// thinking, so bump max_tokens and drop temperature when enabling it.
fn apply_thinking(payload: &mut serde_json::Value, req: &CloudAiRequest) {
    if req.reasoning != Some(true) {
        return;
    }
    let budget = 1024u32;
    let max = payload["max_tokens"].as_u64().unwrap_or(1024) as u32;
    if max <= budget {
        payload["max_tokens"] = json!(budget + 1024);
    }
    if let Some(obj) = payload.as_object_mut() {
        obj.remove("temperature");
    }
    payload["thinking"] = json!({ "type": "enabled", "budget_tokens": budget });
}

pub fn build_body(req: &CloudAiRequest) -> serde_json::Value {
    let mut payload = json!({
        "model": req.model,
        "max_tokens": req.max_tokens.unwrap_or(1024),
        "messages": [{ "role": "user", "content": user_content(req) }],
    });
    if let Some(t) = req.temperature {
        payload["temperature"] = json!(t);
    }
    if let Some(s) = req.system.as_deref().filter(|s| !s.is_empty()) {
        // Claude's Messages API takes the system prompt as a top-level
        // string, not as a message in `messages`. Sending it via the
        // dedicated field improves instruction-following and lets
        // Anthropic's prompt cache pin the system text across turns.
        payload["system"] = json!(s);
    }
    apply_thinking(&mut payload, req);
    payload
}

#[derive(Debug, Deserialize)]
struct RawResponse {
    model: String,
    content: Vec<ContentBlock>,
    stop_reason: Option<String>,
    usage: Usage,
}

#[derive(Debug, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ContentBlock {
    Text {
        text: String,
    },
    #[serde(other)]
    Other,
}

#[derive(Debug, Deserialize)]
struct Usage {
    input_tokens: u32,
    output_tokens: u32,
}

pub fn parse_response(body: &str, latency_ms: u64) -> Result<CloudAiResponse, CloudAiError> {
    let raw: RawResponse = serde_json::from_str(body).map_err(|e| CloudAiError::Parse {
        provider: NAME.into(),
        reason: e.to_string(),
    })?;

    let text = raw
        .content
        .iter()
        .filter_map(|b| match b {
            ContentBlock::Text { text } => Some(text.as_str()),
            ContentBlock::Other => None,
        })
        .collect::<Vec<_>>()
        .join("");

    if text.is_empty() {
        return Err(CloudAiError::Shape {
            provider: NAME.into(),
            detail: "no text content blocks in response".into(),
        });
    }

    Ok(CloudAiResponse {
        provider: NAME.into(),
        model: raw.model,
        text,
        finish_reason: raw.stop_reason.unwrap_or_else(|| "unknown".into()),
        input_tokens: raw.usage.input_tokens,
        output_tokens: raw.usage.output_tokens,
        latency_ms,
    })
}