use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
pub enum ClientRequest {
Chat {
content: String,
#[serde(default)]
context: Option<RequestContext>,
},
QuickAction {
action: QuickActionType,
content: String,
#[serde(default)]
context: Option<RequestContext>,
#[serde(default)]
instructions: Option<String>,
},
NewSession,
Status,
Memory {
operation: MemoryOperation,
},
LoadSession {
session_id: String,
},
ListSessions,
}
#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum QuickActionType {
Explain,
Fix,
GenerateTests,
Refactor,
Optimize,
Document,
Translate,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemoryOperation {
List,
Search { query: String },
Add { content: String, category: Option<String> },
Clear,
Stats,
}
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct RequestContext {
#[serde(default)]
pub workspace: Option<String>,
#[serde(default)]
pub file: Option<String>,
#[serde(default)]
pub language: Option<String>,
#[serde(default)]
pub selection: Option<Selection>,
#[serde(default)]
pub diagnostics: Option<Vec<Diagnostic>>,
#[serde(default)]
pub extra_files: Option<Vec<String>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct Selection {
pub start: Position,
pub end: Position,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct Position {
pub line: u32,
pub character: u32,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct Diagnostic {
pub severity: String,
pub message: String,
pub range: Selection,
#[serde(default)]
pub source: Option<String>,
#[serde(default)]
pub code: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
#[serde(rename_all = "snake_case")]
pub enum StreamEvent {
Text { content: String },
Thinking { content: String },
ToolUse {
id: String,
name: String,
input: serde_json::Value,
},
ToolResult {
tool_use_id: String,
content: String,
#[serde(default)]
success: bool,
},
WebSearchResult {
tool_use_id: String,
content: String,
},
Error {
message: String,
#[serde(default)]
code: Option<String>,
},
Done {
#[serde(default)]
usage: Option<Usage>,
},
SessionStarted {
session_id: String,
#[serde(default)]
memory_count: Option<usize>,
},
StatusResponse {
session_id: Option<String>,
message_count: usize,
total_tokens: u64,
is_streaming: bool,
},
MemoryList {
memories: Vec<MemoryEntry>,
},
MemoryStats {
total: usize,
by_category: HashMap<String, usize>,
},
SessionList {
sessions: Vec<SessionInfo>,
},
MemoryAdded {
category: String,
content: String,
},
Log {
level: String,
message: String,
},
}
#[derive(Debug, Serialize, Deserialize)]
pub struct Usage {
pub input: u64,
pub output: u64,
#[serde(default)]
pub cache_read: Option<u64>,
#[serde(default)]
pub cache_write: Option<u64>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct MemoryEntry {
pub id: String,
pub category: String,
pub content: String,
pub created_at: String,
#[serde(default)]
pub project: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SessionInfo {
pub id: String,
#[serde(default)]
pub name: Option<String>,
pub created_at: String,
pub message_count: usize,
#[serde(default)]
pub last_used: Option<String>,
}
impl StreamEvent {
pub fn text(content: impl Into<String>) -> Self {
StreamEvent::Text { content: content.into() }
}
pub fn thinking(content: impl Into<String>) -> Self {
StreamEvent::Thinking { content: content.into() }
}
pub fn tool_use(id: impl Into<String>, name: impl Into<String>, input: serde_json::Value) -> Self {
StreamEvent::ToolUse {
id: id.into(),
name: name.into(),
input,
}
}
pub fn tool_result(tool_use_id: impl Into<String>, content: impl Into<String>, success: bool) -> Self {
StreamEvent::ToolResult {
tool_use_id: tool_use_id.into(),
content: content.into(),
success,
}
}
pub fn error(message: impl Into<String>) -> Self {
StreamEvent::Error { message: message.into(), code: None }
}
pub fn done(usage: Option<Usage>) -> Self {
StreamEvent::Done { usage }
}
pub fn session_started(session_id: impl Into<String>, memory_count: Option<usize>) -> Self {
StreamEvent::SessionStarted {
session_id: session_id.into(),
memory_count,
}
}
pub fn to_json_line(&self) -> String {
serde_json::to_string(self).unwrap_or_default() + "\n"
}
}
impl Usage {
pub fn new(input: u64, output: u64) -> Self {
Usage { input, output, cache_read: None, cache_write: None }
}
pub fn with_cache(input: u64, output: u64, cache_read: u64, cache_write: u64) -> Self {
Usage {
input,
output,
cache_read: Some(cache_read),
cache_write: Some(cache_write),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_serialize_chat_request() {
let request = ClientRequest::Chat {
content: "Hello".to_string(),
context: None,
};
let json = serde_json::to_string(&request).unwrap();
assert!(json.contains("\"type\":\"chat\""));
assert!(json.contains("\"content\":\"Hello\""));
}
#[test]
fn test_deserialize_chat_request() {
let json = "{\"type\":\"chat\",\"content\":\"Hello\",\"context\":null}";
let request: ClientRequest = serde_json::from_str(json).unwrap();
match request {
ClientRequest::Chat { content, .. } => {
assert_eq!(content, "Hello");
}
_ => panic!("Expected Chat request"),
}
}
#[test]
fn test_serialize_stream_event() {
let event = StreamEvent::text("Hello world");
let json = event.to_json_line();
assert!(json.contains("\"type\":\"text\""));
assert!(json.contains("\"content\":\"Hello world\""));
assert!(json.ends_with("\n"));
}
#[test]
fn test_serialize_tool_use() {
let event = StreamEvent::tool_use("tool_1", "read", serde_json::json!({"path": "src/main.rs"}));
let json = event.to_json_line();
assert!(json.contains("\"type\":\"tool_use\""));
assert!(json.contains("\"id\":\"tool_1\""));
assert!(json.contains("\"name\":\"read\""));
}
#[test]
fn test_request_context_with_file() {
let json = "{\"workspace\":\"/project\",\"file\":\"src/main.rs\",\"language\":\"rust\"}";
let context: RequestContext = serde_json::from_str(json).unwrap();
assert_eq!(context.workspace, Some("/project".to_string()));
assert_eq!(context.file, Some("src/main.rs".to_string()));
assert_eq!(context.language, Some("rust".to_string()));
}
#[test]
fn test_quick_action_request() {
let json = "{\"type\":\"quick_action\",\"action\":\"explain\",\"content\":\"fn main(){}\",\"context\":{\"language\":\"rust\"}}";
let request: ClientRequest = serde_json::from_str(json).unwrap();
match request {
ClientRequest::QuickAction { action, content, .. } => {
assert_eq!(action, QuickActionType::Explain);
assert_eq!(content, "fn main(){}");
}
_ => panic!("Expected QuickAction request"),
}
}
}