test_commands_and_elicitation.py

"""
Unit tests for Commands, UI Elicitation (client→server), and
onElicitationContext (server→client callback) features.

Mirrors the Node.js client.test.ts tests for these features.
"""

import asyncio
from collections.abc import Callable

import pytest

from copilot import CopilotClient, RuntimeConnection
from copilot.session import (
    AutoModeSwitchRequest,
    AutoModeSwitchResponse,
    CommandContext,
    CommandDefinition,
    ElicitationContext,
    ElicitationResult,
    ExitPlanModeRequest,
    ExitPlanModeResult,
    PermissionHandler,
)
from e2e.testharness import CLI_PATH


async def _wait_for(predicate: Callable[[], bool], timeout: float = 2.0) -> None:
    """Poll predicate until True or timeout. Replaces brittle ``asyncio.sleep`` waits.

    Used in unit tests where we dispatch an event and need to wait for the consumer
    coroutine to invoke a handler and (sometimes) for the handler to issue an RPC
    that our mock captures. Polling at 5ms means fast machines exit quickly while
    slow machines still get up to ``timeout`` seconds before the test fails.
    """
    deadline = asyncio.get_event_loop().time() + timeout
    while not predicate():
        if asyncio.get_event_loop().time() >= deadline:
            raise AssertionError(f"Condition not met within {timeout}s")
        await asyncio.sleep(0.005)


# ============================================================================
# Commands
# ============================================================================


