This repository has no description
Something went wrong. Try again.
6.7 kB · 154 lines
Python
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155"""Test-only protocol peer. This is NOT a Python replacement for the Rust host.
It exercises the shipped subprocesses without Rust/provider availability. Ruststorage, worker supervision and SessionRuntime are tested separately by Cargo."""from __future__ import annotations
import asyncioimport contextlibimport jsonimport osimport shutilimport sysimport tempfileimport uuidfrom pathlib import Pathfrom typing import Awaitable, Callable
from klbr_runtime.release import fingerprintfrom 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)