use serde::{Deserialize, Serialize};
use std::fmt;
use std::sync::Arc;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum Content {
#[serde(rename = "text")]
Text { text: String },
#[serde(rename = "image")]
Image {
data: String,
#[serde(rename = "mimeType")]
mime_type: String,
},
#[serde(rename = "thinking")]
#[non_exhaustive]
Thinking {
thinking: String,
#[serde(skip_serializing_if = "Option::is_none")]
signature: Option<String>,
},
#[serde(rename = "toolCall")]
#[non_exhaustive]
ToolCall {
id: String,
name: String,
arguments: serde_json::Value,
#[serde(
default,
skip_serializing_if = "Option::is_none",
rename = "providerMetadata",
alias = "provider_metadata"
)]
provider_metadata: Option<serde_json::Value>,
},
}
impl Content {
pub fn tool_call(
id: impl Into<String>,
name: impl Into<String>,
arguments: serde_json::Value,
) -> Self {
Self::ToolCall {
id: id.into(),
name: name.into(),
arguments,
provider_metadata: None,
}
}
pub fn thinking(text: impl Into<String>) -> Self {
Self::Thinking {
thinking: text.into(),
signature: None,
}
}
pub fn thinking_signed(text: impl Into<String>, signature: impl Into<String>) -> Self {
Self::Thinking {
thinking: text.into(),
signature: Some(signature.into()),
}
}
pub fn tool_call_with_metadata(
id: impl Into<String>,
name: impl Into<String>,
arguments: serde_json::Value,
provider_metadata: serde_json::Value,
) -> Self {
Self::ToolCall {
id: id.into(),
name: name.into(),
arguments,
provider_metadata: Some(provider_metadata),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "role")]
pub enum Message {
#[serde(rename = "user")]
User {
content: Vec<Content>,
timestamp: u64,
},
#[serde(rename = "assistant")]
#[non_exhaustive]
Assistant {
content: Vec<Content>,
#[serde(rename = "stopReason")]
stop_reason: StopReason,
model: String,
provider: String,
usage: Usage,
timestamp: u64,
#[serde(
skip_serializing_if = "Option::is_none",
rename = "errorMessage",
alias = "error_message"
)]
error_message: Option<String>,
},
#[serde(rename = "toolResult")]
ToolResult {
#[serde(rename = "toolCallId")]
tool_call_id: String,
#[serde(rename = "toolName")]
tool_name: String,
content: Vec<Content>,
#[serde(rename = "isError")]
is_error: bool,
timestamp: u64,
},
}
impl Message {
pub fn user(text: impl Into<String>) -> Self {
Self::User {
content: vec![Content::Text { text: text.into() }],
timestamp: now_ms(),
}
}
pub fn assistant(
content: Vec<Content>,
stop_reason: StopReason,
model: impl Into<String>,
provider: impl Into<String>,
usage: Usage,
) -> Self {
Self::Assistant {
content,
stop_reason,
model: model.into(),
provider: provider.into(),
usage,
timestamp: now_ms(),
error_message: None,
}
}
pub fn with_error_message(mut self, msg: impl Into<String>) -> Self {
if let Self::Assistant { error_message, .. } = &mut self {
*error_message = Some(msg.into());
}
self
}
pub fn with_timestamp(mut self, ts: u64) -> Self {
match &mut self {
Self::User { timestamp, .. }
| Self::Assistant { timestamp, .. }
| Self::ToolResult { timestamp, .. } => *timestamp = ts,
}
self
}
pub fn role(&self) -> &str {
match self {
Self::User { .. } => "user",
Self::Assistant { .. } => "assistant",
Self::ToolResult { .. } => "toolResult",
}
}
pub fn is_context_overflow(&self) -> bool {
match self {
Self::Assistant {
stop_reason: StopReason::Error,
error_message: Some(msg),
..
} => crate::provider::is_context_overflow_message(msg),
_ => false,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ExtensionMessage {
pub role: String,
pub kind: String,
pub data: serde_json::Value,
}
impl ExtensionMessage {
pub fn new(kind: impl Into<String>, data: impl Serialize) -> Self {
Self {
role: "extension".into(),
kind: kind.into(),
data: serde_json::to_value(data).unwrap_or(serde_json::Value::Null),
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum AgentMessage {
Llm(Message),
Extension(ExtensionMessage),
}
impl AgentMessage {
pub fn role(&self) -> &str {
match self {
Self::Llm(m) => m.role(),
Self::Extension(ext) => &ext.role,
}
}
pub fn as_llm(&self) -> Option<&Message> {
match self {
Self::Llm(m) => Some(m),
Self::Extension(_) => None,
}
}
}
impl From<Message> for AgentMessage {
fn from(m: Message) -> Self {
Self::Llm(m)
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub enum StopReason {
Stop,
Length,
ToolUse,
Error,
Aborted,
Refusal,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct Usage {
pub input: u64,
pub output: u64,
#[serde(default, rename = "cacheRead", alias = "cache_read")]
pub cache_read: u64,
#[serde(default, rename = "cacheWrite", alias = "cache_write")]
pub cache_write: u64,
#[serde(default, rename = "totalTokens", alias = "total_tokens")]
pub total_tokens: u64,
}
impl Usage {
pub fn cache_hit_rate(&self) -> f64 {
let total_input = self.input + self.cache_read + self.cache_write;
if total_input == 0 {
return 0.0;
}
self.cache_read as f64 / total_input as f64
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct CacheConfig {
pub enabled: bool,
pub strategy: CacheStrategy,
#[serde(default)]
pub session_key: Option<String>,
}
impl CacheConfig {
pub fn new() -> Self {
Self::default()
}
pub fn disabled() -> Self {
Self {
enabled: false,
..Self::default()
}
}
pub fn with_session_key(mut self, key: impl Into<String>) -> Self {
let key = key.into();
self.session_key = if key.trim().is_empty() {
None
} else {
Some(key)
};
self
}
pub fn with_strategy(mut self, strategy: CacheStrategy) -> Self {
self.strategy = strategy;
self
}
pub fn hints_enabled(&self) -> bool {
if !self.enabled {
return false;
}
!matches!(
self.strategy,
CacheStrategy::Disabled
| CacheStrategy::Manual {
cache_system: false,
cache_tools: false,
cache_messages: false,
}
)
}
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
enabled: true,
strategy: CacheStrategy::Auto,
session_key: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum ToolExecutionStrategy {
Sequential,
#[default]
Parallel,
Batched { size: usize },
}
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub enum CacheStrategy {
#[default]
Auto,
Disabled,
Manual {
cache_system: bool,
cache_tools: bool,
cache_messages: bool,
},
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "lowercase")]
pub enum ThinkingLevel {
#[default]
Off,
Minimal,
Low,
Medium,
High,
}
pub type ToolUpdateFn = Arc<dyn Fn(ToolResult) + Send + Sync>;
pub type ProgressFn = Arc<dyn Fn(String) + Send + Sync>;
#[non_exhaustive]
pub struct ToolContext {
pub tool_call_id: String,
pub tool_name: String,
pub cancel: tokio_util::sync::CancellationToken,
pub on_update: Option<ToolUpdateFn>,
pub on_progress: Option<ProgressFn>,
}
impl ToolContext {
pub fn new(tool_call_id: impl Into<String>, tool_name: impl Into<String>) -> Self {
Self {
tool_call_id: tool_call_id.into(),
tool_name: tool_name.into(),
cancel: tokio_util::sync::CancellationToken::new(),
on_update: None,
on_progress: None,
}
}
pub fn with_cancel(mut self, cancel: tokio_util::sync::CancellationToken) -> Self {
self.cancel = cancel;
self
}
pub fn with_on_update(mut self, on_update: ToolUpdateFn) -> Self {
self.on_update = Some(on_update);
self
}
pub fn with_on_progress(mut self, on_progress: ProgressFn) -> Self {
self.on_progress = Some(on_progress);
self
}
}
impl Clone for ToolContext {
fn clone(&self) -> Self {
Self {
tool_call_id: self.tool_call_id.clone(),
tool_name: self.tool_name.clone(),
cancel: self.cancel.clone(),
on_update: self.on_update.clone(),
on_progress: self.on_progress.clone(),
}
}
}
impl std::fmt::Debug for ToolContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ToolContext")
.field("tool_call_id", &self.tool_call_id)
.field("tool_name", &self.tool_name)
.field("cancel", &self.cancel)
.field("on_update", &self.on_update.as_ref().map(|_| "<callback>"))
.field(
"on_progress",
&self.on_progress.as_ref().map(|_| "<callback>"),
)
.finish()
}
}
#[async_trait::async_trait]
pub trait AgentTool: Send + Sync {
fn name(&self) -> &str;
fn label(&self) -> &str;
fn description(&self) -> &str;
fn parameters_schema(&self) -> serde_json::Value;
async fn execute(
&self,
params: serde_json::Value,
ctx: ToolContext,
) -> Result<ToolResult, ToolError>;
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ToolResult {
pub content: Vec<Content>,
#[serde(default)]
pub details: serde_json::Value,
}
#[derive(Debug, thiserror::Error)]
pub enum ToolError {
#[error("{0}")]
Failed(String),
#[error("Tool not found: {0}")]
NotFound(String),
#[error("Invalid arguments: {0}")]
InvalidArgs(String),
#[error("Cancelled")]
Cancelled,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(
tag = "type",
rename_all = "camelCase",
rename_all_fields = "camelCase"
)]
#[non_exhaustive]
pub enum AgentEvent {
AgentStart,
#[non_exhaustive]
AgentEnd {
messages: Vec<AgentMessage>,
#[serde(default)]
stats: SessionStats,
},
TurnStart,
TurnEnd {
message: AgentMessage,
tool_results: Vec<Message>,
},
MessageStart {
message: AgentMessage,
},
MessageUpdate {
message: AgentMessage,
delta: StreamDelta,
},
MessageEnd {
message: AgentMessage,
},
ToolExecutionStart {
tool_call_id: String,
tool_name: String,
args: serde_json::Value,
},
ToolExecutionUpdate {
tool_call_id: String,
tool_name: String,
partial_result: ToolResult,
},
ToolExecutionEnd {
tool_call_id: String,
tool_name: String,
result: ToolResult,
is_error: bool,
},
ProgressMessage {
tool_call_id: String,
tool_name: String,
text: String,
},
InputRejected {
reason: String,
},
#[non_exhaustive]
LoopDetected {
tool_name: String,
repetitions: usize,
aborted: bool,
},
ContextCompacted {
method: CompactionMethod,
messages_before: usize,
messages_after: usize,
tokens_before: usize,
tokens_after: usize,
summary: Option<SummaryStats>,
},
}
impl AgentEvent {
pub fn agent_end(messages: Vec<AgentMessage>, stats: SessionStats) -> Self {
Self::AgentEnd { messages, stats }
}
pub fn loop_detected(tool_name: impl Into<String>, repetitions: usize, aborted: bool) -> Self {
Self::LoopDetected {
tool_name: tool_name.into(),
repetitions,
aborted,
}
}
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub struct SessionStats {
#[serde(default)]
pub usage: Usage,
#[serde(default)]
pub turns: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost_usd: Option<f64>,
#[serde(default)]
pub compactions: u32,
}
impl SessionStats {
pub fn new(usage: Usage, turns: u32, cost_usd: Option<f64>, compactions: u32) -> Self {
Self {
usage,
turns,
cost_usd,
compactions,
}
}
pub fn cache_hit_rate(&self) -> f64 {
self.usage.cache_hit_rate()
}
pub(crate) fn record_turn(
&mut self,
usage: &Usage,
cost: Option<&crate::provider::CostConfig>,
) {
self.usage.input += usage.input;
self.usage.output += usage.output;
self.usage.cache_read += usage.cache_read;
self.usage.cache_write += usage.cache_write;
self.turns += 1;
if let Some(cost) = cost.filter(|c| c.is_configured()) {
*self.cost_usd.get_or_insert(0.0) += cost.cost_usd(usage);
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub struct SummaryStats {
pub messages_summarized: usize,
pub usage: Usage,
#[serde(skip_serializing_if = "Option::is_none")]
pub cost_usd: Option<f64>,
}
impl SummaryStats {
pub fn new(messages_summarized: usize, usage: Usage, cost_usd: Option<f64>) -> Self {
Self {
messages_summarized,
usage,
cost_usd,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub enum CompactionMethod {
Summarized,
Deterministic,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(
tag = "type",
rename_all = "camelCase",
rename_all_fields = "camelCase"
)]
#[non_exhaustive]
pub enum StreamDelta {
Text { delta: String },
Thinking { delta: String },
ToolCallDelta { delta: String },
}
pub struct AgentContext {
pub system_prompt: String,
pub messages: Vec<AgentMessage>,
pub tools: Vec<Box<dyn AgentTool>>,
}
#[derive(Debug, Clone)]
pub enum FilterResult {
Pass,
Warn(String),
Reject(String),
}
pub trait InputFilter: Send + Sync {
fn filter(&self, text: &str) -> FilterResult;
}
#[derive(Debug, Clone)]
pub enum ToolDecision {
Allow,
Modify(serde_json::Value),
Deny(String),
}
#[derive(Debug)]
#[non_exhaustive]
pub struct ToolCallRequest<'a> {
pub tool_call_id: &'a str,
pub tool_name: &'a str,
pub args: &'a serde_json::Value,
}
#[async_trait::async_trait]
pub trait ToolMiddleware: Send + Sync {
async fn before_tool(&self, call: &ToolCallRequest<'_>) -> ToolDecision;
}
pub fn now_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
impl fmt::Display for StopReason {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Stop => write!(f, "stop"),
Self::Length => write!(f, "length"),
Self::ToolUse => write!(f, "toolUse"),
Self::Error => write!(f, "error"),
Self::Aborted => write!(f, "aborted"),
Self::Refusal => write!(f, "refusal"),
}
}
}
#[cfg(test)]
mod wire_tag_freeze {
use super::*;
use std::collections::BTreeSet;
macro_rules! wire_freeze {
($ty:ty, $tag_of:ident, $samples:ident, $($pat:pat_param => $tag:literal = $sample:expr),+ $(,)?) => {
fn $tag_of(v: &$ty) -> &'static str {
match v { $($pat => $tag,)+ }
}
fn $samples() -> Vec<$ty> { vec![$($sample,)+] }
};
}
fn msg() -> AgentMessage {
AgentMessage::Llm(Message::Assistant {
content: vec![Content::Text { text: "hi".into() }],
stop_reason: StopReason::ToolUse,
model: "mock-1".into(),
provider: "mock".into(),
usage: Usage {
input: 11,
output: 22,
cache_read: 33,
cache_write: 44,
total_tokens: 110,
},
timestamp: 7,
error_message: Some("boom".into()),
})
}
fn tool_result() -> ToolResult {
ToolResult {
content: vec![Content::Text { text: "ok".into() }],
details: serde_json::json!({"exitCode": 0}),
}
}
fn tool_result_message() -> Message {
Message::ToolResult {
tool_call_id: "tc-1".into(),
tool_name: "bash".into(),
content: vec![Content::Text { text: "ok".into() }],
is_error: false,
timestamp: 9,
}
}
wire_freeze! {
AgentEvent, expected_event_tag, event_samples,
AgentEvent::AgentStart => "agentStart" = AgentEvent::AgentStart,
AgentEvent::AgentEnd { .. } => "agentEnd" = AgentEvent::agent_end(
vec![msg()],
SessionStats {
usage: Usage {
input: 5,
output: 6,
cache_read: 7,
cache_write: 8,
total_tokens: 26,
},
turns: 3,
cost_usd: Some(0.02),
compactions: 1,
},
),
AgentEvent::TurnStart => "turnStart" = AgentEvent::TurnStart,
AgentEvent::TurnEnd { .. } => "turnEnd" = AgentEvent::TurnEnd {
message: msg(),
tool_results: vec![tool_result_message()],
},
AgentEvent::MessageStart { .. } => "messageStart"
= AgentEvent::MessageStart { message: msg() },
AgentEvent::MessageUpdate { .. } => "messageUpdate" = AgentEvent::MessageUpdate {
message: msg(),
delta: StreamDelta::Text { delta: "hi".into() },
},
AgentEvent::MessageEnd { .. } => "messageEnd"
= AgentEvent::MessageEnd { message: msg() },
AgentEvent::ToolExecutionStart { .. } => "toolExecutionStart"
= AgentEvent::ToolExecutionStart {
tool_call_id: "tc-1".into(),
tool_name: "bash".into(),
args: serde_json::json!({"command": "ls"}),
},
AgentEvent::ToolExecutionUpdate { .. } => "toolExecutionUpdate"
= AgentEvent::ToolExecutionUpdate {
tool_call_id: "tc-1".into(),
tool_name: "bash".into(),
partial_result: tool_result(),
},
AgentEvent::ToolExecutionEnd { .. } => "toolExecutionEnd"
= AgentEvent::ToolExecutionEnd {
tool_call_id: "tc-1".into(),
tool_name: "bash".into(),
result: tool_result(),
is_error: false,
},
AgentEvent::ProgressMessage { .. } => "progressMessage"
= AgentEvent::ProgressMessage {
tool_call_id: "tc-1".into(),
tool_name: "bash".into(),
text: "50% done".into(),
},
AgentEvent::InputRejected { .. } => "inputRejected"
= AgentEvent::InputRejected { reason: "injection detected".into() },
AgentEvent::LoopDetected { .. } => "loopDetected"
= AgentEvent::loop_detected("bash", 3, false),
AgentEvent::ContextCompacted { .. } => "contextCompacted"
= AgentEvent::ContextCompacted {
method: CompactionMethod::Summarized,
messages_before: 40,
messages_after: 13,
tokens_before: 96_500,
tokens_after: 41_200,
summary: Some(SummaryStats::new(
28,
Usage {
input: 54_000,
output: 900,
cache_read: 0,
cache_write: 0,
total_tokens: 54_900,
},
Some(0.17),
)),
},
}
wire_freeze! {
StreamDelta, expected_delta_tag, delta_samples,
StreamDelta::Text { .. } => "text" = StreamDelta::Text { delta: "hi".into() },
StreamDelta::Thinking { .. } => "thinking"
= StreamDelta::Thinking { delta: "hmm".into() },
StreamDelta::ToolCallDelta { .. } => "toolCallDelta"
= StreamDelta::ToolCallDelta { delta: "{}".into() },
}
fn all_keys(v: &serde_json::Value, path: &str, out: &mut Vec<(String, String)>) {
match v {
serde_json::Value::Object(map) => {
for (k, child) in map {
out.push((path.to_string(), k.clone()));
all_keys(child, &format!("{path}.{k}"), out);
}
}
serde_json::Value::Array(items) => {
for (i, child) in items.iter().enumerate() {
all_keys(child, &format!("{path}[{i}]"), out);
}
}
_ => {}
}
}
fn assert_frozen<T>(sample: &T, declared: &'static str, seen: &mut BTreeSet<&'static str>)
where
T: serde::Serialize + serde::de::DeserializeOwned + std::fmt::Debug + PartialEq,
{
let v = serde_json::to_value(sample).expect("serialize");
assert_eq!(
v["type"], declared,
"wire tag drifted: {sample:?} serializes as {} but wire_freeze! declares {declared}. \
Changing a tag breaks every deployed client — if this is intentional it is a \
breaking change, not a test fix",
v["type"]
);
let mut keys = Vec::new();
all_keys(&v, declared, &mut keys);
for (path, key) in &keys {
assert!(
!key.contains('_'),
"payload key {key:?} at {path} is not camelCase. Every struct on this wire \
carries rename_all = \"camelCase\" and TS clients hardcode these names, so a \
snake_case key here means a rename attribute is missing"
);
}
let back: T = serde_json::from_value(v).expect("round-trip deserialize");
assert_eq!(
&back, sample,
"{declared} did not survive a JSON round-trip"
);
assert!(
seen.insert(declared),
"two samples serialize as {declared} — a sample in wire_freeze! does not match \
the pattern on its own line, so some variant has no sample at all"
);
}
#[test]
fn every_event_variant_is_frozen_tagged_and_round_trips() {
let mut seen = BTreeSet::new();
for sample in &event_samples() {
assert_frozen(sample, expected_event_tag(sample), &mut seen);
}
}
#[test]
fn every_delta_variant_is_frozen_tagged_and_round_trips() {
let mut seen = BTreeSet::new();
for sample in &delta_samples() {
assert_frozen(sample, expected_delta_tag(sample), &mut seen);
}
}
}