aura-gui 0.2.2

//! API de Messages de Anthropic: cuerpo de la petición y decodificador del stream.

use crate::ai::provider::Decoder;
use crate::ai::sse::SseEvent;
use crate::ai::types::{AiError, Block, Event, Message, Request, Role, StopReason, Usage};
use serde_json::{Value, json};

fn block(block: &Block) -> Value {
    match block {
        Block::Text(text) => json!({ "type": "text", "text": text }),
        Block::Image { media_type, data_base64, .. } => json!({
            "type": "image",
            "source": { "type": "base64", "media_type": media_type, "data": data_base64 },
        }),
        Block::ToolUse { id, name, input } => json!({ "type": "tool_use", "id": id, "name": name, "input": input }),
        Block::ToolResult { id, content, is_error } => {
            json!({ "type": "tool_result", "tool_use_id": id, "content": content, "is_error": is_error })
        }
    }
}

fn message(message: &Message) -> Value {
    let role = if message.role == Role::User { "user" } else { "assistant" };
    json!({ "role": role, "content": message.blocks.iter().map(block).collect::<Vec<_>>() })
}

/// Marca un punto de caché de prompt: lo anterior se relee a una fracción del precio y sin recalcularlo.
fn cache_mark(value: &mut Value) {
    value["cache_control"] = json!({ "type": "ephemeral" });
}

/// Con tres puntos de caché (herramientas, sistema y último mensaje) cada vuelta del agente solo paga lo nuevo:
/// el prefijo, que crece con cada herramienta llamada, se lee de caché.
pub fn body(request: &Request) -> Value {
    let mut messages: Vec<Value> = request.messages.iter().map(message).collect();
    if let Some(last) = messages.last_mut().and_then(|message| message["content"].as_array_mut()?.last_mut()) {
        cache_mark(last);
    }
    let mut body = json!({
        "model": request.model,
        "max_tokens": request.max_tokens,
        "stream": true,
        "messages": messages,
    });
    if let Some(system) = &request.system {
        let mut block = json!({ "type": "text", "text": system });
        cache_mark(&mut block);
        body["system"] = json!([block]);
    }
    if !request.tools.is_empty() {
        let mut tools: Vec<Value> = request
            .tools
            .iter()
            .map(|tool| json!({ "name": tool.name, "description": tool.description, "input_schema": tool.schema }))
            .collect();
        if let Some(last) = tools.last_mut() {
            cache_mark(last);
        }
        body["tools"] = json!(tools);
    }
    body
}

/// Identidad que el servidor exige a los tokens de suscripción, siempre como primer bloque del sistema.
const SUBSCRIPTION_IDENTITY: &str = "You are Claude Code, Anthropic's official CLI for Claude.";

/// Cuerpo para la suscripción Pro/Max: el mismo que con clave de API más la identidad del cliente.
pub fn subscription_body(request: &Request) -> Value {
    let mut body = body(request);
    let mut system = vec![json!({ "type": "text", "text": SUBSCRIPTION_IDENTITY })];
    if let Some(rest) = body["system"].as_array() {
        system.extend(rest.iter().cloned());
    }
    body["system"] = json!(system);
    body
}

pub fn subscription_headers(token: &str) -> Vec<(String, String)> {
    [
        ("authorization", format!("Bearer {token}")),
        ("anthropic-version", "2023-06-01".to_owned()),
        ("anthropic-beta", "claude-code-20250219,oauth-2025-04-20".to_owned()),
        ("anthropic-dangerous-direct-browser-access", "true".to_owned()),
        ("user-agent", "claude-cli/2.1.280".to_owned()),
        ("x-app", "cli".to_owned()),
    ]
    .into_iter()
    .map(|(name, value)| (name.to_owned(), value))
    .collect()
}

fn stop_reason(reason: &str) -> StopReason {
    match reason {
        "end_turn" | "stop_sequence" => StopReason::EndTurn,
        "tool_use" => StopReason::ToolUse,
        "max_tokens" => StopReason::MaxTokens,
        other => StopReason::Other(other.to_owned()),
    }
}

#[derive(Default)]
pub struct AnthropicDecoder {
    /// Herramienta en construcción: `(id, nombre, JSON parcial de los argumentos)`.
    tool: Option<(String, String, String)>,
    usage: Usage,
    stop: Option<StopReason>,
    finished: bool,
    complete: bool,
}

impl Decoder for AnthropicDecoder {
    fn feed(&mut self, event: &SseEvent, out: &mut Vec<Event>) -> Result<(), AiError> {
        let data: Value = match serde_json::from_str(&event.data) {
            Ok(data) => data,
            Err(_) if event.data.is_empty() => return Ok(()),
            Err(error) => return Err(AiError::Decode(error.to_string())),
        };
        let text = |value: &Value| value.as_str().unwrap_or_default().to_owned();
        match event.event.as_deref().or(data["type"].as_str()) {
            Some("message_start") => {
                // Los tokens leídos o escritos en la caché del prompt también son contexto enviado.
                let usage = &data["message"]["usage"];
                let tokens = |key: &str| usage[key].as_u64().unwrap_or(0) as u32;
                self.usage.input_tokens =
                    tokens("input_tokens") + tokens("cache_read_input_tokens") + tokens("cache_creation_input_tokens");
            }
            Some("content_block_start") if data["content_block"]["type"] == "tool_use" => {
                let start = &data["content_block"];
                self.tool = Some((text(&start["id"]), text(&start["name"]), String::new()));
            }
            Some("content_block_delta") => match data["delta"]["type"].as_str() {
                Some("text_delta") => {
                    let delta = text(&data["delta"]["text"]);
                    if !delta.is_empty() {
                        out.push(Event::Text(delta));
                    }
                }
                Some("input_json_delta") => {
                    if let Some((_, _, json)) = self.tool.as_mut() {
                        json.push_str(data["delta"]["partial_json"].as_str().unwrap_or_default());
                    }
                }
                _ => {}
            },
            Some("content_block_stop") => {
                if let Some((id, name, json)) = self.tool.take() {
                    out.push(Event::ToolUse { id, name, input: crate::ai::wire::tool_input(&json) });
                }
            }
            Some("message_delta") => {
                if let Some(reason) = data["delta"]["stop_reason"].as_str() {
                    self.stop = Some(stop_reason(reason));
                }
                self.usage.output_tokens = data["usage"]["output_tokens"].as_u64().unwrap_or(0) as u32;
            }
            Some("message_stop") => {
                self.complete = true;
                self.finish(out);
            }
            Some("error") => return Err(AiError::Api(text(&data["error"]["message"]))),
            _ => {}
        }
        Ok(())
    }

    fn finish(&mut self, out: &mut Vec<Event>) {
        if std::mem::replace(&mut self.finished, true) {
            return;
        }
        out.push(Event::Usage(self.usage));
        out.push(Event::Done(self.stop.take().unwrap_or(StopReason::EndTurn)));
    }

    fn is_complete(&self) -> bool {
        self.complete
    }
}

#[cfg(test)]
mod tests;