clawmates/crates/cm-llm/src/anthropic.rs

//! Anthropic Messages API provider (cloud target).

use async_stream::try_stream;
use eventsource_stream::Eventsource;
use futures::StreamExt;
use serde_json::{json, Value};

use crate::provider::{
    ChatRequest, ChatRole, ContentPart, EventStream, LlmError, LlmEvent, LlmProvider, StopReason,
};

/// The Claude Code identity line Anthropic requires as the first system block
/// when authenticating with a subscription OAuth token.
const OAUTH_SYSTEM_PREAMBLE: &str = "You are Claude Code, Anthropic's official CLI for Claude.";

/// Setup tokens minted by `claude setup-token` carry this prefix. API keys are
/// `sk-ant-api…`, so the shape is enough to pick the auth scheme.
fn is_setup_token(credential: &str) -> bool {
    credential.trim().starts_with("sk-ant-oat")
}

pub struct AnthropicProvider {
    client: reqwest::Client,
    base_url: String,
    api_key: String,
}

impl AnthropicProvider {
    pub fn new(api_key: String) -> AnthropicProvider {
        AnthropicProvider::with_base_url(api_key, "https://api.anthropic.com".into())
    }

    /// Base URL override exists for self-hosted proxies and tests.
    pub fn with_base_url(api_key: String, base_url: String) -> AnthropicProvider {
        AnthropicProvider {
            client: reqwest::Client::new(),
            base_url,
            api_key,
        }
    }

    /// Whether this provider is authenticating with a subscription token
    /// rather than an API key. Callers that report cost attribution care.
    pub fn is_subscription(&self) -> bool {
        is_setup_token(&self.api_key)
    }

    /// Prepend the Claude Code identity line, unless the caller's system
    /// prompt already opens with it (so repeated wrapping can't stack).
    fn with_oauth_preamble(system: &str) -> String {
        if system.trim_start().starts_with(OAUTH_SYSTEM_PREAMBLE) {
            return system.to_string();
        }
        if system.trim().is_empty() {
            return OAUTH_SYSTEM_PREAMBLE.to_string();
        }
        format!("{OAUTH_SYSTEM_PREAMBLE}\n\n{system}")
    }

    fn wire_messages(request: &ChatRequest) -> Vec<Value> {
        request
            .messages
            .iter()
            .map(|message| {
                let role = match message.role {
                    ChatRole::User => "user",
                    ChatRole::Assistant => "assistant",
                };
                let content: Vec<Value> = message
                    .parts
                    .iter()
                    .map(|part| match part {
                        ContentPart::Text { text } => json!({"type": "text", "text": text}),
                        ContentPart::ToolUse { id, name, input } => {
                            json!({"type": "tool_use", "id": id,
                                   "name": crate::wire_tool_name(name), "input": input})
                        }
                        ContentPart::ToolResult {
                            tool_use_id,
                            content,
                        } => json!({
                            "type": "tool_result",
                            "tool_use_id": tool_use_id,
                            "content": content.to_string(),
                        }),
                    })
                    .collect();
                json!({"role": role, "content": content})
            })
            .collect()
    }
}

fn stop_reason(wire: &str) -> StopReason {
    match wire {
        "tool_use" => StopReason::ToolUse,
        "max_tokens" => StopReason::MaxTokens,
        _ => StopReason::EndTurn,
    }
}

