test_structured_output_e2e.py

"""Structured-output tests against the released runtime and shared provider captures."""

import asyncio
import json

import pytest
from pydantic import BaseModel, ConfigDict, Field

from copilot import define_tool
from copilot.rpc import (
    JSONSchemaResponseFormat,
    ResponseFormat,
    ResponseFormatType,
    SendMessageItem,
    SendMessagesRequest,
)
from copilot.session import PermissionHandler
from copilot.session_events import (
    AssistantMessageData,
    AssistantMessageDeltaData,
    SessionErrorData,
    SessionIdleData,
    UserMessageData,
)

pytestmark = pytest.mark.asyncio(loop_scope="module")


class Inventory(BaseModel):
    model_config = ConfigDict(extra="forbid")
    count: int
    color: str


class Answer(BaseModel):
    model_config = ConfigDict(extra="forbid")
    answer: int


class ToolAnswer(BaseModel):
    model_config = ConfigDict(extra="forbid")
    answer: int
    contract: str


class First(BaseModel):
    model_config = ConfigDict(extra="forbid")
    first: int


class Second(BaseModel):
    model_config = ConfigDict(extra="forbid")
    second: int


def config(ctx):
    return {
        "model": "gpt-4.1",
        "available_tools": [],
        "on_permission_request": PermissionHandler.approve_all,
        "provider": {
            "type": "openai",
            "wire_api": "completions",
            "base_url": ctx.proxy_url,
            "model_id": "gpt-4.1",
            "wire_model": "gpt-4.1",
            "api_key": "fake-token-for-e2e-tests",
            "headers": {
                "Copilot-Integration-Id": "copilot-developer-cli",
                "Copilot-Harness-Id": "copilot-sdk",
                "X-GitHub-Api-Version": "2026-08-01",
            },
        },
    }


async def test_infers_typed_result_after_custom_tool(ctx):
    class AliasedInventory(BaseModel):
        model_config = ConfigDict(extra="forbid", validate_by_alias=False, validate_by_name=True)
        count_value: int = Field(alias="count")
        color_value: str = Field(alias="color")

    calls = 0

    @define_tool("get_inventory", description="Get the current widget inventory.")
    def get_inventory() -> str:
        nonlocal calls
        calls += 1
        return "The inventory contains 42 red widgets."

    async with await ctx.client.create_session(
        **config(ctx), tools=[get_inventory], streaming=True
    ) as session:
        deltas = []
        unsubscribe = session.on(
            lambda event: (
                deltas.append(event.data)
                if isinstance(event.data, AssistantMessageDeltaData)
                else None
            )
        )
        result = await session.send_and_wait_typed(
            "Call get_inventory, then report the widget count and color.",
            AliasedInventory,
            timeout=30,
        )
        assert result.count_value == 42
        assert result.color_value == "red"
        assert calls > 0
        assert deltas
        unsubscribe()
        ordinary = await session.send_and_wait(
            "Now reply with exactly the plain text HELLO, not JSON.", timeout=30
        )
        assert ordinary.data.content.strip() == "HELLO"
        exchanges = await ctx.get_exchanges()
        assert len(exchanges) >= 3
        for exchange in exchanges[:-1]:
            assert exchange["request"]["response_format"] == {
                "type": "json_schema",
                "json_schema": {
                    "name": "response",
                    "strict": True,
                    "schema": AliasedInventory.model_json_schema(),
                },
            }
        assert "response_format" not in exchanges[-1]["request"]


async def test_typed_wait_returns_stop_hook_correction(ctx):
    stops = 0

    def stop_hook(_input, _invocation):
        nonlocal stops
        stops += 1
        if stops == 1:
            return {
                "decision": "block",
                "reason": "Correct the answer to 99, not 42. Do not use tools.",
            }
        return None

    async with await ctx.client.create_session(
        **config(ctx), hooks={"on_agent_stop": stop_hook}
    ) as session:
        replies = []
        unsubscribe = session.on(
            lambda event: (
                replies.append(event.data)
                if not event.agent_id and isinstance(event.data, AssistantMessageData)
                else None
            )
        )
        try:
            result = await session.send_and_wait_typed(
                "What is 19 + 23? Do not use tools.", Answer, timeout=30
            )
            assert result.answer == 99
            assert stops == 2
            assert [Answer.model_validate_json(reply.content).answer for reply in replies] == [
                42,
                99,
            ]
            assert replies[0].originating_message_id
            assert replies[0].originating_message_id == replies[1].originating_message_id
        finally:
            unsubscribe()


