use crate::domain::ActionDisplay;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "provider", rename_all = "snake_case")]
pub enum ProviderContinuation {
Anthropic { signature: String },
MetaResponses { output: Vec<MetaResponseItem> },
}
impl ProviderContinuation {
pub fn anthropic_signature(&self) -> Option<&str> {
match self {
Self::Anthropic { signature } => Some(signature),
Self::MetaResponses { .. } => None,
}
}
pub fn meta_output(&self) -> Option<&[MetaResponseItem]> {
match self {
Self::MetaResponses { output } => Some(output),
Self::Anthropic { .. } => None,
}
}
pub fn retain_meta_function_calls(&mut self, mut keep: impl FnMut(&str) -> bool) {
if let Self::MetaResponses { output } = self {
output.retain(|item| item.function_call_id().is_none_or(&mut keep));
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum MetaResponseItem {
Reasoning {
item: serde_json::Value,
#[serde(with = "crate::utils::serde_base64::string")]
encrypted_content: String,
},
Other {
item: serde_json::Value,
},
}
impl MetaResponseItem {
pub fn from_wire(mut item: serde_json::Value) -> Self {
let is_reasoning =
item.get("type").and_then(serde_json::Value::as_str) == Some("reasoning");
if is_reasoning
&& let Some(encrypted) = item
.as_object_mut()
.and_then(|object| object.remove("encrypted_content"))
.and_then(|value| value.as_str().map(str::to_string))
{
return Self::Reasoning {
item,
encrypted_content: encrypted,
};
}
Self::Other { item }
}
pub fn to_wire(&self) -> serde_json::Value {
match self {
Self::Reasoning {
item,
encrypted_content,
} => {
let mut item = item.clone();
if let Some(object) = item.as_object_mut() {
object.insert(
"encrypted_content".to_string(),
serde_json::Value::String(encrypted_content.clone()),
);
}
item
},
Self::Other { item } => item.clone(),
}
}
pub fn function_call_id(&self) -> Option<&str> {
let item = match self {
Self::Reasoning { item, .. } | Self::Other { item } => item,
};
(item.get("type").and_then(serde_json::Value::as_str) == Some("function_call"))
.then(|| item.get("call_id").and_then(serde_json::Value::as_str))
.flatten()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: MessageRole,
pub content: String,
pub timestamp: chrono::DateTime<chrono::Local>,
#[serde(default)]
pub kind: ChatMessageKind,
#[serde(default)]
pub metadata: Option<serde_json::Value>,
#[serde(default)]
pub actions: Vec<ActionDisplay>,
#[serde(default)]
pub thinking: Option<String>,
#[serde(default)]
pub images: Option<Vec<String>>,
#[serde(default)]
pub image_numbers: Option<Vec<u64>>,
#[serde(default)]
pub tool_calls: Option<Vec<crate::models::tool_call::ToolCall>>,
#[serde(default)]
pub tool_call_id: Option<String>,
#[serde(default)]
pub tool_name: Option<String>,
#[serde(default)]
pub provider_continuation: Option<ProviderContinuation>,
}
impl ChatMessage {
pub fn user(content: impl Into<String>) -> Self {
Self::new(MessageRole::User, content.into())
}
pub fn assistant(content: impl Into<String>) -> Self {
Self::new(MessageRole::Assistant, content.into())
}
pub fn system(content: impl Into<String>) -> Self {
Self::new(MessageRole::System, content.into())
}
pub fn run_summary(content: impl Into<String>) -> Self {
let mut m = Self::new(MessageRole::System, content.into());
m.kind = ChatMessageKind::RunSummary;
m
}
pub fn tool(
tool_call_id: impl Into<String>,
tool_name: impl Into<String>,
content: impl Into<String>,
) -> Self {
Self {
role: MessageRole::Tool,
content: content.into(),
timestamp: chrono::Local::now(),
kind: ChatMessageKind::Normal,
metadata: None,
actions: Vec::new(),
thinking: None,
images: None,
image_numbers: None,
tool_calls: None,
tool_call_id: Some(tool_call_id.into()),
tool_name: Some(tool_name.into()),
provider_continuation: None,
}
}
fn new(role: MessageRole, content: String) -> Self {
Self {
role,
content,
timestamp: chrono::Local::now(),
kind: ChatMessageKind::Normal,
metadata: None,
actions: Vec::new(),
thinking: None,
images: None,
image_numbers: None,
tool_calls: None,
tool_call_id: None,
tool_name: None,
provider_continuation: None,
}
}
pub fn with_images(mut self, images: Vec<String>) -> Self {
self.images = Some(images);
self
}
pub fn with_image_numbers(mut self, numbers: Vec<u64>) -> Self {
self.image_numbers = Some(numbers);
self
}
pub fn with_tool_calls(mut self, tool_calls: Vec<crate::models::tool_call::ToolCall>) -> Self {
self.tool_calls = if tool_calls.is_empty() {
None
} else {
Some(tool_calls)
};
self
}
pub fn with_provider_continuation(mut self, continuation: ProviderContinuation) -> Self {
self.provider_continuation = Some(continuation);
self
}
pub fn extract_thinking(text: &str) -> (Option<String>, String) {
let Some(thinking_start) = text.find("Thinking...") else {
return (None, text.to_string());
};
let content_start = thinking_start + "Thinking...".len();
if let Some(thinking_end) = text.find("...done thinking.") {
let thinking_text = text[content_start..thinking_end].trim().to_string();
let answer_start = thinking_end + "...done thinking.".len();
let answer_text = text[answer_start..].trim().to_string();
return (Some(thinking_text), answer_text);
}
let thinking_text = text[content_start..].trim().to_string();
(Some(thinking_text), String::new())
}
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub enum MessageRole {
User,
Assistant,
System,
Tool,
}
impl<'de> Deserialize<'de> for MessageRole {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let raw = String::deserialize(deserializer)?;
Ok(match raw.as_str() {
"User" => MessageRole::User,
"Assistant" => MessageRole::Assistant,
"System" => MessageRole::System,
"Tool" => MessageRole::Tool,
other => {
tracing::warn!(
role = %other,
"unknown message role in saved conversation; treating as System (version skew?)"
);
MessageRole::System
},
})
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ChatMessageKind {
#[default]
Normal,
ContextCheckpoint,
RunSummary,
Continuation,
RecoveryNudge,
#[serde(other)]
Unknown,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
Stop,
ToolUse,
Length,
ContentFilter,
Other(String),
}
#[derive(Debug, Clone)]
pub struct ModelResponse {
pub content: String,
pub usage: Option<TokenUsage>,
pub model_name: String,
pub thinking: Option<String>,
pub tool_calls: Option<Vec<crate::models::tool_call::ToolCall>>,
pub stop_reason: Option<FinishReason>,
pub provider_continuation: Option<ProviderContinuation>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TokenUsageSource {
#[default]
Provider,
Estimate,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenUsage {
pub prompt_tokens: usize,
pub completion_tokens: usize,
#[serde(default)]
pub cached_input_tokens: usize,
#[serde(default)]
pub cache_creation_input_tokens: usize,
#[serde(default)]
pub reasoning_output_tokens: usize,
#[serde(default)]
pub source: TokenUsageSource,
}
impl TokenUsage {
pub fn provider(prompt_tokens: usize, completion_tokens: usize) -> Self {
Self {
prompt_tokens,
completion_tokens,
cached_input_tokens: 0,
cache_creation_input_tokens: 0,
reasoning_output_tokens: 0,
source: TokenUsageSource::Provider,
}
}
pub fn estimate(prompt_tokens: usize) -> Self {
Self {
prompt_tokens,
completion_tokens: 0,
cached_input_tokens: 0,
cache_creation_input_tokens: 0,
reasoning_output_tokens: 0,
source: TokenUsageSource::Estimate,
}
}
pub fn with_cached_input(mut self, cached_input_tokens: usize) -> Self {
self.cached_input_tokens = cached_input_tokens;
self
}
pub fn with_cache_creation(mut self, cache_creation_input_tokens: usize) -> Self {
self.cache_creation_input_tokens = cache_creation_input_tokens;
self
}
pub fn with_reasoning_output(mut self, reasoning_output_tokens: usize) -> Self {
self.reasoning_output_tokens = reasoning_output_tokens;
self
}
pub fn input_total_tokens(&self) -> usize {
self.prompt_tokens
.saturating_add(self.cached_input_tokens)
.saturating_add(self.cache_creation_input_tokens)
}
pub fn output_total_tokens(&self) -> usize {
self.completion_tokens
.saturating_add(self.reasoning_output_tokens)
}
pub fn total_tokens(&self) -> usize {
self.input_total_tokens()
.saturating_add(self.output_total_tokens())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_message_role_equality() {
let user1 = MessageRole::User;
let user2 = MessageRole::User;
let assistant = MessageRole::Assistant;
assert_eq!(user1, user2, "User roles should be equal");
assert_ne!(user1, assistant, "Different roles should not be equal");
}
#[test]
fn test_chat_message_constructors() {
let user = ChatMessage::user("Hello!");
assert_eq!(user.role, MessageRole::User);
assert_eq!(user.content, "Hello!");
assert!(user.tool_calls.is_none());
let assistant = ChatMessage::assistant("Hi there");
assert_eq!(assistant.role, MessageRole::Assistant);
let system = ChatMessage::system("You are helpful");
assert_eq!(system.role, MessageRole::System);
let tool = ChatMessage::tool("call_1", "read_file", "file contents");
assert_eq!(tool.role, MessageRole::Tool);
assert_eq!(tool.tool_call_id, Some("call_1".to_string()));
assert_eq!(tool.tool_name, Some("read_file".to_string()));
}
#[test]
fn test_chat_message_builders() {
let msg = ChatMessage::user("test").with_images(vec!["base64data".to_string()]);
assert_eq!(msg.images, Some(vec!["base64data".to_string()]));
}
#[test]
fn test_token_usage_structure() {
let usage = TokenUsage::provider(100, 50)
.with_cached_input(25)
.with_cache_creation(5)
.with_reasoning_output(10);
assert_eq!(usage.prompt_tokens, 100);
assert_eq!(usage.completion_tokens, 50);
assert_eq!(usage.cached_input_tokens, 25);
assert_eq!(usage.cache_creation_input_tokens, 5);
assert_eq!(usage.reasoning_output_tokens, 10);
assert_eq!(usage.input_total_tokens(), 130);
assert_eq!(usage.output_total_tokens(), 60);
assert_eq!(usage.total_tokens(), 190);
assert_eq!(usage.source, TokenUsageSource::Provider);
}
#[test]
fn extract_thinking_no_marker_returns_text_unchanged() {
let (thinking, answer) = ChatMessage::extract_thinking("just a plain answer");
assert_eq!(thinking, None);
assert_eq!(answer, "just a plain answer");
}
#[test]
fn extract_thinking_complete_block() {
let raw = "Thinking...\n reasoning here\n...done thinking.\n\nFinal answer";
let (thinking, answer) = ChatMessage::extract_thinking(raw);
assert_eq!(thinking.as_deref(), Some("reasoning here"));
assert_eq!(answer, "Final answer");
}
#[test]
fn provider_continuation_round_trips_through_serde() {
let msg = ChatMessage::assistant("Step 3 lives.").with_provider_continuation(
ProviderContinuation::Anthropic {
signature: "sig_abc123_encrypted_blob".to_string(),
},
);
let json = serde_json::to_string(&msg).expect("serialize");
let back: ChatMessage = serde_json::from_str(&json).expect("deserialize");
assert_eq!(
back.provider_continuation
.as_ref()
.and_then(ProviderContinuation::anthropic_signature),
Some("sig_abc123_encrypted_blob")
);
assert_eq!(back.content, "Step 3 lives.");
}
#[test]
fn provider_continuation_defaults_to_none() {
let pre_step3_json = r#"{
"role": "Assistant",
"content": "hello",
"timestamp": "2026-04-16T12:00:00-04:00"
}"#;
let msg: ChatMessage = serde_json::from_str(pre_step3_json).expect("backward compat");
assert!(msg.provider_continuation.is_none());
}
#[test]
fn meta_encrypted_continuation_survives_persistence_redaction_byte_exact() {
let original = "eyJopaque.reasoning.payload";
let message = ChatMessage::assistant("done").with_provider_continuation(
ProviderContinuation::MetaResponses {
output: vec![MetaResponseItem::from_wire(serde_json::json!({
"type": "reasoning",
"id": "rs_1",
"summary": [],
"encrypted_content": original,
}))],
},
);
let mut persisted = serde_json::to_value(message).unwrap();
crate::utils::redact_json(&mut persisted);
let restored: ChatMessage = serde_json::from_value(persisted).unwrap();
let output = restored
.provider_continuation
.as_ref()
.and_then(ProviderContinuation::meta_output)
.unwrap();
assert_eq!(output[0].to_wire()["encrypted_content"], original);
}
#[test]
fn unknown_message_role_deserializes_to_system() {
let role: MessageRole = serde_json::from_str("\"Developer\"").expect("tolerant");
assert_eq!(role, MessageRole::System);
assert_eq!(
serde_json::from_str::<MessageRole>("\"Tool\"").unwrap(),
MessageRole::Tool
);
}
#[test]
fn unknown_message_kind_deserializes_to_unknown() {
let kind: ChatMessageKind = serde_json::from_str("\"some_future_kind\"").expect("tolerant");
assert_eq!(kind, ChatMessageKind::Unknown);
assert_ne!(kind, ChatMessageKind::Normal);
}
#[test]
fn continuation_kinds_round_trip_through_serde() {
for kind in [
ChatMessageKind::Continuation,
ChatMessageKind::RecoveryNudge,
] {
let json = serde_json::to_string(&kind).unwrap();
let back: ChatMessageKind = serde_json::from_str(&json).unwrap();
assert_eq!(back, kind);
}
assert_eq!(
serde_json::to_string(&ChatMessageKind::Continuation).unwrap(),
"\"continuation\""
);
assert_eq!(
serde_json::to_string(&ChatMessageKind::RecoveryNudge).unwrap(),
"\"recovery_nudge\""
);
}
#[test]
fn chat_message_with_unknown_role_round_trips() {
let json = r#"{
"role": "Developer",
"content": "hi",
"timestamp": "2026-04-16T12:00:00-04:00"
}"#;
let msg: ChatMessage = serde_json::from_str(json).expect("tolerant");
assert_eq!(msg.role, MessageRole::System);
assert_eq!(msg.content, "hi");
}
#[test]
fn extract_thinking_in_progress_no_end_marker() {
let raw = "Thinking...\n partial reasoning so far";
let (thinking, answer) = ChatMessage::extract_thinking(raw);
assert_eq!(thinking.as_deref(), Some("partial reasoning so far"));
assert_eq!(answer, "");
}
#[test]
fn test_model_response_creation() {
let usage = TokenUsage::provider(100, 50);
let response = ModelResponse {
content: "Hello, world!".to_string(),
usage: Some(usage),
model_name: "ollama/tinyllama".to_string(),
thinking: None,
tool_calls: None,
stop_reason: None,
provider_continuation: None,
};
assert_eq!(response.content, "Hello, world!");
assert!(response.usage.is_some());
assert_eq!(response.model_name, "ollama/tinyllama");
assert_eq!(response.usage.unwrap().total_tokens(), 150);
assert!(response.tool_calls.is_none());
}
}