smith

//! Shared fixtures for Smith benchmarks (`SMH-SPEC-SPEC0001`, Performance
//! and portability).
//!
//! Benches live in `benches/`: `instructions` counts callgrind instructions
//! against the pins in `budgets.toml` (`cargo x perf`), `wall_clock` runs
//! criterion for local A/B comparisons (`cargo x bench`). Builders here run
//! in bench setup, outside every measured region, and use only public API,
//! so a bench measures what a caller of the crate gets.

#![forbid(unsafe_code)]

use futures::StreamExt;
use smith::error::Result;
use smith::id::{MessageId, ToolCallId};
use smith::message::{Message, Role};
use smith::stream::{StopReason, StreamEvent};
use smith::{StreamFn, Usage};
use smith_core::agent::Agent;
use smith_core::frame::Frame;
use smith_core::session::{EntryContent, EntryFrame, Session};
use smith_core::tools::ToolSession;
use std::sync::Arc;

/// Plaintext the agent fixture registers as a secret.
pub const SECRET: &str = "sk-bench-0123456789abcdef";

/// A session of `len` chat messages, alternating user and assistant.
///
/// # Errors
///
/// Returns the session error of the first failed append.
pub fn linear_session(len: usize) -> Result<Session> {
    let mut session = Session::new();
    for turn in 0..len {
        let role = if turn % 2 == 0 {
            Role::User
        } else {
            Role::Assistant
        };
        let message = Message::with_text(role, format!("turn {turn}: a short chat message"));
        session.append(&EntryContent::Message(message))?;
    }
    Ok(session)
}

/// A session of `pairs` tool calls, each followed by its result.
///
/// # Errors
///
/// Returns the session error of the first failed append.
pub fn tool_call_session(pairs: usize) -> Result<Session> {
    let mut session = Session::new();
    for call in 0..pairs {
        let call_id = ToolCallId::new();
        session.append(&EntryContent::ToolCall {
            name: "read".to_string(),
            input: serde_json::json!({"path": format!("src/module_{call}.rs"), "limit": 200}),
            call_id,
        })?;
        session.append(&EntryContent::ToolResult {
            ok: true,
            output: format!("pub fn item_{call}() -> u32 {{ {call} }}\n"),
            call_id,
        })?;
    }
    Ok(session)
}

/// The decoded frames a store would read back for `session`: its header,
/// then one entry frame per record of the active branch.
#[must_use]
pub fn frames_of(session: &Session) -> Vec<Frame> {
    let entries = session.active_branch().entries();
    let mut frames = Vec::with_capacity(entries.len() + 1);
    frames.push(Frame::Header {
        header: session.header(),
    });
    frames.extend(entries.iter().map(|entry| Frame::Known {
        entry: EntryFrame {
            id: entry.id,
            parent: entry.parent,
            timestamp_ms: entry.timestamp_ms,
            kind: entry.kind,
            content: session.bytes(entry).to_vec(),
        },
    }));
    frames
}

/// A JSON value `depth` levels deep, four keys per object, every object
/// in reverse key order so canonicalization has to reorder all of them.
#[must_use]
pub fn nested_json(depth: u32) -> serde_json::Value {
    let mut object = serde_json::Map::new();
    for key in (0..4u32).rev() {
        let value = match (depth, key) {
            (0, _) => serde_json::Value::from(format!("leaf {key}")),
            (_, 0) => serde_json::json!([key, "item", {"z": depth, "a": key}]),
            _ => nested_json(depth - 1),
        };
        object.insert(format!("key_{key}"), value);
    }
    serde_json::Value::Object(object)
}

/// An agent over a session of `history` messages with [`SECRET`]
/// registered, whose provider answers every request with one text delta
/// and a clean stop.
///
/// # Errors
///
/// Returns the session error of building the history.
pub fn masked_agent(history: usize) -> Result<Agent> {
    let mut tools = ToolSession::new(".");
    tools.session = linear_session(history)?;
    let mut agent = Agent::new(scripted_reply(), tools, "bench-model");
    agent.secrets_mut().register("api_key", SECRET);
    Ok(agent)
}

/// A provider that streams `ok` and stops, whatever the request.
fn scripted_reply() -> StreamFn {
    Arc::new(|_request, _cancel| {
        let reply = MessageId::new();
        futures::stream::iter([
            Ok(StreamEvent::text_delta(reply, "ok")),
            Ok(StreamEvent::stop(StopReason::EndTurn, Some(Usage::new()))),
        ])
        .boxed()
    })
}

/// SSE bytes: one `data:` event per payload.
fn sse(payloads: &[&str]) -> Vec<u8> {
    let mut out = String::new();
    for payload in payloads {
        out.push_str("data: ");
        out.push_str(payload);
        out.push_str("\n\n");
    }
    out.into_bytes()
}

