smith

//! Replay evidence derived from durable session state (`SMH-SPEC-SPEC0001`,
//! Sessions).
//!
//! A trace is an ordered, provider-free reconstruction of what a session
//! recorded. It is derived only from persisted entry records, never from
//! decoded content, so the same session file always produces the same trace,
//! byte for byte.

use crate::session::{
    COMPACTION_KIND, EntryContent, EntryFrame, EntryKind, SessionHeader, TOOL_MUTATION_KIND,
};
use crate::store::SessionReader;
use serde::Deserialize;
use smith::error::Result;
use smith::id::{SessionId, ToolCallId};

/// Replay-relevant trace record (durable replay evidence).
#[derive(Clone, PartialEq, Eq, Debug, serde::Serialize, serde::Deserialize)]
pub struct TraceRecord {
    /// Timestamp of the recorded event, in Unix milliseconds.
    pub timestamp_ms: i64,
    /// Session the event belongs to.
    pub session_id: SessionId,
    /// Event kind.
    pub kind: TraceKind,
    /// Stable reference to the recorded event.
    pub event_ref: String,
}

/// Kind of replay event.
#[derive(Clone, Copy, PartialEq, Eq, Debug, serde::Serialize, serde::Deserialize)]
pub enum TraceKind {
    /// Session created.
    SessionCreated,
    /// Entry added.
    EntryAdded,
    /// Message streamed.
    MessageStreamed,
    /// Tool call emitted.
    ToolCallEmitted,
    /// Raw tool result observed; not authoritative effect completion.
    ToolResultReceived,
    /// Authoritative tool completion: explicit durable metadata or a replayed
    /// incomplete lifecycle synthesized as `unknown`.
    ToolMutationCompleted,
    /// Compaction occurred.
    Compaction,
    /// Replay started.
    ReplayStarted,
    /// Replay completed.
    ReplayCompleted,
}

/// Visit durable trace records for the selected history in source order.
///
/// Each entry contributes the record kind matching what it recorded, so tool
/// activity and compaction remain visible rather than collapsing into a
/// generic entry event. A raw result is authoritative only when its matching
/// mutation metadata is present; otherwise replay emits a completed `unknown`
/// mutation record. One pending lifecycle is retained while history is streamed
/// through [`SessionReader::replay_at`], never a whole-history call map.
///
/// # Errors
///
/// Returns a reader validation, storage, ancestry, or visitor error.
pub fn trace(
    reader: &SessionReader,
    frame_bytes_max: usize,
    visit: impl FnMut(TraceRecord) -> Result<()>,
) -> Result<()> {
    let header = reader.header(frame_bytes_max)?;
    trace_selected(reader, &header, frame_bytes_max, visit).map(|_| ())
}

/// Replay durable selected history without a provider, yielding each record
/// in source order.
///
/// Replay is framed by exactly one start and completion record. Its completion
/// timestamp and entry count derive from the durable selected tip, never live
/// effective context.
///
/// # Errors
///
/// Returns a reader validation, storage, ancestry, or visitor error.
pub fn replay(
    reader: &SessionReader,
    frame_bytes_max: usize,
    mut visit: impl FnMut(TraceRecord) -> Result<()>,
) -> Result<()> {
    let header = reader.header(frame_bytes_max)?;
    visit(TraceRecord {
        timestamp_ms: 0,
        session_id: header.session_id,
        kind: TraceKind::ReplayStarted,
        event_ref: format!("session:{}", header.session_id),
    })?;
    let (count, timestamp_ms) = trace_selected(reader, &header, frame_bytes_max, &mut visit)?;
    visit(TraceRecord {
        timestamp_ms,
        session_id: header.session_id,
        kind: TraceKind::ReplayCompleted,
        event_ref: format!("entries:{count}"),
    })
}

