from __future__ import annotations import asyncio import unittest from unittest.mock import patch from klbr.remote import ( CancellationState, RemoteError, RemoteUnknown, RemoteValueError, Unsupported, connect, MAX_VALUE_BYTES, _wire_size, ) from klbr_runtime.context import ExecutionScope, current from klbr.rpath import RemotePathBindingError from klbr_runtime.remote_dispatch import MAX_VALUE_BYTES, RemoteDispatchError, RemoteDispatcher class DispatcherTransport: """Shuttle the SDK's opaque operations into a real target dispatcher.""" def __init__(self, dispatcher): self.dispatcher = dispatcher async def request(self, operation, envelope): value = await self.dispatcher.dispatch(operation, envelope["payload"], envelope["request_id"]) return {"status": "ok", "value": value} async def cancel(self, operation_id): return await self.dispatcher.cancel(operation_id) async def close(self): return None class RecordingTransport: def __init__(self): self.requests: list[tuple[str, dict]] = [] self.closed = False self.release = asyncio.Event() async def request(self, operation, envelope): self.requests.append((operation, envelope)) if operation == "path.exists": return {"status": "ok", "value": True} if operation == "path.read_bytes": return {"status": "ok", "value": {"__klbr_bytes__": "aGk="}} if operation == "execute": return {"status": "ok", "value": None} if operation == "slow": await self.release.wait() return {"status": "ok", "value": {"seen": operation}} async def cancel(self, operation_id): return {"state": "requested", "outstanding": operation_id} async def close(self): self.closed = True class RemoteSdkTests(unittest.IsolatedAsyncioTestCase): async def test_host_scope_is_the_default_connection(self): seen = [] class Bridge: async def request(self, execution, request): seen.append(request) if request["method"] == "remote.ensure": return { "target": "linux-build", "workspace": "/srv/work", "namespace_id": "ns-1", "namespace_generation": "gen-1", "persistence": "connection-scoped", "identity": {"machine": "builder", "platform": "Linux", "arch": "x86_64"}, "capabilities": [{"capability": "file_system"}, {"capability": "shell"}], } return {"status": "ok", "value": True} execution = ExecutionScope("e-1", Bridge()) token = current.set(execution) try: remote = await connect("linux-build") self.assertEqual(remote.persistence, "connection-scoped") self.assertTrue(await remote.pathlib.Path("/target-only").exists()) self.assertEqual(seen[0]["method"], "remote.ensure") self.assertEqual(seen[1]["method"], "remote.call") self.assertEqual(seen[1]["namespace_id"], "ns-1") await remote.close() finally: current.reset(token) async def test_remote_path_is_bound_and_local_fspath_refuses(self): transport = RecordingTransport() remote = await connect("target-a", transport=transport, capabilities={"FileSystem": "r1"}) path = remote.pathlib.Path("/target") / "file.txt" self.assertEqual(str(path), "/target/file.txt") with self.assertRaises(TypeError): import os os.fspath(path) self.assertEqual(await path.read_bytes(), b"hi") other = await connect("target-b", transport=RecordingTransport(), capabilities={"FileSystem": "r1"}) with self.assertRaises(RemotePathBindingError) as caught: path.joinpath(other.pathlib.Path("x")) self.assertIn("different target", str(caught.exception)) await remote.close() await other.close() async def test_capability_reflection_and_typed_unsupported(self): remote = await connect("target", transport=RecordingTransport(), capabilities=[{"capability": "closures", "dill_available": False}]) self.assertFalse(remote.control.capabilities["closures"].usable) with self.assertRaises(Unsupported): remote.control.ship(lambda: 1, trusted=True) with self.assertRaises(AttributeError): _ = remote.not_installed await remote.close() async def test_closure_serialization_failure_is_not_source_fallback(self): remote = await connect("target", transport=RecordingTransport(), capabilities=[{"capability": "closures", "dill_available": True}]) with patch("importlib.util.find_spec", return_value=object()): with self.assertRaises(Unsupported): remote.control.ship(lambda: 1, trusted=True) await remote.close() async def test_completed_none_is_not_running(self): remote = await connect("target", transport=RecordingTransport(), capabilities={}) handle = remote.control.execute("pass") self.assertTrue(handle.running) self.assertIsNone(await handle) self.assertTrue(handle.completed) self.assertIsNone(handle.poll()) await remote.close() async def test_timeout_is_typed_unknown_not_retryable_failure(self): transport = RecordingTransport() remote = await connect("target", transport=transport, capabilities={}, call_timeout=0.01) handle = remote._start("slow", {}) with self.assertRaises(RemoteUnknown) as caught: await handle self.assertFalse(caught.exception.can_replay) self.assertIn("deadline", str(caught.exception)) await remote.close() async def test_oversized_error_is_typed(self): class ErrorTransport(RecordingTransport): async def request(self, operation, envelope): return {"status": "error", "error": {"message": "x" * 16_385}} remote = await connect("target", transport=ErrorTransport(), capabilities={}) with self.assertRaises(RemoteValueError): await remote._request("opaque", {}) await remote.close() async def test_encoded_payload_limit_is_typed(self): remote = await connect("target", transport=RecordingTransport(), capabilities={}) length = MAX_VALUE_BYTES while _wire_size({"blob": "x" * length}) > MAX_VALUE_BYTES: length -= 1 self.assertLessEqual(_wire_size({"blob": "x" * length}), MAX_VALUE_BYTES) await remote._request("opaque", {"blob": "x" * length}) with self.assertRaises(RemoteValueError): await remote._request("opaque", {"blob": "x" * (length + 1)}) await remote.close() async def test_sdk_and_target_dispatcher_share_the_production_vocabulary(self): import tempfile from pathlib import Path with tempfile.TemporaryDirectory() as root, tempfile.TemporaryDirectory() as local_root: dispatcher = RemoteDispatcher(root) local_marker = Path(local_root) / "marker.txt" local_marker.write_text("local") class DispatcherTransport: async def request(self, operation, envelope): value = await dispatcher.dispatch(operation, envelope["payload"], envelope["request_id"]) return {"status": "ok", "value": value} async def cancel(self, operation_id): return await dispatcher.cancel(operation_id) async def close(self): return None remote = await connect("target", transport=DispatcherTransport(), capabilities=dispatcher.capabilities) target_path = remote.pathlib.Path(root) / "sdk-target.txt" self.assertEqual(await target_path.write_text("only remote"), len("only remote")) self.assertEqual(await target_path.read_text(), "only remote") self.assertEqual(local_marker.read_text(), "local") await remote.close() async def test_tree_verbs_reach_the_target_with_the_same_shapes(self): import tempfile from pathlib import Path from klbr.workspace import Entry, Listing, RunResult, SearchResult with tempfile.TemporaryDirectory() as root: tree = Path(root) (tree / "src").mkdir() (tree / "src" / "main.py").write_text("print('needle')\n") (tree / "build" / "deep").mkdir(parents=True) (tree / "build" / "deep" / "artifact.o").write_bytes(b"x") remote = await connect("target", transport=DispatcherTransport(RemoteDispatcher(root)), capabilities=RemoteDispatcher(root).capabilities) listing = await remote.ls("src") self.assertIsInstance(listing, Listing) self.assertTrue(all(isinstance(entry, Entry) for entry in listing.entries)) self.assertEqual([entry.relative for entry in listing.entries], ["src/main.py"]) found = await remote.search("needle", ".") self.assertIsInstance(found, SearchResult) self.assertEqual([(match.relative, match.line) for match in found.matches], [("src/main.py", 1)]) file = await remote.open_file("src/main.py") self.assertEqual(await file.text(), "print('needle')\n") await file.write("print('rewritten')\n") await file.replace("rewritten", "replaced") self.assertEqual((tree / "src" / "main.py").read_text(), "print('replaced')\n") self.assertTrue(await file.exists()) ran = await remote.run(("true",)) self.assertIsInstance(ran, RunResult) self.assertTrue(ran.completed) self.assertEqual(ran.exit_code, 0) # A deadline that expires is reported, not raised: the caller learns it timed out. late = await remote.run(("sleep", "5"), timeout=0.05) self.assertTrue(late.timed_out) self.assertFalse(late.completed) # The target's own refusal travels: it is the same rule, not a second copy of it. # The target's refusal is the same rule, and it arrives with its own wording. with self.assertRaises(RemoteError) as refusal: await remote.delete("build") self.assertIn("recursive", str(refusal.exception)) self.assertTrue((tree / "build" / "deep" / "artifact.o").exists()) await remote.delete("build", recursive=True) self.assertFalse((tree / "build").exists()) await remote.close() async def test_target_dispatcher_persists_namespace_and_remote_filesystem(self): import tempfile from pathlib import Path with tempfile.TemporaryDirectory() as root: dispatcher = RemoteDispatcher(root) path = str(Path(root) / "target.txt") await dispatcher.dispatch("path.write_text", {"path": path, "data": "remote", "encoding": "utf-8", "errors": "strict"}) self.assertEqual(await dispatcher.dispatch("path.read_text", {"path": path, "encoding": "utf-8", "errors": "strict"}), "remote") bytes_path = str(Path(root) / "bytes.bin") await dispatcher.dispatch("path.write_bytes", {"path": bytes_path, "data": {"__klbr_bytes__": "AAEC"}}) self.assertEqual(await dispatcher.dispatch("path.read_bytes", {"path": bytes_path}), {"__klbr_bytes__": "AAEC"}) await dispatcher.dispatch("execute", {"code": "answer = 42"}) self.assertEqual(await dispatcher.dispatch("eval", {"expression": "answer"}), 42) await dispatcher.dispatch("os.environ.set", {"key": "KLBR_REMOTE_TEST", "value": "target"}) self.assertEqual(await dispatcher.dispatch("os.environ.get", {"key": "KLBR_REMOTE_TEST"}), "target") subdir = Path(root) / "subdir" subdir.mkdir() await dispatcher.dispatch("os.chdir", {"path": str(subdir)}) self.assertEqual(await dispatcher.dispatch("os.getcwd", {}), str(subdir.resolve())) await dispatcher.dispatch("path.write_text", {"path": "relative.txt", "data": "cwd", "encoding": "utf-8", "errors": "strict"}) self.assertEqual((subdir / "relative.txt").read_text(), "cwd") with self.assertRaises(RemoteDispatchError) as escaped: await dispatcher.dispatch("path.exists", {"path": str(Path(root).parent / "outside")}) self.assertEqual(escaped.exception.code, "path_escape") async def test_target_dispatcher_oversized_error_is_typed(self): dispatcher = RemoteDispatcher() with self.assertRaises(RemoteDispatchError) as oversized: await dispatcher.dispatch("execute", {"code": "raise Exception(" + repr("x" * 20_000) + ")"}) self.assertEqual(oversized.exception.code, "error_too_large") async def test_target_dispatcher_closure_requires_dill_mode(self): dispatcher = RemoteDispatcher() with self.assertRaises(RemoteDispatchError) as unsupported: await dispatcher.dispatch("ship", {"trusted": True, "serialization": "dill", "python_only": True, "function": "eA==", "args": "eA==", "kwargs": "eA=="}) self.assertIn(unsupported.exception.code, {"unsupported", "invalid_payload"}) async def test_target_dispatcher_result_limit_boundary(self): from klbr_runtime.remote_dispatch import _wire_encode dispatcher = RemoteDispatcher() length = MAX_VALUE_BYTES while len(__import__("json").dumps(_wire_encode("x" * length), ensure_ascii=False).encode()) > MAX_VALUE_BYTES: length -= 1 value = await dispatcher.dispatch("eval", {"expression": repr("x") + " * " + str(length)}) self.assertEqual(len(value), length) with self.assertRaises(RemoteDispatchError) as exceeded: await dispatcher.dispatch("eval", {"expression": repr("x") + " * " + str(length + 1)}) self.assertEqual(exceeded.exception.code, "result_too_large") async def test_target_dispatcher_requires_checked_workspace_for_filesystem(self): dispatcher = RemoteDispatcher() with self.assertRaises(RemoteDispatchError) as missing: await dispatcher.dispatch("path.exists", {"path": "."}) self.assertEqual(missing.exception.code, "workspace_unbound") async def test_target_dispatcher_rejects_unbounded_result_and_unknown_operation(self): dispatcher = RemoteDispatcher() with self.assertRaises(RemoteDispatchError) as too_large: await dispatcher.dispatch("eval", {"expression": repr("x") + " * " + str(MAX_VALUE_BYTES)}) self.assertEqual(too_large.exception.code, "result_too_large") with self.assertRaises(RemoteDispatchError) as unknown: await dispatcher.dispatch("not_an_operation", {}) self.assertEqual(unknown.exception.code, "unsupported") if __name__ == "__main__": unittest.main()