clawmates/crates/tc-llm/tests/scripted.rs

use futures::StreamExt;
use serde_json::json;
use tc_llm::{
    ChatMessage, ChatRequest, ChatRole, ContentPart, LlmError, LlmEvent, LlmProvider,
    ScriptedProvider, StopReason,
};

const SCENARIOS: &str = r#"
[[scenario]]
marker = "[[scenario:hello]]"

[[scenario.turns]]
events = [
    { type = "text", text = "Hello! I'm Scout, your research analyst." },
]

[[scenario]]
marker = "[[scenario:tool-time]]"

[[scenario.turns]]
events = [
    { type = "text", text = "Let me check the clock." },
    { type = "tool_use", name = "clock.now", input = {} },
]

[[scenario.turns]]
events = [
    { type = "text", text = "It is exactly noon." },
]
"#;

fn request_with_user_text(text: &str) -> ChatRequest {
    ChatRequest {
        system: "You are Scout.".into(),
        messages: vec![ChatMessage {
            role: ChatRole::User,
            parts: vec![ContentPart::Text(text.into())],
        }],
        tools: vec![],
        model: "scripted".into(),
        max_tokens: 1024,
    }
}

async fn collect(provider: &ScriptedProvider, request: ChatRequest) -> Vec<LlmEvent> {
    let mut stream = provider.stream(request).await.unwrap();
    let mut events = Vec::new();
    while let Some(event) = stream.next().await {
        events.push(event.unwrap());
    }
    events
}

fn joined_text(events: &[LlmEvent]) -> String {
    events
        .iter()
        .filter_map(|e| match e {
            LlmEvent::TextDelta(t) => Some(t.as_str()),
            _ => None,
        })
        .collect()
}

#[tokio::test]
async fn unmatched_prompts_get_a_deterministic_echo() {
    let provider = ScriptedProvider::from_toml(SCENARIOS).unwrap();
    let events = collect(&provider, request_with_user_text("ping")).await;
    assert_eq!(joined_text(&events), "I received: ping");
    assert!(matches!(
        events.last(),
        Some(LlmEvent::Stop(StopReason::EndTurn))
    ));
}

#[tokio::test]
async fn text_is_streamed_as_multiple_deltas() {
    let provider = ScriptedProvider::from_toml(SCENARIOS).unwrap();
    let events = collect(&provider, request_with_user_text("hi [[scenario:hello]]")).await;
    let delta_count = events
        .iter()
        .filter(|e| matches!(e, LlmEvent::TextDelta(_)))
        .count();
    assert!(
        delta_count > 1,
        "expected word-level streaming, got {delta_count}"
    );
    assert_eq!(
        joined_text(&events),
        "Hello! I'm Scout, your research analyst."
    );
}

#[tokio::test]
async fn tool_scenario_emits_tool_use_then_continues_after_result() {
    let provider = ScriptedProvider::from_toml(SCENARIOS).unwrap();

    // First leg: the scripted model asks for the clock tool.
    let first = collect(
        &provider,
        request_with_user_text("what time is it? [[scenario:tool-time]]"),
    )
    .await;
    assert_eq!(joined_text(&first), "Let me check the clock.");
    let tool_use = first
        .iter()
        .find_map(|e| match e {
            LlmEvent::ToolUse { id, name, .. } => Some((id.clone(), name.clone())),
            _ => None,
        })
        .expect("a tool_use event");
    assert_eq!(tool_use.1, "clock.now");
    assert!(matches!(
        first.last(),
        Some(LlmEvent::Stop(StopReason::ToolUse))
    ));

    // Second leg: request now carries the tool result; the scripted model
    // continues with the next turn.
    let mut request = request_with_user_text("what time is it? [[scenario:tool-time]]");
    request.messages.push(ChatMessage {
        role: ChatRole::Assistant,
        parts: vec![ContentPart::ToolUse {
            id: tool_use.0.clone(),
            name: tool_use.1.clone(),
            input: json!({}),
        }],
    });
    request.messages.push(ChatMessage {
        role: ChatRole::User,
        parts: vec![ContentPart::ToolResult {
            tool_use_id: tool_use.0,
            content: json!({"now": "12:00"}),
        }],
    });
    let second = collect(&provider, request).await;
    assert_eq!(joined_text(&second), "It is exactly noon.");
    assert!(matches!(
        second.last(),
        Some(LlmEvent::Stop(StopReason::EndTurn))
    ));
}

#[tokio::test]
async fn exhausted_turns_fall_back_to_end_turn() {
    let provider = ScriptedProvider::from_toml(SCENARIOS).unwrap();
    let mut request = request_with_user_text("x [[scenario:hello]]");
    // Pretend two tool round-trips already happened: only one turn exists.
    for _ in 0..2 {
        request.messages.push(ChatMessage {
            role: ChatRole::User,
            parts: vec![ContentPart::ToolResult {
                tool_use_id: "t1".into(),
                content: json!({}),
            }],
        });
    }
    let events = collect(&provider, request).await;
    assert!(matches!(
        events.last(),
        Some(LlmEvent::Stop(StopReason::EndTurn))
    ));
}

#[test]
fn invalid_scenario_toml_is_a_clear_error() {
    let err = ScriptedProvider::from_toml("not [valid toml").unwrap_err();
    assert!(matches!(err, LlmError::Scenario(_)));
}