use chrono::Utc;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use starweaver_core::{ConversationId, Metadata, RunId, TaskId};
use starweaver_model::{ModelMessage, ModelResponse, ToolCallPart, ToolReturnPart};
use starweaver_usage::Usage;
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum RunStatus {
Starting,
Running,
Waiting,
Completed,
Failed,
Cancelled,
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct AgentRunState {
pub run_id: RunId,
pub conversation_id: ConversationId,
pub message_history: Vec<ModelMessage>,
pub usage: Usage,
pub run_step: usize,
pub status: RunStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub structured_output: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub latest_response: Option<ModelResponse>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub pending_tool_calls: Vec<ToolCallPart>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub pending_tool_returns: Vec<ToolReturnPart>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub pending_approval_tool_returns: Vec<ToolReturnPart>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub deferred_tool_returns: Vec<ToolReturnPart>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub idle_messages: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_run_id: Option<RunId>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_task_id: Option<TaskId>,
#[serde(default, skip_serializing_if = "Metadata::is_empty")]
pub metadata: Metadata,
}
impl AgentRunState {
#[must_use]
pub fn new(run_id: RunId, conversation_id: ConversationId) -> Self {
Self {
run_id,
conversation_id,
message_history: Vec::new(),
usage: Usage::default(),
run_step: 0,
status: RunStatus::Starting,
output: None,
structured_output: None,
latest_response: None,
pending_tool_calls: Vec::new(),
pending_tool_returns: Vec::new(),
pending_approval_tool_returns: Vec::new(),
deferred_tool_returns: Vec::new(),
idle_messages: Vec::new(),
parent_run_id: None,
parent_task_id: None,
metadata: Metadata::default(),
}
}
pub fn apply_model_response(&mut self, mut response: ModelResponse) {
response.run_id.get_or_insert_with(|| self.run_id.clone());
response
.conversation_id
.get_or_insert_with(|| self.conversation_id.clone());
response.timestamp.get_or_insert_with(Utc::now);
self.usage.add_assign(&response.usage);
self.message_history
.push(ModelMessage::Response(response.clone()));
self.latest_response = Some(response);
}
pub fn replace_latest_response(&mut self, response: ModelResponse) {
if let Some(ModelMessage::Response(history_response)) = self.message_history.last_mut() {
*history_response = response.clone();
}
self.latest_response = Some(response);
}
#[must_use]
pub const fn has_pending_hitl(&self) -> bool {
!self.pending_approval_tool_returns.is_empty() || !self.deferred_tool_returns.is_empty()
}
#[must_use]
pub fn pending_approvals(&self) -> &[ToolReturnPart] {
&self.pending_approval_tool_returns
}
#[must_use]
pub fn pending_deferred_tools(&self) -> &[ToolReturnPart] {
&self.deferred_tool_returns
}
pub fn pending_hitl_tool_returns(&self) -> impl Iterator<Item = &ToolReturnPart> {
self.pending_approval_tool_returns
.iter()
.chain(self.deferred_tool_returns.iter())
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct AgentRunResult {
pub output: String,
pub state: AgentRunState,
}
impl AgentRunResult {
#[must_use]
pub const fn has_pending_hitl(&self) -> bool {
self.state.has_pending_hitl()
}
#[must_use]
pub fn pending_approvals(&self) -> &[ToolReturnPart] {
self.state.pending_approvals()
}
#[must_use]
pub fn pending_deferred_tools(&self) -> &[ToolReturnPart] {
self.state.pending_deferred_tools()
}
}