dirigent/crates/dirigent_core/tests/replay_test.rs

//! Integration test: replay archived session into a `MockStream`.
//!
//! Builds a single-backend in-memory (tempdir) archivist, registers a
//! session, appends 10 messages with ascending timestamps, then exercises
//! `replay_session_to_stream` end-to-end.

use std::sync::Arc;

use chrono::{Duration as ChronoDuration, Utc};
use uuid::Uuid;

use dirigent_archivist::{
    Archivist, MessageRecord, RegisterConnectorRequest, RegisterSessionRequest,
    backends::JsonlBackend,
};
use dirigent_core::sharing::{
    MockStream,
    replay::{ReplayOptions, ReplaySpeed, replay_session_to_stream},
};
use dirigent_protocol::streaming::{EventOrigin, SessionStream, StreamScope};

/// Build an in-memory-ish archivist backed by a tempdir + JsonlBackend.
///
/// Matches the pattern used by `dirigent_archivist/tests/integration_tests.rs`.
/// The tempdir is leaked for the duration of the test process — acceptable
/// because the test binary exits immediately after.
async fn build_in_memory_archivist() -> Arc<Archivist> {
    let temp_dir = std::env::temp_dir().join(format!("core_replay_test_{}", Uuid::now_v7()));
    let backend = Arc::new(
        JsonlBackend::new(temp_dir.clone())
            .await
            .expect("JsonlBackend construction"),
    );
    let archivist = Archivist::from_single_backend("main".into(), backend)
        .await
        .expect("Archivist::from_single_backend");
    Arc::new(archivist)
}

/// Register a fresh connector + session and append `n` messages with
/// timestamps one second apart. Returns the scroll_id.
async fn seed_session_with_messages(archivist: &Archivist, n: usize) -> Uuid {
    let connector_resp = archivist
        .register_connector(
            RegisterConnectorRequest {
                r#type: "OpenCode".to_string(),
                title: "Replay Test Connector".to_string(),
                client_native_id: format!("replay-test@{}", Uuid::now_v7()),
                custom_uid: None,
                metadata: serde_json::json!({}),
                fingerprint: None,
            },
            None,
        )
        .await
        .expect("register_connector");

    let session_resp = archivist
        .register_session(
            RegisterSessionRequest {
                connector_uid: connector_resp.connector_uid,
                native_session_id: format!("native-{}", Uuid::now_v7()),
                title: Some("Replay Test Session".to_string()),
                custom_scroll_id: None,
                metadata: serde_json::json!({}),
                completeness: Default::default(),
                parent_scroll_id: None,
                is_subagent: false,
                continuation: None,
                agent_id: None,
                subagent_type: None,
                spawning_tool_use_id: None,
            },
            None,
        )
        .await
        .expect("register_session");

    let scroll_id = session_resp.scroll_id;
    let base_ts = Utc::now();

    let messages: Vec<MessageRecord> = (0..n)
        .map(|i| {
            let role = if i % 2 == 0 { "user" } else { "assistant" };
            MessageRecord {
                version: 1,
                message_id: Uuid::now_v7(),
                session: scroll_id,
                parent_id: None,
                ts: base_ts + ChronoDuration::seconds(i as i64),
                role: role.to_string(),
                author: None,
                content_md: format!("message {i}"),
                content_parts: None,
                attachments: vec![],
                metadata: serde_json::json!({}),
            }
        })
        .collect();

    archivist
        .append_messages(scroll_id, messages, None)
        .await
        .expect("append_messages");

    scroll_id
}

#[tokio::test]
async fn replay_delivers_archived_messages_to_stream() {
    let archivist = build_in_memory_archivist().await;
    let scroll_id = seed_session_with_messages(&archivist, 10).await;

    let mock = MockStream::new("mock", StreamScope::Session { scroll_id });
    let stream: Arc<dyn SessionStream> = mock.clone();

    let report = replay_session_to_stream(
        archivist.as_ref(),
        scroll_id,
        stream,
        ReplayOptions {
            include_meta_events: false,
            speed: ReplaySpeed::AsFastAsPossible,
        },
    )
    .await
    .expect("replay_session_to_stream");

    assert_eq!(report.events_sent, 10, "events_sent");
    assert_eq!(report.failures, 0, "failures");
    assert_eq!(mock.received_count(), 10, "mock received count");

    let received = mock.received.lock().unwrap();
    for evt in received.iter() {
        assert!(
            matches!(evt.origin, EventOrigin::Replay { .. }),
            "every replayed event must carry EventOrigin::Replay"
        );
        assert_eq!(
            evt.routing.scroll_id,
            Some(scroll_id),
            "every replayed event must carry the authoritative scroll_id"
        );
    }
}

#[tokio::test]
async fn replay_continues_on_stream_failure() {
    let archivist = build_in_memory_archivist().await;
    let scroll_id = seed_session_with_messages(&archivist, 10).await;

    let mock = MockStream::new("mock", StreamScope::Session { scroll_id });
    mock.fail_next(3);
    let stream: Arc<dyn SessionStream> = mock.clone();

    let report = replay_session_to_stream(
        archivist.as_ref(),
        scroll_id,
        stream,
        ReplayOptions {
            include_meta_events: false,
            speed: ReplaySpeed::AsFastAsPossible,
        },
    )
    .await
    .expect("replay_session_to_stream");

    // events_sent counts attempted (ok + failed); failures counts Failed only.
    assert_eq!(report.events_sent, 10, "events_sent counts every attempt");
    assert_eq!(report.failures, 3, "first 3 events rejected by mock");
    assert_eq!(
        mock.received_count(),
        7,
        "mock buffer contains the 7 successful events"
    );
}