"""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