agentctl_auth/

claude.rs

//! Claude Messages API client with OAuth, stealth headers, and automatic token rotation.
//!
//! ## Tool Use Support
//!
//! This module supports Claude's tool use API for building agent loops:
//!
//! ```no_run
//! use agentctl_auth::claude::{Client, Tool, ToolHandler, ToolOutput};
//! use anyhow::Result;
//! use async_trait::async_trait;
//!
//! struct MyHandler;
//!
//! #[async_trait]
//! impl ToolHandler for MyHandler {
//!     async fn handle(&self, name: &str, input: &serde_json::Value) -> Result<ToolOutput> {
//!         match name {
//!             "read_file" => Ok(ToolOutput::success("file contents")),
//!             _ => Ok(ToolOutput::error("unknown tool")),
//!         }
//!     }
//! }
//! ```

use crate::pool::AuthPool;
use anyhow::{Context, Result};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::sync::{Arc, Mutex};

/// Claude Messages API client.
pub struct Client {
    http: reqwest::Client,
    pool: Option<Arc<Mutex<AuthPool>>>,
    current_credential: Arc<Mutex<Option<String>>>,
    base_url: String,
}

impl std::fmt::Debug for Client {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Client")
            .field("base_url", &self.base_url)
            .finish_non_exhaustive()
    }
}

/// Claude API message with support for both text and multi-content (tool use).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub role: String,
    #[serde(flatten)]
    pub content: MessageContent,
}

/// Message content — either simple text or array of content blocks.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MessageContent {
    /// Simple text content (common case).
    Text { content: String },
    /// Multi-block content for tool use/results.
    Blocks { content: Vec<ContentBlock> },
}

impl Message {
    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: "user".to_string(),
            content: MessageContent::Text { content: content.into() },
        }
    }

    pub fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: "assistant".to_string(),
            content: MessageContent::Text { content: content.into() },
        }
    }

    /// Create an assistant message with content blocks (for tool_use responses).
    pub fn assistant_blocks(blocks: Vec<ContentBlock>) -> Self {
        Self {
            role: "assistant".to_string(),
            content: MessageContent::Blocks { content: blocks },
        }
    }

    /// Create a user message with tool results.
    pub fn tool_results(results: Vec<ToolResultBlock>) -> Self {
        Self {
            role: "user".to_string(),
            content: MessageContent::Blocks {
                content: results.into_iter().map(|r| ContentBlock::ToolResult { result: r }).collect(),
            },
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Tool Use Types
// ═══════════════════════════════════════════════════════════════════════════════

/// Tool definition for Claude API.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Tool {
    pub name: String,
    pub description: String,
    pub input_schema: serde_json::Value,
}

impl Tool {
    /// Create a new tool definition.
    pub fn new(name: impl Into<String>, description: impl Into<String>, input_schema: serde_json::Value) -> Self {
        Self {
            name: name.into(),
            description: description.into(),
            input_schema,
        }
    }
}

/// A tool use block from Claude's response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolUseBlock {
    pub id: String,
    pub name: String,
    pub input: serde_json::Value,
}

/// A tool result block for user messages.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResultBlock {
    pub tool_use_id: String,
    pub content: String,
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub is_error: bool,
}

impl ToolResultBlock {
    pub fn success(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            tool_use_id: tool_use_id.into(),
            content: content.into(),
            is_error: false,
        }
    }

    pub fn error(tool_use_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            tool_use_id: tool_use_id.into(),
            content: content.into(),
            is_error: true,
        }
    }
}

/// Trait for handling tool calls in agent loops.
#[async_trait]
pub trait ToolHandler: Send + Sync {
    /// Handle a tool call and return the result.
    async fn handle(&self, name: &str, input: &serde_json::Value) -> Result<ToolOutput>;
}

/// Output from a tool handler.
#[derive(Debug, Clone)]
pub struct ToolOutput {
    pub content: String,
    pub is_error: bool,
}

impl ToolOutput {
    pub fn success(content: impl Into<String>) -> Self {
        Self { content: content.into(), is_error: false }
    }

    pub fn error(content: impl Into<String>) -> Self {
        Self { content: content.into(), is_error: true }
    }
}

/// Result of running an agent loop.
#[derive(Debug, Clone)]
pub struct AgentLoopResult {
    /// Final text output from the agent.
    pub final_text: String,
    /// Total input tokens used across all turns.
    pub total_input_tokens: u64,
    /// Total output tokens used across all turns.
    pub total_output_tokens: u64,
    /// Number of conversation turns used.
    pub turns_used: u32,
    /// Names of tools that were called.
    pub tool_calls: Vec<String>,
}

