agent_search_gateway.protocol

  1"""Migration-stable newline-delimited JSON socket protocol."""
  2
  3import asyncio
  4import json
  5from collections.abc import Callable, Mapping
  6from contextlib import suppress
  7from pathlib import Path
  8from typing import TypeAlias
  9
 10from .errors import DaemonUnavailable, ErrorCode, ProtocolFailure
 11from .models import (
 12    ErrorResponse,
 13    KeywordSearchRequest,
 14    LLMSearchRequest,
 15    PaperSearchRequest,
 16    Request,
 17    Response,
 18    ShutdownRequest,
 19    SuccessResponse,
 20    URLFetchRequest,
 21)
 22
 23DecodedRequest: TypeAlias = Request | ErrorResponse
 24RequestParser = Callable[[Mapping[str, object]], Request]
 25
 26_MAX_REQUEST_FRAME_BYTES = 1 << 20
 27_MAX_RESPONSE_FRAME_BYTES = 8 << 20
 28_RESPONSE_TIMEOUT_SECONDS = 300.0
 29
 30
 31def _bad_request(message: str) -> ErrorResponse:
 32    return ErrorResponse(ErrorCode.BAD_REQUEST, message)
 33
 34
 35def _require_exact_keys(payload: Mapping[str, object], expected: set[str]) -> None:
 36    if set(payload) != expected:
 37        raise ValueError("request fields do not match schema")
 38
 39
 40def _parse_keyword(payload: Mapping[str, object]) -> Request:
 41    _require_exact_keys(payload, {"type", "query"})
 42    query = payload["query"]
 43    if not isinstance(query, str):
 44        raise ValueError("query must be a string")
 45    return KeywordSearchRequest(query)
 46
 47
 48def _parse_paper(payload: Mapping[str, object]) -> Request:
 49    _require_exact_keys(payload, {"type", "query"})
 50    query = payload["query"]
 51    if not isinstance(query, str):
 52        raise ValueError("query must be a string")
 53    return PaperSearchRequest(query)
 54
 55
 56def _parse_llm(payload: Mapping[str, object]) -> Request:
 57    keys = set(payload)
 58    if keys not in ({"type", "prompt"}, {"type", "prompt", "scope"}):
 59        raise ValueError("request fields do not match schema")
 60    prompt = payload["prompt"]
 61    if not isinstance(prompt, str):
 62        raise ValueError("prompt must be a string")
 63    scope = payload.get("scope", "web")
 64    if scope == "web":
 65        return LLMSearchRequest(prompt, "web")
 66    if scope == "paper":
 67        return LLMSearchRequest(prompt, "paper")
 68    if scope == "all":
 69        return LLMSearchRequest(prompt, "all")
 70    raise ValueError("scope must be web, paper, or all")
 71
 72
 73def _parse_fetch(payload: Mapping[str, object]) -> Request:
 74    _require_exact_keys(payload, {"type", "url", "focus"})
 75    url = payload["url"]
 76    focus = payload["focus"]
 77    if not isinstance(url, str):
 78        raise ValueError("url must be a string")
 79    if focus is not None and not isinstance(focus, str):
 80        raise ValueError("focus must be a string or null")
 81    return URLFetchRequest(url, focus)
 82
 83
 84def _parse_shutdown(payload: Mapping[str, object]) -> Request:
 85    _require_exact_keys(payload, {"type"})
 86    return ShutdownRequest()
 87
 88
 89_REQUEST_PARSERS: dict[str, RequestParser] = {
 90    "keyword_search": _parse_keyword,
 91    "paper_search": _parse_paper,
 92    "llm_search": _parse_llm,
 93    "url_fetch": _parse_fetch,
 94    "shutdown": _parse_shutdown,
 95}
 96
 97
 98def decode_request_frame(frame: bytes) -> DecodedRequest:
 99    try:
