copilot-sdk/python/test_ffi_runtime_host.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
131
import asyncio
import threading
import time
from unittest.mock import patch

import pytest

from copilot import CopilotClient, RuntimeConnection
from copilot._ffi_runtime_host import FfiRuntimeHost


class _TestLibrary:
    def __init__(self) -> None:
        self.allow_close = False
        self.close_calls = 0
        self.shutdown_calls = 0
        self.shutdown = threading.Event()
        self.ui_calls: list[tuple[int, int | bool]] = []
        self.ui_result = True

    def connection_set_ui_thread_handler(self, connection_id: int, enabled: bool) -> bool:
        self.ui_calls.append((connection_id, enabled))
        return self.ui_result

    def connection_run_ui_action(self, connection_id: int, action_id: int) -> bool:
        self.ui_calls.append((connection_id, action_id))
        return self.ui_result

    def connection_close(self, _connection_id: int) -> bool:
        self.close_calls += 1
        return self.allow_close

    def host_shutdown(self, _server_id: int) -> bool:
        self.shutdown_calls += 1
        self.shutdown.set()
        return True


def test_dispose_retains_callback_until_connection_close_succeeds():
    library = _TestLibrary()
    with (
        patch("copilot._ffi_runtime_host._load_library", return_value=library),
        patch("copilot._ffi_runtime_host._CLEANUP_RETRY_INTERVAL_SECONDS", 0.01),
    ):
        host = FfiRuntimeHost("test-runtime", None)
        callback = object()
        host._server_id = 11
        host._connection_id = 21
        host._outbound_callback = callback

        host.dispose()

        assert host._outbound_callback is callback
        assert host._connection_id == 21
        assert library.close_calls == 1
        assert library.shutdown_calls == 0

        library.allow_close = True
        assert library.shutdown.wait(5), "Deferred native cleanup did not complete"

        assert host._outbound_callback is None
        assert host._connection_id == 0
        assert library.close_calls >= 2
        assert library.shutdown_calls == 1

        close_calls_after_cleanup = library.close_calls
        host.dispose()
        time.sleep(0.05)
        assert library.close_calls == close_calls_after_cleanup
        assert library.shutdown_calls == 1


def test_ui_actions_are_synchronous_connection_scoped_and_reject_after_disposal():
    library = _TestLibrary()
    library.allow_close = True
    with patch("copilot._ffi_runtime_host._load_library", return_value=library):
        host = FfiRuntimeHost("test-runtime", None)
    host._server_id = 11
    host._connection_id = 21
    try:
        host.configure_ui_thread(True)
        host.run_ui_action(42)
        assert library.ui_calls == [(21, True), (21, 42)]
        library.ui_result = False
        with pytest.raises(RuntimeError, match="retired, already executed"):
            host.run_ui_action(42)
    finally:
        host.dispose()
    with pytest.raises(RuntimeError, match="retired, already executed"):
        host.run_ui_action(43)
    assert library.ui_calls == [(21, True), (21, 42), (21, 42)]


@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["stop", "force_stop"])
async def test_client_shutdown_keeps_event_loop_responsive_until_native_cleanup_finishes(method):
    loop = asyncio.get_running_loop()
    loop_response = threading.Event()
    loop_responded_during_shutdown = False

    class BlockingLibrary(_TestLibrary):
        def host_shutdown(self, server_id: int) -> bool:
            nonlocal loop_responded_during_shutdown
            loop.call_soon_threadsafe(loop_response.set)
            loop_responded_during_shutdown = loop_response.wait(5)
            return super().host_shutdown(server_id)

    library = BlockingLibrary()
    library.allow_close = True
    with patch("copilot._ffi_runtime_host._load_library", return_value=library):
        host = FfiRuntimeHost("test-runtime", None)
    host._server_id = 11
    host._connection_id = 21
    host._outbound_callback = object()

    client = CopilotClient(connection=RuntimeConnection.for_stdio(path="copilot"))
    client._ffi_host = host
    client._process = host.process
    client._state = "connected"

    await getattr(client, method)()

    assert loop_responded_during_shutdown, "Native shutdown blocked the asyncio event loop"
    assert library.shutdown.is_set(), "Client shutdown returned before native cleanup"
    assert library.shutdown_calls == 1
    assert host._server_id == 0
    assert host._connection_id == 0
    assert host._outbound_callback is None
    assert client._ffi_host is None
    assert client._process is None
    assert client._state == "disconnected"