rig-core 0.45.0

use super::*;
use crate::message::{AssistantContent, Message};
use crate::streaming::{Item, StreamEvent};
use crate::test_utils::history::{assert_every_variant, assert_restated_agrees, decode};
use crate::wire::Mode;
use serde_json::json;

use super::super::{Interactions, create_request_body};

fn wire() -> Interactions {
    crate::providers::gemini::GeminiConfig::new("test-key").interactions("gemini-2.5-pro")
}

fn frame(value: serde_json::Value) -> WireFrame {
    WireFrame::Text(value.to_string())
}

/// A whole interaction resource holding `steps`.
fn whole(steps: Vec<serde_json::Value>) -> WireFrame {
    frame(json!({"id": "int_1", "status": "completed", "steps": steps}))
}

fn start(index: usize, step: serde_json::Value) -> WireFrame {
    frame(json!({"event_type": "step.start", "index": index, "step": step}))
}

fn delta(index: usize, delta: serde_json::Value) -> WireFrame {
    frame(json!({"event_type": "step.delta", "index": index, "delta": delta}))
}

fn stop(index: usize) -> WireFrame {
    frame(json!({"event_type": "step.stop", "index": index}))
}

fn completed() -> WireFrame {
    frame(json!({
        "event_type": "interaction.completed",
        "interaction": {"id": "int_1", "status": "completed"},
    }))
}

/// The steps the decoded turn of `response` encodes to for the same model,
/// with the tools its calls name declared: what follows the user's prompt,
/// without the results the request boundary adds for its calls.
fn replayed(response: &crate::completion::CompletionResponse) -> Vec<serde_json::Value> {
    let turn = response.message().expect("a turn");
    let tools = match &turn {
        Message::Assistant(turn) => turn
            .tool_calls()
            .map(|call| crate::completion::ToolDefinition {
                name: call.function.name.clone(),
                description: "A tool.".to_owned(),
                parameters: json!({"type": "object", "properties": {}}),
                annotations: Default::default(),
            })
            .map(Into::into)
            .collect(),
        _ => Vec::new(),
    };
    let mut request =
        crate::completion::CompletionRequest::from(vec![Message::user("hello"), turn]);
    request.tools = tools;
    let request = <Completion as crate::wire::Operation>::prepare(
        request,
        &crate::wire::Wire::describe(&wire()),
    )
    .expect("the request is valid");
    let body = create_request_body(&wire(), &request, None).expect("the request builds");
    let mut steps = body
        .get("input")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default();
    steps.remove(0);
    steps.retain(|step| step["type"] != "function_result");
    steps
}

/// How a model output content item decodes, by its `type`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ContentKind {
    Text,
    /// An image block, or an opaque one when it states no data or URI.
    Image,
    /// Audio, video, a document, or a type rig does not know.
    Other,
}

fn content_kind(content: &serde_json::Value) -> ContentKind {
    match content.get("type").and_then(serde_json::Value::as_str) {
        Some("text") => ContentKind::Text,
        Some("image") => ContentKind::Image,
        _ => ContentKind::Other,
    }
}

/// How the decoder classifies the frame `data`.
fn classify_interactions_frame(data: &str) -> WireEvent<InteractionsEvent> {
    InteractionsDecoder::default().classify(WireFrame::Text(data.to_owned()))
}

/// The Interactions wire bound to a transport answering with `frames` as
/// one SSE body, driven through the full normalized path.
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
async fn drive_frames(
    frames: &[&str],
) -> (
    Vec<Result<Item<StreamEvent>, String>>,
    Result<crate::completion::CompletionResponse, ProviderError>,
) {
    use futures::StreamExt;

    let sse_bytes = bytes::Bytes::from(
        frames
            .iter()
            .map(|event| format!("data: {event}\n\n"))
            .collect::<String>(),
    );
    let model =
        crate::driver::Model::new(wire(), crate::test_utils::MockStreamingClient { sse_bytes });
    let mut stream = model
        .stream(crate::completion::CompletionRequest::new("hello"))
        .expect("stream should open");
    let mut items = Vec::new();
    while let Some(item) = stream.next().await {
        items.push(item.map_err(|error| error.to_string()));
    }
    (items, stream.finish().await)
}

#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
fn texts_of(items: &[Result<Item<StreamEvent>, String>]) -> Vec<&str> {
    items
        .iter()
        .filter_map(|item| match item {
            Ok(Item::Event(StreamEvent::Text { text, .. })) => Some(text.as_str()),
            _ => None,
        })
        .collect()
}

