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()