use crate::llm::ContentBlock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use time::OffsetDateTime;
use uuid::Uuid;
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ThreadId(pub String);
impl ThreadId {
#[must_use]
pub fn new() -> Self {
Self(Uuid::new_v4().to_string())
}
#[must_use]
pub fn from_string(s: impl Into<String>) -> Self {
Self(s.into())
}
}
impl Default for ThreadId {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Display for ThreadId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Clone, Debug)]
pub struct AgentConfig {
pub max_turns: Option<usize>,
pub max_tokens: Option<u32>,
pub system_prompt: String,
pub model: String,
pub retry: RetryConfig,
pub streaming: bool,
}
impl Default for AgentConfig {
fn default() -> Self {
Self {
max_turns: None,
max_tokens: None,
system_prompt: String::new(),
model: String::from("claude-sonnet-4-5-20250929"),
retry: RetryConfig::default(),
streaming: false,
}
}
}
#[derive(Clone, Debug)]
pub struct RetryConfig {
pub max_retries: u32,
pub base_delay_ms: u64,
pub max_delay_ms: u64,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_retries: 5,
base_delay_ms: 1000,
max_delay_ms: 120_000,
}
}
}
impl RetryConfig {
#[must_use]
pub const fn no_retry() -> Self {
Self {
max_retries: 0,
base_delay_ms: 0,
max_delay_ms: 0,
}
}
#[must_use]
pub const fn fast() -> Self {
Self {
max_retries: 5,
base_delay_ms: 10,
max_delay_ms: 100,
}
}
}
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct TokenUsage {
pub input_tokens: u32,
pub output_tokens: u32,
}
impl TokenUsage {
pub const fn add(&mut self, other: &Self) {
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ToolResult {
pub success: bool,
pub output: String,
pub data: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub documents: Vec<crate::llm::ContentSource>,
pub duration_ms: Option<u64>,
}
impl ToolResult {
#[must_use]
pub fn success(output: impl Into<String>) -> Self {
Self {
success: true,
output: output.into(),
data: None,
documents: Vec::new(),
duration_ms: None,
}
}
#[must_use]
pub fn success_with_data(output: impl Into<String>, data: serde_json::Value) -> Self {
Self {
success: true,
output: output.into(),
data: Some(data),
documents: Vec::new(),
duration_ms: None,
}
}
#[must_use]
pub fn error(message: impl Into<String>) -> Self {
Self {
success: false,
output: message.into(),
data: None,
documents: Vec::new(),
duration_ms: None,
}
}
#[must_use]
pub const fn with_duration(mut self, duration_ms: u64) -> Self {
self.duration_ms = Some(duration_ms);
self
}
#[must_use]
pub fn with_documents(mut self, documents: Vec<crate::llm::ContentSource>) -> Self {
self.documents = documents;
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ToolTier {
Observe,
Confirm,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AgentState {
pub thread_id: ThreadId,
pub turn_count: usize,
pub total_usage: TokenUsage,
pub metadata: HashMap<String, serde_json::Value>,
#[serde(with = "time::serde::rfc3339")]
pub created_at: OffsetDateTime,
}
impl AgentState {
#[must_use]
pub fn new(thread_id: ThreadId) -> Self {
Self {
thread_id,
turn_count: 0,
total_usage: TokenUsage::default(),
metadata: HashMap::new(),
created_at: OffsetDateTime::now_utc(),
}
}
}
#[derive(Debug, Clone)]
pub struct AgentError {
pub message: String,
pub recoverable: bool,
}
impl AgentError {
#[must_use]
pub fn new(message: impl Into<String>, recoverable: bool) -> Self {
Self {
message: message.into(),
recoverable,
}
}
}
impl std::fmt::Display for AgentError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.message)
}
}
impl std::error::Error for AgentError {}
#[derive(Debug)]
pub enum AgentRunState {
Done {
total_turns: u32,
input_tokens: u64,
output_tokens: u64,
},
Refusal {
total_turns: u32,
input_tokens: u64,
output_tokens: u64,
},
Error(AgentError),
AwaitingConfirmation {
tool_call_id: String,
tool_name: String,
display_name: String,
input: serde_json::Value,
description: String,
continuation: Box<AgentContinuation>,
},
Cancelled {
total_turns: u32,
input_tokens: u64,
output_tokens: u64,
},
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct PendingToolCallInfo {
pub id: String,
pub name: String,
pub display_name: String,
pub input: serde_json::Value,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub listen_context: Option<ListenExecutionContext>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ListenExecutionContext {
pub operation_id: String,
pub revision: u64,
pub snapshot: serde_json::Value,
#[serde(
default,
skip_serializing_if = "Option::is_none",
with = "time::serde::rfc3339::option"
)]
pub expires_at: Option<OffsetDateTime>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AgentContinuation {
pub thread_id: ThreadId,
pub turn: usize,
pub total_usage: TokenUsage,
pub turn_usage: TokenUsage,
pub pending_tool_calls: Vec<PendingToolCallInfo>,
pub awaiting_index: usize,
pub completed_results: Vec<(String, ToolResult)>,
pub state: AgentState,
}
#[derive(Debug)]
pub enum AgentInput {
Text(String),
Message(Vec<ContentBlock>),
Resume {
continuation: Box<AgentContinuation>,
tool_call_id: String,
confirmed: bool,
rejection_reason: Option<String>,
},
Continue,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum ToolOutcome {
Success(ToolResult),
Failed(ToolResult),
InProgress {
operation_id: String,
message: String,
},
}
impl ToolOutcome {
#[must_use]
pub fn success(output: impl Into<String>) -> Self {
Self::Success(ToolResult::success(output))
}
#[must_use]
pub fn failed(message: impl Into<String>) -> Self {
Self::Failed(ToolResult::error(message))
}
#[must_use]
pub fn in_progress(operation_id: impl Into<String>, message: impl Into<String>) -> Self {
Self::InProgress {
operation_id: operation_id.into(),
message: message.into(),
}
}
#[must_use]
pub const fn is_in_progress(&self) -> bool {
matches!(self, Self::InProgress { .. })
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ExecutionStatus {
InFlight,
Completed,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ToolExecution {
pub tool_call_id: String,
pub thread_id: ThreadId,
pub tool_name: String,
pub display_name: String,
pub input: serde_json::Value,
pub status: ExecutionStatus,
pub result: Option<ToolResult>,
pub operation_id: Option<String>,
#[serde(with = "time::serde::rfc3339")]
pub started_at: OffsetDateTime,
#[serde(with = "time::serde::rfc3339::option")]
pub completed_at: Option<OffsetDateTime>,
}
impl ToolExecution {
#[must_use]
pub fn new_in_flight(
tool_call_id: impl Into<String>,
thread_id: ThreadId,
tool_name: impl Into<String>,
display_name: impl Into<String>,
input: serde_json::Value,
started_at: OffsetDateTime,
) -> Self {
Self {
tool_call_id: tool_call_id.into(),
thread_id,
tool_name: tool_name.into(),
display_name: display_name.into(),
input,
status: ExecutionStatus::InFlight,
result: None,
operation_id: None,
started_at,
completed_at: None,
}
}
pub fn complete(&mut self, result: ToolResult) {
self.status = ExecutionStatus::Completed;
self.result = Some(result);
self.completed_at = Some(OffsetDateTime::now_utc());
}
pub fn set_operation_id(&mut self, operation_id: impl Into<String>) {
self.operation_id = Some(operation_id.into());
}
#[must_use]
pub fn is_in_flight(&self) -> bool {
self.status == ExecutionStatus::InFlight
}
#[must_use]
pub fn is_completed(&self) -> bool {
self.status == ExecutionStatus::Completed
}
}
#[derive(Debug)]
pub enum TurnOutcome {
NeedsMoreTurns {
turn: usize,
turn_usage: TokenUsage,
total_usage: TokenUsage,
},
Done {
total_turns: u32,
input_tokens: u64,
output_tokens: u64,
},
AwaitingConfirmation {
tool_call_id: String,
tool_name: String,
display_name: String,
input: serde_json::Value,
description: String,
continuation: Box<AgentContinuation>,
},
Refusal {
total_turns: u32,
input_tokens: u64,
output_tokens: u64,
},
Cancelled {
total_turns: u32,
input_tokens: u64,
output_tokens: u64,
},
Error(AgentError),
}