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.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346"""Lazy kernels: a turn holds no Python kernel until its first tool call."""
import jsonfrom pathlib import Pathimport subprocessimport threadingimport timeimport unittest
from harness import Albedo, Provider, exclusive, python, textfrom kernel_support import editable_daemon
IDLE_SECONDS = 2
def kernels(home): """Pids of the kernels whose run directory is under this daemon's home.""" listing = subprocess.run( ["ps", "-axo", "pid=,command="], capture_output=True, text=True, check=True ).stdout return [ int(line.split(None, 1)[0]) for line in listing.splitlines() if "albedo_kernel.py" in line and f"{home}/run/" in line ]
def wait_until(predicate, timeout=20): deadline = time.monotonic() + timeout while time.monotonic() < deadline: if predicate(): return True time.sleep(0.05) return predicate()
def inputs(record): return record["request"]["input"]
def is_note(item): return item.get("role") == "user" and str(item.get("content", "")).startswith( "<system-note>" )
def prompt_of(request): """The latest user prompt, not a system note that follows a tool result.""" users = [ item for item in request["input"] if item.get("role") == "user" and not is_note(item) ] return str(users[-1].get("content", "")).split("<system-note>")[0] if users else ""
# exclusive: lazy kernels and the idle timeout are daemon-wide, and the tests# count every kernel the daemon runs@exclusiveclass KernelLazyTests(unittest.TestCase): def setUp(self): self.model_gate = threading.Event() self.model_waiting = threading.Event() self.cells = {} self.follow = {}
def script(request): last = request["input"][-1] prompt = prompt_of(request) if last.get("role") == "user" and not is_note(last): if prompt in self.cells: return python(self.cells[prompt]) if prompt.startswith("wait for the model"): self.model_waiting.set() self.model_gate.wait(60) return text("waited") outputs = [ item for item in request["input"] if item.get("type") == "function_call_output" ] follow = self.follow.get(prompt, []) if prompt in self.cells and len(outputs) <= len(follow): # Each follow-up cell waits on the model first when asked to. wait, code = follow[len(outputs) - 1] if outputs else (False, None) if code is not None: if wait: self.model_waiting.set() self.model_gate.wait(60) return python(code) return text("ok")
self.provider = Provider(script) self.addCleanup(self.provider.close)
def start(self, slow_marker=None, crash_marker=None, **env): def prepare(app): # The spare pool stays off unless a test asks for it: the edited # kernel scripts below run only in a kernel booted cold, and the # idle reap counts kernel processes. app.env.update( { "ALBEDO_TEST_TIME_SCALE": "100", "ALBEDO_IDLE_SECONDS": str(IDLE_SECONDS), "ALBEDO_KERNEL_SPARES": "0", **env, } ) if slow_marker is not None or crash_marker is not None: launcher, bundle = editable_daemon(app.root) kernel = bundle / "albedo_kernel.py" prelude = "import os as _os, sys as _sys, time as _time\n" if slow_marker is not None: prelude += f"if _os.path.exists({str(slow_marker)!r}):\n _time.sleep(3)\n" if crash_marker is not None: prelude += f"if _os.path.exists({str(crash_marker)!r}):\n _sys.exit(3)\n" kernel.write_text(prelude + kernel.read_text()) app.env.update(ALBEDO_DAEMON=str(launcher))
app = Albedo(self.provider, protocol="responses", prepare=prepare) app.__enter__() self.addCleanup(app.__exit__, None, None, None) return app
def turn(self, app, session, prompt): app.prompt(session, prompt).close() app.idle(session)
def test_text_turns_and_model_waits_hold_no_kernel(self): app = self.start() existing = set(kernels(app.home)) session = app.session() self.turn(app, session, "just talk") self.assertEqual(set(kernels(app.home)) - existing, set())
self.cells["remember"] = "answer = 7\nrows = [1, 2, 3]" self.turn(app, session, "remember") # The scaled idle window may already have released the kernel the cell # opened, so it is counted by the state its release saved. # Nothing ran before, so the first kernel has no namespace to call lost. self.assertFalse( any(is_note(item) for item in inputs(self.provider.requests[-1])), inputs(self.provider.requests[-1]), ) self.assertTrue( wait_until(lambda: not set(kernels(app.home)) - existing), "idle kernel outlived its limit", ) # The cell's kernel held the namespace: its release saved it. self.assertTrue((app.home / "kernels" / f"{session}.state").exists())
app.prompt(session, "wait for the model").close() self.assertTrue(self.model_waiting.wait(30)) time.sleep(0.5) self.assertEqual( set(kernels(app.home)) - existing, set(), "a turn waiting on its model must not hold a kernel", ) self.model_gate.set() app.idle(session)
self.cells["recall"] = "print('recalled', answer, rows)" self.turn(app, session, "recall") after = inputs(self.provider.requests[-1]) self.assertTrue( any("recalled 7 [1, 2, 3]" in item.get("output", "") for item in after), after, ) # The restore lost nothing, so the model is told nothing. notices = [ json.dumps(item) for item in after if item.get("role") == "user" and "<system-note>" in json.dumps(item) ] self.assertEqual(notices, [], after)
def test_a_restore_with_no_variables_says_nothing_unless_it_lost_one(self): app = self.start() existing = set(kernels(app.home)) session = app.session()
def released_then(prompt, code): self.assertTrue( wait_until(lambda: not set(kernels(app.home)) - existing), "idle kernel outlived its limit", ) self.cells[prompt] = code self.turn(app, session, prompt) return [ str(item.get("content")) for item in inputs(self.provider.requests[-1]) if is_note(item) ]
self.cells["print only"] = "print('nothing to keep')" self.turn(app, session, "print only") # Saved with no variables and nothing missed: a lossless restore. self.assertEqual(released_then("again", "print('again')"), []) # Saved with no variables but one it could not keep: still a loss. self.cells["unpicklable"] = "import threading\nlock = threading.Lock()" self.turn(app, session, "unpicklable") notes = released_then("after the lock", "print('after')") self.assertEqual(len(notes), 1, notes) self.assertIn("lock", notes[0])
def test_interrupt_while_the_kernel_starts_runs_no_cell(self): slow = Path("/tmp") / f"albedo-slow-{time.monotonic_ns()}" self.addCleanup(lambda: slow.unlink(missing_ok=True)) app = self.start(slow_marker=slow) session = app.session() written = app.workspace / "written" self.cells["write"] = f"open({str(written)!r}, 'w').write('x')" slow.touch() start = len(self.provider.requests) with app.prompt(session, "write") as response: input_id = json.load(response)["id"] self.assertTrue(wait_until(lambda: len(self.provider.requests) > start)) with app.api(f"/sessions/{session}?tail=0") as response: captured = json.load(response) # The cell is now waiting on the slow kernel. time.sleep(0.5) interrupted = time.monotonic() app.api( f"/sessions/{session}/interrupt", { "run_id": captured["status"]["run_id"], "through_input_order": captured["input_order"], }, ).close()
def receipt(): with app.api(f"/sessions/{session}/inputs/{input_id}") as response: return json.load(response).get("turn") or {}
# The run ends at once, not when the kernel it was waiting on starts. self.assertTrue( wait_until(lambda: receipt().get("state") == "interrupted", timeout=2) ) self.assertLess(time.monotonic() - interrupted, 2.0) app.idle(session) time.sleep(4) self.assertFalse(written.exists(), "a cancelled turn ran its cell") slow.unlink() self.cells["after"] = "print('after', 6 * 7)" self.turn(app, session, "after") self.assertTrue( any( "after 42" in item.get("output", "") for item in inputs(self.provider.requests[-1]) ) )
def test_a_long_model_wait_mid_turn_releases_and_restores_the_kernel(self): # A zero memory budget is constant pressure, under which a long model # wait releases the kernel; the time scale makes the 3 s rest 30 ms. app = self.start(ALBEDO_KERNEL_BUDGET_MB="0", ALBEDO_KERNEL_SPARES="1") session = app.session() place = app.workspace / "place" place.mkdir() self.cells["work"] = f"import os\nanswer = 7\nos.chdir({str(place)!r})" self.follow["work"] = [ (True, "import os\nprint('kept', answer, os.getcwd())"), (False, "print('again', answer)"), ]
def kernel_state(): with app.api(f"/sessions/{session}?tail=0") as response: return json.load(response)["kernel"]["state"]
app.prompt(session, "work").close() self.assertTrue(self.model_waiting.wait(30)) self.assertTrue( wait_until(lambda: kernel_state() == "none"), "a turn waiting on its model past the rest time kept its kernel", ) self.model_gate.set() app.idle(session) outputs = [ item.get("output", "") for item in inputs(self.provider.requests[-1]) if item.get("type") == "function_call_output" ] self.assertTrue(any("kept 7" in output for output in outputs), outputs) self.assertTrue( any( f"{place.resolve()}" in output or str(place) in output for output in outputs ), outputs, ) self.assertTrue(any("again 7" in output for output in outputs), outputs)
def test_a_kernel_that_dies_while_starting_fails_only_its_call(self): crash = Path("/tmp") / f"albedo-crash-{time.monotonic_ns()}" self.addCleanup(lambda: crash.unlink(missing_ok=True)) app = self.start(crash_marker=crash) session = app.session() self.cells["boom"] = "print('ran', 6 * 7)" crash.touch() self.turn(app, session, "boom") outputs = [ item.get("output", "") for item in inputs(self.provider.requests[-1]) if item.get("type") == "function_call_output" ] self.assertEqual(len(outputs), 1, outputs) self.assertIn("could not start the session python kernel", outputs[0]) self.assertNotIn("ran 42", outputs[0]) crash.unlink() self.cells["again"] = "print('ran', 6 * 7)" self.turn(app, session, "again") self.assertTrue( any( "ran 42" in item.get("output", "") for item in inputs(self.provider.requests[-1]) ) )
def test_a_kernel_lost_mid_turn_is_replaced_at_the_next_call(self): app = self.start() session = app.session() self.cells["die"] = "import os; os._exit(1)" self.follow["die"] = [(False, "print('again', 6 * 7)")] self.turn(app, session, "die") request = inputs(self.provider.requests[-1]) outputs = [ item.get("output", "") for item in request if item.get("type") == "function_call_output" ] self.assertEqual(len(outputs), 2, outputs) self.assertIn("kernel lost", outputs[0]) self.assertIn("again 42", outputs[1]) self.assertTrue( any( "python kernel got reset" in str(item.get("content")) for item in request ), request, )
if __name__ == "__main__": unittest.main()