#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
#[tokio::test]
async fn provider_error_event_ends_the_stream_without_draining_later_frames() {
    let (items, outcome) = drive_frames(&[
        r#"{"event_type":"step.delta","index":0,"delta":{"type":"text","text":"hi"}}"#,
        r#"{"event_type":"error","error":{"code":"internal","message":"boom"}}"#,
        r#"{"event_type":"step.delta","index":0,"delta":{"type":"text","text":"dead"}}"#,
        r#"{"event_type":"something.future","payload":{"x":1}}"#,
        r#"{"event_type":"interaction.completed","interaction":{"id":"int_1","status":"completed"}}"#,
    ])
    .await;
    let error_position = items
        .iter()
        .position(Result::is_err)
        .expect("the provider error must reach the consumer");
    assert_eq!(error_position, items.len() - 1, "got {items:?}");
    assert_eq!(texts_of(&items), ["hi"]);
    assert!(outcome.is_err());
}

/// A stream yields a cursor naming its interaction after
/// `interaction.created` and after each event the API marks with an
/// `event_id`; a stream that resumes an interaction names it without that
/// event, and joins the step that was open.
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
#[tokio::test]
async fn resume_points_yield_cursors() -> anyhow::Result<()> {
    use futures::StreamExt;

    let (items, outcome) = drive_frames(&[
        r#"{"event_type":"interaction.created","interaction":{"id":"int_1","agent":"deep-research-preview-04-2026","status":"in_progress"}}"#,
        r#"{"event_type":"step.start","index":0,"step":{"type":"model_output"}}"#,
        r#"{"event_type":"step.delta","event_id":"e2","index":0,"delta":{"type":"text","text":"hi"}}"#,
        r#"{"event_type":"step.stop","index":0}"#,
        r#"{"event_type":"interaction.completed","interaction":{"id":"int_1","status":"completed"}}"#,
    ])
    .await;
    let cursors: Vec<_> = items
        .iter()
        .filter_map(|item| item.as_ref().ok().and_then(StreamCursor::of))
        .collect();
    let cursor = |event_id: Option<&str>| StreamCursor {
        interaction_id: Some("int_1".to_owned()),
        event_id: event_id.map(str::to_owned),
    };
    anyhow::ensure!(cursors == [cursor(None), cursor(Some("e2"))], "{cursors:?}");
    let other = Item::<StreamEvent>::Unknown(crate::streaming::UnknownPayload::new(
        json!({"type": "other", "event_id": "e9"}),
    ));
    anyhow::ensure!(
        StreamCursor::of(&other).is_none(),
        "only a cursor reads as one"
    );
    anyhow::ensure!(texts_of(&items) == ["hi"]);
    anyhow::ensure!(outcome.is_ok(), "{outcome:?}");

    let resumed_frames = [
        r#"{"event_type":"step.delta","event_id":"e5","index":1,"delta":{"type":"text","text":" again"}}"#,
        r#"{"event_type":"step.stop","index":1}"#,
        r#"{"event_type":"interaction.completed","interaction":{"id":"int_1","status":"completed"}}"#,
    ];
    let sse_bytes = bytes::Bytes::from(
        resumed_frames
            .iter()
            .map(|event| format!("data: {event}\n\n"))
            .collect::<String>(),
    );
    let resumed = crate::driver::Model::new(
        super::super::InteractionResume::new(
            crate::providers::gemini::GeminiConfig::new("test-key"),
            "int_1",
        )
        .after_event("e2"),
        crate::test_utils::MockStreamingClient { sse_bytes },
    );
    let mut stream = resumed.stream(crate::completion::CompletionRequest::new("hello"))?;
    let mut cursors = Vec::new();
    while let Some(item) = stream.next().await {
        cursors.extend(StreamCursor::of(&item?));
    }
    anyhow::ensure!(cursors == [cursor(Some("e5"))], "{cursors:?}");
    let response = stream.finish().await?;
    let text: String = response
        .choice
        .iter()
        .filter_map(|content| match content {
            AssistantContent::Text(text) => Some(text.text.as_str()),
            _ => None,
        })
        .collect();
    anyhow::ensure!(text == " again", "the open step continues: {text:?}");
    Ok(())
}

