Added tests for TCPServer and improved error handling for JsonRpcHandler and FileHandler

Author: m-messer

lf_toolkit/io/file_server.py

@@ -6,6 +6,7 @@
 
 from .base_server import BaseServer
 from .handler import Handler
+from .handler import InvalidSubmissionError
 
 
 class FileHandler(Handler):
@@ -16,8 +17,13 @@ async def dispatch(self, req: str) -> str:
         try:
             result = await self.handle(command, request)
             response = {"command": command, "result": result}
+        except InvalidSubmissionError as e:
+            response = {
+                "command": command,
+                "error": {"message": str(e), "code": "INVALID_SUBMISSION"},
+            }
         except Exception as e:
-            response = {"command": command, "error": str(e)}
+            response = {"command": command, "error": {"message": str(e)}}
         return ujson.dumps(response)
 
 

lf_toolkit/io/handler.py

@@ -14,6 +14,13 @@
 from ..shared import Command
 from ..shared import Params
 
+# JSON-RPC error code for submissions the user handler cannot process.
+INVALID_SUBMISSION_CODE = 422
+
+
+class InvalidSubmissionError(Exception):
+    """Raised when the user handler cannot process the submitted response."""
+
 
 class Handler(ABC):
 
@@ -40,6 +47,8 @@ async def _call_user_handler(self, req: str, *args, **kwargs):
                 return await handler(*args, **kwargs)
             else:
                 return await anyio.to_thread.run_sync(handler, *args, **kwargs)
+        except ValueError as e:
+            raise InvalidSubmissionError(str(e)) from e
         except Exception as e:
             raise ValueError(f"Error calling user handler for '{req}': {e}")
 

lf_toolkit/io/rpc_handler.py

@@ -5,7 +5,9 @@
 from jsonrpcserver import async_dispatch
 
 from ..shared import Command
+from .handler import INVALID_SUBMISSION_CODE
 from .handler import Handler
+from .handler import InvalidSubmissionError
 
 
 class JsonRpcHandler(Handler):
@@ -31,6 +33,8 @@ async def wrapped(req: dict):
         try:
             result = await handler.handle(name, {"params": req})
             return Success(result)
+        except InvalidSubmissionError as e:
+            return Error(INVALID_SUBMISSION_CODE, str(e))
         except Exception as e:
             # Pass only the message: the exception object is not JSON
             # serializable, so sending it as `data` makes serializing the

tests/io/file_server_test.py

@@ -81,7 +81,7 @@ async def test_dispatch_exception_handling(self, handler):
 
         assert response["command"] == "bad_command"
         assert "error" in response
-        assert response["error"] == "Invalid command"
+        assert response["error"] == {"message": "Invalid command"}
         assert "result" not in response
 
     @pytest.mark.asyncio
@@ -195,7 +195,7 @@ async def test_run_with_handler_error(self, temp_files, tmp_path):
 
         assert response["command"] == "failing_command"
         assert "error" in response
-        assert response["error"] == "Processing failed"
+        assert response["error"] == {"message": "Processing failed"}
 
     @pytest.mark.asyncio
     async def test_run_missing_request_file(self, temp_files):
@@ -273,3 +273,37 @@ async def test_run_with_complex_nested_data(self, temp_files, tmp_path):
 
         assert response["command"] == "eval"
         assert response["result"] == {"processed": True}
+
+
+class TestFileHandlerUserErrors:
+    """Tests for errors raised by registered user handlers"""
+
+    @pytest.mark.asyncio
+    async def test_value_error_returns_invalid_submission(self):
+        def preview_fn(response, params):
+            raise ValueError(f"Failed to parse SymPy expression: {response}")
+
+        handler = FileHandler()
+        handler.register("preview", preview_fn)
+
+        request = ujson.dumps({"command": "preview", "params": {"response": "A/(w*"}})
+        response = ujson.loads(await handler.dispatch(request))
+
+        assert response["error"] == {
+            "message": "Failed to parse SymPy expression: A/(w*",
+            "code": "INVALID_SUBMISSION",
+        }
+
+    @pytest.mark.asyncio
+    async def test_unexpected_error_has_no_code(self):
+        def preview_fn(response, params):
+            raise RuntimeError("boom")
+
+        handler = FileHandler()
+        handler.register("preview", preview_fn)
+
+        request = ujson.dumps({"command": "preview", "params": {"response": "x"}})
+        response = ujson.loads(await handler.dispatch(request))
+
+        assert "code" not in response["error"]
+        assert "boom" in response["error"]["message"]

tests/io/rpc_handler_test.py

@@ -69,3 +69,43 @@ def chat_health_fn() -> ChatHealthResponse:
 
         assert response["result"]["status"] == "OK"
         assert response["result"]["capabilities"]["supportsChat"] is True