async def test_send_selects_correlated_response_after_idle(ctx):
    entered = asyncio.Event()
    release = asyncio.Event()

    @define_tool("read_inventory", description="Read the current widget count and color.")
    def read_inventory() -> str:
        return "The inventory contains 42 red widgets."

    async def stop_hook(_input, _invocation):
        entered.set()
        await release.wait()
        return None

    async with await ctx.client.create_session(
        **config(ctx), tools=[read_inventory], hooks={"on_agent_stop": stop_hook}
    ) as session:
        idle = asyncio.get_running_loop().create_future()
        replies = []

        def observe(event):
            if event.agent_id:
                return
            if isinstance(event.data, AssistantMessageData):
                replies.append(event.data)
            elif isinstance(event.data, SessionIdleData) and not idle.done():
                idle.set_result(None)
            elif isinstance(event.data, SessionErrorData) and not idle.done():
                idle.set_exception(RuntimeError(event.data.message))

        unsubscribe = session.on(observe)
        try:
            origin = await session.send(
                "Call read_inventory once, then report the current widget count and color.",
                response_schema=Inventory,
            )
            await asyncio.wait_for(entered.wait(), 30)
            assert not idle.done()
            release.set()
            await asyncio.wait_for(idle, 30)
            final = next(
                reply for reply in reversed(replies) if reply.originating_message_id == origin
            )
            assert final is replies[-1]
            assert not final.tool_requests
            assert any(reply.tool_requests for reply in replies)
            assert Inventory.model_validate_json(final.content) == Inventory(count=42, color="red")
        finally:
            release.set()
            unsubscribe()


async def test_rejects_invalid_formats_before_admission(ctx):
    async with await ctx.client.create_session(**config(ctx)) as session:
        events = []
        unsubscribe = session.on(events.append)
        try:
            with pytest.raises(ValueError, match="immediate"):
                await session.send_and_wait_typed("Must not be admitted", Answer, mode="immediate")
            schema = {"type": "object", "description": "x" * (32 * 1024 * 1024)}
            with pytest.raises(Exception, match="32 MiB"):
                await session.send_and_wait(
                    "Must not be admitted", response_schema=schema, timeout=30
                )
            with pytest.raises(Exception, match="32 MiB"):
                await session.rpc.send_messages(
                    SendMessagesRequest(
                        messages=[],
                        response_format=ResponseFormat(
                            type=ResponseFormatType.JSON_SCHEMA,
                            json_schema=JSONSchemaResponseFormat(name="response", schema=schema),
                        ),
                    )
                )
            assert not (await session.rpc.queue.pending_items()).items
            assert not any(
                isinstance(event.data, (UserMessageData, SessionErrorData)) for event in events
            )
            assert not await ctx.get_exchanges()
        finally:
            unsubscribe()


async def test_sends_explicit_schema_for_message_and_batch(ctx):
    async with await ctx.client.create_session(**config(ctx)) as session:
        idle = asyncio.get_running_loop().create_future()
        replies = []

        def observe(event):
            if event.agent_id:
                return
            if isinstance(event.data, AssistantMessageData):
                replies.append(event.data)
            elif isinstance(event.data, SessionIdleData) and not idle.done():
                idle.set_result(None)
            elif isinstance(event.data, SessionErrorData) and not idle.done():
                idle.set_exception(RuntimeError(event.data.message))

        unsubscribe = session.on(observe)
        schema = Inventory.model_json_schema()
        try:
            accepted = await session.rpc.send_messages(
                SendMessagesRequest(
                    messages=[
                        SendMessageItem(prompt="There are 42 red widgets in stock."),
                        SendMessageItem(prompt="Report the widget count and color."),
                    ],
                    response_format=ResponseFormat(
                        type=ResponseFormatType.JSON_SCHEMA,
                        json_schema=JSONSchemaResponseFormat(
                            name="inventory", schema=schema, strict=True
                        ),
                    ),
                )
            )
            await asyncio.wait_for(idle, 30)
            final = next(
                item
                for item in reversed(replies)
                if item.originating_message_id == accepted.message_ids[-1]
            )
            assert Inventory.model_validate_json(final.content) == Inventory(count=42, color="red")
        finally:
            unsubscribe()
        updated = await session.send_and_wait(
            "The inventory now has 21 blue widgets. Report the new count and color.",
            response_schema=schema,
            timeout=30,
        )
        assert Inventory.model_validate_json(updated.data.content) == Inventory(
            count=21, color="blue"
        )


