use chrono::Utc; use db::sqlx; use model::agent::AgentTraceModel; use serde_json::{Value, json}; use uuid::Uuid; use crate::AppService; use crate::error::AppError; pub struct TraceContext { pub invocation_id: Uuid, pub conversation_id: Uuid, } #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct TraceReplay { pub invocation_id: Uuid, pub conversation_id: Uuid, pub phases: Vec, } #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct TracePhaseRow { pub id: Uuid, pub sequence: i32, pub phase: String, pub label: String, pub content: Option, pub tool_calls: Option, pub tool_results: Option, pub input_tokens: Option, pub output_tokens: Option, pub metadata: Option, pub created_at: chrono::DateTime, } impl AppService { pub async fn trace_record( &self, ctx: &TraceContext, sequence: i32, phase: &str, content: Option<&str>, tool_calls: Option<&Value>, tool_results: Option<&Value>, input_tokens: Option, output_tokens: Option, metadata: Option<&Value>, ) -> Result { let id = Uuid::now_v7(); let now = Utc::now(); sqlx::query( "INSERT INTO agent_trace \ (id, invocation, conversation, sequence, phase, content, \ tool_calls, tool_results, input_tokens, output_tokens, metadata, created_at) \ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)", ) .bind(id) .bind(ctx.invocation_id) .bind(ctx.conversation_id) .bind(sequence) .bind(phase) .bind(content) .bind(tool_calls) .bind(tool_results) .bind(input_tokens) .bind(output_tokens) .bind(metadata) .bind(now) .execute(self.db.writer()) .await .map_err(|e| AppError::DatabaseError(e.to_string()))?; Ok(id) } pub async fn trace_replay_by_invocation( &self, invocation_id: Uuid, ) -> Result { let rows = sqlx::query_as::<_, AgentTraceModel>( "SELECT id, invocation, conversation, sequence, phase, content, \ tool_calls, tool_results, input_tokens, output_tokens, metadata, created_at \ FROM agent_trace WHERE invocation = $1 ORDER BY sequence ASC", ) .bind(invocation_id) .fetch_all(self.db.reader()) .await .map_err(|e| AppError::DatabaseError(e.to_string()))?; let conversation_id = rows.first().map(|r| r.conversation).unwrap_or(Uuid::nil()); Ok(TraceReplay { invocation_id, conversation_id, phases: rows .into_iter() .map(|r| TracePhaseRow { id: r.id, sequence: r.sequence, phase: r.phase.clone(), label: r.phase_label().to_string(), content: r.content, tool_calls: r.tool_calls, tool_results: r.tool_results, input_tokens: r.input_tokens, output_tokens: r.output_tokens, metadata: r.metadata, created_at: r.created_at, }) .collect(), }) } pub async fn trace_replay_by_conversation( &self, conversation_id: Uuid, ) -> Result, AppError> { let rows = sqlx::query_as::<_, AgentTraceModel>( "SELECT id, invocation, conversation, sequence, phase, content, \ tool_calls, tool_results, input_tokens, output_tokens, metadata, created_at \ FROM agent_trace WHERE conversation = $1 ORDER BY invocation, sequence ASC", ) .bind(conversation_id) .fetch_all(self.db.reader()) .await .map_err(|e| AppError::DatabaseError(e.to_string()))?; let mut grouped: std::collections::BTreeMap< Uuid, Vec, > = std::collections::BTreeMap::new(); for row in rows { grouped.entry(row.invocation).or_default().push(row); } Ok(grouped .into_iter() .map(|(invocation_id, rows)| TraceReplay { invocation_id, conversation_id, phases: rows .into_iter() .map(|r| TracePhaseRow { id: r.id, sequence: r.sequence, phase: r.phase.clone(), label: r.phase_label().to_string(), content: r.content, tool_calls: r.tool_calls, tool_results: r.tool_results, input_tokens: r.input_tokens, output_tokens: r.output_tokens, metadata: r.metadata, created_at: r.created_at, }) .collect(), }) .collect()) } } pub struct TraceAccumulator { ctx: TraceContext, seq: i32, think_buf: String, answer_buf: String, think_tokens: i64, answer_tokens: i64, svc: AppService, } impl TraceAccumulator { pub fn new( svc: AppService, invocation_id: Uuid, conversation_id: Uuid, ) -> Self { Self { ctx: TraceContext { invocation_id, conversation_id, }, seq: 0, think_buf: String::new(), answer_buf: String::new(), think_tokens: 0, answer_tokens: 0, svc, } } pub async fn feed_thinking(&mut self, chunk: &str) { self.think_buf.push_str(chunk); self.think_tokens += (chunk.chars().count() as f64 / 2.5).ceil() as i64; } pub async fn feed_text(&mut self, chunk: &str) { if !self.think_buf.is_empty() { self.flush_think().await; } self.answer_buf.push_str(chunk); self.answer_tokens += (chunk.chars().count() as f64 / 2.5).ceil() as i64; } pub async fn feed_tool_call( &mut self, tool_call_id: &str, tool_name: &str, args: &Value, ) { if !self.answer_buf.is_empty() { self.flush_answer().await; } let _ = self.svc.trace_record( &self.ctx, self.seq, "act", None, Some(&json!({ "tool_call_id": tool_call_id, "name": tool_name, "arguments": args })), None, None, None, None, ).await; self.seq += 1; } pub async fn feed_tool_result( &mut self, tool_call_id: &str, tool_name: &str, output: Option<&Value>, error: Option<&str>, elapsed_ms: i64, ) { let _ = self .svc .trace_record( &self.ctx, self.seq, "act", None, None, Some(&json!({ "tool_call_id": tool_call_id, "name": tool_name, "output": output, "error": error, "elapsed_ms": elapsed_ms, })), None, None, None, ) .await; self.seq += 1; } pub async fn finish( &mut self, output: &str, input_tokens: i64, output_tokens: i64, ) { if !self.think_buf.is_empty() { self.flush_think().await; } if !self.answer_buf.is_empty() { self.flush_answer().await; } let _ = self .svc .trace_record( &self.ctx, self.seq, "summarize", Some(output), None, None, Some(input_tokens), Some(output_tokens), None, ) .await; } async fn flush_think(&mut self) { let content = std::mem::take(&mut self.think_buf); let tokens = self.think_tokens; self.think_tokens = 0; let _ = self .svc .trace_record( &self.ctx, self.seq, "think", Some(&content), None, None, Some(tokens), None, None, ) .await; self.seq += 1; } async fn flush_answer(&mut self) { let content = std::mem::take(&mut self.answer_buf); let tokens = self.answer_tokens; self.answer_tokens = 0; let _ = self .svc .trace_record( &self.ctx, self.seq, "answer", Some(&content), None, None, None, Some(tokens), None, ) .await; self.seq += 1; } }