+
+
+class TestJsonRpcHandlerErrors:
+
+    @pytest.fixture
+    def handler(self):
+        return JsonRpcHandler()
+
+    @staticmethod
+    def _eval_request(response):
+        return ujson.dumps({
+            "jsonrpc": "2.0",
+            "method": "eval",
+            "params": [{"response": response, "answer": "x", "params": {}}],
+            "id": 1,
+        })
+
+    @pytest.mark.asyncio
+    async def test_value_error_returns_invalid_submission_code(self, handler):
+        def eval_fn(response, answer, params):
+            raise ValueError(f"Failed to parse SymPy expression: {response}")
+
+        handler.register("eval", eval_fn)
+
+        response = ujson.loads(await handler.dispatch(self._eval_request("A/(w*")))
+
+        assert response["error"]["code"] == 422
+        assert response["error"]["message"] == "Failed to parse SymPy expression: A/(w*"
+
+    @pytest.mark.asyncio
+    async def test_unexpected_error_returns_generic_code(self, handler):
+        def eval_fn(response, answer, params):
+            raise RuntimeError("boom")
+
+        handler.register("eval", eval_fn)
+
+        response = ujson.loads(await handler.dispatch(self._eval_request("x+1")))
+
+        assert response["error"]["code"] == 0
+        assert "boom" in response["error"]["message"]

tests/io/tcp_server_test.py

@@ -0,0 +1,140 @@
+import socket
+
+import anyio
+import pytest
+import ujson
+
+from lf_toolkit.io.stream_io import StreamIO
+from lf_toolkit.io.tcp_server import TCPServer
+
+
+def free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+@pytest.fixture
+def anyio_backend():
+    return "asyncio"
+
+
+def make_newline_message(payload: str) -> bytes:
+    return payload.encode("utf-8") + b"\n"
+
+
+class FakeStreamIO(StreamIO):
+
+    def __init__(self):
+        self._buffer = b""
+        self.responses = []
+        self.close_count = 0
+
+    def feed(self, data: bytes):
+        self._buffer += data
+
+    async def read(self, size: int) -> bytes:
+        if not self._buffer:
+            raise anyio.EndOfStream()
+        chunk = self._buffer[:size]
+        self._buffer = self._buffer[size:]
+        return chunk
+
+    async def write(self, data: bytes):
+        self.responses.append(data)
+
+    async def close(self):
+        self.close_count += 1
+
+
+class TestTCPServer:
+
+    @pytest.fixture
+    def stream(self):
+        return FakeStreamIO()
+
+    @pytest.fixture
+    def server(self):
+        return TCPServer()
+
+    @pytest.mark.anyio
+    async def test_handles_multiple_messages_newline_framed(self, stream, server):
+        server.eval(lambda response, answer, params: {"received": True})
+
+        stream.feed(
+            make_newline_message(
+                ujson.dumps(
+                    {"jsonrpc": "2.0", "method": "eval", "params": [{}], "id": 1}
+                )
+            )
+        )
+        stream.feed(
+            make_newline_message(
+                ujson.dumps(
+                    {"jsonrpc": "2.0", "method": "eval", "params": [{}], "id": 2}
+                )
+            )
+        )
+
+        await server._handle_client(stream)
+
+        assert len(stream.responses) == 2
+        for response in stream.responses:
+            assert response.endswith(b"\n")
+            assert b"Content-Length:" not in response
+
+    @pytest.mark.anyio
+    async def test_closes_connection_once(self, stream, server):
+        stream.feed(
+            make_newline_message(
+                ujson.dumps(
+                    {"jsonrpc": "2.0", "method": "eval", "params": [{}], "id": 1}
+                )
+            )
+        )
+
+        await server._handle_client(stream)
+
+        assert stream.close_count == 1
+
+
+class TestTCPServerLive:
+
+    @pytest.mark.anyio
+    async def test_end_to_end_roundtrip(self):
+        port = free_port()
+        server = TCPServer(address=f"127.0.0.1:{port}")
+        server.eval(lambda response, answer, params: {"is_correct": response == answer})
+
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(server.run)
+
+            request = ujson.dumps(
+                {
+                    "jsonrpc": "2.0",
+                    "method": "eval",
+                    "params": [{"response": "42", "answer": "42"}],
+                    "id": 1,
+                }
+            )
+
+            client = None
+            for _ in range(50):
+                try:
+                    client = await anyio.connect_tcp("127.0.0.1", port)
+                    break
+                except OSError:
+                    await anyio.sleep(0.05)
+            assert client is not None, "TCPServer never started listening"
+
+            await client.send(request.encode("utf-8") + b"\n")
+
+            data = b""
+            while not data.endswith(b"\n"):
+                data += await client.receive(4096)
+            await client.aclose()
+
+            response = ujson.loads(data.decode("utf-8"))
+            assert response["result"]["is_correct"] is True
+
+            tg.cancel_scope.cancel()