#[async_trait::async_trait]
impl LlmProvider for AnthropicProvider {
    async fn stream(&self, request: ChatRequest) -> Result<EventStream, LlmError> {
        let mut tools: Vec<Value> = request
            .tools
            .iter()
            .map(|t| {
                json!({
                    "name": crate::wire_tool_name(&t.name),
                    "description": t.description,
                    "input_schema": t.input_schema,
                })
            })
            .collect();
        if request.web_search {
            // Anthropic server-side web search — the model searches the web itself.
            tools.push(json!({"type": "web_search_20250305", "name": "web_search", "max_uses": 5}));
        }
        // A subscription token authenticates as Claude Code: bearer auth, the
        // Claude Code beta set, and a system prompt whose first line is the
        // Claude Code identity. Sending it as `x-api-key` returns 401.
        let oauth = is_setup_token(&self.api_key);
        let system = if oauth {
            Self::with_oauth_preamble(&request.system)
        } else {
            request.system.clone()
        };
        let body = json!({
            "model": request.model,
            "max_tokens": request.max_tokens,
            "system": system,
            "messages": AnthropicProvider::wire_messages(&request),
            "tools": tools,
            "stream": true,
        });

        let mut req = self
            .client
            .post(format!("{}/v1/messages", self.base_url))
            .header("anthropic-version", "2023-06-01");
        req = if oauth {
            req.header("authorization", format!("Bearer {}", self.api_key))
                .header(
                    "anthropic-beta",
                    "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14",
                )
        } else {
            req.header("x-api-key", &self.api_key)
        };
        let response = req
            .json(&body)
            .send()
            .await
            .map_err(|e| LlmError::Transport(e.to_string()))?;
        if !response.status().is_success() {
            let status = response.status();
            let detail = response.text().await.unwrap_or_default();
            return Err(LlmError::Api(format!("{status}: {detail}")));
        }

        let mut sse = response.bytes_stream().eventsource();
        let stream = try_stream! {
            // tool_use input arrives as accumulated partial JSON between
            // content_block_start and content_block_stop.
            let mut pending_tool: Option<(String, String, String)> = None; // id, name, json
            let mut input_tokens: u32 = 0;
            while let Some(event) = sse.next().await {
                let event = event.map_err(|e| LlmError::Transport(e.to_string()))?;
                let data: Value = serde_json::from_str(&event.data)
                    .map_err(|e| LlmError::Wire(format!("{e}: {}", event.data)))?;
                match data["type"].as_str().unwrap_or_default() {
                    "content_block_start" => {
                        let block = &data["content_block"];
                        if block["type"] == "tool_use" {
                            pending_tool = Some((
                                block["id"].as_str().unwrap_or_default().to_owned(),
                                crate::internal_tool_name(
                                    block["name"].as_str().unwrap_or_default(),
                                ),
                                String::new(),
                            ));
                        }
                    }
                    "content_block_delta" => {
                        let delta = &data["delta"];
                        match delta["type"].as_str().unwrap_or_default() {
                            "text_delta" => {
                                if let Some(text) = delta["text"].as_str() {
                                    yield LlmEvent::TextDelta(text.to_owned());
                                }
                            }
                            "input_json_delta" => {
                                if let Some((_, _, buf)) = pending_tool.as_mut() {
                                    buf.push_str(delta["partial_json"].as_str().unwrap_or_default());
                                }
                            }
                            _ => {}
                        }
                    }
                    "content_block_stop" => {
                        if let Some((id, name, buf)) = pending_tool.take() {
                            let input: Value = if buf.is_empty() {
                                json!({})
                            } else {
                                serde_json::from_str(&buf)
                                    .map_err(|e| LlmError::Wire(format!("tool input: {e}")))?
                            };
                            yield LlmEvent::ToolUse { id, name, input };
                        }
                    }
                    "message_start" => {
                        input_tokens =
                            data["message"]["usage"]["input_tokens"].as_u64().unwrap_or(0) as u32;
                    }
                    "message_delta" => {
                        if let Some(out) = data["usage"]["output_tokens"].as_u64() {
                            yield LlmEvent::Usage {
                                input_tokens,
                                output_tokens: out as u32,
                            };
                        }
                        if let Some(reason) = data["delta"]["stop_reason"].as_str() {
                            yield LlmEvent::Stop(stop_reason(reason));
                        }
                    }
                    "error" => {
                        Err(LlmError::Api(data["error"]["message"]
                            .as_str()
                            .unwrap_or("unknown")
                            .to_owned()))?;
                    }
                    _ => {}
                }
            }
        };
        Ok(Box::pin(stream))
    }
}

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

    #[test]
    fn setup_tokens_are_distinguished_from_api_keys() {
        assert!(is_setup_token("sk-ant-oat01-abc"));
        assert!(is_setup_token("  sk-ant-oat01-abc  "), "trims first");
        assert!(!is_setup_token("sk-ant-api03-abc"));
        assert!(!is_setup_token(""));
    }

    #[test]
    fn oauth_preamble_is_prepended_once() {
        let once = AnthropicProvider::with_oauth_preamble("Judge the condition.");
        assert!(once.starts_with(OAUTH_SYSTEM_PREAMBLE));
        assert!(once.ends_with("Judge the condition."));
        // Re-wrapping must not stack the identity line — the API rejects a
        // system prompt that doesn't *start* with it, and duplicating it is
        // pure token waste on a path whose whole point is being cheap.
        let twice = AnthropicProvider::with_oauth_preamble(&once);
        assert_eq!(once, twice);
        assert_eq!(twice.matches(OAUTH_SYSTEM_PREAMBLE).count(), 1);
    }

    #[test]
    fn oauth_preamble_handles_an_empty_system_prompt() {
        assert_eq!(
            AnthropicProvider::with_oauth_preamble("   "),
            OAUTH_SYSTEM_PREAMBLE
        );
    }

    #[test]
    fn is_subscription_reports_the_credential_kind() {
        assert!(AnthropicProvider::new("sk-ant-oat01-x".into()).is_subscription());
        assert!(!AnthropicProvider::new("sk-ant-api03-x".into()).is_subscription());
    }
}