Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
Python
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191"""A detached session loses its kernel, keeps its variables, and says so.
Runs the daemon with a ten second idle limit; the shipped default is 31 minutes."""
import contextlibimport http.serverimport jsonimport osfrom pathlib import Pathimport subprocessimport tempfileimport threadingimport timeimport urllib.request
ROOT = Path(__file__).resolve().parents[2]IDLE_SECONDS = 10
class Provider(http.server.BaseHTTPRequestHandler): requests = [] lock = threading.Lock()
def log_message(self, *_): pass
def do_POST(self): request = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) with self.lock: self.requests.append(request) prompt = latest_user(request) code = None if request["input"][-1].get("role") == "user": if prompt.startswith("remember"): code = "answer = 7\nrows = [1, 2, 3]\n" elif prompt.startswith("recall"): code = "print('recalled', answer, rows)" output = [{"type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "ok", "annotations": []}]}] if code: output = [{"type": "function_call", "id": "fc-" + str(len(self.requests)), "call_id": "call-" + str(len(self.requests)), "name": "python", "status": "completed", "arguments": json.dumps({"code": code, "timeout_ms": 5000})}] body = ("data: " + json.dumps({"type": "response.completed", "response": { "id": "response", "status": "completed", "output": output, }}) + "\n\n").encode() self.send_response(200) self.send_header("Content-Type", "text/event-stream") self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body)
def latest_user(request): return [item["content"] for item in request["input"] if item.get("role") == "user"][-1]
def kernels(daemon_pid): """Kernels this daemon owns. They are grandchildren: the runtime spawns ports through a helper, and another daemon's kernels are not ours to count.""" listing = subprocess.run(["ps", "-o", "pid=,ppid=,command=", "-ax"], capture_output=True, text=True) processes = {} for line in listing.stdout.splitlines(): fields = line.split(maxsplit=2) if len(fields) == 3 and fields[0].isdigit() and fields[1].isdigit(): processes[int(fields[0])] = (int(fields[1]), fields[2]) owned, frontier = [], [daemon_pid] while frontier: parent = frontier.pop() for pid, (ppid, command) in processes.items(): if ppid != parent: continue frontier.append(pid) if "albedo_kernel.py" in command: owned.append(pid) return owned
def run(endpoint): with tempfile.TemporaryDirectory(prefix="albedo-idle-reap-test-") as directory: home = Path(directory) / "home" workspace = Path(directory) / "workspace" home.mkdir(mode=0o700) workspace.mkdir() (home / "config.json").write_text(json.dumps({ "active": "test", "providers": {"test": { "baseUrl": endpoint, "apiKey": "fixture-key", "model": "fixture-model", "protocol": "responses", }}, })) env = dict(os.environ, ALBEDO_HOME=str(home), ALBEDO_PARENT_PID=str(os.getpid()), ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS)) connection = None
def cli(*args): result = subprocess.run(["node", "cli/bin/albedo.mjs", *args], cwd=ROOT, env=env, text=True, capture_output=True, timeout=45) assert result.returncode == 0, result.stdout + result.stderr return result.stdout
def api(path, body=None): request = urllib.request.Request( f"http://127.0.0.1:{connection['port']}" + path, headers={"Authorization": "Bearer " + connection["token"], "Content-Type": "application/json"}, data=None if body is None else json.dumps(body).encode()) return urllib.request.urlopen(request, timeout=15)
def ready(session_id): deadline = time.monotonic() + 20 while time.monotonic() < deadline: with api(f"/sessions/{session_id}/status") as response: if not json.load(response)["running"]: return time.sleep(0.05) raise AssertionError("session did not settle")
def events(session_id): # From sequence zero: notes live in the session's event ring. A client # with no cursor is served the transcript rebuilt from the database. with api(f"/sessions/{session_id}/stream?after_seq=0") as response: while line := response.readline(): if line.startswith(b"data: "): return json.loads(line[6:])["events"] raise AssertionError("stream ended without a snapshot")
try: cli("sessions") connection = json.loads((home / "daemon.json").read_text()) daemon = connection["pid"]
# A restored session costs nothing until it runs something. session_id = json.loads(cli("new", str(workspace)))["session"] assert kernels(daemon) == [], "a new session opened a kernel before running anything"
with api(f"/sessions/{session_id}/events", {"content": "remember this"}): pass ready(session_id) assert len(kernels(daemon)) == 1, "running a cell did not open exactly one kernel"
# Nothing is attached, so the kernel goes once the limit passes. deadline = time.monotonic() + IDLE_SECONDS * 4 while time.monotonic() < deadline and kernels(daemon): time.sleep(0.5) assert kernels(daemon) == [], "idle kernel outlived its limit"
note = [event for event in events(session_id) if event.get("type") == "note"] assert note and "released" in note[-1]["text"], note assert "2 variables saved to disk" in note[-1]["text"], note[-1]["text"] assert (home / "kernels" / f"{session_id}.state").exists(), "no state file was written"
# Reattaching revives the variables and tells the model what came back. with api(f"/sessions/{session_id}/events", {"content": "recall them"}): pass ready(session_id) prompt = latest_user(Provider.requests[-1]) assert prompt.startswith("recall them<system-note>"), prompt assert "restored from disk: answer, rows" in prompt, prompt assert any("recalled 7 [1, 2, 3]" in item.get("output", "") for item in Provider.requests[-1]["input"]), "restored variables were not usable"
texts = [event["text"] for event in events(session_id) if event.get("type") == "note"] assert any("restored 2 variables from disk" in text for text in texts), texts finally: if connection: with contextlib.suppress(Exception): api("/shutdown", {}).close() deadline = time.monotonic() + 10 while time.monotonic() < deadline: try: os.kill(connection["pid"], 0) except ProcessLookupError: break time.sleep(0.05) else: os.kill(connection["pid"], 9)
if __name__ == "__main__": server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Provider) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() try: run(f"http://127.0.0.1:{server.server_port}/v1") print("idle reap and state restore integration passed") finally: server.shutdown()