async def test_typed_wait_returns_stop_hook_correction_after_terminal_tool(ctx):
    calls = 0
    stops = 0

    @define_tool(
        "lookup_number",
        description="Return the number needed for the calculation.",
        is_terminal=True,
        skip_permission=True,
    )
    def lookup_number() -> int:
        nonlocal calls
        calls += 1
        return 58

    def stop_hook(_input, _invocation):
        nonlocal stops
        stops += 1
        if stops == 1:
            return {
                "decision": "block",
                "reason": "Correct the answer to 99, not 63. Do not use tools.",
            }
        return None

    async with await ctx.client.create_session(
        **config(ctx), tools=[lookup_number], hooks={"on_agent_stop": stop_hook}
    ) as session:
        replies = []
        unsubscribe = session.on(
            lambda event: (
                replies.append(event.data)
                if not event.agent_id
                and isinstance(event.data, AssistantMessageData)
                and not event.data.tool_requests
                else None
            )
        )
        try:
            result = await session.send_and_wait_typed(
                "Call lookup_number exactly once, then add 5 to the returned number. "
                "Do not guess its result.",
                Answer,
                timeout=30,
            )
            assert result.answer == 99
            assert calls == 1
            assert stops == 2
            assert [json.loads(reply.content)["answer"] for reply in replies] == [63, 99]
            assert replies[0].originating_message_id == replies[1].originating_message_id
        finally:
            unsubscribe()


async def test_typed_result_after_terminal_tool_and_steering(ctx):
    calls = 0

    @define_tool(
        "lookup_number",
        description="Return the number needed for the calculation.",
        is_terminal=True,
        skip_permission=True,
    )
    async def lookup_number() -> int:
        nonlocal calls
        calls += 1
        await session.send(
            "Continue with the original calculation. Do not call any more tools.", mode="immediate"
        )
        return 58

    async with await ctx.client.create_session(**config(ctx), tools=[lookup_number]) as session:
        result = await session.send_and_wait_typed(
            "Call lookup_number exactly once, then add 5 to the returned number. "
            "Do not guess its result.",
            ToolAnswer,
            timeout=30,
        )
        assert result == ToolAnswer(answer=63, contract="typed_tool")
        assert calls == 1
        exchanges = await ctx.get_exchanges()
        assert len(exchanges) >= 2
        for exchange in exchanges[1:]:
            assert exchange["request"]["tool_choice"] == "none"
        for exchange in exchanges:
            assert (
                exchange["request"]["response_format"]["json_schema"]["schema"]
                == ToolAnswer.model_json_schema()
            )


async def test_typed_wait_returns_late_steering_response(ctx):
    stops = 0
    steering_id = None

    async def stop_hook(_input, _invocation):
        nonlocal stops, steering_id
        stops += 1
        if stops == 1:
            steering_id = await session.send(
                "Change the answer to 99. Do not use tools.", mode="immediate"
            )
        return None

    async with await ctx.client.create_session(
        **config(ctx), hooks={"on_agent_stop": stop_hook}
    ) as session:
        result = await session.send_and_wait_typed(
            "What is 19 + 23? Do not use tools.", Answer, timeout=30
        )
        assert result.answer == 99
        assert stops == 2
        assert steering_id


async def test_concurrent_typed_sends_return_their_own_results(ctx):
    entered = asyncio.Event()
    release = asyncio.Event()

    @define_tool("first_number", description="Get the number for the first question.")
    async def first_number() -> int:
        entered.set()
        await release.wait()
        return 42

    async with await ctx.client.create_session(**config(ctx), tools=[first_number]) as session:
        first = asyncio.create_task(
            session.send_and_wait_typed(
                "Call first_number exactly once and report its returned number.", First, timeout=30
            )
        )
        try:
            await asyncio.wait_for(entered.wait(), 30)
            second = asyncio.create_task(
                session.send_and_wait_typed("What is 30 + 7? Do not use tools.", Second, timeout=30)
            )
            async with asyncio.timeout(30):
                while not (await session.rpc.queue.pending_items()).items:
                    await asyncio.sleep(0.01)
            release.set()
            assert (await first).first == 42
            assert (await second).second == 37
        finally:
            release.set()
            if not first.done():
                first.cancel()
                await asyncio.gather(first, return_exceptions=True)