/// Claude API request body.
#[derive(Debug, Serialize)]
struct MessagesRequest<'a> {
    model: &'a str,
    messages: &'a [Message],
    max_tokens: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    system: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<&'a [Tool]>,
}

/// Claude API response.
#[derive(Debug, Deserialize)]
pub struct MessagesResponse {
    pub id: String,
    #[serde(rename = "type")]
    pub response_type: String,
    pub role: String,
    pub content: Vec<ContentBlock>,
    pub model: String,
    pub stop_reason: Option<String>,
    pub usage: Usage,
}

/// Content block in Claude API responses.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
    /// Text content.
    Text {
        text: String,
    },
    /// Tool use request from Claude.
    ToolUse {
        id: String,
        name: String,
        input: serde_json::Value,
    },
    /// Tool result from user (for multi-turn).
    ToolResult {
        #[serde(flatten)]
        result: ToolResultBlock,
    },
}

impl ContentBlock {
    /// Get the text content if this is a text block.
    pub fn as_text(&self) -> Option<&str> {
        match self {
            ContentBlock::Text { text } => Some(text),
            _ => None,
        }
    }

    /// Get the tool use info if this is a tool_use block.
    pub fn as_tool_use(&self) -> Option<(&str, &str, &serde_json::Value)> {
        match self {
            ContentBlock::ToolUse { id, name, input } => Some((id, name, input)),
            _ => None,
        }
    }
}

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

impl Client {
    /// Create a new client with a single token (no pool/rotation).
    #[allow(dead_code)]
    pub fn with_token(_token: impl Into<String>) -> Self {
        Self {
            http: Self::build_http_client(),
            pool: None,
            current_credential: Arc::new(Mutex::new(None)),
            base_url: "https://api.anthropic.com".to_string(),
        }
    }

    /// Create a client builder.
    pub fn builder() -> ClientBuilder {
        ClientBuilder::new()
    }

