rho-coding-agent 2.24.0

use base64::Engine;
use pretty_assertions::assert_eq;
use rho_sdk::tool::ToolErrorKind;
use rmcp::model::{CallToolResult, ContentBlock, Resource, ResourceContents};

use super::{
    render, McpImageDelivery, Rendered, ResultExpectation, MAX_RETAINED_IMAGE_BYTES,
    MAX_VALIDATION_DEPTH,
};

const LIMIT: usize = 12_000;

fn encode(bytes: &[u8]) -> String {
    base64::engine::general_purpose::STANDARD.encode(bytes)
}

fn result(content: Vec<ContentBlock>) -> CallToolResult {
    CallToolResult::success(content)
}

fn expect_schema(schema: serde_json::Value) -> ResultExpectation {
    ResultExpectation {
        output_schema: Some(schema),
    }
}

// Covers: each content kind must reach the model as something it can act on.
// Text stays text, binary becomes a short descriptor plus a card asset rather
// than a base64 wall, and resource blocks name their URI.
// Owner: MCP result rendering.
#[test]
fn content_blocks_render_by_kind() {
    let png = [0x89, b'P', b'N', b'G', 0, 1, 2, 3];
    let rendered = render(
        &result(vec![
            ContentBlock::text("first line"),
            ContentBlock::image(encode(&png), "image/png"),
            ContentBlock::audio(encode(&[0u8; 2048]), "audio/wav"),
            ContentBlock::resource(ResourceContents::text("inline body", "file:///notes.md")),
            ContentBlock::resource(ResourceContents::BlobResourceContents {
                uri: "file:///logo.png".into(),
                mime_type: Some("image/png".into()),
                blob: encode(&png),
                meta: None,
            }),
            ContentBlock::ResourceLink(Resource::new("file:///README.md", "readme")),
        ]),
        &ResultExpectation::default(),
        LIMIT,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap();

    assert_eq!(
        rendered.text(),
        "first line\n\n\
         [image image/png, 8 B]\n\n\
         [audio audio/wav, 2.0 KB]\n\n\
         [resource file:///notes.md]\ninline body\n\n\
         [resource file:///logo.png] [resource image/png, 8 B] [not shown: card keeps only the first image]\n\n\
         [resource link file:///README.md \"readme\"]"
    );
    // The card shows one image, so only the first image block is retained.
    assert_eq!(
        rendered
            .assets()
            .iter()
            .map(|asset| (asset.media_type().to_string(), asset.bytes().len()))
            .collect::<Vec<_>>(),
        vec![("image/png".to_string(), 8)]
    );
}

// Covers: structured content must be presented once. A server that mirrors it
// as text for older clients must not cost the context twice, and a server that
// declares an output schema but returns nothing structured must fail loudly.
// Owner: MCP structured-result contract.
#[test]
fn structured_content_is_presented_once_and_required_when_declared() {
    let structured = serde_json::json!({"count": 2});
    let mut mirrored = result(vec![ContentBlock::text(structured.to_string())]);
    mirrored.structured_content = Some(structured.clone());

    assert_eq!(
        render(
            &mirrored,
            &ResultExpectation::default(),
            LIMIT,
            McpImageDelivery::PresentationOnly
        )
        .unwrap(),
        Rendered::new("{\n  \"count\": 2\n}".into(), structured.clone())
    );

    let mut with_prose = result(vec![ContentBlock::text("Found 2 matches.")]);
    with_prose.structured_content = Some(structured.clone());
    assert_eq!(
        render(
            &with_prose,
            &ResultExpectation::default(),
            LIMIT,
            McpImageDelivery::PresentationOnly
        )
        .unwrap()
        .text(),
        "Found 2 matches.\n\n{\n  \"count\": 2\n}"
    );

    let missing = render(
        &result(vec![ContentBlock::text("Found 2 matches.")]),
        &expect_schema(serde_json::json!({
            "type": "object",
            "required": ["count"],
            "properties": {"count": {"type": "integer"}},
        })),
        LIMIT,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap_err();
    assert_eq!(missing.kind(), ToolErrorKind::Execution);
    assert!(missing.message().contains("no structured content"));

    let mut valid = result(Vec::new());
    valid.structured_content = Some(structured);
    assert_eq!(
        render(
            &valid,
            &expect_schema(serde_json::json!({
                "type": "object",
                "required": ["count"],
                "properties": {"count": {"type": "integer"}},
            })),
            LIMIT,
            McpImageDelivery::PresentationOnly,
        )
        .unwrap()
        .text(),
        "{\n  \"count\": 2\n}"
    );
}

// Covers: a declared output schema must reject structured content that does
// not match, instead of handing the model a malformed half-answer.
// Owner: MCP structured-result contract.
#[test]
fn structured_content_must_match_declared_output_schema() {
    let mut invalid = result(Vec::new());
    invalid.structured_content = Some(serde_json::json!({"count": "two"}));
    for limit in [LIMIT, 1] {
        let error = render(
            &invalid,
            &expect_schema(serde_json::json!({
                "type": "object",
                "required": ["count"],
                "properties": {"count": {"type": "integer"}},
            })),
            limit,
            McpImageDelivery::PresentationOnly,
        )
        .unwrap_err();
        assert_eq!(error.kind(), ToolErrorKind::Execution);
    }
}

// Covers: a completed MCP error keeps its readable text, result data, and
// failure status; an empty result still describes the absence of content.
// Owner: MCP result rendering.
#[test]
fn error_results_and_empty_results_stay_readable() {
    let mut failed = result(vec![ContentBlock::text("disk is full")]);
    failed.is_error = Some(true);
    let error = render(
        &failed,
        &ResultExpectation::default(),
        LIMIT,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap();
    assert_eq!(
        (error.is_failure(), error.text(), error.data()),
        (true, "disk is full", None)
    );

    // A structured error result keeps its payload for scripts (Pi resolves
    // MCP error results to their structured content too).
    let payload = serde_json::json!({"free_bytes": 0});
    failed.structured_content = Some(payload.clone());
    let error = render(
        &failed,
        &ResultExpectation::default(),
        LIMIT,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap();
    assert_eq!((error.is_failure(), error.data()), (true, Some(&payload)));

    assert_eq!(
        render(
            &result(Vec::new()),
            &ResultExpectation::default(),
            LIMIT,
            McpImageDelivery::PresentationOnly
        )
        .unwrap()
        .text(),
        "The MCP server returned no content."
    );
}

// Covers: later images and oversized images must not retain unreachable or
// unbounded binary payloads; the model still gets a size descriptor.
// Owner: MCP result rendering.
#[test]
fn image_assets_keep_only_the_first_that_fits_the_budget() {
    let small = [0x89, b'P', b'N', b'G', 0, 1, 2, 3];
    let oversized = vec![0u8; MAX_RETAINED_IMAGE_BYTES + 1];

    let multi = render(
        &result(vec![
            ContentBlock::image(encode(&small), "image/png"),
            ContentBlock::image(encode(&small), "image/png"),
        ]),
        &ResultExpectation::default(),
        LIMIT,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap();
    assert_eq!(multi.assets().len(), 1);
    assert!(multi
        .text()
        .contains("[not shown: card keeps only the first image]"));

    let too_large = render(
        &result(vec![ContentBlock::image(encode(&oversized), "image/png")]),
        &ResultExpectation::default(),
        LIMIT,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap();
    assert!(too_large.assets().is_empty());
    assert!(too_large.text().contains("not retained: exceeds"));
    assert!(!too_large.text().contains(&encode(&oversized[..16])));
}

// Covers: model observations must not depend on which images fit the preview
// card. Delivery remains host-opt-in and preserves typed source payloads.
// Owner: MCP result interpretation, before SDK history commitment.
#[test]
fn model_images_are_independent_of_presentation_retention() {
    use rho_sdk::model::ImageContent;

    let small = b"\x89PNG\r\n\x1a\n".to_vec();
    let mut large = small.clone();
    large.resize(MAX_RETAINED_IMAGE_BYTES + 1, 0);
    let expected = vec![
        ImageContent {
            data: encode(&small),
            mime_type: "image/png".into(),
        },
        ImageContent {
            data: encode(&small),
            mime_type: "image/png".into(),
        },
        ImageContent {
            data: encode(&large),
            mime_type: "image/png".into(),
        },
    ];
    let content = result(vec![
        ContentBlock::image(expected[0].data.clone(), "image/png"),
        ContentBlock::image(expected[1].data.clone(), "image/png"),
        ContentBlock::resource(ResourceContents::BlobResourceContents {
            uri: "file:///observation.png".into(),
            mime_type: Some("image/png".into()),
            blob: expected[2].data.clone(),
            meta: None,
        }),
    ]);
    for (delivery, images) in [
        (McpImageDelivery::PresentationOnly, Vec::new()),
        (McpImageDelivery::ModelAndPresentation, expected),
    ] {
        let rendered = render(&content, &ResultExpectation::default(), LIMIT, delivery).unwrap();
        assert_eq!(rendered.images(), images);
        assert_eq!(rendered.assets().len(), 1);
    }
}

// Covers: raw structured results cannot bypass the MCP output cap, including
// error results and mirrored text; retained payloads at the limit remain usable.
// Owner: MCP result output budget.
#[test]
fn structured_content_obeys_output_budget() {
    for (failed, schema) in [
        (false, None),
        (true, None),
        (false, Some(serde_json::json!({"type": "string"}))),
    ] {
        for size in [LIMIT, LIMIT + 1] {
            // JSON quotes contribute two bytes; construct an exact serialized size.
            let payload = serde_json::json!("x".repeat(size - 2));
            let mut call = result(vec![ContentBlock::text(payload.to_string())]);
            call.structured_content = Some(payload.clone());
            call.is_error = Some(failed);
            let expectation = ResultExpectation {
                output_schema: schema.clone(),
            };
            let rendered = render(
                &call,
                &expectation,
                LIMIT,
                McpImageDelivery::PresentationOnly,
            )
            .unwrap();
            assert_eq!(rendered.is_failure(), failed);
            assert_eq!(rendered.data(), (size <= LIMIT).then_some(&payload));
            assert!(rendered.text().len() <= LIMIT);
            if size <= LIMIT {
                assert_eq!(rendered.text(), payload.to_string());
            } else {
                assert!(rendered.text().starts_with("[structured data truncated:"));
            }
        }
    }
}

// Covers: untrusted output schemas and structured payloads must be rejected
// before compile/validate work can monopolize CPU or memory.
// Owner: MCP structured-result contract.
#[test]
fn structured_validation_rejects_oversize_schemas() {
    // A schema with far more nodes than the budget, built without a huge
    // serialized form by nesting many tiny properties.
    let mut properties = serde_json::Map::new();
    for index in 0..3_000 {
        properties.insert(format!("f{index}"), serde_json::json!({"type": "string"}));
    }
    let schema = serde_json::json!({
        "type": "object",
        "properties": properties,
    });
    let mut call = result(Vec::new());
    call.structured_content = Some(serde_json::json!({}));
    let error = render(
        &call,
        &expect_schema(schema),
        LIMIT,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap_err();
    assert_eq!(error.kind(), ToolErrorKind::Execution);
}

// Covers: large-but-bounded structured results validate. T3 Code's
// `orchestrator_capabilities` returns ~13k nodes in ~190 KB and was rejected by
// a node cap far below what the byte budget already bounds.
// Owner: MCP structured-result contract.
#[test]
fn structured_validation_accepts_many_nodes_within_byte_budget() {
    let payload = serde_json::json!({ "items": vec![0; 20_000] });
    let mut call = result(Vec::new());
    call.structured_content = Some(payload.clone());
    let schema = serde_json::json!({
        "type": "object",
        "properties": { "items": { "type": "array", "items": { "type": "integer" } } },
    });
    let rendered = render(
        &call,
        &expect_schema(schema),
        /*max_output_bytes*/ 128 * 1024,
        McpImageDelivery::PresentationOnly,
    )
    .unwrap();
    assert_eq!(rendered.data(), Some(&payload));
}

// Covers: width is bounded by bytes, depth separately. Nested content past the
// depth budget must fail validation instead of recursing the validator off the
// stack, while content at the budget still validates.
// Owner: MCP structured-result contract.
#[test]
fn structured_validation_bounds_nesting_depth() {
    let schema = serde_json::json!({
        "$defs": { "n": { "anyOf": [
            { "type": "number" },
            { "type": "array", "items": { "$ref": "#/$defs/n" } },
        ] } },
        "$ref": "#/$defs/n",
    });
    for (levels, accepted) in [
        (MAX_VALIDATION_DEPTH, true),
        (MAX_VALIDATION_DEPTH + 1, false),
    ] {
        // `levels` JSON values deep: `levels - 1` arrays around one number.
        let mut payload = serde_json::json!(0);
        for _ in 1..levels {
            payload = serde_json::Value::Array(vec![payload]);
        }
        let mut call = result(Vec::new());
        call.structured_content = Some(payload);
        let rendered = render(
            &call,
            &expect_schema(schema.clone()),
            /*max_output_bytes*/ 128 * 1024,
            McpImageDelivery::PresentationOnly,
        );
        assert_eq!(rendered.is_ok(), accepted, "{levels} levels");
    }
}