okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

Aly Raffauffix(kernel): preserve statement identity during definition capture1c00b60

main
59.5 KiB1643 linesraw
1from __future__ import annotations
2
3import ast
4import albedo_api
5import albedo_output
6import albedo_bundle
7import albedo_state
8from albedo_values import Unencodable, decode, encode
9import albedo_capture
10from albedo_capture import Capture, PREVIEW, MAX_IMAGE_BYTES
11import albedo_link
12import albedo_memory
13from albedo_protocol import parse_incoming
14import albedo_proc
15import albedo_shell
16from collections.abc import Awaitable, Callable, Iterable, Sequence
17from typing import cast
18from types import CodeType, FrameType, ModuleType
19import albedo_trace
20import asyncio
21import builtins
22import codecs
23import collections
24import contextvars
25import dataclasses
26import inspect
27import importlib
28import io
29import json
30import os
31import pathlib
32import reprlib
33import select
34import shutil
35import signal
36import sys
37import tempfile
38import threading
39import time
40import traceback
41import uuid
42import weakref
43
44CLEANUP_DEADLINE = (
45    1.5  # seconds: plugins must finish cleanup inside the supervisor's patience
46)
47LIMITS = {"cell": 16, "job": 64, "native": 1}  # retained captures per kind
48CELL: contextvars.ContextVar[Capture | None] = contextvars.ContextVar(
49    "cell", default=None
50)
51LINK: albedo_link.StdioLink | albedo_link.SocketLink
52RUN_DIR: str | None = None  # a detached kernel's run directory, None over stdio
53WATCH_INTERVAL = 1.0  # seconds between a detached kernel's lifetime checks
54FIRST_ATTACH_GRACE = 60.0  # seconds a detached kernel waits for its first attach
55LOOP = asyncio.new_event_loop()
56asyncio.set_event_loop(LOOP)
57
58albedo_proc.reap_stopped_safely(LOOP)
59QUEUE: asyncio.Queue[albedo_api.Execute | albedo_api.State] = asyncio.Queue()
60PENDING: dict[str, asyncio.Future[albedo_api.HostReply]] = {}
61# Host calls a synchronous API waits for; the reader thread hands their reply
62# over directly, since the loop it would go through is the one waiting.
63WAITING: dict[str, tuple[threading.Event, list[albedo_api.HostReply]]] = {}
64HOST_NOW_TIMEOUT = 30.0
65HOST_SLOTS = asyncio.Semaphore(32)
66OWNER_CALLS = asyncio.Semaphore(32)  # concurrent tool calls from the connection owner
67OWNER_TASKS: dict[str, asyncio.Task[object]] = {}  # interruptable by invoke id
68LIVE: collections.OrderedDict[str, object] = (
69    collections.OrderedDict()
70)  # remote references
71LIVE_LIMIT = 64  # live references retained for the owner; LRU beyond that
72PENDING_OBJECTS: collections.OrderedDict[str, asyncio.Future[tuple[bool, object]]] = (
73    collections.OrderedDict()
74)
75MIRROR_INTERVAL = 0.15  # seconds between output-tail mirror frames while data flows
76MIRROR_TAIL = PREVIEW  # the same bounded tail window as a local job
77BACKGROUND_SECONDS = max(
78    0.01, float(os.environ.get("ALBEDO_CELL_BACKGROUND_SECONDS", "60"))
79)
80CELL_TASKS: dict[str, asyncio.Task[dict[str, object]]] = {}
81CELL_CAPTURES: dict[str, Capture] = {}
82CELL_STARTS: dict[str, int] = {}  # the kernel's footprint as each running cell began
83BACKGROUND_RESULTS: dict[str, dict[str, object]] = {}
84READ_CELLS: set[str] = set()
85CELL_ANNOUNCER: asyncio.Task[None] | None = None
86interrupt_capture: Capture | None = None
87ARCHIVES: collections.OrderedDict[str, Capture] = collections.OrderedDict()
88# Finished executions by cell id: a replayed execute answers with its result
89# instead of running the cell twice. Queued or running ones are in EXECUTING.
90FINISHED: collections.OrderedDict[str, dict[str, object]] = collections.OrderedDict()
91FINISHED_LIMIT = 16
92# Base64 image text kept across finished results; older results drop theirs first.
93FINISHED_IMAGE_BYTES = 8 * 1024 * 1024
94EXECUTING: set[str] = set()
95# Interrupts that arrived while their cell was still queued, by cell id.
96INTERRUPTS: dict[str, str] = {}
97READY = threading.Event()
98NAMESPACE: dict[str, object] = {"__name__": "__main__"}
99DEFINITIONS: collections.OrderedDict[str, str] = collections.OrderedDict()
100CLEANUP: list[albedo_api.Cleanup] = []
101STOP_HANDLERS: list[Callable[[str], Awaitable[dict[str, object] | None]]] = []
102HANDLES: list[type[object]] = [asyncio.Task]
103REPR = reprlib.Repr()
104REPR.maxstring = REPR.maxother = 4000
105REPR.maxdict = REPR.maxlist = REPR.maxtuple = 50
106
107
108def forget_output(id: str) -> None:
109    ARCHIVES.pop(id, None)
110    READ_WATCHERS.pop(id, None)
111
112
113def show(value: object) -> str:
114    return (
115        REPR.repr(value)
116        .encode("utf-8", errors="replace")[:PREVIEW]
117        .decode("utf-8", errors="ignore")
118    )
119
120
121def send(value: dict[str, object]) -> None:
122    LINK.send(value)
123
124
125async def cleanup() -> None:
126    """Run plugin cleanup; a callback may return an awaitable to clean up asynchronously."""
127    for close in reversed(CLEANUP):
128        try:
129            result = close()
130            if inspect.isawaitable(result):
131                await result
132        except BaseException as error:
133            try:
134                send(
135                    {
136                        "type": "cleanup",
137                        "module": getattr(close, "__module__", "plugin"),
138                        "failures": [f"{type(error).__name__}: {error}"],
139                    }
140                )
141            except (OSError, ValueError):
142                pass  # a disconnected owner must not prevent other cleanup callbacks
143
144
145def die() -> None:
146    """Leave the session: clean up while the loop can still run it, then kill our group.
147
148    Called from the reader thread and from the main thread, never from a loop
149    callback, so cleanup may await without deadlocking the interpreter.
150    """
151    try:
152        if LOOP.is_running():
153            asyncio.run_coroutine_threadsafe(cleanup(), LOOP).result(CLEANUP_DEADLINE)
154        else:
155            LOOP.run_until_complete(asyncio.wait_for(cleanup(), CLEANUP_DEADLINE))
156    except BaseException as error:
157        try:
158            send(
159                {
160                    "type": "cleanup",
161                    "module": "kernel",
162                    "failures": [
163                        f"cleanup incomplete: {type(error).__name__}: {error}"
164                    ],
165                }
166            )
167        except (OSError, ValueError):
168            pass  # the owner disconnected; it independently tracks our groups
169    if RUN_DIR is not None:
170        # No socket left behind: an attach from now on learns we are gone.
171        shutil.rmtree(RUN_DIR, ignore_errors=True)
172    # The interpreter and its subprocesses own a separate process group.
173    os.killpg(os.getpid(), signal.SIGKILL)
174
175
176def receive(frame: object) -> None:
177    """Act on one frame from the owner, on the link's reader thread."""
178    global interrupt_capture
179    try:
180        message = parse_incoming(frame)
181    except ValueError:
182        die()
183        return
184    if message["type"] == "shutdown":
185        die()
186    elif message["type"] == "interrupt":
187        capture = CELL_CAPTURES.get(message["id"])
188        if capture is not None:
189            interrupt_capture = capture
190            capture.interruption = message.get("reason", "cancelled")
191            os.kill(os.getpid(), signal.SIGINT)
192            _ = LOOP.call_soon_threadsafe(interrupt_queued, message)
193        else:
194            # The loop takes it after the execute it may name, not started yet.
195            _ = LOOP.call_soon_threadsafe(deliver, message)
196    elif message["type"] == "reply" and message["id"] in WAITING:
197        arrived, answer = WAITING[message["id"]]
198        answer.append(message["value"])
199        arrived.set()
200    else:
201        _ = LOOP.call_soon_threadsafe(deliver, message)
202
203
204def reader() -> None:
205    try:
206        LINK.serve(receive)
207    except Exception:
208        pass
209    die()
210
211
212def watch() -> None:
213    """A detached kernel ends itself: when its run directory is gone (its home
214    was removed), or when nothing attached for the grace period and no job or
215    cell is still running."""
216    assert isinstance(LINK, albedo_link.SocketLink) and RUN_DIR is not None
217    while True:
218        time.sleep(WATCH_INTERVAL)
219        if not os.path.isdir(RUN_DIR):
220            die()
221        busy = bool(CELL_TASKS) or not QUEUE.empty() or LINK.jobs.live() > 0
222        if not busy and LINK.idle_for() > LINK.grace:
223            die()
224
225
226def guard_memory() -> None:
227    """Interrupts every running cell once the kernel's footprint passes its
228    threshold (albedo_memory), so a runaway cell ends instead of the kernel.
229    A thread, like the deadline timer: a synchronous cell holds the loop."""
230    while True:
231        time.sleep(albedo_memory.INTERVAL)
232        if not CELL_STARTS:
233            continue
234        used = albedo_memory.footprint()
235        if used is None:
236            return
237        for id, start in list(CELL_STARTS.items()):
238            capture = CELL_CAPTURES.get(id)
239            if capture is None or capture.interruption == "memory":
240                continue  # already told; it is unwinding
241            if used > albedo_memory.threshold(start):
242                receive({"type": "interrupt", "id": id, "reason": "memory"})
243
244
245def hello() -> dict[str, object]:
246    """What a detached kernel tells each attach about itself."""
247    pgid = os.getpgid(0)
248    return {
249        "pid": os.getpid(),
250        "pgid": pgid if pgid == os.getpid() else None,
251        "leader": albedo_proc.leader_token(os.getpid()),
252        "ready": READY.is_set(),
253        "cells": len(CELL_TASKS),
254    }
255
256
257def deliver(message: albedo_api.Incoming) -> None:
258    if message["type"] == "interrupt":
259        interrupt_queued(message)
260    elif message["type"] == "reply":
261        future = PENDING.pop(message["id"], None)
262        if future is not None and not future.done():
263            future.set_result(message["value"])
264    elif message["type"] == "invoke":
265        _ = LOOP.create_task(serve_invoke(message))
266    elif message["type"] == "introspect":
267        send(
268            {
269                "type": "introspected",
270                "id": message["id"],
271                "names": sorted(
272                    name
273                    for name in NAMESPACE
274                    if name.isidentifier() and not name.startswith("_")
275                ),
276                "handles": [
277                    {"id": key, "repr": show(value)} for key, value in LIVE.items()
278                ],
279            }
280        )
281    elif message["type"] == "release":
282        LIVE.pop(message["handle"], None)
283    elif message["type"] == "job_stop":
284        _ = LOOP.create_task(serve_job_stop(message["id"]))
285    elif message["type"] == "snapshot" and (bool(CELL_TASKS) or not QUEUE.empty()):
286        # A namespace mid-cell is no state to carry: the owner tries again later.
287        send(state_reply(message, {"error": "a cell is still running"}))
288    elif message["type"] == "execute" and message["id"] in FINISHED:
289        send(FINISHED[message["id"]])
290    elif message["type"] == "execute" and message["id"] in EXECUTING:
291        pass  # already queued or running; its done frame answers both
292    else:
293        if message["type"] == "execute":
294            EXECUTING.add(message["id"])
295        QUEUE.put_nowait(cast(albedo_api.Execute | albedo_api.State, message))
296
297
298def interrupt_queued(message: albedo_api.Interrupt) -> None:
299    """An interrupt the reader thread found no running cell for. Here on the
300    loop, the cell it names is still queued, or has started since and waits
301    at an await, where cancelling reaches it."""
302    reason = message.get("reason", "cancelled")
303    task = OWNER_TASKS.get(message["id"])
304    capture = CELL_CAPTURES.get(message["id"])
305    if task is not None:
306        _ = task.cancel()
307    elif capture is not None:
308        capture.interruption = reason
309        _ = CELL_TASKS[message["id"]].cancel()
310    elif message["id"] in EXECUTING:
311        INTERRUPTS[message["id"]] = reason
312
313
314class WorkError(Exception):
315    def __init__(self, code: str, message: str) -> None:
316        super().__init__(message)
317        self.code: str = code
318
319
320async def host(method: str, args: dict[str, object]) -> object:
321    async with HOST_SLOTS:
322        return await _host(method, args)
323
324
325async def _host(method: str, args: dict[str, object]) -> object:
326    key = uuid.uuid4().hex
327    future: asyncio.Future[albedo_api.HostReply] = LOOP.create_future()
328    PENDING[key] = future
329    send({"type": "call", "id": key, "method": method, "args": args})
330    try:
331        answer = await future
332    finally:
333        _ = PENDING.pop(key, None)
334    if answer["ok"]:
335        return answer["value"]
336    raise WorkError(answer["code"], answer["message"])
337
338
339def host_now(method: str, args: dict[str, object]) -> object:
340    """A host call for a synchronous API: blocks the calling thread, the
341    loop's included, until the reply. Only for calls the host answers at
342    once; an interrupt still ends the wait."""
343    key = uuid.uuid4().hex
344    arrived = threading.Event()
345    answers: list[albedo_api.HostReply] = []
346    WAITING[key] = (arrived, answers)
347    send({"type": "call", "id": key, "method": method, "args": args})
348    try:
349        if not arrived.wait(HOST_NOW_TIMEOUT):
350            raise WorkError(
351                "timeout", f"{method} had no answer within {HOST_NOW_TIMEOUT:.0f}s"
352            )
353    finally:
354        _ = WAITING.pop(key, None)
355    answer = answers[0]
356    if answer["ok"]:
357        return answer["value"]
358    raise WorkError(answer["code"], answer["message"])
359
360
361def attach_image(data: bytes) -> str:
362    """Attach image bytes to the running cell; plugins reach this as api.attach_image."""
363    capture = CELL.get()
364    # Spawned tasks inherit context even after their creating cell finishes.
365    if capture is None or capture.kind != "cell" or capture.id not in ACTIVE_CAPTURES:
366        raise RuntimeError(
367            "images attach to a running cell; background work has no result to carry them"
368        )
369    return capture.attach(bytes(data))
370
371
372def show_image(
373    source: bytes | bytearray | memoryview | str | os.PathLike[str],
374) -> albedo_api.Text:
375    """Return an image to yourself with this cell's result. `source` is PNG,
376    JPEG, or WebP bytes, or a path to such a file. At most 4 images and 5 MiB
377    per cell, each within the model's edge limit, or it raises ValueError; the
378    image arrives after the cell finishes, not mid-cell."""
379    if isinstance(source, (bytes, bytearray, memoryview)):
380        data = bytes(source)
381    elif isinstance(source, (str, os.PathLike)):
382        with open(source, "rb") as file:
383            data = file.read(MAX_IMAGE_BYTES + 1)
384    else:
385        raise TypeError(
386            f"show_image takes image bytes or a path, not {type(source).__name__}"
387        )
388    return albedo_api.Text("attached " + attach_image(data))
389
390
391NATIVE = Capture("native", "native")
392NATIVE_LOCK = threading.RLock()
393NATIVE_FD = -1
394DECODER = codecs.getincrementaldecoder("utf-8")("replace")
395NATIVE_SINKS: dict[asyncio.Task[object], Capture] = {}
396ACTIVE_CAPTURES: set[str] = set()
397
398
399def retained(id: str) -> Capture:
400    capture = ARCHIVES.get(id)
401    if capture is None:
402        raise LookupError(
403            f"no retained output for {id!r}: only the 16 most recent cells and 64 most recent jobs "
404            f"keep output, so it may have rolled out; output.list() names what is kept, and a "
405            f"background job keeps its own output on its handle: jobs[{id!r}].tail()"
406        )
407    return capture
408
409
410READ_WATCHERS: dict[str, Callable[[], None] | weakref.WeakMethod] = {}
411
412
413def watch_output(id: str, notify: Callable[[], None]) -> None:
414    """Notify reads while the output is retained, without owning a bound method's job."""
415
416    def expired(reference: weakref.WeakMethod) -> None:
417        if READ_WATCHERS.get(id) is reference:
418            READ_WATCHERS.pop(id, None)
419
420    try:
421        READ_WATCHERS[id] = weakref.WeakMethod(notify, expired)
422    except TypeError:
423        # Functions and other callables need a strong reference to stay usable.
424        READ_WATCHERS[id] = notify
425
426
427class Output:
428    def read(self, id: str, offset: int = 0, limit: int = 4000) -> albedo_api.Text:
429        """Read retained output by cell or job id, at most 64 KiB per call.
430        The 16 most recent cells and 64 most recent jobs keep up to 1 MiB
431        each; page through it with offset."""
432        if limit > PREVIEW:
433            raise ValueError(
434                f"limit is at most {PREVIEW} per read; page through the rest with offset="
435            )
436        capture = retained(id)
437        result = albedo_output.read_page(capture, offset, limit)
438        notify = READ_WATCHERS.get(id)
439        if isinstance(notify, weakref.WeakMethod):
440            notify = notify()
441        if notify is not None:
442            try:
443                notify()
444            except Exception:
445                pass  # a plugin's bookkeeping must not fail the read
446        return result
447
448    def list(self):
449        """Every retained channel: cells, background jobs, and 'native'.
450
451        'native' holds bytes written to fd 1/2 while no cell was running; bytes a
452        cell's subprocesses write land in that cell's own output instead.
453        """
454        return albedo_api.ReadyList(
455            {"id": c.id, "kind": c.kind, "bytes": c.seen, "retained": len(c.data)}
456            for c in ARCHIVES.values()
457        )
458
459
460class Stream(io.TextIOBase):
461    def writable(self) -> bool:
462        return True
463
464    def write(self, text: str) -> int:
465        with NATIVE_LOCK:  # bytes already in the pipe were written first and stay first
466            _ = drain()
467            (CELL.get() or NATIVE).write(text)
468        return len(text)
469
470    def flush(self) -> None:
471        pass
472
473    def fileno(self) -> int:
474        return 1
475
476
477def drain() -> bool:
478    """Fold whatever fd 1/2 already holds into the running cell; False once the pipe closes.
479
480    Subprocesses and C extensions cannot be attributed through contextvars, so
481    they are credited to the cell that is running, the way a terminal does.
482    Both the pump thread and Python-level writes drain, so the two surfaces
483    interleave in the order the bytes were produced.
484    """
485    with NATIVE_LOCK:
486        while True:
487            try:
488                data = os.read(NATIVE_FD, 65536)
489            except BlockingIOError:
490                return True
491            except OSError:
492                return False
493            if not data:
494                return False
495            text = DECODER.decode(data)
496            if text:
497                sinks = list(NATIVE_SINKS.values())
498                (sinks[0] if len(sinks) == 1 else NATIVE).write(text)
499
500
501def native_output(fd: int) -> None:
502    """Wake on output produced while no cell is writing."""
503    try:
504        while select.select([fd], [], [])[0] and drain():
505            pass
506    except OSError:
507        return
508
509
510def interrupt(_signal: int, _frame: FrameType | None) -> None:
511    if CELL.get() is not None and CELL.get() is interrupt_capture:
512        raise KeyboardInterrupt()
513
514
515def compile_cell(
516    source: str, filename: str
517) -> tuple[CodeType, CodeType | None, ast.Module]:
518    """Compile every part before running; return the unmodified tree for state capture."""
519    tree = ast.parse(source, filename=filename)
520    original = ast.Module(body=list(tree.body), type_ignores=tree.type_ignores)
521    last = tree.body[-1] if tree.body else None
522    trailing = None
523    if isinstance(last, ast.Expr):
524        _ = tree.body.pop()
525        trailing = ast.Expression(last.value)
526    prefix = cast(
527        CodeType, compile(tree, filename, "exec", ast.PyCF_ALLOW_TOP_LEVEL_AWAIT)
528    )
529    suffix = (
530        cast(
531            CodeType,
532            compile(trailing, filename, "eval", ast.PyCF_ALLOW_TOP_LEVEL_AWAIT),
533        )
534        if trailing
535        else None
536    )
537    return prefix, suffix, original
538
539
540def _bound_names(node: ast.AST) -> set[str]:
541    names = {
542        item.id
543        for item in ast.walk(node)
544        if isinstance(item, ast.Name) and isinstance(item.ctx, (ast.Store, ast.Del))
545    }
546    for item in ast.walk(node):
547        if isinstance(item, ast.ExceptHandler) and item.name:
548            names.add(item.name)
549        elif isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
550            names.add(item.name)
551    return names
552
553
554def _drop_definition(name: str) -> None:
555    for key in list(DEFINITIONS):
556        if key == name or key.startswith(name + "#"):
557            DEFINITIONS.pop(key, None)
558
559
560def _definition_changes(
561    source: str, tree: ast.Module
562) -> list[tuple[ast.stmt, set[str], str | None, bool]]:
563    changes: list[tuple[ast.stmt, set[str], str | None, bool]] = []
564    for statement in tree.body:
565        names: set[str] = set()
566        definition = isinstance(
567            statement,
568            (
569                ast.FunctionDef,
570                ast.AsyncFunctionDef,
571                ast.ClassDef,
572                ast.Import,
573                ast.ImportFrom,
574            ),
575        )
576        if isinstance(statement, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
577            names.add(statement.name)
578        elif isinstance(statement, ast.Import):
579            names.update(
580                alias.asname or alias.name.split(".")[0] for alias in statement.names
581            )
582        elif isinstance(statement, ast.ImportFrom):
583            names.update(
584                alias.asname or alias.name
585                for alias in statement.names
586                if alias.name != "*"
587            )
588        else:
589            names = _bound_names(statement)
590        if names:
591            changes.append(
592                (
593                    statement,
594                    names,
595                    ast.get_source_segment(source, statement) if definition else None,
596                    isinstance(statement, ast.Import),
597                )
598            )
599    return changes
600
601
602def capture_definitions(
603    source: str, tree: ast.Module, upto_line: int | None = None
604) -> None:
605    for statement, names, segment, is_import in _definition_changes(source, tree):
606        if (
607            upto_line is not None
608            and (statement.end_lineno or statement.lineno) >= upto_line
609        ):
610            continue
611        for name in names:
612            # Plain dotted imports share their root binding and must accumulate.
613            plain_dotted = (
614                is_import
615                and any(
616                    alias.name.startswith(name + ".") and alias.asname is None
617                    for alias in statement.names
618                )
619                if isinstance(statement, ast.Import)
620                else False
621            )
622            if not plain_dotted:
623                _drop_definition(name)
624        if segment is not None:
625            key = next(iter(names), "definition")
626            if (
627                is_import
628                and isinstance(statement, ast.Import)
629                and any(
630                    "." in alias.name and alias.asname is None
631                    for alias in statement.names
632                )
633            ):
634                while key in DEFINITIONS:
635                    key += "#2"
636            elif key in DEFINITIONS:
637                _drop_definition(key)
638            # One source segment is replayed once; its imports bind every listed name.
639            if segment not in DEFINITIONS.values():
640                DEFINITIONS[key] = segment
641
642
643async def evaluate(source: str, cell_id: str, durable: bool = False) -> object:
644    started = False
645    ACTIVE_CAPTURES.add(cell_id)
646    try:
647        prefix, suffix, tree = compile_cell(source, "<albedo:" + cell_id + ">")
648        if durable:
649            _ = await host("cells.started", {"id": cell_id})
650        started = True
651        value = cast(object, eval(prefix, NAMESPACE))
652        if inspect.isawaitable(value):
653            _ = await cast(Awaitable[object], value)
654        if suffix is not None:
655            value = cast(object, eval(suffix, NAMESPACE))
656            if inspect.isawaitable(value) and not isinstance(value, tuple(HANDLES)):
657                value = await cast(Awaitable[object], value)
658            NAMESPACE["_"] = value
659        capture_definitions(source, tree)
660        return value
661    except BaseException as error:
662        filename = "<albedo:" + cell_id + ">"
663        line = next(
664            (
665                lineno
666                for frame, lineno in traceback.walk_tb(error.__traceback__)
667                if frame.f_code.co_filename == filename
668            ),
669            None,
670        )
671        if "tree" in locals() and line is not None:
672            capture_definitions(source, tree, line)
673        if not hasattr(error, "_albedo_cell_id"):
674            setattr(error, "_albedo_cell_id", cell_id)
675            setattr(error, "_albedo_started", started)
676        raise
677    finally:
678        ACTIVE_CAPTURES.discard(cell_id)
679
680
681class Cells:
682    """Immutable saved source. Repairs never automatically repeat side effects."""
683
684    last_id: str | None = None
685
686    async def read(
687        self,
688        id: str,
689        *,
690        start_line: object = 1,
691        end_line: object = None,
692        limit: object = 8000,
693    ) -> str:
694        """Return source lines, capped at `limit` bytes. Narrow with start_line/end_line."""
695        if (
696            not isinstance(start_line, int)
697            or start_line < 1
698            or (
699                end_line is not None
700                and (not isinstance(end_line, int) or end_line < start_line)
701            )
702        ):
703            raise ValueError("expected positive inclusive line numbers")
704        if not isinstance(limit, int) or limit < 1:
705            raise ValueError("limit must be a positive byte count")
706        cell_read(id)
707        cell = cast(albedo_api.SavedCell, await host("cells.read", {"id": id}))
708        lines = cell["source"].splitlines(keepends=True)
709        data = "".join(lines[start_line - 1 : end_line]).encode(
710            "utf-8", errors="replace"
711        )
712        if len(data) <= limit:
713            return data.decode("utf-8", errors="ignore")
714        return data[:limit].decode("utf-8", errors="ignore") + (
715            f"\n[{limit} of {len(data)} bytes; the cell has {len(lines)} lines; "
716            f"narrow with start_line/end_line or raise limit]"
717        )
718
719    async def info(self, id: str) -> albedo_api.Record:
720        """Status (ok, error, interrupted, started, saved, lost, unavailable),
721        parentage, whether the cell started, and how long a finished cell
722        ran, without its source."""
723        cell_read(id)
724        cell = cast(albedo_api.SavedCell, await host("cells.read", {"id": id}))
725        return albedo_api.Record(
726            (key, value) for key, value in cell.items() if key != "source"
727        )
728
729    async def cancel(self, id: str) -> bool:
730        """Request cancellation of a running cell and suppress its completion wake."""
731        task = CELL_TASKS.get(id)
732        if task is None or task.done():
733            cell_read(id)
734            return False
735        READ_CELLS.add(id)
736        CELL_CAPTURES[id].interruption = "cancelled"
737        return task.cancel()
738
739    async def list(self, limit: int = 20) -> albedo_api.ReadyList:
740        """This session's cells, newest first: id, status, parent, and first line.
741        Unlike output.list(), it reaches cells whose output has rolled out."""
742        cells = cast(
743            list[dict[str, object]], await host("cells.list", {"limit": limit})
744        )
745        return albedo_api.ReadyList(albedo_api.Record(cell) for cell in cells)
746
747    async def trace(self, id: str) -> dict[str, object]:
748        """What the cell read, ran and changed: activities plus diffs of files it wrote.
749
750        Diffs can be large; inspect the fields rather than printing the whole trace.
751        """
752        cell_read(id)
753        return cast(dict[str, object], await host("cells.trace", {"id": id}))
754
755    async def run(
756        self,
757        id: str,
758        replacements: Iterable[Sequence[object]] = (),
759        *,
760        allow_partial: bool = False,
761        check: bool = False,
762    ) -> object:
763        """Repair exact unique text and execute a new saved copy.
764
765        check=True compiles the rewrite and returns its compile error, or None when
766        it compiles; nothing is saved and nothing runs.
767        A cell that started needs allow_partial=True after inspecting its effects.
768        Syntax/compile failures are safe to repair without that override.
769        """
770        pairs = [tuple(pair) for pair in replacements]
771        for pair in pairs:
772            if len(pair) != 2 or not all(isinstance(s, str) for s in pair):
773                raise ValueError("replacements must be (old, new) string pairs")
774        if check:
775            source = cast(
776                str, await host("cells.draft", {"id": id, "replacements": pairs})
777            )
778            try:
779                _ = compile_cell(source, "<albedo:" + id + ":check>")
780            except (SyntaxError, ValueError) as error:
781                return "".join(
782                    traceback.format_exception_only(type(error), error)
783                ).strip()
784            return None
785        cell = cast(
786            albedo_api.SavedCell,
787            await host(
788                "cells.prepare",
789                {"id": id, "replacements": pairs, "allow_partial": allow_partial},
790            ),
791        )
792        self.last_id = cell["id"]
793        capture = Capture(cell["id"])
794        remember(capture)
795        outer = CELL.get()
796        token = CELL.set(capture)
797        outer_sink = swap_sink(capture)
798        began = LOOP.time()
799        value: object = None
800        error: BaseException | None = None
801        status = "ok"
802        try:
803            value = await evaluate(cell["source"], cell["id"], True)
804            text = "" if value is None else show(value)
805        except BaseException as failure:
806            error = failure
807            status = (
808                "interrupted"
809                if isinstance(failure, (KeyboardInterrupt, asyncio.CancelledError))
810                else "error"
811            )
812            capture.write(traceback.format_exc())
813            text = ""
814        finally:
815            _ = drain()
816            _ = swap_sink(outer_sink)
817            CELL.reset(token)
818            if outer is not None:
819                outer.write("[cell " + cell["id"] + "]\n" + capture.read(0, PREVIEW))
820                for image in capture.images:
821                    try:
822                        _ = outer.attach(image)
823                    except ValueError as dropped:
824                        outer.write(
825                            f"[image from cell {cell['id']} dropped: {dropped}]\n"
826                        )
827        deliver_trace(capture)
828        _ = await host(
829            "cells.finish",
830            {
831                "id": cell["id"],
832                "outcome": {
833                    "id": cell["id"],
834                    "status": status,
835                    "duration": round(LOOP.time() - began, 3),
836                    "output": capture.preview(status),
837                    "value": text,
838                    "truncated": capture.seen > PREVIEW,
839                    "images": capture.encoded_images(),
840                },
841            },
842        )
843        capture.images.clear()
844        if error is not None:
845            raise error
846        return value
847
848
849INJECTED: set[str] = {"__name__", "__builtins__", "_"}
850
851
852def swap_sink(capture: Capture | None) -> Capture | None:
853    task = cast(asyncio.Task[object], asyncio.current_task())
854    with NATIVE_LOCK:
855        previous = NATIVE_SINKS.pop(task, None)
856        if capture is not None:
857            NATIVE_SINKS[task] = capture
858        return previous
859
860
861def background_capture(id: str) -> Capture:
862    """A retained channel for a plugin's background work, addressable through `output`."""
863    capture = Capture(id, "job")
864    remember(capture)
865    return capture
866
867
868def remember(capture: Capture) -> None:
869    """Keep the newest captures of each kind, so a flood of jobs never evicts cell output."""
870    ARCHIVES[capture.id] = capture
871    ARCHIVES.move_to_end(capture.id)
872    counts = collections.Counter(held.kind for held in ARCHIVES.values())
873    for key, held in list(ARCHIVES.items()):
874        if counts[held.kind] > LIMITS[held.kind] and key not in EXECUTING:
875            forget_output(key)
876            counts[held.kind] -= 1
877
878
879def state_reply(
880    message: albedo_api.State, state: dict[str, object] | None = None
881) -> dict[str, object]:
882    if state is None:
883        state = (
884            albedo_state.save_state(message["path"], NAMESPACE, DEFINITIONS, INJECTED)
885            if message["type"] == "snapshot"
886            else albedo_state.load_state(message["path"], NAMESPACE, DEFINITIONS)
887        )
888    return {
889        "type": "done",
890        "id": message["id"],
891        "status": "error" if "error" in state else "ok",
892        "output": "",
893        "value": "",
894        "truncated": False,
895        "state": state,
896    }
897
898
899MAX_RESULT_BYTES = (
900    4 * 1024 * 1024
901)  # one invoke reply ceiling, below the 8 MiB frame guard
902
903
904def retain(value: object) -> str:
905    """Keep one live object addressable by the owner; identity-stable, oldest rolls out."""
906    for key, held in LIVE.items():
907        if held is value:
908            LIVE.move_to_end(key)
909            return key
910    key = uuid.uuid4().hex
911    LIVE[key] = value
912    LIVE.move_to_end(key)
913    while len(LIVE) > LIVE_LIMIT:
914        LIVE.popitem(last=False)
915    capture = getattr(value, "capture", None)
916    if isinstance(capture, Capture):
917        _ = LOOP.create_task(_mirror(key, value, capture))
918    return key
919
920
921def _mirror_state(obj: object) -> dict[str, object]:
922    """One snapshot of a captured object: what the owner's tail() and poll() answer from."""
923    capture = cast(Capture, getattr(obj, "capture"))
924    tail = capture.tail(MIRROR_TAIL)
925    return {
926        "job": getattr(obj, "id", None),
927        "seen": capture.seen,
928        "tail": tail.decode("utf-8", errors="ignore"),
929        "tail_bytes": len(tail),
930        "exit_code": getattr(obj, "exit_code", None),
931        "timed_out": getattr(obj, "timed_out", False),
932        "duration": getattr(obj, "duration", None),
933    }
934
935
936async def _mirror(key: str, obj: object, capture: Capture) -> None:
937    """Stream one live reference's output tail, so the owner reads it without a round trip.
938
939    Frames are coalesced to the mirror interval; the final frame carries the
940    completion state, and the pump ends once the object has settled and quieted.
941    """
942    seen, quiet = -1, 0
943    while LIVE.get(key) is obj:
944        await asyncio.sleep(MIRROR_INTERVAL)
945        changed = capture.seen != seen
946        if changed:
947            seen, quiet = capture.seen, 0
948        else:
949            quiet += 1
950        finished = getattr(obj, "exit_code", None)
951        if changed or (finished is not None and quiet <= 1):
952            send({"type": "mirror", "handle": key, **_mirror_state(obj)})
953        if finished is not None and quiet >= 3:
954            return
955
956
957def _invoke_reply(
958    call_id: str, result: object, state: dict[str, object] | None = None
959) -> dict[str, object]:
960    """One reference per result: inline the value when it can cross, retain it live.
961
962    Typed composites (dataclasses, list subclasses) get both, so the owner can
963    read them as values and still call methods on the live remote object. Plain
964    data crosses raw: its methods are pure and its state already crossed, so a
965    remote call would be the same computation plus a network fee, and retaining
966    it would evict live references worth keeping.
967    """
968    try:
969        wire = encode(result)
970        encoded = json.dumps(wire, ensure_ascii=True)
971    except (Unencodable, TypeError, ValueError, RecursionError):
972        reply: dict[str, object] = {
973            "type": "invoked",
974            "id": call_id,
975            "ok": True,
976            "handle": retain(result),
977        }
978        if state is not None:
979            reply["state"] = state
980        return reply
981    if len(encoded) > MAX_RESULT_BYTES:
982        return {
983            "type": "invoked",
984            "id": call_id,
985            "ok": False,
986            "error": {
987                "ename": "RemoteValueError",
988                "evalue": f"result is {len(encoded)} bytes, over the "
989                f"{MAX_RESULT_BYTES}-byte ceiling",
990                "traceback": [],
991            },
992        }
993    typed = (
994        isinstance(result, list)
995        and type(result) is not list
996        or dataclasses.is_dataclass(result)
997        and not isinstance(result, type)
998    )
999    if typed:
1000        return {
1001            "type": "invoked",
1002            "id": call_id,
1003            "ok": True,
1004            "handle": retain(result),
1005            "value": wire,
1006        }
1007    if state is not None:
1008        return {
1009            "type": "invoked",
1010            "id": call_id,
1011            "ok": True,
1012            "handle": retain(result),
1013            "state": state,
1014        }
1015    return {"type": "invoked", "id": call_id, "ok": True, "value": wire}
1016
1017
1018async def _owner_args(value: object) -> object:
1019    """Resolve owner argument markers back to what the call expects.
1020
1021    `{"__ref__": id}` names a live reference and `{"pending": id}` a call whose
1022    result this argument uses, resolved in order like a pending target, so an
1023    argument raced ahead of its own reply still lands. `{"__bytes__": b64}` is
1024    byte data, the same marker a reply value uses, decoded by albedo_values.
1025    """
1026    if isinstance(value, dict):
1027        if set(value) == {"__ref__"} and isinstance(value.get("__ref__"), str):
1028            obj = LIVE.get(value["__ref__"])
1029            if obj is None:
1030                raise LookupError(f"live reference {value['__ref__']!r} is gone")
1031            return obj
1032        if set(value) == {"pending"} and isinstance(value.get("pending"), str):
1033            pending = value["pending"]
1034            future = PENDING_OBJECTS.get(pending)
1035            if future is None:
1036                raise LookupError(f"pending call {pending!r} is gone")
1037            ok, resolved = await cast(Awaitable[tuple[bool, object]], future)
1038            if not ok:
1039                raise cast(BaseException, resolved)
1040            return resolved
1041        if set(value) == {"__bytes__"} and isinstance(value.get("__bytes__"), str):
1042            return decode(value)
1043        return {key: await _owner_args(item) for key, item in value.items()}
1044    if isinstance(value, list):
1045        return [await _owner_args(item) for item in value]
1046    return value
1047
1048
1049def _resolve(name: str) -> object:
1050    """The callable a namespace-path invoke addresses; handles resolve before this."""
1051    parts = name.split(".") if name else []
1052    if not all(part.isidentifier() and not part.startswith("_") for part in parts):
1053        raise ValueError(f"unaddressable owner invoke: {name!r}")
1054    obj = NAMESPACE.get(parts[0]) if parts else None
1055    if obj is None:
1056        raise LookupError(
1057            "no remote binding "
1058            + repr(parts[0] if parts else name)
1059            + "; rem.tools() lists what exists"
1060        )
1061    for part in parts[1:]:
1062        obj = getattr(obj, part)
1063    return obj
1064
1065
1066async def _target_object(target: object) -> object | None:
1067    """The object a target names: a live reference, a pending call's result, or nothing.
1068
1069    A `{"pending": id}` target waits for the call that produced the object, so
1070    method calls raced ahead of their own reply still resolve in order.
1071    """
1072    if not isinstance(target, dict):
1073        return None
1074    key = target.get("handle")
1075    if isinstance(key, str):
1076        obj = LIVE.get(key)
1077        if obj is None:
1078            raise LookupError(f"live reference {key!r} is gone")
1079        return obj
1080    pending = target.get("pending")
1081    if isinstance(pending, str):
1082        future = PENDING_OBJECTS.get(pending)
1083        if future is None:
1084            raise LookupError(f"pending call {pending!r} is gone")
1085        ok, value = await cast(Awaitable[tuple[bool, object]], future)
1086        if not ok:
1087            raise cast(BaseException, value)
1088        return value
1089    return None
1090
1091
1092async def serve_invoke(message: albedo_api.Invoke) -> None:
1093    """One owner tool call: resolve, call, await by the cell rule, answer one reference.
1094
1095    The reply inlines the value when it can cross and otherwise retains the
1096    object as a live reference; `value` and `handle` are one concept. Awaitables
1097    run to completion unless their class is a registered handle, the same rule
1098    the local evaluator applies, so a tool's remote shape matches its local one.
1099    An explicit await of a live reference awaits it only when the object is
1100    awaitable (a job); otherwise the object is the answer, as a reference, so
1101    awaiting a call whose result is a plain live object (a spawned Browser)
1102    settles into that reference instead of raising.
1103    A job's reply carries its state, final when awaited, so the owner's poll() is
1104    never racy and the owner can match the job's wake to its reference before
1105    any mirror frame arrives. Every call's result object is retained briefly as
1106    a pending target for method calls that raced ahead of the reply.
1107    """
1108    call_id = message["id"]
1109    OWNER_TASKS[call_id] = cast(asyncio.Task[object], asyncio.current_task())
1110    future: asyncio.Future[tuple[bool, object]] = LOOP.create_future()
1111    PENDING_OBJECTS[call_id] = future
1112    PENDING_OBJECTS.move_to_end(call_id)
1113    while len(PENDING_OBJECTS) > LIVE_LIMIT:
1114        PENDING_OBJECTS.popitem(last=False)
1115    reply: dict[str, object]
1116    result: object = None
1117    state: dict[str, object] | None = None
1118    failure: BaseException | None = None
1119    try:
1120        async with OWNER_CALLS:
1121            base = await _target_object(message.get("target"))
1122            if message.get("await") is True:
1123                if base is None:
1124                    raise ValueError("await needs a live or pending reference target")
1125                if inspect.isawaitable(base):
1126                    result = await cast(Awaitable[object], base)
1127                else:
1128                    # A settled object that cannot be awaited (a spawned
1129                    # Browser) is the result itself; the reply retains it as
1130                    # the same live reference, so nothing is orphaned.
1131                    result = base
1132            else:
1133                name = message.get("name", "")
1134                if base is None:
1135                    call = _resolve(name)
1136                else:
1137                    parts = name.split(".") if name else []
1138                    if not all(
1139                        part.isidentifier() and not part.startswith("_")
1140                        for part in parts
1141                    ):
1142                        raise ValueError(f"unaddressable owner invoke: {name!r}")
1143                    call = base
1144                    for part in parts:
1145                        call = getattr(call, part)
1146                args = [
1147                    await _owner_args(item)
1148                    for item in cast(Sequence[object], message.get("args", ()))
1149                ]
1150                kwargs = {
1151                    key: await _owner_args(item)
1152                    for key, item in message.get("kwargs", {}).items()
1153                }
1154                if not callable(call):
1155                    raise TypeError(f"owner invoke {name!r} is not callable")
1156                result = call(*args, **kwargs)
1157                if inspect.isawaitable(result) and not isinstance(
1158                    result, tuple(HANDLES)
1159                ):
1160                    result = await cast(Awaitable[object], result)
1161            if isinstance(getattr(result, "capture", None), Capture):
1162                state = _mirror_state(result)
1163    except asyncio.CancelledError as error:
1164        failure = error
1165        # The owner cancelled its wait; the remote effect may continue, and the
1166        # live reference (a running job, for instance) stays addressable.
1167        reply = {
1168            "type": "invoked",
1169            "id": call_id,
1170            "ok": False,
1171            "cancelled": True,
1172            "error": {
1173                "ename": "CancelledError",
1174                "evalue": "the owner cancelled this wait; the remote effect "
1175                "may continue and its reference stays addressable",
1176                "traceback": [],
1177            },
1178        }
1179    except BaseException as error:
1180        failure = error
1181        reply = {
1182            "type": "invoked",
1183            "id": call_id,
1184            "ok": False,
1185            "error": {
1186                "ename": type(error).__name__,
1187                "evalue": str(error)[:8192],
1188                "traceback": traceback.format_exc().splitlines()[-8:],
1189            },
1190        }
1191    else:
1192        reply = _invoke_reply(call_id, result, state)
1193    finally:
1194        OWNER_TASKS.pop(call_id, None)
1195        if not future.done():
1196            future.set_result(
1197                (failure is None, failure if failure is not None else result)
1198            )
1199    send(reply)
1200
1201
1202async def serve_job_stop(job_id: str) -> None:
1203    """Stop a background job on behalf of the port owner."""
1204    answer: dict[str, object] = {
1205        "type": "job_stop",
1206        "id": job_id,
1207        "ok": False,
1208        "error": "job not found",
1209    }
1210    for handler in reversed(STOP_HANDLERS):
1211        try:
1212            res = await handler(job_id)
1213            if res is not None:
1214                answer = {"type": "job_stop", "id": job_id, **res}
1215                break
1216        except Exception as error:
1217            answer = {
1218                "type": "job_stop",
1219                "id": job_id,
1220                "ok": False,
1221                "error": str(error).lower() or "could not stop job",
1222            }
1223            break
1224    send(answer)
1225
1226
1227def deliver_trace(capture: Capture) -> None:
1228    payload = capture.trace.finish()
1229    send({"type": "trace", "id": capture.id, "trace": payload})
1230    capture.trace.release()
1231
1232
1233def cell_read(id: str) -> None:
1234    task = CELL_TASKS.get(id)
1235    if id in BACKGROUND_RESULTS or (task is not None and task.done()):
1236        READ_CELLS.add(id)
1237    BACKGROUND_RESULTS.pop(id, None)
1238
1239
1240def first_line(code: str) -> str:
1241    """The cell's first line that is not blank or a comment, for one-line displays."""
1242    for line in code.splitlines():
1243        line = line.strip()
1244        if line and not line.startswith("#"):
1245            return line[:200]
1246    return "python cell"
1247
1248
1249async def announce_cells() -> None:
1250    while BACKGROUND_RESULTS:
1251        batch = dict(BACKGROUND_RESULTS)
1252        display = "\n".join(
1253            f"cell finished (status={result['status']}, ran {result['duration']}s): {result['summary']}"
1254            for result in batch.values()
1255        )
1256        listed = "; ".join(
1257            f"{id}: status={result['status']}, ran {result['duration']}s"
1258            for id, result in batch.items()
1259        )
1260        text = (
1261            "<system-note>background python cells finished with results unread: "
1262            + listed
1263            + ". output.read(id) reads retained output; await cells.info(id) reads "
1264            "the final status. no user sent this message; use the results if needed, "
1265            "otherwise acknowledge briefly and stay idle.</system-note>"
1266        )
1267        try:
1268            await host("cells.completed", {"display": display, "text": text})
1269        except Exception as error:
1270            if getattr(error, "code", "") == "busy":
1271                await asyncio.sleep(2)
1272                continue
1273            for id in batch:
1274                capture = ARCHIVES.get(id)
1275                if capture is not None:
1276                    capture.write(f"\n[completion notice not delivered: {error}]\n")
1277        for id in batch:
1278            BACKGROUND_RESULTS.pop(id, None)
1279
1280
1281async def execute_cell(
1282    message: albedo_api.Execute, capture: Capture
1283) -> dict[str, object]:
1284    token = CELL.set(capture)
1285    status, value = "ok", ""
1286    began = LOOP.time()
1287    _ = swap_sink(capture)
1288    try:
1289        if capture.interruption:
1290            raise asyncio.CancelledError()
1291        result = await evaluate(
1292            message["code"], capture.id, message.get("durable", False)
1293        )
1294        if result is not None:
1295            value = show(result)
1296    except (KeyboardInterrupt, asyncio.CancelledError):
1297        status = "interrupted"
1298        if capture.interruption == "deadline":
1299            capture.write("\n[cell deadline exceeded; execution interrupted]\n")
1300        elif capture.interruption == "memory":
1301            used = albedo_memory.footprint() or 0
1302            capture.write(
1303                f"\n[kernel memory passed its cap: {albedo_memory.mebibytes(used)} held,"
1304                f" cap {albedo_memory.mebibytes(albedo_memory.LIMIT)}; execution interrupted."
1305                " The namespace is kept: del large variables, gc.collect(), or work in smaller pieces.]\n"
1306            )
1307    except BaseException as error:
1308        status = "error"
1309        failed_id = cast(str, getattr(error, "_albedo_cell_id", capture.id))
1310        capture.write(traceback.format_exc())
1311        if message.get("durable", False):
1312            capture.write(
1313                f"\n[cell {failed_id} retained; inspect with await cells.read({failed_id!r}); repair with await cells.run({failed_id!r}, replacements=[(old, new)])]\n"
1314            )
1315    finally:
1316        _ = drain()
1317        _ = swap_sink(None)
1318        CELL.reset(token)
1319    deliver_trace(capture)
1320    return {
1321        "type": "done",
1322        "id": capture.id,
1323        "status": status,
1324        "duration": round(LOOP.time() - began, 3),
1325        "output": capture.preview(status) + capture.spill_note(),
1326        "value": value,
1327        "truncated": capture.seen > PREVIEW,
1328        "images": capture.encoded_images(),
1329    }
1330
1331
1332async def finish_background(
1333    message: albedo_api.Execute, task: asyncio.Task[dict[str, object]]
1334) -> None:
1335    global CELL_ANNOUNCER
1336    done = await task
1337    capture = ARCHIVES[message["id"]]
1338    if done["value"]:
1339        capture.write("\n[result]\n" + str(done["value"]) + "\n")
1340        done["output"] = capture.preview(str(done["status"])) + capture.spill_note()
1341        done["truncated"] = capture.seen > PREVIEW
1342    # Commit before the wake, so cells.info/read see the terminal outcome.
1343    if message.get("durable", False):
1344        try:
1345            await host("cells.finish", {"id": message["id"], "outcome": done})
1346        except Exception as error:
1347            ARCHIVES[message["id"]].write(
1348                f"\n[cell result could not be journaled: {error}]\n"
1349            )
1350    else:
1351        send(done)
1352    retain_cell_result(message["id"], done)
1353    if message["id"] in READ_CELLS:
1354        READ_CELLS.discard(message["id"])
1355        return
1356    BACKGROUND_RESULTS[message["id"]] = {**done, "summary": first_line(message["code"])}
1357    if CELL_ANNOUNCER is None or CELL_ANNOUNCER.done():
1358        CELL_ANNOUNCER = LOOP.create_task(announce_cells())
1359
1360
1361def retain_cell_result(id: str, done: dict[str, object]) -> None:
1362    FINISHED[id] = done
1363    while len(FINISHED) > FINISHED_LIMIT:
1364        old, _ = FINISHED.popitem(last=False)
1365        BACKGROUND_RESULTS.pop(old, None)
1366        READ_CELLS.discard(old)
1367    budget = FINISHED_IMAGE_BYTES
1368    for result in reversed(FINISHED.values()):
1369        images = cast(list[str], result.get("images", []))
1370        budget -= sum(map(len, images))
1371        if budget < 0 and images:
1372            result["images"] = []
1373            result["output"] = (
1374                f"{result['output']}\n[{len(images)} image(s) not kept for a replayed result]"
1375            )
1376    EXECUTING.discard(id)
1377    CELL_TASKS.pop(id, None)
1378    CELL_CAPTURES.pop(id, None)
1379    CELL_STARTS.pop(id, None)
1380    capture = ARCHIVES.get(id)
1381    if capture is not None:
1382        capture.end_spill()
1383        capture.images.clear()
1384        remember(capture)
1385    send({"type": "cells", "live": len(CELL_TASKS)})
1386
1387
1388def cell_timers(
1389    capture: Capture, task: asyncio.Task[dict[str, object]], timeout_ms: int
1390) -> tuple[
1391    threading.Event,
1392    threading.Event,
1393    threading.Lock,
1394    asyncio.Future[None],
1395    threading.Timer,
1396]:
1397    detached = threading.Event()
1398    settled = threading.Event()
1399    response_lock = threading.Lock()
1400    released: asyncio.Future[None] = LOOP.create_future()
1401
1402    def background() -> None:
1403        with response_lock:
1404            if settled.is_set() or task.done():
1405                return
1406            detached.set()
1407            send(background_reply(capture))
1408        LOOP.call_soon_threadsafe(released.set_result, None)
1409
1410    # Threads also fire while synchronous cell code occupies the asyncio loop.
1411    background_timer = threading.Timer(BACKGROUND_SECONDS, background)
1412    deadline_timer = threading.Timer(
1413        timeout_ms / 1000,
1414        receive,
1415        args=({"type": "interrupt", "id": capture.id, "reason": "deadline"},),
1416    )
1417    background_timer.daemon = deadline_timer.daemon = True
1418    background_timer.start()
1419    deadline_timer.start()
1420    task.add_done_callback(lambda _: deadline_timer.cancel())
1421    return detached, settled, response_lock, released, background_timer
1422
1423
1424async def serve():
1425    while True:
1426        message = await QUEUE.get()
1427        if message["type"] != "execute":
1428            send(state_reply(message))
1429            continue
1430        capture = Capture(message["id"], max_edge=message.get("max_edge"))
1431        capture.interruption = ""
1432        remember(capture)
1433        if len(CELL_TASKS) >= LIMITS["cell"]:
1434            done: dict[str, object] = {
1435                "type": "done",
1436                "id": capture.id,
1437                "status": "error",
1438                "duration": 0,
1439                "output": "16 cells are still running; cancel an unused cell with await cells.cancel(id), or end your turn and wait for completion.",
1440                "value": "",
1441                "truncated": False,
1442                "images": [],
1443            }
1444            retain_cell_result(capture.id, done)
1445            send(done)
1446            continue
1447        task = LOOP.create_task(execute_cell(message, capture))
1448        CELL_TASKS[capture.id] = task
1449        CELL_CAPTURES[capture.id] = capture
1450        CELL_STARTS[capture.id] = albedo_memory.footprint() or 0
1451        send({"type": "cells", "live": len(CELL_TASKS)})
1452        interrupted = INTERRUPTS.pop(capture.id, None)
1453        if interrupted is not None:
1454            capture.interruption = interrupted
1455        detached, settled, response_lock, released, background_timer = cell_timers(
1456            capture, task, message.get("timeout_ms", 300_000)
1457        )
1458        _ = await asyncio.wait({task, released}, return_when=asyncio.FIRST_COMPLETED)
1459        with response_lock:
1460            settled.set()
1461            background_timer.cancel()
1462            is_background = detached.is_set()
1463        if is_background:
1464            watch_output(capture.id, lambda id=capture.id: cell_read(id))
1465            _ = LOOP.create_task(finish_background(message, task))
1466        else:
1467            done = task.result()
1468            retain_cell_result(capture.id, done)
1469            READ_CELLS.discard(capture.id)
1470            send(done)
1471
1472
1473def background_reply(capture: Capture) -> dict[str, object]:
1474    return {
1475        "type": "done",
1476        "id": capture.id,
1477        "status": "backgrounded",
1478        "duration": BACKGROUND_SECONDS,
1479        "output": capture.preview("ok")
1480        + capture.spill_note()
1481        + (
1482            f"\n[cell {capture.id} is still running in the background. "
1483            "the session will wake automatically when it finishes; do NOT poll "
1484            "or sleep. meanwhile do other useful work, or give the user a short "
1485            "status report and end your turn. the original timeout_ms deadline "
1486            "still applies. avoid changing variables this cell is using.]\n"
1487        ),
1488        "value": "",
1489        "truncated": capture.seen > PREVIEW,
1490        "images": capture.encoded_images(),
1491    }
1492
1493
1494def flush(timeout: float) -> None:
1495    """Give the owner a moment to attach and take what we sent, before a
1496    detached kernel that failed to start leaves."""
1497    if isinstance(LINK, albedo_link.SocketLink):
1498        deadline = time.monotonic() + timeout
1499        while time.monotonic() < deadline and (
1500            not LINK.attached() or LINK.outbox.entries
1501        ):
1502            time.sleep(0.02)
1503
1504
1505def open_link(argv: list[str]) -> list[str]:
1506    """Stdio for `kernel.py <modules>`; a socket for `kernel.py --run <dir> <modules>`,
1507    whose bridge writes the attach token on our stdin. Returns the module list."""
1508    global LINK, RUN_DIR
1509    if argv[1:2] != ["--run"]:
1510        LINK = albedo_link.StdioLink()
1511        return cast(list[str], json.loads(argv[1]))
1512    RUN_DIR = argv[2]
1513    token = sys.stdin.buffer.readline().decode().strip()
1514    LINK = albedo_link.SocketLink(
1515        RUN_DIR,
1516        token,
1517        hello,
1518        bundle=albedo_bundle.digest(),
1519        grace=FIRST_ATTACH_GRACE,
1520    )
1521    return cast(list[str], json.loads(argv[3]))
1522
1523
1524PRELOADED = (
1525    "asyncio base64 collections datetime functools hashlib itertools json math os re shutil sys textwrap time"
1526).split()
1527
1528
1529def session_builtins(bindings: dict[str, object]) -> dict[str, object]:
1530    """Python's builtins plus the common standard-library modules and the
1531    injected bindings. Cells may shadow any of them and `del` reveals it again;
1532    a binding wins over a module of the same name."""
1533    modules = {name: importlib.import_module(name) for name in PRELOADED}
1534    return vars(builtins) | modules | {"Path": pathlib.Path} | bindings
1535
1536
1537class Bindings(ModuleType):
1538    """`import albedo` and `from albedo import files` answer the injected
1539    bindings, for code written as if they were a library."""
1540
1541    def __init__(self, bindings: dict[str, object]) -> None:
1542        super().__init__("albedo", "The tool bindings every cell already has.")
1543        self.bindings = bindings
1544
1545    def __getattr__(self, name: str) -> object:
1546        try:
1547            return self.bindings[name]
1548        except KeyError:
1549            raise AttributeError(f"albedo has no binding {name!r}") from None
1550
1551    def __dir__(self) -> list[str]:
1552        return sorted(self.bindings)
1553
1554
1555def main():
1556    global NATIVE_FD
1557    if os.getpgrp() != os.getpid():
1558        os.setsid()
1559    modules = open_link(sys.argv)
1560    # Beside the run directories: $ALBEDO_HOME/output, ~/.albedo-remote/output.
1561    albedo_capture.prune_spills(
1562        os.path.join(os.path.dirname(os.path.dirname(RUN_DIR)), "output")
1563        if RUN_DIR is not None
1564        else os.path.join(tempfile.gettempdir(), "albedo-output")
1565    )
1566    _ = os.dup2(os.open(os.devnull, os.O_RDONLY), 0)
1567    read_fd, write_fd = os.pipe()
1568    _ = os.dup2(write_fd, 1)
1569    _ = os.dup2(write_fd, 2)
1570    os.close(write_fd)
1571    os.set_blocking(read_fd, False)
1572    NATIVE_FD = read_fd
1573    remember(NATIVE)
1574    sys.stdin = open(os.devnull)
1575    sys.stdout = sys.stderr = Stream()
1576    threading.Thread(target=native_output, args=(read_fd,), daemon=True).start()
1577    threading.Thread(target=reader, daemon=True).start()
1578    if RUN_DIR is not None:
1579        threading.Thread(target=watch, daemon=True).start()
1580    threading.Thread(target=guard_memory, daemon=True).start()
1581    _ = signal.signal(signal.SIGINT, interrupt)
1582    albedo_shell.install()  # refusals precede observational audit hooks
1583    albedo_trace.install(CELL.get)
1584    api = albedo_api.PythonApi(
1585        version=2,
1586        loop=LOOP,
1587        host=host,
1588        host_now=host_now,
1589        HostError=WorkError,
1590        forget_output=forget_output,
1591        capture=background_capture,
1592        preview=PREVIEW,
1593        send=send,
1594        on_shutdown=CLEANUP.append,
1595        background_handle=HANDLES.append,
1596        modules=modules,
1597        watch_output=watch_output,
1598        attach_image=attach_image,
1599        on_stop_job=STOP_HANDLERS.append,
1600    )
1601    NAMESPACE.update(cells=Cells(), output=Output(), show_image=show_image)
1602    try:
1603        LOOP.run_until_complete(albedo_api.load_plugins(modules, api, NAMESPACE))
1604    except Exception as error:
1605        send({"type": "startup_error", "message": str(error)})
1606        flush(5.0)
1607        die()
1608        return
1609    # Match an interactive Python session only after loading trusted plugins.
1610    sys.path[0] = ""
1611    # Plugin bindings and session objects are rebuilt on every start, never saved.
1612    INJECTED.update(NAMESPACE)
1613    # Cell globals may shadow tools; deletion reveals the original session binding.
1614    bindings = {
1615        name: value for name, value in NAMESPACE.items() if not name.startswith("_")
1616    }
1617    NAMESPACE["__builtins__"] = session_builtins(bindings)
1618    sys.modules["albedo"] = Bindings(bindings)
1619    # Declare our own process group; a group we do not lead is never the supervisor's target.
1620    pgid = os.getpgid(0)
1621    send(
1622        {
1623            "type": "ready",
1624            "pid": os.getpid(),
1625            "pgid": pgid if pgid == os.getpid() else None,
1626            "leader": albedo_proc.leader_token(os.getpid()),
1627        }
1628    )
1629    READY.set()
1630    task = LOOP.create_task(serve())
1631    while not task.done():
1632        try:
1633            LOOP.run_until_complete(task)
1634        except KeyboardInterrupt:
1635            if interrupt_capture is not None:
1636                running = CELL_TASKS.get(interrupt_capture.id)
1637                if running is not None:
1638                    _ = running.cancel()
1639    die()
1640
1641
1642if __name__ == "__main__":
1643    main()