Something went wrong. Try again.
small gleam coding and persistent agent daemon with a detachable cli and a canvas ui forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309"""Editor-side stdio client. Rejects non-JSON-RPC stdout on every exchange.
A paused reader forces pipe backpressure and adapter queue overflow."""
from __future__ import annotations
import itertoolsimport jsonimport osimport subprocessimport threadingimport timefrom typing import Any, NamedTuple
from harness import ROOT
CLI = ROOT / "cli/bin/albedo"
class Exchange(NamedTuple): """A response, and the notifications the adapter wrote before it."""
response: dict[str, Any] notifications: list[dict[str, Any]]
def _key(value: Any) -> str: return json.dumps(value, sort_keys=True)
def _problem(message: Any) -> str | None: """Why a parsed stdout line is not a JSON-RPC 2.0 message, or None.""" if not isinstance(message, dict): return "not an object" if message.get("jsonrpc") != "2.0": return "missing jsonrpc 2.0" if "method" in message: return None if isinstance(message["method"], str) else "method is not a string" if "id" not in message: return "neither a request, a notification, nor a response" if ("result" in message) == ("error" in message): return "a response needs exactly one of result or error" return None
class AcpClient: def __init__(self, env: dict[str, str]): self.env = env # One condition guards every list below. It is an RLock, so helpers # that hold it may call each other. self._cond = threading.Condition() self._messages: list[dict[str, Any]] = [] self._violations: list[str] = [] self._stderr = bytearray() self._eof = False self._started: dict[str, int] = {} self._write_lock = threading.Lock() self._ids = itertools.count(1) self._reading = threading.Event() self._reading.set() self.process = subprocess.Popen( [env.get("ALBEDO_TEST_CLI", str(CLI)), "acp"], cwd=ROOT, env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) self._threads = [ threading.Thread(target=self._read_stdout, daemon=True), threading.Thread(target=self._read_stderr, daemon=True), ] for thread in self._threads: thread.start()
def write(self, data: bytes) -> None: """Write raw bytes to stdin, flushed as one write.""" stream = self.process.stdin assert stream is not None with self._write_lock: try: stream.write(data) stream.flush() except (OSError, ValueError) as error: raise AssertionError( f"adapter no longer reads stdin: {error}\n{self.stderr()}" ) from error
def write_fragmented(self, data: bytes, size: int, pause: float = 0.02) -> None: """Write raw bytes in pieces, so the adapter reads a partial line.""" for offset in range(0, len(data), size): self.write(data[offset : offset + size]) time.sleep(pause)
def expect(self, request_id: Any) -> None: """Mark where a response to ``request_id`` may start, before it is sent.""" with self._cond: self._started[_key(request_id)] = len(self._messages)
def send( self, method: str, params: Any = None, *, id: Any = None, fragment: int = 0, ) -> Any: """Send a request and return its id. ``fragment`` splits the write.""" if id is None: id = next(self._ids) message: dict[str, Any] = {"jsonrpc": "2.0", "id": id, "method": method} if params is not None: message["params"] = params data = json.dumps(message).encode() + b"\n" self.expect(id) if fragment: self.write_fragmented(data, fragment) else: self.write(data) return id
def notify(self, method: str, params: Any = None) -> None: message: dict[str, Any] = {"jsonrpc": "2.0", "method": method} if params is not None: message["params"] = params self.write(json.dumps(message).encode() + b"\n")
def request( self, method: str, params: Any = None, *, id: Any = None, timeout: float = 30 ) -> Exchange: return self.response(self.send(method, params, id=id), timeout=timeout)
def response(self, request_id: Any, timeout: float = 30) -> Exchange: """Wait for the response to a request sent with ``send`` or ``expect``.""" key = _key(request_id) with self._cond: start = self._started.pop(key) deadline = time.monotonic() + timeout with self._cond: while True: self._raise_violations() notifications = [] for message in self._messages[start:]: if "method" in message: if "id" not in message: notifications.append(message) elif _key(message.get("id")) == key: return Exchange(message, notifications) if self._eof: raise AssertionError( f"adapter exited before answering {request_id}\n{self.stderr()}" ) remaining = deadline - time.monotonic() if remaining <= 0: raise AssertionError( f"no response to {request_id} within {timeout}s\n{self.stderr()}" ) self._cond.wait(remaining)
def messages(self) -> list[dict[str, Any]]: """Every stdout message so far, in arrival order (the same objects).""" with self._cond: return list(self._messages)
def notifications(self, method: str | None = None) -> list[dict[str, Any]]: return [ message for message in self.messages() if "method" in message and "id" not in message and (method is None or message["method"] == method) ]
def updates(self, session_id: str) -> list[dict[str, Any]]: """The ``update`` payloads of ``session/update`` for one session.""" return [ message["params"]["update"] for message in self.notifications("session/update") if message["params"].get("sessionId") == session_id ]
def requests_from_adapter(self) -> list[dict[str, Any]]: """Requests the adapter sent to the editor, such as a permission request.""" return [ message for message in self.messages() if "method" in message and "id" in message ]
def wait_for(self, predicate, timeout: float = 30, message: str = "condition"): """Poll ``predicate`` until it returns a truthy value, and return that.""" deadline = time.monotonic() + timeout while True: value = predicate() if value: return value with self._cond: self._raise_violations() if time.monotonic() >= deadline: raise AssertionError( f"timed out after {timeout}s waiting for {message}\n{self.stderr()}" ) time.sleep(0.02)
def stderr(self) -> str: with self._cond: return bytes(self._stderr).decode(errors="replace")
def check_stdout(self) -> None: """Fail if stdout ever carried anything but ACP messages.""" with self._cond: self._raise_violations()
def _accept(self, raw: bytes) -> None: try: message = json.loads(raw) except ValueError: problem = "not JSON" else: problem = _problem(message) with self._cond: if problem is None: self._messages.append(message) else: self._violations.append(f"{problem}: {raw[:300]!r}") self._cond.notify_all()
def _read_stdout(self) -> None: stream = self.process.stdout assert stream is not None buffered = b"" while True: self._reading.wait() chunk = os.read(stream.fileno(), 1 << 16) if not chunk: break buffered += chunk start = 0 while (end := buffered.find(b"\n", start)) >= 0: self._accept(buffered[start:end]) start = end + 1 buffered = buffered[start:] with self._cond: if buffered: self._violations.append(f"stdout ended mid-line: {buffered[:300]!r}") self._eof = True self._cond.notify_all()
def _read_stderr(self) -> None: stream = self.process.stderr assert stream is not None while chunk := os.read(stream.fileno(), 1 << 16): with self._cond: self._stderr.extend(chunk)
def _raise_violations(self) -> None: if self._violations: raise AssertionError( "adapter wrote to stdout outside ACP:\n" + "\n".join(self._violations[:5]) + "\nstderr:\n" + self.stderr() )
def pause_reading(self) -> None: self._reading.clear()
def signal(self, number: int) -> None: self.process.send_signal(number)
def close_stdin(self) -> None: """Close the editor's side so the adapter sees EOF on stdin.""" stream = self.process.stdin if stream is not None and not stream.closed: try: stream.close() except BrokenPipeError: pass
def wait_exit(self, timeout: float = 30) -> int: try: return self.process.wait(timeout=timeout) except subprocess.TimeoutExpired as error: raise AssertionError( f"adapter still running after {timeout}s\n{self.stderr()}" ) from error
def close(self, timeout: float = 30) -> int: """Close stdin, wait for the adapter to exit, and check its stdout.""" self.close_stdin() code = self.wait_exit(timeout) self._join() self.check_stdout() return code
def kill(self) -> None: """Stop the adapter during cleanup.""" if self.process.poll() is None: self.process.kill() self.process.wait(timeout=10) # A paused reader must wake to see EOF and finish. self._reading.set() self._join()
def _join(self) -> None: for thread in self._threads: thread.join(timeout=2)