"""Skill provider unit tests.""" import asyncio import io import json from collections.abc import Awaitable, Callable from typing import Any from unittest.mock import AsyncMock, Mock import pytest from copilot import ( CloudSessionOptions, CopilotClient, RuntimeConnection, SkillProvider, SkillProviderDescriptor, ) from copilot._jsonrpc import JsonRpcClient, JsonRpcError from copilot.session import CopilotSession class _SyncSkillProvider: def __init__( self, *, skills: list[SkillProviderDescriptor] | None = None, markdown: dict[str, str | None] | None = None, ) -> None: self.skills = skills self.markdown = markdown or {} self.list_calls = 0 self.read_calls: list[str] = [] def list_skills(self) -> list[SkillProviderDescriptor] | None: self.list_calls += 1 return self.skills def read_skill(self, name: str) -> str | None: self.read_calls.append(name) return self.markdown.get(name) class _AsyncSkillProvider: def __init__( self, *, skills: list[SkillProviderDescriptor] | None = None, markdown: dict[str, str | None] | None = None, ) -> None: self._sync = _SyncSkillProvider(skills=skills, markdown=markdown) @property def list_calls(self) -> int: return self._sync.list_calls @property def read_calls(self) -> list[str]: return self._sync.read_calls async def list_skills(self) -> list[SkillProviderDescriptor] | None: await asyncio.sleep(0) return self._sync.list_skills() async def read_skill(self, name: str) -> str | None: await asyncio.sleep(0) return self._sync.read_skill(name) class _FailingSkillProvider: def list_skills(self) -> list[SkillProviderDescriptor]: raise RuntimeError("secret-list-token") def read_skill(self, name: str) -> str: raise RuntimeError("secret-read-token") class _WriteOnlyProcess: def __init__(self) -> None: self.stdin = io.BytesIO() self.stdout = io.BytesIO() self.stderr = None def poll(self) -> None: return None def _descriptor(name: str, description: str = "Test skill") -> SkillProviderDescriptor: return SkillProviderDescriptor(name=name, description=description) def _decode_jsonrpc_output(process: _WriteOnlyProcess) -> dict[str, Any]: _header, payload = process.stdin.getvalue().split(b"\r\n\r\n", 1) return json.loads(payload) async def _assert_no_provider_error(session: CopilotSession, provider: _SyncSkillProvider) -> None: list_calls = provider.list_calls read_calls = list(provider.read_calls) with pytest.raises(JsonRpcError) as exc_info: await session._handle_skill_provider_list() assert exc_info.value.code == -32603 assert exc_info.value.message == f"No skill provider for session: {session.session_id}" assert exc_info.value.data is None assert provider.list_calls == list_calls assert provider.read_calls == read_calls def _install_fake_client( client: CopilotClient, handler: Callable[[str, dict, dict[str, Any]], Awaitable[dict[str, Any]]], ) -> None: async def request(method: str, params: dict, **kwargs: Any) -> dict[str, Any]: return await handler(method, params, kwargs) client._client = Mock() client._client.request = request async def _default_request(method: str, params: dict, kwargs: dict[str, Any]) -> dict[str, Any]: if method == "session.create": result = {"sessionId": params.get("sessionId") or "created-session", "workspacePath": None} callback = kwargs.get("on_response_inline") if callback is not None: callback(result) return result if method == "session.resume": return {"sessionId": params["sessionId"], "workspacePath": None} if method == "session.options.update": return {"success": True} return {} def test_skill_provider_is_runtime_checkable_and_exported() -> None: assert isinstance(_SyncSkillProvider(), SkillProvider) assert SkillProviderDescriptor(name="demo", description="Demo").to_dict() == { "description": "Demo", "name": "demo", } @pytest.mark.asyncio async def test_create_and_resume_send_skill_provider_flag_only_when_configured() -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) captured: list[tuple[str, dict]] = [] async def request(method: str, params: dict, kwargs: dict[str, Any]) -> dict[str, Any]: captured.append((method, dict(params))) return await _default_request(method, params, kwargs) _install_fake_client(client, request) provider = _SyncSkillProvider(skills=[_descriptor("provided")]) await client.create_session(session_id="create-with-provider", skill_provider=provider) await client.resume_session("resume-with-provider", skill_provider=provider) await client.create_session(session_id="create-without-provider") await client.resume_session("resume-without-provider") payloads = {(method, params["sessionId"]): params for method, params in captured} assert payloads[("session.create", "create-with-provider")]["hasSkillProvider"] is True assert payloads[("session.resume", "resume-with-provider")]["hasSkillProvider"] is True assert "hasSkillProvider" not in payloads[("session.create", "create-without-provider")] assert "hasSkillProvider" not in payloads[("session.resume", "resume-without-provider")] @pytest.mark.asyncio async def test_empty_mode_keeps_skill_provider_dormant_by_default() -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234"), mode="empty") captured: list[tuple[str, dict]] = [] async def request(method: str, params: dict, kwargs: dict[str, Any]) -> dict[str, Any]: captured.append((method, dict(params))) return await _default_request(method, params, kwargs) _install_fake_client(client, request) provider = _SyncSkillProvider(skills=[_descriptor("provided")]) await client.create_session( session_id="empty-create", available_tools=[], skill_provider=provider, ) await client.resume_session( "empty-resume", available_tools=[], skill_provider=provider, ) payloads = {(method, params["sessionId"]): params for method, params in captured} assert payloads[("session.create", "empty-create")]["enableSkills"] is False assert payloads[("session.resume", "empty-resume")]["enableSkills"] is False assert payloads[("session.create", "empty-create")]["hasSkillProvider"] is True assert payloads[("session.resume", "empty-resume")]["hasSkillProvider"] is True @pytest.mark.asyncio async def test_cloud_create_rejects_skill_provider_before_connecting() -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) start_mock = AsyncMock() setattr(client, "start", start_mock) provider = _SyncSkillProvider(skills=[_descriptor("cloud")]) with pytest.raises(ValueError, match="Skill providers are not supported for cloud sessions\\."): await client.create_session( cloud=CloudSessionOptions(), skill_provider=provider, ) start_mock.assert_not_awaited() assert provider.list_calls == 0 assert provider.read_calls == [] @pytest.mark.asyncio async def test_early_skill_provider_callbacks_are_served_during_create_and_resume() -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) observed: list[dict] = [] provider = _SyncSkillProvider( skills=[_descriptor("early")], markdown={"early": "early skill body"}, ) async def request(method: str, params: dict, kwargs: dict[str, Any]) -> dict[str, Any]: if method in ("session.create", "session.resume"): observed.append( await client._handle_skill_provider_list({"sessionId": params["sessionId"]}) ) observed.append( await client._handle_skill_provider_read( {"sessionId": params["sessionId"], "name": "early"} ) ) return await _default_request(method, params, kwargs) _install_fake_client(client, request) await client.create_session(session_id="early-create", skill_provider=provider) await client.resume_session("early-resume", skill_provider=provider) assert observed == [ {"skills": [{"description": "Test skill", "name": "early"}]}, {"markdown": "early skill body"}, {"skills": [{"description": "Test skill", "name": "early"}]}, {"markdown": "early skill body"}, ] @pytest.mark.parametrize("provider_cls", [_SyncSkillProvider, _AsyncSkillProvider]) @pytest.mark.asyncio async def test_skill_provider_dispatch_serializes_descriptors_and_reads_markdown( provider_cls: type[_SyncSkillProvider] | type[_AsyncSkillProvider], ) -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) provider = provider_cls( skills=[ _descriptor("alpha", "Alpha skill"), SkillProviderDescriptor( name="beta", description="Beta skill", argument_hint="ARG", disable_model_invocation=True, user_invocable=False, ), ], markdown={"alpha": "# Alpha"}, ) session = CopilotSession("session-1", Mock()) session._register_skill_provider(provider) client._sessions["session-1"] = session listed = await client._handle_skill_provider_list({"sessionId": "session-1"}) read = await client._handle_skill_provider_read({"sessionId": "session-1", "name": "alpha"}) assert listed == { "skills": [ {"description": "Alpha skill", "name": "alpha"}, { "argumentHint": "ARG", "description": "Beta skill", "disableModelInvocation": True, "name": "beta", "userInvocable": False, }, ] } assert read == {"markdown": "# Alpha"} @pytest.mark.asyncio async def test_skill_provider_null_list_becomes_empty_catalog() -> None: session = CopilotSession("session-1", Mock()) session._register_skill_provider(_SyncSkillProvider(skills=None)) assert await session._handle_skill_provider_list() == {"skills": []} @pytest.mark.asyncio async def test_skill_provider_read_none_returns_null_markdown() -> None: session = CopilotSession("session-1", Mock()) session._register_skill_provider(_SyncSkillProvider(markdown={"missing": None})) assert await session._handle_skill_provider_read("missing") == {"markdown": None} @pytest.mark.asyncio async def test_skill_provider_failures_are_generic_and_do_not_leak_provider_errors() -> None: session = CopilotSession("session-1", Mock()) session._register_skill_provider(_FailingSkillProvider()) with pytest.raises(JsonRpcError) as list_error: await session._handle_skill_provider_list() with pytest.raises(JsonRpcError) as read_error: await session._handle_skill_provider_read("secret") assert list_error.value.code == -32603 assert list_error.value.message == "Skill provider listSkills failed" assert list_error.value.data is None assert "secret-list-token" not in str(list_error.value) assert read_error.value.code == -32603 assert read_error.value.message == "Skill provider readSkill failed" assert read_error.value.data is None assert "secret-read-token" not in str(read_error.value) @pytest.mark.asyncio async def test_skill_provider_failures_are_logged_with_session_id( caplog: pytest.LogCaptureFixture, ) -> None: session = CopilotSession("session-1", Mock()) session._register_skill_provider(_FailingSkillProvider()) with caplog.at_level("WARNING", logger="copilot.session"): with pytest.raises(JsonRpcError): await session._handle_skill_provider_list() with pytest.raises(JsonRpcError): await session._handle_skill_provider_read("secret") records = [r for r in caplog.records if r.getMessage().startswith("Skill provider")] assert [r.getMessage() for r in records] == [ "Skill provider listSkills failed", "Skill provider readSkill failed", ] assert all(getattr(r, "session_id", None) == "session-1" for r in records) assert [str(r.exc_info[1]) for r in records if r.exc_info] == [ "secret-list-token", "secret-read-token", ] class _SelfCancellingSkillProvider: async def list_skills(self) -> list[SkillProviderDescriptor]: raise asyncio.CancelledError async def read_skill(self, name: str) -> str: raise asyncio.CancelledError @pytest.mark.asyncio async def test_skill_provider_cancelled_error_raised_by_provider_is_a_failure() -> None: session = CopilotSession("session-1", Mock()) session._register_skill_provider(_SelfCancellingSkillProvider()) with pytest.raises(JsonRpcError) as list_error: await session._handle_skill_provider_list() with pytest.raises(JsonRpcError) as read_error: await session._handle_skill_provider_read("cancelled") assert list_error.value.message == "Skill provider listSkills failed" assert read_error.value.message == "Skill provider readSkill failed" class _BlockingSkillProvider: def __init__(self) -> None: self.entered = asyncio.Event() async def list_skills(self) -> list[SkillProviderDescriptor]: self.entered.set() await asyncio.Event().wait() return [] async def read_skill(self, name: str) -> str: return "" @pytest.mark.asyncio async def test_skill_provider_dispatch_cancellation_propagates() -> None: session = CopilotSession("session-1", Mock()) provider = _BlockingSkillProvider() session._register_skill_provider(provider) task = asyncio.create_task(session._handle_skill_provider_list()) await provider.entered.wait() task.cancel() with pytest.raises(asyncio.CancelledError): await task class _JsonRpcErrorSkillProvider: def list_skills(self) -> list[SkillProviderDescriptor]: raise JsonRpcError(-32000, "secret-list-token", data={"code": "provider_secret"}) def read_skill(self, name: str) -> str: raise JsonRpcError(-32000, "secret-read-token", data={"code": "provider_secret"}) @pytest.mark.asyncio async def test_skill_provider_json_rpc_errors_from_provider_are_not_forwarded() -> None: session = CopilotSession("session-1", Mock()) session._register_skill_provider(_JsonRpcErrorSkillProvider()) with pytest.raises(JsonRpcError) as list_error: await session._handle_skill_provider_list() with pytest.raises(JsonRpcError) as read_error: await session._handle_skill_provider_read("secret") assert (list_error.value.code, list_error.value.message, list_error.value.data) == ( -32603, "Skill provider listSkills failed", None, ) assert (read_error.value.code, read_error.value.message, read_error.value.data) == ( -32603, "Skill provider readSkill failed", None, ) @pytest.mark.asyncio async def test_skill_provider_unknown_session_and_no_provider_return_generic_errors() -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) with pytest.raises(JsonRpcError) as unknown_error: await client._handle_skill_provider_list({"sessionId": "missing-session"}) assert (unknown_error.value.code, unknown_error.value.message, unknown_error.value.data) == ( -32603, "Session not found: missing-session", None, ) client._sessions["session-1"] = CopilotSession("session-1", Mock()) with pytest.raises(JsonRpcError) as no_provider_error: await client._handle_skill_provider_read({"sessionId": "session-1", "name": "demo"}) assert ( no_provider_error.value.code, no_provider_error.value.message, no_provider_error.value.data, ) == ( -32603, "No skill provider for session: session-1", None, ) @pytest.mark.parametrize("method", ["list", "read"]) @pytest.mark.asyncio async def test_skill_provider_invalid_params_are_json_rpc_invalid_params(method: str) -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) params = {"sessionId": "session-1"} if method == "read" else {} handler = ( client._handle_skill_provider_read if method == "read" else client._handle_skill_provider_list ) with pytest.raises(JsonRpcError) as exc_info: await handler(params) assert exc_info.value.code == -32602 assert exc_info.value.message == "Invalid params" assert exc_info.value.data is None @pytest.mark.asyncio async def test_skill_provider_read_none_is_written_to_json_rpc_result() -> None: sdk_client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) process = _WriteOnlyProcess() rpc_client = JsonRpcClient(process) sdk_client._client = rpc_client sdk_client._register_session_scoped_request_handlers() session = CopilotSession("session-1", rpc_client) session._register_skill_provider(_SyncSkillProvider(markdown={"missing": None})) sdk_client._sessions["session-1"] = session await rpc_client._dispatch_request( { "jsonrpc": "2.0", "id": 1, "method": "skillProvider.read", "params": {"sessionId": "session-1", "name": "missing"}, }, rpc_client.request_handlers["skillProvider.read"], None, ) response = _decode_jsonrpc_output(process) assert response == { "jsonrpc": "2.0", "id": 1, "result": {"markdown": None}, } @pytest.mark.parametrize("method", ["list", "read"]) @pytest.mark.asyncio async def test_skill_provider_task_is_cancelled_by_runtime_cancel_request( method: str, caplog: pytest.LogCaptureFixture ) -> None: entered = asyncio.Event() cancelled = asyncio.Event() class BlockingProvider: async def _block(self) -> Any: entered.set() try: await asyncio.Event().wait() except asyncio.CancelledError: cancelled.set() raise async def list_skills(self) -> list[SkillProviderDescriptor]: return await self._block() async def read_skill(self, name: str) -> str | None: return await self._block() sdk_client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) process = _WriteOnlyProcess() rpc_client = JsonRpcClient(process) rpc_client._loop = asyncio.get_running_loop() sdk_client._client = rpc_client sdk_client._register_session_scoped_request_handlers() session = CopilotSession("session-1", rpc_client) session._register_skill_provider(BlockingProvider()) sdk_client._sessions["session-1"] = session params = {"sessionId": "session-1", "name": "demo"} with caplog.at_level("WARNING", logger="copilot.session"): rpc_client._handle_request( {"jsonrpc": "2.0", "id": 7, "method": f"skillProvider.{method}", "params": params} ) await asyncio.wait_for(entered.wait(), timeout=5) rpc_client._handle_cancel_request({"id": 7}) await asyncio.wait_for(cancelled.wait(), timeout=5) # The response is written before dispatch cleanup retires the cancellation entry. async with asyncio.timeout(5): while not process.stdin.getvalue() or 7 in rpc_client._incoming_request_cancellations: await asyncio.sleep(0.01) assert _decode_jsonrpc_output(process) == { "jsonrpc": "2.0", "id": 7, "error": {"code": -32800, "message": "Request cancelled", "data": None}, } assert not [r for r in caplog.records if r.getMessage().startswith("Skill provider")] assert rpc_client._incoming_request_cancellations == {} @pytest.mark.parametrize("teardown", ["disconnect", "mark", "force_stop", "delete", "close"]) @pytest.mark.asyncio async def test_skill_provider_is_cleared_on_teardown_paths(teardown: str) -> None: provider = _SyncSkillProvider(skills=[_descriptor("demo")]) client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) session = CopilotSession("session-1", Mock()) session._register_skill_provider(provider) if teardown == "disconnect": session._client.request = AsyncMock(return_value={"success": True}) await session.disconnect() elif teardown == "mark": session._mark_disconnected() elif teardown == "force_stop": client._sessions["session-1"] = session client._client = Mock() client._client.stop = AsyncMock() await client.force_stop() elif teardown == "delete": client._sessions["session-1"] = session client._client = Mock() client._client.request = AsyncMock(return_value={"success": True}) await client.delete_session("session-1") else: client._sessions["session-1"] = session client._client = Mock(_loop=asyncio.get_running_loop()) client._handle_connection_close() await asyncio.sleep(0) await _assert_no_provider_error(session, provider) @pytest.mark.parametrize("operation", ["create", "resume"]) @pytest.mark.asyncio async def test_skill_provider_is_cleared_after_failed_create_or_resume(operation: str) -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) provider = _SyncSkillProvider(skills=[_descriptor("demo")]) captured_session: CopilotSession | None = None async def request(method: str, params: dict, _kwargs: dict[str, Any]) -> dict[str, Any]: nonlocal captured_session if method == f"session.{operation}": captured_session = client._sessions[params["sessionId"]] raise RuntimeError("admission failed") return {} _install_fake_client(client, request) with pytest.raises(RuntimeError, match="admission failed"): if operation == "create": await client.create_session(session_id="failed-session", skill_provider=provider) else: await client.resume_session("failed-session", skill_provider=provider) assert "failed-session" not in client._sessions assert captured_session is not None await _assert_no_provider_error(captured_session, provider) @pytest.mark.asyncio async def test_failed_resume_restores_the_resident_skill_provider() -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) resident = _SyncSkillProvider(skills=[_descriptor("demo")]) replacement = _SyncSkillProvider(skills=[_descriptor("demo")]) async def request(method: str, params: dict, kwargs: dict[str, Any]) -> dict[str, Any]: if method == "session.resume": raise RuntimeError("resume failed") return await _default_request(method, params, kwargs) _install_fake_client(client, request) session = await client.create_session(session_id="session-1", skill_provider=resident) with pytest.raises(RuntimeError, match="resume failed"): await client.resume_session("session-1", skill_provider=replacement) assert client._sessions["session-1"] is session assert await client._handle_skill_provider_list({"sessionId": "session-1"}) == { "skills": [{"name": "demo", "description": "Test skill"}] } assert (resident.list_calls, replacement.list_calls) == (1, 0) @pytest.mark.asyncio async def test_resume_rebinds_or_unbinds_skill_provider() -> None: client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) provider_a = _SyncSkillProvider( skills=[_descriptor("demo")], markdown={"demo": "from A"}, ) provider_b = _SyncSkillProvider( skills=[_descriptor("demo")], markdown={"demo": "from B"}, ) _install_fake_client(client, _default_request) session = CopilotSession("session-1", Mock()) session._register_skill_provider(provider_a) client._sessions["session-1"] = session await client.resume_session("session-1", skill_provider=provider_b) assert await client._handle_skill_provider_read({"sessionId": "session-1", "name": "demo"}) == { "markdown": "from B" } assert provider_a.read_calls == [] await client.resume_session("session-1") with pytest.raises(JsonRpcError) as exc_info: await client._handle_skill_provider_list({"sessionId": "session-1"}) assert (exc_info.value.code, exc_info.value.message, exc_info.value.data) == ( -32603, "No skill provider for session: session-1", None, ) assert provider_b.list_calls == 0