use chrono::{DateTime, Utc};
use serde_json::Value;
use crate::provider::{FinishReason, ModelName, ProviderId, TokenUsage};
use super::event::AgentEvent;
use super::run::{RunId, RunRecord, RunStatus};
use super::store::RunEventRecord;
#[derive(Debug, Clone)]
pub struct RunState {
pub run_id: RunId,
pub session_id: uuid::Uuid,
pub status: RunStatus,
pub provider: ProviderId,
pub model: ModelName,
pub metadata: Value,
pub iteration: usize,
pub total_usage: TokenUsage,
pub last_finish: Option<FinishReason>,
pub last_error: Option<String>,
pub event_count: usize,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
impl RunState {
#[must_use]
pub fn create(record: &RunRecord, events: &[RunEventRecord]) -> Self {
let updated_at = events.last().map_or(record.updated_at, |e| e.timestamp);
let mut state = RunState {
run_id: record.id,
session_id: record.session_id,
status: RunStatus::Pending,
provider: record.provider.clone(),
model: record.model.clone(),
metadata: record.metadata.clone(),
iteration: 0,
total_usage: TokenUsage::new(0, 0),
last_finish: None,
last_error: None,
event_count: events.len(),
created_at: record.created_at,
updated_at,
};
for event_record in events {
state.apply(&event_record.event);
}
state
}
pub fn apply(&mut self, event: &AgentEvent) {
match event {
AgentEvent::RunStarted(e) => {
self.status = RunStatus::SessionLoaded;
self.provider = e.provider.clone();
self.model = e.model.clone();
self.updated_at = e.timestamp;
}
AgentEvent::ModelStarted(e) => {
self.iteration = e.iteration;
self.status = RunStatus::CallingModel;
self.updated_at = e.timestamp;
}
AgentEvent::UsageRecorded(e) => {
self.total_usage = TokenUsage::new(
self.total_usage.input_tokens + e.usage.input_tokens,
self.total_usage.output_tokens + e.usage.output_tokens,
);
self.updated_at = e.timestamp;
}
AgentEvent::RunCompleted(e) => {
self.status = RunStatus::Completed;
self.last_finish = Some(e.finish_reason.clone());
self.updated_at = e.timestamp;
}
AgentEvent::RunFailed(e) => {
self.status = RunStatus::Failed;
self.last_error = Some(e.error.clone());
self.updated_at = e.timestamp;
}
AgentEvent::RunCancelled(e) => {
self.status = RunStatus::Cancelled;
self.updated_at = e.timestamp;
}
AgentEvent::ContextBuilt(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::TextDelta(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::ToolCallStarted(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::ToolCallDelta(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::ToolCallCompleted(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::ToolExecutionStarted(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::ToolExecutionFinished(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::AssistantMessageCommitted(e) | AgentEvent::ToolMessageCommitted(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::DoomLoopDetected(e) => {
self.updated_at = e.timestamp;
}
AgentEvent::CompactionCircuitOpened(e) => {
self.updated_at = e.timestamp;
}
}
}
}
#[cfg(test)]
mod tests {
use super::super::event::{
ModelStarted as ModelStartedEvent, RunCompleted as RunCompletedEvent,
RunFailed as RunFailedEvent, RunStarted as RunStartedEvent,
UsageRecorded as UsageRecordedEvent,
};
use super::*;
fn make_record() -> RunRecord {
RunRecord::new(
RunId::new(),
uuid::Uuid::new_v4(),
ProviderId::new("test-provider"),
ModelName::new("test-model"),
Value::Null,
None,
)
}
#[test]
fn rebuilds_completed_run_from_events() {
let record = make_record();
let run_id = record.id;
let session_id = record.session_id;
let events = vec![
RunEventRecord::new(
0,
run_id,
AgentEvent::RunStarted(RunStartedEvent {
run_id,
session_id,
provider: record.provider.clone(),
model: record.model.clone(),
timestamp: Utc::now(),
}),
),
RunEventRecord::new(
1,
run_id,
AgentEvent::ModelStarted(ModelStartedEvent {
run_id,
provider: record.provider.clone(),
model: record.model.clone(),
iteration: 1,
timestamp: Utc::now(),
}),
),
RunEventRecord::new(
2,
run_id,
AgentEvent::UsageRecorded(UsageRecordedEvent {
run_id,
usage: TokenUsage::new(100, 50),
timestamp: Utc::now(),
}),
),
RunEventRecord::new(
3,
run_id,
AgentEvent::RunCompleted(RunCompletedEvent {
run_id,
finish_reason: FinishReason::Stop,
iterations: 1,
timestamp: Utc::now(),
}),
),
];
let state = RunState::create(&record, &events);
assert_eq!(state.run_id, run_id);
assert_eq!(state.session_id, session_id);
assert_eq!(state.status, RunStatus::Completed);
assert_eq!(state.provider, record.provider);
assert_eq!(state.model, record.model);
assert_eq!(state.iteration, 1);
assert_eq!(state.total_usage, TokenUsage::new(100, 50));
assert_eq!(state.last_finish, Some(FinishReason::Stop));
assert!(state.last_error.is_none());
assert_eq!(state.event_count, 4);
}
#[test]
fn rebuilds_failed_run_from_events() {
let record = make_record();
let run_id = record.id;
let session_id = record.session_id;
let events = vec![
RunEventRecord::new(
0,
run_id,
AgentEvent::RunStarted(RunStartedEvent {
run_id,
session_id,
provider: record.provider.clone(),
model: record.model.clone(),
timestamp: Utc::now(),
}),
),
RunEventRecord::new(
1,
run_id,
AgentEvent::RunFailed(RunFailedEvent {
run_id,
error: "something broke".to_string(),
timestamp: Utc::now(),
}),
),
];
let state = RunState::create(&record, &events);
assert_eq!(state.status, RunStatus::Failed);
assert_eq!(state.last_error, Some("something broke".to_string()));
assert!(state.last_finish.is_none());
}
#[test]
fn empty_events_returns_pending() {
let record = make_record();
let state = RunState::create(&record, &[]);
assert_eq!(state.status, RunStatus::Pending);
assert_eq!(state.iteration, 0);
assert_eq!(state.event_count, 0);
}
}