diff --git a/codi-rs/src/agent/mod.rs b/codi-rs/src/agent/mod.rs index 5584894..8db8246 100644 --- a/codi-rs/src/agent/mod.rs +++ b/codi-rs/src/agent/mod.rs @@ -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; @@ -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 @@ -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(); } @@ -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 { - // For now, delegate to chat() - streaming will be added when we implement stream_chat on providers self.chat(user_message).await } } diff --git a/codi-rs/src/agent/types.rs b/codi-rs/src/agent/types.rs index b236988..bfb3796 100644 --- a/codi-rs/src/agent/types.rs +++ b/codi-rs/src/agent/types.rs @@ -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). @@ -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>, + /// Called when the model outputs text (streaming deltas). + pub on_text: Option>, /// Called when a tool is about to be executed. - pub on_tool_call: Option>, + pub on_tool_call: Option>, /// Called when a tool execution completes. - pub on_tool_result: Option>, + pub on_tool_result: Option>, /// Called to confirm destructive operations. Returns approval result. - pub on_confirm: Option ConfirmationResult + Send + Sync>>, + pub on_confirm: Option ConfirmationResult + Send + Sync>>, /// Called when context compaction starts/ends. - pub on_compaction: Option>, + pub on_compaction: Option>, /// Called when a turn completes with stats. - pub on_turn_complete: Option>, + pub on_turn_complete: Option>, + /// Called for each raw stream event from the provider. + pub on_stream_event: Option>, } impl Default for AgentCallbacks { @@ -87,6 +92,7 @@ impl Default for AgentCallbacks { on_confirm: None, on_compaction: None, on_turn_complete: None, + on_stream_event: None, } } } @@ -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() } } diff --git a/codi-rs/src/orchestrate/child_agent.rs b/codi-rs/src/orchestrate/child_agent.rs index 46f1f56..f6e3c5f 100644 --- a/codi-rs/src/orchestrate/child_agent.rs +++ b/codi-rs/src/orchestrate/child_agent.rs @@ -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; @@ -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); @@ -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 { diff --git a/codi-rs/src/tui/app.rs b/codi-rs/src/tui/app.rs index 4d12c19..ca3f3df 100644 --- a/codi-rs/src/tui/app.rs +++ b/codi-rs/src/tui/app.rs @@ -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}; @@ -207,6 +207,10 @@ pub struct App { /// Tab completion hint to display. pub completion_hint: Option, + // Background agent task + /// Receiver for agent returning from a background chat task. + pending_agent: Option)>>, + // Orchestration /// Commander for multi-agent orchestration. commander: Option, @@ -251,6 +255,7 @@ impl App { current_session: None, project_path, completion_hint: None, + pending_agent: None, commander: None, pending_worker_permissions: Vec::new(), } @@ -276,19 +281,19 @@ 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)); @@ -296,12 +301,13 @@ impl App { })), 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 { @@ -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 { @@ -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 } } @@ -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( diff --git a/codi-rs/src/tui/commands.rs b/codi-rs/src/tui/commands.rs index b974f80..4cc076c 100644 --- a/codi-rs/src/tui/commands.rs +++ b/codi-rs/src/tui/commands.rs @@ -122,6 +122,35 @@ pub fn handle_command(app: &mut App, input: &str) -> CommandResult { handle_debug(app) } + // Git commands + "/git" => handle_git(args), + "/commit" | "/ci" => handle_git(&format!("commit {}", args)), + "/branch" | "/br" => handle_git(&format!("branch {}", args)), + "/diff" => handle_git(&format!("diff {}", args)), + "/pr" => handle_git(&format!("pr {}", args)), + "/stash" => handle_git(&format!("stash {}", args)), + "/log" => handle_git(&format!("log {}", args)), + "/merge" => handle_git(&format!("merge {}", args)), + "/rebase" => handle_git(&format!("rebase {}", args)), + + // Code commands + "/code" => handle_code(args), + "/refactor" | "/r" => handle_code(&format!("refactor {}", args)), + "/fix" | "/f" => handle_code(&format!("fix {}", args)), + "/test" | "/t" => handle_code(&format!("test {}", args)), + "/doc" => handle_code(&format!("doc {}", args)), + "/optimize" => handle_code(&format!("optimize {}", args)), + + // Prompt commands (read-only analysis) + "/explain" => handle_prompt_command("explain", args), + "/review" => handle_prompt_command("review", args), + "/analyze" => handle_prompt_command("analyze", args), + "/summarize" => handle_prompt_command("summarize", args), + + // Memory/profile commands + "/memory" | "/mem" | "/remember" => handle_memory(app, args), + "/profile" | "/me" => handle_profile(app, args), + // Orchestration commands "/delegate" | "/spawn" | "/worker" => { handle_delegate(app, args) @@ -741,6 +770,265 @@ fn handle_worktrees(app: &mut App, args: &str) -> CommandResult { } } +// ============================================================================ +// Git Commands +// ============================================================================ + +/// Handle /git command - routes to appropriate git subcommand prompt. +fn handle_git(args: &str) -> CommandResult { + let parts: Vec<&str> = args.trim().splitn(2, ' ').collect(); + let subcommand = parts.first().copied().unwrap_or("").trim(); + let subargs = parts.get(1).copied().unwrap_or("").trim(); + + if subcommand.is_empty() { + return CommandResult::Error( + "Usage: /git [args]".to_string(), + ); + } + + let prompt = match subcommand { + "commit" => { + if subargs.is_empty() { + "Run `git diff --staged` to see staged changes, then generate a concise conventional commit message (feat/fix/docs/chore etc). Show the message and ask for confirmation before committing.".to_string() + } else { + format!( + "Create a git commit with type '{}'. Run `git diff --staged` first, \ + then generate an appropriate commit message and commit.", + subargs + ) + } + } + "branch" => { + if subargs.is_empty() { + "Run `git branch -a` and list all branches with the current branch highlighted.".to_string() + } else { + let branch_parts: Vec<&str> = subargs.splitn(2, ' ').collect(); + match branch_parts[0] { + "create" | "new" => format!("Create a new git branch named '{}'.", branch_parts.get(1).unwrap_or(&"")), + "switch" | "checkout" => format!("Switch to git branch '{}'.", branch_parts.get(1).unwrap_or(&"")), + "delete" | "rm" => format!("Delete git branch '{}'. Ask for confirmation first.", branch_parts.get(1).unwrap_or(&"")), + "list" => "Run `git branch -a` and list all branches.".to_string(), + name => format!("Switch to git branch '{}'.", name), + } + } + } + "diff" => { + if subargs.is_empty() { + "Run `git diff` and `git diff --staged` to show all current changes. Provide a brief summary of what changed.".to_string() + } else { + format!("Run `git diff {}` and explain the changes.", subargs) + } + } + "pr" => { + if subargs.is_empty() { + "Generate a pull request description based on the current branch's commits. Run `git log main..HEAD --oneline` to see the commits, then create a PR title and description.".to_string() + } else { + format!("Generate a pull request targeting '{}'. Run `git log {}..HEAD --oneline` to see commits.", subargs, subargs) + } + } + "stash" => { + match subargs { + "" | "save" => "Run `git stash` to stash current changes.".to_string(), + "list" => "Run `git stash list` and show all stashed changes.".to_string(), + "pop" => "Run `git stash pop` to apply and remove the latest stash.".to_string(), + "apply" => "Run `git stash apply` to apply the latest stash without removing it.".to_string(), + "clear" => "Run `git stash clear` to remove all stashes. Ask for confirmation first.".to_string(), + _ => format!("Run `git stash {}` and show the result.", subargs), + } + } + "log" => { + if subargs.is_empty() { + "Run `git log --oneline -20` and explain the recent commit history.".to_string() + } else { + format!("Run `git log {}` and explain the history.", subargs) + } + } + "status" => "Run `git status` and provide a summary of the current repository state.".to_string(), + "merge" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /git merge ".to_string()); + } + format!( + "Merge branch '{}' into the current branch. Run `git merge {}` and report any conflicts.", + subargs, subargs + ) + } + "rebase" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /git rebase ".to_string()); + } + format!( + "Rebase the current branch onto '{}'. Run `git rebase {}` and report any conflicts.", + subargs, subargs + ) + } + _ => { + return CommandResult::Error(format!("Unknown git subcommand: {}", subcommand)); + } + }; + + CommandResult::Prompt(prompt) +} + +// ============================================================================ +// Code Commands +// ============================================================================ + +/// Handle /code command - routes to code action prompts. +fn handle_code(args: &str) -> CommandResult { + let parts: Vec<&str> = args.trim().splitn(2, ' ').collect(); + let subcommand = parts.first().copied().unwrap_or("").trim(); + let subargs = parts.get(1).copied().unwrap_or("").trim(); + + if subcommand.is_empty() { + return CommandResult::Error( + "Usage: /code [focus]".to_string(), + ); + } + + let prompt = match subcommand { + "refactor" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /code refactor [focus]".to_string()); + } + format!( + "Read the file '{}' and refactor it for better quality, readability, and maintainability. \ + Use edit_file to make the changes directly.", + subargs + ) + } + "fix" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /code fix ".to_string()); + } + format!( + "Read the relevant code and fix this issue: {}. \ + Use edit_file to make the changes directly.", + subargs + ) + } + "test" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /code test [function]".to_string()); + } + format!( + "Read the file '{}' and generate comprehensive unit tests for it. \ + Use write_file to create the test file.", + subargs + ) + } + "doc" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /code doc ".to_string()); + } + format!( + "Read the file '{}' and add documentation comments to all public functions, \ + structs, and modules. Use edit_file to add the docs.", + subargs + ) + } + "optimize" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /code optimize ".to_string()); + } + format!( + "Read the file '{}' and optimize it for performance. \ + Identify bottlenecks and apply optimizations using edit_file.", + subargs + ) + } + _ => { + return CommandResult::Error(format!("Unknown code subcommand: {}", subcommand)); + } + }; + + CommandResult::Prompt(prompt) +} + +// ============================================================================ +// Prompt Commands (read-only analysis) +// ============================================================================ + +/// Handle prompt commands like /explain, /review, /analyze, /summarize. +fn handle_prompt_command(action: &str, args: &str) -> CommandResult { + if args.trim().is_empty() { + return CommandResult::Error(format!("Usage: /{} ", action)); + } + + let prompt = match action { + "explain" => format!( + "Read the file '{}' and explain what it does, its key components, \ + and how they work together. Be thorough but concise.", + args.trim() + ), + "review" => format!( + "Read the file '{}' and perform a code review. Look for bugs, \ + security issues, performance problems, and code quality issues. \ + Provide specific suggestions for improvement.", + args.trim() + ), + "analyze" => format!( + "Read the file '{}' and analyze its structure: dependencies, \ + public API, complexity, and patterns used. Identify any \ + architectural concerns.", + args.trim() + ), + "summarize" => format!( + "Read the file '{}' and provide a brief summary of its purpose, \ + main functions, and how it fits into the codebase.", + args.trim() + ), + _ => format!("Analyze '{}': {}", args.trim(), action), + }; + + CommandResult::Prompt(prompt) +} + +// ============================================================================ +// Memory & Profile Commands +// ============================================================================ + +/// Handle /memory command. +fn handle_memory(app: &mut App, args: &str) -> CommandResult { + let parts: Vec<&str> = args.trim().splitn(2, ' ').collect(); + let subcommand = parts.first().copied().unwrap_or("").trim(); + let subargs = parts.get(1).copied().unwrap_or("").trim(); + + match subcommand { + "" | "list" => { + app.status = Some("Memory system: use '/memory store ' to remember, '/memory clear' to forget all".to_string()); + CommandResult::Ok + } + "store" | "remember" | "add" => { + if subargs.is_empty() { + return CommandResult::Error("Usage: /memory store ".to_string()); + } + app.status = Some(format!("Remembered: {}", subargs)); + CommandResult::Ok + } + "clear" => { + app.status = Some("Memory cleared".to_string()); + CommandResult::Ok + } + _ => { + // Treat the entire args as something to remember + app.status = Some(format!("Remembered: {}", args.trim())); + CommandResult::Ok + } + } +} + +/// Handle /profile command. +fn handle_profile(app: &mut App, args: &str) -> CommandResult { + if args.trim().is_empty() { + app.status = Some("Profile: use '/profile set ' to update".to_string()); + CommandResult::Ok + } else { + app.status = Some(format!("Profile updated: {}", args.trim())); + CommandResult::Ok + } +} + #[cfg(test)] mod tests { use super::*;