use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::SystemTime;
use crate::tool::ToolCall;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "role", rename_all = "snake_case")]
pub enum AgentMessage {
System {
content: String,
#[serde(default = "default_timestamp", skip_serializing_if = "Option::is_none")]
timestamp: Option<u64>,
},
User {
content: UserContent,
#[serde(default = "default_timestamp", skip_serializing_if = "Option::is_none")]
timestamp: Option<u64>,
},
Assistant {
content: AssistantContent,
stop_reason: StopReason,
#[serde(default, skip_serializing_if = "Option::is_none")]
error_message: Option<String>,
#[serde(default = "default_timestamp", skip_serializing_if = "Option::is_none")]
timestamp: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
usage: Option<Usage>,
},
ToolResult {
tool_call_id: String,
tool_name: String,
content: ToolResultContent,
#[serde(default)]
is_error: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
narration: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
details: Option<Value>,
#[serde(default = "default_timestamp", skip_serializing_if = "Option::is_none")]
timestamp: Option<u64>,
},
Custom {
kind: String,
#[serde(default)]
payload: Value,
#[serde(default = "default_timestamp", skip_serializing_if = "Option::is_none")]
timestamp: Option<u64>,
},
}
fn default_timestamp() -> Option<u64> {
SystemTime::now()
.duration_since(SystemTime::UNIX_EPOCH)
.ok()
.map(|d| d.as_millis() as u64)
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
#[serde(default)]
pub input_tokens: i64,
#[serde(default)]
pub output_tokens: i64,
#[serde(default)]
pub cache_creation_input_tokens: i64,
#[serde(default)]
pub cache_read_input_tokens: i64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
EndTurn,
ToolUse,
MaxTokens,
Error,
Aborted,
Other,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum UserContent {
Text(String),
Blocks(Vec<UserBlock>),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum UserBlock {
Text(TextContent),
Image(ImageContent),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct AssistantContent {
pub blocks: Vec<AssistantBlock>,
}
impl AssistantContent {
pub fn text(text: impl Into<String>) -> Self {
Self {
blocks: vec![AssistantBlock::Text(TextContent { text: text.into() })],
}
}
pub fn with_tool_calls(text: Option<String>, tool_calls: Vec<ToolCall>) -> Self {
let mut blocks = Vec::new();
if let Some(t) = text.filter(|s| !s.trim().is_empty()) {
blocks.push(AssistantBlock::Text(TextContent { text: t }));
}
for call in tool_calls {
blocks.push(AssistantBlock::ToolCall(call));
}
Self { blocks }
}
pub fn plain_text(&self) -> String {
self.blocks
.iter()
.filter_map(|b| match b {
AssistantBlock::Text(t) => Some(t.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("")
}
pub fn tool_calls(&self) -> Vec<&ToolCall> {
self.blocks
.iter()
.filter_map(|b| match b {
AssistantBlock::ToolCall(c) => Some(c),
_ => None,
})
.collect()
}
pub fn thinking_text(&self) -> String {
self.blocks
.iter()
.filter_map(|b| match b {
AssistantBlock::Thinking(t) => Some(t.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("\n")
}
pub fn reasoning_text(&self) -> String {
self.blocks
.iter()
.filter_map(|b| match b {
AssistantBlock::Reasoning(t) => Some(t.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("")
}
pub fn reasoning_details_values(&self) -> Vec<Value> {
self.blocks
.iter()
.filter_map(|b| match b {
AssistantBlock::ReasoningDetails(d) => Some(d.details.as_slice()),
_ => None,
})
.flatten()
.cloned()
.collect()
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AssistantBlock {
Text(TextContent),
Thinking(TextContent),
Reasoning(TextContent),
ReasoningDetails(ReasoningDetailsContent),
ToolCall(ToolCall),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ReasoningDetailsContent {
pub details: Vec<Value>,
}
impl ReasoningDetailsContent {
pub fn new(details: Vec<Value>) -> Self {
Self { details }
}
pub fn as_items(&self) -> Vec<crate::reasoning::ReasoningItem> {
self.details
.iter()
.filter_map(crate::reasoning::ReasoningItem::from_openrouter_value)
.collect()
}
pub fn from_items(items: &[crate::reasoning::ReasoningItem]) -> Self {
Self {
details: items
.iter()
.map(crate::reasoning::ReasoningItem::to_openrouter_value)
.collect(),
}
}
pub fn has_signed_payload(&self) -> bool {
self.as_items()
.iter()
.any(crate::reasoning::ReasoningItem::carries_signed_payload)
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ToolResultContent {
pub blocks: Vec<ToolResultBlock>,
}
impl ToolResultContent {
pub fn text(text: impl Into<String>) -> Self {
Self {
blocks: vec![ToolResultBlock::Text(TextContent { text: text.into() })],
}
}
pub fn plain_text(&self) -> String {
self.blocks
.iter()
.filter_map(|b| match b {
ToolResultBlock::Text(t) => Some(t.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("")
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ToolResultBlock {
Text(TextContent),
Image(ImageContent),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TextContent {
pub text: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ImageContent {
pub source: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub media_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub alt: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct RunIdentity {
pub run_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_run_id: Option<String>,
#[serde(default)]
pub depth: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub conversation_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub deadline_unix_ms: Option<u64>,
}
impl RunIdentity {
pub fn root() -> Self {
Self {
run_id: uuid::Uuid::new_v4().to_string(),
parent_run_id: None,
depth: 0,
conversation_id: None,
deadline_unix_ms: None,
}
}
pub fn child_of(parent: &Self) -> Self {
Self {
run_id: uuid::Uuid::new_v4().to_string(),
parent_run_id: Some(parent.run_id.clone()),
depth: parent.depth + 1,
conversation_id: parent.conversation_id.clone(),
deadline_unix_ms: parent.deadline_unix_ms,
}
}
pub fn with_run_id(mut self, id: impl Into<String>) -> Self {
self.run_id = id.into();
self
}
pub fn with_conversation_id(mut self, id: impl Into<String>) -> Self {
self.conversation_id = Some(id.into());
self
}
pub fn with_deadline_unix_ms(mut self, ms: u64) -> Self {
self.deadline_unix_ms = Some(ms);
self
}
}
#[derive(Debug, Clone)]
pub struct AgentContext {
pub system_prompt: String,
pub messages: Vec<AgentMessage>,
pub identity: Option<RunIdentity>,
}
impl AgentContext {
pub fn new(system_prompt: impl Into<String>) -> Self {
Self {
system_prompt: system_prompt.into(),
messages: Vec::new(),
identity: None,
}
}
pub fn with_messages(mut self, messages: Vec<AgentMessage>) -> Self {
self.messages = messages;
self
}
pub fn with_identity(mut self, identity: RunIdentity) -> Self {
self.identity = Some(identity);
self
}
pub fn spawn_child(&self, system_prompt: impl Into<String>) -> Self {
let parent_identity = self.identity.clone().unwrap_or_else(RunIdentity::root);
Self {
system_prompt: system_prompt.into(),
messages: Vec::new(),
identity: Some(RunIdentity::child_of(&parent_identity)),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn user_text_round_trip() {
let msg = AgentMessage::User {
content: UserContent::Text("hello".into()),
timestamp: Some(0),
};
let json = serde_json::to_value(&msg).unwrap();
assert_eq!(json["role"], "user");
assert_eq!(json["content"], "hello");
}
#[test]
fn assistant_with_tool_call_blocks() {
let content = AssistantContent::with_tool_calls(
Some("calling…".into()),
vec![ToolCall {
id: "call_1".into(),
name: "shell".into(),
arguments: serde_json::json!({"cmd": "ls"}),
}],
);
assert_eq!(content.tool_calls().len(), 1);
assert_eq!(content.plain_text(), "calling…");
}
#[test]
fn custom_message_passthrough() {
let msg = AgentMessage::Custom {
kind: "ui_notification".into(),
payload: serde_json::json!({"text": "build started"}),
timestamp: None,
};
let json = serde_json::to_value(&msg).unwrap();
assert_eq!(json["role"], "custom");
assert_eq!(json["kind"], "ui_notification");
}
}