atif-sql 0.2.0__py3-none-any.whl

atif_cli/lake.py

@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""The ``atif-sql lake`` subcommand group: rebuild, verify, status, compact.
+
+The lake is the one DuckLake every corpus is queried through
+(:mod:`atif_duck.infrastructure.lake`). ``materialize`` keeps it current one
+batch of published sessions at a time; these commands build it, check it,
+report on it, and keep its file count down.
+
+Lean by construction like :mod:`atif_cli.cron`: this module loads eagerly with
+``atif_cli.app`` (sub-apps must exist at registration time), so everything
+beyond cyclopts and the output helpers is imported inside the command bodies.
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated
+
+import cyclopts
+
+from atif_cli.errors import EXIT_CODES, ClassifiedError
+from atif_cli.output import OutputFormat, emit_error, emit_json, resolve_format
+
+if TYPE_CHECKING:
+    from atif_duck.infrastructure.lake import LakeCorpus, LakeLayout
+    from atif_duck.infrastructure.lake_settings import LakeSettings
+
+lake_app = cyclopts.App(
+    name="lake",
+    help="Build, check and maintain the DuckLake every corpus is queried through.",
+)
+
+
+def lake_layout(lake_root: Path | None) -> tuple[LakeLayout, LakeSettings]:
+    """The lake at ``lake_root`` (default ``ATIF_SQL_LAKE_ROOT``), and the settings it came from."""
+    from atif_duck.infrastructure.lake import LakeLayout
+    from atif_duck.infrastructure.lake_settings import LakeSettings
+
+    settings = LakeSettings()
+    return LakeLayout(lake_root if lake_root is not None else settings.lake_root), settings
+
+
+def writer_memory_limit() -> int:
+    """The same host- and cgroup-derived cap ``query`` runs under."""
+    from atif_cli.app import query_memory_limit_bytes
+
+    return query_memory_limit_bytes()
+
+
+def _fail(kind: str, message: str, hint: str, fmt: OutputFormat) -> None:
+    err = ClassifiedError(kind=kind, exit_code=EXIT_CODES[kind], message=message, hint=hint)
+    emit_error(err, fmt)
+    raise SystemExit(err.exit_code)
+
+
+def discover_corpora(
+    layout: LakeLayout, settings: LakeSettings, roots: list[Path] | None
+) -> list[LakeCorpus]:
+    """The corpora to act on: the named roots, else what the lake holds plus every corpus under the base."""
+    from atif_duck.infrastructure.lake import LakeCorpus, corpus_agent, registered_corpora
+
+    if roots:
+        return [LakeCorpus(root=root, agent=corpus_agent(root)) for root in roots]
+    found: dict[Path, LakeCorpus] = {
+        corpus.root: corpus
+        for corpus in registered_corpora(layout).values()
+        if (corpus.root / "sessions").is_dir()
+    }
+    base = settings.corpus_base
+    if base.is_dir():
+        for child in sorted(base.iterdir()):
+            if (child / "sessions").is_dir() and child not in found:
+                found[child] = LakeCorpus(root=child, agent=corpus_agent(child))
+    return sorted(found.values(), key=lambda corpus: corpus.name)
+
+
+@lake_app.command
+def rebuild(
+    *,
+    corpus_root: Annotated[list[Path] | None, cyclopts.Parameter(consume_multiple=False)] = None,
+    lake_root: Path | None = None,
+    fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
+) -> None:
+    """Load every corpus's per-session artifacts into a fresh lake, then swap it into place.
+
+    Parameters
+    ----------
+    corpus_root
+        A corpus to load; repeat for more. Default: every corpus the current
+        lake holds, plus every directory under ``ATIF_SQL_CORPUS_BASE``
+        (default ``~/.atif-sql/corpus``) that holds ``sessions/``.
+    lake_root
+        The lake to rebuild (default ``ATIF_SQL_LAKE_ROOT``, else
+        ``~/.atif-sql/lake``).
+    fmt
+        ``auto`` = human lines on a TTY, JSON on a pipe.
+
+    Installs the ``ducklake`` DuckDB extension first when it is missing (the
+    one network fetch the lake needs; ``query`` never installs it). The new
+    lake is built beside the old one under the writer lock and renamed into
+    place when complete, so a reader sees one lake or the other.
+    """
+    from atif_duck.infrastructure.lake import (
+        LakeCorpusConflictError,
+        install_ducklake_extension,
+        rebuild_lake,
+    )
+
+    layout, settings = lake_layout(lake_root)
+    corpora = discover_corpora(layout, settings, corpus_root)
+    if not corpora:
+        _fail(
+            "invalid_input",
+            "no corpus to load",
+            "materialize a corpus first, or pass --corpus-root",
+            fmt,
+        )
+    install_ducklake_extension()
+    try:
+        report = rebuild_lake(
+            layout,
+            corpora,
+            batch_size=settings.lake_load_batch_size,
+            lock_timeout_seconds=settings.lake_lock_timeout_seconds,
+            memory_limit_bytes=writer_memory_limit(),
+            stage_workers=settings.lake_stage_workers,
+        )
+    except LakeCorpusConflictError as exc:
+        _fail("invalid_input", str(exc), "give each corpus a distinct directory name", fmt)
+        return
+    payload = {
+        "lake_root": str(report.root),
+        "corpora": [
+            {"corpus": name, "agent": agent, "sessions": sessions}
+            for name, agent, sessions in report.corpora
+        ],
+        "data_files": report.data_files,
+        "seconds": round(report.seconds, 3),
+    }
+    if resolve_format(fmt) is OutputFormat.TABLE:
+        for name, agent, sessions in report.corpora:
+            print(f"loaded {name} ({agent}): {sessions} sessions")
+        print(f"lake: {report.root}  data files: {report.data_files}  in {report.seconds:.1f}s")
+    else:
+        emit_json(payload, fmt)
+
+
+@lake_app.command
+def verify(
+    *,
+    corpus_root: Annotated[list[Path] | None, cyclopts.Parameter(consume_multiple=False)] = None,
+    lake_root: Path | None = None,
+    limit: int = 20,
+    fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
+) -> None:
+    """Compare every session's lake rows with its per-session artifacts; exit 65 on a difference.
+
+    Per table and session: the row count and an order-free content hash of
+    the lake's rows against the same over what the per-session path reads.
+    Exits 0 when all match, 65 (``lake_mismatch``) when any differs, 78
+    (``lake_unavailable``) when there is no lake or its schema is stale.
+
+    Parameters
+    ----------
+    corpus_root
+        Verify only this corpus; repeat for more. Default: every corpus the lake holds.
+    lake_root
+        The lake to verify (default ``ATIF_SQL_LAKE_ROOT``).
+    limit
+        How many differences to list (all are counted).
+    fmt
+        ``auto`` = human lines on a TTY, JSON on a pipe.
+    """
+    from atif_duck.infrastructure.lake import (
+        LakeCorpus,
+        LakeError,
+        corpus_agent,
+        verify_lake,
+    )
+
+    layout, settings = lake_layout(lake_root)
+    corpora = (
+        [LakeCorpus(root=root, agent=corpus_agent(root)) for root in corpus_root]
+        if corpus_root
+        else None
+    )
+    try:
+        report = verify_lake(
+            layout,
+            corpora,
+            memory_limit_bytes=writer_memory_limit(),
+            batch_size=settings.lake_load_batch_size,
+            stage_workers=settings.lake_stage_workers,
+        )
+    except LakeError as exc:
+        _fail("lake_unavailable", str(exc), "run `atif-sql lake rebuild`", fmt)
+        return
+    if report.stale:
+        _fail(
+            "lake_unavailable",
+            f"cannot verify {layout.root}: {', '.join(report.stale)}",
+            "run `atif-sql lake rebuild`",
+            fmt,
+        )
+    sessions = sorted({(m.corpus, m.session_id) for m in report.mismatches})
+    payload = {
+        "lake_root": str(layout.root),
+        "clean": report.clean,
+        "corpora": [{"corpus": name, "sessions": n} for name, n in report.corpora],
+        "mismatched_sessions": len(sessions),
+        "mismatches": [
+            {
+                "corpus": m.corpus,
+                "table": m.table,
+                "session_id": m.session_id,
+                "artifact_rows": m.artifact_rows,
+                "lake_rows": m.lake_rows,
+            }
+            for m in report.mismatches[: max(0, limit)]
+        ],
+    }
+    if resolve_format(fmt) is OutputFormat.TABLE:
+        for name, n in report.corpora:
+            print(f"{name}: {n} sessions checked")
+        for m in report.mismatches[: max(0, limit)]:
+            print(
+                f"  DIFFERS {m.corpus}/{m.session_id} {m.table}: "
+                f"{m.artifact_rows} artifact row(s), {m.lake_rows} lake row(s)",
+                file=sys.stderr,
+            )
+        print("clean" if report.clean else f"{len(sessions)} session(s) differ")
+    else:
+        emit_json(payload, fmt)
+    if not report.clean:
+        raise SystemExit(EXIT_CODES["lake_mismatch"])
+
+
+@lake_app.command
+def status(
+    *,
+    lake_root: Path | None = None,
+    fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
+) -> None:
+    """Report the lake: present, schema current, corpora, snapshots, files, last write.
+
+    Read from the published reader catalog, so it never waits on a writer.
+
+    Parameters
+    ----------
+    lake_root
+        The lake to report on (default ``ATIF_SQL_LAKE_ROOT``).
+    fmt
+        ``auto`` = human lines on a TTY, JSON on a pipe.
+    """
+    from atif_duck.infrastructure.lake import lake_status
+
+    layout, _ = lake_layout(lake_root)
+    state = lake_status(layout).as_dict()
+    if resolve_format(fmt) is OutputFormat.TABLE:
+        print(f"lake root:    {state['root']}")
+        print(
+            f"present:      {state['present']}  (ducklake extension: {state['extension_installed']})"
+        )
+        if state["present"]:
+            print(
+                f"schema:       {'current' if state['schema_current'] else 'STALE ' + ', '.join(state['stale'])}"
+            )
+            for corpus in state["corpora"]:
+                print(
+                    f"corpus:       {corpus['corpus']} ({corpus['agent']}): {corpus['sessions']} sessions"
+                )
+            print(f"snapshots:    {state['snapshots']}")
+            print(
+                f"data files:   {state['data_files']} ({state['data_bytes']:,} bytes), {state['delete_files']} delete files"
+            )
+            print(f"last write:   {state['last_write']}")
+    else:
+        emit_json(state, fmt)
+
+
+@lake_app.command
+def compact(
+    *,
+    expire_older_than_days: int | None = None,
+    memory_limit: str | None = None,
+    lake_root: Path | None = None,
+    fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
+) -> None:
+    """Merge small files, expire old snapshots, and remove the files nothing references.
+
+    Parameters
+    ----------
+    expire_older_than_days
+        Expire snapshots older than this (default ``ATIF_SQL_LAKE_EXPIRE_DAYS``,
+        else 30). Files only expired snapshots referenced are then removed,
+        once an hour old, so a reader in flight keeps every file it names.
+    memory_limit
+        DuckDB's memory budget for the run, as a size (``2GiB``, ``1500MB``),
+        instead of the one derived from the host and cgroup. The writer's own
+        2 GiB ceiling still applies. The nightly refresh lane passes one that
+        fits its memory scope, so a sizing mistake can't shrink the budget
+        until the merge runs out of memory.
+    lake_root
+        The lake to compact (default ``ATIF_SQL_LAKE_ROOT``).
+    fmt
+        ``auto`` = human lines on a TTY, JSON on a pipe.
+    """
+    from atif_duck.infrastructure.lake import (
+        DEFAULT_WRITER_MEMORY_BYTES,
+        LakeError,
+        compact_lake,
+    )
+
+    layout, settings = lake_layout(lake_root)
+    days = (
+        expire_older_than_days if expire_older_than_days is not None else settings.lake_expire_days
+    )
+    if days < 0:
+        _fail("invalid_input", "--expire-older-than-days must be >= 0", "pass 0 or more", fmt)
+    budget = writer_memory_limit()
+    if memory_limit is not None:
+        from atif_cli.app import parse_size
+
+        try:
+            budget = parse_size(memory_limit)
+        except ValueError as exc:
+            _fail("invalid_input", str(exc), "pass a size such as 2GiB or 1500MB", fmt)
+        if budget <= 0:
+            _fail("invalid_input", "--memory-limit must be a positive size", "e.g. 2GiB", fmt)
+    import duckdb
+
+    try:
+        report = compact_lake(
+            layout,
+            expire_older_than_days=days,
+            lock_timeout_seconds=settings.lake_lock_timeout_seconds,
+            memory_limit_bytes=budget,
+        )
+    except LakeError as exc:
+        _fail("lake_unavailable", str(exc), "run `atif-sql lake rebuild`", fmt)
+        return
+    except duckdb.Error as exc:
+        # An out-of-memory merge is the likely one. Nothing was published, so
+        # readers keep the previous catalog; the next run starts over.
+        _fail(
+            "runtime_error",
+            str(exc).splitlines()[0],
+            "raise --memory-limit (the writer's ceiling is 2GiB) or run it again",
+            fmt,
+        )
+        return
+    payload = {
+        "lake_root": str(layout.root),
+        "data_files_before": report.files_before,
+        "data_files_after": report.files_after,
+        "snapshots_before": report.snapshots_before,
+        "snapshots_after": report.snapshots_after,
+        "memory_limit_bytes": min(budget, DEFAULT_WRITER_MEMORY_BYTES),
+        "seconds": round(report.seconds, 3),
+    }
+    if resolve_format(fmt) is OutputFormat.TABLE:
+        print(
+            f"data files: {report.files_before} -> {report.files_after}  "
+            f"snapshots: {report.snapshots_before} -> {report.snapshots_after}  "
+            f"in {report.seconds:.1f}s"
+        )
+    else:
+        emit_json(payload, fmt)
+
+
+__all__ = ["lake_app"]

