From a9882cf79fe72b749df678db75c379a4e00d490c Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Mon, 9 Nov 2026 19:19:21 +0800 Subject: [PATCH] feat: read tapes from any Bub tape store in plugin mode As `bub walkman`, enter Bub's framework lifespan and serve the tape store from `provide_tape_store`. Bub's FileTapeStore keeps the file path (byte offsets, file watching); any other sync or async store is read through the store API: followed tapes are polled, the list shows entry counts instead of sizes, and deleting resets the tape and its sidecars in the store. Co-Authored-By: Claude Opus 5.5 --- README.md | 1 + src/bub_walkman/cli.py | 25 ++- src/bub_walkman/server.py | 109 ++++------ src/bub_walkman/sources.py | 395 +++++++++++++++++++++++++++++++++++ src/bub_walkman/tapes.py | 32 ++- tests/test_store_source.py | 134 ++++++++++++ web/src/App.tsx | 2 +- web/src/components/Tapes.tsx | 23 +- web/src/hooks.ts | 14 +- 9 files changed, 634 insertions(+), 101 deletions(-) create mode 100644 src/bub_walkman/sources.py create mode 100644 tests/test_store_source.py diff --git a/README.md b/README.md index 19e48b6..27b0fd0 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ uvx bub-walkman ~/.bub/tapes/.jsonl When installed into Bub's environment (for example `bub install bub-walkman`), `bub-walkman` also works as a Bub plugin via `bub walkman `. - **Tape directory**: `bub.home / "tapes"` when Bub is installed, otherwise `$BUB_HOME/tapes` (default `~/.bub/tapes`). +- **Other tape stores**: as `bub walkman`, tapes come from the tape store Bub is configured with (the `provide_tape_store` hook), so tapes kept by a store plugin (a database, for example) work too. Such a store has no files: the list shows entry counts instead of sizes, a followed tape is re-read every second, and deleting a tape resets it and its sidecars in the store. Bub's default file store is still read from the tape directory. - **Server**: starts a local server at `http://localhost:7357` by default (or a free port if `7357` is taken) and opens the page in the browser. - **Options**: `--host`, `--port`, `--no-open`, `--version`. diff --git a/src/bub_walkman/cli.py b/src/bub_walkman/cli.py index 5afa486..0dced76 100644 --- a/src/bub_walkman/cli.py +++ b/src/bub_walkman/cli.py @@ -58,10 +58,11 @@ def _version(value: bool) -> None: def walkman( + ctx: typer.Context, tape: Annotated[ Path | None, typer.Argument( - help="Tape .jsonl file, or the name of a tape in Bub's tape directory. Without one, the page lists the tapes " + help="Tape .jsonl file, or the name of a tape in Bub's tape store. Without one, the page lists the tapes " "(running ones first) to pick from.", show_default=False, ), @@ -69,30 +70,38 @@ def walkman( host: Annotated[str, typer.Option(help="Bind address.")] = "127.0.0.1", port: Annotated[int, typer.Option(help="Port; falls back to a free one when taken.")] = 7357, open_browser: Annotated[bool, typer.Option("--open/--no-open", help="Open the page in a browser.")] = True, - version: Annotated[ - bool | None, typer.Option("--version", callback=_version, is_eager=True, help="Show the version and exit.") - ] = None, + version: Annotated[bool | None, typer.Option("--version", callback=_version, is_eager=True, help="Show the version and exit.")] = None, ) -> None: """Render a Bub tape as an agent timeline in the browser. Live following and replay are controlled from the web page.""" import uvicorn from .server import create_app + from .sources import BubSource, FileSource, TapeSource + # As `bub walkman`, Bub passes its framework: read tapes from whatever tape store Bub is configured with. + framework = ctx.obj if hasattr(ctx.obj, "get_tape_store") and hasattr(ctx.obj, "running") else None directory = tapes_dir() path = None + name = None if tape is not None: path = resolve_tape(tape) - if path is None: + if path is not None: + name = path.stem + elif framework is not None and len(tape.parts) == 1: + name = tape.name.removesuffix(".jsonl") # a tape in the store; the page says so if there is none + else: typer.echo(f"bub-walkman: tape not found: {tape} (also looked in {directory})", err=True) raise typer.Exit(2) + source: TapeSource = BubSource(framework, directory, path) if framework is not None else FileSource(directory, path) sock = _bind(host, port) bound = sock.getsockname()[1] shown = "localhost" if host in ("127.0.0.1", "0.0.0.0", "::") else host - url = f"http://{shown}:{bound}/" + (f"?tape={quote(path.stem)}" if path else "") - typer.echo(f"bub-walkman {path.name if path else directory}\n {url}\n Ctrl+C to stop") + url = f"http://{shown}:{bound}/" + (f"?tape={quote(name)}" if name else "") + where = path.name if path else name or ("Bub's tape store" if framework is not None else directory) + typer.echo(f"bub-walkman {where}\n {url}\n Ctrl+C to stop") - app = create_app(directory, path) + app = create_app(source) class Server(uvicorn.Server): @override diff --git a/src/bub_walkman/server.py b/src/bub_walkman/server.py index 0f5ab49..7e96de0 100644 --- a/src/bub_walkman/server.py +++ b/src/bub_walkman/server.py @@ -1,4 +1,4 @@ -"""Starlette app serving the tapes of a directory: static UI, tape list, JSON snapshots and SSE tail streams.""" +"""Starlette app serving tapes: static UI, tape list, JSON snapshots and SSE tail streams.""" import asyncio import json @@ -17,55 +17,32 @@ ) from starlette.routing import Mount, Route from starlette.staticfiles import StaticFiles -from watchfiles import awatch -from .tapes import TapeDirectory, TapeReader, read_all, read_spill, sidecar_path +from .sources import FileSource, TapeSource STATIC_DIR = Path(__file__).parent / "static" HEARTBEAT_INTERVAL = 15.0 -FALLBACK_POLL = 2.0 - - -def tape_info(path: Path) -> dict[str, Any]: - try: - st = path.stat() - size, mtime = st.st_size, st.st_mtime - except OSError: - size, mtime = 0, 0.0 - return {"name": path.stem, "path": str(path), "size": size, "mtime": mtime} def sse(event: str, data: Any) -> str: return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n" -async def _changes(path: Path, stop: asyncio.Event) -> AsyncIterator[None]: - """Yield whenever ``path`` may have changed: on filesystem events, or every few seconds.""" - watcher = awatch( - path.parent, - stop_event=stop, - debounce=50, - step=50, - rust_timeout=int(FALLBACK_POLL * 1000), - yield_on_timeout=True, - watch_filter=lambda _, p: Path(p).name == path.name, - ) - async for _ in watcher: - yield None - +def create_app(source: TapeSource | Path, tape: Path | None = None) -> Starlette: + """Serve the tapes of ``source``: a tape store, or a tape directory (plus ``tape``, which may live elsewhere). -def create_app(tapes_dir: Path, tape: Path | None = None) -> Starlette: - """Serve the tapes in ``tapes_dir`` (plus ``tape``, which may live elsewhere). API calls pick one with ``?tape=``.""" - directory = TapeDirectory(tapes_dir, tape) - home = str(Path.home()) - shown_dir = "~" + str(tapes_dir)[len(home) :] if str(tapes_dir).startswith(home + "/") else str(tapes_dir) + API calls pick one tape with ``?tape=``. + """ + if isinstance(source, Path): + source = FileSource(source, tape) + tapes_source = source - def pick(request: Request) -> Path | Response: + async def pick(request: Request) -> tuple[str, dict[str, Any]] | Response: name = request.query_params.get("tape") if not name: return JSONResponse({"error": "missing ?tape="}, 400) - path = directory.resolve(name) - return path if path is not None else JSONResponse({"error": f"tape not found: {name}"}, 404) + info = await tapes_source.info(name) + return (name, info) if info is not None else JSONResponse({"error": f"tape not found: {name}"}, 404) # Stop events of open SSE streams; set them all on shutdown so streams end and the server exits promptly. streams: set[asyncio.Event] = set() @@ -83,8 +60,8 @@ async def favicon(request: Request) -> Response: return FileResponse(icon, media_type="image/svg+xml", headers={"Cache-Control": "max-age=86400"}) async def tapes(request: Request) -> Response: - items = await asyncio.to_thread(directory.summaries) - return JSONResponse({"directory": shown_dir, "tapes": items}, headers={"Cache-Control": "no-store"}) + items = await tapes_source.summaries() + return JSONResponse({"directory": tapes_source.label, "tapes": items}, headers={"Cache-Control": "no-store"}) async def delete(request: Request) -> Response: # Deleting files is only for the person at this machine: a loopback client, our own page (not a page from @@ -104,64 +81,50 @@ async def delete(request: Request) -> Response: names = [str(n) for n in body["names"]] except (ValueError, KeyError, TypeError): return JSONResponse({"error": 'expected {"names": [...], "force": false}'}, 400) - deleted, errors = await asyncio.to_thread(directory.delete, names, bool(body.get("force"))) + deleted, errors = await tapes_source.delete(names, bool(body.get("force"))) return JSONResponse({"deleted": deleted, "errors": errors}, 200 if not errors else 409) async def snapshot(request: Request) -> Response: - tape = pick(request) - if isinstance(tape, Response): - return tape - entries, offset = await asyncio.to_thread(read_all, tape) - return JSONResponse({"tape": tape_info(tape), "entries": entries, "offset": offset}) + picked = await pick(request) + if isinstance(picked, Response): + return picked + name, info = picked + entries, offset = await tapes_source.snapshot(name) + return JSONResponse({"tape": info, "entries": entries, "offset": offset}) async def spill(request: Request) -> Response: - tape = pick(request) - if isinstance(tape, Response): - return tape + picked = await pick(request) + if isinstance(picked, Response): + return picked handle = request.path_params["handle"] if not handle.isalnum(): return JSONResponse({"error": "bad handle"}, 400) - path = sidecar_path(tape, "spill") - if not path.is_file(): - return JSONResponse({"error": "spill sidecar not found"}, 404) - data = await asyncio.to_thread(read_spill, path, handle) + data = await tapes_source.spill(picked[0], handle) + if isinstance(data, str): + return JSONResponse({"error": data}, 404) if not data["chunks"]: return JSONResponse({"error": "spill handle not found (sidecar may have been reset)"}, 404) return JSONResponse(data) async def stream(request: Request) -> Response: - picked = pick(request) + picked = await pick(request) if isinstance(picked, Response): return picked - tape = picked + name = picked[0] offset = int(request.query_params.get("offset") or 0) async def events() -> AsyncIterator[str]: - reader = TapeReader(tape) - try: - st = tape.stat() - except OSError: - st = None - if st is not None and 0 < offset <= st.st_size: - reader.resume(offset, st) # continue right after the client's snapshot - elif offset: - yield sse("reset", {"offset": 0}) - yield sse("hello", {"tape": tape_info(tape), "offset": reader.offset}) - stop = asyncio.Event() streams.add(stop) loop = asyncio.get_running_loop() last_beat = loop.time() try: - async for _ in _changes(tape, stop): - if await request.is_disconnected(): - break - reset, entries = await asyncio.to_thread(reader.poll) - if reset: - yield sse("reset", {"offset": 0}) - if entries: - yield sse("entries", {"entries": entries, "offset": reader.offset}) + async for event in tapes_source.follow(name, offset, stop): + if event is not None: + yield sse(*event) last_beat = loop.time() + elif await request.is_disconnected(): + break elif loop.time() - last_beat > HEARTBEAT_INTERVAL: yield ": beat\n\n" last_beat = loop.time() @@ -176,6 +139,7 @@ async def events() -> AsyncIterator[str]: ) app = Starlette( + lifespan=lambda app: tapes_source.lifespan(), routes=[ Route("/", index), Route("/favicon.svg", favicon), @@ -185,8 +149,9 @@ async def events() -> AsyncIterator[str]: Route("/api/stream", stream), Route("/api/spill/{handle}", spill), Mount("/assets", StaticFiles(directory=STATIC_DIR / "assets", check_dir=False), name="assets"), - ] + ], ) + def close_streams() -> None: for stop in list(streams): stop.set() diff --git a/src/bub_walkman/sources.py b/src/bub_walkman/sources.py new file mode 100644 index 0000000..61d9968 --- /dev/null +++ b/src/bub_walkman/sources.py @@ -0,0 +1,395 @@ +"""Where tapes come from: JSONL files in a directory, or any Bub tape store (when running as a Bub plugin). + +Every source speaks the same small async protocol, so the server does not care which one it serves. Offsets are +opaque to the page: a byte offset for files, an entry count for stores. +""" + +import asyncio +import contextlib +import dataclasses +import inspect +import time +from collections.abc import AsyncIterator, Callable, Iterable +from datetime import datetime +from pathlib import Path +from typing import Any, Protocol + +from watchfiles import awatch + +from .tapes import ( + SIDECAR_MARK, + Entry, + TapeDirectory, + TapeReader, + TapeSummary, + collect_spill, + read_all, + read_spill, + scan_entries, + sidecar_path, + sort_summaries, + tape_status, +) + +FALLBACK_POLL = 2.0 +STORE_POLL = 1.0 +STORE_LIST_TTL = 3.0 +STORE_CONCURRENCY = 8 + +type Event = tuple[str, dict[str, Any]] | None +"""An SSE event for the page (``hello``, ``reset``, ``entries``), or None for "nothing new yet".""" + + +class TapeSource(Protocol): + @property + def label(self) -> str: + """Where the tapes live, shown on the tape list (a directory, or the store's name).""" + ... + + def lifespan(self) -> contextlib.AbstractAsyncContextManager[None]: ... + + async def summaries(self) -> list[TapeSummary]: ... + + async def info(self, name: str) -> dict[str, Any] | None: + """``{name, path, size, mtime}`` of a tape, or None when there is no such tape.""" + ... + + async def snapshot(self, name: str) -> tuple[list[Entry], int]: + """All entries of a tape, and the offset to follow it from.""" + ... + + def follow(self, name: str, offset: int, stop: asyncio.Event) -> AsyncIterator[Event]: + """Events for a tape from ``offset`` on, until ``stop`` is set. Yields None now and then while idle.""" + ... + + async def spill(self, name: str, handle: str) -> dict[str, Any] | str: + """The spilled tool output ``handle`` of a tape, or an error message.""" + ... + + async def delete(self, names: list[str], force: bool) -> tuple[list[str], dict[str, str]]: ... + + +# -- files ---------------------------------------------------------------------------- + + +async def _changes(path: Path, stop: asyncio.Event) -> AsyncIterator[None]: + """Yield whenever ``path`` may have changed: on filesystem events, or every few seconds.""" + watcher = awatch( + path.parent, + stop_event=stop, + debounce=50, + step=50, + rust_timeout=int(FALLBACK_POLL * 1000), + yield_on_timeout=True, + watch_filter=lambda _, p: Path(p).name == path.name, + ) + async for _ in watcher: + yield None + + +class FileSource: + """Tapes stored as ``.jsonl`` files in Bub's tape directory (plus one tape file given from elsewhere).""" + + def __init__(self, directory: Path, extra: Path | None = None) -> None: + self.directory = TapeDirectory(directory, extra) + home = str(Path.home()) + text = str(directory) + self.label = "~" + text[len(home) :] if text.startswith(home + "/") else text + + @contextlib.asynccontextmanager + async def lifespan(self) -> AsyncIterator[None]: + yield + + async def summaries(self) -> list[TapeSummary]: + return await asyncio.to_thread(self.directory.summaries) + + @staticmethod + def _info(path: Path) -> dict[str, Any]: + try: + st = path.stat() + size, mtime = st.st_size, st.st_mtime + except OSError: + size, mtime = 0, 0.0 + return {"name": path.stem, "path": str(path), "size": size, "mtime": mtime} + + async def info(self, name: str) -> dict[str, Any] | None: + path = self.directory.resolve(name) + return None if path is None else self._info(path) + + async def snapshot(self, name: str) -> tuple[list[Entry], int]: + path = self.directory.resolve(name) + return await asyncio.to_thread(read_all, path) if path is not None else ([], 0) + + async def follow(self, name: str, offset: int, stop: asyncio.Event) -> AsyncIterator[Event]: + path = self.directory.resolve(name) + if path is None: + return + reader = TapeReader(path) + try: + st = path.stat() + except OSError: + st = None + if st is not None and 0 < offset <= st.st_size: + reader.resume(offset, st) # continue right after the page's snapshot + elif offset: + yield "reset", {"offset": 0} + yield "hello", {"tape": self._info(path), "offset": reader.offset} + async for _ in _changes(path, stop): + reset, entries = await asyncio.to_thread(reader.poll) + if reset: + yield "reset", {"offset": 0} + yield ("entries", {"entries": entries, "offset": reader.offset, "size": reader.offset}) if entries else None + + async def spill(self, name: str, handle: str) -> dict[str, Any] | str: + path = self.directory.resolve(name) + sidecar = sidecar_path(path, "spill") if path is not None else None + if sidecar is None or not sidecar.is_file(): + return "spill sidecar not found" + return await asyncio.to_thread(read_spill, sidecar, handle) + + async def delete(self, names: list[str], force: bool) -> tuple[list[str], dict[str, str]]: + return await asyncio.to_thread(self.directory.delete, names, force) + + +# -- Bub tape stores ------------------------------------------------------------------ + + +def _entry(entry: Any) -> Entry: + if isinstance(entry, dict): + return entry + if dataclasses.is_dataclass(entry) and not isinstance(entry, type): + return dataclasses.asdict(entry) + return {k: getattr(entry, k, None) for k in ("id", "kind", "payload", "meta", "date")} + + +def _timestamp(date: str) -> float: + try: + return datetime.fromisoformat(date).timestamp() + except (TypeError, ValueError): + return 0.0 + + +def _bub_query(tape: str, store: Any) -> Any: + from bub.core.store import TapeQuery + + return TapeQuery(tape=tape, store=store) + + +class StoreSource: + """Tapes in a Bub ``TapeStore`` or ``AsyncTapeStore``, whatever it persists to. + + A store has no change notifications or byte offsets, so a followed tape is re-read every second and its offset is + the number of entries already sent. Deleting a tape resets it (and its sidecars) in the store. + """ + + def __init__( + self, + store: Any, + *, + sidecars: Iterable[str] = (), + label: str | None = None, + query: Callable[[str, Any], Any] = _bub_query, + poll: float = STORE_POLL, + ) -> None: + self.store = store + self.label = label or f"{type(store).__name__} (Bub tape store)" + self.sidecar_names = sorted({"spill", *sidecars}) + self._query = query + self._poll = poll + self._cache: dict[str, tuple[tuple[int, Any], TapeSummary]] = {} + self._listed: tuple[float, list[TapeSummary]] | None = None + self._limit = asyncio.Semaphore(STORE_CONCURRENCY) + + @contextlib.asynccontextmanager + async def lifespan(self) -> AsyncIterator[None]: + yield + + async def _call(self, method: str, *args: Any) -> Any: + fn = getattr(self.store, method) + if inspect.iscoroutinefunction(fn): + return await fn(*args) + return await asyncio.to_thread(fn, *args) + + async def _read(self, name: str) -> list[Entry]: + async with self._limit: + entries = await self._call("fetch_all", self._query(name, self.store)) + return [_entry(e) for e in entries or ()] + + async def _names(self) -> list[str]: + return [str(n) for n in await self._call("list_tapes")] + + async def _summary(self, name: str) -> TapeSummary | None: + entries = await self._read(name) + if not entries: + return None # never written, or reset: nothing to show + key = (len(entries), entries[-1].get("id")) + cached = self._cache.get(name) + if cached is None or cached[0] != key: + sidecars = [] + for sidecar in self.sidecar_names: + tape = f"{name}{SIDECAR_MARK}{sidecar}" + if count := len(await self._read(tape)): + sidecars.append({"name": tape, "size": None, "entries": count}) + scan = scan_entries(entries) + cached = (key, {**scan, "entries": len(entries), "sidecars": sidecars}) + self._cache[name] = cached + scan = cached[1] + mtime = _timestamp(scan["last"]) + return { + "name": name, + "size": None, + "entries": scan["entries"], + "mtime": mtime, + "status": tape_status(scan["open"], time.time() - mtime), + "prompt": scan["prompt"], + "last": scan["last"], + "sidecars": scan["sidecars"], + "deletable": True, + } + + async def summaries(self) -> list[TapeSummary]: + # Listing reads every tape; the page asks every few seconds, so share one listing between close requests. + if self._listed is not None and time.monotonic() - self._listed[0] < STORE_LIST_TTL: + return self._listed[1] + mains = [n for n in await self._names() if SIDECAR_MARK not in n] + found = await asyncio.gather(*(self._summary(n) for n in mains)) + for gone in self._cache.keys() - set(mains): + del self._cache[gone] + result = sort_summaries([s for s in found if s is not None]) + self._listed = (time.monotonic(), result) + return result + + def _valid(self, name: str) -> bool: + return bool(name) and SIDECAR_MARK not in name + + async def info(self, name: str) -> dict[str, Any] | None: + if not self._valid(name): + return None + entries = await self._read(name) + if not entries and name not in await self._names(): + return None + return self._info(name, entries) + + def _info(self, name: str, entries: list[Entry]) -> dict[str, Any]: + last = str(entries[-1].get("date") or "") if entries else "" + return {"name": name, "path": f"{self.label}: {name}", "size": None, "entries": len(entries), "mtime": _timestamp(last)} + + async def snapshot(self, name: str) -> tuple[list[Entry], int]: + entries = await self._read(name) + return entries, len(entries) + + async def follow(self, name: str, offset: int, stop: asyncio.Event) -> AsyncIterator[Event]: + entries = await self._read(name) + sent = offset + if not 0 <= offset <= len(entries): + yield "reset", {"offset": 0} + sent = 0 + # The id of the last entry sent tells a reset apart from a tape that only grew. + last_id = entries[sent - 1].get("id") if sent else None + yield "hello", {"tape": self._info(name, entries[:sent]), "offset": sent} + while not stop.is_set(): + if len(entries) < sent or (sent and entries[sent - 1].get("id") != last_id): + yield "reset", {"offset": 0} + sent, last_id = 0, None + if len(entries) > sent: + new = entries[sent:] + sent, last_id = len(entries), new[-1].get("id") + yield "entries", {"entries": new, "offset": sent, "size": None} + else: + yield None + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(stop.wait(), self._poll) + if stop.is_set(): + break + entries = await self._read(name) + + async def spill(self, name: str, handle: str) -> dict[str, Any] | str: + entries = await self._read(f"{name}{SIDECAR_MARK}spill") + if not entries: + return "spill sidecar not found" + return collect_spill(entries, handle) + + async def delete(self, names: list[str], force: bool) -> tuple[list[str], dict[str, str]]: + self._listed = None + status = {t["name"]: t for t in await self.summaries()} + deleted: list[str] = [] + errors: dict[str, str] = {} + for name in names: + summary = status.get(name) + if summary is None: + errors[name] = "tape not found" + continue + if summary["status"] == "running" and not force: + errors[name] = "tape is running" + continue + for tape in [*(s["name"] for s in summary["sidecars"]), name]: # the tape itself last + try: + await self._call("reset", tape) + except Exception as exc: # any store error: report it for this tape and go on + errors[name] = f"{tape}: {exc}" + break + deleted.append(tape) + self._cache.pop(name, None) + self._listed = None + return deleted, errors + + +# -- Bub plugin ----------------------------------------------------------------------- + + +class BubSource: + """The tape store of a running Bub framework, picked when the server starts. + + Bub's file store keeps using the directory directly (byte offsets, file watching); any other store is read through + the store API. + """ + + def __init__(self, framework: Any, fallback_dir: Path, extra: Path | None = None) -> None: + self.framework = framework + self.fallback_dir = fallback_dir + self.extra = extra + self.inner: TapeSource = FileSource(fallback_dir, extra) + + @property + def label(self) -> str: + return self.inner.label + + @contextlib.asynccontextmanager + async def lifespan(self) -> AsyncIterator[None]: + async with self.framework.running(): + self.inner = pick_source(self.framework, self.fallback_dir, self.extra) + async with self.inner.lifespan(): + yield + + async def summaries(self) -> list[TapeSummary]: + return await self.inner.summaries() + + async def info(self, name: str) -> dict[str, Any] | None: + return await self.inner.info(name) + + async def snapshot(self, name: str) -> tuple[list[Entry], int]: + return await self.inner.snapshot(name) + + def follow(self, name: str, offset: int, stop: asyncio.Event) -> AsyncIterator[Event]: + return self.inner.follow(name, offset, stop) + + async def spill(self, name: str, handle: str) -> dict[str, Any] | str: + return await self.inner.spill(name, handle) + + async def delete(self, names: list[str], force: bool) -> tuple[list[str], dict[str, str]]: + return await self.inner.delete(names, force) + + +def pick_source(framework: Any, fallback_dir: Path, extra: Path | None = None) -> TapeSource: + """A source for the framework's tape store: files for Bub's ``FileTapeStore``, the store API for anything else.""" + from bub.core.store import FileTapeStore + + store = framework.get_tape_store() + if store is None: + return FileSource(fallback_dir, extra) + if isinstance(store, FileTapeStore): + directory = getattr(store, "_directory", None) + return FileSource(Path(directory) if directory is not None else fallback_dir, extra) + sidecars = [s.name for s in framework.get_tape_sidecars()] + inner = getattr(store, "_store", store) # name the store an AsyncTapeStoreAdapter wraps, not the adapter + return StoreSource(store, sidecars=sidecars, label=f"{type(inner).__name__} (Bub tape store)") diff --git a/src/bub_walkman/tapes.py b/src/bub_walkman/tapes.py index 97207b2..41d91b1 100644 --- a/src/bub_walkman/tapes.py +++ b/src/bub_walkman/tapes.py @@ -81,11 +81,15 @@ def read_all(path: Path) -> tuple[list[Entry], int]: def read_spill(path: Path, handle: str) -> dict[str, Any]: - """Collect the chunks of one spilled tool output from a spill sidecar tape.""" + """Collect the chunks of one spilled tool output from a spill sidecar tape file.""" + return collect_spill(read_all(path)[0], handle) + + +def collect_spill(entries: list[Entry], handle: str) -> dict[str, Any]: + """Collect the chunks of one spilled tool output from the entries of a spill sidecar tape.""" chunks: dict[int, str] = {} current: int | None = None prefix = f"spill/{handle}/chunk/" - entries, _ = read_all(path) for entry in entries: kind = entry.get("kind") payload = entry.get("payload") or {} @@ -150,11 +154,15 @@ def _short_prompt(prompt: str) -> str: def _scan(path: Path, size: int) -> TapeSummary: + return scan_entries(_tail_entries(path, size)) + + +def scan_entries(entries: list[Entry]) -> TapeSummary: """What the tail of a tape says about its latest turn: whether it is still open, and its prompt.""" open_turn: bool | None = None prompt = "" last = "" - for entry in _tail_entries(path, size): + for entry in entries: last = str(entry.get("date") or last) if entry.get("kind") != "event": continue @@ -173,6 +181,17 @@ def _scan(path: Path, size: int) -> TapeSummary: return {"open": bool(open_turn), "prompt": prompt, "last": last} +def tape_status(open_turn: bool, age: float) -> str: + """running: an open turn written to recently; recent: written to within the hour; else idle.""" + return "running" if open_turn and age < RUNNING_WINDOW else "recent" if age < RECENT_WINDOW else "idle" + + +def sort_summaries(summaries: list[TapeSummary]) -> list[TapeSummary]: + rank = {"running": 0, "recent": 1, "idle": 2} + summaries.sort(key=lambda t: (rank[t["status"]], -t["mtime"])) + return summaries + + class TapeDirectory: """Lists the tapes in Bub's tape directory, with a cached tail scan per file version.""" @@ -221,8 +240,7 @@ def summaries(self, now: float | None = None) -> list[TapeSummary]: continue self._cache[path] = cached scan = cached[1] - age = now - st.st_mtime - status = "running" if scan["open"] and age < RUNNING_WINDOW else "recent" if age < RECENT_WINDOW else "idle" + status = tape_status(scan["open"], now - st.st_mtime) result.append( { "name": path.stem, @@ -238,9 +256,7 @@ def summaries(self, now: float | None = None) -> list[TapeSummary]: ) for gone in self._cache.keys() - seen: del self._cache[gone] - rank = {"running": 0, "recent": 1, "idle": 2} - result.sort(key=lambda t: (rank[t["status"]], -t["mtime"])) - return result + return sort_summaries(result) def delete(self, names: list[str], force: bool = False) -> tuple[list[str], dict[str, str]]: """Delete tapes with all their sidecars. Running tapes are skipped unless ``force``. diff --git a/tests/test_store_source.py b/tests/test_store_source.py new file mode 100644 index 0000000..5c220d1 --- /dev/null +++ b/tests/test_store_source.py @@ -0,0 +1,134 @@ +import asyncio +from datetime import UTC, datetime +from types import SimpleNamespace +from typing import Any + +import pytest +from starlette.testclient import TestClient + +from bub_walkman.server import create_app +from bub_walkman.sources import StoreSource + + +def entry(id: int, kind: str = "event", meta: dict | None = None, **payload: Any) -> dict: + return {"id": id, "kind": kind, "payload": payload, "meta": meta or {}, "date": datetime.now(UTC).isoformat()} + + +class MemoryStore: + """A minimal Bub-style TapeStore that is not file based.""" + + def __init__(self) -> None: + self.tapes: dict[str, list[dict]] = {} + + def list_tapes(self) -> list[str]: + return sorted(self.tapes) + + def reset(self, tape: str) -> None: + self.tapes.pop(tape, None) + + def fetch_all(self, query: Any) -> list[dict]: + return list(self.tapes.get(query.tape, [])) + + def append(self, tape: str, item: dict) -> None: + self.tapes.setdefault(tape, []).append(item) + + +class AsyncMemoryStore(MemoryStore): + async def list_tapes(self) -> list[str]: # type: ignore[override] + return super().list_tapes() + + async def reset(self, tape: str) -> None: # type: ignore[override] + super().reset(tape) + + async def fetch_all(self, query: Any) -> list[dict]: # type: ignore[override] + return super().fetch_all(query) + + +def query(tape: str, store: Any) -> Any: + return SimpleNamespace(tape=tape, store=store) + + +def fill(store: MemoryStore) -> None: + store.tapes["a__1"] = [entry(1, name="loop.start", data={"prompt": "fix it"}), entry(2, name="loop.step.start")] + store.tapes["a__1__sidecar__spill"] = [ + entry(1, "anchor", name="spill/h1/chunk/0"), + entry(2, "tool_result", meta={"spill_handle": "h1", "spill_chunk": 0}, results=["hello"]), + ] + store.tapes["b__2"] = [entry(1, name="loop.start"), entry(2, name="loop.finish")] + store.tapes["empty__3"] = [] + + +@pytest.mark.parametrize("cls", [MemoryStore, AsyncMemoryStore]) +def test_store_endpoints(cls: type[MemoryStore]) -> None: + store = cls() + fill(store) + client = TestClient(create_app(StoreSource(store, query=query)), base_url="http://localhost", client=("127.0.0.1", 1)) + + listing = client.get("/api/tapes").json() + assert listing["directory"] == f"{cls.__name__} (Bub tape store)" + tapes = {t["name"]: t for t in listing["tapes"]} + assert list(tapes) == ["a__1", "b__2"] # running first; sidecars and empty tapes are not listed + assert tapes["a__1"]["status"] == "running" and tapes["a__1"]["prompt"] == "fix it" + assert tapes["a__1"]["size"] is None and tapes["a__1"]["entries"] == 2 + assert tapes["a__1"]["sidecars"] == [{"name": "a__1__sidecar__spill", "size": None, "entries": 2}] + + snap = client.get("/api/tape?tape=a__1").json() + assert [e["id"] for e in snap["entries"]] == [1, 2] and snap["offset"] == 2 + assert client.get("/api/tape?tape=nope").status_code == 404 + assert client.get("/api/tape?tape=a__1__sidecar__spill").status_code == 404 + assert client.get("/api/spill/h1?tape=a__1").json()["content"] == "hello" + assert client.get("/api/spill/h1?tape=b__2").status_code == 404 + + delete = lambda names, force=False: client.post( # noqa: E731 + "/api/tapes/delete", json={"names": names, "force": force}, headers={"X-Walkman": "1"} + ) + res = delete(["a__1"]) + assert res.status_code == 409 and res.json()["errors"] == {"a__1": "tape is running"} + assert delete(["a__1", "b__2"], force=True).json() == { + "deleted": ["a__1__sidecar__spill", "a__1", "b__2"], + "errors": {}, + } + assert set(store.tapes) == {"empty__3"} + assert client.get("/api/tapes").json()["tapes"] == [] + + +def test_store_follow_appends_and_resets() -> None: + store = MemoryStore() + fill(store) + source = StoreSource(store, query=query, poll=0.01) + + async def run() -> list[tuple[str, Any]]: + stop = asyncio.Event() + seen: list[tuple[str, Any]] = [] + steps = iter( + [ + lambda: store.append("a__1", entry(3, "message", role="assistant", content="done")), + lambda: store.tapes.__setitem__("a__1", [entry(1, name="loop.start")]), # reset, then rewritten + stop.set, + ] + ) + async for event in source.follow("a__1", 2, stop): + if event is None: + next(steps)() + else: + seen.append((event[0], [e["id"] for e in event[1].get("entries", [])] or event[1].get("offset"))) + return seen + + assert asyncio.run(run()) == [("hello", 2), ("entries", [3]), ("reset", 0), ("entries", [1])] + + +def test_store_follow_from_stale_offset_resets() -> None: + store = MemoryStore() + fill(store) + + async def first_events() -> list[str]: + stop = asyncio.Event() + names = [] + async for event in StoreSource(store, query=query, poll=0.01).follow("a__1", 9, stop): + if event is None: + stop.set() + else: + names.append(event[0]) + return names + + assert asyncio.run(first_events()) == ["reset", "hello", "entries"] diff --git a/web/src/App.tsx b/web/src/App.tsx index a7ae55a..16fa499 100644 --- a/web/src/App.tsx +++ b/web/src/App.tsx @@ -421,7 +421,7 @@ export default function App({ name: tapeName, onNavigate }: { name: string; onNa {tokens(tl.stats.output)} out tokens {tl.stats.first != null && tl.stats.last != null && {duration(tl.stats.last - tl.stats.first)} span} - {tape.info && {bytes(tape.info.size)}} + {tape.info?.size != null && {bytes(tape.info.size)}} @@ -128,8 +135,9 @@ function DeleteDialog({ tapes, onClose, onDeleted }: { tapes: TapeSummary[]; onC const [busy, setBusy] = useState(false); const [errors, setErrors] = useState | null>(null); const running = tapes.filter((t) => t.status === "running"); - const files = tapes.flatMap((t) => [{ name: `${t.name}.jsonl`, size: t.size }, ...t.sidecars]); - const total = files.reduce((n, f) => n + f.size, 0); + // Files on disk, or tapes in a Bub tape store (no size; they are reset there). + const inStore = tapes.some((t) => t.size == null); + const files = tapes.flatMap((t) => [{ name: inStore ? t.name : `${t.name}.jsonl`, size: t.size, entries: t.entries }, ...t.sidecars]); const one = tapes.length === 1; const confirm = async () => { setBusy(true); @@ -158,8 +166,9 @@ function DeleteDialog({ tapes, onClose, onDeleted }: { tapes: TapeSummary[]; onC {one ? "Delete this tape?" : `Delete ${tapes.length} tapes?`} - {one ? "The tape and its sidecars" : "These tapes and their sidecars"} are removed from disk ({files.length} file{files.length === 1 ? "" : "s"},{" "} - {bytes(total)}). This cannot be undone. + {one ? "The tape and its sidecars" : "These tapes and their sidecars"}{" "} + {inStore ? `are reset in the tape store (${files.length} tape${files.length === 1 ? "" : "s"}, ${amount(files)})` : `are removed from disk (${files.length} file${files.length === 1 ? "" : "s"}, ${amount(files)})`}. + This cannot be undone. {running.length > 0 && (

@@ -171,7 +180,7 @@ function DeleteDialog({ tapes, onClose, onDeleted }: { tapes: TapeSummary[]; onC {files.map((f) => (

  • {f.name} - {bytes(f.size)} + {amount([f])}
    • ))} diff --git a/web/src/hooks.ts b/web/src/hooks.ts index 0528b0a..a5f8a0a 100644 --- a/web/src/hooks.ts +++ b/web/src/hooks.ts @@ -4,7 +4,9 @@ import type { Entry } from "./model"; export interface TapeInfo { name: string; path: string; - size: number; + /** Bytes on disk; null for a tape in a (non-file) Bub tape store. */ + size: number | null; + entries?: number; mtime: number; } @@ -32,7 +34,7 @@ export function apiUrl(path: string, params: Record = { return `${path}?${q}`; } -/** Loads the tape snapshot, then tails it over SSE (reconnecting from the last byte offset). */ +/** Loads the tape snapshot, then tails it over SSE (reconnecting from the last offset: bytes for files, entries for stores). */ export function useTape(name: string): TapeData { const entries = useRef([]); const offset = useRef(0); @@ -82,7 +84,7 @@ export function useTape(name: string): TapeData { if (append(data.entries)) { setState((s) => ({ ...s, - info: s.info && { ...s.info, size: data.offset }, + info: s.info && { ...s.info, size: data.size ?? null }, version: s.version + 1, lastUpdate: Date.now(), })); @@ -171,13 +173,15 @@ export type TapeStatus = "running" | "recent" | "idle"; export interface TapeSummary { name: string; - size: number; + /** Bytes on disk; null for a tape in a (non-file) Bub tape store, which counts `entries` instead. */ + size: number | null; + entries?: number; mtime: number; status: TapeStatus; /** Prompt of the latest turn (truncated). */ prompt: string; last: string; - sidecars: Array<{ name: string; size: number }>; + sidecars: Array<{ name: string; size: number | null; entries?: number }>; /** False for a tape given on the command line from outside the tape directory. */ deletable: boolean; }