Agent Call Tracing

{
  "started_at": 1704067200000,
  "step_name": "execute-task",
  "prompt_preview": "You are planning how to implement a task. The file tree of the repository is provided as previous step output.\n\nTask: add health check endpoint\n\nCreate a brief plan:",
  "events": [
    {
      "kind": "text",
      "content": "I'll add a /health endpoint",
      "elapsed_ms": 1200
    },
    {
      "kind": "tool_request",
      "content": "read_file { path: \"src/main.rs\" }",
      "elapsed_ms": 1250
    },
    {
      "kind": "tool_response",
      "content": "fn main() { ... }",
      "elapsed_ms": 1350
    }
  ],
  "response_preview": "I'll add a /health endpoint that returns 200 OK with JSON: { \"status\": \"ok\" }. The endpoint will be added to the existing Axum router in src/main.rs.",
  "duration_ms": 45320,
  "tool_call_count": 8
}

pub struct TraceBuilder {
    step_name: String,
    prompt_preview: String,
    events: Vec<TraceEvent>,
    tool_call_count: u32,
    start: Instant,
    started_at: u64,
    /// Buffer for Text events — flushed on newlines or non-Text events.
    text_buf: String,
}

impl TraceBuilder {
    pub fn new(step_name: &str, prompt: &str) -> Self {
        let started_at = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;

        let prompt_preview: String = prompt.chars().take(200).collect();

        if is_trace_verbose() {
            eprintln!("[{step_name}] CALL: {:?}", truncate(prompt, 120));
        }

        Self {
            step_name: step_name.to_string(),
            prompt_preview,
            events: Vec::new(),
            tool_call_count: 0,
            start: Instant::now(),
            started_at,
            text_buf: String::new(),
        }
    }
}

pub fn record_event(&mut self, kind: EventKind, content: &str) {
    let elapsed_ms = self.start.elapsed().as_millis() as u64;

    if is_trace_verbose() {
        match &kind {
            EventKind::Text => {
                // Buffer text and flush complete lines — avoids per-token noise.
                self.text_buf.push_str(content);
                while let Some(nl) = self.text_buf.find('\n') {
                    let line = self.text_buf[..nl].to_string();
                    self.text_buf.drain(..=nl);
                    if !line.is_empty() {
                        eprintln!("[{}] TEXT: {}", self.step_name, truncate(&line, 120));
                    }
                }
            }
            _ => {
                // Flush any buffered text before printing a non-Text event.
                self.flush_text_buf();
                let label = match &kind {
                    EventKind::ToolRequest => "TOOL_REQ",
                    EventKind::ToolResponse => "TOOL_RES",
                    EventKind::Thinking => "THINK",
                    EventKind::Error => "ERROR",
                    _ => unreachable!(),
                };
                eprintln!("[{}] {}: {}", self.step_name, label, truncate(content, 120));
            }
        }
    }

    if matches!(kind, EventKind::ToolRequest) {
        self.tool_call_count += 1;
    }

    self.events.push(TraceEvent {
        kind,
        content: content.to_string(),
        elapsed_ms,
    });
}

pub fn write_trace(trace: &AgentCallTrace, trace_dir: &Path) -> Result<()> {
    use std::io::Write;

    std::fs::create_dir_all(trace_dir)?;

    let date_str = date_from_unix_ms(trace.started_at);
    let file_path = trace_dir.join(format!("magpie-trace-{date_str}.jsonl"));

    let line = serde_json::to_string(trace)?;

    let mut file = std::fs::OpenOptions::new()
        .create(true)
        .append(true)
        .open(&file_path)?;
    writeln!(file, "{line}")?;

    debug!(path = %file_path.display(), "trace written");
    Ok()
}

static TRACE_VERBOSE: AtomicBool = AtomicBool::new(false);

pub fn set_trace_verbose(enabled: bool) {
    TRACE_VERBOSE.store(enabled, Ordering::Relaxed);
}

pub fn is_trace_verbose() -> bool {
    TRACE_VERBOSE.load(Ordering::Relaxed)
}

async fn claude_call(prompt: &str, step_name: &str, trace_dir: Option<&PathBuf>) -> Result<String> {
    let mut tb = TraceBuilder::new(step_name, prompt);

    let output = tokio::process::Command::new("claude")
        .args(["-p", prompt])
        .env_remove("CLAUDECODE")
        .output()
        .await
        .context("failed to run `claude` CLI — is it installed and on PATH?")?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        tb.record_event(EventKind::Error, &stderr);
        let trace = tb.finish("");
        if let Some(dir) = trace_dir {
            let _ = trace::write_trace(&trace, dir);
        }
        anyhow::bail!("claude CLI exited with {}: {}", output.status, stderr.trim());
    }

    let response = String::from_utf8_lossy(&output.stdout).trim().to_string();
    tb.record_event(EventKind::Text, &response);

    let trace = tb.finish(&response);
    if let Some(dir) = trace_dir {
        let _ = trace::write_trace(&trace, dir);
    }

    Ok(response)
}