ai-agent-sdk 0.5.0

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Message {
    pub role: MessageRole,
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub attachments: Option<Vec<Attachment>>,
    /// Tool call ID for tool role messages (required by OpenAI API)
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub tool_call_id: Option<String>,
    /// Tool calls for assistant messages (required to pair with tool results)
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub tool_calls: Option<Vec<ToolCall>>,
}

/// A tool call from the model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
    pub id: String,
    pub name: String,
    pub arguments: serde_json::Value,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum MessageRole {
    #[default]
    User,
    Assistant,
    #[serde(rename = "tool")]
    Tool,
}

/// Attachments for messages - files, images, etc.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum Attachment {
    /// File attachment (at-mentioned file)
    File {
        path: String,
    },
    /// Reference to a file previously read
    AlreadyReadFile {
        path: String,
        content: String,
    },
    /// PDF reference
    PdfReference {
        path: String,
    },
    /// Text file that was edited
    EditedTextFile {
        filename: String,
        snippet: String,
    },
    /// Image file that was edited
    EditedImageFile {
        filename: String,
    },
    /// Directory listing
    Directory {
        path: String,
        content: String,
        display_path: String,
    },
    /// Selected lines in IDE
    SelectedLinesInIde {
        ide_name: String,
        filename: String,
        start_line: u32,
        end_line: u32,
    },
    /// Memory file reference
    MemoryFile {
        path: String,
    },
    /// Skill listing attachment
    SkillListing {
        skills: Vec<SkillInfo>,
    },
    /// Invoked skills attachment
    InvokedSkills {
        skills: Vec<InvokedSkill>,
    },
    /// Task status
    TaskStatus {
        task_id: String,
        description: String,
        status: String,
    },
    /// Plan file reference
    PlanFileReference {
        path: String,
    },
    /// MCP tool resources
    McpResources {
        tools: Vec<String>,
    },
    /// Deferred tools delta
    DeferredTools {
        tools: Vec<String>,
    },
    /// Agent listing
    AgentListing {
        agents: Vec<String>,
    },
    /// Custom attachment
    Custom {
        name: String,
        content: serde_json::Value,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SkillInfo {
    pub name: String,
    pub description: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InvokedSkill {
    pub name: String,
    pub path: String,
    pub content: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cache_creation_input_tokens: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cache_read_input_tokens: Option<u64>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub input_schema: ToolInputSchema,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolInputSchema {
    #[serde(rename = "type")]
    pub schema_type: String,
    pub properties: serde_json::Value,
    pub required: Option<Vec<String>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ToolContext {
    pub cwd: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub abort_signal: Option<()>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolResult {
    #[serde(rename = "type")]
    pub result_type: String,
    pub tool_use_id: String,
    pub content: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub is_error: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentOptions {
    pub model: Option<String>,
    pub api_key: Option<String>,
    pub base_url: Option<String>,
    pub cwd: Option<String>,
    pub system_prompt: Option<String>,
    pub max_turns: Option<u32>,
    pub max_budget_usd: Option<f64>,
    pub max_tokens: Option<u32>,
    #[serde(default)]
    pub tools: Vec<ToolDefinition>,
    #[serde(default)]
    pub allowed_tools: Vec<String>,
    #[serde(default)]
    pub disallowed_tools: Vec<String>,
    /// MCP server configurations
    #[serde(default)]
    pub mcp_servers: Option<std::collections::HashMap<String, McpServerConfig>>,
}

impl Default for AgentOptions {
    fn default() -> Self {
        Self {
            model: None,
            api_key: None,
            base_url: None,
            cwd: None,
            system_prompt: None,
            max_turns: None,
            max_budget_usd: None,
            max_tokens: None,
            tools: Vec::new(),
            allowed_tools: Vec::new(),
            disallowed_tools: Vec::new(),
            mcp_servers: None,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueryResult {
    pub text: String,
    pub usage: TokenUsage,
    pub num_turns: u32,
    pub duration_ms: u64,
}

// --------------------------------------------------------------------------
// MCP Types
// --------------------------------------------------------------------------

/// MCP server configuration (union type)
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum McpServerConfig {
    Stdio(McpStdioConfig),
    Sse(McpSseConfig),
    Http(McpHttpConfig),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpStdioConfig {
    #[serde(default = "default_stdio_type")]
    pub transport_type: Option<String>,
    pub command: String,
    pub args: Option<Vec<String>>,
    pub env: Option<std::collections::HashMap<String, String>>,
}

fn default_stdio_type() -> Option<String> {
    Some("stdio".to_string())
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpSseConfig {
    pub transport_type: String,
    pub url: String,
    pub headers: Option<std::collections::HashMap<String, String>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct McpHttpConfig {
    pub transport_type: String,
    pub url: String,
    pub headers: Option<std::collections::HashMap<String, String>>,
}

/// MCP connection status
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum McpConnectionStatus {
    Connected,
    Disconnected,
    Error,
}

/// MCP tool representation from server
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpTool {
    pub name: String,
    pub description: Option<String>,
    #[serde(rename = "inputSchema")]
    pub input_schema: Option<serde_json::Value>,
}