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(default, skip_serializing_if = "Option::is_none")]
redacted: Option<String>,
#[serde(
default,
skip_serializing_if = "Option::is_none",
rename = "redactedProtocol",
alias = "redacted_protocol",
deserialize_with = "lenient_protocol"
)]
redacted_protocol: Option<crate::provider::ApiProtocol>,
},
#[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>,
},
}
fn lenient_protocol<'de, D>(
deserializer: D,
) -> Result<Option<crate::provider::ApiProtocol>, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = <Option<serde_json::Value> as serde::Deserialize>::deserialize(deserializer)?;
Ok(value.and_then(|v| serde_json::from_value(v).ok()))
}
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,
redacted: None,
redacted_protocol: None,
}
}
pub fn thinking_redacted(
protocol: crate::provider::ApiProtocol,
data: impl Into<String>,
) -> Self {
Self::Thinking {
thinking: String::new(),
signature: None,
redacted: Some(data.into()),
redacted_protocol: Some(protocol),
}
}
pub(crate) fn redacted_for(&self, protocol: crate::provider::ApiProtocol) -> Option<&str> {
match self {
Self::Thinking {
redacted: Some(data),
redacted_protocol: Some(p),
..
} if *p == protocol => Some(data),
_ => None,
}
}
pub fn thinking_signed(text: impl Into<String>, signature: impl Into<String>) -> Self {
Self::Thinking {
thinking: text.into(),
signature: Some(signature.into()),
redacted: None,
redacted_protocol: None,
}
}
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),
..
} => {
msg.starts_with(crate::provider::traits::CONTEXT_OVERFLOW_DISPLAY_PREFIX)
|| 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")]
#[non_exhaustive]
pub enum ThinkingLevel {
#[default]
Off,
Minimal,
Low,
Medium,
High,
XHigh,
Max,
}
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>,
pub(crate) sub_agent_report: Option<SubAgentReport>,
}
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,
sub_agent_report: 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
}
pub fn report_delegated_run(&self, stats: SessionStats) {
if let Some(report) = &self.sub_agent_report {
report.lock().unwrap_or_else(|e| e.into_inner()).push(stats);
}
}
}
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(),
sub_agent_report: self.sub_agent_report.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()
}
}
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
pub trait AgentTool: crate::rt::MaybeSend + crate::rt::MaybeSync {
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]
ProviderRetry {
attempt: usize,
max_attempts: usize,
error: String,
delay_ms: u64,
},
#[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 provider_retry(
attempt: usize,
max_attempts: usize,
error: impl Into<String>,
delay: std::time::Duration,
) -> Self {
Self::ProviderRetry {
attempt,
max_attempts,
error: error.into(),
delay_ms: u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
}
}
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,
#[serde(default, skip_serializing_if = "SubAgentSpend::is_empty")]
pub sub_agents: SubAgentSpend,
#[serde(default, skip_serializing_if = "DecisionStats::is_empty")]
pub decision: DecisionStats,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub struct DecisionStats {
#[serde(default)]
pub requests: u32,
#[serde(default)]
pub failures: u32,
#[serde(default)]
pub timeouts: u32,
#[serde(default)]
pub usage: Usage,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost_usd: Option<f64>,
#[serde(default)]
pub unpriced: u32,
}
impl DecisionStats {
pub fn is_empty(&self) -> bool {
self.requests == 0 && usage_is_zero(&self.usage) && self.cost_usd.is_none()
}
pub fn is_unpriced(&self) -> bool {
self.unpriced > 0 || is_unpriced(&self.usage, self.cost_usd)
}
pub fn merge(&mut self, other: &DecisionStats) {
self.cost_usd = if self.unpriced > 0 || other.unpriced > 0 {
None
} else {
combine_cost(&self.usage, self.cost_usd, &other.usage, other.cost_usd)
};
self.usage = add_usage(&self.usage, &other.usage);
self.requests = self.requests.saturating_add(other.requests);
self.failures = self.failures.saturating_add(other.failures);
self.timeouts = self.timeouts.saturating_add(other.timeouts);
self.unpriced = self.unpriced.saturating_add(other.unpriced);
}
#[cfg_attr(not(feature = "decision"), allow(dead_code))]
pub(crate) fn record_success(&mut self, input: u64, output: u64, cost_usd: Option<f64>) {
self.merge(&DecisionStats {
requests: 1,
usage: Usage {
input,
output,
..Default::default()
},
cost_usd,
unpriced: u32::from(cost_usd.is_none()),
..Default::default()
});
}
#[cfg_attr(not(feature = "decision"), allow(dead_code))]
pub(crate) fn record_billed(&mut self, input: u64, output: u64, cost_usd: Option<f64>) {
self.merge(&DecisionStats {
usage: Usage {
input,
output,
..Default::default()
},
cost_usd,
unpriced: u32::from(cost_usd.is_none()),
..Default::default()
});
}
#[cfg_attr(not(feature = "decision"), allow(dead_code))]
pub(crate) fn record_failure(&mut self, timed_out: bool) {
self.requests = self.requests.saturating_add(1);
self.failures = self.failures.saturating_add(1);
if timed_out {
self.timeouts = self.timeouts.saturating_add(1);
}
}
}
impl SessionStats {
pub fn new(usage: Usage, turns: u32, cost_usd: Option<f64>, compactions: u32) -> Self {
Self {
usage,
turns,
cost_usd,
compactions,
sub_agents: SubAgentSpend::default(),
decision: DecisionStats::default(),
}
}
pub fn is_unpriced(&self) -> bool {
is_unpriced(&self.usage, self.cost_usd)
}
pub fn total_usage(&self) -> Usage {
add_usage(&self.usage, &self.sub_agents.usage)
}
pub fn total_cost_usd(&self) -> Option<f64> {
let own_and_delegated = combine_cost(
&self.usage,
self.cost_usd,
&self.sub_agents.usage,
self.sub_agents.cost_usd,
);
if self.decision.is_unpriced() {
return None;
}
let llm_usage = add_usage(&self.usage, &self.sub_agents.usage);
combine_cost(
&llm_usage,
own_and_delegated,
&self.decision.usage,
self.decision.cost_usd,
)
}
pub(crate) fn merge(&mut self, other: &SessionStats) {
self.cost_usd = combine_cost(&self.usage, self.cost_usd, &other.usage, other.cost_usd);
self.usage = add_usage(&self.usage, &other.usage);
self.turns = self.turns.saturating_add(other.turns);
self.compactions = self.compactions.saturating_add(other.compactions);
self.sub_agents.merge(&other.sub_agents);
self.decision.merge(&other.decision);
}
pub fn from_sub_agent_result(result: &ToolResult) -> Option<SessionStats> {
serde_json::from_value(result.details.get(SUB_AGENT_STATS_KEY)?.clone()).ok()
}
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>,
) {
let turn_cost = cost.map(|c| c.cost_usd(usage));
self.cost_usd = combine_cost(&self.usage, self.cost_usd, usage, turn_cost);
self.usage = add_usage(&self.usage, usage);
self.turns = self.turns.saturating_add(1);
}
}
pub const SUB_AGENT_STATS_KEY: &str = "sub_agent_stats";
fn add_usage(a: &Usage, b: &Usage) -> Usage {
Usage {
input: a.input.saturating_add(b.input),
output: a.output.saturating_add(b.output),
cache_read: a.cache_read.saturating_add(b.cache_read),
cache_write: a.cache_write.saturating_add(b.cache_write),
total_tokens: 0,
}
}
fn usage_is_zero(u: &Usage) -> bool {
u.input == 0 && u.output == 0 && u.cache_read == 0 && u.cache_write == 0
}
fn is_unpriced(usage: &Usage, cost: Option<f64>) -> bool {
cost.is_none() && !usage_is_zero(usage)
}
pub(crate) fn combine_cost(
a_usage: &Usage,
a_cost: Option<f64>,
b_usage: &Usage,
b_cost: Option<f64>,
) -> Option<f64> {
if is_unpriced(a_usage, a_cost) || is_unpriced(b_usage, b_cost) {
return None;
}
match (a_cost, b_cost) {
(None, None) => None,
(a, b) => Some(a.unwrap_or(0.0) + b.unwrap_or(0.0)),
}
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub struct SubAgentSpend {
#[serde(default)]
pub usage: Usage,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cost_usd: Option<f64>,
#[serde(default)]
pub runs: u32,
}
impl SubAgentSpend {
pub fn is_empty(&self) -> bool {
self.runs == 0 && usage_is_zero(&self.usage) && self.cost_usd.is_none()
}
pub fn is_unpriced(&self) -> bool {
is_unpriced(&self.usage, self.cost_usd)
}
pub fn merge(&mut self, other: &SubAgentSpend) {
self.cost_usd = combine_cost(&self.usage, self.cost_usd, &other.usage, other.cost_usd);
self.usage = add_usage(&self.usage, &other.usage);
self.runs = self.runs.saturating_add(other.runs);
}
pub(crate) fn record_run(&mut self, child: &SessionStats) {
let mut subtree = SubAgentSpend {
usage: child.usage.clone(),
cost_usd: child.cost_usd,
runs: 1,
};
subtree.merge(&child.sub_agents);
self.merge(&subtree);
}
}
pub(crate) type SubAgentReport = Arc<std::sync::Mutex<Vec<SessionStats>>>;
#[cfg(test)]
mod spend_rollup_tests {
use super::*;
use crate::provider::CostConfig;
fn u(input: u64) -> Usage {
Usage {
input,
..Usage::default()
}
}
fn priced() -> CostConfig {
CostConfig::new(1.0, 0.0)
}
#[test]
fn record_turn_unpriced_then_priced_is_unknown() {
let mut stats = SessionStats::default();
stats.record_turn(&u(1_000_000), None);
assert_eq!(stats.cost_usd, None);
stats.record_turn(&u(1_000_000), Some(&priced()));
assert_eq!(stats.cost_usd, None, "a partial sum would under-report");
assert!(stats.is_unpriced());
assert_eq!(stats.usage, u(2_000_000));
assert_eq!(stats.turns, 2);
}
#[test]
fn record_turn_priced_then_unpriced_is_unknown() {
let mut stats = SessionStats::default();
stats.record_turn(&u(1_000_000), Some(&priced()));
assert_eq!(stats.cost_usd, Some(1.0));
stats.record_turn(&u(1), None);
assert_eq!(stats.cost_usd, None);
}
#[test]
fn record_turn_free_model_prices_at_zero() {
let mut stats = SessionStats::default();
stats.record_turn(&u(1_000_000), Some(&priced()));
stats.record_turn(&u(1_000_000), Some(&CostConfig::default()));
assert_eq!(stats.cost_usd, Some(1.0));
assert!(!stats.is_unpriced());
}
#[test]
fn record_turn_sums_priced_turns_and_ignores_empty_unpriced_ones() {
let mut stats = SessionStats::default();
stats.record_turn(&Usage::default(), None);
assert_eq!(stats.cost_usd, None);
assert!(!stats.is_unpriced(), "nothing spent is not unpriced");
stats.record_turn(&u(1_000_000), Some(&priced()));
stats.record_turn(&u(2_000_000), Some(&priced()));
assert_eq!(stats.cost_usd, Some(3.0));
assert!(!stats.is_unpriced());
}
#[test]
fn usage_sums_saturate_instead_of_overflowing() {
let mut a = SubAgentSpend {
usage: u(u64::MAX - 1),
cost_usd: Some(1.0),
runs: u32::MAX,
};
a.merge(&a.clone());
assert_eq!(a.usage.input, u64::MAX);
assert_eq!(a.runs, u32::MAX);
}
}
#[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: crate::rt::MaybeSend + crate::rt::MaybeSync {
fn filter(&self, text: &str) -> FilterResult;
fn as_async(&self) -> Option<&dyn AsyncInputFilter> {
None
}
}
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
pub trait AsyncInputFilter: crate::rt::MaybeSend + crate::rt::MaybeSync {
async fn filter(&self, text: &str) -> FilterResult;
}
pub struct AsyncFilter<F>(F);
impl<F: AsyncInputFilter> AsyncFilter<F> {
pub fn new(filter: F) -> Self {
Self(filter)
}
}
impl<F: AsyncInputFilter> InputFilter for AsyncFilter<F> {
fn filter(&self, _text: &str) -> FilterResult {
FilterResult::Reject(
"an async input filter was called synchronously; await it via InputFilter::as_async"
.into(),
)
}
fn as_async(&self) -> Option<&dyn AsyncInputFilter> {
Some(&self.0)
}
}
pub fn is_loop_injected(text: &str) -> bool {
[
crate::context::SUMMARY_PREFIX,
crate::context::COMPACTION_MARKER,
crate::llm_compaction::SUMMARY_MARKER,
crate::agent_loop::AGENT_STOPPED_PREFIX,
crate::agent_loop::LOOP_ABORT_PREFIX,
crate::agent_loop::LOOP_NUDGE_PREFIX,
]
.iter()
.any(|prefix| text.starts_with(prefix))
}
fn text_blocks(content: &[Content]) -> Option<String> {
let text: Vec<&str> = content
.iter()
.filter_map(|c| match c {
Content::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
_ => None,
})
.collect();
(!text.is_empty()).then(|| text.join("\n"))
}
fn real_user_text(m: &Message) -> Option<String> {
match m {
Message::User { content, .. } if is_real_user(m) => text_blocks(content),
_ => None,
}
}
fn is_real_user(m: &Message) -> bool {
match m {
Message::User { content, .. } => !content
.iter()
.find_map(|c| match c {
Content::Text { text } => Some(text.as_str()),
_ => None,
})
.is_some_and(is_loop_injected),
_ => false,
}
}
fn is_compaction_boundary(m: &Message) -> bool {
let Message::User { content, .. } = m else {
return false;
};
content.iter().any(|c| {
matches!(c, Content::Text { text } if
text.starts_with(crate::context::COMPACTION_MARKER)
|| text.starts_with(crate::llm_compaction::SUMMARY_MARKER)
|| text.starts_with(crate::context::SUMMARY_PREFIX))
})
}
const SHORT_REPLY_CHARS: usize = 40;
pub(crate) fn ends_with_question(text: &str) -> bool {
let mut prose = String::new();
let mut in_fence = false;
for line in text.lines() {
if line.trim_start().starts_with("```") {
in_fence = !in_fence;
continue;
}
if in_fence {
continue;
}
for (i, part) in line.split('`').enumerate() {
if i % 2 == 0 {
prose.push_str(part);
}
}
prose.push('\n');
}
let trimmed = prose
.trim_end_matches(|c: char| c.is_whitespace() || matches!(c, '*' | '_' | ')' | '"' | '\''));
trimmed.ends_with('?')
}
pub(crate) fn latest_user_text_of(messages: &[&Message]) -> Option<String> {
let start = messages
.iter()
.rposition(|m| is_compaction_boundary(m))
.map_or(0, |b| b + 1);
messages[start..]
.iter()
.rev()
.find(|m| is_real_user(m))
.and_then(|m| real_user_text(m))
}
pub(crate) fn user_request_of(messages: &[&Message], run_prompts: &[Message]) -> Option<String> {
user_request_parts_of(messages, run_prompts).map(|p| p.to_prose())
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct UserRequestParts {
pub latest: String,
pub reply: Option<ReplyContext>,
pub source: UserRequestSource,
pub run_prompts: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct ReplyContext {
pub question: String,
pub earlier_request: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum UserRequestSource {
Conversation,
RunPrompts,
}
impl UserRequestParts {
fn to_prose(&self) -> String {
if self.source == UserRequestSource::RunPrompts && self.run_prompts.len() > 1 {
return self
.run_prompts
.iter()
.map(|t| format!("User: {t}"))
.collect::<Vec<_>>()
.join("\n\n");
}
let Some(reply) = &self.reply else {
return self.latest.clone();
};
let mut parts = Vec::new();
if let Some(e) = &reply.earlier_request {
parts.push(format!("Earlier user request: {e}"));
}
parts.push(format!("Assistant: {}", reply.question));
parts.push(format!("Latest user message: {}", self.latest));
parts.join("\n\n")
}
}
pub(crate) fn user_request_parts_of(
messages: &[&Message],
run_prompts: &[Message],
) -> Option<UserRequestParts> {
let from_prompts = || {
let texts: Vec<String> = run_prompts.iter().filter_map(real_user_text).collect();
texts.last().cloned().map(|latest| UserRequestParts {
latest,
reply: None,
source: UserRequestSource::RunPrompts,
run_prompts: texts,
})
};
let start = messages
.iter()
.rposition(|m| is_compaction_boundary(m))
.map_or(0, |b| b + 1);
let window = &messages[start..];
let Some(at) = window.iter().rposition(|m| is_real_user(m)) else {
return from_prompts();
};
let Some(latest) = real_user_text(window[at]) else {
return from_prompts();
};
let short = latest.chars().count() < SHORT_REPLY_CHARS;
let mut parts = UserRequestParts {
latest,
reply: None,
source: UserRequestSource::Conversation,
run_prompts: Vec::new(),
};
if !short {
return Some(parts);
}
let mut asked: Option<(usize, String)> = None;
for (i, m) in window[..at].iter().enumerate().rev() {
if is_real_user(m) {
break;
}
if let Message::Assistant { content, .. } = m {
if let Some(text) = text_blocks(content) {
asked = Some((i, text));
break;
}
}
}
let Some((asked_at, question)) = asked.filter(|(_, a)| ends_with_question(a)) else {
return Some(parts);
};
let earlier_request = window[..asked_at]
.iter()
.rev()
.find(|m| is_real_user(m))
.and_then(|m| real_user_text(m));
parts.reply = Some(ReplyContext {
question,
earlier_request,
});
Some(parts)
}
#[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,
pub messages: &'a [AgentMessage],
pub run_prompts: &'a [Message],
}
impl<'a> ToolCallRequest<'a> {
pub fn new(tool_call_id: &'a str, tool_name: &'a str, args: &'a serde_json::Value) -> Self {
Self {
tool_call_id,
tool_name,
args,
messages: &[],
run_prompts: &[],
}
}
pub fn with_messages(mut self, messages: &'a [AgentMessage]) -> Self {
self.messages = messages;
self
}
pub fn with_run_prompts(mut self, run_prompts: &'a [Message]) -> Self {
self.run_prompts = run_prompts;
self
}
fn llm_messages(&self) -> Vec<&Message> {
self.messages
.iter()
.filter_map(AgentMessage::as_llm)
.collect()
}
pub fn latest_user_text(&self) -> Option<String> {
latest_user_text_of(&self.llm_messages())
}
pub fn user_request(&self) -> Option<String> {
user_request_of(&self.llm_messages(), self.run_prompts)
}
pub fn user_request_parts(&self) -> Option<UserRequestParts> {
user_request_parts_of(&self.llm_messages(), self.run_prompts)
}
}
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
pub trait ToolMiddleware: crate::rt::MaybeSend + crate::rt::MaybeSync {
async fn before_tool(&self, call: &ToolCallRequest<'_>) -> ToolDecision;
}
#[derive(Debug)]
#[non_exhaustive]
pub struct TurnContext<'a> {
pub system_prompt: &'a str,
pub messages: &'a [Message],
pub tools: &'a [crate::provider::ToolDefinition],
pub model: &'a str,
pub run_prompts: &'a [Message],
}
impl<'a> TurnContext<'a> {
pub fn new(
system_prompt: &'a str,
messages: &'a [Message],
tools: &'a [crate::provider::ToolDefinition],
model: &'a str,
) -> Self {
Self {
system_prompt,
messages,
tools,
model,
run_prompts: &[],
}
}
pub fn with_run_prompts(mut self, run_prompts: &'a [Message]) -> Self {
self.run_prompts = run_prompts;
self
}
pub fn latest_user_text(&self) -> Option<String> {
latest_user_text_of(&self.messages.iter().collect::<Vec<_>>())
}
pub fn user_request(&self) -> Option<String> {
user_request_of(&self.messages.iter().collect::<Vec<_>>(), self.run_prompts)
}
pub fn user_request_parts(&self) -> Option<UserRequestParts> {
user_request_parts_of(&self.messages.iter().collect::<Vec<_>>(), self.run_prompts)
}
}
#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
pub trait TurnHook: crate::rt::MaybeSend + crate::rt::MaybeSync {
async fn before_turn(&self, turn: &TurnContext<'_>) -> Option<String>;
}
pub fn now_ms() -> u64 {
web_time::SystemTime::now()
.duration_since(web_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,
sub_agents: SubAgentSpend {
usage: Usage {
input: 50,
output: 60,
cache_read: 70,
cache_write: 80,
total_tokens: 0,
},
cost_usd: Some(0.2),
runs: 2,
},
decision: DecisionStats {
requests: 3,
failures: 1,
timeouts: 1,
usage: Usage {
input: 900,
..Default::default()
},
cost_usd: Some(0.0000378),
unpriced: 0,
},
},
),
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::ProviderRetry { .. } => "providerRetry"
= AgentEvent::provider_retry(1, 4, "rate limited", std::time::Duration::from_millis(1500)),
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);
}
}
}
#[cfg(test)]
mod thinking_level_tests {
use super::ThinkingLevel;
#[test]
fn serde_names_are_lowercase_and_old_values_still_load() {
for (level, name) in [
(ThinkingLevel::Off, "off"),
(ThinkingLevel::Minimal, "minimal"),
(ThinkingLevel::Low, "low"),
(ThinkingLevel::Medium, "medium"),
(ThinkingLevel::High, "high"),
(ThinkingLevel::XHigh, "xhigh"),
(ThinkingLevel::Max, "max"),
] {
let json = serde_json::to_string(&level).unwrap();
assert_eq!(json, format!("\"{name}\""));
let back: ThinkingLevel = serde_json::from_str(&json).unwrap();
assert_eq!(back, level);
}
}
}