test_session_e2e.py

"""E2E Session Tests"""

import base64
import os
import uuid
from datetime import datetime

import pytest

from copilot import CopilotClient, RuntimeConnection
from copilot.session import PermissionHandler
from copilot.session_events import SessionModelChangeData
from copilot.tools import Tool, ToolResult

from .testharness import (
    DEFAULT_GITHUB_TOKEN,
    E2ETestContext,
    get_next_event_of_type,
    wait_for_condition,
)

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


class TestSessions:
    async def test_should_create_and_disconnect_sessions(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all, model="claude-sonnet-5"
        )
        assert session.session_id

        messages = await session.get_events()
        assert len(messages) > 0
        assert messages[0].type.value == "session.start"
        assert messages[0].data.session_id == session.session_id
        assert messages[0].data.selected_model == "claude-sonnet-5"

        await session.disconnect()

        with pytest.raises(Exception, match="Session not found"):
            await session.get_events()

    async def test_should_have_stateful_conversation(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )

        assistant_message = await session.send_and_wait("What is 1+1?")
        assert assistant_message is not None
        assert "2" in assistant_message.data.content

        second_message = await session.send_and_wait("Now if you double that, what do you get?")
        assert second_message is not None
        assert "4" in second_message.data.content

    async def test_should_create_a_session_with_appended_systemMessage_config(
        self, ctx: E2ETestContext
    ):
        system_message_suffix = "End each response with the phrase 'Have a nice day!'"
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            system_message={"mode": "append", "content": system_message_suffix},
        )

        assistant_message = await session.send_and_wait("What is your full name?", timeout=10.0)
        assert assistant_message is not None
        assert "GitHub" in assistant_message.data.content
        assert "Have a nice day!" in assistant_message.data.content

        # Also validate the underlying traffic
        traffic = await ctx.get_exchanges()
        system_message = _get_system_message(traffic[0])
        assert "GitHub" in system_message
        assert system_message_suffix in system_message

    async def test_should_create_a_session_with_replaced_systemMessage_config(
        self, ctx: E2ETestContext
    ):
        test_system_message = "You are an assistant called Testy McTestface. Reply succinctly."
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            system_message={"mode": "replace", "content": test_system_message},
        )

        assistant_message = await session.send_and_wait("What is your full name?", timeout=10.0)
        assert assistant_message is not None
        assert "GitHub" not in assistant_message.data.content
        assert "Testy" in assistant_message.data.content

        # Also validate the underlying traffic
        traffic = await ctx.get_exchanges()
        system_message = _get_system_message(traffic[0])
        assert system_message == test_system_message  # Exact match

    async def test_should_create_a_session_with_customized_systemMessage_config(
        self, ctx: E2ETestContext
    ):
        custom_tone = "Respond in a warm, professional tone. Be thorough in explanations."
        appended_content = "Always mention quarterly earnings."
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            system_message={
                "mode": "customize",
                "sections": {
                    "tone": {"action": "replace", "content": custom_tone},
                    "code_change_rules": {"action": "remove"},
                },
                "content": appended_content,
            },
        )

        try:
            await session.send("Who are you?")

            # Validate the system message sent to the model
            traffic = await ctx.wait_for_exchanges()
            system_message = _get_system_message(traffic[0])
            assert custom_tone in system_message
            assert appended_content in system_message
            assert "<code_change_instructions>" not in system_message
        finally:
            await session.disconnect()

    async def test_should_create_a_session_with_availableTools(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            available_tools=["view", "edit"],
        )

        try:
            await session.send("What is 1+1?")

            # It only tells the model about the specified tools and no others
            traffic = await ctx.wait_for_exchanges()
            tools = traffic[0]["request"]["tools"]
            tool_names = [t["function"]["name"] for t in tools]
            assert len(tool_names) == 2
            assert "view" in tool_names
            assert "edit" in tool_names
        finally:
            await session.disconnect()

    async def test_should_create_a_session_with_excludedTools(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all, excluded_tools=["view"]
        )

        try:
            await session.send("What is 1+1?")

            # It has other tools, but not the one we excluded
            traffic = await ctx.wait_for_exchanges()
            tools = traffic[0]["request"]["tools"]
            tool_names = [t["function"]["name"] for t in tools]
            assert "edit" in tool_names
            assert "grep" in tool_names
            assert "view" not in tool_names
        finally:
            await session.disconnect()

    async def test_should_create_a_session_with_defaultAgent_excludedTools(
        self, ctx: E2ETestContext
    ):
        secret_tool = Tool(
            name="secret_tool",
            description="A secret tool hidden from the default agent",
            handler=lambda args: "SECRET",
            parameters={
                "type": "object",
                "properties": {"input": {"type": "string"}},
            },
        )

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            tools=[secret_tool],
            default_agent={"excluded_tools": ["secret_tool"]},
        )

        try:
            await session.send("What is 1+1?")

            # The real assertion: verify the runtime excluded the tool from the CAPI request
            traffic = await ctx.wait_for_exchanges()
            tools = traffic[0]["request"]["tools"]
            tool_names = [t["function"]["name"] for t in tools]
            assert "secret_tool" not in tool_names
        finally:
            await session.disconnect()

    # TODO: This test shows there's a race condition inside client.ts. If createSession
    # is called concurrently and autoStart is on, it may start multiple child processes.
    # This needs to be fixed. Right now it manifests as being unable to delete the temp
    # directories during afterAll even though we stopped all the clients.
    @pytest.mark.skip(reason="Known race condition - see TypeScript test")
    async def test_should_handle_multiple_concurrent_sessions(self, ctx: E2ETestContext):
        import asyncio

        s1, s2, s3 = await asyncio.gather(
            ctx.client.create_session(on_permission_request=PermissionHandler.approve_all),
            ctx.client.create_session(on_permission_request=PermissionHandler.approve_all),
            ctx.client.create_session(on_permission_request=PermissionHandler.approve_all),
        )

        # All sessions should have unique IDs
        session_ids = {s1.session_id, s2.session_id, s3.session_id}
        assert len(session_ids) == 3

        # All are connected
        for s in [s1, s2, s3]:
            messages = await s.get_events()
            assert len(messages) > 0
            assert messages[0].type.value == "session.start"
            assert messages[0].data.session_id == s.session_id

        # All can be disconnected
        await asyncio.gather(s1.disconnect(), s2.disconnect(), s3.disconnect())
        for s in [s1, s2, s3]:
            with pytest.raises(Exception, match="Session not found"):
                await s.get_events()

    async def test_should_resume_a_session_using_the_same_client(self, ctx: E2ETestContext):
        # Create initial session
        session1 = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        session_id = session1.session_id
        answer = await session1.send_and_wait("What is 1+1?")
        assert answer is not None
        assert "2" in answer.data.content

        # Resume using the same client
        session2 = await ctx.client.resume_session(
            session_id, on_permission_request=PermissionHandler.approve_all
        )
        assert session2.session_id == session_id
        # The completed turn's assistant message is durable; session.idle is not.
        messages = await session2.get_events()
        assert not any(message.type.value == "session.error" for message in messages)
        answers = [message for message in messages if message.type.value == "assistant.message"]
        assert answers
        assert "2" in answers[-1].data.content

        # Can continue the conversation statefully
        answer3 = await session2.send_and_wait("Now if you double that, what do you get?")
        assert answer3 is not None
        assert "4" in answer3.data.content

    async def test_should_resume_a_session_using_a_new_client(self, ctx: E2ETestContext):
        # Create initial session
        session1 = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        session_id = session1.session_id
        answer = await session1.send_and_wait("What is 1+1?")
        assert answer is not None
        assert "2" in answer.data.content

        # Resume using a new client
        github_token = (
            "fake-token-for-e2e-tests" if os.environ.get("GITHUB_ACTIONS") == "true" else None
        )
        new_client = CopilotClient(
            connection=RuntimeConnection.for_stdio(path=ctx.cli_path),
            working_directory=ctx.work_dir,
            env=ctx.get_env(),
            github_token=github_token,
        )

        try:
            session2 = await new_client.resume_session(
                session_id, on_permission_request=PermissionHandler.approve_all
            )
            assert session2.session_id == session_id

            messages = await session2.get_events()
            message_types = [m.type.value for m in messages]
            assert "user.message" in message_types
            assert "session.resume" in message_types

            # Can continue the conversation statefully
            answer2 = await session2.send_and_wait("Now if you double that, what do you get?")
            assert answer2 is not None
            assert "4" in answer2.data.content
        finally:
            await new_client.force_stop()

    async def test_resumes_a_persisted_session_from_a_new_client_when_an_mcp_oauth_handler_is_configured(  # noqa: E501
        self, ctx: E2ETestContext
    ):
        def on_mcp_auth_request(_request, _invocation):
            return {"kind": "cancelled"}

        session1 = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            on_mcp_auth_request=on_mcp_auth_request,
        )
        session_id = session1.session_id
        answer = await session1.send_and_wait("What is 1+1?")
        assert answer is not None
        assert "2" in answer.data.content

        github_token = DEFAULT_GITHUB_TOKEN if os.environ.get("GITHUB_ACTIONS") == "true" else None
        new_client = CopilotClient(
            connection=RuntimeConnection.for_stdio(path=ctx.cli_path),
            working_directory=ctx.work_dir,
            env=ctx.get_env(),
            github_token=github_token,
        )

        try:
            session2 = await new_client.resume_session(
                session_id,
                on_permission_request=PermissionHandler.approve_all,
                on_mcp_auth_request=on_mcp_auth_request,
            )
            assert session2.session_id == session_id
            await session2.disconnect()
        finally:
            await new_client.force_stop()

    async def test_should_recover_marker_after_cold_resume_with_explicit_session_id(
        self, ctx: E2ETestContext
    ):
        session_id = f"e2e-cold-resume-{uuid.uuid4()}"
        github_token = DEFAULT_GITHUB_TOKEN if os.environ.get("GITHUB_ACTIONS") == "true" else None

        client1 = CopilotClient(
            connection=RuntimeConnection.for_stdio(path=ctx.cli_path),
            working_directory=ctx.work_dir,
            env=ctx.get_env(),
            github_token=github_token,
        )
        try:
            session1 = await client1.create_session(
                on_permission_request=PermissionHandler.approve_all,
                session_id=session_id,
            )
            assert session1.session_id == session_id

            answer = await session1.send_and_wait(
                "Please remember this exact secret marker for later - MARKER-7f3ac21e. "
                'Reply with only the single word "Acknowledged".'
            )
            assert answer is not None
            assert "Acknowledged" in answer.data.content

            await session1.disconnect()
        finally:
            await client1.force_stop()

        client2 = CopilotClient(
            connection=RuntimeConnection.for_stdio(path=ctx.cli_path),
            working_directory=ctx.work_dir,
            env=ctx.get_env(),
            github_token=github_token,
        )
        try:
            session2 = await client2.resume_session(
                session_id, on_permission_request=PermissionHandler.approve_all
            )
            assert session2.session_id == session_id

            answer2 = await session2.send_and_wait(
                "What was the exact secret marker I asked you to remember earlier? "
                "Reply with only that marker value and nothing else."
            )
            assert answer2 is not None
            assert "MARKER-7f3ac21e" in answer2.data.content
        finally:
            await client2.force_stop()

    async def test_should_throw_error_resuming_nonexistent_session(self, ctx: E2ETestContext):
        with pytest.raises(Exception):
            await ctx.client.resume_session(
                "non-existent-session-id", on_permission_request=PermissionHandler.approve_all
            )

    async def test_should_list_sessions(self, ctx: E2ETestContext):
        import asyncio

        # Create a couple of sessions and send messages to persist them
        session1 = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        await session1.send_and_wait("Say hello")
        session2 = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        await session2.send_and_wait("Say goodbye")

        # Small delay to ensure session files are written to disk
        await asyncio.sleep(0.2)

        # List sessions and verify they're included
        sessions = await ctx.client.list_sessions()
        assert isinstance(sessions, list)

        session_ids = [s.session_id for s in sessions]
        assert session1.session_id in session_ids
        assert session2.session_id in session_ids

        # Verify session metadata structure
        for session_data in sessions:
            assert hasattr(session_data, "session_id")
            assert hasattr(session_data, "start_time")
            assert hasattr(session_data, "modified_time")
            assert hasattr(session_data, "is_remote")
            # summary is optional
            assert isinstance(session_data.session_id, str)
            assert isinstance(session_data.start_time, datetime)
            assert isinstance(session_data.modified_time, datetime)
            assert isinstance(session_data.is_remote, bool)

        # Verify context field is present
        for session_data in sessions:
            assert hasattr(session_data, "context")
            if session_data.context is not None:
                assert hasattr(session_data.context, "working_directory")
                assert isinstance(session_data.context.working_directory, str)

    async def test_should_delete_session(self, ctx: E2ETestContext):
        import asyncio

        # Create a session and send a message to persist it
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        await session.send_and_wait("Hello")
        session_id = session.session_id

        # Small delay to ensure session file is written to disk
        await asyncio.sleep(0.2)

        # Verify session exists in the list
        sessions = await ctx.client.list_sessions()
        session_ids = [s.session_id for s in sessions]
        assert session_id in session_ids

        # Delete the session
        await ctx.client.delete_session(session_id)

        # Verify session no longer exists in the list
        sessions_after = await ctx.client.list_sessions()
        session_ids_after = [s.session_id for s in sessions_after]
        assert session_id not in session_ids_after

        # Verify we cannot resume the deleted session
        with pytest.raises(Exception):
            await ctx.client.resume_session(
                session_id, on_permission_request=PermissionHandler.approve_all
            )

    async def test_should_get_session_metadata(self, ctx: E2ETestContext):
        # Create a session and send a message to persist it
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        await session.send_and_wait("Say hello")

        metadata = None

        async def metadata_is_available() -> bool:
            nonlocal metadata
            metadata = await ctx.client.get_session_metadata(session.session_id)
            return metadata is not None

        await wait_for_condition(
            metadata_is_available,
            timeout=10.0,
            timeout_message="Timed out waiting for session metadata to persist.",
        )

        # Get metadata for the session we just created
        assert metadata is not None
        assert metadata.session_id == session.session_id
        assert isinstance(metadata.start_time, datetime)
        assert isinstance(metadata.modified_time, datetime)
        assert isinstance(metadata.is_remote, bool)

        # Verify context field is present
        if metadata.context is not None:
            assert hasattr(metadata.context, "working_directory")
            assert isinstance(metadata.context.working_directory, str)

        # Verify non-existent session returns None
        not_found = await ctx.client.get_session_metadata("non-existent-session-id")
        assert not_found is None

    async def test_should_get_last_session_id(self, ctx: E2ETestContext):
        import asyncio

        # Create a session and send a message to persist it
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        await session.send_and_wait("Say hello")

        # Small delay to ensure session data is flushed to disk
        await asyncio.sleep(0.5)

        last_session_id = await ctx.client.get_last_session_id()
        assert last_session_id == session.session_id

        await session.disconnect()

    async def test_should_create_session_with_custom_tool(self, ctx: E2ETestContext):
        # This test uses the low-level Tool() API to show that Pydantic is optional
        def get_secret_number_handler(invocation):
            key = invocation.arguments.get("key", "") if invocation.arguments else ""
            return ToolResult(
                text_result_for_llm="54321" if key == "ALPHA" else "unknown",
                result_type="success",
            )

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            tools=[
                Tool(
                    name="get_secret_number",
                    description="Gets the secret number",
                    handler=get_secret_number_handler,
                    parameters={
                        "type": "object",
                        "properties": {"key": {"type": "string", "description": "Key"}},
                        "required": ["key"],
                    },
                )
            ],
        )

        answer = await session.send_and_wait("What is the secret number for key ALPHA?")
        assert answer is not None
        assert "54321" in answer.data.content

    async def test_should_create_session_with_custom_provider(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            provider={
                "type": "openai",
                "base_url": "https://api.openai.com/v1",
                "api_key": "fake-key",
            },
        )
        assert session.session_id

    async def test_should_create_session_with_azure_provider(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            provider={
                "type": "azure",
                "base_url": "https://my-resource.openai.azure.com",
                "api_key": "fake-key",
                "azure": {
                    "api_version": "2024-02-15-preview",
                },
            },
        )
        assert session.session_id

    async def test_should_resume_session_with_custom_provider(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )
        session_id = session.session_id

        # Resume the session with a provider
        session2 = await ctx.client.resume_session(
            session_id,
            on_permission_request=PermissionHandler.approve_all,
            provider={
                "type": "openai",
                "base_url": "https://api.openai.com/v1",
                "api_key": "fake-key",
            },
        )

        assert session2.session_id == session_id

    async def test_should_abort_a_session(self, ctx: E2ETestContext):
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )

        # Set up event listeners BEFORE sending to avoid race conditions
        wait_for_tool_start = get_next_event_of_type(session, "tool.execution_start", timeout=60.0)
        wait_for_session_idle = get_next_event_of_type(session, "session.idle", timeout=30.0)

        try:
            # Send a message that will trigger a long-running shell command
            await session.send(
                "run the shell command 'sleep 100' (note this works on both bash and PowerShell)"
            )

            # Wait for the tool to start executing
            _ = await wait_for_tool_start

            # Abort the session while the tool is running
            await session.abort()

            # Wait for session to become idle after abort
            _ = await wait_for_session_idle
        finally:
            wait_for_tool_start.cancel()
            wait_for_session_idle.cancel()
            await asyncio.gather(wait_for_tool_start, wait_for_session_idle, return_exceptions=True)

        # The session should still be alive and usable after abort
        messages = await session.get_events()
        assert len(messages) > 0

        # Verify an abort event exists in messages
        abort_events = [m for m in messages if m.type.value == "abort"]
        assert len(abort_events) > 0, "Expected an abort event in messages"

        # We should be able to send another message
        answer = await session.send_and_wait("What is 2+2?", timeout=60.0)
        assert answer is not None
        assert "4" in answer.data.content

    async def test_should_receive_session_events(self, ctx: E2ETestContext):
        import asyncio

        # Use on_event to capture events dispatched after session creation begins.
        # session.start is emitted during or shortly after the session.create RPC;
        # if the session weren't registered in the sessions map before the RPC,
        # the event would be dropped.
        early_events = []
        session_start_event = asyncio.Event()

        def capture_early(event):
            early_events.append(event)
            if event.type.value == "session.start":
                session_start_event.set()

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            on_event=capture_early,
        )

        try:
            await asyncio.wait_for(session_start_event.wait(), timeout=10)
        except TimeoutError:
            pytest.fail("Timed out waiting for session.start event")
        assert any(e.type.value == "session.start" for e in early_events)

        received_events = []
        idle_event = asyncio.Event()

        def on_event(event):
            received_events.append(event)
            if event.type.value == "session.idle":
                idle_event.set()

        session.on(on_event)

        # Send a message to trigger events
        await session.send("What is 100+200?")

        # Wait for session to become idle
        try:
            await asyncio.wait_for(idle_event.wait(), timeout=60)
        except TimeoutError:
            pytest.fail("Timed out waiting for session.idle")

        # Should have received multiple events
        assert len(received_events) > 0
        event_types = [e.type.value for e in received_events]
        assert "user.message" in event_types
        assert "assistant.message" in event_types
        assert "session.idle" in event_types

        # Idle was observed live, so inspect the messages captured for this turn.
        assert "session.error" not in event_types
        assistant_messages = [
            event for event in received_events if event.type.value == "assistant.message"
        ]
        assert assistant_messages
        assert "300" in assistant_messages[-1].data.content

    async def test_should_create_session_with_custom_config_dir(self, ctx: E2ETestContext):
        import os

        custom_config_dir = os.path.join(ctx.home_dir, "custom-config")
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all, config_directory=custom_config_dir
        )

        assert session.session_id

        # Session should work normally with custom config dir
        assistant_message = await session.send_and_wait("What is 1+1?", timeout=10.0)
        assert assistant_message is not None
        assert "2" in assistant_message.data.content

    async def test_session_log_emits_events_at_all_levels(self, ctx: E2ETestContext):
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )

        received_events = []

        def on_event(event):
            if event.type.value in ("session.info", "session.warning", "session.error"):
                received_events.append(event)

        session.on(on_event)

        await session.log("Info message")
        await session.log("Warning message", level="warning")
        await session.log("Error message", level="error")
        await session.log("Ephemeral message", ephemeral=True)

        # Poll until all 4 notification events arrive
        deadline = asyncio.get_event_loop().time() + 10
        while len(received_events) < 4:
            if asyncio.get_event_loop().time() > deadline:
                pytest.fail(
                    f"Timed out waiting for 4 notification events, got {len(received_events)}"
                )
            await asyncio.sleep(0.1)

        by_message = {e.data.message: e for e in received_events}

        assert by_message["Info message"].type.value == "session.info"
        assert by_message["Info message"].data.info_type == "notification"

        assert by_message["Warning message"].type.value == "session.warning"
        assert by_message["Warning message"].data.warning_type == "notification"

        assert by_message["Error message"].type.value == "session.error"
        assert by_message["Error message"].data.error_type == "notification"

        assert by_message["Ephemeral message"].type.value == "session.info"
        assert by_message["Ephemeral message"].data.info_type == "notification"

    async def test_should_set_model_with_reasoning_effort(self, ctx: E2ETestContext):
        """Test that setModel passes reasoningEffort and it appears in the model_change event."""
        import asyncio

        isolated_ctx = E2ETestContext()
        await isolated_ctx.setup()
        try:
            await isolated_ctx.configure_for_test(
                "session", "should_set_model_with_reasoningeffort"
            )
            session = await isolated_ctx.client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )

            model_change_event = asyncio.get_event_loop().create_future()

            def on_event(event):
                if model_change_event.done():
                    return

                match event.data:
                    case SessionModelChangeData() as data:
                        model_change_event.set_result(data)

            session.on(on_event)

            await session.set_model("gpt-5.4", reasoning_effort="high")

            data = await asyncio.wait_for(model_change_event, timeout=30)
            assert data.new_model == "gpt-5.4"
            assert data.reasoning_effort == "high"
            await session.disconnect()
        finally:
            await isolated_ctx.teardown()

    async def test_should_accept_blob_attachments(self, ctx: E2ETestContext):
        # Write the image to disk so the model can view it
        pixel_png = (
            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAY"
            "AAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhg"
            "GAWjR9awAAAABJRU5ErkJggg=="
        )
        png_path = os.path.join(ctx.work_dir, "test-pixel.png")
        with open(png_path, "wb") as f:
            f.write(base64.b64decode(pixel_png))

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )

        await session.send_and_wait(
            "Describe this image",
            attachments=[
                {
                    "type": "blob",
                    "data": pixel_png,
                    "mimeType": "image/png",
                    "displayName": "test-pixel.png",
                },
            ],
        )

        await session.disconnect()

    async def test_should_send_with_file_attachment(self, ctx: E2ETestContext):
        from copilot.session_events import UserMessageData

        file_path = os.path.join(ctx.work_dir, "attached-file.txt")
        with open(file_path, "w", encoding="utf-8") as f:
            f.write("FILE_ATTACHMENT_SENTINEL")

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        await session.send_and_wait(
            "Read the attached file and reply with its contents.",
            attachments=[
                {
                    "type": "file",
                    "displayName": "attached-file.txt",
                    "path": file_path,
                    "lineRange": {"start": 1, "end": 1},  # type: ignore[typeddict-unknown-key]
                },
            ],
        )

        messages = await session.get_events()
        user_messages = [m for m in messages if isinstance(m.data, UserMessageData)]
        assert user_messages
        attachments = user_messages[-1].data.attachments
        assert attachments is not None and len(attachments) == 1
        attachment = attachments[0]
        assert attachment.type == "file"
        assert attachment.display_name == "attached-file.txt"
        assert attachment.path == file_path
        assert attachment.line_range is not None
        assert attachment.line_range.start == 1
        assert attachment.line_range.end == 1

        await session.disconnect()

    async def test_should_send_with_directory_attachment(self, ctx: E2ETestContext):
        from copilot.session_events import UserMessageData

        directory_path = os.path.join(ctx.work_dir, "attached-directory")
        os.makedirs(directory_path, exist_ok=True)
        with open(os.path.join(directory_path, "readme.txt"), "w", encoding="utf-8") as f:
            f.write("DIRECTORY_ATTACHMENT_SENTINEL")

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        await session.send_and_wait(
            "List the attached directory.",
            attachments=[
                {
                    "type": "directory",
                    "displayName": "attached-directory",
                    "path": directory_path,
                },
            ],
        )

        messages = await session.get_events()
        user_messages = [m for m in messages if isinstance(m.data, UserMessageData)]
        assert user_messages
        attachments = user_messages[-1].data.attachments
        assert attachments is not None and len(attachments) == 1
        attachment = attachments[0]
        assert attachment.type == "directory"
        assert attachment.display_name == "attached-directory"
        assert attachment.path == directory_path

        await session.disconnect()

    async def test_should_send_with_selection_attachment(self, ctx: E2ETestContext):
        from copilot.session_events import UserMessageData

        file_path = os.path.join(ctx.work_dir, "selected-file.cs")
        with open(file_path, "w", encoding="utf-8") as f:
            f.write('class C { string Value = "SELECTION_SENTINEL"; }')

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        await session.send_and_wait(
            "Summarize the selected code.",
            attachments=[
                {
                    "type": "selection",
                    "displayName": "selected-file.cs",
                    "filePath": file_path,
                    "text": 'string Value = "SELECTION_SENTINEL";',
                    "selection": {
                        "start": {"line": 1, "character": 10},
                        "end": {"line": 1, "character": 45},
                    },
                },
            ],
        )

        messages = await session.get_events()
        user_messages = [m for m in messages if isinstance(m.data, UserMessageData)]
        assert user_messages
        attachments = user_messages[-1].data.attachments
        assert attachments is not None and len(attachments) == 1
        attachment = attachments[0]
        assert attachment.type == "selection"
        assert attachment.display_name == "selected-file.cs"
        assert attachment.file_path == file_path
        assert attachment.text == 'string Value = "SELECTION_SENTINEL";'
        assert attachment.selection is not None
        assert attachment.selection.start.line == 1
        assert attachment.selection.start.character == 10
        assert attachment.selection.end.line == 1
        assert attachment.selection.end.character == 45

        await session.disconnect()

    async def test_should_send_with_custom_requestheaders(self, ctx: E2ETestContext):
        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        await session.send_and_wait(
            "What is 1+1?",
            request_headers={"x-copilot-sdk-test-header": "python-request-headers"},
        )

        exchanges = await ctx.get_exchanges()
        assert exchanges
        last_headers = exchanges[-1].get("requestHeaders") or {}
        normalized = {k.lower(): str(v) for k, v in last_headers.items()}
        header_value = normalized.get("x-copilot-sdk-test-header", "")
        assert "python-request-headers" in header_value

        await session.disconnect()

    async def test_should_list_sessions_with_context(self, ctx: E2ETestContext):
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )
        await session.send_and_wait("Say OK.")

        # Allow the session to flush metadata to disk before reading it back.
        our_session = None
        for _ in range(50):
            sessions = await ctx.client.list_sessions()
            our_session = next((s for s in sessions if s.session_id == session.session_id), None)
            if our_session is not None:
                break
            await asyncio.sleep(0.1)
        assert our_session is not None

        all_sessions = await ctx.client.list_sessions()
        assert all_sessions

        if our_session.context is not None:
            assert (
                isinstance(our_session.context.working_directory, str)
                and our_session.context.working_directory
            )

        await session.disconnect()

    async def test_should_get_session_metadata_by_id(self, ctx: E2ETestContext):
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )
        await session.send_and_wait("Say hello")

        metadata = None
        for _ in range(50):
            metadata = await ctx.client.get_session_metadata(session.session_id)
            if metadata is not None:
                break
            await asyncio.sleep(0.1)
        assert metadata is not None
        assert metadata.session_id == session.session_id
        assert isinstance(metadata.start_time, datetime)
        assert isinstance(metadata.modified_time, datetime)

        not_found = await ctx.client.get_session_metadata("non-existent-session-id")
        assert not_found is None

        await session.disconnect()

    async def test_send_returns_immediately_while_events_stream_in_background(
        self, ctx: E2ETestContext
    ):
        """`send` returns before the session goes idle; events are streamed."""
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )
        events = []

        def on_event(event):
            events.append(event)

        unsubscribe = session.on(on_event)
        idle_task = get_next_event_of_type(session, "session.idle", timeout=10.0)
        try:
            # Use a slow command so we can verify send() returns before completion
            await session.send("Run 'sleep 2 && echo done'")

            # send() should return before turn completes (no session.idle yet)
            assert not any(event.type.value == "session.idle" for event in events)

            await idle_task
            messages = [event for event in events if event.type.value == "assistant.message"]
            assert messages
            assert "done" in messages[-1].data.content
            assert any(event.type.value == "session.idle" for event in events)
        finally:
            idle_task.cancel()
            await asyncio.gather(idle_task, return_exceptions=True)
            unsubscribe()
            await session.disconnect()

    async def test_sendandwait_blocks_until_session_idle_and_returns_final_assistant_message(
        self, ctx: E2ETestContext
    ):
        """The synchronous root-idle listener finishes before `send_and_wait` returns."""
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )
        events: list[str] = []
        ordering: list[str] = []
        waiter_pending: list[bool] = []
        listener_loops = []
        loop = asyncio.get_running_loop()

        def on_event(event):
            events.append(event.type.value)
            if event.type.value == "session.idle" and not event.agent_id:
                ordering.append("listener entered")
                listener_loops.append(asyncio.get_running_loop())
                waiter_pending.append(not send_task.done())
                # Dispatch is synchronous on the asyncio loop, even though the
                # JSON-RPC reader is a thread. A blocking gate would deadlock it.
                ordering.append("listener completed")

        unsubscribe = session.on(on_event)
        try:
            send_task = asyncio.create_task(session.send_and_wait("What is 2+2?"))
            response = await send_task
            ordering.append("wait returned")
            assert ordering == ["listener entered", "listener completed", "wait returned"]
            assert waiter_pending == [True]
            assert listener_loops == [loop]
            assert response is not None
            assert response.type.value == "assistant.message"
            assert "4" in (response.data.content or "")
            assert "session.idle" in events
            assert "assistant.message" in events
        finally:
            unsubscribe()
            await session.disconnect()

    async def test_sendandwait_throws_on_timeout(self, ctx: E2ETestContext):
        """`send_and_wait` raises TimeoutError when the session does not become idle."""
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        # Subscribe before sending so even an idle emitted before the abort reply is captured.
        idle_task = get_next_event_of_type(session, "session.idle", timeout=30.0)
        try:
            with pytest.raises(TimeoutError) as exc_info:
                await session.send_and_wait(
                    "Run 'sleep 2 && echo done'",
                    timeout=0.1,
                )
            assert "Timeout" in str(exc_info.value) or "timed out" in str(exc_info.value).lower()

            # The timeout only cancels the client-side wait; abort the agent and wait for idle
            # so leftover requests don't leak into subsequent tests.
            await session.abort()
            await idle_task
        finally:
            idle_task.cancel()
            await asyncio.gather(idle_task, return_exceptions=True)
            await session.disconnect()

    async def test_sendandwait_throws_operationcanceledexception_when_token_cancelled(
        self, ctx: E2ETestContext
    ):
        """`send_and_wait` raises CancelledError when the surrounding task is cancelled."""
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        tool_start_task = get_next_event_of_type(session, "tool.execution_start", timeout=60.0)
        idle_task = get_next_event_of_type(session, "session.idle", timeout=30.0)

        send_task = asyncio.create_task(
            session.send_and_wait(
                "run the shell command 'sleep 10' (note this works on both bash and PowerShell)",
                timeout=120.0,
            )
        )

        try:
            # Wait for the tool to begin executing before cancelling.
            await tool_start_task

            send_task.cancel()
            with pytest.raises((asyncio.CancelledError, BaseException)):
                await send_task

            # Cancelling only cancels the client-side wait; abort and wait for idle.
            await session.abort()
            await idle_task
        finally:
            tool_start_task.cancel()
            idle_task.cancel()
            send_task.cancel()
            await asyncio.gather(tool_start_task, idle_task, send_task, return_exceptions=True)
            await session.disconnect()

    async def test_should_set_model_on_existing_session(self, ctx: E2ETestContext):
        """`set_model` emits a session.model_change event with the new model."""
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )

        model_change_event: asyncio.Future[SessionModelChangeData] = (
            asyncio.get_event_loop().create_future()
        )

        def on_event(event):
            if model_change_event.done():
                return
            match event.data:
                case SessionModelChangeData() as data:
                    model_change_event.set_result(data)

        session.on(on_event)

        await session.set_model("gpt-4.1")

        data = await asyncio.wait_for(model_change_event, timeout=30)
        assert data.new_model == "gpt-4.1"

        await session.disconnect()

    async def test_handler_exception_does_not_halt_event_delivery(self, ctx: E2ETestContext):
        """A throwing handler does not stop subsequent events from being delivered."""
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        event_count = 0
        idle_event = asyncio.Event()

        def handler(event):
            nonlocal event_count
            event_count += 1
            if event_count == 1:
                raise RuntimeError("boom")
            if event.type.value == "session.idle":
                idle_event.set()

        session.on(handler)

        await session.send("What is 1+1?")

        try:
            await asyncio.wait_for(idle_event.wait(), timeout=30.0)
        except TimeoutError:
            pytest.fail("Timed out waiting for session.idle after handler exception")

        # Handler saw more than just the first (throwing) event.
        assert event_count > 1

        await session.disconnect()

    async def test_disposeasync_from_handler_does_not_deadlock(self, ctx: E2ETestContext):
        """Calling `disconnect` from inside a handler must not deadlock.

        Named to match the C# snapshot file `disposeasync_from_handler_does_not_deadlock.yaml`.
        """
        import asyncio

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        disposed = asyncio.Event()
        disconnect_started = False

        def handler(event):
            nonlocal disconnect_started
            # Disconnect once the assistant.message has arrived (CAPI has completed),
            # so we don't leak in-flight CAPI requests into a sibling test's snapshot.
            if event.type.value == "assistant.message" and not disconnect_started:
                disconnect_started = True

                async def _disconnect():
                    try:
                        await session.disconnect()
                    finally:
                        disposed.set()

                asyncio.get_event_loop().create_task(_disconnect())

        session.on(handler)

        await session.send("What is 1+1?")

        try:
            await asyncio.wait_for(disposed.wait(), timeout=10.0)
        except TimeoutError:
            pytest.fail("disconnect from within handler appears to have deadlocked")

    async def test_should_send_with_mode_property(self, ctx: E2ETestContext):
        """Per-message `agent_mode` is forwarded and echoed back on user.message."""
        from copilot.session_events import (
            UserMessageAgentMode,
            UserMessageData,
        )

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
        )

        await session.send_and_wait(
            "Say mode ok.",
            agent_mode="plan",
        )

        messages = await session.get_events()
        user_messages = [m for m in messages if isinstance(m.data, UserMessageData)]
        assert user_messages
        last = user_messages[-1].data
        assert last.content == "Say mode ok."
        assert last.agent_mode == UserMessageAgentMode.PLAN

        await session.disconnect()


def _get_system_message(exchange: dict) -> str:
    messages = exchange.get("request", {}).get("messages", [])
    for msg in messages:
        if msg.get("role") == "system":
            return msg.get("content", "")
    return ""