procyon 0.1.2

//! The `/v1/messages` wire format: request, response and stream event shapes.
//!
//! Separated from the domain vocabulary in `agent` because the two change for different reasons.
//! `ContentPart` is what the session log stores, so a resumed log must deserialize the shape it
//! was written with; the types here follow whatever Anthropic's API asks for, cache breakpoints
//! included. Keeping them in one file with the domain meant every protocol detail looked like
//! part of the agent's own vocabulary.

use serde::{Deserialize, Serialize};

use crate::agent::{ContentPart, Message, Role, ToolDefinition};

#[derive(Debug, Serialize)]
pub struct MessagesRequest<'a> {
    model: String,
    max_tokens: u32,
    messages: Vec<WireMessage<'a>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    system: Option<Vec<SystemBlock<'a>>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<WireTool<'a>>>,
}

impl<'a> MessagesRequest<'a> {
    pub fn build(
        model: String,
        max_tokens: u32,
        history: &'a [Message],
        tools: Option<&'a [ToolDefinition]>,
        system: Option<&'a str>,
        stream: bool,
    ) -> Self {
        Self {
            model,
            max_tokens,
            messages: to_wire_messages(history),
            system: system.map(to_wire_system),
            stream: Some(stream),
            tools: tools.map(to_wire_tools),
        }
    }
}

// Anthropic's prompt cache is opt-in: without an explicit breakpoint the whole prefix is billed
// at full rate on every request, however stable it is. Three markers cover an agent loop — the
// tool definitions and system prompt never change, and the newest user message is the boundary a
// single turn re-reads on each tool round trip. The limit is four.
#[derive(Debug, Serialize)]
struct CacheControl {
    #[serde(rename = "type")]
    kind: &'static str,
}

impl CacheControl {
    fn ephemeral() -> Self {
        Self { kind: "ephemeral" }
    }
}

#[derive(Debug, Serialize)]
struct SystemBlock<'a> {
    #[serde(rename = "type")]
    kind: &'static str,
    text: &'a str,
    #[serde(skip_serializing_if = "Option::is_none")]
    cache_control: Option<CacheControl>,
}

// The wire shapes mirror `Message`/`ContentPart`/`ToolDefinition` with a cache marker added. They
// exist so the breakpoints never reach the durable types: `ContentPart` is what the session log
// stores, and a resumed log must deserialize the same shape it was written with.
#[derive(Debug, Serialize)]
struct WireMessage<'a> {
    role: &'a str,
    content: Vec<WirePart<'a>>,
}

#[derive(Debug, Serialize)]
#[serde(tag = "type")]
enum WirePart<'a> {
    #[serde(rename = "text")]
    Text {
        text: &'a str,
        #[serde(skip_serializing_if = "Option::is_none")]
        cache_control: Option<CacheControl>,
    },
    #[serde(rename = "tool_use")]
    ToolUse {
        id: &'a str,
        name: &'a str,
        input: &'a serde_json::Value,
    },
    #[serde(rename = "tool_result")]
    ToolResult {
        tool_use_id: &'a str,
        content: &'a str,
        #[serde(skip_serializing_if = "Option::is_none")]
        cache_control: Option<CacheControl>,
    },
}

#[derive(Debug, Serialize)]
struct WireTool<'a> {
    name: &'a str,
    description: &'a str,
    input_schema: &'a serde_json::Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    cache_control: Option<CacheControl>,
}

fn to_wire_system(system: &str) -> Vec<SystemBlock<'_>> {
    vec![SystemBlock {
        kind: "text",
        text: system,
        cache_control: Some(CacheControl::ephemeral()),
    }]
}

// One marker on the last definition covers the whole block: a breakpoint caches everything
// before it, and tool definitions precede the conversation.
fn to_wire_tools(tools: &[ToolDefinition]) -> Vec<WireTool<'_>> {
    let last = tools.len().saturating_sub(1);
    tools
        .iter()
        .enumerate()
        .map(|(i, tool)| WireTool {
            name: &tool.name,
            description: &tool.description,
            input_schema: &tool.input_schema,
            cache_control: (i == last).then(CacheControl::ephemeral),
        })
        .collect()
}

fn to_wire_part(part: &ContentPart, cache: bool) -> WirePart<'_> {
    match part {
        ContentPart::Text { text } => WirePart::Text {
            text,
            cache_control: cache.then(CacheControl::ephemeral),
        },
        // A marker on a tool_use block would cache a prefix that is about to be extended by its
        // own result, so the call blocks are never the breakpoint.
        ContentPart::ToolUse { id, name, input } => WirePart::ToolUse { id, name, input },
        ContentPart::ToolResult {
            tool_use_id,
            content,
        } => WirePart::ToolResult {
            tool_use_id,
            content,
            cache_control: cache.then(CacheControl::ephemeral),
        },
    }
}