100        text = frame.decode("utf-8")
101        payload = json.loads(text)
102    except (UnicodeDecodeError, json.JSONDecodeError):
103        return _bad_request("Request must be valid UTF-8 JSON")
104    if not isinstance(payload, dict):
105        return _bad_request("Request must be a JSON object")
106    request_type = payload.get("type")
107    if not isinstance(request_type, str):
108        return _bad_request("Request type must be a string")
109    parser = _REQUEST_PARSERS.get(request_type)
110    if parser is None:
111        return _bad_request("Unknown request type")
112    try:
113        return parser(payload)
114    except (KeyError, ValueError):
115        return _bad_request("Request fields do not match schema")
116
117
118class NDJSONDecoder:
119    def __init__(self) -> None:
120        self._buffer = bytearray()
121        self._discarding_oversized_frame = False
122
123    def feed(self, data: bytes) -> list[DecodedRequest]:
124        decoded: list[DecodedRequest] = []
125        remaining = data
126        while remaining:
127            if self._discarding_oversized_frame:
128                boundary = remaining.find(b"\n")
129                if boundary < 0:
130                    return decoded
131                self._discarding_oversized_frame = False
132                remaining = remaining[boundary + 1 :]
133                continue
134
135            boundary = remaining.find(b"\n")
136            if boundary < 0:
137                capacity = _MAX_REQUEST_FRAME_BYTES - len(self._buffer)
138                if len(remaining) > capacity:
139                    self._buffer.clear()
140                    self._discarding_oversized_frame = True
141                    decoded.append(_bad_request("Request frame is too large"))
142                    return decoded
143                self._buffer.extend(remaining)
144                return decoded
145
146            if len(self._buffer) + boundary > _MAX_REQUEST_FRAME_BYTES:
147                self._buffer.clear()
148                decoded.append(_bad_request("Request frame is too large"))
149                remaining = remaining[boundary + 1 :]
150                continue
151
152            self._buffer.extend(remaining[:boundary])
153            frame = bytes(self._buffer)
154            self._buffer.clear()
155            decoded.append(decode_request_frame(frame))
156            remaining = remaining[boundary + 1 :]
157        return decoded
158
159
160def _request_payload(request: Request) -> dict[str, object]:
161    if isinstance(request, KeywordSearchRequest):
162        return {"type": "keyword_search", "query": request.query}
163    if isinstance(request, PaperSearchRequest):
164        return {"type": "paper_search", "query": request.query}
165    if isinstance(request, LLMSearchRequest):
166        payload: dict[str, object] = {"type": "llm_search", "prompt": request.prompt}
167        if request.scope != "web":
168            payload["scope"] = request.scope
169        return payload
170    if isinstance(request, URLFetchRequest):
171        return {"type": "url_fetch", "url": request.url, "focus": request.focus}
172    return {"type": "shutdown"}
173
174
175def _response_payload(response: Response) -> dict[str, object]:
176    if isinstance(response, SuccessResponse):
177        return {"ok": True, "text": response.text}
178    return {"ok": False, "error": response.error.value, "message": response.message}
179
180
181def _encode(payload: Mapping[str, object]) -> bytes:
182    return json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + b"\n"
183
184
185def encode_request(request: Request) -> bytes:
186    return _encode(_request_payload(request))
187
188
189def encode_response(response: Response) -> bytes:
190    return _encode(_response_payload(response))
191
192
193def parse_response_frame(frame: bytes) -> Response:
194    try:
195        text = frame.decode("utf-8")
196        payload = json.loads(text)
197    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
198        raise _protocol_error("Response must be valid UTF-8 JSON") from exc
199    if not isinstance(payload, dict) or not isinstance(payload.get("ok"), bool):
200        raise _protocol_error("Response must contain boolean ok")
201    if payload["ok"] is True:
202        if set(payload) != {"ok", "text"} or not isinstance(payload.get("text"), str):
203            raise _protocol_error("Success response fields do not match schema")
204        return SuccessResponse(payload["text"])
205    if set(payload) != {"ok", "error", "message"}:
206        raise _protocol_error("Error response fields do not match schema")
207    error = payload.get("error")
208    message = payload.get("message")
209    if not isinstance(error, str) or not isinstance(message, str):
210        raise _protocol_error("Error response fields have invalid types")
211    try:
212        code = ErrorCode(error)
213    except ValueError as exc:
214        raise _protocol_error("Error response contains unknown error code") from exc
215    return ErrorResponse(code, message)
216
217
218def _protocol_error(message: str) -> ProtocolFailure:
219    return ProtocolFailure(ErrorCode.PROTOCOL_ERROR, message)
220
221
222async def send_request(
223    socket_path: Path,
224    request: Request,
225    *,
226    response_timeout_seconds: float = _RESPONSE_TIMEOUT_SECONDS,
227) -> Response:
228    try:
229        reader, writer = await asyncio.open_unix_connection(
230            path=socket_path,
231            limit=_MAX_RESPONSE_FRAME_BYTES,
232        )
233    except (FileNotFoundError, ConnectionRefusedError) as exc:
234        raise DaemonUnavailable(str(socket_path)) from exc
235
236    try:
237        writer.write(encode_request(request))
238        await writer.drain()
239        try:
240            line = await asyncio.wait_for(reader.readline(), timeout=response_timeout_seconds)
241        except TimeoutError as exc:
242            raise _protocol_error("Timed out waiting for daemon response") from exc
243        except ValueError as exc:
244            raise _protocol_error("Daemon response frame is too large") from exc
245        if not line or not line.endswith(b"\n"):
246            raise _protocol_error("Daemon response ended before newline")
247        return parse_response_frame(line[:-1])
248    finally:
249        writer.close()
250        with suppress(ConnectionError, BrokenPipeError):
251            await writer.wait_closed()