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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933"""Shared fixtures for offline end-to-end tests, run through ``run.py``.
A provider is a scripted model (``fake_provider.py``, re-exported here): itscallback receives each decoded request and returns ``text(...)``,``python(...)``, ``error(...)``, or a custom ``Reply``. Both OpenAI streamingprotocols are served. ``Provider(catalog=...)`` also servesa local models catalog. ``Albedo`` gives each fixture a separate workspace andprovider route on one shared HTTP server, on the daemon the runner gave itstest. Tests share one daemon and keep to their own profile; an ``@exclusive``test, which may change global settings or restart the daemon, gets a fresh onethat boots after its fixture's ``prepare(app)`` and is discarded afterwards, soit never restores what it changed. ``providers={}`` leaves it unconfigured.Setting ``app.daemon.open_files = (soft, hard)`` in ``prepare`` boots thatdaemon under an open-file limit. Boot and restart use ``albedo daemon`` tovalidate readiness without listing or formatting every stored session; scenariosthat exercise session listing do so explicitly.``app.delete_tree(session)`` deletes a session and its descendants and returnsthe outcome. ``write_extensions(settings)`` replaces extensions.json withoutletting the daemon reach the network, and ``store_secrets(section, value)`` writes acreds.json section such as the OAuth ``accounts`` or ``mcp`` server secrets.A response ``Reply(..., usage=None)`` omits provider usage.A ``python(...)`` cell that starts with ``LOOPBACK_REMOTE`` has ``rem`` bound toa remote connection whose ssh runs everything on this machine.
Example::
from harness import Albedo, Provider, text
provider = Provider(lambda request: text("hello")) try: with Albedo(provider) as app: session = app.session() app.prompt(session, "say hello").close() app.idle(session) assert provider.requests finally: provider.close()"""
from __future__ import annotations
from typing import Any
import atexitimport contextlibimport errnoimport http.clientimport ioimport jsonimport osfrom pathlib import Pathimport shutilimport secretsimport signalimport subprocessimport sysimport tempfileimport threadingimport timeimport urllib.errorimport urllib.parseimport uuid
import fake_providerfrom fake_provider import Reply, error, python, text # noqa: F401
ROOT = Path(__file__).resolve().parents[2]sys.path.insert(0, str(ROOT / "test"))import scratch # noqa: E402
scratch.claim("e2e")# A test daemon needs two schedulers, and busy-waiting ones make every boot# pin all cores of the machine.TEST_VM_FLAGS = "+S 2:2 +SDcpu 2:2 +sbwt none +sbwtdcpu none +sbwtdio none"# Credentials a provider reads from the environment. The daemon must not see# the developer's, or a catalog reload reaches the live provider.PROVIDER_ENVIRONMENT = ("ALIBABA_API_KEY", "DASHSCOPE_API_KEY", "ALIBABA_BASE_URL")# Extension settings that keep a daemon off the network: the models catalog# never refreshes by itself and the cache-TTL table fetches no remote copy.OFFLINE = {"models": {"refreshHours": 0}, "cacheTtl": {"url": None}}
def operation_id(timestamp=None): """Generate one UUIDv7 action identity, optionally at a fixture time.""" milliseconds = int(time.time() * 1000) if timestamp is None else timestamp bits = (milliseconds << 80) | (7 << 76) | secrets.randbits(76) bits = (bits & ~(3 << 62)) | (2 << 62) return str(uuid.UUID(int=bits))
def exclusive(target): """Run the test, or every test of the class, on a daemon of its own that boots after the fixture's prepare step and is discarded afterwards: it may change global settings or restart the daemon.""" target._e2e_exclusive = True return target
_daemons: list[Daemon] = []# The launcher of this run's daemon build; see snapshot_daemon()._executable = None_daemons_lock = threading.RLock()_current = threading.local()daemon_boots = 0restart_seconds = 0.0# mist closes a keep-alive connection after 10 idle seconds; one idle for less# than this is reused, so a request never races the daemon closing it.KEEP_ALIVE_SECONDS = 5
class Provider(fake_provider.Provider): """A scripted model the running test owns: the runner releases its route once the test's sessions stop."""
def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) getattr(_current, "providers", []).append(self)
class _Connection(http.client.HTTPConnection): """A daemon connection that notes when it can carry the next request."""
def __init__(self, port, token): super().__init__("127.0.0.1", port, timeout=20) self.token = token self.response_class = _Response self.used_at = time.monotonic() self.reusable = True self.response = None
def idle(self): """Whether the next request may go out on this connection: its last response is done with and the daemon has not dropped it yet.""" return ( self.sock is not None and self.reusable and (self.response is None or self.response.isclosed()) and time.monotonic() - self.used_at < KEEP_ALIVE_SECONDS )
class _Response(http.client.HTTPResponse): """Closing drains a short unread body so the connection can be reused; an event stream or a long body gives the connection up instead."""
connection = None
def close(self): if self.headers is None: # begin() died before any status line: there is no body to drain, # and reading headers here would mask the error that stopped it. super().close() return connection = self.connection streaming = self.getheader("content-type", "").startswith("text/event-stream") if connection is not None and streaming: connection.reusable = False if connection is not None and not self.isclosed(): if streaming or self.length is None or self.length > 1 << 20: connection.reusable = False else: try: self.read() except (OSError, http.client.HTTPException): connection.reusable = False super().close() if connection is not None: connection.used_at = time.monotonic() if self.will_close or not connection.reusable: connection.close()
def log_tail(path, offset=0): try: with path.open("rb") as log: log.seek(0, os.SEEK_END) log.seek(max(offset, log.tell() - 8000)) return log.read(8000).decode(errors="replace") except FileNotFoundError: return ""
class Daemon: """One daemon and its hermetic home. A concurrent daemon is shared by tests that each keep to their own provider profile; any other serves one test, which may change its global state or restart it."""
def __init__(self, name, *, concurrent=False): if not hasattr(sys.modules["__main__"], "__file__"): # A daemon lives until the process that booted it exits, which a # long-lived interpreter or agent kernel never does. raise RuntimeError( "run E2E tests through test/e2e/run.py, not an interactive interpreter" ) self.concurrent = concurrent self.lock = threading.RLock() # Resolved, so a path the daemon reports through a symlink (macOS /tmp) matches. self.root = Path(tempfile.mkdtemp(prefix=f"{name}-")).resolve() self.home = self.root / "home" self.home.mkdir() (self.root / "user-home").mkdir() # Nothing of the developer's albedo setup reaches the daemon. self.env = dict( { key: value for key, value in os.environ.items() if key not in PROVIDER_ENVIRONMENT and key not in ("HOME", "ERL_FLAGS") and not key.startswith("ALBEDO_") }, HOME=str(self.root / "user-home"), ALBEDO_HOME=str(self.home), ALBEDO_ROOT=str(ROOT), ALBEDO_NO_BROWSER="1", ALBEDO_PARENT_PID=str(os.getpid()), ERL_FLAGS=TEST_VM_FLAGS, # Test daemons do not need the package-upgrade safety preload; the # daemon preload test removes this opt-out before it boots. ALBEDO_PRELOAD="0", ALBEDO_IDLE_SECONDS="10", # Kernels outlive a stopped daemon only this long; one whose home # is removed leaves at once. ALBEDO_KERNEL_GRACE_SECONDS="30", # An OS-picked intake port, so test daemons never take 8088 from # the developer's daemon. ALBEDO_WEBHOOK_PORT="0", ALBEDO_MCP_SECRET="configured-secret", ALBEDO_MCP_CLOSED=str(self.root / "closed"), ALBEDO_MCP_AMBIENT="must-not-reach-the-server", ) if _executable: self.env["ALBEDO_DAEMON"] = _executable if "ALBEDO_TEST_CLI" in os.environ: self.env["ALBEDO_TEST_CLI"] = os.environ["ALBEDO_TEST_CLI"] # Opt-in: NAME=VALUE pairs, space separated, for every test daemon, so # the whole suite can run under a non-default daemon setting. A test's # own `prepare` still overrides them. for pair in os.environ.get("ALBEDO_E2E_DAEMON_ENV", "").split(): name, _, value = pair.partition("=") self.env[name] = value # A (soft, hard) open-file limit the CLI, and so the daemon it boots, # starts under; None passes on the test process's own. self.open_files = None self.connection = None self.base = None self._pid = None self._local = threading.local() self.boot_action = None self._launchers = [] self._commands = [] with _daemons_lock: _daemons.append(self)
@property def booted(self): return self._pid is not None
def boot(self): """Start the daemon; the first fixture boots it after its prepare.""" extensions = self.home / "extensions.json" if not extensions.exists(): extensions.write_text(json.dumps(OFFLINE)) if self.boot_action is not None: self.boot_action(self) self.cli("daemon") self._refresh()
def start_launcher(self): """Run the real launcher directly, retaining its log and process owner.""" with (self.home / "daemon.log").open("ab") as log: process = subprocess.Popen( [self.env["ALBEDO_DAEMON"]], cwd=ROOT, env=self.env, stdout=log, stderr=log, ) self._launchers.append(process) # Direct launches are our children. Reap them as soon as they exit so # another CLI can observe that the old daemon is gone during upgrade. threading.Thread(target=process.wait, daemon=True).start() return process
def await_ready(self, timeout=20): """Adopt the published endpoint once its authenticated API answers.""" deadline = time.monotonic() + timeout while time.monotonic() < deadline: try: self._refresh() with self.api("/server") as response: if json.load(response)["state"] == "ready": return except (OSError, ValueError, http.client.HTTPException): pass time.sleep(0.02) self._fail("launcher did not publish an authenticated endpoint")
def start_cli(self, *args, **streams): """Start a CLI command whose input or cancellation the fixture controls.""" process = subprocess.Popen( [self.env.get("ALBEDO_TEST_CLI", str(ROOT / "cli/bin/albedo")), *args], cwd=ROOT, env=self.env, **streams, ) self._commands.append(process) return process
def _refresh(self): global daemon_boots self.connection = json.loads((self.home / "daemon.json").read_text()) self.base = f"http://127.0.0.1:{self.connection['port']}" if self.connection["pid"] != self._pid: # A new PID is valid only after the old daemon has exited. if self._pid is not None and alive(self._pid): raise AssertionError(f"a second daemon started in {self.home}") self._pid = self.connection["pid"] with _daemons_lock: daemon_boots += 1
def restart(self, *, crash=False, prepare=None): """Replace the daemon under the same home, preparing storage between.""" global restart_seconds started = time.monotonic() if crash: assert self._pid is not None os.kill(self._pid, signal.SIGKILL) # SIGKILL is asynchronous: a CLI started before the process is gone # reads the dead daemon's record and finds nothing listening. if not self._gone(timeout=15): self._fail("daemon did not die") elif not self._stop(timeout=15): self._fail("daemon did not stop") if prepare is not None: prepare() self.cli("daemon") self._refresh() with _daemons_lock: restart_seconds += time.monotonic() - started
def _stop(self, timeout): try: with self.api("/server") as response: instance_id = json.load(response)["instance_id"] self.api( "/server/shutdown", { "instance_id": instance_id, "timeout_ms": min(timeout * 1000, 30000), }, ).close() except (OSError, http.client.HTTPException): pass return self._gone(timeout)
def _gone(self, timeout): deadline = time.monotonic() + timeout while alive(self._pid): if time.monotonic() > deadline: return False time.sleep(0.05) return True
def shutdown(self): """Stop the daemon, killing it if it will not stop, and drop the home.""" with self.lock: # A failed CLI boot may already have launched this private home's daemon. path = self.home / "daemon.json" if path.exists(): self.connection = json.loads(path.read_text()) self.base = f"http://127.0.0.1:{self.connection['port']}" self._pid = self.connection["pid"] if self._pid is not None: if alive(self._pid) and not self._stop(timeout=10): try: os.kill(self._pid, signal.SIGKILL) except ProcessLookupError: pass self._pid = None for launcher in self._launchers: if launcher.poll() is None: launcher.kill() launcher.wait() for command in self._commands: if command.poll() is None: command.kill() command.wait() shutil.rmtree(self.root, ignore_errors=True) with _daemons_lock: if self in _daemons: _daemons.remove(self)
def kill(self): """SIGKILL the daemon without waiting for its lock, even mid-stop.""" pid = self._pid if pid is not None: try: os.kill(pid, signal.SIGKILL) except ProcessLookupError: pass
def discard(self): """Discard a finished test's daemon before its worker admits another.""" self.kill() self.shutdown()
def _fail(self, message): tail = log_tail(self.home / "daemon.log") raise AssertionError(message + ("\n" + tail if tail else ""))
def _connection(self): connection = getattr(self._local, "connection", None) assert self.connection is not None port, token = self.connection["port"], self.connection["token"] # A restarted daemon has a new token and may be back on the same port. if connection is None or connection.token != token or not connection.idle(): if connection is not None: # A response still being read, such as an event stream, keeps # its socket; closing it gives the connection up. connection.reusable = False if connection.response is None or connection.response.isclosed(): connection.close() connection = self._local.connection = _Connection(port, token) return connection
def api(self, path, body=None, *, method=None, headers=None): """One daemon request over this thread's kept-alive connection. A socket per request, polled at test speed, leaves thousands of ports in TIME_WAIT; the host runs out and unrelated connections fail.""" connection = self._connection() assert self.connection is not None connection.used_at = time.monotonic() payload = None if body is None else json.dumps(body).encode() try: connection.request( method or ("GET" if payload is None else "POST"), path, payload, { "Authorization": "Bearer " + self.connection["token"], "Content-Type": "application/merge-patch+json" if method == "PATCH" else "application/json", **(headers or {}), }, ) response = connection.getresponse() except BaseException: connection.close() raise response.connection, connection.response = connection, response if response.status < 400: return response payload = response.read() response.close() raise urllib.error.HTTPError( self.base + path, response.status, f"{response.reason} ({method} {path}): " + payload.decode(errors="replace"), response.headers, io.BytesIO(payload), )
def cli(self, *args): command = [self.env.get("ALBEDO_TEST_CLI", str(ROOT / "cli/bin/albedo")), *args] if self.open_files: command = [ sys.executable, "-c", _WITH_OPEN_FILES, *map(str, self.open_files), *command, ] result = subprocess.run( command, cwd=ROOT, env=self.env, text=True, capture_output=True, timeout=45, ) if result.returncode: self._fail(result.stdout + result.stderr) return result.stdout
def install_shared_features(self): """Enable the additive global gateways concurrent tests rely on, once, before they start: none of them may change global state itself.""" settings = json.loads((self.home / "extensions.json").read_text()) settings.setdefault("enabled", {})["proxy"] = True (self.home / "extensions.json").write_text(json.dumps(settings)) with self.api("/settings?group=extensions") as response: validator = response.getheader("ETag") response.read() self.api( "/settings?group=extensions", {"defaults": {"webhooks": True}}, method="PATCH", headers={"If-Match": validator}, ).close()
# Runs a command under a (soft, hard) open-file limit, as a shell started with# that `ulimit -n` would._WITH_OPEN_FILES = """import os, resource, sysresource.setrlimit(resource.RLIMIT_NOFILE, (int(sys.argv[1]), int(sys.argv[2])))os.execv(sys.argv[3], sys.argv[3:])"""
def latest_user(request): """Return the last user's raw content from either provider protocol.""" messages = request.get("messages", request.get("input", [])) return next( ( item.get("content", "") for item in reversed(messages) if item.get("role") == "user" ), "", )
def prompt_text(request): """The last user's text without a trailing `<system-note>`: a daemon note (a restore after a kernel release, a switch) rides on the user turn, and a scripted provider keyed on the prompt must still match it. A note sent on its own, after a cell's result, is not the prompt.""" messages = request.get("messages", request.get("input", [])) return next( ( str(item.get("content", "")).split("<system-note>")[0] for item in reversed(messages) if item.get("role") == "user" and not is_note(item) ), "", )
def is_note(item): """Whether a request item is a daemon note sent on its own.""" return item.get("role") == "user" and str(item.get("content", "")).startswith( "<system-note>" )
def turn_notes(request): """The daemon notes since the last prompt: those riding on it, and those sent after a cell's result, such as the one a kernel's arrival leaves.""" messages = request.get("messages", request.get("input", [])) users = [ str(item.get("content", "")) for item in messages if item.get("role") == "user" ] prompts = [ i for i, text in enumerate(users) if not text.startswith("<system-note>") ] since = users[prompts[-1] :] if prompts else users return [ "<system-note>" + part for text in since for part in text.split("<system-note>")[1:] ]
def wait_until(condition, seconds, message): deadline = time.monotonic() + seconds while time.monotonic() < deadline: if condition(): return time.sleep(0.1) raise AssertionError(message)
def release_fifo(path): """Let a job blocked reading one byte from the fifo at `path` go on. True once a reader had it open and got its byte; False while none has it open yet, so callers retry: `wait_until(lambda: release_fifo(gate), ...)`. A reader that closes between the open and the write (a released job exiting while a cleanup releases it again) got no byte either.""" try: descriptor = os.open(path, os.O_WRONLY | os.O_NONBLOCK) except OSError as failure: if failure.errno == errno.ENXIO: return False raise try: os.write(descriptor, b"x") except BrokenPipeError: return False finally: os.close(descriptor) return True
# An ssh that runs its command right here, under a login shell that does the# same, so a cell's remote.connect("loopback") boots, stages and talks to a# real remote kernel._LOOPBACK_SSH = """#!/bin/shwhile [ $# -gt 0 ]; do case "$1" in -o) shift 2 ;; -*) shift ;; *) break ;; esacdoneshiftexec sh -c "$*""""_LOOPBACK_LOGIN = '#!/bin/sh\n[ "$1" = -l ] && shift\nexec /bin/sh "$@"\n'# The start of a cell that binds `rem` to a loopback remote connection.LOOPBACK_REMOTE = f"""import osos.makedirs("bin", exist_ok=True)for name, script in (("ssh", {_LOOPBACK_SSH!r}), ("loginsh", {_LOOPBACK_LOGIN!r})): with open(f"bin/{{name}}", "w") as file: file.write(script) os.chmod(f"bin/{{name}}", 0o755)os.environ["SHELL"] = os.path.abspath("bin/loginsh")os.environ["PATH"] = os.path.abspath("bin") + os.pathsep + os.environ["PATH"]rem = await remote.connect("loopback")"""
def alive(pid): try: os.kill(pid, 0) except ProcessLookupError: return False return True
def snapshot_daemon(): """Compile the daemon once for the run and boot every test daemon from a copy, not through `gleam run` (see test/snapshot-daemon.sh). test.sh snapshots once for every suite and passes it as ALBEDO_TEST_DAEMON.""" global _executable _executable = os.environ.get("ALBEDO_TEST_DAEMON") if _executable: return result = subprocess.run( [str(ROOT / "test/snapshot-daemon.sh"), tempfile.mkdtemp(prefix="daemon-")], text=True, capture_output=True, ) if result.returncode: raise RuntimeError("building the daemon failed:\n" + result.stderr) _executable = result.stdout.strip()
def current_daemon(): """The daemon the runner gave this thread.""" daemon = getattr(_current, "daemon", None) if daemon is None: raise RuntimeError("run E2E tests through test/e2e/run.py") return daemon
@contextlib.contextmanagerdef on_daemon(daemon): """Run fixtures created on this thread against daemon.""" previous = getattr(_current, "daemon", None) _current.daemon = daemon _current.fixtures = [] _current.providers = [] try: yield daemon finally: _current.daemon = previous _current.fixtures = [] _current.providers = []
def finish_test(): """Stop owned session trees before releasing their scripted providers.""" daemon = current_daemon() failures = [] try: if daemon.concurrent: for fixture in _current.fixtures: try: fixture.cleanup() except Exception as error: failures.append(error) else: daemon.discard() finally: for provider in _current.providers: provider.release() if failures: raise ExceptionGroup("fixture teardown failed", failures)
def shutdown(): """Stop every daemon, then the provider server.""" # The run is over and no test watches these daemons stop, so none is # drained: a graceful stop sits out OTP's one-second pause for buffered # output (user_sup:terminate/2), and the shared daemon's drain of every # session takes longer still. Their kernels exit with them. for daemon in list(_daemons): daemon.kill() for daemon in list(_daemons): daemon.shutdown() fake_provider.close()
atexit.register(shutdown)
def _write_config(path, value): with tempfile.NamedTemporaryFile(mode="w", dir=path.parent, delete=False) as output: json.dump(value, output) temporary = output.name os.replace(temporary, path)
class Albedo: """A test's workspace and provider profile on the daemon the runner gave it."""
def __init__( self, provider: Provider | None = None, *, protocol="chat_completions", providers=None, prepare=None, ): self.provider = provider or Provider(lambda _request: text("ok")) self._owns_provider = provider is None self.protocol = protocol self.providers = providers self.prepare = prepare self._sessions = set()
def __enter__(self): self.daemon = daemon = current_daemon() with daemon.lock: self.root, self.home, self.env = daemon.root, daemon.home, daemon.env self.workspace = Path(tempfile.mkdtemp(prefix="workspace-", dir=self.root)) _current.fixtures.append(self) self._concurrent = daemon.concurrent if self.prepare: self.prepare(self) if not (self.home / "extensions.json").exists(): self.write_extensions({}) default_name = ( f"fixture-{self.provider.route}" if self._concurrent else "fixture" ) configured: dict[str, Any] = ( self.providers if self.providers is not None else { default_name: { "extension": "openai", "baseUrl": self.provider.url, "apiKey": "fixture-key", "model": "fixture-model", "protocol": self.protocol, } } ) self.profile = next(iter(configured), None) if configured: current: dict[str, Any] = ( json.loads((self.home / "config.json").read_text()) if self._concurrent and (self.home / "config.json").exists() else {"providers": {}} ) current["providers"].update(configured) current["active"] = ( current.get("active", self.profile) if self._concurrent else self.profile ) _write_config(self.home / "config.json", current) elif not self._concurrent and (self.home / "config.json").exists(): (self.home / "config.json").unlink() if not daemon.booted: daemon.boot() return self
@property def connection(self): return self.daemon.connection
@property def base(self): return self.daemon.base
def restart(self, *, crash=False, prepare=None): """Restart the daemon; optionally prepare offline storage fixtures.""" self.daemon.restart( crash=crash, prepare=None if prepare is None else lambda: prepare(self) )
def _fail(self, message): self.daemon._fail(message)
def api(self, path, body=None, *, method=None, headers=None): response = self.daemon.api(path, body, method=method, headers=headers) if method == "PUT" and isinstance(body, dict) and body.get("kind") == "new": self._sessions.add(path.split("?")[0].rsplit("/", 1)[-1]) return response
def cli(self, *args): return self.daemon.cli(*args)
def write_extensions(self, settings): """Replace extensions.json, keeping the daemon offline in whatever catalog settings leave out.""" merged = { name: {**OFFLINE.get(name, {}), **value} if name in OFFLINE and isinstance(value, dict) else value for name, value in {**OFFLINE, **settings}.items() } (self.home / "extensions.json").write_text(json.dumps(merged))
def store_secrets(self, section, value): """Replaces one creds.json section, where the daemon keeps every secret.""" path = self.home / "creds.json" creds = json.loads(path.read_text()) if path.exists() else {} creds[section] = value path.write_text(json.dumps(creds)) path.chmod(0o600)
def session(self, workspace=None): session_id = operation_id() with self.api( f"/sessions/{session_id}", { "kind": "new", "workspace": str(workspace or self.workspace), "provider_profile": self.profile, }, method="PUT", headers={"If-None-Match": "*"}, ) as response: created = json.load(response) assert created["id"] == session_id, created assert created["provider_profile"] == self.profile, created return session_id
def prompt(self, session_id, content): return self.api( f"/sessions/{session_id}/inputs/{operation_id()}", {"kind": "message", "text": content}, method="PUT", )
def idle(self, session_id, timeout=30): deadline = time.monotonic() + timeout snapshot = None while time.monotonic() < deadline: try: with self.api(f"/sessions/{session_id}?tail=0") as response: snapshot = json.load(response) except urllib.error.HTTPError as failure: if failure.code != 503: raise snapshot = json.load(failure) else: if ( snapshot["status"]["phase"] == "idle" and not snapshot["pending_inputs"] ): return snapshot time.sleep(0.05) self._fail(f"session did not settle: {snapshot}")
def stream_page(self, session_id, after=None, *, tail=None): query = {} if after is not None: query = { "after_generation": after["generation"], "after_seq": after["cursor"], } if tail is not None: query["tail"] = tail suffix = "?" + urllib.parse.urlencode(query) if query else "" with self.api( f"/sessions/{session_id}{suffix}", headers={"Accept": "text/event-stream"}, ) as response: for line in response: if line.startswith(b"data: "): return json.loads(line[6:]) self._fail("stream closed before its first batch")
def history(self, session_id): with self.api(f"/sessions/{session_id}/history") as response: return json.load(response)
def family_revision(self, session_id): """The family revision a workspace move or subtree deletion is conditional on. The configuration resource carries none; the session representation does.""" with self.api(f"/sessions/{session_id}?tail=0") as response: return json.load(response)["family_revision"]
def delete_tree(self, session): """Delete a session and its descendants; the outcome, or None when the session is already gone.""" try: with self.api(f"/sessions/{session}?view=configuration") as response: response.read() etag = response.headers["ETag"] family = self.family_revision(session) with self.api( f"/sessions/{session}?view=configuration&scope=subtree" f"&family_revision={family}", method="DELETE", headers={"If-Match": etag}, ) as response: return json.load(response) except urllib.error.HTTPError as refused: if refused.code in (404, 410): return None raise
def cleanup(self): failures = [] for session in self._sessions: try: outcome = self.delete_tree(session) if outcome and outcome["state"] != "complete": raise AssertionError(f"fixture deletion incomplete: {outcome}") except Exception as failure: failures.append(failure) if failures: raise ExceptionGroup("session teardown failed", failures) shutil.rmtree(self.workspace, ignore_errors=True)
def __exit__(self, *_exc): if self._owns_provider: self.provider.close()