amkernel 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
amkernel/protocol.py ADDED
@@ -0,0 +1,1464 @@
1
+ """Versioned private protocol shared by the AMKernel host and Runner."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import base64
6
+ import binascii
7
+ import hashlib
8
+ import json
9
+ import math
10
+ import posixpath
11
+ import re
12
+ import struct
13
+ import unicodedata
14
+ from collections.abc import Callable, Mapping
15
+ from dataclasses import dataclass, field
16
+ from typing import Any, BinaryIO, Final, Literal, TypeAlias, cast
17
+
18
+ PROTOCOL_VERSION: Final = 2
19
+
20
+ # Shared byte, count, frame, result, resource, and image bounds from ADR-0006.
21
+ MAX_CONTROL_FRAME_BYTES: Final = 8_388_608
22
+ MAX_CODE_BYTES: Final = 1_048_576
23
+ MAX_RESULT_BYTES: Final = 1_048_576
24
+ MAX_ERROR_BYTES: Final = 65_536
25
+ MAX_PROTOCOL_ERROR_MESSAGE_BYTES: Final = 4_096
26
+ MAX_PATH_BYTES: Final = 16_384
27
+ MAX_PYTHON_VERSION_BYTES: Final = 128
28
+ MAX_ERROR_TYPE_BYTES: Final = 256
29
+ MAX_ERROR_MESSAGE_BYTES: Final = 16_384
30
+ MAX_IMPLEMENTATION_BYTES: Final = 256
31
+ MAX_PLATFORM_BYTES: Final = 256
32
+ MAX_PACKAGE_COUNT: Final = 64
33
+ MAX_PACKAGE_NAME_BYTES: Final = 128
34
+ MAX_PACKAGE_VERSION_BYTES: Final = 1_024
35
+ MAX_LOGICAL_NAME_CODEPOINTS: Final = 255
36
+ MAX_MIME_BYTES: Final = 255
37
+ MAX_OUTPUT_SEQUENCE: Final = 128
38
+ MAX_OUTPUT_FRAMES: Final = 4_288
39
+ MAX_ARTIFACT_TRANSACTIONS: Final = 64
40
+ MAX_JSON_DISPLAYS: Final = 64
41
+ MAX_IMAGE_TRANSACTIONS: Final = 16
42
+ MAX_ARTIFACT_CHUNKS: Final = 4_096
43
+ MAX_ARTIFACT_BYTES: Final = 268_435_456
44
+ MAX_JSON_DISPLAY_ITEM_BYTES: Final = 1_048_576
45
+ MAX_JSON_DISPLAY_BYTES: Final = 4_194_304
46
+ MAX_ARTIFACT_CHUNK_BYTES: Final = 65_536
47
+ MAX_ARTIFACT_CHUNK_ENCODED_BYTES: Final = 87_384
48
+ STREAM_TAIL_BYTES: Final = 262_144
49
+ MAX_STREAM_BYTES: Final = 268_435_456
50
+ MAX_STORED_EXPLICIT_ITEMS: Final = 64
51
+ MAX_STORED_STREAM_ITEMS: Final = 2
52
+ MAX_STORED_ITEMS: Final = 66
53
+ MAX_STRUCTURED_RESULT_BYTES: Final = 4_194_304
54
+ EXECUTION_BUNDLE_RESERVATION_BYTES: Final = 570_425_344
55
+ ARTIFACT_STORE_QUOTA_BYTES: Final = 2_147_483_648
56
+ ARTIFACT_RESOURCE_PART_BYTES: Final = 8_388_608
57
+ MAX_ARTIFACT_PARTS: Final = 32
58
+ MAX_EXECUTION_PARTS: Final = 128
59
+ MAX_RESOURCE_READS: Final = 1
60
+ MAX_INLINE_IMAGE_ITEM_BYTES: Final = 4_194_304
61
+ MAX_INLINE_IMAGE_BYTES: Final = 8_388_608
62
+ MAX_IMAGE_DIMENSION: Final = 32_768
63
+ MAX_IMAGE_FRAMES: Final = 256
64
+ MAX_IMAGE_PIXEL_FRAMES: Final = 40_000_000
65
+ MAX_LIVE_MESSAGE_BYTES: Final = 8_192
66
+ MAX_LIVE_UPDATES_PER_SECOND: Final = 10
67
+
68
+ REQUEST_FD_ARGUMENT: Final = "--request-fd"
69
+ RESPONSE_FD_ARGUMENT: Final = "--response-fd"
70
+ DIAGNOSTIC_FD_ARGUMENT: Final = "--diagnostic-fd"
71
+
72
+ SAFE_MIME_TYPES: Final = frozenset(
73
+ {
74
+ "text/plain; charset=utf-8",
75
+ "application/json",
76
+ "image/png",
77
+ "image/jpeg",
78
+ "image/gif",
79
+ "image/webp",
80
+ }
81
+ )
82
+ IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg", "image/gif", "image/webp"})
83
+ _CORRELATION_ID = re.compile(r"^[0-9a-f]{32}$")
84
+
85
+ _UUID_HEX = re.compile(r"^[0-9a-f]{12}4[0-9a-f]{3}[89ab][0-9a-f]{15}$")
86
+ _SHA256 = re.compile(r"^[0-9a-f]{64}$")
87
+ _PACKAGE_NAME = re.compile(r"^[A-Za-z0-9]+(?:[._-]+[A-Za-z0-9]+)*$")
88
+ _MIME_TOKEN = re.compile(r"^[!#$%&'*+.^_`|~0-9A-Za-z-]+$")
89
+ _LENGTH = struct.Struct(">I")
90
+ _JSON_DUMPS = json.dumps
91
+ _JSON_LOADS = json.loads
92
+ _ISFINITE = math.isfinite
93
+ _NORMPATH = posixpath.normpath
94
+
95
+ EnvironmentKind: TypeAlias = Literal["project_venv", "active_venv", "system"]
96
+ DoneOutcome: TypeAlias = Literal["completed", "failed", "interrupted"]
97
+ ArtifactKind: TypeAlias = Literal["artifact", "image"]
98
+ PackageError: TypeAlias = Literal["not installed", "metadata unavailable", "invalid metadata"]
99
+ ProtocolErrorCode: TypeAlias = Literal[
100
+ "version_mismatch",
101
+ "invalid_sequence",
102
+ "invalid_request",
103
+ "runtime_diagnostics_error",
104
+ "internal_protocol_error",
105
+ ]
106
+ JsonValue: TypeAlias = bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"] | None
107
+
108
+
109
+ class ProtocolViolation(ValueError):
110
+ """A malformed, oversized, or schema-invalid private protocol message."""
111
+
112
+ def __init__(
113
+ self,
114
+ message: str,
115
+ *,
116
+ code: ProtocolErrorCode = "invalid_request",
117
+ correlation_id: str | None = None,
118
+ ) -> None:
119
+ super().__init__(message)
120
+ self.code = code
121
+ self.correlation_id = correlation_id
122
+
123
+
124
+ @dataclass(frozen=True, slots=True)
125
+ class RuntimePackage:
126
+ """One requested distribution diagnostic."""
127
+
128
+ name: str
129
+ version: str | None = None
130
+ error: PackageError | None = None
131
+
132
+
133
+ @dataclass(frozen=True, slots=True)
134
+ class RuntimeDiagnostics:
135
+ """Interpreter facts reported by the Runner."""
136
+
137
+ executable: str
138
+ python_version: str
139
+ working_directory: str
140
+ environment_kind: EnvironmentKind
141
+ implementation: str = "CPython"
142
+ platform: str = "unknown"
143
+ packages: tuple[RuntimePackage, ...] = ()
144
+
145
+
146
+ @dataclass(frozen=True, slots=True)
147
+ class ErrorInfo:
148
+ """Bounded user-code error information."""
149
+
150
+ type: str
151
+ message: str
152
+ traceback: tuple[str, ...]
153
+
154
+
155
+ @dataclass(frozen=True, slots=True)
156
+ class InitializeRequest:
157
+ """Initialize one fresh Runner."""
158
+
159
+ correlation_id: str
160
+ launch_working_directory: str
161
+ environment_kind: EnvironmentKind
162
+
163
+
164
+ @dataclass(frozen=True, slots=True)
165
+ class ExecuteRequest:
166
+ """Execute one Python cell."""
167
+
168
+ correlation_id: str
169
+ code: str
170
+ packages: tuple[str, ...] = ()
171
+
172
+
173
+ @dataclass(frozen=True, slots=True)
174
+ class ShutdownRequest:
175
+ """Shut down an idle Runner."""
176
+
177
+ correlation_id: str
178
+
179
+
180
+ Request: TypeAlias = InitializeRequest | ExecuteRequest | ShutdownRequest
181
+
182
+
183
+ @dataclass(frozen=True, slots=True)
184
+ class ReadyResponse:
185
+ """A fresh Runner is initialized and idle."""
186
+
187
+ correlation_id: str
188
+ runtime: RuntimeDiagnostics
189
+ execution_count: int = 0
190
+
191
+
192
+ @dataclass(frozen=True, slots=True)
193
+ class StartedResponse:
194
+ """The Runner accepted one cell and committed its execution count."""
195
+
196
+ correlation_id: str
197
+ execution_count: int
198
+
199
+
200
+ @dataclass(frozen=True, slots=True)
201
+ class DoneResponse:
202
+ """The Runner returned to an idle request loop after a cell."""
203
+
204
+ correlation_id: str
205
+ outcome: DoneOutcome
206
+ execution_count: int
207
+ runtime: RuntimeDiagnostics
208
+ has_result: bool
209
+ result: JsonValue = None
210
+ error: ErrorInfo | None = None
211
+
212
+
213
+ @dataclass(frozen=True, slots=True)
214
+ class DisplayJsonResponse:
215
+ """One ordered strict-JSON display event."""
216
+
217
+ correlation_id: str
218
+ execution_count: int
219
+ sequence: int
220
+ value: JsonValue
221
+
222
+
223
+ @dataclass(frozen=True, slots=True)
224
+ class ArtifactBeginResponse:
225
+ """Begin one noninterleaved artifact transaction."""
226
+
227
+ correlation_id: str
228
+ execution_count: int
229
+ sequence: int
230
+ transaction_id: str
231
+ kind: ArtifactKind
232
+ name: str
233
+ declared_mime_type: str
234
+ size: int
235
+
236
+
237
+ @dataclass(frozen=True, slots=True)
238
+ class ArtifactChunkResponse:
239
+ """Carry one canonical-base64 artifact chunk as decoded bytes."""
240
+
241
+ correlation_id: str
242
+ execution_count: int
243
+ sequence: int
244
+ transaction_id: str
245
+ chunk_index: int
246
+ data: bytes
247
+
248
+
249
+ @dataclass(frozen=True, slots=True)
250
+ class ArtifactEndResponse:
251
+ """Commit one complete artifact transaction."""
252
+
253
+ correlation_id: str
254
+ execution_count: int
255
+ sequence: int
256
+ transaction_id: str
257
+ chunk_count: int
258
+ sha256: str
259
+
260
+
261
+ @dataclass(frozen=True, slots=True)
262
+ class ProtocolErrorResponse:
263
+ """A fatal correlated protocol error emitted before Runner exit."""
264
+
265
+ correlation_id: str
266
+ code: ProtocolErrorCode
267
+ message: str
268
+
269
+
270
+ Response: TypeAlias = (
271
+ ReadyResponse
272
+ | StartedResponse
273
+ | DisplayJsonResponse
274
+ | ArtifactBeginResponse
275
+ | ArtifactChunkResponse
276
+ | ArtifactEndResponse
277
+ | DoneResponse
278
+ | ProtocolErrorResponse
279
+ )
280
+
281
+
282
+ @dataclass(slots=True)
283
+ class _ArtifactTransaction:
284
+ transaction_id: str
285
+ sequence: int
286
+ size: int
287
+ chunk_count: int = 0
288
+ received_bytes: int = 0
289
+ digest: object = field(default_factory=hashlib.sha256)
290
+
291
+
292
+ @dataclass(slots=True)
293
+ class ExecutionResponseState:
294
+ """Validate one started execution's correlation, sequence, and transactions."""
295
+
296
+ correlation_id: str
297
+ execution_count: int
298
+ next_sequence: int = 1
299
+ output_frames: int = 0
300
+ artifact_transactions: int = 0
301
+ image_transactions: int = 0
302
+ json_displays: int = 0
303
+ artifact_chunks: int = 0
304
+ artifact_bytes: int = 0
305
+ json_bytes: int = 0
306
+ transaction_ids: set[str] = field(default_factory=set)
307
+ active_transaction: _ArtifactTransaction | None = None
308
+ terminal: bool = False
309
+
310
+ def accept(self, response: Response) -> None:
311
+ """Accept one decoded response or raise on any v2 transition violation."""
312
+ if self.terminal:
313
+ raise ProtocolViolation("response after terminal frame")
314
+ if response.correlation_id != self.correlation_id:
315
+ raise ProtocolViolation("uncorrelated execution response")
316
+ if isinstance(response, ProtocolErrorResponse):
317
+ self.terminal = True
318
+ return
319
+ if isinstance(response, DoneResponse):
320
+ self._require_execution_count(response.execution_count)
321
+ if self.active_transaction is not None:
322
+ raise ProtocolViolation("done during incomplete artifact transaction")
323
+ self.terminal = True
324
+ return
325
+ if not isinstance(
326
+ response,
327
+ (
328
+ DisplayJsonResponse,
329
+ ArtifactBeginResponse,
330
+ ArtifactChunkResponse,
331
+ ArtifactEndResponse,
332
+ ),
333
+ ):
334
+ raise ProtocolViolation("invalid response after started")
335
+
336
+ self.output_frames += 1
337
+ if self.output_frames > MAX_OUTPUT_FRAMES:
338
+ raise ProtocolViolation("output frame count exceeds maximum")
339
+ self._require_execution_count(response.execution_count)
340
+
341
+ if isinstance(response, DisplayJsonResponse):
342
+ self._accept_display(response)
343
+ elif isinstance(response, ArtifactBeginResponse):
344
+ self._accept_begin(response)
345
+ elif isinstance(response, ArtifactChunkResponse):
346
+ self._accept_chunk(response)
347
+ else:
348
+ self._accept_end(response)
349
+
350
+ def _accept_display(self, response: DisplayJsonResponse) -> None:
351
+ if self.active_transaction is not None:
352
+ raise ProtocolViolation("display interleaves artifact transaction")
353
+ self._require_sequence(response.sequence)
354
+ self.json_displays += 1
355
+ if self.json_displays > MAX_JSON_DISPLAYS:
356
+ raise ProtocolViolation("JSON display count exceeds maximum")
357
+ self.json_bytes += json_size(response.value, MAX_JSON_DISPLAY_BYTES)
358
+ if self.json_bytes > MAX_JSON_DISPLAY_BYTES:
359
+ raise ProtocolViolation("JSON display bytes exceed maximum")
360
+ self.next_sequence += 1
361
+
362
+ def _accept_begin(self, response: ArtifactBeginResponse) -> None:
363
+ if self.active_transaction is not None:
364
+ raise ProtocolViolation("artifact transactions cannot interleave")
365
+ self._require_sequence(response.sequence)
366
+ if response.transaction_id in self.transaction_ids:
367
+ raise ProtocolViolation("duplicate artifact transactionId")
368
+ self.transaction_ids.add(response.transaction_id)
369
+ self.artifact_transactions += 1
370
+ if self.artifact_transactions > MAX_ARTIFACT_TRANSACTIONS:
371
+ raise ProtocolViolation("artifact transaction count exceeds maximum")
372
+ if response.kind == "image":
373
+ self.image_transactions += 1
374
+ if self.image_transactions > MAX_IMAGE_TRANSACTIONS:
375
+ raise ProtocolViolation("image transaction count exceeds maximum")
376
+ self.artifact_bytes += response.size
377
+ if self.artifact_bytes > MAX_ARTIFACT_BYTES:
378
+ raise ProtocolViolation("artifact bytes exceed maximum")
379
+ self.active_transaction = _ArtifactTransaction(
380
+ response.transaction_id, response.sequence, response.size
381
+ )
382
+
383
+ def _accept_chunk(self, response: ArtifactChunkResponse) -> None:
384
+ transaction = self.active_transaction
385
+ if transaction is None or response.transaction_id != transaction.transaction_id:
386
+ raise ProtocolViolation("artifact chunk has no matching transaction")
387
+ if response.sequence != transaction.sequence:
388
+ raise ProtocolViolation("artifact chunk sequence mismatch")
389
+ if response.chunk_index != transaction.chunk_count:
390
+ raise ProtocolViolation("artifact chunkIndex is not contiguous")
391
+ self.artifact_chunks += 1
392
+ if self.artifact_chunks > MAX_ARTIFACT_CHUNKS:
393
+ raise ProtocolViolation("artifact chunk count exceeds maximum")
394
+ transaction.chunk_count += 1
395
+ transaction.received_bytes += len(response.data)
396
+ if transaction.received_bytes > transaction.size:
397
+ raise ProtocolViolation("artifact chunks exceed declared size")
398
+ digest = cast(Any, transaction.digest)
399
+ digest.update(response.data)
400
+
401
+ def _accept_end(self, response: ArtifactEndResponse) -> None:
402
+ transaction = self.active_transaction
403
+ if transaction is None or response.transaction_id != transaction.transaction_id:
404
+ raise ProtocolViolation("artifact end has no matching transaction")
405
+ if response.sequence != transaction.sequence:
406
+ raise ProtocolViolation("artifact end sequence mismatch")
407
+ if response.chunk_count != transaction.chunk_count:
408
+ raise ProtocolViolation("artifact chunkCount mismatch")
409
+ if transaction.received_bytes != transaction.size:
410
+ raise ProtocolViolation("artifact size mismatch")
411
+ digest = cast(Any, transaction.digest)
412
+ if response.sha256 != digest.hexdigest():
413
+ raise ProtocolViolation("artifact sha256 mismatch")
414
+ self.active_transaction = None
415
+ self.next_sequence += 1
416
+
417
+ def _require_execution_count(self, execution_count: int) -> None:
418
+ if execution_count != self.execution_count:
419
+ raise ProtocolViolation("output executionCount mismatch")
420
+
421
+ def _require_sequence(self, sequence: int) -> None:
422
+ if sequence != self.next_sequence or sequence > MAX_OUTPUT_SEQUENCE:
423
+ raise ProtocolViolation("output sequence gap or duplicate")
424
+
425
+
426
+ def encode_request(request: Request) -> bytes:
427
+ """Encode one validated host request including its length prefix."""
428
+ message_type, payload = _request_parts(request)
429
+ return _encode_envelope(message_type, request.correlation_id, payload, parse_request_payload)
430
+
431
+
432
+ def encode_response(response: Response) -> bytes:
433
+ """Encode one validated Runner response including its length prefix."""
434
+ message_type, payload = _response_parts(response)
435
+ return _encode_envelope(message_type, response.correlation_id, payload, parse_response_payload)
436
+
437
+
438
+ def parse_request_payload(payload: bytes) -> Request:
439
+ """Parse a request and retain correlation for schema-valid envelopes."""
440
+ try:
441
+ return _parse_request_payload_inner(payload)
442
+ except ProtocolViolation as error:
443
+ if error.correlation_id is None:
444
+ error.correlation_id = _peek_correlation_id(payload)
445
+ raise
446
+
447
+
448
+ def _parse_request_payload_inner(payload: bytes) -> Request:
449
+ """Parse and validate one request JSON payload without a length prefix."""
450
+ message_type, correlation_id, body = _parse_envelope(payload)
451
+
452
+ if message_type == "initialize":
453
+ _require_exact_keys(body, {"launchWorkingDirectory", "environmentKind"})
454
+ return InitializeRequest(
455
+ correlation_id=correlation_id,
456
+ launch_working_directory=_require_path(body, "launchWorkingDirectory", canonical=True),
457
+ environment_kind=_require_environment_kind(body, "environmentKind"),
458
+ )
459
+ if message_type == "execute":
460
+ _require_exact_keys(body, {"code", "packages"})
461
+ return ExecuteRequest(
462
+ correlation_id=correlation_id,
463
+ code=_require_string(body, "code", maximum_bytes=MAX_CODE_BYTES),
464
+ packages=_parse_package_names(body, "packages"),
465
+ )
466
+ if message_type == "shutdown":
467
+ _require_exact_keys(body, set())
468
+ return ShutdownRequest(correlation_id=correlation_id)
469
+
470
+ raise ProtocolViolation(f"unknown request type: {message_type}", correlation_id=correlation_id)
471
+
472
+
473
+ def parse_response_payload(payload: bytes) -> Response:
474
+ """Parse a response and retain correlation for schema-valid envelopes."""
475
+ try:
476
+ return _parse_response_payload_inner(payload)
477
+ except ProtocolViolation as error:
478
+ if error.correlation_id is None:
479
+ error.correlation_id = _peek_correlation_id(payload)
480
+ raise
481
+
482
+
483
+ def _parse_response_payload_inner(payload: bytes) -> Response:
484
+ """Parse and validate one response JSON payload without a length prefix."""
485
+ message_type, correlation_id, body = _parse_envelope(payload)
486
+
487
+ if message_type == "ready":
488
+ _require_exact_keys(body, {"executionCount", "runtime"})
489
+ execution_count = _require_integer(body, "executionCount", minimum=0)
490
+ if execution_count != 0:
491
+ raise ProtocolViolation("ready executionCount must be 0", correlation_id=correlation_id)
492
+ return ReadyResponse(
493
+ correlation_id=correlation_id,
494
+ execution_count=execution_count,
495
+ runtime=_parse_runtime(_require_object(body, "runtime"), correlation_id),
496
+ )
497
+ if message_type == "started":
498
+ _require_exact_keys(body, {"executionCount"})
499
+ return StartedResponse(
500
+ correlation_id=correlation_id,
501
+ execution_count=_require_integer(body, "executionCount", minimum=1),
502
+ )
503
+ if message_type == "display_json":
504
+ _require_exact_keys(body, {"executionCount", "sequence", "value"})
505
+ value = cast(JsonValue, body["value"])
506
+ if json_size(value, MAX_JSON_DISPLAY_ITEM_BYTES) > MAX_JSON_DISPLAY_ITEM_BYTES:
507
+ raise ProtocolViolation("display_json value exceeds maximum size")
508
+ return DisplayJsonResponse(
509
+ correlation_id=correlation_id,
510
+ execution_count=_require_integer(body, "executionCount", minimum=1),
511
+ sequence=_require_bounded_integer(body, "sequence", 1, MAX_OUTPUT_SEQUENCE),
512
+ value=snapshot_json(value, MAX_JSON_DISPLAY_ITEM_BYTES),
513
+ )
514
+ if message_type == "artifact_begin":
515
+ _require_exact_keys(
516
+ body,
517
+ {
518
+ "executionCount",
519
+ "sequence",
520
+ "transactionId",
521
+ "kind",
522
+ "name",
523
+ "declaredMimeType",
524
+ "size",
525
+ },
526
+ )
527
+ return ArtifactBeginResponse(
528
+ correlation_id=correlation_id,
529
+ execution_count=_require_integer(body, "executionCount", minimum=1),
530
+ sequence=_require_bounded_integer(body, "sequence", 1, MAX_OUTPUT_SEQUENCE),
531
+ transaction_id=_require_uuid(body, "transactionId"),
532
+ kind=_require_artifact_kind(body, "kind"),
533
+ name=_require_logical_name(body, "name"),
534
+ declared_mime_type=_require_mime(body, "declaredMimeType"),
535
+ size=_require_bounded_integer(body, "size", 0, MAX_ARTIFACT_BYTES),
536
+ )
537
+ if message_type == "artifact_chunk":
538
+ _require_exact_keys(
539
+ body, {"executionCount", "sequence", "transactionId", "chunkIndex", "data"}
540
+ )
541
+ return ArtifactChunkResponse(
542
+ correlation_id=correlation_id,
543
+ execution_count=_require_integer(body, "executionCount", minimum=1),
544
+ sequence=_require_bounded_integer(body, "sequence", 1, MAX_OUTPUT_SEQUENCE),
545
+ transaction_id=_require_uuid(body, "transactionId"),
546
+ chunk_index=_require_bounded_integer(body, "chunkIndex", 0, MAX_ARTIFACT_CHUNKS - 1),
547
+ data=_require_canonical_base64(body, "data"),
548
+ )
549
+ if message_type == "artifact_end":
550
+ _require_exact_keys(
551
+ body, {"executionCount", "sequence", "transactionId", "chunkCount", "sha256"}
552
+ )
553
+ return ArtifactEndResponse(
554
+ correlation_id=correlation_id,
555
+ execution_count=_require_integer(body, "executionCount", minimum=1),
556
+ sequence=_require_bounded_integer(body, "sequence", 1, MAX_OUTPUT_SEQUENCE),
557
+ transaction_id=_require_uuid(body, "transactionId"),
558
+ chunk_count=_require_bounded_integer(body, "chunkCount", 0, MAX_ARTIFACT_CHUNKS),
559
+ sha256=_require_sha256(body, "sha256"),
560
+ )
561
+ if message_type == "done":
562
+ return _parse_done(body, correlation_id)
563
+ if message_type == "protocol_error":
564
+ _require_exact_keys(body, {"code", "message"})
565
+ return ProtocolErrorResponse(
566
+ correlation_id=correlation_id,
567
+ code=_require_protocol_error_code(body, "code"),
568
+ message=_require_string(
569
+ body, "message", maximum_bytes=MAX_PROTOCOL_ERROR_MESSAGE_BYTES, nonempty=True
570
+ ),
571
+ )
572
+
573
+ raise ProtocolViolation(f"unknown response type: {message_type}", correlation_id=correlation_id)
574
+
575
+
576
+ def read_request(stream: BinaryIO) -> Request | None:
577
+ """Read one framed request, returning None only for clean prefix-boundary EOF."""
578
+ payload = read_payload(stream)
579
+ return None if payload is None else parse_request_payload(payload)
580
+
581
+
582
+ def read_response(stream: BinaryIO) -> Response | None:
583
+ """Read one framed response, returning None only for clean prefix-boundary EOF."""
584
+ payload = read_payload(stream)
585
+ return None if payload is None else parse_response_payload(payload)
586
+
587
+
588
+ def read_payload(stream: BinaryIO) -> bytes | None:
589
+ """Read one bounded length-prefixed JSON payload."""
590
+ prefix = _read_exact(stream, _LENGTH.size, allow_initial_eof=True)
591
+ if prefix is None:
592
+ return None
593
+
594
+ (length,) = _LENGTH.unpack(prefix)
595
+ if length == 0:
596
+ raise ProtocolViolation("zero-length protocol frame")
597
+ if length > MAX_CONTROL_FRAME_BYTES:
598
+ raise ProtocolViolation(f"protocol frame exceeds {MAX_CONTROL_FRAME_BYTES} bytes")
599
+
600
+ payload = _read_exact(stream, length, allow_initial_eof=False)
601
+ if payload is None: # pragma: no cover - forbidden by allow_initial_eof=False
602
+ raise ProtocolViolation("truncated protocol frame")
603
+ return payload
604
+
605
+
606
+ def write_frame(stream: BinaryIO, frame: bytes) -> None:
607
+ """Write and flush one already encoded frame, handling partial writes."""
608
+ view = memoryview(frame)
609
+ while view:
610
+ try:
611
+ written = stream.write(view)
612
+ except InterruptedError:
613
+ continue
614
+ if written is None or written <= 0:
615
+ raise BrokenPipeError("protocol stream accepted no bytes")
616
+ view = view[written:]
617
+ stream.flush()
618
+
619
+
620
+ def json_size(value: JsonValue, maximum_bytes: int | None = None) -> int:
621
+ """Return compact JSON size without materializing the complete encoding."""
622
+ if maximum_bytes is not None and maximum_bytes < 0:
623
+ raise ValueError("maximum_bytes must be non-negative")
624
+ return _json_encoded_size(value, maximum_bytes, set())
625
+
626
+
627
+ def snapshot_json(value: JsonValue, maximum_bytes: int) -> JsonValue:
628
+ """Detach one JSON value after proving its complete encoding is bounded."""
629
+ if json_size(value, maximum_bytes) > maximum_bytes:
630
+ raise ProtocolViolation(f"JSON value exceeds {maximum_bytes} bytes")
631
+ encoded = _dump_json(value)
632
+ if len(encoded) > maximum_bytes:
633
+ raise ProtocolViolation(f"JSON value exceeds {maximum_bytes} bytes")
634
+ try:
635
+ parsed = cast(JsonValue, _JSON_LOADS(encoded))
636
+ except (RecursionError, ValueError) as error:
637
+ raise ProtocolViolation(f"bounded JSON snapshot failed: {error}") from error
638
+ _validate_json_value(parsed)
639
+ return parsed
640
+
641
+
642
+ def _json_encoded_size(
643
+ value: JsonValue,
644
+ maximum_bytes: int | None,
645
+ active: set[int],
646
+ ) -> int:
647
+ if value is None:
648
+ return _bounded_size(4, maximum_bytes)
649
+ if type(value) is bool:
650
+ return _bounded_size(4 if value else 5, maximum_bytes)
651
+ if type(value) is int:
652
+ if maximum_bytes is not None and value:
653
+ estimated_digits = (abs(value).bit_length() * 30_103) // 100_000 + 1
654
+ estimated_digits += value < 0
655
+ if estimated_digits > maximum_bytes:
656
+ return maximum_bytes + 1
657
+ try:
658
+ return _bounded_size(len(str(value)), maximum_bytes)
659
+ except ValueError as error:
660
+ raise ProtocolViolation(f"integer cannot be encoded as JSON: {error}") from error
661
+ if type(value) is float:
662
+ if not _ISFINITE(value):
663
+ raise ProtocolViolation("non-finite JSON number")
664
+ return _bounded_size(len(_dump_json(value)), maximum_bytes)
665
+ if type(value) is str:
666
+ return _json_string_size(value, maximum_bytes)
667
+ if type(value) is list:
668
+ identity = id(value)
669
+ if identity in active:
670
+ raise ProtocolViolation("cyclic JSON array")
671
+ active.add(identity)
672
+ try:
673
+ total = _bounded_size(2, maximum_bytes)
674
+ if maximum_bytes is not None and total > maximum_bytes:
675
+ return total
676
+ for index, item in enumerate(value):
677
+ if index:
678
+ total += 1
679
+ remaining = None if maximum_bytes is None else max(0, maximum_bytes - total)
680
+ total += _json_encoded_size(item, remaining, active)
681
+ if maximum_bytes is not None and total > maximum_bytes:
682
+ return maximum_bytes + 1
683
+ return total
684
+ finally:
685
+ active.remove(identity)
686
+ if type(value) is dict:
687
+ identity = id(value)
688
+ if identity in active:
689
+ raise ProtocolViolation("cyclic JSON object")
690
+ active.add(identity)
691
+ try:
692
+ total = _bounded_size(2, maximum_bytes)
693
+ if maximum_bytes is not None and total > maximum_bytes:
694
+ return total
695
+ for index, (key, item) in enumerate(value.items()):
696
+ if type(key) is not str:
697
+ raise ProtocolViolation("JSON object keys must be strings")
698
+ if index:
699
+ total += 1
700
+ remaining = None if maximum_bytes is None else max(0, maximum_bytes - total)
701
+ total += _json_string_size(key, remaining)
702
+ total += 1
703
+ remaining = None if maximum_bytes is None else max(0, maximum_bytes - total)
704
+ total += _json_encoded_size(item, remaining, active)
705
+ if maximum_bytes is not None and total > maximum_bytes:
706
+ return maximum_bytes + 1
707
+ return total
708
+ finally:
709
+ active.remove(identity)
710
+ raise ProtocolViolation(f"unsupported JSON value type: {type(value).__name__}")
711
+
712
+
713
+ def _json_string_size(value: str, maximum_bytes: int | None) -> int:
714
+ total = 2
715
+ short_escapes = {"\b", "\t", "\n", "\f", "\r"}
716
+ for character in value:
717
+ codepoint = ord(character)
718
+ if character in {'"', "\\"} or character in short_escapes:
719
+ total += 2
720
+ elif codepoint < 0x20:
721
+ total += 6
722
+ elif 0xD800 <= codepoint <= 0xDFFF:
723
+ raise ProtocolViolation("JSON string contains invalid Unicode surrogate")
724
+ elif codepoint < 0x80:
725
+ total += 1
726
+ elif codepoint < 0x800:
727
+ total += 2
728
+ elif codepoint < 0x10000:
729
+ total += 3
730
+ else:
731
+ total += 4
732
+ if maximum_bytes is not None and total > maximum_bytes:
733
+ return maximum_bytes + 1
734
+ return total
735
+
736
+
737
+ def _bounded_size(size: int, maximum_bytes: int | None) -> int:
738
+ if maximum_bytes is not None and size > maximum_bytes:
739
+ return maximum_bytes + 1
740
+ return size
741
+
742
+
743
+ def _encode_envelope(
744
+ message_type: str,
745
+ correlation_id: str,
746
+ payload: Mapping[str, object],
747
+ validator: Callable[[bytes], object],
748
+ ) -> bytes:
749
+ envelope: dict[str, object] = {
750
+ "protocolVersion": PROTOCOL_VERSION,
751
+ "type": message_type,
752
+ "correlationId": correlation_id,
753
+ "payload": dict(payload),
754
+ }
755
+ encoded = _dump_json(envelope)
756
+ if not encoded:
757
+ raise ProtocolViolation("empty encoded protocol frame")
758
+ if len(encoded) > MAX_CONTROL_FRAME_BYTES:
759
+ raise ProtocolViolation(
760
+ f"encoded protocol frame exceeds {MAX_CONTROL_FRAME_BYTES} bytes",
761
+ correlation_id=correlation_id if _valid_correlation_id(correlation_id) else None,
762
+ )
763
+ validator(encoded)
764
+ return _LENGTH.pack(len(encoded)) + encoded
765
+
766
+
767
+ def _request_parts(request: Request) -> tuple[str, Mapping[str, object]]:
768
+ if isinstance(request, InitializeRequest):
769
+ return (
770
+ "initialize",
771
+ {
772
+ "launchWorkingDirectory": request.launch_working_directory,
773
+ "environmentKind": request.environment_kind,
774
+ },
775
+ )
776
+ if isinstance(request, ExecuteRequest):
777
+ return "execute", {"code": request.code, "packages": list(request.packages)}
778
+ return "shutdown", {}
779
+
780
+
781
+ def _response_parts(response: Response) -> tuple[str, Mapping[str, object]]:
782
+ if isinstance(response, ReadyResponse):
783
+ return (
784
+ "ready",
785
+ {
786
+ "executionCount": response.execution_count,
787
+ "runtime": _runtime_to_wire(response.runtime),
788
+ },
789
+ )
790
+ if isinstance(response, StartedResponse):
791
+ return "started", {"executionCount": response.execution_count}
792
+ if isinstance(response, DisplayJsonResponse):
793
+ return "display_json", {
794
+ "executionCount": response.execution_count,
795
+ "sequence": response.sequence,
796
+ "value": response.value,
797
+ }
798
+ if isinstance(response, ArtifactBeginResponse):
799
+ return "artifact_begin", {
800
+ "executionCount": response.execution_count,
801
+ "sequence": response.sequence,
802
+ "transactionId": response.transaction_id,
803
+ "kind": response.kind,
804
+ "name": response.name,
805
+ "declaredMimeType": response.declared_mime_type,
806
+ "size": response.size,
807
+ }
808
+ if isinstance(response, ArtifactChunkResponse):
809
+ return "artifact_chunk", {
810
+ "executionCount": response.execution_count,
811
+ "sequence": response.sequence,
812
+ "transactionId": response.transaction_id,
813
+ "chunkIndex": response.chunk_index,
814
+ "data": base64.b64encode(response.data).decode("ascii"),
815
+ }
816
+ if isinstance(response, ArtifactEndResponse):
817
+ return "artifact_end", {
818
+ "executionCount": response.execution_count,
819
+ "sequence": response.sequence,
820
+ "transactionId": response.transaction_id,
821
+ "chunkCount": response.chunk_count,
822
+ "sha256": response.sha256,
823
+ }
824
+ if isinstance(response, DoneResponse):
825
+ if not response.has_result and response.result is not None:
826
+ raise ProtocolViolation("hasResult=false requires result=None")
827
+ payload: dict[str, object] = {
828
+ "outcome": response.outcome,
829
+ "executionCount": response.execution_count,
830
+ "runtime": _runtime_to_wire(response.runtime),
831
+ "hasResult": response.has_result,
832
+ }
833
+ if response.has_result:
834
+ payload["result"] = response.result
835
+ if response.error is not None:
836
+ payload["error"] = _error_to_wire(response.error)
837
+ return "done", payload
838
+ return "protocol_error", {"code": response.code, "message": response.message}
839
+
840
+
841
+ def _parse_envelope(payload: bytes) -> tuple[str, str, dict[str, object]]:
842
+ raw = _load_json_object(payload)
843
+ correlation_id = raw.get("correlationId")
844
+ valid_correlation = (
845
+ correlation_id
846
+ if isinstance(correlation_id, str) and _valid_correlation_id(correlation_id)
847
+ else None
848
+ )
849
+ try:
850
+ _validate_json_value(cast(JsonValue, raw))
851
+ except RecursionError as error:
852
+ raise ProtocolViolation(f"invalid JSON payload: {error}") from error
853
+ _require_exact_keys(raw, {"protocolVersion", "type", "correlationId", "payload"})
854
+
855
+ version = _require_integer(raw, "protocolVersion", minimum=0)
856
+ if version != PROTOCOL_VERSION:
857
+ raise ProtocolViolation(
858
+ f"protocol version {version} does not match {PROTOCOL_VERSION}",
859
+ code="version_mismatch",
860
+ correlation_id=valid_correlation,
861
+ )
862
+
863
+ parsed_correlation = _require_string(raw, "correlationId", maximum_bytes=32)
864
+ if not _valid_correlation_id(parsed_correlation):
865
+ raise ProtocolViolation("invalid correlationId")
866
+
867
+ message_type = _require_string(raw, "type", maximum_bytes=64, nonempty=True)
868
+ body = _require_object(raw, "payload")
869
+ return message_type, parsed_correlation, body
870
+
871
+
872
+ def _parse_done(body: dict[str, object], correlation_id: str) -> DoneResponse:
873
+ allowed = {
874
+ "outcome",
875
+ "executionCount",
876
+ "runtime",
877
+ "hasResult",
878
+ "result",
879
+ "error",
880
+ }
881
+ _require_allowed_keys(body, allowed)
882
+ required = {"outcome", "executionCount", "runtime", "hasResult"}
883
+ missing = required - body.keys()
884
+ if missing:
885
+ raise ProtocolViolation(
886
+ f"missing done fields: {sorted(missing)}", correlation_id=correlation_id
887
+ )
888
+
889
+ outcome = _require_done_outcome(body, "outcome")
890
+ execution_count = _require_integer(body, "executionCount", minimum=1)
891
+ runtime = _parse_runtime(_require_object(body, "runtime"), correlation_id)
892
+ has_result = _require_boolean(body, "hasResult")
893
+ has_result_field = "result" in body
894
+ has_error_field = "error" in body
895
+
896
+ if outcome == "completed":
897
+ if has_result != has_result_field:
898
+ raise ProtocolViolation(
899
+ "completed result presence must match hasResult",
900
+ correlation_id=correlation_id,
901
+ )
902
+ if has_error_field:
903
+ raise ProtocolViolation(
904
+ "completed response cannot contain error", correlation_id=correlation_id
905
+ )
906
+ else:
907
+ if has_result or has_result_field:
908
+ raise ProtocolViolation(
909
+ f"{outcome} response cannot contain result", correlation_id=correlation_id
910
+ )
911
+ if not has_error_field:
912
+ raise ProtocolViolation(
913
+ f"{outcome} response requires error", correlation_id=correlation_id
914
+ )
915
+
916
+ result: JsonValue = None
917
+ if has_result_field:
918
+ result = cast(JsonValue, body["result"])
919
+ if json_size(result, MAX_RESULT_BYTES) > MAX_RESULT_BYTES:
920
+ raise ProtocolViolation(
921
+ "done result exceeds maximum size", correlation_id=correlation_id
922
+ )
923
+
924
+ error = None
925
+ if has_error_field:
926
+ error = _parse_error(_require_object(body, "error"), correlation_id)
927
+ if outcome == "interrupted" and error.type != "KeyboardInterrupt":
928
+ raise ProtocolViolation(
929
+ "interrupted error type must be KeyboardInterrupt",
930
+ correlation_id=correlation_id,
931
+ )
932
+
933
+ return DoneResponse(
934
+ correlation_id=correlation_id,
935
+ outcome=outcome,
936
+ execution_count=execution_count,
937
+ runtime=runtime,
938
+ has_result=has_result,
939
+ result=result,
940
+ error=error,
941
+ )
942
+
943
+
944
+ def _parse_runtime(body: dict[str, object], correlation_id: str) -> RuntimeDiagnostics:
945
+ _require_exact_keys(
946
+ body,
947
+ {
948
+ "executable",
949
+ "pythonVersion",
950
+ "workingDirectory",
951
+ "environmentKind",
952
+ "implementation",
953
+ "platform",
954
+ "packages",
955
+ },
956
+ )
957
+ return RuntimeDiagnostics(
958
+ executable=_require_path(body, "executable", canonical=False),
959
+ python_version=_require_string(
960
+ body, "pythonVersion", maximum_bytes=MAX_PYTHON_VERSION_BYTES, nonempty=True
961
+ ),
962
+ working_directory=_require_path(body, "workingDirectory", canonical=True),
963
+ environment_kind=_require_environment_kind(body, "environmentKind"),
964
+ implementation=_require_control_free_string(
965
+ body, "implementation", maximum_bytes=MAX_IMPLEMENTATION_BYTES
966
+ ),
967
+ platform=_require_control_free_string(body, "platform", maximum_bytes=MAX_PLATFORM_BYTES),
968
+ packages=_parse_runtime_packages(body, correlation_id),
969
+ )
970
+
971
+
972
+ def _parse_error(body: dict[str, object], correlation_id: str) -> ErrorInfo:
973
+ _require_exact_keys(body, {"type", "message", "traceback"})
974
+ error_type = _require_string(body, "type", maximum_bytes=MAX_ERROR_TYPE_BYTES, nonempty=True)
975
+ message = _require_string(body, "message", maximum_bytes=MAX_ERROR_MESSAGE_BYTES)
976
+ traceback_value = body.get("traceback")
977
+ if not isinstance(traceback_value, list) or not all(
978
+ isinstance(line, str) for line in traceback_value
979
+ ):
980
+ raise ProtocolViolation(
981
+ "error traceback must be a string array", correlation_id=correlation_id
982
+ )
983
+ error = ErrorInfo(error_type, message, tuple(cast(list[str], traceback_value)))
984
+ if json_size(cast(JsonValue, _error_to_wire(error)), MAX_ERROR_BYTES) > MAX_ERROR_BYTES:
985
+ raise ProtocolViolation("encoded error exceeds maximum size", correlation_id=correlation_id)
986
+ return error
987
+
988
+
989
+ def _runtime_to_wire(runtime: RuntimeDiagnostics) -> dict[str, object]:
990
+ return {
991
+ "executable": runtime.executable,
992
+ "pythonVersion": runtime.python_version,
993
+ "workingDirectory": runtime.working_directory,
994
+ "environmentKind": runtime.environment_kind,
995
+ "implementation": runtime.implementation,
996
+ "platform": runtime.platform,
997
+ "packages": [_runtime_package_to_wire(package) for package in runtime.packages],
998
+ }
999
+
1000
+
1001
+ def _error_to_wire(error: ErrorInfo) -> dict[str, object]:
1002
+ return {
1003
+ "type": error.type,
1004
+ "message": error.message,
1005
+ "traceback": list(error.traceback),
1006
+ }
1007
+
1008
+
1009
+ _INVALID_JSON_INTEGER = object()
1010
+
1011
+
1012
+ def _parse_json_integer(value: str) -> int | object:
1013
+ try:
1014
+ return int(value)
1015
+ except ValueError:
1016
+ return _INVALID_JSON_INTEGER
1017
+
1018
+
1019
+ def _load_json_object(payload: bytes) -> dict[str, object]:
1020
+ if not payload:
1021
+ raise ProtocolViolation("empty JSON payload")
1022
+ if len(payload) > MAX_CONTROL_FRAME_BYTES:
1023
+ raise ProtocolViolation("JSON payload exceeds maximum size")
1024
+ try:
1025
+ text = payload.decode("utf-8")
1026
+ parsed = _JSON_LOADS(
1027
+ text,
1028
+ object_pairs_hook=_reject_duplicate_keys,
1029
+ parse_constant=_reject_nonfinite_constant,
1030
+ parse_int=_parse_json_integer,
1031
+ )
1032
+ except ProtocolViolation:
1033
+ raise
1034
+ except (UnicodeDecodeError, RecursionError, ValueError) as error:
1035
+ raise ProtocolViolation(f"invalid JSON payload: {error}") from error
1036
+ if not isinstance(parsed, dict):
1037
+ raise ProtocolViolation("protocol payload must be a JSON object")
1038
+ result = cast(dict[str, object], parsed)
1039
+ return result
1040
+
1041
+
1042
+ def _peek_correlation_id(payload: bytes) -> str | None:
1043
+ try:
1044
+ parsed = _JSON_LOADS(
1045
+ payload.decode("utf-8"),
1046
+ object_pairs_hook=_reject_duplicate_keys,
1047
+ parse_int=_parse_json_integer,
1048
+ )
1049
+ except (UnicodeDecodeError, RecursionError, ValueError):
1050
+ return None
1051
+ if not isinstance(parsed, dict):
1052
+ return None
1053
+ correlation_id = parsed.get("correlationId")
1054
+ if isinstance(correlation_id, str) and _valid_correlation_id(correlation_id):
1055
+ return correlation_id
1056
+ return None
1057
+
1058
+
1059
+ def _dump_json(value: object) -> bytes:
1060
+ try:
1061
+ return _JSON_DUMPS(
1062
+ value,
1063
+ allow_nan=False,
1064
+ ensure_ascii=False,
1065
+ separators=(",", ":"),
1066
+ ).encode("utf-8")
1067
+ except (TypeError, ValueError, OverflowError, RecursionError, UnicodeEncodeError) as error:
1068
+ raise ProtocolViolation(f"value is not strict JSON: {error}") from error
1069
+
1070
+
1071
+ class _DuplicateKeyError(ValueError):
1072
+ pass
1073
+
1074
+
1075
+ def _reject_duplicate_keys(pairs: list[tuple[str, object]]) -> dict[str, object]:
1076
+ result: dict[str, object] = {}
1077
+ for key, value in pairs:
1078
+ if key in result:
1079
+ raise _DuplicateKeyError(f"duplicate key: {key}")
1080
+ result[key] = value
1081
+ return result
1082
+
1083
+
1084
+ def _reject_nonfinite_constant(value: str) -> object:
1085
+ raise ProtocolViolation(f"non-finite JSON constant: {value}")
1086
+
1087
+
1088
+ def _validate_json_value(value: JsonValue, active: set[int] | None = None) -> None:
1089
+ if value is None or type(value) in {bool, int, str}:
1090
+ return
1091
+ if type(value) is float:
1092
+ if not _ISFINITE(value):
1093
+ raise ProtocolViolation("non-finite JSON number")
1094
+ return
1095
+ if active is None:
1096
+ active = set()
1097
+ if type(value) is list:
1098
+ identity = id(value)
1099
+ if identity in active:
1100
+ raise ProtocolViolation("cyclic JSON array")
1101
+ active.add(identity)
1102
+ try:
1103
+ for item in value:
1104
+ _validate_json_value(item, active)
1105
+ finally:
1106
+ active.remove(identity)
1107
+ return
1108
+ if type(value) is dict:
1109
+ identity = id(value)
1110
+ if identity in active:
1111
+ raise ProtocolViolation("cyclic JSON object")
1112
+ active.add(identity)
1113
+ try:
1114
+ for key, item in value.items():
1115
+ if type(key) is not str:
1116
+ raise ProtocolViolation("JSON object keys must be strings")
1117
+ _validate_json_value(item, active)
1118
+ finally:
1119
+ active.remove(identity)
1120
+ return
1121
+ raise ProtocolViolation(f"unsupported JSON value type: {type(value).__name__}")
1122
+
1123
+
1124
+ def _require_exact_keys(body: Mapping[str, object], expected: set[str]) -> None:
1125
+ actual = set(body)
1126
+ if actual != expected:
1127
+ missing = sorted(expected - actual)
1128
+ unknown = sorted(actual - expected)
1129
+ raise ProtocolViolation(f"schema keys differ: missing={missing}, unknown={unknown}")
1130
+
1131
+
1132
+ def _require_allowed_keys(body: Mapping[str, object], allowed: set[str]) -> None:
1133
+ unknown = sorted(set(body) - allowed)
1134
+ if unknown:
1135
+ raise ProtocolViolation(f"unknown schema keys: {unknown}")
1136
+
1137
+
1138
+ def _require_object(body: Mapping[str, object], key: str) -> dict[str, object]:
1139
+ value = body.get(key)
1140
+ if not isinstance(value, dict):
1141
+ raise ProtocolViolation(f"{key} must be an object")
1142
+ return cast(dict[str, object], value)
1143
+
1144
+
1145
+ def _require_string(
1146
+ body: Mapping[str, object],
1147
+ key: str,
1148
+ *,
1149
+ maximum_bytes: int,
1150
+ nonempty: bool = False,
1151
+ ) -> str:
1152
+ value = body.get(key)
1153
+ if type(value) is not str:
1154
+ raise ProtocolViolation(f"{key} must be an exact string")
1155
+ try:
1156
+ size = len(value.encode("utf-8"))
1157
+ except UnicodeEncodeError as error:
1158
+ raise ProtocolViolation(f"{key} contains invalid Unicode") from error
1159
+ if nonempty and not value:
1160
+ raise ProtocolViolation(f"{key} must not be empty")
1161
+ if size > maximum_bytes:
1162
+ raise ProtocolViolation(f"{key} exceeds {maximum_bytes} UTF-8 bytes")
1163
+ return value
1164
+
1165
+
1166
+ def normalize_package_name(value: str) -> str:
1167
+ """Return one validated PEP 503 normalized distribution name."""
1168
+ if not (1 <= len(value) <= MAX_PACKAGE_NAME_BYTES) or _PACKAGE_NAME.fullmatch(value) is None:
1169
+ raise ValueError("invalid Python distribution name")
1170
+ return re.sub(r"[-_.]+", "-", value).lower()
1171
+
1172
+
1173
+ def validate_package_names(values: list[str] | tuple[str, ...]) -> tuple[str, ...]:
1174
+ """Validate the bounded package request and reject normalized duplicates."""
1175
+ if len(values) > MAX_PACKAGE_COUNT:
1176
+ raise ValueError(f"packages exceeds {MAX_PACKAGE_COUNT} items")
1177
+ normalized: set[str] = set()
1178
+ result: list[str] = []
1179
+ for value in values:
1180
+ if type(value) is not str:
1181
+ raise ValueError("package names must be exact strings")
1182
+ key = normalize_package_name(value)
1183
+ if key in normalized:
1184
+ raise ValueError("duplicate normalized package name")
1185
+ normalized.add(key)
1186
+ result.append(value)
1187
+ return tuple(result)
1188
+
1189
+
1190
+ def validate_logical_name(value: str) -> str:
1191
+ """Validate one public artifact logical name."""
1192
+ if not (1 <= len(value) <= MAX_LOGICAL_NAME_CODEPOINTS):
1193
+ raise ValueError("invalid artifact name length")
1194
+ if "/" in value or "\\" in value:
1195
+ raise ValueError("artifact name contains a path separator")
1196
+ if any(unicodedata.category(character) in {"Cc", "Cf"} for character in value):
1197
+ raise ValueError("artifact name contains a control or format character")
1198
+ return value
1199
+
1200
+
1201
+ def validate_mime_type(value: str) -> str:
1202
+ """Validate one printable-ASCII media type with optional parameters."""
1203
+ if not (1 <= len(value) <= MAX_MIME_BYTES) or any(
1204
+ ord(character) < 0x20 or ord(character) > 0x7E for character in value
1205
+ ):
1206
+ raise ValueError("invalid MIME length or character")
1207
+ parts = _split_mime_parts(value)
1208
+ media_type = parts[0].strip()
1209
+ if media_type.count("/") != 1:
1210
+ raise ValueError("invalid MIME media type")
1211
+ major, minor = media_type.split("/", 1)
1212
+ if _MIME_TOKEN.fullmatch(major) is None or _MIME_TOKEN.fullmatch(minor) is None:
1213
+ raise ValueError("invalid MIME media type")
1214
+ for parameter in parts[1:]:
1215
+ item = parameter.strip()
1216
+ if "=" not in item:
1217
+ raise ValueError("invalid MIME parameter")
1218
+ name, raw_value = item.split("=", 1)
1219
+ if _MIME_TOKEN.fullmatch(name.strip()) is None or not _valid_mime_parameter_value(
1220
+ raw_value.strip()
1221
+ ):
1222
+ raise ValueError("invalid MIME parameter")
1223
+ return value
1224
+
1225
+
1226
+ def _split_mime_parts(value: str) -> list[str]:
1227
+ parts: list[str] = []
1228
+ start = 0
1229
+ quoted = False
1230
+ escaped = False
1231
+ for index, character in enumerate(value):
1232
+ if escaped:
1233
+ escaped = False
1234
+ elif quoted and character == "\\":
1235
+ escaped = True
1236
+ elif character == '"':
1237
+ quoted = not quoted
1238
+ elif character == ";" and not quoted:
1239
+ parts.append(value[start:index])
1240
+ start = index + 1
1241
+ if quoted or escaped:
1242
+ raise ValueError("unterminated MIME quoted string")
1243
+ parts.append(value[start:])
1244
+ return parts
1245
+
1246
+
1247
+ def _valid_mime_parameter_value(value: str) -> bool:
1248
+ if _MIME_TOKEN.fullmatch(value) is not None:
1249
+ return True
1250
+ if len(value) < 2 or value[0] != '"' or value[-1] != '"':
1251
+ return False
1252
+ escaped = False
1253
+ for character in value[1:-1]:
1254
+ if escaped:
1255
+ escaped = False
1256
+ elif character == "\\":
1257
+ escaped = True
1258
+ elif character == '"':
1259
+ return False
1260
+ return not escaped
1261
+
1262
+
1263
+ def _parse_package_names(body: Mapping[str, object], key: str) -> tuple[str, ...]:
1264
+ value = body.get(key)
1265
+ if not isinstance(value, list):
1266
+ raise ProtocolViolation(f"{key} must be an array")
1267
+ try:
1268
+ return validate_package_names(cast(list[str], value))
1269
+ except ValueError as error:
1270
+ raise ProtocolViolation(str(error)) from error
1271
+
1272
+
1273
+ def _parse_runtime_packages(
1274
+ body: Mapping[str, object], correlation_id: str
1275
+ ) -> tuple[RuntimePackage, ...]:
1276
+ raw = body.get("packages")
1277
+ if not isinstance(raw, list) or len(raw) > MAX_PACKAGE_COUNT:
1278
+ raise ProtocolViolation("runtime packages must be a bounded array")
1279
+ result: list[RuntimePackage] = []
1280
+ names: list[str] = []
1281
+ for item in raw:
1282
+ if not isinstance(item, dict):
1283
+ raise ProtocolViolation("runtime package must be an object")
1284
+ package = cast(dict[str, object], item)
1285
+ _require_allowed_keys(package, {"name", "version", "error"})
1286
+ if set(package) not in ({"name", "version"}, {"name", "error"}):
1287
+ raise ProtocolViolation("runtime package requires exactly one of version or error")
1288
+ name = _require_string(package, "name", maximum_bytes=MAX_PACKAGE_NAME_BYTES, nonempty=True)
1289
+ names.append(name)
1290
+ if "version" in package:
1291
+ version = _require_control_free_string(
1292
+ package, "version", maximum_bytes=MAX_PACKAGE_VERSION_BYTES
1293
+ )
1294
+ result.append(RuntimePackage(name=name, version=version))
1295
+ else:
1296
+ result.append(RuntimePackage(name=name, error=_require_package_error(package, "error")))
1297
+ try:
1298
+ validate_package_names(names)
1299
+ except ValueError as error:
1300
+ raise ProtocolViolation(str(error), correlation_id=correlation_id) from error
1301
+ return tuple(result)
1302
+
1303
+
1304
+ def _runtime_package_to_wire(package: RuntimePackage) -> dict[str, object]:
1305
+ result: dict[str, object] = {"name": package.name}
1306
+ if (package.version is None) == (package.error is None):
1307
+ raise ProtocolViolation("runtime package requires exactly one of version or error")
1308
+ if package.version is not None:
1309
+ result["version"] = package.version
1310
+ else:
1311
+ result["error"] = package.error
1312
+ return result
1313
+
1314
+
1315
+ def _require_control_free_string(
1316
+ body: Mapping[str, object], key: str, *, maximum_bytes: int
1317
+ ) -> str:
1318
+ value = _require_string(body, key, maximum_bytes=maximum_bytes, nonempty=True)
1319
+ if any(unicodedata.category(character) in {"Cc", "Cf"} for character in value):
1320
+ raise ProtocolViolation(f"{key} contains a control or format character")
1321
+ return value
1322
+
1323
+
1324
+ def _require_bounded_integer(
1325
+ body: Mapping[str, object], key: str, minimum: int, maximum: int
1326
+ ) -> int:
1327
+ value = _require_integer(body, key, minimum=minimum)
1328
+ if value > maximum:
1329
+ raise ProtocolViolation(f"{key} must be at most {maximum}")
1330
+ return value
1331
+
1332
+
1333
+ def _require_uuid(body: Mapping[str, object], key: str) -> str:
1334
+ value = _require_string(body, key, maximum_bytes=32)
1335
+ if _UUID_HEX.fullmatch(value) is None:
1336
+ raise ProtocolViolation(f"invalid {key}")
1337
+ return value
1338
+
1339
+
1340
+ def _require_sha256(body: Mapping[str, object], key: str) -> str:
1341
+ value = _require_string(body, key, maximum_bytes=64)
1342
+ if _SHA256.fullmatch(value) is None:
1343
+ raise ProtocolViolation(f"invalid {key}")
1344
+ return value
1345
+
1346
+
1347
+ def _require_artifact_kind(body: Mapping[str, object], key: str) -> ArtifactKind:
1348
+ value = _require_string(body, key, maximum_bytes=8)
1349
+ if value not in {"artifact", "image"}:
1350
+ raise ProtocolViolation(f"invalid {key}")
1351
+ return cast(ArtifactKind, value)
1352
+
1353
+
1354
+ def _require_package_error(body: Mapping[str, object], key: str) -> PackageError:
1355
+ value = _require_string(body, key, maximum_bytes=32)
1356
+ if value not in {"not installed", "metadata unavailable", "invalid metadata"}:
1357
+ raise ProtocolViolation(f"invalid {key}")
1358
+ return cast(PackageError, value)
1359
+
1360
+
1361
+ def _require_logical_name(body: Mapping[str, object], key: str) -> str:
1362
+ value = _require_string(body, key, maximum_bytes=MAX_CONTROL_FRAME_BYTES, nonempty=True)
1363
+ try:
1364
+ return validate_logical_name(value)
1365
+ except ValueError as error:
1366
+ raise ProtocolViolation(str(error)) from error
1367
+
1368
+
1369
+ def _require_mime(body: Mapping[str, object], key: str) -> str:
1370
+ value = _require_string(body, key, maximum_bytes=MAX_MIME_BYTES, nonempty=True)
1371
+ try:
1372
+ return validate_mime_type(value)
1373
+ except ValueError as error:
1374
+ raise ProtocolViolation(str(error)) from error
1375
+
1376
+
1377
+ def _require_canonical_base64(body: Mapping[str, object], key: str) -> bytes:
1378
+ value = _require_string(body, key, maximum_bytes=MAX_ARTIFACT_CHUNK_ENCODED_BYTES)
1379
+ try:
1380
+ encoded = value.encode("ascii")
1381
+ decoded = base64.b64decode(encoded, validate=True)
1382
+ except (UnicodeEncodeError, binascii.Error) as error:
1383
+ raise ProtocolViolation(f"{key} must be canonical base64") from error
1384
+ if (
1385
+ not decoded
1386
+ or len(decoded) > MAX_ARTIFACT_CHUNK_BYTES
1387
+ or base64.b64encode(decoded) != encoded
1388
+ ):
1389
+ raise ProtocolViolation(f"{key} must be canonical bounded base64")
1390
+ return decoded
1391
+
1392
+
1393
+ def _require_path(body: Mapping[str, object], key: str, *, canonical: bool) -> str:
1394
+ value = _require_string(body, key, maximum_bytes=MAX_PATH_BYTES, nonempty=True)
1395
+ if not value.startswith("/"):
1396
+ raise ProtocolViolation(f"{key} must be an absolute POSIX path")
1397
+ if canonical and _NORMPATH(value) != value:
1398
+ raise ProtocolViolation(f"{key} must be canonical")
1399
+ return value
1400
+
1401
+
1402
+ def _require_integer(body: Mapping[str, object], key: str, *, minimum: int) -> int:
1403
+ value = body.get(key)
1404
+ if isinstance(value, bool) or not isinstance(value, int):
1405
+ raise ProtocolViolation(f"{key} must be an integer")
1406
+ if value < minimum:
1407
+ raise ProtocolViolation(f"{key} must be at least {minimum}")
1408
+ return value
1409
+
1410
+
1411
+ def _require_boolean(body: Mapping[str, object], key: str) -> bool:
1412
+ value = body.get(key)
1413
+ if not isinstance(value, bool):
1414
+ raise ProtocolViolation(f"{key} must be a boolean")
1415
+ return value
1416
+
1417
+
1418
+ def _require_environment_kind(body: Mapping[str, object], key: str) -> EnvironmentKind:
1419
+ value = _require_string(body, key, maximum_bytes=32)
1420
+ if value not in {"project_venv", "active_venv", "system"}:
1421
+ raise ProtocolViolation(f"invalid {key}: {value}")
1422
+ return cast(EnvironmentKind, value)
1423
+
1424
+
1425
+ def _require_done_outcome(body: Mapping[str, object], key: str) -> DoneOutcome:
1426
+ value = _require_string(body, key, maximum_bytes=32)
1427
+ if value not in {"completed", "failed", "interrupted"}:
1428
+ raise ProtocolViolation(f"invalid {key}: {value}")
1429
+ return cast(DoneOutcome, value)
1430
+
1431
+
1432
+ def _require_protocol_error_code(body: Mapping[str, object], key: str) -> ProtocolErrorCode:
1433
+ value = _require_string(body, key, maximum_bytes=64)
1434
+ allowed = {
1435
+ "version_mismatch",
1436
+ "invalid_sequence",
1437
+ "invalid_request",
1438
+ "runtime_diagnostics_error",
1439
+ "internal_protocol_error",
1440
+ }
1441
+ if value not in allowed:
1442
+ raise ProtocolViolation(f"invalid {key}: {value}")
1443
+ return cast(ProtocolErrorCode, value)
1444
+
1445
+
1446
+ def _valid_correlation_id(value: str) -> bool:
1447
+ return _CORRELATION_ID.fullmatch(value) is not None
1448
+
1449
+
1450
+ def _read_exact(stream: BinaryIO, size: int, *, allow_initial_eof: bool) -> bytes | None:
1451
+ chunks = bytearray()
1452
+ while len(chunks) < size:
1453
+ try:
1454
+ chunk = stream.read(size - len(chunks))
1455
+ except InterruptedError:
1456
+ continue
1457
+ if chunk is None:
1458
+ raise ProtocolViolation("protocol stream returned no bytes")
1459
+ if not chunk:
1460
+ if allow_initial_eof and not chunks:
1461
+ return None
1462
+ raise ProtocolViolation("truncated protocol frame")
1463
+ chunks.extend(chunk)
1464
+ return bytes(chunks)