Add streaming support fo JsonRpcServer

Author: RaoulSchaffranek

pyk/src/pyk/rpc/rpc.py

@@ -6,15 +6,15 @@
 from dataclasses import dataclass
 from functools import partial
 from http.server import BaseHTTPRequestHandler, HTTPServer
-from typing import TYPE_CHECKING, Any, Final, NamedTuple
-
+from typing import TYPE_CHECKING, NamedTuple
 from typing_extensions import Protocol
 
 from ..cli.cli import Options
 
 if TYPE_CHECKING:
     from collections.abc import Callable
     from pathlib import Path
+    from typing import Any, Final, Iterator
 
 
 _LOGGER: Final = logging.getLogger(__name__)
@@ -86,7 +86,8 @@ class JsonRpcBatchRequest(NamedTuple):
 class JsonRpcResult(ABC):
 
     @abstractmethod
-    def encode(self) -> bytes: ...
+    def encode(self) -> Iterator[bytes]:
+        ...
 
 
 @dataclass(frozen=True)
@@ -96,7 +97,7 @@ class JsonRpcError(JsonRpcResult):
     message: str
     id: str | int | None
 
-    def to_json(self) -> dict[str, Any]:
+    def wrap_response(self) -> dict[str, Any]:
         return {
             'jsonrpc': JsonRpcServer.JSONRPC_VERSION,
             'error': {
@@ -106,32 +107,40 @@ def to_json(self) -> dict[str, Any]:
             'id': self.id,
         }
 
-    def encode(self) -> bytes:
-        return json.dumps(self.to_json()).encode('ascii')
+    def encode(self) -> Iterator[bytes]:
+        yield json.dumps(self.wrap_response()).encode('ascii')
 
 
 @dataclass(frozen=True)
 class JsonRpcSuccess(JsonRpcResult):
     payload: Any
     id: Any
 
-    def to_json(self) -> dict[str, Any]:
-        return {
-            'jsonrpc': JsonRpcServer.JSONRPC_VERSION,
-            'result': self.payload,
-            'id': self.id,
-        }
-
-    def encode(self) -> bytes:
-        return json.dumps(self.to_json()).encode('ascii')
+    def encode(self) -> Iterator[bytes]:
+        yield f'{{"jsonrpc":"2.0", "id": {self.id}, "result": '.encode('ascii')
+        if isinstance(self.payload, Iterator):
+            for chunk in self.payload:
+                yield chunk.encode('ascii')
+        else:
+            yield json.dumps(self.payload).encode('ascii')
+        yield b'}'
 
 
 @dataclass(frozen=True)
 class JsonRpcBatchResult(JsonRpcResult):
     results: tuple[JsonRpcError | JsonRpcSuccess, ...]
 
-    def encode(self) -> bytes:
-        return json.dumps([result.to_json() for result in self.results]).encode('ascii')
+    def encode(self) -> Iterator[bytes]:
+        yield b'['
+        first = True
+        for result in self.results:
+            if not first:
+                yield b','
+            else:
+                first = False
+            for chunk in result.encode():
+                yield chunk
+        yield b']'
 
 
 class JsonRpcRequestHandler(BaseHTTPRequestHandler):
@@ -143,8 +152,10 @@ def __init__(self, methods: dict[str, JsonRpcMethod], *args: Any, **kwargs: Any)
 
     def _send_response(self, response: JsonRpcResult) -> None:
         self.send_response_headers()
-        response_bytes = response.encode()
-        self.wfile.write(response_bytes)
+        response_body = response.encode()
+        for chunk in response_body:
+            self.wfile.write(chunk)
+        self.wfile.flush()
 
     def send_response_headers(self) -> None:
         self.send_response(200)

pyk/src/tests/integration/test_json_rpc.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+from collections.abc import Iterator
 import json
 from http.client import HTTPConnection
 from threading import Thread
 from time import sleep
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Iterator
 
 from pyk.cterm import CTerm
 from pyk.kast.inner import KApply, KSequence, KSort, KToken
@@ -154,6 +155,7 @@ def __init__(self, options: ServeRpcOptions) -> None:
         self.register_method('set_x', self.exec_set_x)
         self.register_method('set_y', self.exec_set_y)
         self.register_method('add', self.exec_add)
+        self.register_method('streaming', self.exec_streaming)
 
     def exec_get_x(self) -> int:
         return self.x
@@ -169,6 +171,11 @@ def exec_set_y(self, n: int) -> None:
 
     def exec_add(self) -> int:
         return self.x + self.y
+    
+    def exec_streaming(self) -> Iterator[bytes]:
+        yield b'{'
+        yield b'"foo": "bar"'
+        yield b'}'
 
 
 class TestJsonRPCServer(KRunTest):
@@ -221,6 +228,9 @@ def wait_until_ready() -> None:
         assert len(res) == 3
         assert res[2]['result'] == 1 + 2
 
+        res = rpc_client.request('streaming', [])
+        assert res == {'foo': 'bar'}
+
         server.shutdown()
         thread.join()