copilot-sdk/python/test_subagent_hooks.py at main · 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
# Copyright (c) Microsoft Corporation. All rights reserved.

from datetime import UTC, datetime

import pytest

from copilot import (
    SessionHooks,
    SubagentStartHookInput,
    SubagentStartHookOutput,
    SubagentStopHookInput,
    SubagentStopHookOutput,
)
from copilot.session import CopilotSession


@pytest.mark.asyncio
async def test_subagent_start_dispatches_normalized_input_and_context():
    captured: list[tuple[SubagentStartHookInput, dict[str, str]]] = []

    async def on_start(
        input_data: SubagentStartHookInput, invocation: dict[str, str]
    ) -> SubagentStartHookOutput:
        captured.append((input_data, invocation))
        return {"additionalContext": "Report the entire file"}

    hooks: SessionHooks = {"on_subagent_start": on_start}
    session = CopilotSession("parent", None)
    session._register_hooks(hooks)

    output = await session._handle_hooks_invoke(
        "subagentStart",
        {
            "sessionId": "parent",
            "timestamp": 1_700_000_000_000,
            "cwd": "work",
            "transcriptPath": "transcript.jsonl",
            "agentName": "explore",
            "agentDisplayName": "Explorer",
            "agentDescription": "Reads files",
        },
    )

    assert output == {"additionalContext": "Report the entire file"}
    assert captured == [
        (
            {
                "sessionId": "parent",
                "timestamp": datetime.fromtimestamp(1_700_000_000, tz=UTC),
                "workingDirectory": "work",
                "transcriptPath": "transcript.jsonl",
                "agentName": "explore",
                "agentDisplayName": "Explorer",
                "agentDescription": "Reads files",
            },
            {"session_id": "parent"},
        )
    ]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "hook_output",
    [
        {"decision": "block", "reason": "Read the rest of the file"},
        {"modifiedResponse": "Verified: Hello from subagent test!"},
        {"decision": "allow", "modifiedResponse": "Verified: Hello from subagent test!"},
    ],
)
async def test_subagent_stop_dispatches_normalized_input_and_verdict(
    hook_output: SubagentStopHookOutput,
):
    captured: list[tuple[SubagentStopHookInput, dict[str, str]]] = []

    def on_stop(
        input_data: SubagentStopHookInput, invocation: dict[str, str]
    ) -> SubagentStopHookOutput:
        captured.append((input_data, invocation))
        return hook_output

    hooks: SessionHooks = {"on_subagent_stop": on_stop}
    session = CopilotSession("parent", None)
    session._register_hooks(hooks)

    output = await session._handle_hooks_invoke(
        "subagentStop",
        {
            "sessionId": "parent",
            "timestamp": 1_700_000_000_000,
            "cwd": "work",
            "transcriptPath": "transcript.jsonl",
            "agentName": "explore",
            "agentType": "explore",
            "stopReason": "end_turn",
            "response": "Hello from subagent test!",
        },
    )

    assert output == hook_output
    assert captured == [
        (
            {
                "sessionId": "parent",
                "timestamp": datetime.fromtimestamp(1_700_000_000, tz=UTC),
                "workingDirectory": "work",
                "transcriptPath": "transcript.jsonl",
                "agentName": "explore",
                "agentType": "explore",
                "stopReason": "end_turn",
                "response": "Hello from subagent test!",
            },
            {"session_id": "parent"},
        )
    ]


@pytest.mark.asyncio
async def test_subagent_stop_without_handler_returns_none():
    session = CopilotSession("parent", None)
    session._register_hooks({"on_subagent_start": lambda _input, _invocation: None})

    assert await session._handle_hooks_invoke("subagentStop", {}) is None


@pytest.mark.asyncio
async def test_subagent_start_without_handler_returns_none():
    session = CopilotSession("parent", None)
    session._register_hooks({"on_subagent_stop": lambda _input, _invocation: None})

    assert await session._handle_hooks_invoke("subagentStart", {}) is None