anda_engine 0.16.6

//! Provider raw-history classification and pruning.
//!
//! `CompletionRequest::raw_history` carries provider-native message JSON
//! (OpenAI Chat, OpenAI Responses, Anthropic, and Gemini shapes all differ), so
//! reclaiming context-window budget from it requires recognizing tool-call and
//! tool-output items in each provider's wire format. That knowledge belongs to
//! the model layer: this module is the single home for it, and
//! [`CompletionFeaturesDyn`](crate::model::CompletionFeaturesDyn) exposes it as
//! overridable default methods so a custom provider whose wire format is not
//! covered here can supply its own classification without any change to the
//! completion runner.
//!
//! The implementation here is a conservative union over all built-in provider
//! shapes. Items are rewritten in place rather than round-tripped through
//! [`Message`](anda_core::Message), which would drop or corrupt
//! provider-specific fields such as reasoning signatures.

use anda_core::Json;
use serde_json::json;

/// Removes unanswered tool-call requests from `raw_history[start..]`.
///
/// Used after an interrupt (steering, discard, stop) so the next request does
/// not inherit a tool-call requirement the provider would reject as unanswered.
/// Visible text and reasoning stay untouched; wrapper items left without
/// meaningful content are removed entirely.
pub(crate) fn prune_unanswered_tool_calls(raw_history: &mut Vec<Json>, start: usize) {
    if start >= raw_history.len() {
        return;
    }

    let tail: Vec<Json> = raw_history.drain(start..).collect();
    raw_history.extend(prune_items(tail, is_tool_call_item));
}

/// Removes completed tool interactions (calls and results) from provider raw history.
pub(crate) fn prune_tool_interactions(raw_history: &mut Vec<Json>) {
    let items = std::mem::take(raw_history);
    *raw_history = prune_items(items, |value| {
        is_tool_call_item(value) || is_tool_output_item(value)
    });
}

/// Replaces inline media payloads in provider raw history with a short text note.
///
/// Inline media is anything that carries its bytes in the request: base64 or text sources,
/// `data:` URLs, and audio or file data fields. Media referenced by a remote URL or a provider
/// file id is kept, and so are messages the model wrote: media it generated is its answer, not
/// an attachment, and Gemini signs it. The note takes the text shape of the provider the item
/// belongs to, so the surrounding message keeps valid, non-empty content.
pub(crate) fn prune_inline_media(raw_history: &mut [Json]) {
    for item in raw_history {
        if !matches!(
            item.get("role").and_then(Json::as_str),
            Some("assistant" | "model")
        ) {
            replace_inline_media(item, false);
        }
    }
}

fn replace_inline_media(value: &mut Json, responses: bool) {
    match value {
        Json::Array(items) => {
            // OpenAI Chat and Responses share the `input_audio` shape but not their text shape;
            // a Responses item among the siblings settles which one this array uses.
            let responses = responses || items.iter().any(is_responses_content_item);
            for item in items {
                match inline_media_note(item, responses) {
                    Some(note) => *item = note,
                    None => replace_inline_media(item, responses),
                }
            }
        }
        Json::Object(map) => {
            let responses = responses
                || matches!(
                    map.get("type").and_then(Json::as_str),
                    Some("message" | "function_call_output")
                );
            for (key, value) in map.iter_mut() {
                if matches!(key.as_str(), "functionResponse" | "function_response") {
                    drop_function_response_media(value);
                }
                replace_inline_media(value, responses);
            }
        }
        _ => {}
    }
}

