copilot-sdk/python/e2e/test_mcp_and_agents.py at exp-core · github/copilot-sdk · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
"""
Tests for MCP servers and custom agents functionality
"""

import pytest

from copilot import CustomAgentConfig, MCPServerConfig

from .testharness import E2ETestContext, get_final_assistant_message

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


class TestMCPServers:
    async def test_should_accept_mcp_server_configuration_on_session_create(
        self, ctx: E2ETestContext
    ):
        """Test that MCP server configuration is accepted on session create"""
        mcp_servers: dict[str, MCPServerConfig] = {
            "test-server": {
                "type": "local",
                "command": "echo",
                "args": ["hello"],
                "tools": ["*"],
            }
        }

        session = await ctx.client.create_session({"mcp_servers": mcp_servers})

        assert session.session_id is not None

        # Simple interaction to verify session works
        message = await session.send_and_wait({"prompt": "What is 2+2?"})
        assert message is not None
        assert "4" in message.data.content

        await session.destroy()

    async def test_should_accept_mcp_server_configuration_on_session_resume(
        self, ctx: E2ETestContext
    ):
        """Test that MCP server configuration is accepted on session resume"""
        # Create a session first
        session1 = await ctx.client.create_session()
        session_id = session1.session_id
        await session1.send_and_wait({"prompt": "What is 1+1?"})

        # Resume with MCP servers
        mcp_servers: dict[str, MCPServerConfig] = {
            "test-server": {
                "type": "local",
                "command": "echo",
                "args": ["hello"],
                "tools": ["*"],
            }
        }

        session2 = await ctx.client.resume_session(session_id, {"mcp_servers": mcp_servers})

        assert session2.session_id == session_id

        message = await session2.send_and_wait({"prompt": "What is 3+3?"})
        assert message is not None
        assert "6" in message.data.content

        await session2.destroy()


class TestCustomAgents:
    async def test_should_accept_custom_agent_configuration_on_session_create(
        self, ctx: E2ETestContext
    ):
        """Test that custom agent configuration is accepted on session create"""
        custom_agents: list[CustomAgentConfig] = [
            {
                "name": "test-agent",
                "display_name": "Test Agent",
                "description": "A test agent for SDK testing",
                "prompt": "You are a helpful test agent.",
                "infer": True,
            }
        ]

        session = await ctx.client.create_session({"custom_agents": custom_agents})

        assert session.session_id is not None

        # Simple interaction to verify session works
        message = await session.send_and_wait({"prompt": "What is 5+5?"})
        assert message is not None
        assert "10" in message.data.content

        await session.destroy()

    async def test_should_accept_custom_agent_configuration_on_session_resume(
        self, ctx: E2ETestContext
    ):
        """Test that custom agent configuration is accepted on session resume"""
        # Create a session first
        session1 = await ctx.client.create_session()
        session_id = session1.session_id
        await session1.send_and_wait({"prompt": "What is 1+1?"})

        # Resume with custom agents
        custom_agents: list[CustomAgentConfig] = [
            {
                "name": "resume-agent",
                "display_name": "Resume Agent",
                "description": "An agent added on resume",
                "prompt": "You are a resume test agent.",
            }
        ]

        session2 = await ctx.client.resume_session(session_id, {"custom_agents": custom_agents})

        assert session2.session_id == session_id

        message = await session2.send_and_wait({"prompt": "What is 6+6?"})
        assert message is not None
        assert "12" in message.data.content

        await session2.destroy()


class TestCombinedConfiguration:
    async def test_should_accept_both_mcp_servers_and_custom_agents(self, ctx: E2ETestContext):
        """Test that both MCP servers and custom agents can be configured together"""
        mcp_servers: dict[str, MCPServerConfig] = {
            "shared-server": {
                "type": "local",
                "command": "echo",
                "args": ["shared"],
                "tools": ["*"],
            }
        }

        custom_agents: list[CustomAgentConfig] = [
            {
                "name": "combined-agent",
                "display_name": "Combined Agent",
                "description": "An agent using shared MCP servers",
                "prompt": "You are a combined test agent.",
            }
        ]

        session = await ctx.client.create_session(
            {"mcp_servers": mcp_servers, "custom_agents": custom_agents}
        )

        assert session.session_id is not None

        await session.send({"prompt": "What is 7+7?"})
        message = await get_final_assistant_message(session)
        assert "14" in message.data.content

        await session.destroy()