    fn build_http_client() -> reqwest::Client {
        reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(120))
            .build()
            .expect("Failed to build HTTP client")
    }

    /// Send a Messages API request.
    pub async fn message(
        &self,
        model: &str,
        messages: &[Message],
        max_tokens: u32,
    ) -> Result<MessagesResponse> {
        self.message_with_system(model, messages, max_tokens, None)
            .await
    }

    /// Send a Messages API request with system prompt.
    pub async fn message_with_system(
        &self,
        model: &str,
        messages: &[Message],
        max_tokens: u32,
        system: Option<&str>,
    ) -> Result<MessagesResponse> {
        self.message_with_tools(model, messages, max_tokens, system, None).await
    }

    /// Send a Messages API request with system prompt and tools.
    pub async fn message_with_tools(
        &self,
        model: &str,
        messages: &[Message],
        max_tokens: u32,
        system: Option<&str>,
        tools: Option<&[Tool]>,
    ) -> Result<MessagesResponse> {
        let body = MessagesRequest {
            model,
            messages,
            max_tokens,
            system,
            tools,
        };

        let mut attempts = 0;
        let max_attempts = if self.pool.is_some() { 3 } else { 1 };

        loop {
            attempts += 1;
            let (token, cred_name) = self.get_current_token()?;

            let response = self
                .http
                .post(format!("{}/v1/messages", self.base_url))
                .header("x-api-key", &token)
                .header("anthropic-version", "2023-06-01")
                .header("content-type", "application/json")
                // Stealth headers (from Pi-AI SDK)
                .header("anthropic-beta", "claude-code-20250219,oauth-2025-04-20")
                .header("user-agent", "claude-cli/2.1.39 (external, cli)")
                .header("x-app", "cli")
                .header("anthropic-dangerous-direct-browser-access", "true")
                .json(&body)
                .send()
                .await
                .context("Failed to send request to Claude API")?;

            let status = response.status();

            if status.is_success() {
                // Success — record usage and return
                if let Some(ref pool) = self.pool {
                    if let Some(ref name) = cred_name {
                        pool.lock().unwrap().record_usage(name, true);
                    }
                }
                let result: MessagesResponse = response
                    .json()
                    .await
                    .context("Failed to parse Claude API response")?;
                return Ok(result);
            } else if status.as_u16() == 429 {
                // Rate limit — record failure and rotate
                tracing::warn!(
                    credential = cred_name.as_deref().unwrap_or("<unknown>"),
                    "Claude API 429 rate limit, rotating credential"
                );

                if let Some(ref pool) = self.pool {
                    let mut pool_guard = pool.lock().unwrap();
                    if let Some(ref name) = cred_name {
                        pool_guard.record_usage(name, false);

                        // Try to get next credential
                        if let Some((next_name, _next_cred)) =
                            pool_guard.next_credential("anthropic", name)
                        {
                            tracing::info!(next = next_name, "Rotating to next credential");
                            *self.current_credential.lock().unwrap() =
                                Some(next_name.to_string());
                            if attempts < max_attempts {
                                continue;
                            }
                        }
                    }
                }

                // Out of credentials or retries
                anyhow::bail!("Claude API rate limit (429) and no more credentials to rotate");
            } else {
                // Other error
                let error_text = response
                    .text()
                    .await
                    .unwrap_or_else(|_| "<failed to read error>".to_string());
                anyhow::bail!("Claude API error {}: {}", status, error_text);
            }
        }
    }

    /// Run an agent loop with tool use support.
    ///
    /// This method sends the initial message, handles tool calls via the handler,
    /// and continues the conversation until Claude stops calling tools or max_turns
    /// is reached.
    ///
    /// # Arguments
    /// - `model`: The model to use (e.g., "claude-sonnet-4-5")
    /// - `system`: System prompt
    /// - `initial_message`: The user's initial request
    /// - `tools`: Tool definitions
    /// - `max_turns`: Maximum conversation turns (each turn = one API call)
    /// - `tool_handler`: Handler for executing tools
    ///
    /// # Returns
    /// An [`AgentLoopResult`] with final text, total tokens, turns used, and tool calls.
    pub async fn run_agent_loop(
        &self,
        model: &str,
        system: &str,
        initial_message: &str,
        tools: &[Tool],
        max_turns: u32,
        tool_handler: &dyn ToolHandler,
    ) -> Result<AgentLoopResult> {
        let mut messages = vec![Message::user(initial_message)];
        let mut total_input_tokens: u64 = 0;
        let mut total_output_tokens: u64 = 0;
        let mut turns_used: u32 = 0;
        let mut tool_calls: Vec<String> = Vec::new();
        let mut final_text = String::new();

        loop {
            if turns_used >= max_turns {
                tracing::warn!(turns = turns_used, max = max_turns, "Agent loop hit max turns");
                break;
            }

            turns_used += 1;
            tracing::debug!(turn = turns_used, "Agent loop turn");

            let response = self
                .message_with_tools(model, &messages, 16384, Some(system), Some(tools))
                .await?;

            total_input_tokens += response.usage.input_tokens as u64;
            total_output_tokens += response.usage.output_tokens as u64;

            // Collect tool uses and text from response
            let mut pending_tool_uses: Vec<(String, String, serde_json::Value)> = Vec::new();
            let mut response_text = String::new();

            for block in &response.content {
                match block {
                    ContentBlock::Text { text } => {
                        response_text.push_str(text);
                    }
                    ContentBlock::ToolUse { id, name, input } => {
                        pending_tool_uses.push((id.clone(), name.clone(), input.clone()));
                        tool_calls.push(name.clone());
                    }
                    ContentBlock::ToolResult { .. } => {
                        // Shouldn't appear in response, but ignore
                    }
                }
            }

            final_text = response_text;

            // Check stop reason
            let stop_reason = response.stop_reason.as_deref().unwrap_or("");
            if stop_reason == "end_turn" && pending_tool_uses.is_empty() {
                // Normal completion
                tracing::debug!("Agent loop completed normally");
                break;
            }

            if pending_tool_uses.is_empty() {
                // No tool calls and not end_turn — unexpected but treat as complete
                tracing::debug!(stop_reason, "Agent loop ended (no tool calls)");
                break;
            }

            // Add assistant message with tool uses to history
            messages.push(Message::assistant_blocks(response.content.clone()));

            // Execute tools and collect results
            let mut results: Vec<ToolResultBlock> = Vec::new();
            for (id, name, input) in pending_tool_uses {
                tracing::debug!(tool = %name, "Executing tool");
                let output = tool_handler.handle(&name, &input).await;
                match output {
                    Ok(out) => {
                        if out.is_error {
                            results.push(ToolResultBlock::error(&id, out.content));
                        } else {
                            results.push(ToolResultBlock::success(&id, out.content));
                        }
                    }
                    Err(e) => {
                        results.push(ToolResultBlock::error(&id, format!("Tool error: {}", e)));
                    }
                }
            }

            // Add tool results as user message
            messages.push(Message::tool_results(results));
        }

        Ok(AgentLoopResult {
            final_text,
            total_input_tokens,
            total_output_tokens,
            turns_used,
            tool_calls,
        })
    }

    fn get_current_token(&self) -> Result<(String, Option<String>)> {
        if let Some(ref pool) = self.pool {
            let pool_guard = pool.lock().unwrap();

            // If current_credential is set, use it
            let current_lock = self.current_credential.lock().unwrap();
            if let Some(ref name) = *current_lock {
                if let Some(cred) = pool_guard.get(name) {
                    let token = cred
                        .resolved_token()
                        .ok_or_else(|| anyhow::anyhow!("Credential '{}' has no token", name))?
                        .to_string();
                    return Ok((token, Some(name.clone())));
                }
            }
            drop(current_lock);

            // Otherwise, get default
            if let Some((name, cred)) = pool_guard.get_default("anthropic") {
                let token = cred
                    .resolved_token()
                    .ok_or_else(|| anyhow::anyhow!("Credential '{}' has no token", name))?
                    .to_string();
                *self.current_credential.lock().unwrap() = Some(name.to_string());
                return Ok((token, Some(name.to_string())));
            }

            anyhow::bail!("No anthropic credentials in pool");
        } else {
            // No pool — must have been constructed with with_token (not currently implemented)
            anyhow::bail!("No pool configured and with_token not yet implemented");
        }
    }
}