/// An in-band error carrying an HTTP status code gets the verdict the Gemini
/// REST stream gives the same envelope: a 503 or 429 is retryable.
///
/// Not a cassette test: an overload or quota hit is a transient condition
/// that cannot be recorded on demand, and the defect is the verdict the
/// envelope gets, which a hand-written frame pins exactly.
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
#[tokio::test]
async fn an_in_band_error_with_a_status_code_is_classified_like_the_rest_stream()
-> anyhow::Result<()> {
    use futures::StreamExt;

    for code in [429_u16, 503] {
        let frame = json!({
            "event_type": "error",
            "error": { "code": code, "message": "The model is overloaded.", "status": "UNAVAILABLE" }
        })
        .to_string();
        let model = crate::driver::Model::new(
            wire(),
            crate::test_utils::MockStreamingClient {
                sse_bytes: bytes::Bytes::from(format!("data: {frame}\n\n")),
            },
        );
        let mut stream = model.stream(crate::completion::CompletionRequest::new("hello"))?;
        let mut errors = Vec::new();
        while let Some(item) = stream.next().await {
            if let Err(error) = item {
                errors.push(error);
            }
        }
        anyhow::ensure!(errors.len() == 1, "one terminal failure: {errors:?}");
        anyhow::ensure!(
            errors[0].is_retryable(),
            "code {code} must be retryable: {:?}",
            errors[0]
        );
    }
    Ok(())
}

/// A whole thought, a model output with annotated text, and a hosted
/// search: the unary resource and its restatement as a stream fold into
/// the same turn.
#[test]
fn a_whole_interaction_and_its_stream_agree() {
    let thought = json!({"type": "thought", "signature": "sig", "summary": [{"type": "text", "text": "plan"}]});
    let search = json!({"type": "google_search_call", "id": "call_1", "signature": "s", "arguments": {"queries": ["euro 2024"]}, "search_type": "web_search"});
    let output = json!({"type": "model_output", "content": [{"type": "text", "text": "Spain won.", "annotations": [{"start_index": 0, "end_index": 5, "source": "https://example.com"}]}]});
    assert_restated_agrees(
        &wire(),
        [whole(vec![thought, search.clone(), output])],
        [
            start(0, json!({"type": "thought"})),
            delta(
                0,
                json!({"type": "thought_summary", "content": {"type": "text", "text": "plan"}}),
            ),
            delta(0, json!({"type": "thought_signature", "signature": "sig"})),
            stop(0),
            start(
                1,
                json!({"type": "google_search_call", "id": "call_1", "signature": ""}),
            ),
            delta(1, search),
            stop(1),
            start(2, json!({"type": "model_output"})),
            delta(
                2,
                json!({"type": "text", "text": "Spain", "annotations": [{"start_index": 0, "end_index": 5, "source": "https://example.com"}]}),
            ),
            delta(2, json!({"type": "text", "text": " won."})),
            stop(2),
            completed(),
        ],
    );
}

/// A thought `step.start` that already states its summary writes it, as a
/// whole reply does.
#[test]
fn a_thought_start_with_its_summary_agrees_with_the_whole() {
    let thought = json!({"type": "thought", "signature": "sig", "summary": [{"type": "text", "text": "plan"}]});
    assert_restated_agrees(
        &wire(),
        [whole(vec![thought.clone()])],
        [start(0, thought), stop(0), completed()],
    );
}

/// Every model output content kind decodes to its block: text to text, an
/// image with a source to an image, and audio, video, documents, a
/// sourceless image and invented items to opaque blocks that replay.
#[test]
fn every_content_kind_decodes() {
    #[deny(clippy::wildcard_enum_match_arm)]
    fn variant_index(kind: ContentKind) -> usize {
        match kind {
            ContentKind::Text => 0,
            ContentKind::Image => 1,
            ContentKind::Other => 2,
        }
    }
    let samples = [
        json!({"type": "text", "text": "hello"}),
        json!({"type": "image", "uri": "https://example.com/a.png", "mime_type": "image/png"}),
        json!({"type": "image", "mime_type": "image/png"}),
        json!({"type": "audio", "data": "YQ==", "mime_type": "audio/wav"}),
        json!({"type": "video", "uri": "https://example.com/a.mp4"}),
        json!({"type": "document", "data": "YQ==", "mime_type": "application/pdf"}),
        json!({"type": "x_rig_invented"}),
    ];
    assert_every_variant(&samples, |item| variant_index(content_kind(item)), 3);
    for item in &samples {
        let step = json!({"type": "model_output", "content": [item]});
        let response =
            decode(&wire(), Mode::Unary, [whole(vec![step.clone()])]).expect("the sample decodes");
        let block = response.choice.first();
        let expected = match (content_kind(item), item.get("uri").or(item.get("data"))) {
            (ContentKind::Text, _) => matches!(block, Some(AssistantContent::Text(_))),
            (ContentKind::Image, Some(_)) => matches!(block, Some(AssistantContent::Image(_))),
            (ContentKind::Image, None) | (ContentKind::Other, _) => {
                matches!(block, Some(AssistantContent::Opaque(opaque)) if opaque.replay && opaque.item == step)
            }
        };
        assert!(expected, "{item} decoded to {block:?}");
        assert_eq!(replayed(&response), [step]);
    }
}

