"""Shared fixtures for offline end-to-end tests, run through ``run.py``. A provider is a scripted model: its callback receives each decoded request and returns ``text(...)``, ``python(...)``, ``error(...)``, or a custom ``Reply``. Both OpenAI streaming protocols are served. ``Provider(catalog=...)`` also serves a local models catalog. ``Albedo`` gives each fixture a separate workspace and provider route on one shared HTTP server, on the daemon the runner gave its test. 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 one that boots after its fixture's ``prepare(app)`` and is discarded afterwards, so it never restores what it changed. ``providers={}`` leaves it unconfigured. Setting ``app.daemon.open_files = (soft, hard)`` in ``prepare`` boots that daemon under an open-file limit. ``write_extensions(settings)`` replaces extensions.json without letting the daemon reach the network, and ``store_secrets(section, value)`` writes a creds.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 to a 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 collections.abc import Iterable from typing import Any import atexit import contextlib import errno from dataclasses import dataclass, field import http.client import http.server import io import itertools import json import os from pathlib import Path import shutil import secrets import signal import subprocess import sys import tempfile import threading import time import urllib.error import urllib.parse import uuid from provider_events import chat_events, responses_events 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}} @dataclass class Reply: """A scripted assistant response and optional transport behavior.""" kind: str value: str = "" status: int = 200 reasoning: str | None = None usage: dict | None = field( default_factory=lambda: {"input_tokens": 10, "output_tokens": 20} ) delay: float = 0 hang: bool = False tool_name: str = "python" tool_arguments: dict | None = None events: Iterable[dict] | None = None def text( value: str, *, reasoning: str | None = None, usage: dict | None = None, delay: float = 0, hang: bool = False, ) -> Reply: return Reply( "text", value, reasoning=reasoning, usage=usage or {"input_tokens": 10, "output_tokens": 20}, delay=delay, hang=hang, ) def python( code: str, *, reasoning: str | None = None, tool_name: str = "python" ) -> Reply: return Reply("python", code, reasoning=reasoning, tool_name=tool_name) def error(status: int, body: str = "fixture error") -> Reply: return Reply("error", body, status=status) # Guards the provider registry and the shared provider server. _config_lock = threading.RLock() 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 _providers = {} _provider_ids = itertools.count(1) _server = None def _provider_server(): global _server if _server is None: class Handler(http.server.BaseHTTPRequestHandler): def log_message(self, format, *args): pass def do_GET(self): self.provider()._get(self) def do_POST(self): self.provider()._post(self) def provider(self): return _providers[self.path.split("/", 3)[2]] _server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) _server.daemon_threads = True threading.Thread(target=_server.serve_forever, daemon=True).start() return _server class Provider: """Threaded fake provider that records requests and emits protocol SSE.""" def __init__(self, script, *, chunk_size: int = 20, catalog=None): self.catalog = catalog self.script = script self.chunk_size = chunk_size self.requests: list[dict] = [] self._lock = threading.Lock() self.route = str(next(_provider_ids)) with _config_lock: _providers[self.route] = self self.server = _provider_server() self.url = f"http://127.0.0.1:{self.server.server_port}/t/{self.route}" def _get(owner, self): payload = json.dumps(owner.catalog).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.end_headers() self.wfile.write(payload) def _post(owner, self): request = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) record = { "path": self.path, "authorization": self.headers.get("Authorization"), "model": request.get("model"), "request": request, } with owner._lock: owner.requests.append(record) index = len(owner.requests) - 1 reply = owner.script(request) if isinstance(reply, (list, tuple)): reply = reply[min(index, len(reply) - 1)] if reply.kind == "error": payload = json.dumps({"error": {"message": reply.value}}).encode() self.send_response(reply.status) self.send_header("Content-Type", "application/json") self.end_headers() self.wfile.write(payload) return self.send_response(200) self.send_header("Content-Type", "text/event-stream") self.end_headers() def emit(value): if reply.delay: time.sleep(reply.delay) self.wfile.write(("data: " + json.dumps(value) + "\n\n").encode()) self.wfile.flush() try: if reply.hang: time.sleep(60) return if reply.events is not None: for event in reply.events: emit(event) self.wfile.write(b"data: [DONE]\n\n") self.wfile.flush() return call_id = f"fixture-call-{index + 1}" arguments = json.dumps( reply.tool_arguments if reply.tool_arguments is not None else {"pattern": reply.value} if reply.tool_name == "lcm_grep" else {"code": reply.value, "timeout_ms": 60000} ) events = ( chat_events(reply, arguments, call_id, owner.chunk_size) if self.path.endswith("chat/completions") else responses_events(reply, arguments, call_id, owner.chunk_size) ) for event in events: emit(event) self.wfile.write(b"data: [DONE]\n\n") self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass def close(self): # Sessions can finish background turns after their test returns. Keep # their route alive until the shared daemon has stopped. pass _daemons: list[Daemon] = [] # The launcher of this run's daemon build; see snapshot_daemon(). _executable = None _daemons_lock = threading.RLock() _discarding: list[threading.Thread] = [] _current = threading.local() daemon_boots = 0 restart_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 _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() 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, 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", 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"] # 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("sessions") 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("sessions") 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): """Shut the daemon down in the background; the test that used it is done, and nothing waits on its teardown until the run ends.""" thread = threading.Thread(target=self.shutdown) thread.start() with _daemons_lock: _discarding.append(thread) def _fail(self, message): log = self.home / "daemon.log" raise AssertionError(message + ("\n" + log.read_text() if log.exists() 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, sys resource.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 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), ...)`.""" try: descriptor = os.open(path, os.O_WRONLY | os.O_NONBLOCK) except OSError as error: if error.errno == errno.ENXIO: return False raise try: os.write(descriptor, b"x") 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/sh while [ $# -gt 0 ]; do case "$1" in -o) shift 2 ;; -*) shift ;; *) break ;; esac done shift exec 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 os os.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.contextmanager def on_daemon(daemon): """Run fixtures created on this thread against daemon.""" previous = getattr(_current, "daemon", None) _current.daemon = daemon try: yield daemon finally: _current.daemon = previous def shutdown(): """Stop every daemon, then the provider server.""" global _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 thread in list(_discarding): thread.join() _discarding.clear() for daemon in list(_daemons): daemon.shutdown() if _server is not None: _server.shutdown() _server.server_close() _server = None with _config_lock: _providers.clear() 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 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)) 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): return self.daemon.api(path, body, method=method, headers=headers) 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 error: if error.code != 503: raise snapshot = json.load(error) 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 __exit__(self, *_exc): if self._owns_provider: self.provider.close()