"""Test-only protocol peer. This is NOT a Python replacement for the Rust host.
It exercises the shipped subprocesses without Rust/provider availability. Rust
storage, worker supervision and SessionRuntime are tested separately by Cargo.
"""
from __future__ import annotations
import asyncio
import contextlib
import json
import os
import shutil
import sys
import tempfile
import uuid
from pathlib import Path
from typing import Awaitable, Callable
from klbr_runtime.release import fingerprint
from klbr_runtime.wire import Wire
ROOT = Path(__file__).resolve().parents[3]
RUNNER = ROOT / "python/klbr-runtime/src/klbr_runtime/__main__.py"
Service = Callable[[dict], Awaitable[dict]]
def behavior_copy(target: Path, delivery: str | None = None) -> Path:
shutil.copytree(ROOT / "behavior", target)
if delivery is not None:
with (target / "klbr_hooks/defaults.py").open("a") as out:
out.write("\n" + delivery + "\n")
return target
class Peer:
def __init__(self, role: str = "workbench", behavior: Path | None = None,
skills: Path | None = None):
self.role, self.behavior, self.skills = role, behavior, skills
self.session, self.generation = "test-session", str(uuid.uuid4())
self.token = str(uuid.uuid4())
self.directory = tempfile.TemporaryDirectory(prefix="klbr-test-")
self.process = None
self.wire = None
self.server = None
self.drains = []
self.raw = {"stdout": bytearray(), "stderr": bytearray()}
self.requests = []
async def start(self) -> Peer:
loop = asyncio.get_running_loop()
accepted = loop.create_future()
def accept(reader, writer):
if accepted.done():
writer.close()
else:
accepted.set_result((reader, writer))
socket_path = str(Path(self.directory.name) / "control.sock")
self.server = await asyncio.start_unix_server(accept, socket_path)
args = [sys.executable, "-I", "-B", "-u", str(RUNNER), "--socket", socket_path,
"--role", self.role, "--session", self.session, "--generation", self.generation,
"--token", self.token]
revision = None
if self.behavior is not None:
revision = fingerprint(self.behavior)
args += ["--behavior", str(self.behavior), "--revision", revision]
skills_revision = None
if self.skills is not None:
skills_revision = json.loads(Path(self.skills).read_text())["revision"]
args += ["--skills", str(self.skills)]
environment = {key: os.environ[key] for key in
("PATH", "HOME", "LANG", "LC_ALL", "TMPDIR", "LD_LIBRARY_PATH", "DYLD_LIBRARY_PATH")
if key in os.environ}
self.process = await asyncio.create_subprocess_exec(*args, env=environment, stdin=asyncio.subprocess.DEVNULL,
stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
async def drain(name, stream):
while chunk := await stream.read(4096):
self.raw[name] += chunk
del self.raw[name][:-8192]
self.drains = [asyncio.create_task(drain("stdout", self.process.stdout)),
asyncio.create_task(drain("stderr", self.process.stderr))]
exited = asyncio.create_task(self.process.wait())
try:
done, _ = await asyncio.wait([accepted, exited], timeout=3, return_when=asyncio.FIRST_COMPLETED)
if accepted not in done:
raise RuntimeError("worker failed before hello: " + self.raw["stderr"].decode(errors="replace"))
reader, writer = accepted.result()
self.wire = Wire(reader, writer, self.session, self.generation)
hello = await asyncio.wait_for(self.wire.receive(), 3)
assert hello == {"kind":"hello", "token":self.token, "role":self.role,
"revision":revision,"slots":["attention.plan","delivery.review"] if revision else [],
"skills":skills_revision}, hello
except BaseException:
await self.close()
raise
finally:
exited.cancel()
with contextlib.suppress(asyncio.CancelledError):
await exited
return self
async def close(self) -> None:
if self.wire:
self.wire.writer.close()
with contextlib.suppress(ConnectionError):
await self.wire.writer.wait_closed()
if self.process and self.process.returncode is None:
self.process.kill()
await self.process.wait()
if self.server:
self.server.close()
await self.server.wait_closed()
for task in self.drains:
with contextlib.suppress(asyncio.CancelledError):
await task
self.directory.cleanup()
async def __aenter__(self):
return await self.start()
async def __aexit__(self, *_):
await self.close()
async def begin(self, code: str) -> str:
operation_id = str(uuid.uuid4())
await self.wire.send({"kind":"execute", "id":operation_id, "code":code})
return operation_id
async def finish(self, operation_id: str, service: Service | None = None) -> dict:
while True:
message = await self.wire.receive()
if message["kind"] == "completed":
assert message["id"] == operation_id
return message["outcome"]
assert message["kind"] == "host_request", message
assert message["execution_id"] == operation_id
self.requests.append(message)
if service is None:
raise AssertionError(f"unexpected host request: {message}")
reply = await service(message)
await self.wire.send({"kind":"host_reply", "id":message["id"],
"execution_id":operation_id, "reply":reply})
async def execute(self, code: str, service: Service | None = None, timeout: float = 3) -> dict:
operation_id = await self.begin(code)
return await asyncio.wait_for(self.finish(operation_id, service), timeout)
async def invoke(self, slot: str, value: dict, timeout: float = 3) -> dict:
operation_id = str(uuid.uuid4())
await self.wire.send({"kind":"invoke","id":operation_id,"slot":slot,"input":value})
return await asyncio.wait_for(self.finish(operation_id), timeout)
def output(outcome: dict, stream: str | None = None) -> str:
return "".join(record["text"] for record in outcome.get("output", []) if stream is None or record["stream"] == stream)