smith

//! Durable trace snapshot (`SMH-SPEC-SPEC0001`, Sessions).
//!
//! Entries are persisted through real tool and agent records so compaction
//! prunes live context while this snapshot covers durable replay history.
//! Session, branch, entry, message, tool-call ids, and timestamps are redacted.

use futures::StreamExt;
use serde_json::Value;
use smith::config::Config;
use smith::provider::StreamFn;
use smith::stream::{StopReason, StreamEvent, Usage};
use smith_core::agent::{Agent, CompactionSummary};
use smith_core::session::EntryKind;
use smith_core::store::{SessionReader, open_session};
use smith_core::tools::ToolSession;
use smith_core::trace::trace;

/// One JSON line per durable trace record, with run-specific values replaced.
fn redacted_trace(reader: &SessionReader, frame_bytes_max: usize) -> String {
    let header = reader.header(frame_bytes_max).unwrap();
    let mut placeholders = vec![
        (header.session_id.to_string(), "<session>".to_string()),
        (header.branch_id.to_string(), "<branch>".to_string()),
    ];
    if let Some(tip) = reader.selected_tip(frame_bytes_max).unwrap() {
        let mut entry_count = 0;
        let mut call_count = 0;
        reader
            .replay_at(tip, frame_bytes_max, |entry| {
                let index = entry_count;
                entry_count += 1;
                placeholders.push((entry.id.to_string(), format!("<entry-{index}>")));
                match entry.kind {
                    EntryKind::Message { id, .. } => {
                        placeholders.push((id.to_string(), format!("<message-{index}>")));
                    }
                    EntryKind::ToolCall { call_id } => {
                        placeholders.push((call_id.to_string(), format!("<call-{call_count}>")));
                        call_count += 1;
                    }
                    EntryKind::ToolResult { .. } | EntryKind::Meta { .. } => {}
                }
                Ok(())
            })
            .unwrap();
    }
    let redact = |text: &str| {
        placeholders
            .iter()
            .fold(text.to_string(), |text, (id, placeholder)| {
                text.replace(id, placeholder)
            })
    };
    let mut lines = Vec::new();
    trace(reader, frame_bytes_max, |record| {
        let mut line = serde_json::to_value(record).unwrap();
        if line["timestamp_ms"] != 0 {
            line["timestamp_ms"] = Value::from("<timestamp>");
        }
        lines.push(redact(&line.to_string()));
        Ok(())
    })
    .unwrap();
    lines.join("\n")
}

fn final_responses() -> StreamFn {
    let responses = std::sync::Arc::new(std::sync::Mutex::new(vec![
        vec![
            StreamEvent::text_delta(smith::id::MessageId::new(), "listed"),
            StreamEvent::stop(StopReason::EndTurn, Some(Usage::new())),
        ],
        vec![
            StreamEvent::text_delta(smith::id::MessageId::new(), "still listed"),
            StreamEvent::stop(StopReason::EndTurn, Some(Usage::new())),
        ],
    ]));
    std::sync::Arc::new(move |_, _| {
        futures::stream::iter(responses.lock().unwrap().remove(0).into_iter().map(Ok)).boxed()
    })
}

#[test]
fn trace_of_message_tool_round_trip_and_compaction() {
    let dir = std::env::temp_dir().join(format!("smith_trace_snapshot_{}", std::process::id()));
    let _ = std::fs::remove_dir_all(&dir);
    std::fs::create_dir_all(&dir).unwrap();
    let path = dir.join("trace.smh");
    let config = Config::default_valid();
    let opened = open_session(&path, config.clone()).unwrap();
    let mut tools = ToolSession::with_session(&dir, opened.session, opened.writer);
    tools
        .invoke(
            "write",
            &smith::ToolArguments::from_serializable(
                &serde_json::json!({"path": "a.txt", "content": "a"}),
            )
            .unwrap(),
        )
        .unwrap();
    let mut agent = Agent::new(final_responses(), tools, "mock-smith");
    agent.run_turn("list files").unwrap();
    agent
        .compact(&CompactionSummary {
            text: "listed a.txt".to_string(),
            complete: true,
        })
        .unwrap();
    agent.run_turn("and after compaction?").unwrap();

    let reader = SessionReader::open(&path).unwrap();
    insta::assert_snapshot!(redacted_trace(&reader, config.frame_bytes_max));
    drop(reader);
    drop(agent);
    let _ = std::fs::remove_dir_all(&dir);
}