Images returned from MCP tool results are silently dropped when using a BYOK provider

"""Minimal reproducer: images returned from an MCP tool result never reach the model."""

from __future__ import annotations

import asyncio
import base64
import os
import random
import struct
import sys
import zlib

import copilot
import uvicorn
from copilot.generated import rpc
from fastmcp import FastMCP
from fastmcp.utilities.types import Image

PORT = int(os.environ.get("REPRO_PORT", "8899"))

PALETTE = {
    "red": (220, 30, 30),
    "green": (30, 180, 60),
    "blue": (40, 70, 220),
    "yellow": (240, 220, 40),
    "purple": (140, 40, 180),
    "orange": (245, 140, 20),
}


def make_quadrant_png(colors: list[tuple[int, int, int]], size: int = 512) -> bytes:
    """A PNG split into four solid quadrants: TL, TR, BL, BR. No image library needed."""

    def chunk(tag: bytes, data: bytes) -> bytes:
        return (
            struct.pack(">I", len(data))
            + tag
            + data
            + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)
        )

    half = size // 2
    raw = bytearray()
    for y in range(size):
        raw.append(0)
        top = y < half
        for x in range(size):
            left = x < half
            idx = 0 if (top and left) else 1 if top else 2 if left else 3
            raw += bytes(colors[idx])

    return (
        b"\x89PNG\r\n\x1a\n"
        + chunk(b"IHDR", struct.pack(">IIBBBBB", size, size, 8, 2, 0, 0, 0))
        + chunk(b"IDAT", zlib.compress(bytes(raw), 6))
        + chunk(b"IEND", b"")
    )


IMAGE: bytes = b""
TOOL_CALLED = False

mcp = FastMCP("repro")


@mcp.tool
async def get_test_image() -> Image:
    """Return a test image for visual analysis."""
    global TOOL_CALLED
    TOOL_CALLED = True
    return Image(data=IMAGE, format="png")


app = mcp.http_app(transport="streamable-http", path="/mcp/")

SYSTEM = (
    "You are a visual analysis assistant. When asked about an image, you MUST "
    "look at the actual image. Never guess. Answer exactly as instructed."
)

INSTRUCTION = (
    "The image is divided into four solid-colour quadrants. Report the colour of "
    "each quadrant, choosing from: red, green, blue, yellow, purple, orange.\n\n"
    "Reply with EXACTLY four lowercase words separated by single spaces, in this "
    "order: top-left, top-right, bottom-left, bottom-right. No punctuation, no "
    "explanation, no other text."
)

MCP_QUESTION = "Call the get_test_image tool, then look at the image it returns. " + INSTRUCTION


def session_kwargs() -> dict:
    """Default Copilot provider, or BYOK if REPRO_AZURE_* are set."""
    base_url = os.environ.get("REPRO_AZURE_BASE_URL")
    if not base_url:
        return {}
    model = os.environ["REPRO_AZURE_MODEL"]
    token = os.environ["REPRO_AZURE_TOKEN"]
    return {
        "model": model,
        "provider": {
            "type": "azure",
            "wire_api": os.environ.get("REPRO_WIRE_API", "responses"),
            "base_url": base_url,
            "bearer_token_provider": lambda _a: token,
            "model_id": model,
        },
    }


async def ask(*, mcp_servers=None, attachments=None, question: str) -> str:
    client = copilot.CopilotClient(log_level="error")
    await client.start()
    try:
        extra = {}
        if mcp_servers:
            extra["mcp_servers"] = mcp_servers
            extra["available_tools"] = copilot.ToolSet().add_mcp("*")
        session = await client.create_session(
            system_message={"mode": "replace", "content": SYSTEM},
            on_permission_request=lambda *_a: rpc.PermissionDecisionApproveOnce(),
            streaming=False,
            **session_kwargs(),
            **extra,
        )
        if mcp_servers:
            await asyncio.sleep(2.0)  # allow MCP negotiation to finish
        send = session.send_and_wait(question, attachments=attachments) if attachments \
            else session.send_and_wait(question)
        response = await asyncio.wait_for(send, timeout=180)
    finally:
        await client.stop()
    data = getattr(response, "data", None)
    return (getattr(data, "content", "") or "").strip().lower()


async def main() -> None:
    global IMAGE

    names = random.sample(list(PALETTE), 4)
    IMAGE = make_quadrant_png([PALETTE[n] for n in names])

    print(f"\nRandom arrangement (1 of 360): {' '.join(names)}")
    print(f"PNG: {len(IMAGE):,} bytes\n")

    server = uvicorn.Server(
        uvicorn.Config(app, host="127.0.0.1", port=PORT, log_level="warning")
    )
    server_task = asyncio.create_task(server.serve())
    for _ in range(100):
        if getattr(server, "started", False):
            break
        await asyncio.sleep(0.1)

    results = {}
    try:
        print("control -- PNG as a blob attachment on the user turn")
        answer = await ask(
            question=INSTRUCTION,
            attachments=[
                {
                    "type": "blob",
                    "data": base64.b64encode(IMAGE).decode("ascii"),
                    "mimeType": "image/png",
                    "displayName": "test.png",
                }
            ],
        )
        ok = answer.split()[:4] == names
        results["control (attachment)"] = (ok, answer)
        print(f"  -> {answer[:60]!r}  {'CORRECT' if ok else 'WRONG'}\n")

        print("mcp -- PNG returned from an MCP tool result")
        answer = await ask(
            question=MCP_QUESTION,
            mcp_servers={
                "repro": {
                    "type": "http",
                    "url": f"http://127.0.0.1:{PORT}/mcp/",
                    "tools": ["*"],
                }
            },
        )
        ok = answer.split()[:4] == names
        results["mcp (tool result)"] = (ok, answer)
        print(f"  -> {answer[:60]!r}  tool_called={TOOL_CALLED}  "
              f"{'CORRECT' if ok else 'WRONG'}\n")
    finally:
        server.should_exit = True
        await asyncio.sleep(0.3)
        server_task.cancel()

    print(f"expected: {' '.join(names)}")
    for name, (ok, answer) in results.items():
        print(f"  {name}: {'PASS' if ok else 'FAIL'} -- {answer[:60]!r}")
    sys.exit(0 if all(ok for ok, _ in results.values()) else 1)


if __name__ == "__main__":
    asyncio.run(main())