use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use futures_timer::Delay;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
macro_rules! id_newtype {
($(#[$meta:meta])* $name:ident) => {
$(#[$meta])*
#[derive(
Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize,
)]
pub struct $name(pub String);
impl $name {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
}
impl fmt::Display for $name {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
impl From<&str> for $name {
fn from(value: &str) -> Self {
Self::new(value)
}
}
impl From<String> for $name {
fn from(value: String) -> Self {
Self(value)
}
}
};
}
id_newtype!(
SessionId
);
id_newtype!(
TurnId
);
id_newtype!(
MessageId
);
id_newtype!(
ToolCallId
);
id_newtype!(
ToolResultId
);
id_newtype!(
TaskId
);
id_newtype!(
ProviderMessageId
);
id_newtype!(
ArtifactId
);
id_newtype!(
PartId
);
id_newtype!(
ApprovalId
);
pub type MetadataMap = BTreeMap<String, Value>;
#[derive(Default)]
struct CancellationState {
generation: AtomicU64,
}
#[derive(Clone, Default)]
pub struct CancellationController {
state: Arc<CancellationState>,
}
#[derive(Clone, Default)]
pub struct CancellationHandle {
state: Arc<CancellationState>,
}
#[derive(Clone, Default)]
pub struct TurnCancellation {
handle: CancellationHandle,
generation: u64,
}
impl CancellationController {
pub fn new() -> Self {
Self::default()
}
pub fn handle(&self) -> CancellationHandle {
CancellationHandle {
state: Arc::clone(&self.state),
}
}
pub fn interrupt(&self) -> u64 {
self.state.generation.fetch_add(1, Ordering::SeqCst) + 1
}
}
impl CancellationHandle {
pub fn generation(&self) -> u64 {
self.state.generation.load(Ordering::SeqCst)
}
pub fn checkpoint(&self) -> TurnCancellation {
TurnCancellation {
handle: self.clone(),
generation: self.generation(),
}
}
pub fn is_cancelled_since(&self, generation: u64) -> bool {
self.generation() != generation
}
pub async fn cancelled_since(&self, generation: u64) {
while !self.is_cancelled_since(generation) {
Delay::new(Duration::from_millis(10)).await;
}
}
}
impl TurnCancellation {
pub fn new(handle: CancellationHandle) -> Self {
handle.checkpoint()
}
pub fn generation(&self) -> u64 {
self.generation
}
pub fn is_cancelled(&self) -> bool {
self.handle.is_cancelled_since(self.generation)
}
pub async fn cancelled(&self) {
self.handle.cancelled_since(self.generation).await;
}
pub fn handle(&self) -> &CancellationHandle {
&self.handle
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Item {
pub id: Option<MessageId>,
pub kind: ItemKind,
pub parts: Vec<Part>,
pub metadata: MetadataMap,
}
impl Item {
pub fn new(kind: ItemKind, parts: Vec<Part>) -> Self {
Self {
id: None,
kind,
parts,
metadata: MetadataMap::new(),
}
}
pub fn text(kind: ItemKind, text: impl Into<String>) -> Self {
Self::new(kind, vec![Part::Text(TextPart::new(text))])
}
pub fn with_id(mut self, id: impl Into<MessageId>) -> Self {
self.id = Some(id.into());
self
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
pub fn push_part(mut self, part: Part) -> Self {
self.parts.push(part);
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub enum ItemKind {
System,
Developer,
User,
Assistant,
Tool,
Context,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum Part {
Text(TextPart),
Media(MediaPart),
File(FilePart),
Structured(StructuredPart),
Reasoning(ReasoningPart),
ToolCall(ToolCallPart),
ToolResult(ToolResultPart),
Custom(CustomPart),
}
impl Part {
pub fn text(text: impl Into<String>) -> Self {
Self::Text(TextPart::new(text))
}
pub fn media(modality: Modality, mime_type: impl Into<String>, data: DataRef) -> Self {
Self::Media(MediaPart::new(modality, mime_type, data))
}
pub fn file(data: DataRef) -> Self {
Self::File(FilePart::new(data))
}
pub fn structured(value: Value) -> Self {
Self::Structured(StructuredPart::new(value))
}
pub fn reasoning(summary: impl Into<String>) -> Self {
Self::Reasoning(ReasoningPart::summary(summary))
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum PartKind {
Text,
Media,
File,
Structured,
Reasoning,
ToolCall,
ToolResult,
Custom,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TextPart {
pub text: String,
pub metadata: MetadataMap,
}
impl TextPart {
pub fn new(text: impl Into<String>) -> Self {
Self {
text: text.into(),
metadata: MetadataMap::new(),
}
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct MediaPart {
pub modality: Modality,
pub mime_type: String,
pub data: DataRef,
pub metadata: MetadataMap,
}
impl MediaPart {
pub fn new(modality: Modality, mime_type: impl Into<String>, data: DataRef) -> Self {
Self {
modality,
mime_type: mime_type.into(),
data,
metadata: MetadataMap::new(),
}
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum Modality {
Audio,
Image,
Video,
Binary,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum DataRef {
InlineText(String),
InlineBytes(Vec<u8>),
Uri(String),
Handle(ArtifactId),
}
impl DataRef {
pub fn inline_text(text: impl Into<String>) -> Self {
Self::InlineText(text.into())
}
pub fn inline_bytes(bytes: impl Into<Vec<u8>>) -> Self {
Self::InlineBytes(bytes.into())
}
pub fn uri(uri: impl Into<String>) -> Self {
Self::Uri(uri.into())
}
pub fn handle(id: impl Into<ArtifactId>) -> Self {
Self::Handle(id.into())
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct FilePart {
pub name: Option<String>,
pub mime_type: Option<String>,
pub data: DataRef,
pub metadata: MetadataMap,
}
impl FilePart {
pub fn new(data: DataRef) -> Self {
Self {
name: None,
mime_type: None,
data,
metadata: MetadataMap::new(),
}
}
pub fn named(name: impl Into<String>, data: DataRef) -> Self {
Self::new(data).with_name(name)
}
pub fn with_name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
pub fn with_mime_type(mut self, mime_type: impl Into<String>) -> Self {
self.mime_type = Some(mime_type.into());
self
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct StructuredPart {
pub value: Value,
pub schema: Option<Value>,
pub metadata: MetadataMap,
}
impl StructuredPart {
pub fn new(value: Value) -> Self {
Self {
value,
schema: None,
metadata: MetadataMap::new(),
}
}
pub fn with_schema(mut self, schema: Value) -> Self {
self.schema = Some(schema);
self
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ReasoningPart {
pub summary: Option<String>,
pub data: Option<DataRef>,
pub redacted: bool,
pub metadata: MetadataMap,
}
impl ReasoningPart {
pub fn summary(summary: impl Into<String>) -> Self {
Self {
summary: Some(summary.into()),
data: None,
redacted: false,
metadata: MetadataMap::new(),
}
}
pub fn redacted_summary(summary: impl Into<String>) -> Self {
Self::summary(summary).with_redacted(true)
}
pub fn with_data(mut self, data: DataRef) -> Self {
self.data = Some(data);
self
}
pub fn with_redacted(mut self, redacted: bool) -> Self {
self.redacted = redacted;
self
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ToolCallPart {
pub id: ToolCallId,
pub name: String,
pub input: Value,
pub metadata: MetadataMap,
}
impl ToolCallPart {
pub fn new(id: impl Into<ToolCallId>, name: impl Into<String>, input: Value) -> Self {
Self {
id: id.into(),
name: name.into(),
input,
metadata: MetadataMap::new(),
}
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ToolResultPart {
pub call_id: ToolCallId,
pub output: ToolOutput,
pub is_error: bool,
pub metadata: MetadataMap,
}
impl ToolResultPart {
pub fn success(call_id: impl Into<ToolCallId>, output: ToolOutput) -> Self {
Self {
call_id: call_id.into(),
output,
is_error: false,
metadata: MetadataMap::new(),
}
}
pub fn error(call_id: impl Into<ToolCallId>, output: ToolOutput) -> Self {
Self::success(call_id, output).with_is_error(true)
}
pub fn with_is_error(mut self, is_error: bool) -> Self {
self.is_error = is_error;
self
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum ToolOutput {
Text(String),
Structured(Value),
Parts(Vec<Part>),
Files(Vec<FilePart>),
}
impl ToolOutput {
pub fn text(text: impl Into<String>) -> Self {
Self::Text(text.into())
}
pub fn structured(value: Value) -> Self {
Self::Structured(value)
}
pub fn parts(parts: Vec<Part>) -> Self {
Self::Parts(parts)
}
pub fn files(files: Vec<FilePart>) -> Self {
Self::Files(files)
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CustomPart {
pub kind: String,
pub data: Option<DataRef>,
pub value: Option<Value>,
pub metadata: MetadataMap,
}
impl CustomPart {
pub fn new(kind: impl Into<String>) -> Self {
Self {
kind: kind.into(),
data: None,
value: None,
metadata: MetadataMap::new(),
}
}
pub fn with_data(mut self, data: DataRef) -> Self {
self.data = Some(data);
self
}
pub fn with_value(mut self, value: Value) -> Self {
self.value = Some(value);
self
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum Delta {
BeginPart {
part_id: PartId,
kind: PartKind,
},
AppendText {
part_id: PartId,
chunk: String,
},
AppendBytes {
part_id: PartId,
chunk: Vec<u8>,
},
ReplaceStructured {
part_id: PartId,
value: Value,
},
SetMetadata {
part_id: PartId,
metadata: MetadataMap,
},
CommitPart {
part: Part,
},
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct Usage {
pub tokens: Option<TokenUsage>,
pub cost: Option<CostUsage>,
pub metadata: MetadataMap,
}
impl Usage {
pub fn new(tokens: TokenUsage) -> Self {
Self {
tokens: Some(tokens),
cost: None,
metadata: MetadataMap::new(),
}
}
pub fn with_cost(mut self, cost: CostUsage) -> Self {
self.cost = Some(cost);
self
}
pub fn with_metadata(mut self, metadata: MetadataMap) -> Self {
self.metadata = metadata;
self
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenUsage {
pub input_tokens: u64,
pub output_tokens: u64,
pub reasoning_tokens: Option<u64>,
pub cached_input_tokens: Option<u64>,
pub cache_write_input_tokens: Option<u64>,
}
impl TokenUsage {
pub fn new(input_tokens: u64, output_tokens: u64) -> Self {
Self {
input_tokens,
output_tokens,
reasoning_tokens: None,
cached_input_tokens: None,
cache_write_input_tokens: None,
}
}
pub fn with_reasoning_tokens(mut self, reasoning_tokens: u64) -> Self {
self.reasoning_tokens = Some(reasoning_tokens);
self
}
pub fn with_cached_input_tokens(mut self, cached_input_tokens: u64) -> Self {
self.cached_input_tokens = Some(cached_input_tokens);
self
}
pub fn with_cache_write_input_tokens(mut self, cache_write_input_tokens: u64) -> Self {
self.cache_write_input_tokens = Some(cache_write_input_tokens);
self
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CostUsage {
pub amount: f64,
pub currency: String,
pub provider_amount: Option<String>,
}
impl CostUsage {
pub fn new(amount: f64, currency: impl Into<String>) -> Self {
Self {
amount,
currency: currency.into(),
provider_amount: None,
}
}
pub fn with_provider_amount(mut self, provider_amount: impl Into<String>) -> Self {
self.provider_amount = Some(provider_amount.into());
self
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum FinishReason {
Completed,
ToolCall,
MaxTokens,
Cancelled,
Blocked,
Error,
Other(String),
}
pub trait ItemView {
fn kind(&self) -> ItemKind;
fn parts(&self) -> &[Part];
fn metadata(&self) -> &MetadataMap;
}
impl ItemView for Item {
fn kind(&self) -> ItemKind {
self.kind
}
fn parts(&self) -> &[Part] {
&self.parts
}
fn metadata(&self) -> &MetadataMap {
&self.metadata
}
}
#[derive(Debug, Error)]
pub enum NormalizeError {
#[error("unsupported content shape: {0}")]
Unsupported(String),
}
#[derive(Debug, Error)]
pub enum ProtocolError {
#[error("invalid protocol state: {0}")]
InvalidState(String),
}
#[derive(Debug, Error)]
pub enum AgentError {
#[error(transparent)]
Normalize(#[from] NormalizeError),
#[error(transparent)]
Protocol(#[from] ProtocolError),
}