216 lines
8.1 KiB
Python
216 lines
8.1 KiB
Python
"""Bounded newline-delimited JSON protocol primitives for dictate protocol v2."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import json
|
|
from typing import Any, Mapping
|
|
|
|
PROTOCOL_VERSION = 2
|
|
MAX_FRAME_BYTES = 64 * 1024
|
|
MAX_TRANSCRIPT_BYTES = 32 * 1024
|
|
MAX_ERROR_DETAIL_BYTES = 2 * 1024
|
|
MAX_ID_BYTES = 128
|
|
|
|
|
|
class ProtocolError(ValueError):
|
|
"""A stable protocol error suitable for returning to a client."""
|
|
|
|
def __init__(
|
|
self,
|
|
code: str,
|
|
detail: str,
|
|
*,
|
|
close_connection: bool = False,
|
|
request_id: str = "",
|
|
generation_id: str | None = None,
|
|
) -> None:
|
|
super().__init__(detail)
|
|
self.code = code
|
|
self.detail = detail
|
|
self.close_connection = close_connection
|
|
self.request_id = request_id
|
|
self.generation_id = generation_id
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Request:
|
|
version: int
|
|
type: str
|
|
request_id: str
|
|
generation_id: str | None = None
|
|
|
|
|
|
_REQUEST_FIELDS = {
|
|
"hello": frozenset({"version", "type", "requestId"}),
|
|
"ping": frozenset({"version", "type", "requestId"}),
|
|
"start": frozenset({"version", "type", "requestId", "generationId"}),
|
|
"stop": frozenset({"version", "type", "requestId", "generationId"}),
|
|
"cancel": frozenset({"version", "type", "requestId", "generationId"}),
|
|
}
|
|
|
|
|
|
def _bounded_utf8(value: str, limit: int) -> str:
|
|
# Normalize replacement characters even when the original string contains
|
|
# lone surrogates. This keeps error output serializable and bounded.
|
|
raw = value.encode("utf-8", errors="replace")
|
|
return raw[:limit].decode("utf-8", errors="ignore")
|
|
|
|
|
|
def _validate_id(value: Any, field: str) -> str:
|
|
if not isinstance(value, str) or not value:
|
|
raise ProtocolError("invalid_request", f"{field} must be a non-empty string")
|
|
try:
|
|
encoded = value.encode("utf-8")
|
|
except UnicodeEncodeError as error:
|
|
raise ProtocolError("invalid_request", f"{field} must contain valid Unicode") from error
|
|
if len(encoded) > MAX_ID_BYTES:
|
|
raise ProtocolError("invalid_request", f"{field} exceeds {MAX_ID_BYTES} UTF-8 bytes")
|
|
return value
|
|
|
|
|
|
def _safe_id(value: Any) -> str:
|
|
try:
|
|
return _validate_id(value, "identifier")
|
|
except ProtocolError:
|
|
return ""
|
|
|
|
|
|
def parse_request(frame: Mapping[str, Any]) -> Request:
|
|
"""Validate a decoded request and reject all fields outside its exact schema."""
|
|
|
|
if not isinstance(frame, Mapping):
|
|
raise ProtocolError("invalid_request", "request must be a JSON object")
|
|
|
|
request_id_value = frame.get("requestId")
|
|
request_id = _safe_id(request_id_value)
|
|
version = frame.get("version")
|
|
if not isinstance(version, int) or isinstance(version, bool):
|
|
raise ProtocolError("invalid_request", "version must be an integer", request_id=request_id)
|
|
if version != PROTOCOL_VERSION:
|
|
raise ProtocolError(
|
|
"unsupported_version",
|
|
f"protocol version {version} is unsupported",
|
|
close_connection=True,
|
|
request_id=request_id,
|
|
)
|
|
|
|
request_type = frame.get("type")
|
|
if not isinstance(request_type, str) or request_type not in _REQUEST_FIELDS:
|
|
raise ProtocolError("invalid_request", "unknown request type", request_id=request_id)
|
|
|
|
expected_fields = _REQUEST_FIELDS[request_type]
|
|
try:
|
|
actual_fields = frozenset(frame)
|
|
except TypeError as error:
|
|
raise ProtocolError("invalid_request", "request contains invalid field names", request_id=request_id) from error
|
|
if actual_fields != expected_fields:
|
|
raise ProtocolError(
|
|
"invalid_request",
|
|
f"{request_type} requires exactly: {', '.join(sorted(expected_fields))}",
|
|
request_id=request_id,
|
|
)
|
|
|
|
request_id = _validate_id(request_id_value, "requestId")
|
|
generation_id = None
|
|
if "generationId" in expected_fields:
|
|
generation_id = _validate_id(frame.get("generationId"), "generationId")
|
|
|
|
return Request(version, request_type, request_id, generation_id)
|
|
|
|
|
|
class FrameDecoder:
|
|
"""Incrementally decode bounded newline-delimited JSON objects."""
|
|
|
|
def __init__(self) -> None:
|
|
self._buffer = bytearray()
|
|
self._failed = False
|
|
|
|
def feed(self, data: bytes) -> list[dict[str, Any]]:
|
|
if self._failed:
|
|
raise ProtocolError("connection_closed", "decoder is closed", close_connection=True)
|
|
if not isinstance(data, bytes):
|
|
self._fail("malformed_json", "frame data is invalid")
|
|
frames: list[dict[str, Any]] = []
|
|
remaining = data
|
|
|
|
# Consume each newline-delimited frame before considering the next one.
|
|
# The size limit applies to an individual frame, not to a recv batch:
|
|
# clients may pipeline any number of bounded frames in one feed call.
|
|
while remaining:
|
|
newline = remaining.find(b"\n")
|
|
if newline < 0:
|
|
if len(self._buffer) + len(remaining) > MAX_FRAME_BYTES:
|
|
self._fail("frame_too_large", f"frame exceeds {MAX_FRAME_BYTES} bytes")
|
|
self._buffer.extend(remaining)
|
|
return frames
|
|
if len(self._buffer) + newline > MAX_FRAME_BYTES:
|
|
self._fail("frame_too_large", f"frame exceeds {MAX_FRAME_BYTES} bytes")
|
|
|
|
self._buffer.extend(remaining[:newline])
|
|
raw = bytes(self._buffer)
|
|
self._buffer.clear()
|
|
remaining = remaining[newline + 1 :]
|
|
try:
|
|
value = json.loads(raw.decode("utf-8"))
|
|
except (UnicodeDecodeError, ValueError, RecursionError):
|
|
self._fail("malformed_json", "frame contains malformed JSON")
|
|
if not isinstance(value, dict):
|
|
self._fail("malformed_json", "frame must contain a JSON object")
|
|
frames.append(value)
|
|
|
|
return frames
|
|
|
|
def _fail(self, code: str, detail: str) -> None:
|
|
self._failed = True
|
|
self._buffer.clear()
|
|
raise ProtocolError(code, detail, close_connection=True)
|
|
|
|
|
|
def encode_frame(frame: Mapping[str, Any]) -> bytes:
|
|
"""Encode one bounded service frame, applying the error-detail output cap."""
|
|
|
|
payload = dict(frame)
|
|
if payload.get("type") == "result":
|
|
transcript = payload.get("transcript")
|
|
if not isinstance(transcript, str):
|
|
raise ProtocolError("invalid_event", "result transcript must be a string")
|
|
try:
|
|
if len(transcript.encode("utf-8")) > MAX_TRANSCRIPT_BYTES:
|
|
raise ProtocolError("transcript_too_large", "transcript exceeds output limit")
|
|
except UnicodeEncodeError as error:
|
|
raise ProtocolError("invalid_event", "result transcript contains invalid Unicode") from error
|
|
if payload.get("type") == "error":
|
|
detail = payload.get("detail", "")
|
|
if not isinstance(detail, str):
|
|
detail = str(detail)
|
|
payload["detail"] = _bounded_utf8(detail, MAX_ERROR_DETAIL_BYTES)
|
|
payload["requestId"] = _safe_id(payload.get("requestId"))
|
|
generation_id = _safe_id(payload.get("generationId"))
|
|
if generation_id:
|
|
payload["generationId"] = generation_id
|
|
else:
|
|
payload.pop("generationId", None)
|
|
|
|
try:
|
|
encoded = json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
|
except (TypeError, ValueError, UnicodeError, RecursionError, OverflowError) as error:
|
|
raise ProtocolError("invalid_event", "event contains unsupported values") from error
|
|
if len(encoded) > MAX_FRAME_BYTES:
|
|
raise ProtocolError("frame_too_large", "encoded frame exceeds output limit")
|
|
return encoded + b"\n"
|
|
|
|
|
|
def error_event(error: ProtocolError) -> dict[str, Any]:
|
|
event: dict[str, Any] = {
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "error",
|
|
"requestId": _safe_id(error.request_id),
|
|
"code": _bounded_utf8(error.code if isinstance(error.code, str) else "service_error", MAX_ID_BYTES),
|
|
"detail": _bounded_utf8(error.detail if isinstance(error.detail, str) else str(error.detail), MAX_ERROR_DETAIL_BYTES),
|
|
}
|
|
generation_id = _safe_id(error.generation_id)
|
|
if generation_id:
|
|
event["generationId"] = generation_id
|
|
return event
|