/// Gemini `FunctionResponsePart`s accept media only, so a text note in their place would be
/// rejected: inline media leaves `parts`, and its note joins the response text instead.
fn drop_function_response_media(response: &mut Json) {
    let Some(response) = response.as_object_mut() else {
        return;
    };
    let Some(Json::Array(parts)) = response.get_mut("parts") else {
        return;
    };
    let mut notes = String::new();
    parts.retain(|part| match inline_media_note(part, false) {
        Some(note) => {
            notes.push('\n');
            notes.push_str(note.get("text").and_then(Json::as_str).unwrap_or_default());
            false
        }
        None => true,
    });
    if parts.is_empty() {
        response.remove("parts");
    }
    // The adapter sends a tool's presentation text as the `output` (or `error`) string.
    if let Some(Json::Object(result)) = response.get_mut("response") {
        let key = if result.contains_key("output") {
            "output"
        } else {
            "error"
        };
        if let Some(Json::String(text)) = result.get_mut(key) {
            text.push_str(&notes);
        }
    }
}

fn is_responses_content_item(value: &Json) -> bool {
    matches!(
        value.get("type").and_then(Json::as_str),
        Some("input_text" | "input_image" | "input_file" | "output_text" | "refusal")
    )
}

/// Returns the replacement for `item` when it is an inline media block.
fn inline_media_note(item: &Json, responses: bool) -> Option<Json> {
    let map = item.as_object()?;

    // Gemini parts carry no `type`; the payload sits under `inlineData`, or under `fileData`
    // as a `data:` URI. (A Responses `input_file` has a string `file_data` instead.)
    if let Some(data) = map.get("inlineData").or_else(|| map.get("inline_data")) {
        let mime = str_at(data.get("mimeType").or_else(|| data.get("mime_type")));
        return Some(json!({ "text": inline_media_text(mime) }));
    }
    if let Some(data) = map
        .get("fileData")
        .or_else(|| map.get("file_data"))
        .filter(|data| data.is_object())
    {
        let uri = str_at(data.get("fileUri").or_else(|| data.get("file_uri")))?;
        return Some(json!({ "text": inline_media_text(Some(&data_url_mime(uri)?)) }));
    }

    let kind = str_at(map.get("type"))?;
    let mime = match kind {
        // Anthropic image and document blocks.
        "image" | "document" => {
            let source = map.get("source")?;
            match str_at(source.get("type"))? {
                "base64" | "text" => str_at(source.get("media_type")).map(str::to_string),
                "url" => Some(data_url_mime(str_at(source.get("url"))?)?),
                _ => return None,
            }
        }
        // OpenAI Chat Completions.
        "image_url" | "video_url" => Some(data_url_mime(str_at(map.get(kind)?.get("url"))?)?),
        "file" => Some(data_url_mime(str_at(map.get("file")?.get("file_data"))?)?),
        // OpenAI Responses (and Chat Completions for `input_audio`).
        "input_image" => Some(data_url_mime(str_at(map.get("image_url"))?)?),
        // The adapter sends a `data:` `FileData` as `file_url`.
        "input_file" => Some(
            ["file_data", "file_url"]
                .into_iter()
                .find_map(|key| data_url_mime(str_at(map.get(key))?))?,
        ),
        "input_audio" => {
            let audio = map.get("input_audio")?;
            str_at(audio.get("data"))?;
            str_at(audio.get("format")).map(|format| format!("audio/{format}"))
        }
        _ => return None,
    };

    let text_kind = match kind {
        "input_image" | "input_file" => "input_text",
        "input_audio" if responses => "input_text",
        _ => "text",
    };
    Some(json!({ "type": text_kind, "text": inline_media_text(mime.as_deref()) }))
}

fn str_at(value: Option<&Json>) -> Option<&str> {
    value.and_then(Json::as_str)
}

/// Returns the MIME type of a `data:` URL, or `None` for any other URL.
fn data_url_mime(url: &str) -> Option<String> {
    let header = url.strip_prefix("data:")?;
    let mime = header
        .split([';', ','])
        .next()
        .filter(|mime| !mime.is_empty())
        .unwrap_or("application/octet-stream");
    Some(mime.to_string())
}

