use serde::{Deserialize, Serialize};
use crate::tool::ToolProvenance;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum SystemCacheType {
Ephemeral,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct SystemCacheMarker {
pub offset: usize,
pub length: usize,
pub cache_type: SystemCacheType,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub input: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct MessageToolResult {
pub tool_use_id: String,
pub content: String,
pub is_error: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ReasoningBlock {
Thinking {
text: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
signature: Option<String>,
},
Redacted { data: String },
Plain { text: String },
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct AssistantReasoning {
pub provider: String,
pub model: String,
pub blocks: Vec<ReasoningBlock>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ContentDigest([u8; 32]);
impl ContentDigest {
pub const fn from_raw(bytes: [u8; 32]) -> Self {
Self(bytes)
}
pub const fn as_bytes(&self) -> &[u8; 32] {
&self.0
}
pub fn to_hex(&self) -> String {
let mut out = String::with_capacity(64);
for byte in self.0 {
out.push_str(&format!("{byte:02x}"));
}
out
}
pub fn from_hex(s: &str) -> Result<Self, String> {
if s.len() != 64 {
return Err(format!(
"content_digest must be 64 hex chars, got {}",
s.len()
));
}
let mut bytes = [0u8; 32];
for (i, chunk) in s.as_bytes().chunks(2).enumerate() {
let hex = std::str::from_utf8(chunk).map_err(|e| e.to_string())?;
bytes[i] = u8::from_str_radix(hex, 16).map_err(|e| e.to_string())?;
}
Ok(Self(bytes))
}
}
impl std::fmt::Display for ContentDigest {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.to_hex())
}
}
impl serde::Serialize for ContentDigest {
fn serialize<S: serde::Serializer>(&self, s: S) -> Result<S::Ok, S::Error> {
s.serialize_str(&self.to_hex())
}
}
impl<'de> serde::Deserialize<'de> for ContentDigest {
fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
let s = String::deserialize(d)?;
Self::from_hex(&s).map_err(serde::de::Error::custom)
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentPart {
Text {
text: String,
},
Image {
path: std::path::PathBuf,
mime: String,
byte_count: u64,
#[serde(default)]
content_digest: ContentDigest,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "role", rename_all = "snake_case")]
pub enum Message {
System {
content: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
cache_markers: Vec<SystemCacheMarker>,
},
Context {
content: String,
},
User {
content: String,
},
Multimodal {
parts: Vec<ContentPart>,
},
Assistant {
content: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
tool_calls: Vec<ToolCall>,
#[serde(default, skip_serializing_if = "Option::is_none")]
reasoning: Option<AssistantReasoning>,
},
Tool {
result: MessageToolResult,
},
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum StopReason {
EndTurn,
ToolUse,
MaxTokens,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct Usage {
pub input_tokens: u32,
pub output_tokens: u32,
#[serde(default)]
pub cache_read_tokens: u32,
#[serde(default)]
pub cache_write_tokens: u32,
#[serde(default)]
pub reasoning_tokens: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(tag = "type", rename_all = "snake_case")]
#[non_exhaustive]
pub enum AgentEvent {
TurnStart,
TextDelta {
delta: String,
},
ThinkingDelta {
delta: String,
},
ReasoningComplete {
blocks: Vec<ReasoningBlock>,
},
ToolCallStarted {
name: String,
},
ToolCall {
call: ToolCall,
provenance: ToolProvenance,
summary_fields: Vec<String>,
},
ToolResult {
result: MessageToolResult,
},
TurnEnd {
stop_reason: StopReason,
usage: Usage,
},
Error {
message: String,
},
}
#[cfg(test)]
#[allow(warnings)]
#[allow(warnings)]
#[allow(warnings)]
#[allow(warnings)]
mod tests {
use super::*;
#[test]
fn message_roundtrip_user() {
let msg = Message::User {
content: "Hello, world!".into(),
};
let json = serde_json::to_string(&msg).expect("operation should succeed");
let decoded: Message = serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(msg, decoded);
}
#[test]
fn message_roundtrip_assistant() {
let msg = Message::Assistant {
content: "I can help with that.".into(),
tool_calls: vec![ToolCall {
id: "call_1".into(),
name: "read_file".into(),
input: serde_json::json!({"path": "src/main.rs"}),
}],
reasoning: None,
};
let json = serde_json::to_string(&msg).expect("operation should succeed");
let decoded: Message = serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(msg, decoded);
}
#[test]
fn message_roundtrip_tool() {
let msg = Message::Tool {
result: MessageToolResult {
tool_use_id: "call_1".into(),
content: "fn main() {}".into(),
is_error: false,
},
};
let json = serde_json::to_string(&msg).expect("operation should succeed");
let decoded: Message = serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(msg, decoded);
}
#[test]
fn event_roundtrip() {
let events = vec![
AgentEvent::TurnStart,
AgentEvent::TextDelta {
delta: "Hello".into(),
},
AgentEvent::ToolCall {
call: ToolCall {
id: "c1".into(),
name: "bash".into(),
input: serde_json::json!({"command": "ls"}),
},
provenance: ToolProvenance::Native,
summary_fields: vec![],
},
AgentEvent::ToolResult {
result: MessageToolResult {
tool_use_id: "c1".into(),
content: "file.rs".into(),
is_error: false,
},
},
AgentEvent::TurnEnd {
stop_reason: StopReason::EndTurn,
usage: Usage {
input_tokens: 100,
output_tokens: 50,
cache_read_tokens: 80,
cache_write_tokens: 20,
reasoning_tokens: 0,
},
},
AgentEvent::Error {
message: "something failed".into(),
},
];
for event in events {
let json = serde_json::to_string(&event).expect("operation should succeed");
let decoded: AgentEvent =
serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(event, decoded);
}
}
#[test]
fn extract_tool_calls_preserves_id_from_json_tool_block() {
let text = r#"I'll run that for you.
```json-tool
{"id":"call_abc123","args":{"command":"ls"},"name":"bash"}
```
Done."#;
let calls = extract_tool_calls_from_text(text);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].id, "call_abc123");
assert_eq!(calls[0].name, "bash");
assert_eq!(calls[0].input, serde_json::json!({"command": "ls"}));
}
#[test]
fn extract_tool_calls_falls_back_to_synthetic_id_when_missing() {
let text = r#"```json-tool
{"args":{"command":"ls"},"name":"bash"}
```"#;
let calls = extract_tool_calls_from_text(text);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].id, "tc_0");
assert_eq!(calls[0].name, "bash");
}
#[test]
fn extract_tool_calls_falls_back_when_id_is_empty() {
let text = r#"```json-tool
{"id":"","args":{"command":"ls"},"name":"bash"}
```"#;
let calls = extract_tool_calls_from_text(text);
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].id, "tc_0");
}
#[test]
fn content_part_text_roundtrip() {
let part = ContentPart::Text {
text: "Hello, image!".into(),
};
let json = serde_json::to_string(&part).expect("operation should succeed");
let decoded: ContentPart = serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(part, decoded);
}
#[test]
fn content_part_image_roundtrip() {
let part = ContentPart::Image {
path: "/tmp/test.png".into(),
mime: "image/png".into(),
byte_count: 12345,
content_digest: ContentDigest::from_raw([7u8; 32]),
};
let json = serde_json::to_string(&part).expect("operation should succeed");
let decoded: ContentPart = serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(part, decoded);
}
#[test]
fn message_multimodal_roundtrip() {
let msg = Message::Multimodal {
parts: vec![
ContentPart::Text {
text: "What is in this image?".into(),
},
ContentPart::Image {
path: "/tmp/screenshot.png".into(),
mime: "image/png".into(),
byte_count: 67890,
content_digest: ContentDigest::from_raw([9u8; 32]),
},
],
};
let json = serde_json::to_string(&msg).expect("operation should succeed");
let decoded: Message = serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(msg, decoded);
}
#[test]
fn message_user_still_works_after_multimodal_addition() {
let msg = Message::User {
content: "text only".into(),
};
let json = serde_json::to_string(&msg).expect("operation should succeed");
let decoded: Message = serde_json::from_str(&json).expect("operation should succeed");
assert_eq!(msg, decoded);
}
}
pub fn extract_tool_calls_from_text(text: &str) -> Vec<ToolCall> {
let mut calls = Vec::new();
let mut remaining = text;
while let Some(start) = remaining.find("```json-tool") {
let inner_start = start + "```json-tool".len();
let inner = remaining[inner_start..].trim_start();
let end = inner.find("```").unwrap_or(inner.len());
let content = inner[..end].trim();
if let Ok(obj) = serde_json::from_str::<serde_json::Value>(content)
&& let (Some(name), Some(args)) = (obj["name"].as_str(), Some(obj["args"].clone()))
{
let id = obj["id"]
.as_str()
.filter(|s| !s.is_empty())
.map(String::from)
.unwrap_or_else(|| format!("tc_{}", calls.len()));
calls.push(ToolCall {
id,
name: name.to_string(),
input: args,
});
}
remaining = &inner[end..];
if end + 3 < remaining.len() {
remaining = &remaining[3..];
} else {
break;
}
}
calls
}
pub fn strip_tool_syntax(text: &str) -> String {
let mut result = text.to_string();
while let Some(start) = result.find("```json-tool") {
let inner_start = start + "```json-tool".len();
let inner = &result[inner_start..];
let end = inner_start + inner.find("```").unwrap_or(inner.len()) + 3;
result.replace_range(start..end, "");
}
result.trim().to_string()
}
pub fn project_displayable_reasoning(ar: &AssistantReasoning) -> Option<String> {
let mut parts = Vec::new();
for block in &ar.blocks {
match block {
ReasoningBlock::Thinking { text, .. } if !text.is_empty() => parts.push(text.clone()),
ReasoningBlock::Plain { text } if !text.is_empty() => parts.push(text.clone()),
_ => {}
}
}
if parts.is_empty() {
None
} else {
Some(parts.join("\n"))
}
}