/// Marks the last content part of the newest user message.
///
/// That position is what makes an agent loop cheap: one user turn expands into many
/// assistant/tool round trips, and each of those re-sends the whole history. A breakpoint here
/// means every request after the first in a turn reads the prefix from cache.
fn to_wire_messages(history: &[Message]) -> Vec<WireMessage<'_>> {
    let breakpoint = history.iter().rposition(|m| m.role == Role::User);

    history
        .iter()
        .enumerate()
        .map(|(i, message)| {
            let last_part = message.content.len().saturating_sub(1);
            WireMessage {
                role: match message.role {
                    Role::User => "user",
                    Role::Assistant => "assistant",
                },
                content: message
                    .content
                    .iter()
                    .enumerate()
                    .map(|(j, part)| to_wire_part(part, breakpoint == Some(i) && j == last_part))
                    .collect(),
            }
        })
        .collect()
}

// Used by the non-streaming send_message path, kept for callers that do not want deltas.
#[allow(dead_code)]
#[derive(Debug, Deserialize)]
pub struct MessagesResponse {
    pub content: Vec<ContentPart>,
    pub stop_reason: Option<String>,
}

// `index` is part of the wire format and must be declared to deserialize, but block
// reassembly is positional today rather than index-keyed.
#[allow(dead_code)]
#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
pub enum StreamEvent {
    #[serde(rename = "message_start")]
    MessageStart { message: MessageStartData },
    #[serde(rename = "content_block_start")]
    ContentBlockStart {
        index: usize,
        content_block: ContentPart,
    },
    #[serde(rename = "content_block_delta")]
    ContentBlockDelta { index: usize, delta: Delta },
    #[serde(rename = "content_block_stop")]
    ContentBlockStop { index: usize },
    #[serde(rename = "message_delta")]
    MessageDelta {
        delta: MessageDeltaData,
        usage: Option<UsageData>,
    },
    #[serde(rename = "message_stop")]
    MessageStop {},
    #[serde(rename = "ping")]
    Ping {},
    #[serde(rename = "error")]
    Error { error: ErrorData },
}

#[derive(Debug, Deserialize)]
#[serde(tag = "type")]
pub enum Delta {
    #[serde(rename = "text_delta")]
    TextDelta { text: String },
    #[serde(rename = "input_json_delta")]
    InputJsonDelta { partial_json: String },
}

#[derive(Debug, Deserialize)]
pub struct MessageDeltaData {
    pub stop_reason: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct MessageStartData {
    pub usage: Option<UsageData>,
}

#[derive(Debug, Default, Deserialize)]
pub struct UsageData {
    pub input_tokens: Option<usize>,
    pub cache_creation_input_tokens: Option<usize>,
    pub cache_read_input_tokens: Option<usize>,
    pub output_tokens: Option<usize>,
}

#[derive(Debug, Deserialize)]
pub struct ErrorData {
    pub message: String,
}

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

    fn request(system: Option<&str>) -> serde_json::Value {
        let history = vec![Message::user("oi")];
        let req = MessagesRequest::build("m".to_string(), 4096, &history, None, system, true);
        serde_json::to_value(&req).unwrap()
    }

    #[test]
    fn explain_off_omits_the_system_field_entirely() {
        let body = request(None);
        assert!(
            body.get("system").is_none(),
            "an absent system prompt must not be sent as null: {}",
            body
        );
    }

    #[test]
    fn explain_on_sends_the_system_prompt() {
        let body = request(Some(EXPLAIN_SYSTEM_PROMPT));
        assert_eq!(body["system"][0]["text"], EXPLAIN_SYSTEM_PROMPT);
    }

    // The prefix is only billed once if a breakpoint says where it ends; without these markers
    // every request pays full rate for a system prompt and tool block that never change.
    #[test]
    fn the_system_prompt_carries_a_cache_breakpoint() {
        let body = request(Some(EXPLAIN_SYSTEM_PROMPT));
        assert_eq!(body["system"][0]["cache_control"]["type"], "ephemeral");
    }

    #[test]
    fn only_the_last_tool_definition_carries_a_breakpoint() {
        let tools: Vec<ToolDefinition> = ["a", "b", "c"]
            .iter()
            .map(|name| ToolDefinition {
                name: name.to_string(),
                description: "t".to_string(),
                input_schema: serde_json::json!({}),
            })
            .collect();

        let wire = serde_json::to_value(to_wire_tools(&tools)).unwrap();

        assert!(wire[0].get("cache_control").is_none());
        assert!(wire[1].get("cache_control").is_none());
        assert_eq!(wire[2]["cache_control"]["type"], "ephemeral");
    }

