generated from nhcarrigan/template
naomi
136f95cd1a
Also includes cached tokens in cost calculations to provide more accurate billing estimates. Reviewed-on: #110 Co-authored-by: Naomi Carrigan <commits@nhcarrigan.com> Co-committed-by: Naomi Carrigan <commits@nhcarrigan.com>
399 lines
11 KiB
Rust
399 lines
11 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct UsageInfo {
|
|
pub input_tokens: u64,
|
|
pub output_tokens: u64,
|
|
#[serde(default)]
|
|
pub cache_creation_input_tokens: Option<u64>,
|
|
#[serde(default)]
|
|
pub cache_read_input_tokens: Option<u64>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum CharacterState {
|
|
#[default]
|
|
Idle,
|
|
Thinking,
|
|
Typing,
|
|
Searching,
|
|
Coding,
|
|
Mcp,
|
|
Permission,
|
|
Success,
|
|
Error,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum ConnectionStatus {
|
|
#[default]
|
|
Disconnected,
|
|
Connecting,
|
|
Connected,
|
|
Error,
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct TerminalLine {
|
|
pub id: String,
|
|
#[serde(rename = "type")]
|
|
pub line_type: String,
|
|
pub content: String,
|
|
pub timestamp: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub tool_name: Option<String>,
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PermissionRequest {
|
|
pub id: String,
|
|
pub tool: String,
|
|
pub description: String,
|
|
pub input: serde_json::Value,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PermissionDenial {
|
|
pub tool_name: String,
|
|
pub tool_use_id: String,
|
|
pub tool_input: serde_json::Value,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(tag = "type")]
|
|
pub enum ClaudeMessage {
|
|
#[serde(rename = "system")]
|
|
System {
|
|
subtype: String,
|
|
#[serde(default)]
|
|
session_id: Option<String>,
|
|
#[serde(default)]
|
|
cwd: Option<String>,
|
|
#[serde(default)]
|
|
tools: Option<Vec<String>>,
|
|
},
|
|
#[serde(rename = "assistant")]
|
|
Assistant {
|
|
message: AssistantMessageContent,
|
|
#[serde(default)]
|
|
parent_tool_use_id: Option<String>,
|
|
},
|
|
#[serde(rename = "user")]
|
|
User { message: UserMessageContent },
|
|
#[serde(rename = "stream_event")]
|
|
StreamEvent { event: StreamEventData },
|
|
#[serde(rename = "result")]
|
|
Result {
|
|
subtype: String,
|
|
#[serde(default)]
|
|
result: Option<String>,
|
|
#[serde(default)]
|
|
duration_ms: Option<u64>,
|
|
#[serde(default)]
|
|
num_turns: Option<u32>,
|
|
#[serde(default)]
|
|
permission_denials: Option<Vec<PermissionDenial>>,
|
|
#[serde(default)]
|
|
usage: Option<UsageInfo>,
|
|
},
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct AssistantMessageContent {
|
|
pub content: Vec<ContentBlock>,
|
|
#[serde(default)]
|
|
pub model: Option<String>,
|
|
#[serde(default)]
|
|
pub stop_reason: Option<String>,
|
|
#[serde(default)]
|
|
pub usage: Option<UsageInfo>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct UserMessageContent {
|
|
pub content: Vec<ContentBlock>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
#[serde(tag = "type")]
|
|
pub enum ContentBlock {
|
|
#[serde(rename = "text")]
|
|
Text { text: String },
|
|
#[serde(rename = "thinking")]
|
|
Thinking { thinking: String },
|
|
#[serde(rename = "tool_use")]
|
|
ToolUse {
|
|
id: String,
|
|
name: String,
|
|
input: serde_json::Value,
|
|
},
|
|
#[serde(rename = "tool_result")]
|
|
ToolResult {
|
|
tool_use_id: String,
|
|
content: serde_json::Value,
|
|
#[serde(default)]
|
|
is_error: Option<bool>,
|
|
},
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct StreamEventData {
|
|
#[serde(rename = "type")]
|
|
pub event_type: String,
|
|
#[serde(default)]
|
|
pub index: Option<u32>,
|
|
#[serde(default)]
|
|
pub content_block: Option<ContentBlockStart>,
|
|
#[serde(default)]
|
|
pub delta: Option<DeltaContent>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ContentBlockStart {
|
|
#[serde(rename = "type")]
|
|
pub block_type: String,
|
|
#[serde(default)]
|
|
pub id: Option<String>,
|
|
#[serde(default)]
|
|
pub name: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct DeltaContent {
|
|
#[serde(rename = "type")]
|
|
pub delta_type: String,
|
|
#[serde(default)]
|
|
pub text: Option<String>,
|
|
#[serde(default)]
|
|
pub thinking: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct StateChangeEvent {
|
|
pub state: CharacterState,
|
|
pub tool_name: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub conversation_id: Option<String>,
|
|
}
|
|
|
|
/// Cost information for a message
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct MessageCost {
|
|
pub input_tokens: u64,
|
|
pub output_tokens: u64,
|
|
pub cost_usd: f64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct OutputEvent {
|
|
pub line_type: String,
|
|
pub content: String,
|
|
pub tool_name: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub conversation_id: Option<String>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub cost: Option<MessageCost>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct PermissionPromptEvent {
|
|
pub id: String,
|
|
pub tool_name: String,
|
|
pub tool_input: serde_json::Value,
|
|
pub description: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub conversation_id: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ConnectionEvent {
|
|
pub status: ConnectionStatus,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub conversation_id: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SessionEvent {
|
|
pub session_id: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub conversation_id: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct WorkingDirectoryEvent {
|
|
pub directory: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub conversation_id: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct QuestionOption {
|
|
pub label: String,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub description: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct UserQuestionEvent {
|
|
pub id: String,
|
|
pub question: String,
|
|
pub header: Option<String>,
|
|
pub options: Vec<QuestionOption>,
|
|
pub multi_select: bool,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub conversation_id: Option<String>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_character_state_default() {
|
|
let state = CharacterState::default();
|
|
assert_eq!(state, CharacterState::Idle);
|
|
}
|
|
|
|
#[test]
|
|
fn test_connection_status_default() {
|
|
let status = ConnectionStatus::default();
|
|
matches!(status, ConnectionStatus::Disconnected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_character_state_serialization() {
|
|
let state = CharacterState::Thinking;
|
|
let serialized = serde_json::to_string(&state).unwrap();
|
|
assert_eq!(serialized, "\"thinking\"");
|
|
|
|
let deserialized: CharacterState = serde_json::from_str(&serialized).unwrap();
|
|
assert_eq!(deserialized, CharacterState::Thinking);
|
|
}
|
|
|
|
#[test]
|
|
fn test_all_character_states_serialize() {
|
|
let states = vec![
|
|
(CharacterState::Idle, "\"idle\""),
|
|
(CharacterState::Thinking, "\"thinking\""),
|
|
(CharacterState::Typing, "\"typing\""),
|
|
(CharacterState::Searching, "\"searching\""),
|
|
(CharacterState::Coding, "\"coding\""),
|
|
(CharacterState::Mcp, "\"mcp\""),
|
|
(CharacterState::Permission, "\"permission\""),
|
|
(CharacterState::Success, "\"success\""),
|
|
(CharacterState::Error, "\"error\""),
|
|
];
|
|
|
|
for (state, expected) in states {
|
|
let serialized = serde_json::to_string(&state).unwrap();
|
|
assert_eq!(serialized, expected, "Failed for state: {:?}", state);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_terminal_line_serialization() {
|
|
let line = TerminalLine {
|
|
id: "test-123".to_string(),
|
|
line_type: "assistant".to_string(),
|
|
content: "Hello, world!".to_string(),
|
|
timestamp: "2024-01-01T00:00:00Z".to_string(),
|
|
tool_name: None,
|
|
};
|
|
|
|
let serialized = serde_json::to_string(&line).unwrap();
|
|
assert!(serialized.contains("\"type\":\"assistant\""));
|
|
assert!(serialized.contains("\"content\":\"Hello, world!\""));
|
|
assert!(!serialized.contains("tool_name"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_terminal_line_with_tool_name() {
|
|
let line = TerminalLine {
|
|
id: "test-456".to_string(),
|
|
line_type: "tool".to_string(),
|
|
content: "Reading file...".to_string(),
|
|
timestamp: "2024-01-01T00:00:00Z".to_string(),
|
|
tool_name: Some("Read".to_string()),
|
|
};
|
|
|
|
let serialized = serde_json::to_string(&line).unwrap();
|
|
assert!(serialized.contains("\"tool_name\":\"Read\""));
|
|
}
|
|
|
|
#[test]
|
|
fn test_content_block_text() {
|
|
let block = ContentBlock::Text {
|
|
text: "Hello!".to_string(),
|
|
};
|
|
|
|
let serialized = serde_json::to_string(&block).unwrap();
|
|
assert!(serialized.contains("\"type\":\"text\""));
|
|
assert!(serialized.contains("\"text\":\"Hello!\""));
|
|
}
|
|
|
|
#[test]
|
|
fn test_content_block_tool_use() {
|
|
let block = ContentBlock::ToolUse {
|
|
id: "tool-123".to_string(),
|
|
name: "Read".to_string(),
|
|
input: serde_json::json!({"file_path": "/test.txt"}),
|
|
};
|
|
|
|
let serialized = serde_json::to_string(&block).unwrap();
|
|
assert!(serialized.contains("\"type\":\"tool_use\""));
|
|
assert!(serialized.contains("\"name\":\"Read\""));
|
|
}
|
|
|
|
#[test]
|
|
fn test_state_change_event() {
|
|
let event = StateChangeEvent {
|
|
state: CharacterState::Coding,
|
|
tool_name: Some("Edit".to_string()),
|
|
conversation_id: None,
|
|
};
|
|
|
|
let serialized = serde_json::to_string(&event).unwrap();
|
|
assert!(serialized.contains("\"state\":\"coding\""));
|
|
assert!(serialized.contains("\"tool_name\":\"Edit\""));
|
|
}
|
|
|
|
#[test]
|
|
fn test_output_event() {
|
|
let event = OutputEvent {
|
|
line_type: "assistant".to_string(),
|
|
content: "Test output".to_string(),
|
|
tool_name: None,
|
|
conversation_id: None,
|
|
cost: None,
|
|
};
|
|
|
|
let serialized = serde_json::to_string(&event).unwrap();
|
|
assert!(serialized.contains("\"line_type\":\"assistant\""));
|
|
assert!(serialized.contains("\"content\":\"Test output\""));
|
|
}
|
|
|
|
#[test]
|
|
fn test_output_event_with_cost() {
|
|
let event = OutputEvent {
|
|
line_type: "assistant".to_string(),
|
|
content: "Test output".to_string(),
|
|
tool_name: None,
|
|
conversation_id: Some("conv-123".to_string()),
|
|
cost: Some(MessageCost {
|
|
input_tokens: 100,
|
|
output_tokens: 50,
|
|
cost_usd: 0.005,
|
|
}),
|
|
};
|
|
|
|
let serialized = serde_json::to_string(&event).unwrap();
|
|
assert!(serialized.contains("\"cost\":"));
|
|
assert!(serialized.contains("\"input_tokens\":100"));
|
|
assert!(serialized.contains("\"output_tokens\":50"));
|
|
}
|
|
}
|