⚡️ Replace unbounded channels with bounded channels and cancellation tokens

Switch content update channels from unbounded to bounded (capacity 100) for
better backpressure handling. Replace AtomicBool task coordination with
CancellationToken for cleaner async task lifecycle management.

The status polling task now uses tokio::select! with the cancellation token
and an interval timer, while the content update listener uses select! for
zero-latency message handling instead of polling with try_recv.

Author: hyperb1iss

Cargo.lock

@@ -56,6 +56,7 @@ thiserror = "2.0.4"
 tiktoken-rs = "0.6.0"
 tokio = { version = "1.44.2", features = ["full"] }
 tokio-retry = "0.3.0"
+tokio-util = "0.7"
 toml = "0.8.19"
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.18", features = ["env-filter", "fmt", "time", "json", "chrono"] }

Cargo.toml

@@ -31,13 +31,16 @@ pub enum ContentUpdate {
     Review { content: String },
 }
 
+/// Channel capacity for content updates
+pub const CONTENT_UPDATE_CHANNEL_CAPACITY: usize = 100;
+
 /// Sender for content updates - passed to tools and listened to by Studio
-pub type ContentUpdateSender = mpsc::UnboundedSender<ContentUpdate>;
-pub type ContentUpdateReceiver = mpsc::UnboundedReceiver<ContentUpdate>;
+pub type ContentUpdateSender = mpsc::Sender<ContentUpdate>;
+pub type ContentUpdateReceiver = mpsc::Receiver<ContentUpdate>;
 
-/// Create a new content update channel
+/// Create a new bounded content update channel
 pub fn create_content_update_channel() -> (ContentUpdateSender, ContentUpdateReceiver) {
-    mpsc::unbounded_channel()
+    mpsc::channel(CONTENT_UPDATE_CHANNEL_CAPACITY)
 }
 
 // ═══════════════════════════════════════════════════════════════════════════════
@@ -98,7 +101,7 @@ impl Tool for UpdateCommitTool {
             message: args.message.clone().unwrap_or_default(),
         };
 
-        self.sender.send(update).map_err(|e| {
+        self.sender.try_send(update).map_err(|e| {
             tracing::error!("Failed to send content update: {}", e);
             ContentUpdateError(format!("Failed to send update: {}", e))
         })?;
@@ -161,7 +164,7 @@ impl Tool for UpdatePRTool {
         };
 
         self.sender
-            .send(update)
+            .try_send(update)
             .map_err(|e| ContentUpdateError(format!("Failed to send update: {}", e)))?;
 
         let result = json!({
@@ -219,7 +222,7 @@ impl Tool for UpdateReviewTool {
         };
 
         self.sender
-            .send(update)
+            .try_send(update)
             .map_err(|e| ContentUpdateError(format!("Failed to send update: {}", e)))?;
 
         let result = json!({

src/agents/tools/content_update.rs

@@ -743,8 +743,7 @@ impl StudioApp {
         use crate::agents::status::IRIS_STATUS;
         use crate::agents::tools::{ContentUpdate, create_content_update_channel};
         use crate::studio::state::{ChatMessage, ChatRole};
-        use std::sync::Arc;
-        use std::sync::atomic::{AtomicBool, Ordering};
+        use tokio_util::sync::CancellationToken;
 
         let Some(agent) = self.agent_service.clone() else {
             let tx = self.iris_result_tx.clone();
@@ -754,7 +753,7 @@ impl StudioApp {
             return;
         };
 
-        // Create content update channel for tool-based updates
+        // Create bounded content update channel for tool-based updates
         let (content_tx, mut content_rx) = create_content_update_channel();
 
         // Capture context before spawning async task
@@ -772,44 +771,50 @@ impl StudioApp {
             .current_content
             .or_else(|| self.get_current_content_for_chat());
 
-        // Flag to signal when the main task is done
-        let is_done = Arc::new(AtomicBool::new(false));
-        let is_done_clone = is_done.clone();
-        let is_done_updates = is_done.clone();
+        // Cancellation token to signal when the main task is done
+        let cancel_token = CancellationToken::new();
+        let cancel_status = cancel_token.clone();
+        let cancel_updates = cancel_token.clone();
 
-        // Spawn a status polling task
+        // Spawn a status polling task (polls global state, so still uses interval)
         tokio::spawn(async move {
             use crate::agents::status::IrisPhase;
             let mut last_tool: Option<String> = None;
-
-            while !is_done_clone.load(Ordering::Relaxed) {
-                let status = IRIS_STATUS.get_current();
-
-                // Check if we're in a tool execution phase
-                if let IrisPhase::ToolExecution {
-                    ref tool_name,
-                    ref reason,
-                } = status.phase
-                {
-                    // Only send if it's a new tool
-                    if last_tool.as_ref() != Some(tool_name) {
-                        let _ = tx_status.send(IrisTaskResult::ToolStatus {
-                            tool_name: tool_name.clone(),
-                            message: reason.clone(),
-                        });
-                        last_tool = Some(tool_name.clone());
+            let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(100));
+
+            loop {
+                tokio::select! {
+                    _ = cancel_status.cancelled() => break,
+                    _ = interval.tick() => {
+                        let status = IRIS_STATUS.get_current();
+
+                        // Check if we're in a tool execution phase
+                        if let IrisPhase::ToolExecution {
+                            ref tool_name,
+                            ref reason,
+                        } = status.phase
+                        {
+                            // Only send if it's a new tool
+                            if last_tool.as_ref() != Some(tool_name) {
+                                let _ = tx_status.send(IrisTaskResult::ToolStatus {
+                                    tool_name: tool_name.clone(),
+                                    message: reason.clone(),
+                                });
+                                last_tool = Some(tool_name.clone());
+                            }
+                        }
                     }
                 }
-
-                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
             }
         });
 
-        // Spawn a task to listen for content updates from tools
+        // Spawn a task to listen for content updates from tools (uses select! for zero latency)
         tokio::spawn(async move {
-            while !is_done_updates.load(Ordering::Relaxed) {
-                match content_rx.try_recv() {
-                    Ok(update) => {
+            loop {
+                tokio::select! {
+                    _ = cancel_updates.cancelled() => break,
+                    update = content_rx.recv() => {
+                        let Some(update) = update else { break };
                         let chat_update = match update {
                             ContentUpdate::Commit {
                                 emoji,
@@ -834,10 +839,6 @@ impl StudioApp {
                         };
                         let _ = tx_updates.send(IrisTaskResult::ChatUpdate(chat_update));
                     }
-                    Err(tokio::sync::mpsc::error::TryRecvError::Empty) => {
-                        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
-                    }
-                    Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => break,
                 }
             }
         });
@@ -926,8 +927,8 @@ Simply call the appropriate tool with the new content. Do NOT echo back the full
                 }
             }
 
-            // Signal that we're done so the status polling task stops
-            is_done.store(true, Ordering::Relaxed);
+            // Signal that we're done so the helper tasks stop
+            cancel_token.cancel();
         });
     }