atif_cli/lake_sessions.py

@@ -0,0 +1,270 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""The lake behind atif-analytics' session port: what ``analyze`` reads when the lake holds the corpus.
+
+atif-analytics may import no package of ours but atif-models, so it declares
+the port (:class:`atif_analytics.domain.ports.SessionSource`) and atif-cli,
+the composition root, adapts atif-duck's :class:`~atif_duck.infrastructure.lake_sessions.LakeSessionReader`
+to it here. The rows become :class:`~atif_analytics.domain.transcript.StepEvent`
+through :func:`~atif_analytics.domain.transcript.step_event`, the projection
+the file source uses too, so both sources apply one set of rules, and the
+author label comes from atif-analytics' rule table (pinned to atif-duck's
+``step_author`` by ``test_authorship_twin_pin.py``).
+
+Per session, the lake is used only when it holds the session's CURRENT
+artifacts: its ``session_meta`` row must carry the ``materialized_at`` and
+``source_mtime_ns`` the session's ``meta.json`` carries now. Anything else (a
+session the lake has not loaded, one a failed or skipped lake write left
+behind, one whose ``meta.json`` cannot be read) is read from its files, so a
+lagging lake can never hand the pipelines an older transcript than the files
+hold. Enumeration and the ``trajectory.json`` mtime bound come from the
+session directories exactly as on the file path, which keeps the checkpoint's
+bounds identical whichever source ran.
+
+Order inside that per-session check matters: the trajectory is stat'ed before
+``meta.json`` is read, so a materialize swapping the session in between leaves
+an OLD mtime beside whatever was read, and the checkpoint re-admits the
+session next run rather than skipping a transcript it never analyzed.
+
+A lake read that fails mid-run (a ``lake rebuild`` swapped the lake out from
+under the open attach) logs one warning, and that read and every later one
+come from the files, so a rebuild never takes a stage down.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import json
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any
+
+from loguru import logger
+
+from atif_analytics.domain.transcript import StepEvent, step_event
+from atif_analytics.infrastructure.corpus_reader import (
+    META_FILENAME,
+    TrajectoryFileSource,
+    complete_session_dirs,
+    last_edge_ts,
+    trajectory_mtime,
+)
+from atif_duck.infrastructure.lake_sessions import LakeReadError
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from datetime import datetime
+    from pathlib import Path
+
+    from atif_analytics.domain.ports import SessionBounds
+    from atif_duck.infrastructure.lake import LakeLayout, LakeUnavailable
+    from atif_duck.infrastructure.lake_sessions import LakeSessionReader, SessionFreshness
+
+#: Sessions per gate-level lake read. One read costs about the same for one
+#: session as for this many (it scans the same row groups' metadata), and the
+#: gates walk the newest-first order, so the reader reads ahead by this much.
+TURNS_BATCH_SIZE: int = 64
+
+#: Sessions per full read (the tool payloads). Smaller: a batch of results can
+#: be large, and the renderer only reads what it admits.
+STEPS_BATCH_SIZE: int = 4
+
+
+def _json(text: str | None) -> Any:
+    return None if text is None else json.loads(text)
+
+
+def _is_current(session_dir: Path, recorded: SessionFreshness | None) -> bool:
+    """True when the lake's row for this session is its current ``meta.json``'s."""
+    if recorded is None:
+        return False
+    try:
+        meta = json.loads((session_dir / META_FILENAME).read_text())
+    except (OSError, ValueError):
+        return False
+    if not isinstance(meta, dict):
+        return False
+    return bool(
+        meta.get("materialized_at") == recorded.materialized_at
+        and meta.get("source_mtime_ns") == recorded.source_mtime_ns
+    )
+
+
+class LakeSessionSource:
+    """atif-analytics' ``SessionSource`` over the lake, falling back to the files per session."""
+
+    name: str = "lake"
+    turns_batch_size: int = TURNS_BATCH_SIZE
+    steps_batch_size: int = STEPS_BATCH_SIZE
+    turns_are_complete: bool = False
+
+    def __init__(self, reader: LakeSessionReader, corpus_root: Path) -> None:
+        self._reader = reader
+        self._sessions_dir = corpus_root / "sessions"
+        self._files = TrajectoryFileSource(corpus_root)
+        self._current: set[str] | None = None
+        #: Sessions the last enumeration found on disk but read from files.
+        self.from_files: tuple[str, ...] = ()
+        #: True once a lake read failed; the rest of the run reads files.
+        self.failed: bool = False
+
+    @classmethod
+    def open(
+        cls, layout: LakeLayout, corpus_root: Path, *, memory_limit_bytes: int | None = None
+    ) -> LakeSessionSource | LakeUnavailable:
+        """Attach the lake for ``corpus_root``, or return why it cannot serve it."""
+        from atif_duck.infrastructure.lake_sessions import LakeSessionReader
+
+        reader = LakeSessionReader.open(layout, corpus_root, memory_limit_bytes=memory_limit_bytes)
+        if not isinstance(reader, LakeSessionReader):
+            return reader
+        return cls(reader, corpus_root)
+
+    def close(self) -> None:
+        """Release the lake connection."""
+        self._reader.close()
+
+    # ------------------------------------------------------------------
+    # Enumeration
+    # ------------------------------------------------------------------
+
+    def session_bounds(self) -> SessionBounds:
+        """Every complete session's bounds; the lake supplies the newest-record bound when current."""
+        freshness: dict[str, SessionFreshness] = {}
+        newest: dict[str, datetime | None] = {}
+        if not self.failed:
+            try:
+                freshness = self._reader.freshness()
+                newest = self._reader.last_edge_ts()
+            except LakeReadError as exc:
+                self._fail(exc)
+                freshness.clear()
+                newest.clear()
+        bounds: SessionBounds = {}
+        current: set[str] = set()
+        from_files: list[str] = []
+        for session_dir in complete_session_dirs(self._sessions_dir):
+            sid = session_dir.name
+            mtime = trajectory_mtime(session_dir)
+            if _is_current(session_dir, freshness.get(sid)):
+                current.add(sid)
+                bounds[sid] = (newest.get(sid), mtime)
+            else:
+                from_files.append(sid)
+                bounds[sid] = (last_edge_ts(session_dir), mtime)
+        self._current = current
+        self.from_files = tuple(from_files)
+        return bounds
+
+    def _fail(self, exc: LakeReadError) -> None:
+        """Stop reading the lake for the rest of the run; every later read comes from files."""
+        if not self.failed:
+            logger.warning(
+                "analyze: a lake read failed ({}); reading the per-session artifacts for the rest of the run",
+                exc,
+            )
+        self.failed = True
+        self._current = set()
+
+    def _split(self, session_ids: Sequence[str]) -> tuple[list[str], list[str]]:
+        """``(read from the lake, read from files)``, enumerating first if nothing has."""
+        if self._current is None:
+            self.session_bounds()
+        current = self._current or set()
+        in_lake = [sid for sid in session_ids if sid in current]
+        on_disk = [sid for sid in session_ids if sid not in current]
+        return in_lake, on_disk
+
+    def batchable(self, session_id: str) -> bool:
+        """True for a session read from the lake; a file fallback is parsed only on request."""
+        return self._current is not None and session_id in self._current
+
+    # ------------------------------------------------------------------
+    # Steps
+    # ------------------------------------------------------------------
+
+    def load_turns(self, session_ids: Sequence[str]) -> dict[str, list[StepEvent]]:
+        """Each session's steps without tool payloads, error flags included."""
+        in_lake, on_disk = self._split(session_ids)
+        try:
+            from_lake = self._lake_turns(in_lake)
+        except LakeReadError as exc:
+            self._fail(exc)
+            from_lake, on_disk = {}, list(session_ids)
+        # A fallback parse yields the payloads; the turns memo holds many
+        # sessions, so keep it as small as the lake's rows.
+        out = {
+            sid: [dataclasses.replace(s, tool_calls=[], tool_results=[]) for s in steps]
+            for sid, steps in self._files.load_turns(on_disk).items()
+        }
+        out.update(from_lake)
+        return out
+
+    def _lake_turns(self, in_lake: list[str]) -> dict[str, list[StepEvent]]:
+        errors = self._reader.error_steps(in_lake)
+        grouped: dict[str, list[StepEvent]] = defaultdict(list)
+        for sid, step_id, ts, source, message, sidechain, compact, uuids in self._reader.turns(
+            in_lake
+        ):
+            grouped[str(sid)].append(
+                step_event(
+                    ts=ts,
+                    source=source,
+                    text=message,
+                    source_uuids=_json(uuids),
+                    is_sidechain=bool(sidechain),
+                    is_compact_summary=bool(compact),
+                    has_error_result=(str(sid), int(step_id)) in errors,
+                )
+            )
+        return {sid: grouped.get(sid, []) for sid in in_lake}
+
+    def load_steps(self, session_ids: Sequence[str]) -> dict[str, list[StepEvent]]:
+        """Each session's steps with every tool call and result, in source order."""
+        in_lake, on_disk = self._split(session_ids)
+        try:
+            from_lake = self._lake_steps(in_lake)
+        except LakeReadError as exc:
+            self._fail(exc)
+            from_lake, on_disk = {}, list(session_ids)
+        out: dict[str, list[StepEvent]] = dict(self._files.load_steps(on_disk))
+        out.update(from_lake)
+        return out
+
+    def _lake_steps(self, in_lake: list[str]) -> dict[str, list[StepEvent]]:
+        calls: dict[tuple[str, int], list[tuple[str | None, object]]] = defaultdict(list)
+        for sid, step_id, name, tool_input in self._reader.tool_calls(in_lake):
+            calls[(str(sid), int(step_id))].append((name, _json(tool_input)))
+        results: dict[tuple[str, int], list[tuple[str | None, object, bool]]] = defaultdict(list)
+        for sid, step_id, call_id, content, is_error in self._reader.tool_results(in_lake):
+            results[(str(sid), int(step_id))].append((call_id, _json(content), is_error is True))
+        grouped: dict[str, list[StepEvent]] = defaultdict(list)
+        for sid, step_id, ts, source, message, sidechain, compact, uuids in self._reader.turns(
+            in_lake
+        ):
+            key = (str(sid), int(step_id))
+            grouped[key[0]].append(
+                step_event(
+                    ts=ts,
+                    source=source,
+                    text=message,
+                    source_uuids=_json(uuids),
+                    is_sidechain=bool(sidechain),
+                    is_compact_summary=bool(compact),
+                    tool_calls=calls.get(key, ()),
+                    tool_results=results.get(key, ()),
+                )
+            )
+        return {sid: grouped.get(sid, []) for sid in in_lake}
+
+    def edges_uuids(self, session_id: str) -> set[str] | None:
+        """The session's raw-record uuids, from the lake when it holds the session's current rows."""
+        in_lake, _ = self._split([session_id])
+        if in_lake:
+            try:
+                return self._reader.edge_uuids(session_id)
+            except LakeReadError as exc:
+                self._fail(exc)
+        return self._files.edges_uuids(session_id)
+
+
+__all__ = ["STEPS_BATCH_SIZE", "TURNS_BATCH_SIZE", "LakeSessionSource"]