/// An `OpenAI` chat-completions stream: text, fragmented tool arguments,
/// usage, done; the shapes of the adapter's conformance fixtures.
#[must_use]
pub fn openai_sse() -> Vec<u8> {
    sse(&[
        r#"{"choices":[{"delta":{"role":"assistant","content":""}}]}"#,
        r#"{"choices":[{"delta":{"content":"Hello"}}]}"#,
        r#"{"choices":[{"delta":{"content":" world"}}]}"#,
        r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"a","function":{"name":"write","arguments":""}}]}}]}"#,
        r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"path\""}}]}}]}"#,
        r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":":\"f.txt\",\"content\":\"hi\"}"}}]}}]}"#,
        r#"{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}"#,
        r#"{"choices":[],"usage":{"prompt_tokens":9,"completion_tokens":2,"prompt_tokens_details":{"cached_tokens":4}}}"#,
        "[DONE]",
    ])
}

/// An Anthropic Messages stream: text, thinking, assembled tool input,
/// usage; the adapter's conformance fixture.
#[must_use]
pub fn anthropic_sse() -> Vec<u8> {
    sse(&[
        r#"{"type":"message_start","message":{"usage":{"input_tokens":12}}}"#,
        r#"{"type":"content_block_start","index":0,"content_block":{"type":"text"}}"#,
        r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Look"}}"#,
        r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ing"}}"#,
        r#"{"type":"content_block_stop","index":0}"#,
        r#"{"type":"content_block_start","index":1,"content_block":{"type":"thinking"}}"#,
        r#"{"type":"content_block_delta","index":1,"delta":{"type":"thinking_delta","thinking":"hmm"}}"#,
        r#"{"type":"content_block_stop","index":1}"#,
        r#"{"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"t","name":"write"}}"#,
        r#"{"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"path\""}}"#,
        r#"{"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":":[1,2]}"}}"#,
        r#"{"type":"content_block_stop","index":2}"#,
        r#"{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":7,"cache_read_input_tokens":3}}"#,
        r#"{"type":"message_stop"}"#,
    ])
}

/// A Gemini stream: text around a complete function call, usage, stop;
/// the adapter's conformance fixture.
#[must_use]
pub fn gemini_sse() -> Vec<u8> {
    sse(&[
        r#"{"candidates":[{"content":{"parts":[{"text":"Writing"}]}}]}"#,
        r#"{"candidates":[{"content":{"parts":[{"functionCall":{"name":"write","args":{"path":"f","content":"hi"}}}]}}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":9,"cachedContentTokenCount":2}}"#,
        r#"{"candidates":[{"content":{"parts":[{"text":" now"}]}}]}"#,
        r#"{"candidates":[{"content":{},"finishReason":"STOP"}]}"#,
    ])
}

#[cfg(test)]
#[expect(
    clippy::unwrap_used,
    reason = "tests may panic on invariant violations"
)]
mod tests {
    use super::*;
    use smith_ai::{AnthropicDecoder, GeminiDecoder, OpenAiDecoder, StreamDecoder};

    #[test]
    fn fixtures_have_the_sizes_benches_name() {
        let linear = linear_session(1000).unwrap();
        assert_eq!(linear.active_branch().entries().len(), 1000);
        assert_eq!(frames_of(&linear).len(), 1001);
        let calls = tool_call_session(500).unwrap();
        assert_eq!(calls.active_branch().entries().len(), 1000);
    }

    #[test]
    fn nested_json_is_not_canonical() {
        let value = nested_json(2);
        let keys: Vec<&String> = value.as_object().unwrap().keys().collect();
        assert_eq!(keys, ["key_3", "key_2", "key_1", "key_0"]);
    }

    #[test]
    fn every_sse_fixture_decodes_to_a_clean_stop() {
        let cases: [(Box<dyn StreamDecoder>, Vec<u8>); 3] = [
            (Box::new(OpenAiDecoder::new()), openai_sse()),
            (Box::new(AnthropicDecoder::new()), anthropic_sse()),
            (Box::new(GeminiDecoder::new()), gemini_sse()),
        ];
        for (mut decoder, bytes) in cases {
            let mut events = decoder.push(&bytes).unwrap();
            events.extend(decoder.finish().unwrap());
            assert!(
                matches!(events.last(), Some(StreamEvent::Stop { .. })),
                "{events:?}"
            );
        }
    }

    #[test]
    fn masked_agent_answers_a_turn() {
        let mut agent = masked_agent(10).unwrap();
        agent.run_turn(format!("use {SECRET}")).unwrap();
        let entries = agent.session().active_branch().entries();
        assert_eq!(entries.len(), 12);
        let input = agent.session().content(&entries[10]).unwrap();
        assert!(!format!("{input:?}").contains(SECRET), "{input:?}");
    }
}