/// Builder for Claude client.
pub struct ClientBuilder {
    pool: Option<Arc<Mutex<AuthPool>>>,
    base_url: Option<String>,
}

impl ClientBuilder {
    pub fn new() -> Self {
        Self {
            pool: None,
            base_url: None,
        }
    }

    /// Use an auth pool for automatic token rotation.
    pub fn pool(mut self, pool: &AuthPool) -> Self {
        self.pool = Some(Arc::new(Mutex::new(pool.clone())));
        self
    }

    /// Set a custom base URL (e.g., for proxies).
    #[allow(dead_code)]
    pub fn base_url(mut self, url: impl Into<String>) -> Self {
        self.base_url = Some(url.into());
        self
    }

    /// Build the client.
    pub fn build(self) -> Result<Client> {
        let pool = self
            .pool
            .ok_or_else(|| anyhow::anyhow!("Pool is required (use .pool())"))?;

        Ok(Client {
            http: Client::build_http_client(),
            pool: Some(pool),
            current_credential: Arc::new(Mutex::new(None)),
            base_url: self
                .base_url
                .unwrap_or_else(|| "https://api.anthropic.com".to_string()),
        })
    }
}

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

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

    #[test]
    fn test_message_construction() {
        let msg = Message::user("Hello!");
        assert_eq!(msg.role, "user");
        match msg.content {
            MessageContent::Text { content } => assert_eq!(content, "Hello!"),
            _ => panic!("Expected text content"),
        }

        let msg = Message::assistant("Hi there");
        assert_eq!(msg.role, "assistant");
        match msg.content {
            MessageContent::Text { content } => assert_eq!(content, "Hi there"),
            _ => panic!("Expected text content"),
        }
    }

    #[test]
    fn test_tool_result_block() {
        let success = ToolResultBlock::success("id-123", "file contents");
        assert_eq!(success.tool_use_id, "id-123");
        assert_eq!(success.content, "file contents");
        assert!(!success.is_error);

        let error = ToolResultBlock::error("id-456", "not found");
        assert!(error.is_error);
    }

    #[test]
    fn test_tool_definition() {
        let tool = Tool::new(
            "read_file",
            "Read a file's contents",
            serde_json::json!({
                "type": "object",
                "properties": {
                    "path": { "type": "string" }
                },
                "required": ["path"]
            }),
        );
        assert_eq!(tool.name, "read_file");
        assert_eq!(tool.description, "Read a file's contents");
    }

    #[test]
    fn test_content_block_helpers() {
        let text = ContentBlock::Text { text: "hello".to_string() };
        assert_eq!(text.as_text(), Some("hello"));
        assert!(text.as_tool_use().is_none());

        let tool_use = ContentBlock::ToolUse {
            id: "id-1".to_string(),
            name: "bash".to_string(),
            input: serde_json::json!({"command": "ls"}),
        };
        assert!(tool_use.as_text().is_none());
        let (id, name, input) = tool_use.as_tool_use().unwrap();
        assert_eq!(id, "id-1");
        assert_eq!(name, "bash");
        assert_eq!(input["command"], "ls");
    }

    #[tokio::test]
    async fn test_client_requires_pool() {
        let result = Client::builder().build();
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("Pool is required"));
    }
}