trae_agent_rs_core/output/
mod.rs

1//! Output abstraction layer for the Trae Agent core
2//!
3//! This module provides an abstract interface for outputting agent execution information,
4//! allowing different implementations for CLI, API, logging, etc.
5
6use crate::tools::{ToolCall, ToolResult};
7use async_trait::async_trait;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11// Core only provides abstractions - implementations are in calling modules
12
13/// Null output handler that discards all events (useful for testing and backward compatibility)
14pub struct NullOutput;
15
16#[async_trait]
17impl AgentOutput for NullOutput {
18    async fn emit_event(
19        &self,
20        _event: AgentEvent,
21    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
22        Ok(())
23    }
24}
25
26/// Convenience module for backward compatibility
27pub mod events {
28    pub use super::NullOutput;
29}
30
31/// Status of tool execution
32#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
33pub enum ToolExecutionStatus {
34    /// Tool is currently executing
35    Executing,
36    /// Tool completed successfully
37    Success,
38    /// Tool failed with an error
39    Error,
40}
41
42/// Rich tool execution information
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct ToolExecutionInfo {
45    /// Unique identifier for this tool execution
46    pub execution_id: String,
47    /// Tool name (e.g., "bash", "str_replace_based_edit_tool")
48    pub tool_name: String,
49    /// Tool parameters/arguments
50    pub parameters: HashMap<String, serde_json::Value>,
51    /// Current execution status
52    pub status: ToolExecutionStatus,
53    /// Tool result (if completed)
54    pub result: Option<ToolResult>,
55    /// Timestamp of status change
56    pub timestamp: chrono::DateTime<chrono::Utc>,
57    /// Additional metadata for tool-specific information
58    pub metadata: HashMap<String, serde_json::Value>,
59}
60
61/// Agent execution step information
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct AgentStepInfo {
64    /// Step number in the execution sequence
65    pub step_number: usize,
66    /// Current task description
67    pub task: String,
68    /// LLM thinking/reasoning (if available)
69    pub thinking: Option<String>,
70    /// Tool executions in this step
71    pub tool_executions: Vec<ToolExecutionInfo>,
72    /// Step completion status
73    pub completed: bool,
74}
75
76/// Token usage statistics
77#[derive(Debug, Clone, Serialize, Deserialize, Default)]
78pub struct TokenUsage {
79    /// Total input tokens consumed
80    pub input_tokens: u32,
81    /// Total output tokens generated
82    pub output_tokens: u32,
83    /// Total tokens (input + output)
84    pub total_tokens: u32,
85}
86
87/// Agent execution context information
88#[derive(Debug, Clone, Serialize, Deserialize)]
89pub struct AgentExecutionContext {
90    /// Agent configuration name or identifier
91    pub agent_id: String,
92    /// Current task being executed
93    pub task: String,
94    /// Project path or working directory
95    pub project_path: String,
96    /// Maximum allowed steps
97    pub max_steps: usize,
98    /// Current step number
99    pub current_step: usize,
100    /// Total execution time so far
101    pub execution_time: std::time::Duration,
102    /// Token usage statistics
103    pub token_usage: TokenUsage,
104}
105
106/// Events that can be emitted during agent execution
107#[derive(Debug, Clone, Serialize, Deserialize)]
108pub enum AgentEvent {
109    /// Agent execution started
110    ExecutionStarted { context: AgentExecutionContext },
111    /// Agent execution completed
112    ExecutionCompleted {
113        context: AgentExecutionContext,
114        success: bool,
115        summary: String,
116    },
117    /// New step started
118    StepStarted { step_info: AgentStepInfo },
119    /// Step completed
120    StepCompleted { step_info: AgentStepInfo },
121    /// Tool execution started
122    ToolExecutionStarted { tool_info: ToolExecutionInfo },
123    /// Tool execution status updated
124    ToolExecutionUpdated { tool_info: ToolExecutionInfo },
125    /// Tool execution completed
126    ToolExecutionCompleted { tool_info: ToolExecutionInfo },
127    /// Agent thinking/reasoning
128    AgentThinking {
129        step_number: usize,
130        thinking: String,
131    },
132    /// Token usage updated (emitted after each LLM call)
133    TokenUsageUpdated { token_usage: TokenUsage },
134    /// Agent status update (for interactive mode status reporting)
135    StatusUpdate {
136        status: String,
137        metadata: HashMap<String, serde_json::Value>,
138    },
139    /// General message or log
140    Message {
141        level: MessageLevel,
142        content: String,
143        metadata: HashMap<String, serde_json::Value>,
144    },
145}
146
147/// Message severity levels
148#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
149pub enum MessageLevel {
150    Debug,
151    Info,
152    Normal,
153    Warning,
154    Error,
155}
156
157/// Abstract output interface for agent execution
158#[async_trait]
159pub trait AgentOutput: Send + Sync {
160    /// Emit an agent event
161    async fn emit_event(
162        &self,
163        event: AgentEvent,
164    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>>;
165
166    /// Emit a message with specified level
167    async fn emit_message(
168        &self,
169        level: MessageLevel,
170        content: &str,
171    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
172        self.emit_event(AgentEvent::Message {
173            level,
174            content: content.to_string(),
175            metadata: HashMap::new(),
176        })
177        .await
178    }
179
180    /// Emit debug message
181    async fn debug(&self, content: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
182        self.emit_message(MessageLevel::Debug, content).await
183    }
184
185    /// Emit token usage update
186    async fn emit_token_update(
187        &self,
188        token_usage: TokenUsage,
189    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
190        self.emit_event(AgentEvent::TokenUsageUpdated { token_usage })
191            .await
192    }
193
194    /// Emit info message
195    async fn info(&self, content: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
196        self.emit_message(MessageLevel::Info, content).await
197    }
198
199    /// Emit warning message
200    async fn warning(&self, content: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
201        self.emit_message(MessageLevel::Warning, content).await
202    }
203
204    /// Emit error message
205    async fn error(&self, content: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
206        self.emit_message(MessageLevel::Error, content).await
207    }
208
209    /// Emit normal text message
210    async fn normal(&self, content: &str) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
211        self.emit_message(MessageLevel::Normal, content).await
212    }
213
214    /// Emit status update (for interactive mode status reporting)
215    async fn emit_status_update(
216        &self,
217        status: &str,
218    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
219        self.emit_event(AgentEvent::StatusUpdate {
220            status: status.to_string(),
221            metadata: HashMap::new(),
222        })
223        .await
224    }
225
226    /// Check if this output handler supports real-time updates
227    fn supports_realtime_updates(&self) -> bool {
228        false
229    }
230
231    /// Flush any buffered output (for implementations that buffer)
232    async fn flush(&self) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
233        Ok(())
234    }
235}
236
237/// Helper trait for creating tool execution info
238pub trait ToolExecutionInfoBuilder {
239    fn create_tool_execution_info(
240        tool_call: &ToolCall,
241        status: ToolExecutionStatus,
242        result: Option<&ToolResult>,
243    ) -> ToolExecutionInfo;
244}
245
246impl ToolExecutionInfoBuilder for ToolExecutionInfo {
247    fn create_tool_execution_info(
248        tool_call: &ToolCall,
249        status: ToolExecutionStatus,
250        result: Option<&ToolResult>,
251    ) -> ToolExecutionInfo {
252        let parameters = if let serde_json::Value::Object(map) = &tool_call.parameters {
253            map.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
254        } else {
255            let mut map = HashMap::new();
256            map.insert("raw_parameters".to_string(), tool_call.parameters.clone());
257            map
258        };
259
260        ToolExecutionInfo {
261            execution_id: tool_call.id.clone(),
262            tool_name: tool_call.name.clone(),
263            parameters,
264            status,
265            result: result.cloned(),
266            timestamp: chrono::Utc::now(),
267            metadata: HashMap::new(),
268        }
269    }
270}
trae_agent_rs_core/output/mod.rs

trae_agent_rs_core/output/
mod.rs