/// A tagged frame is an event, never a whole interaction, even with a
/// `status` key; a step event without its index fails the reply.
#[test]
fn a_tagged_frame_never_passes_for_a_whole_interaction() {
    let event = classify_interactions_frame(r#"{"event_type":"step.stop","status":"completed"}"#);
    assert!(matches!(event, WireEvent::Known(InteractionsEvent::Sse(_))));
    let outcome = decode(
        &wire(),
        Mode::Streaming,
        [frame(
            json!({"event_type": "step.stop", "status": "completed"}),
        )],
    );
    assert!(outcome.is_err(), "{outcome:?}");
}

/// A text item's annotations cite bytes of its text, past a multi-byte
/// dash: a URL, a file page and a place. Annotations streamed after the
/// text as `text_annotation_delta` join the text item, so the stream cites
/// and replays what the whole interaction does. Other annotations cite
/// nothing.
#[test]
fn annotations_cite_their_text_item_in_both_modes() {
    let text = "Spain won 2\u{2013}1 in Berlin.";
    let at = |part: &str| text.find(part).expect("the text holds it");
    let span = |part: &str| json!({"start_index": at(part), "end_index": at(part) + part.len()});
    let annotation = |part: &str, fields: serde_json::Value| {
        let mut annotation = span(part);
        if let (Some(annotation), serde_json::Value::Object(fields)) =
            (annotation.as_object_mut(), fields)
        {
            annotation.extend(fields);
        }
        annotation
    };
    let annotations = json!([
        annotation(
            "in Berlin.",
            json!({"type": "url_citation", "url": "https://example.com", "title": "example.com"})
        ),
        annotation(
            "Spain won",
            json!({"type": "file_citation", "document_uri": "files/report", "file_name": "report.pdf", "page_number": 3})
        ),
        annotation(
            "Berlin",
            json!({"type": "place_citation", "place_id": "places/berlin", "name": "Berlin"})
        ),
        annotation("Spain", json!({"type": "word_info"})),
    ]);
    let output = json!({"type": "model_output", "content": [{"type": "text", "text": text, "annotations": annotations}]});
    let streamed = [
        start(0, json!({"type": "model_output"})),
        delta(0, json!({"type": "text", "text": &text[..at("1 in")]})),
        delta(0, json!({"type": "text", "text": &text[at("1 in")..]})),
        delta(
            0,
            json!({"type": "text_annotation_delta", "annotations": annotations}),
        ),
        stop(0),
        completed(),
    ];
    assert_restated_agrees(&wire(), [whole(vec![output.clone()])], streamed.clone());
    let response = decode(&wire(), Mode::Streaming, streamed).expect("the stream decodes");
    let [AssistantContent::Text(block)] = response.choice.as_slice() else {
        panic!("one text block: {:?}", response.choice);
    };
    let cited: Vec<_> = block
        .citations()
        .iter()
        .map(|citation| (block.cited(citation), citation.sources.clone()))
        .collect();
    let url = Source::new(SourceLocation::Url {
        url: "https://example.com".to_owned(),
    })
    .title("example.com");
    let file = Source::new(SourceLocation::Document {
        index: None,
        id: Some("files/report".to_owned()),
        within: Some(DocumentRange::Pages(3..4)),
    })
    .title("report.pdf");
    let place = Source::new(SourceLocation::Document {
        index: None,
        id: Some("places/berlin".to_owned()),
        within: None,
    })
    .title("Berlin");
    assert_eq!(
        cited,
        [
            (Some("in Berlin."), vec![url]),
            (Some("Spain won"), vec![file]),
            (Some("Berlin"), vec![place]),
        ]
    );
    assert_eq!(replayed(&response), [output]);
}

/// An arguments delta may carry the arguments as JSON rather than text, or
/// none at all; the call reads them either way.
#[test]
fn a_call_reads_arguments_a_delta_sends_as_json_or_omits() {
    let frames = [
        start(
            0,
            json!({"type": "function_call", "id": "c1", "name": "add", "arguments": {}}),
        ),
        delta(0, json!({"type": "arguments_delta", "arguments": null})),
        delta(
            0,
            json!({"type": "arguments_delta", "arguments": {"x": 17}}),
        ),
        stop(0),
        completed(),
    ];
    let response = decode(&wire(), Mode::Streaming, frames).expect("the stream decodes");
    let arguments: Vec<_> = response
        .tool_calls()
        .map(|call| serde_json::Value::Object(call.function.arguments.clone()))
        .collect();
    assert_eq!(arguments, [json!({"x": 17})]);
}