poe2_agent/

trace.rs

//! Agent reasoning trace — a structured record of everything that happens
//! during a single `respond()` call.
//!
//! The trace captures the full internal reasoning loop: user input, each
//! LLM round's tool calls and results, streamed text output, and final
//! token usage. Consumers (e.g. the web backend) decide how to persist it
//! (JSON file, database, etc.).

use serde::Serialize;

use crate::llm::Usage;

/// Complete trace of a single agent response.
#[derive(Debug, Clone, Serialize)]
pub struct AgentTrace {
    /// ISO-8601 timestamp when the response started.
    pub started_at: String,
    /// The user message that triggered this response.
    pub user_message: String,
    /// Conversation history provided as context (user and assistant turns).
    pub history: Vec<TraceMessage>,
    /// Ordered list of events that occurred during the response.
    pub events: Vec<TraceEvent>,
    /// Cumulative token usage across all LLM rounds.
    pub usage: Option<TraceUsage>,
}

/// A message from conversation history.
#[derive(Debug, Clone, Serialize)]
pub struct TraceMessage {
    pub role: String,
    pub content: String,
}

/// A single event in the agent reasoning loop.
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum TraceEvent {
    /// The LLM decided to call a tool.
    #[serde(rename = "tool_call")]
    ToolCall {
        /// Tool name (e.g. "search_trade", "get_build_stats").
        name: String,
        /// Raw JSON arguments the LLM passed.
        arguments: String,
    },
    /// A tool returned a result.
    #[serde(rename = "tool_result")]
    ToolResult {
        /// Tool name.
        name: String,
        /// The full result content (JSON string).
        content: String,
    },
    /// A chunk of streamed text from the LLM's final answer.
    #[serde(rename = "text_delta")]
    TextDelta {
        /// The text fragment.
        text: String,
    },
    /// The agent produced a build mutation (item creation/modification).
    #[serde(rename = "build_mutation")]
    BuildMutation {
        /// Label describing the mutation.
        label: String,
    },
}

/// Token usage summary, mirroring [`Usage`] in a serializable form.
#[derive(Debug, Clone, Serialize)]
pub struct TraceUsage {
    pub input_tokens: u32,
    pub output_tokens: u32,
    pub cached_tokens: u32,
    pub total_tokens: u32,
}

impl From<&Usage> for TraceUsage {
    fn from(u: &Usage) -> Self {
        Self {
            input_tokens: u.input_tokens,
            output_tokens: u.output_tokens,
            cached_tokens: u.cached_tokens(),
            total_tokens: u.total_tokens,
        }
    }
}

/// Builder that accumulates trace events during an agent response.
///
/// Created at the start of `respond()`, written to as events occur,
/// then finalized into an [`AgentTrace`].
pub(crate) struct TraceBuilder {
    started_at: String,
    user_message: String,
    history: Vec<TraceMessage>,
    events: Vec<TraceEvent>,
}

impl TraceBuilder {
    /// Start a new trace for a response to `user_message`.
    pub fn new(user_message: &str, history: impl IntoIterator<Item = TraceMessage>) -> Self {
        Self {
            started_at: now_iso8601(),
            user_message: user_message.to_owned(),
            history: history.into_iter().collect(),
            events: Vec::new(),
        }
    }

    /// Record a tool call.
    pub fn tool_call(&mut self, name: &str, arguments: &str) {
        self.events.push(TraceEvent::ToolCall {
            name: name.to_owned(),
            arguments: arguments.to_owned(),
        });
    }

    /// Record a tool result.
    pub fn tool_result(&mut self, name: &str, content: &str) {
        self.events.push(TraceEvent::ToolResult {
            name: name.to_owned(),
            content: content.to_owned(),
        });
    }

    /// Record a text delta from the final streamed answer.
    pub fn text_delta(&mut self, text: &str) {
        self.events.push(TraceEvent::TextDelta {
            text: text.to_owned(),
        });
    }

    /// Record a build mutation.
    pub fn build_mutation(&mut self, label: &str) {
        self.events.push(TraceEvent::BuildMutation {
            label: label.to_owned(),
        });
    }

    /// Finalize the trace with usage stats.
    pub fn finish(self, usage: &Usage) -> AgentTrace {
        AgentTrace {
            started_at: self.started_at,
            user_message: self.user_message,
            history: self.history,
            events: self.events,
            usage: Some(TraceUsage::from(usage)),
        }
    }
}

fn now_iso8601() -> String {
    // Use std SystemTime — no chrono dependency needed.
    let now = std::time::SystemTime::now();
    let duration = now
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default();
    let secs = duration.as_secs();

    // Format as UTC: YYYY-MM-DDTHH:MM:SSZ
    let days = secs / 86400;
    let time_secs = secs % 86400;
    let hours = time_secs / 3600;
    let minutes = (time_secs % 3600) / 60;
    let seconds = time_secs % 60;

    // Days since epoch to Y-M-D (civil calendar).
    let (y, m, d) = days_to_ymd(days);
    format!("{y:04}-{m:02}-{d:02}T{hours:02}:{minutes:02}:{seconds:02}Z")
}

/// Convert days since Unix epoch to (year, month, day).
fn days_to_ymd(days: u64) -> (u64, u64, u64) {
    // Algorithm from Howard Hinnant's date library (public domain).
    let z = days + 719468;
    let era = z / 146097;
    let doe = z - era * 146097;
    let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
    let y = yoe + era * 400;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let d = doy - (153 * mp + 2) / 5 + 1;
    let m = if mp < 10 { mp + 3 } else { mp - 9 };
    let y = if m <= 2 { y + 1 } else { y };
    (y, m, d)
}