copilot-sdk/python/e2e/test_pre_mcp_tool_call_hook_e2e.py at roji-python-windows-in-process-flakiness · 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
"""
E2E tests for the preMcpToolCall hook, verifying meta manipulation scenarios:
setting meta, replacing meta, and removing meta.
"""

from __future__ import annotations

from datetime import datetime
from pathlib import Path

import pytest

from copilot.session import MCPServerConfig, PermissionHandler

from .testharness import E2ETestContext

TEST_MCP_META_ECHO_SERVER = str(
    (Path(__file__).parents[2] / "test" / "harness" / "test-mcp-meta-echo-server.mjs").resolve()
)
TEST_HARNESS_DIR = str((Path(__file__).parents[2] / "test" / "harness").resolve())

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


def meta_echo_mcp_config() -> dict[str, MCPServerConfig]:
    return {
        "meta-echo": {
            "command": "node",
            "args": [TEST_MCP_META_ECHO_SERVER],
            "working_directory": TEST_HARNESS_DIR,
            "tools": ["*"],
        }
    }


class TestPreMcpToolCallHook:
    async def test_should_set_meta_via_premcptoolcall_hook(self, ctx: E2ETestContext):
        inputs: list[dict] = []

        async def on_pre_mcp_tool_call(input_data, invocation):
            inputs.append(input_data)
            return {"metaToUse": {"injected": "by-hook", "source": "test"}}

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            mcp_servers=meta_echo_mcp_config(),
            hooks={"on_pre_mcp_tool_call": on_pre_mcp_tool_call},
        )
        try:
            response = await session.send_and_wait(
                "Use the meta-echo/echo_meta tool with value 'test-set'."
                " Reply with just the raw tool result."
            )
            assert response is not None
            assert "injected" in (response.data.content or "")
            assert "by-hook" in (response.data.content or "")

            assert inputs
            assert inputs[0].get("serverName") == "meta-echo"
            assert inputs[0].get("toolName") == "echo_meta"
            assert inputs[0].get("workingDirectory")
            assert isinstance(inputs[0].get("timestamp"), datetime)
        finally:
            await session.disconnect()

    async def test_should_replace_meta_via_premcptoolcall_hook(self, ctx: E2ETestContext):
        inputs: list[dict] = []

        async def on_pre_mcp_tool_call(input_data, invocation):
            inputs.append(input_data)
            return {"metaToUse": {"completely": "replaced"}}

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            mcp_servers=meta_echo_mcp_config(),
            hooks={"on_pre_mcp_tool_call": on_pre_mcp_tool_call},
        )
        try:
            response = await session.send_and_wait(
                "Use the meta-echo/echo_meta tool with value 'test-replace'."
                " Reply with just the raw tool result."
            )
            assert response is not None
            assert "completely" in (response.data.content or "")
            assert "replaced" in (response.data.content or "")

            assert inputs
            assert inputs[0].get("serverName") == "meta-echo"
            assert inputs[0].get("toolName") == "echo_meta"
        finally:
            await session.disconnect()

    async def test_should_remove_meta_via_premcptoolcall_hook(self, ctx: E2ETestContext):
        inputs: list[dict] = []

        async def on_pre_mcp_tool_call(input_data, invocation):
            inputs.append(input_data)
            return {"metaToUse": None}

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            mcp_servers=meta_echo_mcp_config(),
            hooks={"on_pre_mcp_tool_call": on_pre_mcp_tool_call},
        )
        try:
            response = await session.send_and_wait(
                "Use the meta-echo/echo_meta tool with value 'test-remove'."
                " Reply with just the raw tool result."
            )
            assert response is not None
            assert '"meta":null' in (response.data.content or "") or '"meta": null' in (
                response.data.content or ""
            )
            assert "test-remove" in (response.data.content or "")

            assert inputs
            assert inputs[0].get("serverName") == "meta-echo"
            assert inputs[0].get("toolName") == "echo_meta"
        finally:
            await session.disconnect()