fn trace_selected(
    reader: &SessionReader,
    header: &SessionHeader,
    frame_bytes_max: usize,
    mut visit: impl FnMut(TraceRecord) -> Result<()>,
) -> Result<(usize, i64)> {
    visit(TraceRecord {
        timestamp_ms: 0,
        session_id: header.session_id,
        kind: TraceKind::SessionCreated,
        event_ref: format!("branch:{}", header.branch_id),
    })?;
    let Some(tip) = reader.selected_tip(frame_bytes_max)? else {
        return Ok((0, 0));
    };
    let mut count = 0;
    let mut timestamp_ms = 0;
    // `(call_id, has_result)`: exactly one live tool lifecycle.
    // ToolSession serializes effects and refuses subsequent effects after a
    // completion-recording failure, so another lifecycle closes the prior one.
    let mut pending = None;
    reader.replay_at(tip, frame_bytes_max, |entry| {
        timestamp_ms = entry.timestamp_ms;
        count += 1;
        match entry.kind {
            EntryKind::ToolCall { call_id } => {
                flush_incomplete(
                    header.session_id,
                    entry.timestamp_ms,
                    &mut pending,
                    &mut visit,
                )?;
                visit(record_for(header.session_id, &entry)?)?;
                pending = Some((call_id, false));
            }
            EntryKind::ToolResult { call_id, .. } => {
                if pending
                    .is_some_and(|(pending_id, has_result)| pending_id != call_id || has_result)
                {
                    flush_incomplete(
                        header.session_id,
                        entry.timestamp_ms,
                        &mut pending,
                        &mut visit,
                    )?;
                }
                visit(record_for(header.session_id, &entry)?)?;
                pending = Some((call_id, true));
            }
            EntryKind::Meta { .. } if mutation_call_id(&entry)? == pending.map(|(id, _)| id) => {
                visit(record_for(header.session_id, &entry)?)?;
                pending = None;
            }
            EntryKind::Meta { .. } | EntryKind::Message { .. } => {
                visit(record_for(header.session_id, &entry)?)?;
            }
        }
        Ok(())
    })?;
    flush_incomplete(header.session_id, timestamp_ms, &mut pending, &mut visit)?;
    Ok((count, timestamp_ms))
}

fn flush_incomplete(
    session_id: SessionId,
    timestamp_ms: i64,
    pending: &mut Option<(smith::id::ToolCallId, bool)>,
    visit: &mut impl FnMut(TraceRecord) -> Result<()>,
) -> Result<()> {
    if let Some((call_id, _)) = pending.take() {
        visit(TraceRecord {
            timestamp_ms,
            session_id,
            kind: TraceKind::ToolMutationCompleted,
            event_ref: format!("mutation:{call_id}:unknown"),
        })?;
    }
    Ok(())
}

fn mutation_call_id(entry: &EntryFrame) -> Result<Option<ToolCallId>> {
    let EntryContent::Meta { kind, detail } = EntryContent::decode(&entry.content)? else {
        return Ok(None);
    };
    if kind != TOOL_MUTATION_KIND {
        return Ok(None);
    }
    let detail: MutationDetail =
        detail
            .deserialize_as()
            .map_err(|error| smith::error::SmithError::Session {
                code: smith::error::SessionCode::EntryDecode,
                message: format!("invalid tool mutation state metadata: {error}"),
            })?;
    Ok(Some(detail.call_id))
}

#[derive(Deserialize)]
struct MutationDetail {
    call_id: ToolCallId,
    mutation_state: smith::error::MutationState,
}

fn record_for(session_id: SessionId, entry: &EntryFrame) -> Result<TraceRecord> {
    let (kind, event_ref) = match entry.kind {
        EntryKind::Message { id, .. } => (TraceKind::MessageStreamed, format!("message:{id}")),
        EntryKind::ToolCall { call_id } => (TraceKind::ToolCallEmitted, format!("call:{call_id}")),
        EntryKind::ToolResult { call_id, ok } => (
            TraceKind::ToolResultReceived,
            format!("result:{call_id}:{}", if ok { "ok" } else { "failed" }),
        ),
        EntryKind::Meta { compaction: true } => (
            TraceKind::Compaction,
            format!("meta:{COMPACTION_KIND}:{}", entry.id),
        ),
        EntryKind::Meta { compaction: false } => match EntryContent::decode(&entry.content)? {
            EntryContent::Meta { kind, detail } if kind == TOOL_MUTATION_KIND => {
                let detail: MutationDetail =
                    detail
                        .deserialize_as()
                        .map_err(|error| smith::error::SmithError::Session {
                            code: smith::error::SessionCode::EntryDecode,
                            message: format!("invalid tool mutation state metadata: {error}"),
                        })?;
                (
                    TraceKind::ToolMutationCompleted,
                    format!("mutation:{}:{}", detail.call_id, detail.mutation_state),
                )
            }
            _ => (TraceKind::EntryAdded, format!("meta:{}", entry.id)),
        },
    };
    Ok(TraceRecord {
        timestamp_ms: entry.timestamp_ms,
        session_id,
        kind,
        event_ref,
    })
}

