PicoBot/src/protocol.rs

447 lines
15 KiB
Rust

use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionSummary {
pub session_id: String,
pub title: String,
pub channel_name: String,
pub chat_id: String,
pub message_count: i64,
pub last_active_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub archived_at: Option<i64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SlashCommandInfo {
pub name: String,
pub description: String,
pub aliases: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct UploadDescriptor {
pub upload_id: String,
pub name: String,
pub media_type: String,
pub mime_type: String,
pub size: u64,
pub expires_at: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct MessageAttachment {
pub index: u32,
pub name: String,
pub media_type: String,
pub mime_type: String,
}
impl MessageAttachment {
pub fn from_media_ref(index: usize, media_ref: &crate::bus::MediaRef) -> Self {
let name = std::path::Path::new(&media_ref.path)
.file_name()
.map(|name| name.to_string_lossy().into_owned())
.filter(|name| !name.is_empty())
.unwrap_or_else(|| "attachment".to_string());
let mime_type = mime_guess::from_path(&name)
.first_or_octet_stream()
.to_string();
Self {
index: u32::try_from(index).unwrap_or(u32::MAX),
name,
media_type: media_ref.media_type.clone(),
mime_type,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoryMessage {
pub id: String,
pub seq: i64,
pub role: String,
pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>,
#[serde(default)]
pub completion_status: crate::bus::CompletionStatus,
pub created_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<crate::providers::ToolCall>>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub attachments: Vec<MessageAttachment>,
}
impl From<crate::bus::CommittedMessage> for HistoryMessage {
fn from(message: crate::bus::CommittedMessage) -> Self {
let attachments = message
.media_refs
.iter()
.enumerate()
.map(|(index, media_ref)| MessageAttachment::from_media_ref(index, media_ref))
.collect();
Self {
id: message.id,
seq: message.seq,
role: message.role,
content: message.content,
reasoning_content: message.reasoning_content,
completion_status: message.completion_status,
created_at: message.created_at,
tool_call_id: message.tool_call_id,
tool_name: message.tool_name,
tool_calls: message.tool_calls,
attachments,
}
}
}
impl HistoryMessage {
pub fn from_message_meta(message: crate::storage::message::MessageMeta) -> Self {
let attachments = message
.media_refs
.as_deref()
.and_then(|refs| serde_json::from_str::<Vec<crate::bus::MediaRef>>(refs).ok())
.unwrap_or_default()
.iter()
.enumerate()
.map(|(index, media_ref)| MessageAttachment::from_media_ref(index, media_ref))
.collect();
Self {
id: message.id,
seq: message.seq,
role: message.role,
content: message.content,
reasoning_content: message.reasoning_content,
completion_status: message.completion_status,
created_at: message.created_at,
tool_call_id: message.tool_call_id,
tool_name: message.tool_name,
tool_calls: message
.tool_calls
.and_then(|calls| serde_json::from_str(&calls).ok()),
attachments,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum WsInbound {
#[serde(rename = "user_input")]
UserInput {
content: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
upload_ids: Vec<String>,
/// Stable id generated by the client for optimistic-message
/// reconciliation. It is optional for older clients and channels.
#[serde(default, skip_serializing_if = "Option::is_none")]
client_message_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
channel: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
chat_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
sender_id: Option<String>,
},
#[serde(rename = "clear_history")]
ClearHistory {
#[serde(default, skip_serializing_if = "Option::is_none")]
chat_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
#[serde(rename = "create_session")]
CreateSession {
#[serde(default, skip_serializing_if = "Option::is_none")]
title: Option<String>,
},
#[serde(rename = "list_sessions")]
ListSessions {
#[serde(default)]
include_archived: bool,
},
#[serde(rename = "load_session")]
LoadSession { session_id: String },
#[serde(rename = "get_session_history")]
GetSessionHistory {
session_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
limit: Option<u32>,
},
#[serde(rename = "get_session_plan")]
GetSessionPlan { session_id: String },
#[serde(rename = "get_session_stats")]
GetSessionStats { session_id: String },
#[serde(rename = "rename_session")]
RenameSession {
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
title: String,
},
#[serde(rename = "archive_session")]
ArchiveSession {
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
#[serde(rename = "delete_session")]
DeleteSession {
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
#[serde(rename = "get_slash_commands")]
GetSlashCommands,
#[serde(rename = "ping")]
Ping,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum WsOutbound {
#[serde(rename = "turn_updated")]
TurnUpdated {
snapshot: crate::session::TurnSnapshot,
},
#[serde(rename = "turn_committed")]
TurnCommitted {
session_id: String,
history_revision: i64,
messages: Vec<HistoryMessage>,
},
#[serde(rename = "assistant_response")]
AssistantResponse {
id: String,
content: String,
role: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
attachments: Vec<MessageAttachment>,
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
#[serde(rename = "error")]
Error { code: String, message: String },
#[serde(rename = "session_established")]
SessionEstablished {
session_id: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
capabilities: Vec<String>,
},
#[serde(rename = "session_created")]
SessionCreated { session_id: String, title: String },
#[serde(rename = "session_list")]
SessionList {
sessions: Vec<SessionSummary>,
#[serde(default, skip_serializing_if = "Option::is_none")]
current_session_id: Option<String>,
},
#[serde(rename = "session_loaded")]
SessionLoaded {
session_id: String,
title: String,
message_count: i64,
},
#[serde(rename = "session_history")]
SessionHistory {
session_id: String,
messages: Vec<HistoryMessage>,
},
#[serde(rename = "session_plan")]
SessionPlan {
session_id: String,
plan: Option<crate::work::TaskPlan>,
},
#[serde(rename = "session_stats")]
SessionStats { stats: crate::session::SessionStats },
#[serde(rename = "plan_updated")]
PlanUpdated {
session_id: String,
reason: String,
changed_item_ids: Vec<String>,
plan: Option<crate::work::TaskPlan>,
},
#[serde(rename = "session_renamed")]
SessionRenamed { session_id: String, title: String },
#[serde(rename = "session_archived")]
SessionArchived { session_id: String },
#[serde(rename = "session_deleted")]
SessionDeleted { session_id: String },
#[serde(rename = "history_cleared")]
HistoryCleared { session_id: String },
#[serde(rename = "slash_commands_list")]
SlashCommandsList { commands: Vec<SlashCommandInfo> },
#[serde(rename = "pong")]
Pong,
#[serde(rename = "command_executed")]
CommandExecuted { message: String },
#[serde(rename = "system_notification")]
SystemNotification {
content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
session_id: Option<String>,
},
}
pub fn parse_inbound(raw: &str) -> Result<WsInbound, serde_json::Error> {
serde_json::from_str(raw)
}
pub fn serialize_inbound(msg: &WsInbound) -> Result<String, serde_json::Error> {
serde_json::to_string(msg)
}
pub fn serialize_outbound(msg: &WsOutbound) -> Result<String, serde_json::Error> {
serde_json::to_string(msg)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::session::{TurnId, TurnPhase, TurnState, TurnStatus};
#[test]
fn turn_updated_serializes_as_one_complete_snapshot_frame() {
let frame = WsOutbound::TurnUpdated {
snapshot: TurnState {
id: TurnId("turn-1".into()),
session_id: "cli_chat:client:dialog".into(),
message_id: "message-1".into(),
revision: 7,
status: TurnStatus::Running,
phase: TurnPhase::Responding,
blocks: Vec::new(),
usage: None,
error: None,
},
};
let json = serialize_outbound(&frame).unwrap();
let value: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(value["type"], "turn_updated");
assert_eq!(value["snapshot"]["revision"], 7);
assert_eq!(value["snapshot"]["status"], "running");
}
#[test]
fn turn_committed_serializes_revision_and_durable_delta() {
let frame = WsOutbound::TurnCommitted {
session_id: "session".to_string(),
history_revision: 7,
messages: vec![HistoryMessage {
id: "message".to_string(),
seq: 7,
role: "assistant".to_string(),
content: "done".to_string(),
reasoning_content: None,
completion_status: crate::bus::CompletionStatus::Completed,
created_at: 1,
tool_call_id: None,
tool_name: None,
tool_calls: None,
attachments: Vec::new(),
}],
};
let value = serde_json::to_value(frame).unwrap();
assert_eq!(value["type"], "turn_committed");
assert_eq!(value["history_revision"], 7);
assert_eq!(value["messages"][0]["id"], "message");
}
#[test]
fn user_input_preserves_optional_client_message_id() {
let inbound = parse_inbound(
r#"{"type":"user_input","content":"hello","client_message_id":"client-1"}"#,
)
.unwrap();
match inbound {
WsInbound::UserInput {
client_message_id,
upload_ids,
..
} => {
assert_eq!(client_message_id.as_deref(), Some("client-1"));
assert!(upload_ids.is_empty());
}
other => panic!("unexpected frame: {other:?}"),
}
let serialized = serialize_inbound(&WsInbound::UserInput {
content: "hello".to_string(),
upload_ids: Vec::new(),
client_message_id: Some("client-1".to_string()),
channel: None,
chat_id: None,
sender_id: None,
})
.unwrap();
assert!(serialized.contains(r#""client_message_id":"client-1""#));
}
#[test]
fn history_defaults_new_reasoning_fields_for_old_frames() {
let message: HistoryMessage = serde_json::from_value(serde_json::json!({
"id": "message",
"seq": 1,
"role": "assistant",
"content": "answer",
"created_at": 1,
"attachments": []
}))
.unwrap();
assert_eq!(message.reasoning_content, None);
assert_eq!(
message.completion_status,
crate::bus::CompletionStatus::Completed
);
}
#[test]
fn session_stats_request_and_response_use_structured_frames() {
let inbound =
parse_inbound(r#"{"type":"get_session_stats","session_id":"cli_chat:client:dialog"}"#)
.unwrap();
assert!(matches!(inbound, WsInbound::GetSessionStats { .. }));
let stats = crate::session::SessionStats {
session_id: "cli_chat:client:dialog".into(),
title: "stats".into(),
provider: "provider".into(),
model: "model".into(),
user_message_count: 1,
history_message_count: 2,
lifetime_usage: crate::session::LifetimeUsage {
input_tokens: 100,
output_tokens: 20,
total_tokens: 120,
cached_input_tokens: Some(40),
request_count: 1,
turn_count: 1,
tracked_since: Some(1),
},
context: crate::session::ContextUsage {
configured_window_tokens: 128_000,
effective_window_tokens: 128_000,
used_tokens: 100,
remaining_tokens: 127_900,
compression_threshold_tokens: 89_600,
source: crate::session::ContextUsageSource::Hybrid,
last_observed_prompt_tokens: Some(90),
observed_at: Some(1),
},
created_at: 1,
last_active_at: 2,
updated_at: 3,
};
let value = serde_json::to_value(WsOutbound::SessionStats { stats }).unwrap();
assert_eq!(value["type"], "session_stats");
assert_eq!(value["stats"]["context"]["source"], "hybrid");
assert_eq!(value["stats"]["lifetime_usage"]["input_tokens"], 100);
}
}