From a7980b7f728d54a472ceb90e8c8b85e4fa333a5c Mon Sep 17 00:00:00 2001 From: xiaoxixi Date: Fri, 17 Jul 2026 14:52:56 +0800 Subject: [PATCH] feat: add authoritative turn state machine --- src/agent/mod.rs | 2 + src/agent/turn_event.rs | 67 +++++ src/providers/traits.rs | 2 +- src/session/mod.rs | 5 + src/session/turn.rs | 592 ++++++++++++++++++++++++++++++++++++++++ 5 files changed, 667 insertions(+), 1 deletion(-) create mode 100644 src/agent/turn_event.rs create mode 100644 src/session/turn.rs diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 64591be..6e825b2 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -3,6 +3,7 @@ pub mod context_compressor; pub mod media_handler; pub mod sub_agent; pub mod system_prompt; +pub mod turn_event; pub use agent_loop::{AgentError, AgentLoop, AgentProcessResult}; pub use context_compressor::{ContextCompressor, estimate_tokens}; @@ -14,3 +15,4 @@ pub use system_prompt::{ PromptContext, PromptSection, SystemPromptBuilder, build_sub_agent_system_prompt, build_system_prompt, }; +pub use turn_event::{TurnEmitError, TurnEmitter, TurnEvent}; diff --git a/src/agent/turn_event.rs b/src/agent/turn_event.rs new file mode 100644 index 0000000..7f7c1cc --- /dev/null +++ b/src/agent/turn_event.rs @@ -0,0 +1,67 @@ +use std::sync::Arc; + +use thiserror::Error; + +use crate::providers::ToolCall; + +/// Presentation facts emitted while AgentLoop processes one model turn. +/// +/// Events contain no persistence or channel-delivery decisions. Session owns +/// the lifecycle around these facts and reduces them into authoritative state. +#[derive(Debug, Clone)] +pub enum TurnEvent { + ReasoningDelta { + iteration: u32, + delta: String, + }, + TextDelta { + iteration: u32, + delta: String, + }, + TextSegmentFinished { + iteration: u32, + }, + ToolStarted { + iteration: u32, + call: ToolCall, + }, + ToolFinished { + iteration: u32, + call_id: String, + success: bool, + preview: Option, + }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum TurnEmitError { + #[error("turn is no longer active")] + Inactive, + #[error("tool call {0} already exists in this turn")] + DuplicateTool(String), + #[error("tool call {0} does not exist in this turn")] + UnknownTool(String), +} + +type EmitFn = dyn Fn(TurnEvent) -> Result<(), TurnEmitError> + Send + Sync; + +/// Cheap cloneable handle used by AgentLoop to report presentation facts. +#[derive(Clone)] +pub struct TurnEmitter { + emit: Arc, +} + +impl TurnEmitter { + pub(crate) fn new(emit: F) -> Self + where + F: Fn(TurnEvent) -> Result<(), TurnEmitError> + Send + Sync + 'static, + { + Self { + emit: Arc::new(emit), + } + } + + pub fn emit(&self, event: TurnEvent) -> Result<(), TurnEmitError> { + (self.emit)(event) + } +} diff --git a/src/providers/traits.rs b/src/providers/traits.rs index 84ce6ab..f547b65 100644 --- a/src/providers/traits.rs +++ b/src/providers/traits.rs @@ -105,7 +105,7 @@ pub struct ChatCompletionResponse { pub usage: Usage, } -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct Usage { pub prompt_tokens: u32, pub completion_tokens: u32, diff --git a/src/session/mod.rs b/src/session/mod.rs index 0445afa..9ea385d 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -7,9 +7,14 @@ mod persistence; #[allow(clippy::module_inception)] pub mod session; pub mod session_id; +pub mod turn; pub use commands::SessionCommand; pub use error::SessionError; pub use events::{DialogInfo, SessionEvent}; pub use session::{SLASH_COMMANDS, Session, SessionManager, SlashCommand}; pub use session_id::UnifiedSessionId; +pub use turn::{ + BlockId, ToolStatus, TurnBlock, TurnController, TurnId, TurnPhase, TurnSnapshot, TurnState, + TurnStatus, +}; diff --git a/src/session/turn.rs b/src/session/turn.rs new file mode 100644 index 0000000..89d042b --- /dev/null +++ b/src/session/turn.rs @@ -0,0 +1,592 @@ +use std::sync::{Arc, Mutex, MutexGuard, Weak}; + +use serde::{Deserialize, Serialize}; +use tokio::sync::watch; + +use crate::agent::{TurnEmitError, TurnEmitter, TurnEvent}; +use crate::providers::Usage; + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct TurnId(pub String); + +impl TurnId { + pub fn new() -> Self { + Self(uuid::Uuid::new_v4().to_string()) + } +} + +impl Default for TurnId { + fn default() -> Self { + Self::new() + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct BlockId(pub String); + +impl BlockId { + fn new() -> Self { + Self(uuid::Uuid::new_v4().to_string()) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TurnStatus { + Running, + Completed, + Cancelled, + Failed, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum TurnPhase { + Queued, + Reasoning, + Responding, + Acting, + Finalizing, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolStatus { + Running, + Completed, + Failed, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum TurnBlock { + Reasoning { + id: BlockId, + iteration: u32, + text: String, + }, + Assistant { + id: BlockId, + iteration: u32, + text: String, + }, + Tool { + id: String, + iteration: u32, + name: String, + arguments: serde_json::Value, + status: ToolStatus, + preview: Option, + }, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct TurnState { + pub id: TurnId, + pub session_id: String, + pub message_id: String, + pub revision: u64, + pub status: TurnStatus, + pub phase: TurnPhase, + pub blocks: Vec, + pub usage: Option, + pub error: Option, +} + +pub type TurnSnapshot = TurnState; + +struct TurnControllerInner { + state: TurnState, + snapshots: watch::Sender>, + text_segment_open: bool, +} + +impl TurnControllerInner { + fn emit(&mut self, event: TurnEvent) -> Result<(), TurnEmitError> { + if self.state.status != TurnStatus::Running { + return Err(TurnEmitError::Inactive); + } + + let changed = match event { + TurnEvent::ReasoningDelta { iteration, delta } => { + if delta.is_empty() { + false + } else { + self.state.phase = TurnPhase::Reasoning; + match self.state.blocks.last_mut() { + Some(TurnBlock::Reasoning { + iteration: current, + text, + .. + }) if *current == iteration => text.push_str(&delta), + _ => self.state.blocks.push(TurnBlock::Reasoning { + id: BlockId::new(), + iteration, + text: delta, + }), + } + true + } + } + TurnEvent::TextDelta { iteration, delta } => { + if delta.is_empty() { + false + } else { + self.state.phase = TurnPhase::Responding; + if self.text_segment_open { + match self.state.blocks.last_mut() { + Some(TurnBlock::Assistant { + iteration: current, + text, + .. + }) if *current == iteration => text.push_str(&delta), + _ => { + self.push_text_block(iteration, delta); + } + } + } else { + self.push_text_block(iteration, delta); + } + self.text_segment_open = true; + true + } + } + TurnEvent::TextSegmentFinished { .. } => { + self.text_segment_open = false; + false + } + TurnEvent::ToolStarted { iteration, call } => { + if self + .state + .blocks + .iter() + .any(|block| matches!(block, TurnBlock::Tool { id, .. } if id == &call.id)) + { + return Err(TurnEmitError::DuplicateTool(call.id)); + } + self.text_segment_open = false; + self.state.phase = TurnPhase::Acting; + self.state.blocks.push(TurnBlock::Tool { + id: call.id, + iteration, + name: call.name, + arguments: call.arguments, + status: ToolStatus::Running, + preview: None, + }); + true + } + TurnEvent::ToolFinished { + iteration, + call_id, + success, + preview, + } => { + let Some(TurnBlock::Tool { + status, + preview: current_preview, + .. + }) = self.state.blocks.iter_mut().find(|block| { + matches!(block, TurnBlock::Tool { id, iteration: current, .. } if id == &call_id && *current == iteration) + }) + else { + return Err(TurnEmitError::UnknownTool(call_id)); + }; + *status = if success { + ToolStatus::Completed + } else { + ToolStatus::Failed + }; + *current_preview = preview; + true + } + }; + + if changed { + self.publish(); + } + Ok(()) + } + + fn push_text_block(&mut self, iteration: u32, text: String) { + self.state.blocks.push(TurnBlock::Assistant { + id: BlockId::new(), + iteration, + text, + }); + } + + fn publish(&mut self) { + self.state.revision = self.state.revision.wrapping_add(1); + self.snapshots.send_replace(Arc::new(self.state.clone())); + } + + fn transition_terminal( + &mut self, + status: TurnStatus, + usage: Option, + error: Option, + ) -> bool { + if self.state.status != TurnStatus::Running { + return false; + } + self.text_segment_open = false; + self.state.status = status; + self.state.phase = TurnPhase::Finalizing; + self.state.usage = usage; + self.state.error = error; + self.publish(); + true + } +} + +/// The sole writer for one running turn's presentation state. +/// +/// Mutation is synchronous and bounded to a small in-memory reduction. This +/// lets AgentLoop emit facts without creating an unbounded token queue or a +/// reducer background task. +pub struct TurnController { + inner: Arc>, +} + +impl TurnController { + pub fn start( + session_id: impl Into, + message_id: impl Into, + ) -> (Self, TurnEmitter, watch::Receiver>) { + let initial = TurnState { + id: TurnId::new(), + session_id: session_id.into(), + message_id: message_id.into(), + revision: 0, + status: TurnStatus::Running, + phase: TurnPhase::Queued, + blocks: Vec::new(), + usage: None, + error: None, + }; + let (snapshots, receiver) = watch::channel(Arc::new(initial.clone())); + let inner = Arc::new(Mutex::new(TurnControllerInner { + state: initial, + snapshots, + text_segment_open: false, + })); + let weak: Weak> = Arc::downgrade(&inner); + let emitter = TurnEmitter::new(move |event| { + let Some(inner) = weak.upgrade() else { + return Err(TurnEmitError::Inactive); + }; + lock_unpoisoned(&inner).emit(event) + }); + (Self { inner }, emitter, receiver) + } + + pub fn snapshot(&self) -> Arc { + Arc::new(lock_unpoisoned(&self.inner).state.clone()) + } + + pub fn begin_finalizing(&self) -> bool { + let mut inner = lock_unpoisoned(&self.inner); + if inner.state.status != TurnStatus::Running || inner.state.phase == TurnPhase::Finalizing { + return false; + } + inner.text_segment_open = false; + inner.state.phase = TurnPhase::Finalizing; + inner.publish(); + true + } + + pub fn complete(&self, usage: Option) -> bool { + lock_unpoisoned(&self.inner).transition_terminal(TurnStatus::Completed, usage, None) + } + + pub fn cancel(&self, reason: Option) -> bool { + lock_unpoisoned(&self.inner).transition_terminal(TurnStatus::Cancelled, None, reason) + } + + pub fn fail(&self, error: impl Into) -> bool { + lock_unpoisoned(&self.inner).transition_terminal( + TurnStatus::Failed, + None, + Some(error.into()), + ) + } +} + +fn lock_unpoisoned(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::providers::ToolCall; + + fn start() -> ( + TurnController, + TurnEmitter, + watch::Receiver>, + ) { + TurnController::start("cli:test:dialog", "assistant-message") + } + + #[test] + fn ordered_blocks_preserve_reasoning_text_tool_and_iterations() { + let (controller, emitter, _) = start(); + emitter + .emit(TurnEvent::ReasoningDelta { + iteration: 0, + delta: "plan ".into(), + }) + .unwrap(); + emitter + .emit(TurnEvent::ReasoningDelta { + iteration: 0, + delta: "step".into(), + }) + .unwrap(); + emitter + .emit(TurnEvent::TextDelta { + iteration: 0, + delta: "checking".into(), + }) + .unwrap(); + emitter + .emit(TurnEvent::ToolStarted { + iteration: 0, + call: ToolCall { + id: "call-1".into(), + name: "bash".into(), + arguments: serde_json::json!({"cmd":"pwd"}), + }, + }) + .unwrap(); + emitter + .emit(TurnEvent::ToolFinished { + iteration: 0, + call_id: "call-1".into(), + success: true, + preview: Some("/tmp".into()), + }) + .unwrap(); + emitter + .emit(TurnEvent::ReasoningDelta { + iteration: 1, + delta: "done thinking".into(), + }) + .unwrap(); + emitter + .emit(TurnEvent::TextDelta { + iteration: 1, + delta: "final".into(), + }) + .unwrap(); + + let snapshot = controller.snapshot(); + assert_eq!(snapshot.revision, 7); + assert_eq!(snapshot.phase, TurnPhase::Responding); + assert_eq!(snapshot.blocks.len(), 5); + assert!(matches!( + &snapshot.blocks[0], + TurnBlock::Reasoning { iteration: 0, text, .. } if text == "plan step" + )); + assert!(matches!( + &snapshot.blocks[1], + TurnBlock::Assistant { iteration: 0, text, .. } if text == "checking" + )); + assert!(matches!( + &snapshot.blocks[2], + TurnBlock::Tool { + id, + status: ToolStatus::Completed, + preview: Some(preview), + .. + } if id == "call-1" && preview == "/tmp" + )); + assert!(matches!( + &snapshot.blocks[3], + TurnBlock::Reasoning { iteration: 1, .. } + )); + assert!(matches!( + &snapshot.blocks[4], + TurnBlock::Assistant { iteration: 1, .. } + )); + } + + #[test] + fn explicit_segment_boundary_prevents_text_coalescing() { + let (controller, emitter, _) = start(); + emitter + .emit(TurnEvent::TextDelta { + iteration: 0, + delta: "before".into(), + }) + .unwrap(); + emitter + .emit(TurnEvent::TextSegmentFinished { iteration: 0 }) + .unwrap(); + emitter + .emit(TurnEvent::TextDelta { + iteration: 0, + delta: "after".into(), + }) + .unwrap(); + + let snapshot = controller.snapshot(); + assert_eq!(snapshot.revision, 2); + assert_eq!(snapshot.blocks.len(), 2); + assert!( + matches!(&snapshot.blocks[0], TurnBlock::Assistant { text, .. } if text == "before") + ); + assert!( + matches!(&snapshot.blocks[1], TurnBlock::Assistant { text, .. } if text == "after") + ); + } + + #[test] + fn revisions_are_monotonic_and_watch_is_latest_wins() { + let (_controller, emitter, receiver) = start(); + for delta in ["a", "b", "c"] { + emitter + .emit(TurnEvent::TextDelta { + iteration: 0, + delta: delta.into(), + }) + .unwrap(); + } + + let latest = receiver.borrow().clone(); + assert_eq!(latest.revision, 3); + assert!(matches!(&latest.blocks[0], TurnBlock::Assistant { text, .. } if text == "abc")); + } + + #[test] + fn terminal_state_rejects_late_events_and_is_idempotent() { + let (controller, emitter, receiver) = start(); + emitter + .emit(TurnEvent::TextDelta { + iteration: 0, + delta: "saved".into(), + }) + .unwrap(); + assert!(controller.begin_finalizing()); + assert!(controller.complete(None)); + assert!(!controller.complete(None)); + assert_eq!( + emitter.emit(TurnEvent::TextDelta { + iteration: 0, + delta: "late".into(), + }), + Err(TurnEmitError::Inactive) + ); + + let latest = receiver.borrow().clone(); + assert_eq!(latest.status, TurnStatus::Completed); + assert_eq!(latest.phase, TurnPhase::Finalizing); + assert_eq!(latest.revision, 3); + assert!(matches!(&latest.blocks[0], TurnBlock::Assistant { text, .. } if text == "saved")); + } + + #[test] + fn invalid_tool_transitions_do_not_publish() { + let (_controller, emitter, receiver) = start(); + let unknown = emitter.emit(TurnEvent::ToolFinished { + iteration: 0, + call_id: "missing".into(), + success: false, + preview: None, + }); + assert_eq!(unknown, Err(TurnEmitError::UnknownTool("missing".into()))); + assert_eq!(receiver.borrow().revision, 0); + + let call = ToolCall { + id: "same".into(), + name: "bash".into(), + arguments: serde_json::json!({}), + }; + emitter + .emit(TurnEvent::ToolStarted { + iteration: 0, + call: call.clone(), + }) + .unwrap(); + assert_eq!( + emitter.emit(TurnEvent::ToolStarted { iteration: 0, call }), + Err(TurnEmitError::DuplicateTool("same".into())) + ); + assert_eq!(receiver.borrow().revision, 1); + } + + #[test] + fn parallel_tools_update_independently() { + let (controller, emitter, _) = start(); + for id in ["first", "second"] { + emitter + .emit(TurnEvent::ToolStarted { + iteration: 0, + call: ToolCall { + id: id.into(), + name: "bash".into(), + arguments: serde_json::json!({"cmd": id}), + }, + }) + .unwrap(); + } + emitter + .emit(TurnEvent::ToolFinished { + iteration: 0, + call_id: "second".into(), + success: false, + preview: Some("failed".into()), + }) + .unwrap(); + + let snapshot = controller.snapshot(); + assert!(matches!( + &snapshot.blocks[0], + TurnBlock::Tool { id, status: ToolStatus::Running, .. } if id == "first" + )); + assert!(matches!( + &snapshot.blocks[1], + TurnBlock::Tool { id, status: ToolStatus::Failed, .. } if id == "second" + )); + } + + #[test] + fn empty_deltas_are_noops_and_cancel_reason_is_terminal() { + let (controller, emitter, receiver) = start(); + emitter + .emit(TurnEvent::ReasoningDelta { + iteration: 0, + delta: String::new(), + }) + .unwrap(); + emitter + .emit(TurnEvent::TextDelta { + iteration: 0, + delta: String::new(), + }) + .unwrap(); + assert_eq!(receiver.borrow().revision, 0); + + assert!(controller.cancel(Some("stopped by user".into()))); + let snapshot = controller.snapshot(); + assert_eq!(snapshot.status, TurnStatus::Cancelled); + assert_eq!(snapshot.error.as_deref(), Some("stopped by user")); + assert_eq!(snapshot.revision, 1); + } + + #[test] + fn failure_is_published_as_structured_terminal_state() { + let (controller, _emitter, _) = start(); + assert!(controller.fail("provider disconnected")); + let snapshot = controller.snapshot(); + assert_eq!(snapshot.status, TurnStatus::Failed); + assert_eq!(snapshot.error.as_deref(), Some("provider disconnected")); + assert_eq!(snapshot.phase, TurnPhase::Finalizing); + } +}