clawmates/crates/cm-llm/src/scripted.rs

//! The deterministic scenario provider.
//!
//! A real, config-selectable production provider (`provider = "scripted"`):
//! it powers TDD, Playwright E2E, and the air-gapped smoke-test mode, so
//! tests exercise the exact seam production uses. Scenarios are TOML; a
//! prompt that contains a scenario's `marker` plays that scenario's turns,
//! anything else gets a deterministic echo.

use futures::stream;
use serde::Deserialize;
use serde_json::Value;

use crate::provider::{
    ChatRequest, ContentPart, EventStream, LlmError, LlmEvent, LlmProvider, StopReason,
};

#[derive(Debug, Deserialize)]
struct ScenarioFile {
    #[serde(default)]
    scenario: Vec<Scenario>,
}

#[derive(Debug, Deserialize)]
struct Scenario {
    marker: String,
    #[serde(default)]
    turns: Vec<Turn>,
}

#[derive(Debug, Deserialize)]
struct Turn {
    #[serde(default)]
    events: Vec<ScriptedEvent>,
}

#[derive(Debug, Clone, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum ScriptedEvent {
    Text {
        text: String,
    },
    ToolUse {
        name: String,
        #[serde(default)]
        input: Value,
    },
}

#[derive(Debug)]
pub struct ScriptedProvider {
    scenarios: Vec<Scenario>,
}

impl ScriptedProvider {
    pub fn from_toml(toml_source: &str) -> Result<ScriptedProvider, LlmError> {
        let file: ScenarioFile =
            toml::from_str(toml_source).map_err(|e| LlmError::Scenario(e.to_string()))?;
        Ok(ScriptedProvider {
            scenarios: file.scenario,
        })
    }

    pub fn from_path(path: &std::path::Path) -> Result<ScriptedProvider, LlmError> {
        let source = std::fs::read_to_string(path)
            .map_err(|e| LlmError::Scenario(format!("read {}: {e}", path.display())))?;
        ScriptedProvider::from_toml(&source)
    }

    /// Streams a fixed text as word-level deltas to exercise real streaming
    /// behavior in every consumer.
    fn text_deltas(text: &str, out: &mut Vec<Result<LlmEvent, LlmError>>) {
        let mut rest = text;
        while !rest.is_empty() {
            let cut = rest
                .char_indices()
                .skip_while(|(_, c)| *c == ' ')
                .find(|(_, c)| *c == ' ')
                .map(|(i, _)| i)
                .unwrap_or(rest.len());
            let (chunk, tail) = rest.split_at(cut.max(1));
            out.push(Ok(LlmEvent::TextDelta(chunk.to_owned())));
            rest = tail;
        }
    }
}

#[async_trait::async_trait]
impl LlmProvider for ScriptedProvider {
    async fn stream(&self, request: ChatRequest) -> Result<EventStream, LlmError> {
        // Scenario selection keys on the MOST RECENT text carrying any
        // marker: earlier turns keep their markers in session history, and
        // the latest user intent must win.
        let prompt_text: String = request
            .messages
            .iter()
            .rev()
            .flat_map(|m| m.parts.iter())
            .filter_map(|p| match p {
                ContentPart::Text { text } => Some(text.as_str()),
                _ => None,
            })
            .find(|text| self.scenarios.iter().any(|s| text.contains(&s.marker)))
            .unwrap_or_default()
            .to_owned();
        // Which leg of a multi-tool conversation is this? One ToolResult in
        // the request means turn 0 already played; play turn 1, and so on.
        let turn_index = request
            .messages
            .iter()
            .flat_map(|m| m.parts.iter())
            .filter(|p| matches!(p, ContentPart::ToolResult { .. }))
            .count();

        let mut events: Vec<Result<LlmEvent, LlmError>> = Vec::new();
        let scenario = self
            .scenarios
            .iter()
            .find(|s| prompt_text.contains(&s.marker));

        match scenario {
            Some(scenario) => {
                match scenario.turns.get(turn_index) {
                    Some(turn) => {
                        let mut stopped_for_tool = false;
                        for (i, event) in turn.events.iter().enumerate() {
                            match event {
                                ScriptedEvent::Text { text } => {
                                    Self::text_deltas(text, &mut events);
                                }
                                ScriptedEvent::ToolUse { name, input } => {
                                    events.push(Ok(LlmEvent::ToolUse {
                                        id: format!("scripted-tool-{turn_index}-{i}"),
                                        name: name.clone(),
                                        input: input.clone(),
                                    }));
                                    stopped_for_tool = true;
                                }
                            }
                        }
                        events.push(Ok(LlmEvent::Stop(if stopped_for_tool {
                            StopReason::ToolUse
                        } else {
                            StopReason::EndTurn
                        })));
                    }
                    // More tool round-trips than scripted turns: end cleanly.
                    None => events.push(Ok(LlmEvent::Stop(StopReason::EndTurn))),
                }
            }
            None => {
                let last_user_text = request
                    .messages
                    .iter()
                    .rev()
                    .flat_map(|m| m.parts.iter())
                    .find_map(|p| match p {
                        ContentPart::Text { text } => Some(text.clone()),
                        _ => None,
                    })
                    .unwrap_or_default();
                Self::text_deltas(&format!("I received: {last_user_text}"), &mut events);
                events.push(Ok(LlmEvent::Stop(StopReason::EndTurn)));
            }
        }

        // Deterministic accounting: one "token" per whitespace word in and
        // out, so billing tests can predict exact charges.
        let input_tokens = request
            .messages
            .iter()
            .flat_map(|m| m.parts.iter())
            .filter_map(|p| match p {
                ContentPart::Text { text } => Some(text.split_whitespace().count()),
                _ => None,
            })
            .sum::<usize>() as u32;
        let output_tokens = events
            .iter()
            .filter_map(|e| match e {
                Ok(LlmEvent::TextDelta(t)) => Some(t.split_whitespace().count()),
                _ => None,
            })
            .sum::<usize>() as u32;
        let stop_index = events.len().saturating_sub(1);
        events.insert(
            stop_index,
            Ok(LlmEvent::Usage {
                input_tokens,
                output_tokens,
            }),
        );
        Ok(Box::pin(stream::iter(events)))
    }
}