pub(crate) fn inline_media_text(mime: Option<&str>) -> String {
    format!("[inline {} data omitted]", mime.unwrap_or("media"))
}

/// Drops every item `is_pruned` classifies, recursing into wrapper items so a
/// wrapper keeps its non-tool context and disappears only once nothing
/// meaningful is left.
///
/// Walks backwards so an OpenAI Responses `reasoning` item can see whether the item it
/// must immediately precede survived; the API rejects a reasoning item whose required
/// following item is missing, so orphaned reasoning follows its pruned sibling out.
fn prune_items(items: Vec<Json>, is_pruned: impl Fn(&Json) -> bool + Copy) -> Vec<Json> {
    let mut retained: Vec<Json> = Vec::with_capacity(items.len());
    let mut next_kept = true;
    for value in items.into_iter().rev() {
        if !next_kept && value.get("type").and_then(|v| v.as_str()) == Some("reasoning") {
            continue;
        }
        match prune_item(value, is_pruned) {
            Some(value) => {
                retained.push(value);
                next_kept = true;
            }
            None => next_kept = false,
        }
    }
    retained.reverse();
    retained
}

fn prune_item(mut value: Json, is_pruned: impl Fn(&Json) -> bool + Copy) -> Option<Json> {
    if is_pruned(&value) {
        return None;
    }

    prune_nested(&mut value, is_pruned);
    if item_has_context(&value) {
        Some(value)
    } else {
        None
    }
}

fn prune_nested(value: &mut Json, is_pruned: impl Fn(&Json) -> bool + Copy) {
    match value {
        Json::Array(items) => {
            let retained: Vec<Json> = items
                .drain(..)
                .filter_map(|item| prune_item(item, is_pruned))
                .collect();
            *items = retained;
        }
        Json::Object(map) => {
            // OpenAI Chat Completions keeps tool calls as fields on an assistant message.
            // Remove the unanswered calls, but keep any text/reasoning fields on the same
            // message. Provider raw history may also wrap output items under arbitrary fields,
            // so recurse through every remaining value instead of only common content arrays.
            map.remove("tool_calls");
            map.remove("function_call");
            map.remove("functionCall");

            for value in map.values_mut() {
                prune_nested(value, is_pruned);
            }
        }
        _ => {}
    }
}

fn is_tool_call_item(value: &Json) -> bool {
    let Some(map) = value.as_object() else {
        return false;
    };

    if matches!(
        map.get("type").and_then(|v| v.as_str()),
        Some(
            "function_call"
                | "custom_tool_call"
                | "computer_call"
                | "tool_search_call"
                | "local_shell_call"
                | "shell_call"
                | "apply_patch_call"
                | "mcp_approval_request"
                | "tool_call"
                | "tool_use"
                | "ToolCall"
                | "toolCall"
        )
    ) {
        return true;
    }

    // Gemini function-call parts do not use a `type` field. Treat only the part itself as a
    // tool call; wrapper objects that also carry text or metadata should be pruned recursively
    // so their non-tool context survives.
    map.contains_key("functionCall")
        && map.keys().all(|key| {
            matches!(
                key.as_str(),
                "functionCall" | "thought" | "thoughtSignature"
            )
        })
}

fn is_tool_output_item(value: &Json) -> bool {
    let Some(map) = value.as_object() else {
        return false;
    };

    // OpenAI Chat Completions carries tool results as whole messages with the "tool" role.
    if map.get("role").and_then(|v| v.as_str()) == Some("tool") {
        return true;
    }

    if matches!(
        map.get("type").and_then(|v| v.as_str()),
        Some(
            "function_call_output"
                | "custom_tool_call_output"
                | "computer_call_output"
                | "local_shell_call_output"
                | "shell_call_output"
                | "apply_patch_call_output"
                | "mcp_approval_response"
                | "tool_result"
                | "tool_output"
                | "ToolOutput"
                | "toolOutput"
        )
    ) {
        return true;
    }

    // Gemini function-response parts do not use a `type` field; mirror the
    // `is_tool_call_item` handling for `functionCall` parts.
    map.contains_key("functionResponse")
        && map.keys().all(|key| {
            matches!(
                key.as_str(),
                "functionResponse" | "thought" | "thoughtSignature"
            )
        })
}

