talos-core 0.2.0

//! Core message types and event protocol.

use serde::{Deserialize, Serialize};

use crate::tool::ToolProvenance;

/// Provider-side caching behavior for a system prompt range.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum SystemCacheType {
    /// Cache this prompt range ephemerally when the provider supports it.
    Ephemeral,
}

/// A byte range in the system prompt that is stable enough for provider caching.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SystemCacheMarker {
    /// Starting byte offset in the system prompt content.
    pub offset: usize,
    /// Length of the cacheable range in bytes.
    pub length: usize,
    /// Cache behavior requested for this range.
    pub cache_type: SystemCacheType,
}

/// A tool call requested by the assistant.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolCall {
    /// Unique identifier for this tool call.
    pub id: String,
    /// Name of the tool to invoke.
    pub name: String,
    /// JSON-encoded arguments for the tool.
    pub input: serde_json::Value,
}

/// Result of a tool execution (message-layer).
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MessageToolResult {
    /// ID of the tool call this result corresponds to.
    pub tool_use_id: String,
    /// Text output from the tool.
    pub content: String,
    /// Whether the tool execution failed.
    pub is_error: bool,
}

/// A message in the conversation.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "role", rename_all = "snake_case")]
pub enum Message {
    /// System-level instruction (identity, rules, tool guide).
    System {
        /// System prompt content.
        content: String,
        /// Stable prompt ranges suitable for provider-side caching.
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        cache_markers: Vec<SystemCacheMarker>,
    },
    /// Workspace context (AGENTS.md, history summary, retrieved files).
    Context {
        /// Context content.
        content: String,
    },
    /// Message from the user.
    User {
        /// The user's message text.
        content: String,
    },
    /// Response from the assistant.
    Assistant {
        /// The assistant's response text.
        content: String,
        /// Tool calls requested by the assistant.
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        tool_calls: Vec<ToolCall>,
    },
    /// Result of a tool execution.
    Tool {
        /// The tool result.
        result: MessageToolResult,
    },
}

/// Reason the assistant stopped generating.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
    /// Assistant finished its response.
    EndTurn,
    /// Assistant wants to call a tool.
    ToolUse,
    /// Reached the maximum token limit.
    MaxTokens,
}

/// Token usage statistics for a turn.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct Usage {
    /// Tokens in the input prompt.
    pub input_tokens: u32,
    /// Tokens generated by the model.
    pub output_tokens: u32,
    /// Tokens read from cache.
    #[serde(default)]
    pub cache_read_tokens: u32,
    /// Tokens written to cache.
    #[serde(default)]
    pub cache_write_tokens: u32,
}

/// Events emitted during a turn for streaming.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
#[non_exhaustive]
pub enum AgentEvent {
    /// Turn has started.
    TurnStart,
    /// A text delta was received from the provider.
    TextDelta {
        /// The text chunk.
        delta: String,
    },
    /// Tool call detected: parameters still streaming.
    ToolCallStarted {
        /// Name of the tool being called.
        name: String,
    },
    /// A tool call was requested.
    ToolCall {
        /// The tool call details.
        call: ToolCall,
        /// The provenance of the tool being called.
        provenance: ToolProvenance,
        /// Fields to display in the TUI summary (from tool summary_fields()).
        summary_fields: Vec<String>,
    },
    /// A tool call completed.
    ToolResult {
        /// The tool result.
        result: MessageToolResult,
    },
    /// Turn has ended.
    TurnEnd {
        /// Why the turn ended.
        stop_reason: StopReason,
        /// Token usage for this turn.
        usage: Usage,
    },
    /// An error occurred.
    Error {
        /// Error message.
        message: String,
    },
}

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

    #[test]
    fn message_roundtrip_user() {
        let msg = Message::User {
            content: "Hello, world!".into(),
        };
        let json = serde_json::to_string(&msg).unwrap();
        let decoded: Message = serde_json::from_str(&json).unwrap();
        assert_eq!(msg, decoded);
    }

    #[test]
    fn message_roundtrip_assistant() {
        let msg = Message::Assistant {
            content: "I can help with that.".into(),
            tool_calls: vec![ToolCall {
                id: "call_1".into(),
                name: "read_file".into(),
                input: serde_json::json!({"path": "src/main.rs"}),
            }],
        };
        let json = serde_json::to_string(&msg).unwrap();
        let decoded: Message = serde_json::from_str(&json).unwrap();
        assert_eq!(msg, decoded);
    }

    #[test]
    fn message_roundtrip_tool() {
        let msg = Message::Tool {
            result: MessageToolResult {
                tool_use_id: "call_1".into(),
                content: "fn main() {}".into(),
                is_error: false,
            },
        };
        let json = serde_json::to_string(&msg).unwrap();
        let decoded: Message = serde_json::from_str(&json).unwrap();
        assert_eq!(msg, decoded);
    }

    #[test]
    fn event_roundtrip() {
        let events = vec![
            AgentEvent::TurnStart,
            AgentEvent::TextDelta {
                delta: "Hello".into(),
            },
            AgentEvent::ToolCall {
                call: ToolCall {
                    id: "c1".into(),
                    name: "bash".into(),
                    input: serde_json::json!({"command": "ls"}),
                },
                provenance: ToolProvenance::Native,
                summary_fields: vec![],
            },
            AgentEvent::ToolResult {
                result: MessageToolResult {
                    tool_use_id: "c1".into(),
                    content: "file.rs".into(),
                    is_error: false,
                },
            },
            AgentEvent::TurnEnd {
                stop_reason: StopReason::EndTurn,
                usage: Usage {
                    input_tokens: 100,
                    output_tokens: 50,
                    cache_read_tokens: 80,
                    cache_write_tokens: 20,
                },
            },
            AgentEvent::Error {
                message: "something failed".into(),
            },
        ];
        for event in events {
            let json = serde_json::to_string(&event).unwrap();
            let decoded: AgentEvent = serde_json::from_str(&json).unwrap();
            assert_eq!(event, decoded);
        }
    }
}

pub fn extract_tool_calls_from_text(text: &str) -> Vec<ToolCall> {
    let mut calls = Vec::new();
    let mut remaining = text;

    while let Some(start) = remaining.find("```json-tool") {
        let inner_start = start + "```json-tool".len();
        let inner = remaining[inner_start..].trim_start();
        let end = inner.find("```").unwrap_or(inner.len());
        let content = inner[..end].trim();

        if let Ok(obj) = serde_json::from_str::<serde_json::Value>(content)
            && let (Some(name), Some(args)) = (obj["name"].as_str(), Some(obj["args"].clone()))
        {
            calls.push(ToolCall {
                id: format!("tc_{}", calls.len()),
                name: name.to_string(),
                input: args,
            });
        }

        remaining = &inner[end..];
        if end + 3 < remaining.len() {
            remaining = &remaining[3..];
        } else {
            break;
        }
    }

    calls
}

pub fn strip_tool_syntax(text: &str) -> String {
    let mut result = text.to_string();
    while let Some(start) = result.find("```json-tool") {
        let inner_start = start + "```json-tool".len();
        let inner = &result[inner_start..];
        let end = inner_start + inner.find("```").unwrap_or(inner.len()) + 3;
        result.replace_range(start..end, "");
    }
    result.trim().to_string()
}