diff --git a/crates/ironclaw_engine/src/executor/loop_engine.rs b/crates/ironclaw_engine/src/executor/loop_engine.rs index 422542d7..1a50f92c 100644 --- a/crates/ironclaw_engine/src/executor/loop_engine.rs +++ b/crates/ironclaw_engine/src/executor/loop_engine.rs @@ -33,6 +33,8 @@ pub struct ExecutionLoop { policy: Arc, signal_rx: SignalReceiver, user_id: String, + /// Optional broadcast sender for live event streaming. + event_tx: Option>, } impl ExecutionLoop { @@ -53,9 +55,29 @@ impl ExecutionLoop { policy, signal_rx, user_id, + event_tx: None, } } + /// Set the event broadcast sender for live status updates. + pub fn with_event_tx( + mut self, + tx: tokio::sync::broadcast::Sender, + ) -> Self { + self.event_tx = Some(tx); + self + } + + /// Add an event to the thread and broadcast it for live status updates. + fn emit_event(&mut self, kind: EventKind) { + let event = crate::types::event::ThreadEvent::new(self.thread.id, kind); + if let Some(ref tx) = self.event_tx { + let _ = tx.send(event.clone()); + } + self.thread.events.push(event); + self.thread.updated_at = chrono::Utc::now(); + } + /// Run the execution loop to completion. pub async fn run(&mut self) -> Result { // Transition to Running @@ -168,7 +190,7 @@ impl ExecutionLoop { // 6. Create step let mut step = Step::new(self.thread.id, iteration + 1); step.status = StepStatus::LlmCalling; - self.thread.add_event(EventKind::StepStarted { + self.emit_event(EventKind::StepStarted { step_id: step.id, }); @@ -199,7 +221,7 @@ impl ExecutionLoop { .add_message(ThreadMessage::assistant(text)); step.status = StepStatus::Completed; step.completed_at = Some(chrono::Utc::now()); - self.thread.add_event(EventKind::StepCompleted { + self.emit_event(EventKind::StepCompleted { step_id: step.id, tokens: step.tokens_used, }); @@ -231,7 +253,7 @@ impl ExecutionLoop { step.status = StepStatus::Completed; step.completed_at = Some(chrono::Utc::now()); - self.thread.add_event(EventKind::StepCompleted { + self.emit_event(EventKind::StepCompleted { step_id: step.id, tokens: step.tokens_used, }); @@ -245,7 +267,7 @@ impl ExecutionLoop { step.status = StepStatus::Completed; step.completed_at = Some(chrono::Utc::now()); - self.thread.add_event(EventKind::StepCompleted { + self.emit_event(EventKind::StepCompleted { step_id: step.id, tokens: step.tokens_used, }); @@ -292,7 +314,7 @@ impl ExecutionLoop { // Record events for event_kind in batch.events { - self.thread.add_event(event_kind); + self.emit_event(event_kind); } // Add action results as messages @@ -307,7 +329,7 @@ impl ExecutionLoop { step.action_results = batch.results; step.status = StepStatus::Completed; step.completed_at = Some(chrono::Utc::now()); - self.thread.add_event(EventKind::StepCompleted { + self.emit_event(EventKind::StepCompleted { step_id: step.id, tokens: step.tokens_used, }); @@ -365,7 +387,7 @@ impl ExecutionLoop { // Record events for event_kind in code_result.events { - self.thread.add_event(event_kind); + self.emit_event(event_kind); } // Add action results as messages @@ -388,7 +410,7 @@ impl ExecutionLoop { step.status = StepStatus::Completed; step.completed_at = Some(chrono::Utc::now()); - self.thread.add_event(EventKind::StepCompleted { + self.emit_event(EventKind::StepCompleted { step_id: step.id, tokens: step.tokens_used, }); diff --git a/crates/ironclaw_engine/src/runtime/manager.rs b/crates/ironclaw_engine/src/runtime/manager.rs index d00944c4..9a0cae72 100644 --- a/crates/ironclaw_engine/src/runtime/manager.rs +++ b/crates/ironclaw_engine/src/runtime/manager.rs @@ -38,6 +38,8 @@ pub struct ThreadManager { pub policy: Arc, tree: RwLock, running: RwLock>, + /// Broadcast channel for thread events (for live status updates). + event_tx: tokio::sync::broadcast::Sender, } impl ThreadManager { @@ -49,6 +51,7 @@ impl ThreadManager { leases: Arc, policy: Arc, ) -> Self { + let (event_tx, _) = tokio::sync::broadcast::channel(256); Self { llm, effects, @@ -58,9 +61,15 @@ impl ThreadManager { policy, tree: RwLock::new(ThreadTree::new()), running: RwLock::new(HashMap::new()), + event_tx, } } + /// Subscribe to thread events for live status updates. + pub fn subscribe_events(&self) -> tokio::sync::broadcast::Receiver { + self.event_tx.subscribe() + } + /// Spawn a new thread and start executing it. /// /// Grants default capability leases for all registered capabilities. @@ -142,7 +151,8 @@ impl ThreadManager { let leases = Arc::clone(&self.leases); let policy = Arc::clone(&self.policy); - let exec_loop = ExecutionLoop::new(thread, llm, effects, leases, policy, rx, user_id); + let exec_loop = ExecutionLoop::new(thread, llm, effects, leases, policy, rx, user_id) + .with_event_tx(self.event_tx.clone()); // Spawn background task let handle = tokio::spawn(async move { diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 44d27391..771ae471 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -175,7 +175,7 @@ pub struct AgentDeps { pub struct Agent { pub(super) config: AgentConfig, pub(super) deps: AgentDeps, - pub(super) channels: Arc, + pub(crate) channels: Arc, pub(super) context_manager: Arc, pub(super) scheduler: Arc, pub(super) router: Router, diff --git a/src/bridge/router.rs b/src/bridge/router.rs index 954209d8..f5a7bce0 100644 --- a/src/bridge/router.rs +++ b/src/bridge/router.rs @@ -14,7 +14,7 @@ use crate::agent::Agent; use crate::bridge::effect_adapter::EffectBridgeAdapter; use crate::bridge::llm_adapter::LlmBridgeAdapter; use crate::bridge::store_adapter::InMemoryStore; -use crate::channels::IncomingMessage; +use crate::channels::{IncomingMessage, StatusUpdate}; use crate::error::Error; /// Check if the engine v2 is enabled via `ENGINE_V2=true` environment variable. @@ -144,6 +144,16 @@ pub async fn handle_with_engine( "engine v2: handling message" ); + // Send "Thinking..." status to the channel + let _ = agent + .channels + .send_status( + &message.channel, + StatusUpdate::Thinking("Processing...".into()), + &message.metadata, + ) + .await; + // Get or create conversation for this channel+user let conv_id = state .conversation_manager @@ -168,19 +178,43 @@ pub async fn handle_with_engine( }) })?; - debug!(thread_id = %thread_id, "engine v2: thread active, waiting for completion"); + debug!(thread_id = %thread_id, "engine v2: thread spawned"); - // Wait for the thread to complete - let outcome = state - .thread_manager - .join_thread(thread_id) - .await - .map_err(|e| { - crate::error::Error::from(crate::error::JobError::ContextError { - id: uuid::Uuid::nil(), - reason: format!("engine v2 join error: {e}"), - }) - })?; + // Subscribe to live events for progress updates + let mut event_rx = state.thread_manager.subscribe_events(); + let channels = &agent.channels; + let channel_name = &message.channel; + let metadata = &message.metadata; + + // Forward events to the channel while waiting for thread completion + loop { + tokio::select! { + // Check for events from the execution loop + event = event_rx.recv() => { + match event { + Ok(ref evt) if evt.thread_id == thread_id => { + forward_event_to_channel(evt, channels, channel_name, metadata).await; + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => break, + _ => {} // Event for a different thread, or lagged + } + } + // Also check if the thread has finished (in case we miss the events) + _ = tokio::time::sleep(std::time::Duration::from_millis(500)) => { + if !state.thread_manager.is_running(thread_id).await { + break; + } + } + } + } + + // Join the thread to get the outcome + let outcome = state.thread_manager.join_thread(thread_id).await.map_err(|e| { + crate::error::Error::from(crate::error::JobError::ContextError { + id: uuid::Uuid::nil(), + reason: format!("engine v2 join error: {e}"), + }) + })?; // Record outcome in conversation state @@ -208,3 +242,70 @@ pub async fn handle_with_engine( ))), } } + +/// Forward an engine ThreadEvent to the channel as a StatusUpdate. +async fn forward_event_to_channel( + event: &ironclaw_engine::ThreadEvent, + channels: &std::sync::Arc, + channel_name: &str, + metadata: &serde_json::Value, +) { + use ironclaw_engine::EventKind; + + match &event.kind { + EventKind::StepStarted { .. } => { + let _ = channels + .send_status( + channel_name, + StatusUpdate::Thinking("Thinking...".into()), + metadata, + ) + .await; + } + EventKind::ActionExecuted { + action_name, + .. + } => { + let _ = channels + .send_status( + channel_name, + StatusUpdate::ToolCompleted { + name: action_name.clone(), + success: true, + error: None, + parameters: None, + }, + metadata, + ) + .await; + } + EventKind::ActionFailed { + action_name, + error, + .. + } => { + let _ = channels + .send_status( + channel_name, + StatusUpdate::ToolCompleted { + name: action_name.clone(), + success: false, + error: Some(error.clone()), + parameters: None, + }, + metadata, + ) + .await; + } + EventKind::StepCompleted { .. } => { + let _ = channels + .send_status( + channel_name, + StatusUpdate::Thinking("Processing results...".into()), + metadata, + ) + .await; + } + _ => {} // Other events don't need channel status + } +}