fn item_has_context(value: &Json) -> bool {
    match value {
        Json::Null => false,
        Json::Bool(_) | Json::Number(_) => true,
        Json::String(text) => !text.is_empty(),
        Json::Array(items) => items.iter().any(item_has_context),
        Json::Object(map) => map.iter().any(|(key, value)| {
            !matches!(
                key.as_str(),
                "role" | "type" | "name" | "id" | "status" | "phase" | "timestamp"
            ) && item_has_context(value)
        }),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn prunes_contextless_raw_history_items() {
        let mut raw_history = vec![
            json!(null),
            json!(""),
            json!({"role": "assistant", "id": "meta-only", "status": "ok"}),
            json!({"role": "assistant", "content": []}),
            json!({
                "role": "assistant",
                "content": [
                    {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}},
                    {"text": "kept text"}
                ],
                "tool_calls": [{"id": "call_1"}],
                "function_call": {"name": "lookup"}
            }),
            json!({"type": "function_call", "call_id": "call_1"}),
            json!(42),
        ];

        prune_unanswered_tool_calls(&mut raw_history, 0);

        assert_eq!(raw_history.len(), 2);
        assert_eq!(raw_history[0]["content"], json!([{"text": "kept text"}]));
        assert!(raw_history[0].get("tool_calls").is_none());
        assert!(raw_history[0].get("function_call").is_none());
        assert_eq!(raw_history[1], json!(42));
    }

    #[test]
    fn prunes_deeply_nested_raw_tool_calls() {
        let sentinel = json!({"role": "user", "content": "prior"});
        let mut raw_history = vec![
            sentinel.clone(),
            json!({
                "role": "assistant",
                "output": [
                    {
                        "type": "message",
                        "content": [
                            {"type": "output_text", "text": "kept nested text"}
                        ]
                    },
                    {
                        "type": "function_call",
                        "call_id": "call_nested_function",
                        "name": "lookup",
                        "arguments": "{}"
                    },
                    {
                        "type": "local_shell_call",
                        "id": "lsh_1",
                        "action": {"type": "exec", "command": ["pwd"]},
                        "call_id": "call_nested_shell",
                        "status": "completed"
                    },
                    {
                        "type": "apply_patch_call",
                        "id": "ap_1",
                        "call_id": "call_nested_patch",
                        "operation": {"type": "update", "path": "src/lib.rs", "diff": "@@"},
                        "status": "completed"
                    }
                ],
                "metadata": {
                    "items": [
                        {
                            "type": "custom_tool_call",
                            "call_id": "call_nested_custom",
                            "input": "select 1",
                            "name": "sql"
                        },
                        {"note": "keep nested metadata"}
                    ]
                },
                "tool_calls": [{"id": "call_nested_chat"}],
                "function_call": {"name": "legacy"},
                "functionCall": {"name": "gemini"}
            }),
        ];

        prune_unanswered_tool_calls(&mut raw_history, 1);

        assert_eq!(raw_history.len(), 2);
        assert_eq!(raw_history[0], sentinel);
        let pruned = serde_json::to_string(&raw_history[1]).unwrap();
        assert!(pruned.contains("kept nested text"));
        assert!(pruned.contains("keep nested metadata"));
        assert!(!pruned.contains("call_nested"));
        assert!(raw_history[1].get("tool_calls").is_none());
        assert!(raw_history[1].get("function_call").is_none());
        assert!(raw_history[1].get("functionCall").is_none());
    }

    #[test]
    fn prunes_completed_tool_interactions_from_raw_history() {
        let mut raw_history = vec![
            // OpenAI Chat Completions shapes.
            json!({"role": "user", "content": "question"}),
            json!({
                "role": "assistant",
                "content": null,
                "tool_calls": [{"id": "call_1", "function": {"name": "lookup", "arguments": "{}"}}]
            }),
            json!({"role": "tool", "tool_call_id": "call_1", "content": "chat tool result"}),
            // OpenAI Responses shapes: orphaned reasoning must follow its pruned call out.
            json!({"type": "reasoning", "id": "rs_orphan", "encrypted_content": "opaque"}),
            json!({"type": "function_call", "call_id": "call_2", "name": "lookup", "arguments": "{}"}),
            json!({"type": "function_call_output", "call_id": "call_2", "output": "ok"}),
            // Anthropic shapes.
            json!({
                "role": "assistant",
                "content": [
                    {"type": "thinking", "thinking": "kept thinking", "signature": "sig"},
                    {"type": "text", "text": "anthropic text"},
                    {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}}
                ]
            }),
            json!({
                "role": "user",
                "content": [
                    {"type": "tool_result", "tool_use_id": "toolu_1", "content": "anthropic result"}
                ]
            }),
            // Gemini shapes.
            json!({
                "role": "model",
                "parts": [
                    {"text": "gemini text"},
                    {"functionCall": {"name": "lookup", "args": {}}}
                ]
            }),
            json!({
                "role": "user",
                "parts": [{"functionResponse": {"name": "lookup", "response": {"ok": true}}}]
            }),
            // Reasoning followed by a surviving item stays.
            json!({"type": "reasoning", "id": "rs_kept", "encrypted_content": "opaque"}),
            json!({
                "type": "message",
                "role": "assistant",
                "content": [{"type": "output_text", "text": "final answer"}]
            }),
        ];

        prune_tool_interactions(&mut raw_history);

        let pruned = serde_json::to_string(&raw_history).unwrap();
        assert_eq!(raw_history.len(), 5);
        assert_eq!(raw_history[0]["content"], json!("question"));
        assert!(pruned.contains("kept thinking"));
        assert!(pruned.contains("anthropic text"));
        assert!(pruned.contains("gemini text"));
        assert!(pruned.contains("rs_kept"));
        assert!(pruned.contains("final answer"));
        assert!(!pruned.contains("call_1"));
        assert!(!pruned.contains("call_2"));
        assert!(!pruned.contains("rs_orphan"));
        assert!(!pruned.contains("toolu_1"));
        assert!(!pruned.contains("chat tool result"));
        assert!(!pruned.contains("anthropic result"));
        assert!(!pruned.contains("functionCall"));
        assert!(!pruned.contains("functionResponse"));
    }

    #[test]
    fn prunes_only_unanswered_raw_tool_call_items() {
        let sentinel = json!({"role": "user", "content": "prior"});
        let mut raw_history = vec![
            sentinel.clone(),
            json!({
                "role": "assistant",
                "content": [
                    {"type": "text", "text": "anthropic text"},
                    {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {}}
                ]
            }),
            json!({
                "role": "model",
                "parts": [
                    {"text": "gemini text"},
                    {"functionCall": {"name": "lookup", "args": {}}}
                ]
            }),
            json!({"type": "function_call", "call_id": "call_1"}),
            json!({"type": "custom_tool_call", "call_id": "call_2"}),
        ];

        prune_unanswered_tool_calls(&mut raw_history, 1);

        assert_eq!(raw_history.len(), 3);
        assert_eq!(raw_history[0], sentinel);
        assert_eq!(raw_history[1]["content"].as_array().unwrap().len(), 1);
        assert_eq!(raw_history[1]["content"][0]["text"], "anthropic text");
        assert_eq!(raw_history[2]["parts"].as_array().unwrap().len(), 1);
        assert_eq!(raw_history[2]["parts"][0]["text"], "gemini text");
    }

    #[test]
    fn prunes_reasoning_orphaned_by_an_unanswered_tool_call() {
        // OpenAI Responses emits `[reasoning, function_call]` for a tool turn; dropping the
        // unanswered call must take its reasoning sibling with it, or the next request is
        // rejected for a reasoning item without its required following item.
        let sentinel = json!({"role": "user", "content": "prior"});
        let mut raw_history = vec![
            sentinel.clone(),
            json!({"type": "reasoning", "id": "rs_kept", "encrypted_content": "opaque"}),
            json!({
                "type": "message",
                "role": "assistant",
                "content": [{"type": "output_text", "text": "planning"}]
            }),
            json!({"type": "reasoning", "id": "rs_orphan", "encrypted_content": "opaque"}),
            json!({"type": "function_call", "call_id": "call_1", "name": "lookup"}),
        ];

        prune_unanswered_tool_calls(&mut raw_history, 1);

        let pruned = serde_json::to_string(&raw_history).unwrap();
        assert_eq!(raw_history.len(), 3);
        assert_eq!(raw_history[0], sentinel);
        assert!(pruned.contains("rs_kept"));
        assert!(pruned.contains("planning"));
        assert!(!pruned.contains("rs_orphan"));
        assert!(!pruned.contains("call_1"));
    }

    #[test]
    fn replaces_inline_media_in_every_provider_shape() {
        let note = |mime: &str| format!("[inline {mime} data omitted]");
        let mut raw_history = vec![
            // Anthropic: base64 image, text document, and an image inside a tool result.
            json!({"role": "user", "content": [
                {"type": "text", "text": "look"},
                {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AAAA"}},
                {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "body"}},
                {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}},
                {"type": "tool_result", "tool_use_id": "t1", "content": [
                    {"type": "image", "source": {"type": "url", "url": "data:image/jpeg;base64,BBBB"}}
                ]}
            ]}),
            // OpenAI Chat Completions.
            json!({"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": "data:image/webp;base64,CCCC"}},
                {"type": "image_url", "image_url": {"url": "https://example.com/b.png"}},
                {"type": "input_audio", "input_audio": {"data": "DDDD", "format": "wav"}},
                {"type": "file", "file": {"file_data": "data:application/pdf;base64,EEEE", "filename": "a.pdf"}}
            ]}),
            // OpenAI Responses: `input_audio` takes the Responses text shape here.
            json!({"type": "message", "role": "user", "content": [
                {"type": "input_image", "image_url": "data:image/gif;base64,FFFF", "detail": "auto"},
                {"type": "input_image", "file_id": "file_1", "detail": "auto"},
                {"type": "input_file", "file_data": "data:text/csv;base64,GGGG"},
                {"type": "input_audio", "input_audio": {"data": "HHHH", "format": "mp3"}}
            ]}),
            // Gemini.
            json!({"role": "user", "parts": [
                {"inlineData": {"mimeType": "video/mp4", "data": "IIII"}},
                {"fileData": {"fileUri": "gs://bucket/v.mp4", "mimeType": "video/mp4"}},
                {"text": "kept"}
            ]}),
        ];

        prune_inline_media(&mut raw_history);

        assert_eq!(
            raw_history[0]["content"],
            json!([
                {"type": "text", "text": "look"},
                {"type": "text", "text": note("image/png")},
                {"type": "text", "text": note("text/plain")},
                {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}},
                {"type": "tool_result", "tool_use_id": "t1", "content": [
                    {"type": "text", "text": note("image/jpeg")}
                ]}
            ])
        );
        assert_eq!(
            raw_history[1]["content"],
            json!([
                {"type": "text", "text": note("image/webp")},
                {"type": "image_url", "image_url": {"url": "https://example.com/b.png"}},
                {"type": "text", "text": note("audio/wav")},
                {"type": "text", "text": note("application/pdf")}
            ])
        );
        assert_eq!(
            raw_history[2]["content"],
            json!([
                {"type": "input_text", "text": note("image/gif")},
                {"type": "input_image", "file_id": "file_1", "detail": "auto"},
                {"type": "input_text", "text": note("text/csv")},
                {"type": "input_text", "text": note("audio/mp3")}
            ])
        );
        assert_eq!(
            raw_history[3]["parts"],
            json!([
                {"text": note("video/mp4")},
                {"fileData": {"fileUri": "gs://bucket/v.mp4", "mimeType": "video/mp4"}},
                {"text": "kept"}
            ])
        );

        // Pruning again finds nothing left to replace.
        let pruned = raw_history.clone();
        prune_inline_media(&mut raw_history);
        assert_eq!(raw_history, pruned);
    }

    #[test]
    fn replaces_data_uri_file_references() {
        let note = |mime: &str| format!("[inline {mime} data omitted]");
        let mut raw_history = vec![
            // OpenAI Responses sends a `data:` `FileData` as `file_url`.
            json!({"type": "message", "role": "user", "content": [
                {"type": "input_file", "file_url": "data:application/pdf;base64,AAAA"},
                {"type": "input_file", "file_url": "https://example.com/a.pdf"}
            ]}),
            // Gemini sends it as `fileData`.
            json!({"role": "user", "parts": [
                {"fileData": {"fileUri": "data:application/pdf;base64,BBBB", "mimeType": "application/pdf"}},
                {"fileData": {"fileUri": "https://example.com/b.pdf", "mimeType": "application/pdf"}}
            ]}),
        ];

        prune_inline_media(&mut raw_history);

        assert_eq!(
            raw_history[0]["content"],
            json!([
                {"type": "input_text", "text": note("application/pdf")},
                {"type": "input_file", "file_url": "https://example.com/a.pdf"}
            ])
        );
        assert_eq!(
            raw_history[1]["parts"],
            json!([
                {"text": note("application/pdf")},
                {"fileData": {"fileUri": "https://example.com/b.pdf", "mimeType": "application/pdf"}}
            ])
        );
    }

    #[test]
    fn gemini_function_response_media_is_dropped_not_turned_into_text() {
        let mut raw_history = vec![json!({"role": "user", "parts": [
            {"functionResponse": {
                "name": "screenshot",
                "response": {"output": "captured"},
                "parts": [{"inlineData": {"mimeType": "image/png", "data": "AAAA"}}]
            }},
            {"functionResponse": {
                "name": "render",
                "response": {"output": "rendered"},
                "parts": [
                    {"inlineData": {"mimeType": "image/jpeg", "data": "BBBB"}},
                    {"fileData": {"fileUri": "gs://bucket/c.png", "mimeType": "image/png"}}
                ]
            }}
        ]})];

        prune_inline_media(&mut raw_history);

        // `FunctionResponsePart` accepts media only, so a text note there would be rejected.
        assert_eq!(
            raw_history[0]["parts"],
            json!([
                {"functionResponse": {
                    "name": "screenshot",
                    "response": {"output": "captured\n[inline image/png data omitted]"}
                }},
                {"functionResponse": {
                    "name": "render",
                    "response": {"output": "rendered\n[inline image/jpeg data omitted]"},
                    "parts": [{"fileData": {"fileUri": "gs://bucket/c.png", "mimeType": "image/png"}}]
                }}
            ])
        );
    }

    #[test]
    fn keeps_media_the_model_generated() {
        let generated = vec![
            json!({"role": "model", "parts": [
                {"inlineData": {"mimeType": "image/png", "data": "AAAA"}, "thoughtSignature": "sig"}
            ]}),
            json!({"role": "assistant", "content": [
                {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "BBBB"}}
            ]}),
        ];
        let mut raw_history = generated.clone();

        prune_inline_media(&mut raw_history);

        assert_eq!(raw_history, generated);
    }
}