diff --git a/codi-rs/src/orchestrate/child_agent.rs b/codi-rs/src/orchestrate/child_agent.rs index e2b5fff..a06bdca 100644 --- a/codi-rs/src/orchestrate/child_agent.rs +++ b/codi-rs/src/orchestrate/child_agent.rs @@ -15,13 +15,21 @@ //! 5. Report completion or error use std::path::Path; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex}; use std::time::Instant; use tokio::sync::Mutex; use tracing::{error, info, warn}; -use crate::agent::{Agent, AgentCallbacks, AgentConfig, AgentOptions, ConfirmationResult, ToolConfirmation}; +use crate::agent::{ + Agent, + AgentCallbacks, + AgentConfig, + AgentOptions, + ConfirmationResult, + ToolConfirmation, + TurnStats, +}; use crate::tools::ToolRegistry; use crate::types::TokenUsage; use crate::providers::create_provider_from_env; @@ -129,22 +137,33 @@ impl ChildAgent { // Report completion match &result { - Ok(response) => { + Ok((response, stats)) => { let mut ipc = ipc.lock().await; // Get git stats let commits = child_agent.get_commits().await.unwrap_or_default(); let files_changed = child_agent.get_changed_files().await.unwrap_or_default(); + let (tool_count, usage) = match stats { + Some(stats) => { + let usage = TokenUsage { + input_tokens: stats.input_tokens as u32, + output_tokens: stats.output_tokens as u32, + ..Default::default() + }; + (stats.tool_call_count as u32, Some(usage)) + } + None => (0, None), + }; let worker_result = WorkerResult { success: true, response: response.clone(), - tool_count: 0, // TODO: Track this + tool_count, duration_ms, commits, files_changed, branch: Some(child_agent.workspace.branch().to_string()), - usage: None, // TODO: Track this + usage, }; ipc.send_task_complete(worker_result.clone()).await?; @@ -173,7 +192,10 @@ impl ChildAgent { } /// Execute the task using an agent. - async fn execute_task(&mut self, task: &str) -> Result { + async fn execute_task( + &mut self, + task: &str, + ) -> Result<(String, Option), ChildAgentError> { // Create provider let provider = create_provider_from_env()?; @@ -183,6 +205,8 @@ impl ChildAgent { // Create agent with IPC-based confirmation let ipc = Arc::clone(&self.ipc); let auto_approve = self.auto_approve.clone(); + let turn_stats: Arc>> = Arc::new(StdMutex::new(None)); + let turn_stats_capture = Arc::clone(&turn_stats); let callbacks = AgentCallbacks { on_confirm: Some(Arc::new(move |confirmation: ToolConfirmation| { @@ -232,7 +256,11 @@ impl ChildAgent { })), on_tool_result: None, on_compaction: None, - on_turn_complete: None, + on_turn_complete: Some(Arc::new(move |stats: &TurnStats| { + if let Ok(mut guard) = turn_stats_capture.lock() { + *guard = Some(stats.clone()); + } + })), on_stream_event: None, }; @@ -265,8 +293,9 @@ impl ChildAgent { // Execute chat let response = agent.chat(task).await .map_err(|e| ChildAgentError::Agent(crate::error::AgentError::InvalidState(e.to_string())))?; + let stats = turn_stats.lock().ok().and_then(|guard| guard.clone()); - Ok(response) + Ok((response, stats)) } /// Get commits made in this workspace.