1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
/// Agent state management
use crate::types::{StopReason, ToolResult};
use oxi_ai::{ContentBlock, Message, TextContent};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
/// Agent execution state
///
/// Tracks the full lifecycle of an agent conversation including messages,
/// token usage, tool results, and iteration progress.
///
/// Derives `Serialize`/`Deserialize` for session persistence and
/// cross-process state transfer (e.g. oxios supervisor serialization).
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct AgentState {
/// Conversation message history (user, assistant, and tool-result messages).
pub messages: Vec<Message>,
/// Current agent loop iteration (incremented after each assistant turn).
pub iteration: usize,
/// The reason the last turn stopped, if any.
pub stop_reason: Option<StopReason>,
/// Accumulated results from tool executions in the current conversation.
pub tool_results: Vec<ToolResult>,
/// Cumulative token count (input + output) across all turns.
pub total_tokens: usize,
/// Cumulative prompt / input tokens across all turns.
pub input_tokens: usize,
/// Cumulative completion / output tokens across all turns.
pub output_tokens: usize,
}
impl AgentState {
/// Create a new, default-initialized agent state.
pub fn new() -> Self {
Self::default()
}
/// Add a user message
pub fn add_user_message(&mut self, content: String) {
self.messages
.push(Message::User(oxi_ai::UserMessage::new(content)));
}
/// Add an assistant message
pub fn add_assistant_message(&mut self, content: String) {
let mut assistant =
oxi_ai::AssistantMessage::new(oxi_ai::Api::AnthropicMessages, "agent", "agent-model");
assistant.content = vec![ContentBlock::Text(TextContent::new(content))];
self.messages.push(Message::Assistant(assistant));
}
/// Add a tool result message to both the message history and the tool results list.
pub fn add_tool_result(&mut self, tool_call_id: String, content: String) {
let content_for_result = content.clone();
let tool_result_msg = oxi_ai::ToolResultMessage::new(
tool_call_id.clone(),
"tool",
vec![ContentBlock::Text(TextContent::new(content))],
);
self.messages
.push(oxi_ai::Message::ToolResult(tool_result_msg));
self.tool_results
.push(ToolResult::success(tool_call_id, content_for_result));
}
/// Increment the iteration counter after an assistant turn completes.
pub fn increment_iteration(&mut self) {
self.iteration += 1;
}
/// Record the reason the last turn stopped.
pub fn set_stop_reason(&mut self, reason: StopReason) {
self.stop_reason = Some(reason);
}
/// Accumulate token usage from a completed LLM call.
pub fn record_usage(&mut self, input: usize, output: usize) {
self.input_tokens += input;
self.output_tokens += output;
self.total_tokens += input + output;
}
/// Clear all state, resetting for a new conversation.
pub fn clear(&mut self) {
self.messages.clear();
self.iteration = 0;
self.stop_reason = None;
self.tool_results.clear();
self.total_tokens = 0;
self.input_tokens = 0;
self.output_tokens = 0;
}
/// Replace the entire message history (used after context compaction).
pub fn replace_messages(&mut self, messages: Vec<Message>) {
self.messages = messages;
}
/// Rough token-count estimate based on the serialized message JSON length.
pub fn estimate_tokens(&self) -> usize {
let json = serde_json::to_string(&self.messages).unwrap_or_default();
json.len() / 4 // Rough approximation
}
/// Returns `true` if the agent has signaled a stop reason.
pub fn is_complete(&self) -> bool {
self.stop_reason.is_some()
}
}
/// Thread-safe agent state wrapper.
#[derive(Default, Clone)]
pub struct SharedState {
state: Arc<RwLock<AgentState>>,
}
impl SharedState {
/// Create a new SharedState with default (empty) agent state.
pub fn new() -> Self {
Self::default()
}
/// Obtain a snapshot of the current agent state.
pub fn get_state(&self) -> AgentState {
self.state.read().clone()
}
/// Mutably update the agent state under a write lock.
pub fn update<F>(&self, f: F)
where
F: FnOnce(&mut AgentState),
{
let mut state = self.state.write();
f(&mut state);
}
/// Reset the state for a new conversation (delegates to [`AgentState::clear`]).
pub fn reset(&self) {
let mut state = self.state.write();
state.clear();
}
}