rab/agent/
types.rs

1use chrono::Utc;
2use serde::{Deserialize, Serialize};
3
4/// Whether the agent executes tool calls in parallel (default) or sequentially.
5#[derive(Debug, Clone, Copy, PartialEq, Default)]
6pub enum ToolExecutionMode {
7    /// Execute all tool calls concurrently after sequential preflight.
8    #[default]
9    Parallel,
10    /// Execute tool calls one at a time in order.
11    Sequential,
12}
13
14/// How queued messages are drained from a pending message queue.
15#[derive(Debug, Clone, Copy, PartialEq)]
16pub enum QueueMode {
17    /// Drain all queued messages at once.
18    All,
19    /// Drain one message at a time.
20    OneAtATime,
21}
22
23/// A pending message queue with a configurable drain mode.
24/// Used for steering (mid-stream) and follow-up (post-agent) message delivery.
25#[derive(Debug)]
26pub struct PendingMessageQueue {
27    messages: Vec<AgentMessage>,
28    mode: QueueMode,
29}
30
31impl PendingMessageQueue {
32    pub fn new(mode: QueueMode) -> Self {
33        Self {
34            messages: Vec::new(),
35            mode,
36        }
37    }
38
39    /// Add a message to the back of the queue.
40    pub fn enqueue(&mut self, msg: AgentMessage) {
41        self.messages.push(msg);
42    }
43
44    /// Drain messages according to the current mode.
45    pub fn drain(&mut self) -> Vec<AgentMessage> {
46        match self.mode {
47            QueueMode::All => self.messages.drain(..).collect(),
48            QueueMode::OneAtATime => {
49                if self.messages.is_empty() {
50                    vec![]
51                } else {
52                    vec![self.messages.remove(0)]
53                }
54            }
55        }
56    }
57
58    /// Drain all messages regardless of mode.
59    /// Used for dequeue operations that need to restore all messages.
60    pub fn drain_all(&mut self) -> Vec<AgentMessage> {
61        self.messages.drain(..).collect()
62    }
63
64    pub fn is_empty(&self) -> bool {
65        self.messages.is_empty()
66    }
67
68    pub fn len(&self) -> usize {
69        self.messages.len()
70    }
71
72    pub fn clear(&mut self) {
73        self.messages.clear();
74    }
75}
76
77/// Role of a message in the conversation.
78#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
79#[serde(rename_all = "camelCase")]
80pub enum Role {
81    User,
82    Assistant,
83    ToolResult,
84}
85
86/// A tool call requested by the assistant.
87#[derive(Debug, Clone, Serialize, Deserialize)]
88pub struct ToolCall {
89    pub id: String,
90    pub name: String,
91    pub arguments: serde_json::Value,
92}
93
94/// Token usage information for an assistant response.
95#[derive(Debug, Clone, Default, Serialize, Deserialize)]
96pub struct Usage {
97    pub input_tokens: Option<i32>,
98    pub output_tokens: Option<i32>,
99    pub cache_tokens: Option<i32>,
100}
101
102/// A universal message type in the conversation.
103#[derive(Debug, Clone, Serialize, Deserialize)]
104#[serde(rename_all = "camelCase")]
105pub struct AgentMessage {
106    pub id: String,
107    pub parent_id: Option<String>,
108    pub role: Role,
109    pub content: String,
110    #[serde(default, skip_serializing_if = "Vec::is_empty")]
111    pub tool_calls: Vec<ToolCall>,
112    #[serde(skip_serializing_if = "Option::is_none")]
113    pub tool_call_id: Option<String>,
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub usage: Option<Usage>,
116    pub is_error: bool,
117    pub timestamp: i64,
118}
119
120impl AgentMessage {
121    pub fn user(content: impl Into<String>) -> Self {
122        Self {
123            id: uuid::Uuid::new_v4().to_string(),
124            parent_id: None,
125            role: Role::User,
126            content: content.into(),
127            tool_calls: vec![],
128            tool_call_id: None,
129            usage: None,
130            is_error: false,
131            timestamp: Utc::now().timestamp_millis(),
132        }
133    }
134
135    pub fn tool_result(
136        tool_call_id: impl Into<String>,
137        content: impl Into<String>,
138        is_error: bool,
139    ) -> Self {
140        Self {
141            id: uuid::Uuid::new_v4().to_string(),
142            parent_id: None,
143            role: Role::ToolResult,
144            content: content.into(),
145            tool_calls: vec![],
146            tool_call_id: Some(tool_call_id.into()),
147            usage: None,
148            is_error,
149            timestamp: Utc::now().timestamp_millis(),
150        }
151    }
152}
rab/agent/types.rs

rab/agent/
types.rs
