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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351"""Run an interactive shell in a kernel pseudo-terminal over 4-byte-framed stdio.
The first stdin frame is a UTF-8 JSON object: {"cwd": "/path", "cols": 80, "rows": 24, "shell": "/bin/zsh", "credit": 8}
Outbound frames (stdout): 0x00 + utf8-json: {"pid": int, "cwd": str, "shell": str} 0x01 + bytes: raw PTY output bytes 0x02 + int32-be: process exit status (negative for signals)
Inbound frames (stdin): 0x01 + bytes: write bytes to PTY master 0x02 + uint16-be*2: resize PTY to (cols, rows) 0x03: terminate the PTY process group 0x04 + uint8: send signal number to foreground process group 0x05 + uint16-be: allow that many more output frames
With "credit": N in the spec, at most N output frames are sent before thedaemon grants more, and while none are allowed the PTY is not read, so a fastproducer blocks on the PTY instead of filling the daemon's mailbox. Withoutit (a daemon from before flow control) output is never held back."""
from __future__ import annotations
import contextlibimport errnoimport fcntlimport jsonimport osimport ptyimport selectimport shutilimport signalimport structimport subprocessimport sysimport termiosimport threadingimport time
FRAME_READY = 0FRAME_DATA = 1FRAME_RESIZE = 2FRAME_EXIT = 2FRAME_STOP = 3FRAME_SIGNAL = 4FRAME_CREDIT = 5
STRIPPED_ENV = ( "ALBEDO_DAEMON_TOKEN", "ALBEDO_INSPECT", "ALBEDO_PARENT_PID", "ALBEDO_TEST_DAEMON", "ALBEDO_PRELOAD", "PYTHONUNBUFFERED",)
write_lock = threading.Lock()
class Credit: """Output frames the daemon allows; None means unlimited."""
def __init__(self, frames: int | None) -> None: self.frames = frames self.closed = False self.changed = threading.Condition()
def grant(self, frames: int) -> None: with self.changed: if self.frames is not None: self.frames += frames self.changed.notify()
def close(self) -> None: with self.changed: self.closed = True self.changed.notify()
def take(self) -> bool: with self.changed: while self.frames == 0 and not self.closed: self.changed.wait() if self.closed: return False if self.frames is not None: self.frames -= 1 return True
def send_frame(kind: int, payload: bytes) -> bool: frame = struct.pack("!IB", len(payload) + 1, kind) + payload with write_lock: view = memoryview(frame) while view: try: written = os.write(1, view) except OSError: return False if written <= 0: return False view = view[written:] return True
def read_exact(fd: int, size: int) -> bytes | None: chunks = bytearray() while len(chunks) < size: try: chunk = os.read(fd, size - len(chunks)) except OSError as error: if error.errno == errno.EINTR: continue return None if not chunk: return None chunks.extend(chunk) return bytes(chunks)
def read_frame(fd: int) -> bytes | None: header = read_exact(fd, 4) if header is None: return None (length,) = struct.unpack("!I", header) if length == 0 or length > 1_048_576: return None return read_exact(fd, length)
def resolve_cwd(requested: object) -> str: if isinstance(requested, str) and requested.strip(): expanded = os.path.abspath(os.path.expanduser(requested.strip())) if os.path.isdir(expanded): return expanded home = os.path.expanduser("~") return home if os.path.isdir(home) else "/"
def resolve_shell(requested: object) -> str: candidates: list[str] = [] if isinstance(requested, str) and requested.strip(): candidates.append(requested.strip()) env_shell = os.environ.get("SHELL", "").strip() if env_shell: candidates.append(env_shell) candidates.extend(("/bin/zsh", "/bin/bash", "/bin/sh", "sh")) for candidate in candidates: resolved = ( candidate if os.path.isabs(candidate) and os.access(candidate, os.X_OK) else shutil.which(candidate) ) if resolved and os.access(resolved, os.X_OK): return resolved return "/bin/sh"
def build_env(cwd: str, shell: str) -> dict[str, str]: env = {k: v for k, v in os.environ.items() if k not in STRIPPED_ENV} env["TERM"] = "xterm-256color" env["COLORTERM"] = "truecolor" env["TERM_PROGRAM"] = "albedo" env["PWD"] = cwd env["SHELL"] = shell env.setdefault("HOME", os.path.expanduser("~")) env.setdefault("LANG", "en_US.UTF-8") return env
def set_winsize(fd: int, cols: int, rows: int) -> None: c = max(2, min(500, cols)) r = max(1, min(200, rows)) with contextlib.suppress(OSError): fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", r, c, 0, 0))
def terminate_group(pid: int) -> None: for sig, wait_s in ( (signal.SIGHUP, 0.08), (signal.SIGTERM, 0.12), (signal.SIGKILL, 0.05), ): try: os.killpg(pid, sig) except ProcessLookupError: return except PermissionError: with contextlib.suppress(OSError): os.kill(pid, sig) deadline = time.monotonic() + wait_s while time.monotonic() < deadline: try: waited, _ = os.waitpid(pid, os.WNOHANG) if waited == pid: return except ChildProcessError: return time.sleep(0.01)
def write_all(fd: int, data: bytes) -> None: view = memoryview(data) while view: try: written = os.write(fd, view) except OSError as error: if error.errno in (errno.EINTR, errno.EAGAIN): time.sleep(0.005) continue return if written <= 0: return view = view[written:]
def read_output(master_fd: int, proc: subprocess.Popen[bytes]) -> bytes: """Return the next PTY output, or b"" once the shell is gone.""" while True: try: ready, _, _ = select.select([master_fd], [], [], 0.5) except (OSError, ValueError): return b"" if not ready: if proc.poll() is not None: return b"" continue try: return os.read(master_fd, 16_384) except OSError as error: if error.errno != errno.EINTR: return b""
def pump_output(master_fd: int, proc: subprocess.Popen[bytes], credit: Credit) -> None: try: while credit.take(): chunk = read_output(master_fd, proc) if not chunk or not send_frame(FRAME_DATA, chunk): break finally: try: code = proc.wait(timeout=0.5) except subprocess.TimeoutExpired: terminate_group(proc.pid) code = proc.wait() terminate_group(proc.pid) send_frame(FRAME_EXIT, struct.pack("!i", int(code)))
def main() -> int: init_raw = read_frame(0) if init_raw is None: return 1 try: spec = json.loads(init_raw.decode("utf-8")) except Exception: return 1 if not isinstance(spec, dict): return 1
cwd = resolve_cwd(spec.get("cwd")) shell = resolve_shell(spec.get("shell")) cols = spec.get("cols") if isinstance(spec.get("cols"), int) else 80 rows = spec.get("rows") if isinstance(spec.get("rows"), int) else 24
master_fd, slave_fd = pty.openpty() set_winsize(master_fd, cols, rows)
def child_setup() -> None: os.setsid() with contextlib.suppress(OSError): fcntl.ioctl(0, termios.TIOCSCTTY, 0)
try: proc = subprocess.Popen( [shell], stdin=slave_fd, stdout=slave_fd, stderr=slave_fd, cwd=cwd, env=build_env(cwd, shell), preexec_fn=child_setup, close_fds=True, ) except Exception: os.close(master_fd) os.close(slave_fd) return 1 finally: with contextlib.suppress(OSError): os.close(slave_fd)
ready_payload = json.dumps( {"pid": proc.pid, "cwd": cwd, "shell": os.path.basename(shell)} ).encode("utf-8") if not send_frame(FRAME_READY, ready_payload): terminate_group(proc.pid) os.close(master_fd) return 1
window = spec.get("credit") credit = Credit(window if isinstance(window, int) and window > 0 else None) reader = threading.Thread( target=pump_output, args=(master_fd, proc, credit), daemon=True ) reader.start()
try: while proc.poll() is None: try: ready, _, _ = select.select([0], [], [], 0.2) except (OSError, ValueError): break if not ready: continue frame = read_frame(0) if frame is None: break kind = frame[0] body = frame[1:] if kind == FRAME_DATA and body: write_all(master_fd, body) elif kind == FRAME_RESIZE and len(body) == 4: new_cols, new_rows = struct.unpack("!HH", body) set_winsize(master_fd, new_cols, new_rows) elif kind == FRAME_STOP: break elif kind == FRAME_CREDIT and len(body) == 2: credit.grant(struct.unpack("!H", body)[0]) elif kind == FRAME_SIGNAL and len(body) == 1: sig = body[0] try: fg_pgid = os.tcgetpgrp(master_fd) except OSError: fg_pgid = proc.pid with contextlib.suppress(OSError): os.killpg(fg_pgid, sig) finally: credit.close() terminate_group(proc.pid) with contextlib.suppress(OSError): os.close(master_fd) reader.join(timeout=1.0) return 0
if __name__ == "__main__": sys.exit(main())