{{ message }}
Repository navigation
Expand file tree
/
Copy pathtest_host.py
More file actions
492 lines (432 loc) · 17.6 KB
/
Copy pathtest_host.py
File metadata and controls
492 lines (432 loc) · 17.6 KB
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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
# Copyright (c) Microsoft Corporation. All rights reserved.
import asyncio
from unittest.mock import AsyncMock, Mock
import pytest
from copilot import AhpHostOptions, CopilotClient, RuntimeConnection
from copilot._jsonrpc import JsonRpcClient
from copilot.generated.rpc import ServerRpc
from copilot.rpc import HostGitHubEnvironmentOptions, HostLocalServerOptions
def client_fixture():
client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234"))
async def request(method, params, **_):
if method == "host.start":
return {"hostId": params["hostId"], "url": "ws://127.0.0.1:12345", "token": "secret"}
if method == "host.publishSession":
return {
"sessionId": params["sessionId"],
"sessionUri": f"copilot:/{params['sessionId']}",
}
if method == "host.listSessions":
return {"sessions": []}
if method in ("session.create", "session.resume"):
return {"sessionId": params["sessionId"]}
return {}
rpc = Mock(request=AsyncMock(side_effect=request))
client._client = rpc
client._rpc = ServerRpc(rpc)
client._state = "connected"
return client, rpc
def handoff(host, *, resume=False, **config):
return {
"hostId": host.host_id,
"handoffId": "participation",
"resume": resume,
"config": {"sessionId": "session", "workingDirectory": "/workspace", **config},
}
def release(client, host):
client._ahp_hosts.notification(
"host.sessionReleased", {"hostId": host.host_id, "handoffId": "participation"}
)
@pytest.mark.parametrize("options", [None, AhpHostOptions()])
async def test_requires_explicit_transport(options):
client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234"))
client.start = AsyncMock()
with pytest.raises(ValueError, match="At least one"):
await client.start_ahp_host(options)
client.start.assert_not_called()
@pytest.mark.parametrize("local", [False, True])
async def test_github_hosting_with_optional_local_transport(local):
client, rpc = client_fixture()
async def start(method, params):
assert method == "host.start"
return {
"hostId": params["hostId"],
"environmentId": "environment-1",
**({"url": "ws://127.0.0.1:12345", "token": "local-token"} if local else {}),
}
rpc.request.side_effect = start
host = await client.start_ahp_host(
AhpHostOptions(
compute_id="stable-installation-id",
github_environment=HostGitHubEnvironmentOptions(
name="Application", compute_id="stable-installation-id"
),
local_server=HostLocalServerOptions() if local else None,
)
)
assert rpc.request.call_args.args[1] == {
"hostId": host.host_id,
"computeId": "stable-installation-id",
"githubEnvironment": {"name": "Application", "computeId": "stable-installation-id"},
**({"localServer": {}} if local else {}),
}
assert host.environment_id == "environment-1"
assert host.url == ("ws://127.0.0.1:12345" if local else None)
assert host.token == ("local-token" if local else None)
assert host.pid is None
@pytest.mark.parametrize("compute_id", [None, "", "stable-local-compute"])
async def test_local_compute_identity_is_forwarded_without_sdk_defaults(compute_id):
client, rpc = client_fixture()
host = await client.start_ahp_host(
AhpHostOptions(local_server=HostLocalServerOptions(), compute_id=compute_id)
)
assert rpc.request.call_args.args[1] == {
"hostId": host.host_id,
"localServer": {},
**({"computeId": compute_id} if compute_id is not None else {}),
}
async def test_handle_uses_original_transport_and_forwards_every_disposal():
client, rpc = client_fixture()
exit_received = asyncio.Event()
exited = Mock(side_effect=lambda _: exit_received.set())
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(port=0, require_connection_token=False),
on_exit=exited,
)
)
assert host.pid is None
params = rpc.request.call_args_list[0].args[1]
assert params == {
"hostId": host.host_id,
"localServer": {"port": 0, "requireConnectionToken": False},
}
replacement = Mock(request=AsyncMock())
client._client = replacement
published = await host.publish_session("resident")
assert published.session_id == "resident"
sessions = await host.list_sessions()
assert sessions.sessions == []
assert rpc.request.call_args_list[-1].args[:2] == (
"host.listSessions",
{"hostId": host.host_id},
)
assert await asyncio.gather(host.dispose(), host.dispose()) == [None, None]
assert await host.dispose() is None
assert sum(call.args[0] == "host.dispose" for call in rpc.request.call_args_list) == 3
replacement.request.assert_not_called()
for _ in range(2):
client._ahp_hosts.notification(
"host.exited", {"hostId": host.host_id, "reason": "disposed"}
)
client._ahp_hosts.disconnect()
await asyncio.wait_for(exit_received.wait(), timeout=1)
exited.assert_called_once()
assert exited.call_args.args[0].reason.value == "disposed"
async def test_unpublication_forwards_results_and_errors_over_original_transport():
client, rpc = client_fixture()
host = await client.start_ahp_host(AhpHostOptions(local_server=HostLocalServerOptions()))
replacement = Mock(request=AsyncMock())
client._client = replacement
rpc.request.reset_mock()
rpc.request.side_effect = [
{"unpublished": True},
{"unpublished": False},
RuntimeError("nonapplication publication"),
]
assert (await host.unpublish_session("dormant")).unpublished is True
assert (await host.unpublish_session("dormant")).unpublished is False
with pytest.raises(RuntimeError, match="nonapplication publication"):
await host.unpublish_session("dormant")
assert len(rpc.request.call_args_list) == 3
for call in rpc.request.call_args_list:
assert call.args == (
"host.unpublishSession",
{"hostId": host.host_id, "sessionId": "dormant"},
)
replacement.request.assert_not_called()
@pytest.mark.parametrize("resume", [False, True])
async def test_factories_preserve_configuration_and_release_exact_original_once(resume):
client, _ = client_fixture()
release_received = asyncio.Event()
released = Mock(side_effect=lambda _: release_received.set())
originals = []
cancellations = []
async def create(request):
cancellations.append(request.cancellation_event)
original = await client.create_session(**request.config)
originals.append(original)
return original
async def restore(request):
cancellations.append(request.cancellation_event)
original = await client.resume_session(request.session_id, **request.config)
originals.append(original)
return original
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(),
create_session=create,
resume_session=restore,
on_session_released=released,
)
)
config = (
{"continuePendingWork": False, "suppressResumeEvent": True}
if resume
else {"streaming": True}
)
result = await client._ahp_hosts.materialize(
handoff(
host,
resume=resume,
model="test-model",
configDir="/config",
mcpOAuthTokenStorage="in-memory",
gitHubToken="test-auth",
additionalDirectories=[],
enableMcpApps=True,
enableExperimentalMode=False,
infiniteSessions={"enabled": False, "backgroundCompactionThreshold": 0.7},
featureFlags={"Arbitrary.MixedCase": True},
**config,
)
)
assert result == {"sessionId": "session"}
client._ahp_hosts.notification(
"host.sessionReleased", {"hostId": "wrong-host", "handoffId": "participation"}
)
assert not cancellations[0].is_set()
release(client, host)
release(client, host)
client._ahp_hosts.disconnect()
await asyncio.wait_for(release_received.wait(), timeout=1)
assert cancellations[0].is_set()
released.assert_called_once_with(originals[0])
assert client._get_session("session") is originals[0]
async def test_resume_can_return_a_retained_original_without_reconfiguration():
client, rpc = client_fixture()
original = await client.create_session(session_id="session", working_directory="/workspace")
factory = AsyncMock(return_value=original)
host = await client.start_ahp_host(
AhpHostOptions(local_server=HostLocalServerOptions(), resume_session=factory)
)
rpc.request.reset_mock()
assert await client._ahp_hosts.materialize(handoff(host, resume=True)) == {
"sessionId": "session"
}
rpc.request.assert_not_called()
@pytest.mark.parametrize("wrong", ["settings", "identity", "owner"])
async def test_invalid_factory_result_is_rejected_and_released_without_destruction(wrong):
client, rpc = client_fixture()
other, _ = client_fixture()
release_received = asyncio.Event()
released = Mock(side_effect=lambda _: release_received.set())
originals = []
async def create(request):
config = dict(request.config)
if wrong == "settings":
config["working_directory"] = "/wrong"
if wrong == "identity":
config["session_id"] = "wrong"
original = await (other if wrong == "owner" else client).create_session(**config)
originals.append(original)
return original
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(),
create_session=create,
on_session_released=released,
)
)
with pytest.raises(ValueError, match="AHP callback must"):
await client._ahp_hosts.materialize(handoff(host))
await asyncio.wait_for(release_received.wait(), timeout=1)
released.assert_called_once_with(originals[0])
assert not any(
call.args[0] in ("session.destroy", "session.detach") for call in rpc.request.call_args_list
)
async def test_cancelled_handoff_unblocks_before_late_factory_and_releases_it_once():
client, _ = client_fixture()
started, finish, delivered = asyncio.Event(), asyncio.Event(), asyncio.Event()
requests, originals, released = [], [], []
async def create(request):
requests.append(request)
originals.append(await client.create_session(**request.config))
started.set()
await finish.wait()
return originals[0]
def on_released(session):
released.append(session)
delivered.set()
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(),
create_session=create,
on_session_released=on_released,
)
)
task = asyncio.create_task(client._ahp_hosts.materialize(handoff(host)))
await asyncio.wait_for(started.wait(), 1)
with pytest.raises(ValueError, match="handoff already exists"):
await client._ahp_hosts.materialize(handoff(host))
release(client, host)
with pytest.raises(RuntimeError, match="handoff ended"):
await asyncio.wait_for(task, 1)
assert requests[0].cancellation_event.is_set()
assert released == []
finish.set()
await asyncio.wait_for(delivered.wait(), 1)
client._ahp_hosts.disconnect()
assert released == originals
assert client._get_session("session") is originals[0]
@pytest.mark.parametrize("resume", [False, True])
async def test_cancelled_factory_returns_rpc_error_and_releases_handoff(resume):
client, _ = client_fixture()
requests = []
async def factory(request):
requests.append(request)
work = asyncio.create_task(asyncio.sleep(0))
work.cancel()
await work
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(),
create_session=factory,
resume_session=factory,
)
)
rpc = JsonRpcClient(Mock())
rpc._send_message = AsyncMock()
await asyncio.wait_for(
rpc._dispatch_request(
{
"jsonrpc": "2.0",
"id": "materialize",
"method": "host.materializeSession",
"params": handoff(host, resume=resume),
},
client._ahp_hosts.materialize,
cancellation_id=None,
),
1,
)
rpc._send_message.assert_awaited_once()
response = rpc._send_message.call_args.args[0]
assert response["id"] == "materialize"
assert response["error"]["code"] == -32603
assert response["error"]["message"] == "AHP session factory was cancelled"
assert requests[0].cancellation_event.is_set()
assert not client._ahp_hosts._handoffs
async def test_start_failure_removes_factory_and_exit_callback():
client, rpc = client_fixture()
exited = Mock()
rpc.request.side_effect = RuntimeError("bind failed")
with pytest.raises(RuntimeError, match="bind failed"):
await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(), create_session=AsyncMock(), on_exit=exited
)
)
client._ahp_hosts.disconnect()
await asyncio.sleep(0)
exited.assert_not_called()
@pytest.mark.parametrize("startup_fails", [False, True])
@pytest.mark.parametrize("cancel_mode", ["cancel", "timeout"])
async def test_cancelled_start_settles_before_cleanup_on_original_connection(
startup_fails, cancel_mode
):
client, rpc = client_fixture()
accepted, finish = asyncio.Event(), asyncio.Event()
disposed, finish_dispose = asyncio.Event(), asyncio.Event()
host_ids = []
async def request(method, params, **_):
if method == "host.start":
host_ids.append(params["hostId"])
accepted.set()
await finish.wait()
if startup_fails:
raise RuntimeError("bind failed")
return {"hostId": params["hostId"], "url": "ws://127.0.0.1:12345"}
assert method == "host.dispose"
assert params == {"hostId": host_ids[0]}
disposed.set()
await finish_dispose.wait()
return {}
rpc.request.side_effect = request
task = asyncio.create_task(
client.start_ahp_host(
AhpHostOptions(local_server=HostLocalServerOptions(), create_session=AsyncMock())
)
)
await asyncio.wait_for(accepted.wait(), 1)
if cancel_mode == "cancel":
task.cancel()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, 1)
else:
with pytest.raises(TimeoutError):
await asyncio.wait_for(task, 0)
assert host_ids[0] in client._ahp_hosts._hosts
replacement = Mock(request=AsyncMock())
client._client = replacement
finish.set()
if not startup_fails:
await asyncio.wait_for(disposed.wait(), 1)
assert host_ids[0] in client._ahp_hosts._hosts
finish_dispose.set()
await asyncio.wait_for(asyncio.gather(*client._ahp_hosts._tasks), 1)
assert not client._ahp_hosts._hosts
assert disposed.is_set() is not startup_fails
replacement.request.assert_not_called()
@pytest.mark.parametrize("resume", [False, True])
async def test_failed_factory_does_not_leave_capture_on_retained_original(resume):
client, rpc = client_fixture()
originals = []
async def create(request):
originals.append(await client.create_session(**request.config))
raise RuntimeError("application setup failed")
async def restore(request):
if originals:
return originals[0]
originals.append(await client.resume_session(request.session_id, **request.config))
raise RuntimeError("application setup failed")
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(), create_session=create, resume_session=restore
)
)
with pytest.raises(RuntimeError, match="application setup failed"):
await client._ahp_hosts.materialize(handoff(host, resume=resume))
rpc.request.reset_mock()
assert await client._ahp_hosts.materialize(
handoff(host, resume=True, continuePendingWork=False, suppressResumeEvent=True)
) == {"sessionId": "session"}
rpc.request.assert_not_called()
assert client._get_session("session") is originals[0]
async def test_factory_error_propagates_and_release_callback_errors_are_logged(caplog):
client, _ = client_fixture()
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(),
create_session=AsyncMock(side_effect=ValueError("factory failed")),
)
)
with pytest.raises(ValueError, match="factory failed"):
await client._ahp_hosts.materialize(handoff(host))
release_called = asyncio.Event()
def fail_release(_):
release_called.set()
raise ValueError("release failed")
async def create(request):
return await client.create_session(**request.config)
host = await client.start_ahp_host(
AhpHostOptions(
local_server=HostLocalServerOptions(),
create_session=create,
on_session_released=fail_release,
)
)
await client._ahp_hosts.materialize(handoff(host))
release(client, host)
await asyncio.wait_for(release_called.wait(), timeout=1)
assert "AHP session release callback failed" in caplog.text
You can’t perform that action at this time.