class TestCommands:
    @pytest.mark.asyncio
    async def test_forwards_commands_in_session_create_rpc(self):
        """Verifies that commands (name + description) are serialized in session.create payload."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            captured: dict = {}
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                captured[method] = params
                return await original_request(method, params)

            client._client.request = mock_request

            await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                commands=[
                    CommandDefinition(
                        name="deploy",
                        description="Deploy the app",
                        handler=lambda ctx: None,
                    ),
                    CommandDefinition(
                        name="rollback",
                        handler=lambda ctx: None,
                    ),
                ],
            )

            payload = captured["session.create"]
            assert payload["commands"] == [
                {"name": "deploy", "description": "Deploy the app"},
                {"name": "rollback", "description": ""},
            ]
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_forwards_commands_in_session_resume_rpc(self):
        """Verifies that commands are serialized in session.resume payload."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )

            captured: dict = {}

            async def mock_request(method, params, **kwargs):
                captured[method] = params
                if method == "session.resume":
                    return {"sessionId": params["sessionId"]}
                raise RuntimeError(f"Unexpected method: {method}")

            client._client.request = mock_request

            await client.resume_session(
                session.session_id,
                on_permission_request=PermissionHandler.approve_all,
                commands=[
                    CommandDefinition(
                        name="deploy",
                        description="Deploy",
                        handler=lambda ctx: None,
                    ),
                ],
            )

            payload = captured["session.resume"]
            assert payload["commands"] == [{"name": "deploy", "description": "Deploy"}]
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_routes_command_execute_event_to_correct_handler(self):
        """Verifies the command dispatch works for command.execute events."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            handler_calls: list[CommandContext] = []

            async def deploy_handler(ctx: CommandContext) -> None:
                handler_calls.append(ctx)

            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                commands=[
                    CommandDefinition(name="deploy", handler=deploy_handler),
                ],
            )

            # Mock the RPC so handlePendingCommand doesn't fail
            rpc_calls: list[tuple] = []
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                if method == "session.commands.handlePendingCommand":
                    rpc_calls.append((method, params))
                    return {"success": True}
                return await original_request(method, params)

            client._client.request = mock_request

            # Simulate a command.execute broadcast event
            from copilot.session_events import (
                CommandExecuteData,
                SessionEvent,
                SessionEventType,
            )

            event = SessionEvent(
                data=CommandExecuteData(
                    request_id="req-1",
                    command="/deploy production",
                    command_name="deploy",
                    args="production",
                ),
                id="evt-1",
                timestamp="2025-01-01T00:00:00Z",
                type=SessionEventType.COMMAND_EXECUTE,
                ephemeral=True,
                parent_id=None,
            )
            session._dispatch_event(event)

            # Wait for the consumer coroutine to invoke the handler and the handler
            # to issue the handlePendingCommand RPC that our mock captures.
            await _wait_for(lambda: len(handler_calls) >= 1 and len(rpc_calls) >= 1)

            assert len(handler_calls) == 1
            assert handler_calls[0].session_id == session.session_id
            assert handler_calls[0].command == "/deploy production"
            assert handler_calls[0].command_name == "deploy"
            assert handler_calls[0].args == "production"

            # Verify handlePendingCommand was called
            assert len(rpc_calls) >= 1
            assert rpc_calls[0][1]["requestId"] == "req-1"
            # No error key means success
            assert "error" not in rpc_calls[0][1] or rpc_calls[0][1].get("error") is None
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_sends_error_when_command_handler_throws(self):
        """Verifies error is sent via RPC when a command handler raises."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:

            def fail_handler(ctx: CommandContext) -> None:
                raise RuntimeError("deploy failed")

            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                commands=[
                    CommandDefinition(name="fail", handler=fail_handler),
                ],
            )

            rpc_calls: list[tuple] = []
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                if method == "session.commands.handlePendingCommand":
                    rpc_calls.append((method, params))
                    return {"success": True}
                return await original_request(method, params)

            client._client.request = mock_request

            from copilot.session_events import (
                CommandExecuteData,
                SessionEvent,
                SessionEventType,
            )

            event = SessionEvent(
                data=CommandExecuteData(
                    request_id="req-2",
                    command="/fail",
                    command_name="fail",
                    args="",
                ),
                id="evt-2",
                timestamp="2025-01-01T00:00:00Z",
                type=SessionEventType.COMMAND_EXECUTE,
                ephemeral=True,
                parent_id=None,
            )
            session._dispatch_event(event)

            await _wait_for(lambda: len(rpc_calls) >= 1)

            assert len(rpc_calls) >= 1
            assert rpc_calls[0][1]["requestId"] == "req-2"
            assert "deploy failed" in rpc_calls[0][1]["error"]
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_sends_error_for_unknown_command(self):
        """Verifies error is sent via RPC for an unrecognized command."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                commands=[
                    CommandDefinition(name="deploy", handler=lambda ctx: None),
                ],
            )

            rpc_calls: list[tuple] = []
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                if method == "session.commands.handlePendingCommand":
                    rpc_calls.append((method, params))
                    return {"success": True}
                return await original_request(method, params)

            client._client.request = mock_request

            from copilot.session_events import (
                CommandExecuteData,
                SessionEvent,
                SessionEventType,
            )

            event = SessionEvent(
                data=CommandExecuteData(
                    request_id="req-3",
                    command="/unknown",
                    command_name="unknown",
                    args="",
                ),
                id="evt-3",
                timestamp="2025-01-01T00:00:00Z",
                type=SessionEventType.COMMAND_EXECUTE,
                ephemeral=True,
                parent_id=None,
            )
            session._dispatch_event(event)

            await _wait_for(lambda: len(rpc_calls) >= 1)

            assert len(rpc_calls) >= 1
            assert rpc_calls[0][1]["requestId"] == "req-3"
            assert "Unknown command" in rpc_calls[0][1]["error"]
        finally:
            await client.force_stop()


# ============================================================================
# UI Elicitation (client → server)
# ============================================================================


class TestUiElicitation:
    @pytest.mark.asyncio
    async def test_reads_capabilities_from_session_create_response(self):
        """Verifies capabilities are parsed from session.create response."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                if method == "session.create":
                    result = await original_request(method, params)
                    return {**result, "capabilities": {"ui": {"elicitation": True}}}
                return await original_request(method, params)

            client._client.request = mock_request

            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )
            assert session.capabilities == {"ui": {"elicitation": True}}
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_defaults_capabilities_when_not_injected(self):
        """Verifies capabilities default to empty when server returns none."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )
            # CLI returns actual capabilities; in headless mode, elicitation is
            # either False or absent. Just verify we don't crash.
            ui_caps = session.capabilities.get("ui", {})
            assert ui_caps.get("elicitation") in (False, None, True)
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_elicitation_throws_when_capability_is_missing(self):
        """Verifies that UI methods throw when elicitation is not supported."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )
            # Force capabilities to not support elicitation
            session._set_capabilities({})

            with pytest.raises(RuntimeError, match="not supported"):
                await session.ui.elicitation(
                    {
                        "message": "Enter name",
                        "requestedSchema": {
                            "type": "object",
                            "properties": {"name": {"type": "string", "minLength": 1}},
                            "required": ["name"],
                        },
                    }
                )
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_confirm_throws_when_capability_is_missing(self):
        """Verifies confirm throws when elicitation is not supported."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )
            session._set_capabilities({})

            with pytest.raises(RuntimeError, match="not supported"):
                await session.ui.confirm("Deploy?")
        finally:
            await client.force_stop()


# ============================================================================
# onElicitationContext (server → client callback)
# ============================================================================


class TestOnElicitationContext:
    @pytest.mark.asyncio
    async def test_sends_request_elicitation_flag_when_handler_provided(self):
        """Verifies requestElicitation=true is sent when onElicitationContext is provided."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            captured: dict = {}
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                captured[method] = params
                return await original_request(method, params)

            client._client.request = mock_request

            async def elicitation_handler(
                context: ElicitationContext,
            ) -> ElicitationResult:
                return {"action": "accept", "content": {}}

            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                on_elicitation_request=elicitation_handler,
            )
            assert session is not None

            payload = captured["session.create"]
            assert payload["requestElicitation"] is True
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_does_not_send_request_elicitation_when_no_handler(self):
        """Verifies requestElicitation=false when no handler is provided."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            captured: dict = {}
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                captured[method] = params
                return await original_request(method, params)

            client._client.request = mock_request

            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
            )
            assert session is not None

            payload = captured["session.create"]
            assert payload["requestElicitation"] is False
            assert payload["requestExitPlanMode"] is False
            assert payload["requestAutoModeSwitch"] is False
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_sends_mode_callback_flags_when_handlers_provided(self):
        """Verifies mode callback flags are sent when handlers are provided."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            captured: dict = {}
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                captured[method] = params
                return await original_request(method, params)

            client._client.request = mock_request

            def exit_handler(
                request: ExitPlanModeRequest, invocation: dict[str, str]
            ) -> ExitPlanModeResult:
                return {"approved": True}

            def auto_handler(
                request: AutoModeSwitchRequest, invocation: dict[str, str]
            ) -> AutoModeSwitchResponse:
                return "yes_always"

            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                on_exit_plan_mode_request=exit_handler,
                on_auto_mode_switch_request=auto_handler,
            )
            assert session is not None

            payload = captured["session.create"]
            assert payload["requestExitPlanMode"] is True
            assert payload["requestAutoModeSwitch"] is True
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_sends_mode_callback_flags_on_resume_when_handlers_provided(self):
        """Verifies mode callback flags are sent on session.resume."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )
            captured: dict = {}

            async def mock_request(method, params, **kwargs):
                captured[method] = params
                if method == "session.resume":
                    return {"sessionId": params["sessionId"]}
                raise RuntimeError(f"Unexpected method: {method}")

            client._client.request = mock_request

            await client.resume_session(
                session.session_id,
                on_permission_request=PermissionHandler.approve_all,
                on_exit_plan_mode_request=lambda request, invocation: {"approved": True},
                on_auto_mode_switch_request=lambda request, invocation: "yes",
            )

            payload = captured["session.resume"]
            assert payload["requestExitPlanMode"] is True
            assert payload["requestAutoModeSwitch"] is True
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_dispatches_mode_callback_requests_to_registered_handlers(self):
        """Verifies direct mode requests are dispatched to registered handlers."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:

            async def exit_handler(
                request: ExitPlanModeRequest, invocation: dict[str, str]
            ) -> ExitPlanModeResult:
                assert invocation["session_id"] == session.session_id
                assert request["summary"] == "Review the plan"
                assert request["planContent"] == "Plan body"
                assert request["actions"] == ["interactive", "autopilot"]
                assert request["recommendedAction"] == "autopilot"
                return {
                    "approved": True,
                    "selectedAction": "interactive",
                    "feedback": "Looks good",
                }

            async def auto_handler(
                request: AutoModeSwitchRequest, invocation: dict[str, str]
            ) -> AutoModeSwitchResponse:
                assert invocation["session_id"] == session.session_id
                assert request["errorCode"] == "user_weekly_rate_limited"
                assert request["retryAfterSeconds"] == 3600
                return "yes_always"

            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all,
                on_exit_plan_mode_request=exit_handler,
                on_auto_mode_switch_request=auto_handler,
            )

            exit_result = await client._handle_exit_plan_mode_request(
                {
                    "sessionId": session.session_id,
                    "summary": "Review the plan",
                    "planContent": "Plan body",
                    "actions": ["interactive", "autopilot"],
                    "recommendedAction": "autopilot",
                }
            )
            assert exit_result == {
                "approved": True,
                "selectedAction": "interactive",
                "feedback": "Looks good",
            }

            auto_result = await client._handle_auto_mode_switch_request(
                {
                    "sessionId": session.session_id,
                    "errorCode": "user_weekly_rate_limited",
                    "retryAfterSeconds": 3600,
                }
            )
            assert auto_result == {"response": "yes_always"}
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_sends_cancel_when_elicitation_handler_throws(self):
        """Verifies auto-cancel when the elicitation handler raises."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:

            async def bad_handler(
                context: ElicitationContext,
            ) -> ElicitationResult:
                raise RuntimeError("handler exploded")

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

            rpc_calls: list[tuple] = []
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                if method == "session.ui.handlePendingElicitation":
                    rpc_calls.append((method, params))
                    return {"success": True}
                return await original_request(method, params)

            client._client.request = mock_request

            # Call _handle_elicitation_request directly (as Node.js test does)
            await session._handle_elicitation_request(
                {"session_id": session.session_id, "message": "Pick a color"}, "req-123"
            )

            assert len(rpc_calls) >= 1
            cancel_call = next(
                (call for call in rpc_calls if call[1].get("result", {}).get("action") == "cancel"),
                None,
            )
            assert cancel_call is not None
            assert cancel_call[1]["requestId"] == "req-123"
            assert cancel_call[1]["result"]["action"] == "cancel"
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_dispatches_elicitation_requested_event_to_handler(self):
        """Verifies that an elicitation.requested event dispatches to the handler."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            handler_calls: list = []

            async def elicitation_handler(
                context: ElicitationContext,
            ) -> ElicitationResult:
                handler_calls.append(context)
                return {"action": "accept", "content": {"color": "blue"}}

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

            rpc_calls: list[tuple] = []
            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                if method == "session.ui.handlePendingElicitation":
                    rpc_calls.append((method, params))
                    return {"success": True}
                return await original_request(method, params)

            client._client.request = mock_request

            from copilot.session_events import (
                ElicitationRequestedData,
                SessionEvent,
                SessionEventType,
            )

            event = SessionEvent(
                data=ElicitationRequestedData(
                    request_id="req-elicit-1",
                    message="Pick a color",
                ),
                id="evt-elicit-1",
                timestamp="2025-01-01T00:00:00Z",
                type=SessionEventType.ELICITATION_REQUESTED,
                ephemeral=True,
                parent_id=None,
            )
            session._dispatch_event(event)

            await _wait_for(lambda: len(handler_calls) >= 1 and len(rpc_calls) >= 1)

            assert len(handler_calls) == 1
            assert handler_calls[0]["message"] == "Pick a color"

            assert len(rpc_calls) >= 1
            assert rpc_calls[0][1]["requestId"] == "req-elicit-1"
            assert rpc_calls[0][1]["result"]["action"] == "accept"
        finally:
            await client.force_stop()

    @pytest.mark.asyncio
    async def test_elicitation_handler_receives_full_schema(self):
        """Verifies that requestedSchema passes type, properties, and required to handler."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            handler_calls: list = []

            async def elicitation_handler(
                context: ElicitationContext,
            ) -> ElicitationResult:
                handler_calls.append(context)
                return {"action": "cancel"}

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

            original_request = client._client.request

            async def mock_request(method, params, **kwargs):
                if method == "session.ui.handlePendingElicitation":
                    return {"success": True}
                return await original_request(method, params)

            client._client.request = mock_request

            from copilot.session_events import (
                ElicitationRequestedData,
                ElicitationRequestedSchema,
                SessionEvent,
                SessionEventType,
            )

            event = SessionEvent(
                data=ElicitationRequestedData(
                    request_id="req-schema-1",
                    message="Fill in your details",
                    requested_schema=ElicitationRequestedSchema(
                        type="object",
                        properties={
                            "name": {"type": "string"},
                            "age": {"type": "number"},
                        },
                        required=["name", "age"],
                    ),
                ),
                id="evt-schema-1",
                timestamp="2025-01-01T00:00:00Z",
                type=SessionEventType.ELICITATION_REQUESTED,
                ephemeral=True,
                parent_id=None,
            )
            session._dispatch_event(event)

            await _wait_for(lambda: len(handler_calls) >= 1)

            assert len(handler_calls) == 1
            schema = handler_calls[0].get("requestedSchema")
            assert schema is not None, "Expected requestedSchema in handler call"
            assert schema["type"] == "object"
            assert "name" in schema["properties"]
            assert "age" in schema["properties"]
            assert schema["required"] == ["name", "age"]
        finally:
            await client.force_stop()


# ============================================================================
# Capabilities changed event
# ============================================================================


class TestCapabilitiesChanged:
    @pytest.mark.asyncio
    async def test_capabilities_changed_event_updates_session(self):
        """Verifies that a capabilities.changed event updates session capabilities."""
        client = CopilotClient(connection=RuntimeConnection.for_stdio(path=CLI_PATH))
        await client.start()

        try:
            session = await client.create_session(
                on_permission_request=PermissionHandler.approve_all
            )
            session._set_capabilities({})

            from copilot.session_events import (
                CapabilitiesChangedData,
                CapabilitiesChangedUI,
                SessionEvent,
                SessionEventType,
            )

            event = SessionEvent(
                data=CapabilitiesChangedData(ui=CapabilitiesChangedUI(elicitation=True)),
                id="evt-cap-1",
                timestamp="2025-01-01T00:00:00Z",
                type=SessionEventType.CAPABILITIES_CHANGED,
                ephemeral=True,
                parent_id=None,
            )
            session._dispatch_event(event)

            assert session.capabilities.get("ui", {}).get("elicitation") is True
        finally:
            await client.force_stop()