codex-cli-sdk 0.0.1

//! Scenario builder for constructing multi-turn test scenarios.

use crate::testing::mock_transport::MockTransport;
use serde_json::Value;

// ── Exchange ──────────────────────────────────────────────────

/// A single exchange within a scenario: a sequence of events with associated
/// token usage counts.
///
/// Use [`Exchange::new`] for the default token counts (100 input / 0 cached / 50 output),
/// or [`Exchange::with_usage`] to control them precisely — useful for
/// `max_budget_tokens` tests.
pub struct Exchange {
    events: Vec<Value>,
    input_tokens: u64,
    cached_tokens: u64,
    output_tokens: u64,
}

impl Exchange {
    /// Create an exchange with default token counts (100 input / 0 cached / 50 output).
    pub fn new(events: Vec<Value>) -> Self {
        Self {
            events,
            input_tokens: 100,
            cached_tokens: 0,
            output_tokens: 50,
        }
    }

    /// Override the token counts for this exchange.
    pub fn with_usage(mut self, input: u64, cached: u64, output: u64) -> Self {
        self.input_tokens = input;
        self.cached_tokens = cached;
        self.output_tokens = output;
        self
    }
}

// ── ScenarioBuilder ───────────────────────────────────────────

/// Fluent builder for constructing multi-turn test scenarios.
///
/// Automatically wraps exchanges with `thread.started`, `turn.started`, and
/// `turn.completed` bookend events so tests only need to specify the
/// interesting parts.
///
/// # Example
///
/// ```rust,no_run
/// use codex_cli_sdk::testing::{ScenarioBuilder, builders};
/// use codex_cli_sdk::testing::scenario::Exchange;
///
/// let transport = ScenarioBuilder::new("thread-1")
///     .exchange(vec![
///         builders::agent_message_completed("msg-1", "Hello!"),
///     ])
///     .exchange_typed(
///         Exchange::new(vec![builders::agent_message_completed("msg-2", "Turn 2")])
///             .with_usage(200, 0, 500),
///     )
///     .build();
/// ```
pub struct ScenarioBuilder {
    thread_id: String,
    exchanges: Vec<Exchange>,
}

impl ScenarioBuilder {
    /// Create a new scenario with the given thread ID.
    #[must_use]
    pub fn new(thread_id: impl Into<String>) -> Self {
        Self {
            thread_id: thread_id.into(),
            exchanges: Vec::new(),
        }
    }

    /// Add an exchange with default token counts (100 input / 0 cached / 50 output).
    #[must_use]
    pub fn exchange(self, events: Vec<Value>) -> Self {
        self.exchange_typed(Exchange::new(events))
    }

    /// Add an exchange with full control over events and token usage.
    ///
    /// Use this when testing `max_budget_tokens` enforcement, where the exact
    /// `output_tokens` per exchange matters.
    #[must_use]
    pub fn exchange_typed(mut self, exchange: Exchange) -> Self {
        self.exchanges.push(exchange);
        self
    }

    /// Build a [`MockTransport`] pre-loaded with the scenario events.
    ///
    /// The transport is loaded with:
    /// 1. `thread.started` with the scenario's thread ID
    /// 2. For each exchange: `turn.started`, the exchange events, `turn.completed`
    ///    (using the per-exchange token counts)
    #[must_use]
    pub fn build(self) -> MockTransport {
        use crate::testing::builders;

        let transport = MockTransport::new();

        // Thread started
        transport.enqueue_event(builders::thread_started(&self.thread_id));

        // Each exchange gets turn bookends with its own usage counts.
        for exchange in &self.exchanges {
            transport.enqueue_event(builders::turn_started());
            for event in &exchange.events {
                transport.enqueue_event(event.clone());
            }
            transport.enqueue_event(builders::turn_completed(
                exchange.input_tokens,
                exchange.cached_tokens,
                exchange.output_tokens,
            ));
        }

        transport
    }

    /// Build a [`MockTransport`] without the auto-generated bookend events.
    ///
    /// Use this when you want full control over the event sequence.
    #[must_use]
    pub fn build_raw(self) -> MockTransport {
        let transport = MockTransport::new();
        for exchange in &self.exchanges {
            for event in &exchange.events {
                transport.enqueue_event(event.clone());
            }
        }
        transport
    }
}

// ── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::testing::builders;

    #[test]
    fn empty_build_has_thread_started_only() {
        let transport = ScenarioBuilder::new("t1").build();
        // thread.started only (no exchanges → no turn bookends)
        assert_eq!(transport.queued_count(), 1);
    }

    #[test]
    fn single_exchange() {
        let transport = ScenarioBuilder::new("t1")
            .exchange(vec![builders::agent_message_completed("m1", "Hello!")])
            .build();
        // thread.started + turn.started + agent_message + turn.completed = 4
        assert_eq!(transport.queued_count(), 4);
    }

    #[test]
    fn multiple_exchanges() {
        let transport = ScenarioBuilder::new("t1")
            .exchange(vec![builders::agent_message_completed("m1", "Turn 1")])
            .exchange(vec![builders::agent_message_completed("m2", "Turn 2")])
            .build();
        // thread.started + 2 * (turn.started + message + turn.completed) = 1 + 6 = 7
        assert_eq!(transport.queued_count(), 7);
    }

    #[test]
    fn raw_mode_no_bookends() {
        let transport = ScenarioBuilder::new("t1")
            .exchange(vec![builders::agent_message_completed("m1", "raw")])
            .build_raw();
        // Only the exchange event, no bookends
        assert_eq!(transport.queued_count(), 1);
    }

    #[test]
    fn exchange_typed_custom_usage_accepted() {
        // Smoke test: exchange_typed with custom usage builds without panic.
        let transport = ScenarioBuilder::new("t1")
            .exchange_typed(
                Exchange::new(vec![builders::agent_message_completed("m1", "big turn")])
                    .with_usage(500, 100, 1000),
            )
            .build();
        // thread.started + turn.started + message + turn.completed = 4
        assert_eq!(transport.queued_count(), 4);
    }
}