Skip to content

feat: fix streaming hang and integrate missing slash commands #260

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laynepenney merged 1 commit into from
Feb 3, 2026
Merged

feat: fix streaming hang and integrate missing slash commands #260

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions codi-rs/src/agent/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ use std::time::{Duration, Instant};

use crate::error::{AgentError, Result};
use crate::types::{
BoxedProvider, ContentBlock, Message, Role,
BoxedProvider, ContentBlock, Message, Role, StreamEvent,
ToolCall, ToolDefinition, ToolResult,
};
use crate::tools::ToolRegistry;
Expand Down Expand Up @@ -315,11 +315,27 @@ impl Agent {
let tools = self.get_tool_definitions();
let system_context = self.build_system_context();

// Call the provider
let response = self.provider.chat(
// Clone callbacks for the streaming closure (Arc clones are cheap)
let on_text = self.callbacks.on_text.clone();
let on_stream_event = self.callbacks.on_stream_event.clone();

// Call the provider with streaming
let response = self.provider.stream_chat(
&self.state.messages,
tools.as_deref(),
Some(&system_context),
Box::new(move |event| {
// Forward raw stream events
if let Some(ref cb) = on_stream_event {
cb(&event);
}
// Fire on_text for text deltas
if let StreamEvent::TextDelta(ref text) = event {
if let Some(ref cb) = on_text {
cb(text);
}
}
}),
).await?;

// Update token stats
Expand All @@ -329,11 +345,8 @@ impl Agent {
turn_stats.total_tokens = turn_stats.input_tokens + turn_stats.output_tokens;
}

// Stream text to callback
// Store final response text
if !response.content.is_empty() {
if let Some(ref on_text) = self.callbacks.on_text {
on_text(&response.content);
}
final_response = response.content.clone();
}

Expand Down Expand Up @@ -407,10 +420,9 @@ impl Agent {

/// Chat with streaming output.
///
/// Similar to `chat()` but streams text output via the `on_text` callback
/// as it's received from the model.
/// Alias for `chat()` - streaming is now built into the main chat loop
/// via `provider.stream_chat()`.
pub async fn stream_chat(&mut self, user_message: &str) -> Result<String> {
// For now, delegate to chat() - streaming will be added when we implement stream_chat on providers
self.chat(user_message).await
}
}
Expand Down
23 changes: 15 additions & 8 deletions codi-rs/src/agent/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

use std::sync::Arc;

use crate::types::{BoxedProvider, Message};
use crate::types::{BoxedProvider, Message, StreamEvent};
use crate::tools::ToolRegistry;

/// Statistics for a single turn (user message -> final response).
Expand Down Expand Up @@ -63,19 +63,24 @@ pub enum ConfirmationResult {
}

/// Callbacks for agent events.
///
/// Uses `Arc` instead of `Box` so callbacks can be cloned into streaming
/// closures and background tasks without lifetime issues.
pub struct AgentCallbacks {
/// Called when the model outputs text.
pub on_text: Option<Box<dyn Fn(&str) + Send + Sync>>,
/// Called when the model outputs text (streaming deltas).
pub on_text: Option<Arc<dyn Fn(&str) + Send + Sync>>,
/// Called when a tool is about to be executed.
pub on_tool_call: Option<Box<dyn Fn(&str, &serde_json::Value) + Send + Sync>>,
pub on_tool_call: Option<Arc<dyn Fn(&str, &serde_json::Value) + Send + Sync>>,
/// Called when a tool execution completes.
pub on_tool_result: Option<Box<dyn Fn(&str, &str, bool) + Send + Sync>>,
pub on_tool_result: Option<Arc<dyn Fn(&str, &str, bool) + Send + Sync>>,
/// Called to confirm destructive operations. Returns approval result.
pub on_confirm: Option<Box<dyn Fn(ToolConfirmation) -> ConfirmationResult + Send + Sync>>,
pub on_confirm: Option<Arc<dyn Fn(ToolConfirmation) -> ConfirmationResult + Send + Sync>>,
/// Called when context compaction starts/ends.
pub on_compaction: Option<Box<dyn Fn(bool) + Send + Sync>>,
pub on_compaction: Option<Arc<dyn Fn(bool) + Send + Sync>>,
/// Called when a turn completes with stats.
pub on_turn_complete: Option<Box<dyn Fn(&TurnStats) + Send + Sync>>,
pub on_turn_complete: Option<Arc<dyn Fn(&TurnStats) + Send + Sync>>,
/// Called for each raw stream event from the provider.
pub on_stream_event: Option<Arc<dyn Fn(&StreamEvent) + Send + Sync>>,
}

impl Default for AgentCallbacks {
Expand All @@ -87,6 +92,7 @@ impl Default for AgentCallbacks {
on_confirm: None,
on_compaction: None,
on_turn_complete: None,
on_stream_event: None,
}
}
}
Expand All @@ -100,6 +106,7 @@ impl std::fmt::Debug for AgentCallbacks {
.field("on_confirm", &self.on_confirm.is_some())
.field("on_compaction", &self.on_compaction.is_some())
.field("on_turn_complete", &self.on_turn_complete.is_some())
.field("on_stream_event", &self.on_stream_event.is_some())
.finish()
}
}
Expand Down
5 changes: 3 additions & 2 deletions codi-rs/src/orchestrate/child_agent.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ impl ChildAgent {
let auto_approve = self.auto_approve.clone();

let callbacks = AgentCallbacks {
on_confirm: Some(Box::new(move |confirmation: ToolConfirmation| {
on_confirm: Some(Arc::new(move |confirmation: ToolConfirmation| {
// Check auto-approve list
if auto_approve.contains(&confirmation.tool_name) {
return ConfirmationResult::Approve;
Expand Down Expand Up @@ -212,7 +212,7 @@ impl ChildAgent {
}
})),
on_text: None,
on_tool_call: Some(Box::new({
on_tool_call: Some(Arc::new({
let ipc = Arc::clone(&self.ipc);
move |tool_name: &str, _input: &serde_json::Value| {
let ipc = Arc::clone(&ipc);
Expand All @@ -229,6 +229,7 @@ impl ChildAgent {
on_tool_result: None,
on_compaction: None,
on_turn_complete: None,
on_stream_event: None,
};

let agent_config = AgentConfig {
Expand Down
90 changes: 72 additions & 18 deletions codi-rs/src/tui/app.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use crate::agent::{
Agent, AgentCallbacks, AgentConfig, AgentOptions,
ConfirmationResult, ToolConfirmation, TurnStats,
};
use crate::error::ToolError;
use crate::error::{Result as CodiResult, ToolError};
use crate::completion::{complete_line, get_completion_matches};
use crate::orchestrate::{Commander, CommanderConfig, WorkerConfig, WorkerStatus, WorkspaceInfo, PermissionResult};
use crate::session::{Session, SessionInfo, SessionService};
Expand Down Expand Up @@ -207,6 +207,10 @@ pub struct App {
/// Tab completion hint to display.
pub completion_hint: Option<String>,

// Background agent task
/// Receiver for agent returning from a background chat task.
pending_agent: Option<tokio::sync::oneshot::Receiver<(Agent, CodiResult<String>)>>,

// Orchestration
/// Commander for multi-agent orchestration.
commander: Option<Commander>,
Expand Down Expand Up @@ -251,6 +255,7 @@ impl App {
current_session: None,
project_path,
completion_hint: None,
pending_agent: None,
commander: None,
pending_worker_permissions: Vec::new(),
}
Expand All @@ -276,32 +281,33 @@ impl App {
let event_tx = self.event_tx.clone().unwrap();

let callbacks = AgentCallbacks {
on_text: Some(Box::new({
on_text: Some(Arc::new({
let tx = event_tx.clone();
move |text: &str| {
let _ = tx.send(AppEvent::TextDelta(text.to_string()));
}
})),
on_tool_call: Some(Box::new({
on_tool_call: Some(Arc::new({
let tx = event_tx.clone();
move |name: &str, input: &serde_json::Value| {
let _ = tx.send(AppEvent::ToolStart(name.to_string(), input.clone()));
}
})),
on_tool_result: Some(Box::new({
on_tool_result: Some(Arc::new({
let tx = event_tx.clone();
move |name: &str, result: &str, is_error: bool| {
let _ = tx.send(AppEvent::ToolResult(name.to_string(), result.to_string(), is_error));
}
})),
on_confirm: None, // Handled via channel-based approach
on_compaction: None,
on_turn_complete: Some(Box::new({
on_turn_complete: Some(Arc::new({
let tx = event_tx.clone();
move |stats: &TurnStats| {
let _ = tx.send(AppEvent::TurnComplete(stats.clone()));
}
})),
on_stream_event: None,
};

self.agent = Some(Agent::new(AgentOptions {
Expand Down Expand Up @@ -350,6 +356,36 @@ impl App {

/// Process any pending app events from agent callbacks.
fn process_app_events(&mut self) {
// Check if the background agent task has completed
if let Some(ref mut rx) = self.pending_agent {
match rx.try_recv() {
Ok((agent, result)) => {
self.agent = Some(agent);
self.pending_agent = None;
match result {
Ok(_) => {
// Response was streamed via callbacks; TurnComplete will finalize
}
Err(e) => {
self.status = Some(format!("Error: {}", e));
self.mode = AppMode::Normal;
self.finalize_streaming();
}
}
}
Err(tokio::sync::oneshot::error::TryRecvError::Closed) => {
// Task panicked or was dropped
self.pending_agent = None;
self.status = Some("Agent task failed unexpectedly".to_string());
self.mode = AppMode::Normal;
self.finalize_streaming();
}
Err(tokio::sync::oneshot::error::TryRecvError::Empty) => {
// Still running, keep waiting
}
}
}

// Collect events first to avoid borrow issues
let mut events = Vec::new();
if let Some(ref mut rx) = self.event_rx {
Expand Down Expand Up @@ -757,7 +793,25 @@ impl App {
// Execute async command
let _ = execute_async_command(self, cmd).await;
}
CommandResult::Ok | CommandResult::Error(_) | CommandResult::Prompt(_) => {
CommandResult::Prompt(prompt) => {
// Command generated a prompt to send to the AI
self.messages.push(Message::user(&prompt));
self.scroll_to_bottom();

if let Some(mut agent) = self.agent.take() {
self.mode = AppMode::Waiting;
self.status = Some("Thinking...".to_string());

let (tx, rx) = tokio::sync::oneshot::channel();
self.pending_agent = Some(rx);

tokio::spawn(async move {
let result = agent.chat(&prompt).await;
let _ = tx.send((agent, result));
});
}
}
CommandResult::Ok | CommandResult::Error(_) => {
// Already handled synchronously
}
}
Expand All @@ -768,21 +822,21 @@ impl App {
self.messages.push(Message::user(&input));
self.scroll_to_bottom();

// Get AI response
if let Some(ref mut agent) = self.agent {
// Get AI response - spawn on background task so the event loop stays responsive
if let Some(mut agent) = self.agent.take() {
self.mode = AppMode::Waiting;
self.status = Some("Thinking...".to_string());

// Call the agent
match agent.chat(&input).await {
Ok(_response) => {
// Response is handled via callbacks
}
Err(e) => {
self.status = Some(format!("Error: {}", e));
self.mode = AppMode::Normal;
}
}
// Create a oneshot channel to get the agent back when done
let (tx, rx) = tokio::sync::oneshot::channel();
self.pending_agent = Some(rx);

// Spawn the agent chat on a background task
tokio::spawn(async move {
let result = agent.chat(&input).await;
// Send the agent and result back (ignore error if receiver dropped)
let _ = tx.send((agent, result));
});
} else {
// No agent, just echo
self.messages.push(Message::assistant(
Expand Down
Loading