#[cfg(test)]
#[expect(
    clippy::unwrap_used,
    reason = "tests may panic on invariant violations"
)]
mod tests {
    use super::*;
    use crate::session::EntryContent;
    use crate::store::{SessionReader, open_session};
    use crate::tools::ToolSession;
    use smith::config::Config;
    use smith::error::{MutationState, SmithError, ToolCode};
    use smith::id::ToolCallId;
    use smith::message::{Message, Role};
    use smith::{StructuredData, ToolArguments};

    fn kinds(trace: &[TraceRecord]) -> Vec<TraceKind> {
        trace.iter().map(|record| record.kind).collect()
    }

    fn replayed(reader: &SessionReader, frame_bytes_max: usize) -> Vec<TraceRecord> {
        let mut records = Vec::new();
        replay(reader, frame_bytes_max, |record| {
            records.push(record);
            Ok(())
        })
        .unwrap();
        records
    }

    #[test]
    fn durable_replay_and_fork_keep_pre_compaction_history() {
        let dir =
            std::env::temp_dir().join(format!("smith_trace_compaction_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("source.smh");
        let fork_path = dir.join("fork.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);
        let message = tools
            .record(&EntryContent::Message(Message::with_text(
                Role::User,
                "before",
            )))
            .unwrap();
        let call = tools
            .record(&EntryContent::ToolCall {
                name: "read".to_string(),
                input: ToolArguments::from_serializable(&serde_json::json!({"path": "note.txt"}))
                    .unwrap(),
                call_id: ToolCallId::from_u128(1),
            })
            .unwrap();
        let result = tools
            .record(&EntryContent::ToolResult {
                ok: true,
                output: "note".to_string(),
                call_id: ToolCallId::from_u128(1),
            })
            .unwrap();
        let compaction = tools
            .record(&EntryContent::Meta {
                kind: COMPACTION_KIND.to_string(),
                detail: StructuredData::from_serializable(
                    &serde_json::json!({"summary": "before"}),
                )
                .unwrap(),
            })
            .unwrap();
        let after = tools
            .record(&EntryContent::Message(Message::with_text(
                Role::Assistant,
                "after",
            )))
            .unwrap();
        assert_eq!(tools.session.active_branch().entries().len(), 2);

        let reader = SessionReader::open(&path).unwrap();
        let mut records = Vec::new();
        replay(&reader, config.frame_bytes_max, |record| {
            records.push(record);
            Ok(())
        })
        .unwrap();
        assert_eq!(
            kinds(&records),
            vec![
                TraceKind::ReplayStarted,
                TraceKind::SessionCreated,
                TraceKind::MessageStreamed,
                TraceKind::ToolCallEmitted,
                TraceKind::ToolResultReceived,
                TraceKind::Compaction,
                TraceKind::MessageStreamed,
                TraceKind::ToolMutationCompleted,
                TraceKind::ReplayCompleted,
            ]
        );
        assert_eq!(
            records
                .iter()
                .filter(|record| record.kind == TraceKind::ReplayStarted)
                .count(),
            1
        );
        assert_eq!(
            records
                .iter()
                .filter(|record| record.kind == TraceKind::ReplayCompleted)
                .count(),
            1
        );
        assert_eq!(records.last().unwrap().timestamp_ms, after.timestamp_ms);
        assert_eq!(records.last().unwrap().event_ref, "entries:5");

        let forked = reader
            .fork_at(after.id, &fork_path, config.clone())
            .unwrap();
        assert_eq!(forked.session.selected_id(), Some(after.id));
        let fork_reader = SessionReader::open(&fork_path).unwrap();
        let mut forked_ids = Vec::new();
        fork_reader
            .replay_at(after.id, config.frame_bytes_max, |entry| {
                forked_ids.push(entry.id);
                Ok(())
            })
            .unwrap();
        assert_eq!(
            forked_ids,
            vec![message.id, call.id, result.id, compaction.id, after.id]
        );
        drop(forked);
        drop(tools);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn replay_marks_result_without_durable_completion_unknown_after_third_append_fails() {
        let dir = std::env::temp_dir().join(format!(
            "smith_trace_incomplete_completion_{}",
            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 writer = opened.writer;
        let mut appends = 0;
        let mut tools = ToolSession::with_persist(&dir, opened.session, move |entry, bytes| {
            appends += 1;
            if appends == 3 {
                return Err(SmithError::Tool {
                    code: ToolCode::TestRecord,
                    message: "third append rejected".to_string(),
                });
            }
            writer.append_entry(entry, bytes)
        })
        .with_config(config.clone());
        let call_id = ToolCallId::from_u128(7);

        let error = tools
            .invoke_local(
                call_id,
                "write",
                &ToolArguments::from_serializable(
                    &serde_json::json!({"path": "made.txt", "content": "made"}),
                )
                .unwrap(),
                &ToolArguments::from_serializable(
                    &serde_json::json!({"path": "made.txt", "content": "made"}),
                )
                .unwrap(),
                str::to_string,
            )
            .unwrap_err();
        assert!(matches!(
            error,
            SmithError::RecordFailed {
                mutation_state: MutationState::Applied,
                ..
            }
        ));
        assert_eq!(
            std::fs::read_to_string(dir.join("made.txt")).unwrap(),
            "made"
        );
        drop(tools);

        let reader = SessionReader::open(&path).unwrap();
        let tip = reader
            .selected_tip(config.frame_bytes_max)
            .unwrap()
            .unwrap();
        let mut entries = Vec::new();
        reader
            .replay_at(tip, config.frame_bytes_max, |entry| {
                entries.push(entry);
                Ok(())
            })
            .unwrap();
        assert_eq!(entries.len(), 2);
        let header = reader.header(config.frame_bytes_max).unwrap();
        let records = replayed(&reader, config.frame_bytes_max);
        assert_eq!(
            records,
            vec![
                TraceRecord {
                    timestamp_ms: 0,
                    session_id: header.session_id,
                    kind: TraceKind::ReplayStarted,
                    event_ref: format!("session:{}", header.session_id),
                },
                TraceRecord {
                    timestamp_ms: 0,
                    session_id: header.session_id,
                    kind: TraceKind::SessionCreated,
                    event_ref: format!("branch:{}", header.branch_id),
                },
                TraceRecord {
                    timestamp_ms: entries[0].timestamp_ms,
                    session_id: header.session_id,
                    kind: TraceKind::ToolCallEmitted,
                    event_ref: format!("call:{call_id}"),
                },
                TraceRecord {
                    timestamp_ms: entries[1].timestamp_ms,
                    session_id: header.session_id,
                    kind: TraceKind::ToolResultReceived,
                    event_ref: format!("result:{call_id}:ok"),
                },
                TraceRecord {
                    timestamp_ms: entries[1].timestamp_ms,
                    session_id: header.session_id,
                    kind: TraceKind::ToolMutationCompleted,
                    event_ref: format!("mutation:{call_id}:unknown"),
                },
                TraceRecord {
                    timestamp_ms: entries[1].timestamp_ms,
                    session_id: header.session_id,
                    kind: TraceKind::ReplayCompleted,
                    event_ref: "entries:2".to_string(),
                },
            ]
        );
        drop(reader);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn replay_marks_call_without_result_unknown() {
        let dir = std::env::temp_dir().join(format!(
            "smith_trace_incomplete_call_{}",
            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 session_id = opened.session.id();
        let branch_id = opened.session.active_branch().id();
        let mut tools = ToolSession::with_session(&dir, opened.session, opened.writer);
        let call_id = ToolCallId::from_u128(8);
        let call = tools
            .record(&EntryContent::ToolCall {
                name: "write".to_string(),
                input: ToolArguments::from_serializable(
                    &serde_json::json!({"path": "never-retry.txt", "content": "x"}),
                )
                .unwrap(),
                call_id,
            })
            .unwrap();
        drop(tools);

        let reader = SessionReader::open(&path).unwrap();
        let mut records = Vec::new();
        trace(&reader, config.frame_bytes_max, |record| {
            records.push(record);
            Ok(())
        })
        .unwrap();
        assert_eq!(
            records,
            vec![
                TraceRecord {
                    timestamp_ms: 0,
                    session_id,
                    kind: TraceKind::SessionCreated,
                    event_ref: format!("branch:{branch_id}"),
                },
                TraceRecord {
                    timestamp_ms: call.timestamp_ms,
                    session_id,
                    kind: TraceKind::ToolCallEmitted,
                    event_ref: format!("call:{call_id}"),
                },
                TraceRecord {
                    timestamp_ms: call.timestamp_ms,
                    session_id,
                    kind: TraceKind::ToolMutationCompleted,
                    event_ref: format!("mutation:{call_id}:unknown"),
                },
            ]
        );
        drop(reader);
        assert!(!dir.join("never-retry.txt").exists());
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn replay_looks_past_unrelated_entries_for_completion_metadata() {
        let dir = std::env::temp_dir().join(format!(
            "smith_trace_completion_lookahead_{}",
            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 session_id = opened.session.id();
        let branch_id = opened.session.active_branch().id();
        let mut tools = ToolSession::with_session(&dir, opened.session, opened.writer);
        let call_id = ToolCallId::from_u128(9);
        let call = tools
            .record(&EntryContent::ToolCall {
                name: "write".to_string(),
                input: ToolArguments::from_serializable(
                    &serde_json::json!({"path": "note.txt", "content": "x"}),
                )
                .unwrap(),
                call_id,
            })
            .unwrap();
        let result = tools
            .record(&EntryContent::ToolResult {
                ok: true,
                output: "wrote note.txt".to_string(),
                call_id,
            })
            .unwrap();
        let message = Message::with_text(Role::Assistant, "interleaved");
        let message_id = message.id;
        let interleaved = tools.record(&EntryContent::Message(message)).unwrap();
        let completion = tools
            .record(&EntryContent::Meta {
                kind: TOOL_MUTATION_KIND.to_string(),
                detail: StructuredData::from_serializable(
                    &serde_json::json!({"call_id": call_id, "mutation_state": "applied"}),
                )
                .unwrap(),
            })
            .unwrap();
        drop(tools);

        let reader = SessionReader::open(&path).unwrap();
        let mut records = Vec::new();
        trace(&reader, config.frame_bytes_max, |record| {
            records.push(record);
            Ok(())
        })
        .unwrap();
        assert_eq!(
            records,
            vec![
                TraceRecord {
                    timestamp_ms: 0,
                    session_id,
                    kind: TraceKind::SessionCreated,
                    event_ref: format!("branch:{branch_id}")
                },
                TraceRecord {
                    timestamp_ms: call.timestamp_ms,
                    session_id,
                    kind: TraceKind::ToolCallEmitted,
                    event_ref: format!("call:{call_id}")
                },
                TraceRecord {
                    timestamp_ms: result.timestamp_ms,
                    session_id,
                    kind: TraceKind::ToolResultReceived,
                    event_ref: format!("result:{call_id}:ok")
                },
                TraceRecord {
                    timestamp_ms: interleaved.timestamp_ms,
                    session_id,
                    kind: TraceKind::MessageStreamed,
                    event_ref: format!("message:{message_id}")
                },
                TraceRecord {
                    timestamp_ms: completion.timestamp_ms,
                    session_id,
                    kind: TraceKind::ToolMutationCompleted,
                    event_ref: format!("mutation:{call_id}:applied")
                },
            ]
        );
        drop(reader);
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn replaying_a_persisted_session_is_byte_stable() {
        let dir = std::env::temp_dir().join(format!("smith_trace_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join("trace.smh");
        let cfg = Config::default_valid();
        {
            let opened = open_session(&path, cfg.clone()).unwrap();
            let mut tools = ToolSession::with_session(&dir, opened.session, opened.writer)
                .with_config(cfg.clone());
            tools
                .invoke(
                    "write",
                    &ToolArguments::from_serializable(
                        &serde_json::json!({"path": "note.txt", "content": "hi"}),
                    )
                    .unwrap(),
                )
                .unwrap();
            tools
                .invoke(
                    "read",
                    &ToolArguments::from_serializable(&serde_json::json!({"path": "note.txt"}))
                        .unwrap(),
                )
                .unwrap();
        }
        let first = SessionReader::open(&path).unwrap();
        let second = SessionReader::open(&path).unwrap();
        let left = replayed(&first, cfg.frame_bytes_max);
        let right = replayed(&second, cfg.frame_bytes_max);
        assert_eq!(left, right);
        assert_eq!(
            kinds(&replayed(&first, cfg.frame_bytes_max)),
            vec![
                TraceKind::ReplayStarted,
                TraceKind::SessionCreated,
                TraceKind::ToolCallEmitted,
                TraceKind::ToolResultReceived,
                TraceKind::ToolMutationCompleted,
                TraceKind::ToolCallEmitted,
                TraceKind::ToolResultReceived,
                TraceKind::ToolMutationCompleted,
                TraceKind::ReplayCompleted,
            ]
        );
        let _ = std::fs::remove_dir_all(&dir);
    }
}