    // A turn expands into many assistant/tool round trips that each re-send the whole history,
    // so this breakpoint is what makes the second request onward cheap.
    #[test]
    fn the_newest_user_message_carries_a_breakpoint() {
        let history = vec![
            Message::user("first"),
            Message::assistant(vec![ContentPart::Text {
                text: "reply".to_string(),
            }]),
            Message::user("second"),
        ];

        let wire = serde_json::to_value(to_wire_messages(&history)).unwrap();

        assert!(wire[0]["content"][0].get("cache_control").is_none());
        assert!(wire[1]["content"][0].get("cache_control").is_none());
        assert_eq!(
            wire[2]["content"][0]["cache_control"]["type"], "ephemeral",
            "the newest user message must be the breakpoint: {}",
            wire
        );
    }

    // Anthropic allows four breakpoints; tools, system and the newest user message use three, so
    // a history full of tool traffic must not add more.
    #[test]
    fn a_tool_heavy_history_stays_within_one_message_breakpoint() {
        let history = vec![
            Message::user("go"),
            Message::assistant(vec![ContentPart::ToolUse {
                id: "1".to_string(),
                name: "read".to_string(),
                input: serde_json::json!({}),
            }]),
            Message::tool_results(vec![("1".to_string(), "ok".to_string())]),
        ];

        let wire = serde_json::to_value(to_wire_messages(&history)).unwrap();
        let marked = wire
            .as_array()
            .unwrap()
            .iter()
            .flat_map(|m| m["content"].as_array().unwrap())
            .filter(|part| part.get("cache_control").is_some())
            .count();

        assert_eq!(marked, 1, "expected exactly one breakpoint: {}", wire);
    }

    // A tool_result message is a user turn, so it becomes the newest breakpoint as a turn
    // progresses — which is what keeps the growing history cached mid-turn.
    #[test]
    fn a_tool_result_can_be_the_breakpoint() {
        let history = vec![
            Message::user("go"),
            Message::assistant(vec![ContentPart::ToolUse {
                id: "1".to_string(),
                name: "read".to_string(),
                input: serde_json::json!({}),
            }]),
            Message::tool_results(vec![("1".to_string(), "ok".to_string())]),
        ];

        let wire = serde_json::to_value(to_wire_messages(&history)).unwrap();

        assert_eq!(wire[2]["content"][0]["cache_control"]["type"], "ephemeral");
        assert_eq!(wire[2]["content"][0]["type"], "tool_result");
    }

    // A marker here would cache a prefix that its own result is about to extend, spending a
    // breakpoint on a position that can never be read back.
    #[test]
    fn tool_use_blocks_never_carry_a_breakpoint() {
        let part = ContentPart::ToolUse {
            id: "1".to_string(),
            name: "read".to_string(),
            input: serde_json::json!({}),
        };

        let wire = serde_json::to_value(to_wire_part(&part, true)).unwrap();

        assert!(wire.get("cache_control").is_none());
    }

    fn parse_event(json: &str) -> StreamEvent {
        serde_json::from_str(json).unwrap_or_else(|e| panic!("failed on {}: {}", json, e))
    }

    #[test]
    fn message_start_usage_is_parsed() {
        let event = parse_event(
            r#"{"type":"message_start","message":{"id":"m","usage":{"input_tokens":1200,"cache_creation_input_tokens":30,"cache_read_input_tokens":400,"output_tokens":2}}}"#,
        );
        match event {
            StreamEvent::MessageStart { message } => {
                let usage = message.usage.expect("usage present");
                assert_eq!(usage.input_tokens, Some(1200));
                assert_eq!(usage.cache_read_input_tokens, Some(400));
                assert_eq!(usage.cache_creation_input_tokens, Some(30));
            }
            other => panic!("wrong variant: {:?}", other),
        }
    }

    #[test]
    fn message_delta_usage_is_parsed_from_the_event_not_the_delta() {
        let event = parse_event(
            r#"{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":915}}"#,
        );
        match event {
            StreamEvent::MessageDelta { delta, usage } => {
                assert_eq!(delta.stop_reason.as_deref(), Some("end_turn"));
                assert_eq!(usage.unwrap().output_tokens, Some(915));
            }
            other => panic!("wrong variant: {:?}", other),
        }
    }

    #[test]
    fn a_message_start_without_usage_still_parses() {
        let event = parse_event(r#"{"type":"message_start","message":{"id":"m"}}"#);
        assert!(matches!(event, StreamEvent::MessageStart { .. }));
    }
}