feat(session): add conformance checker and pluggable storage resolution #341

src/tau_agent/session/__init__.py

@@ -2,6 +2,11 @@
 
 from __future__ import annotations
 
+from tau_agent.session.conformance import (
+    SessionStorageConformanceError,
+    conformance_entries,
+    verify_session_storage,
+)
 from tau_agent.session.entries import (
     BaseSessionEntry,
     BranchSummaryEntry,
@@ -40,11 +45,14 @@
     "SessionJsonlError",
     "SessionState",
     "SessionStorage",
+    "SessionStorageConformanceError",
     "SessionTreeError",
     "ThinkingLevelChangeEntry",
+    "conformance_entries",
     "entries_by_id",
     "entries_from_json_lines",
     "entry_from_json_line",
     "entry_to_json_line",
     "path_to_entry",
+    "verify_session_storage",
 ]

src/tau_agent/session/conformance.py

@@ -0,0 +1,184 @@
+"""Conformance checks for third-party SessionStorage implementations.
+
+`SessionStorage` is a small protocol (`append` / `read_all`), which makes it a
+natural extension point: sessions can live in a database, an object store, or
+any other durable backend instead of local JSONL files. This module gives
+those implementations a way to verify they honor the same contract as
+`JsonlSessionStorage`, so the rest of Tau (replay, branching, compaction,
+export) keeps working unchanged.
+
+Typical usage from an external package's test suite:
+
+    from tau_agent.session import verify_session_storage
+
+    async def test_my_storage_conforms(tmp_path):
+        await verify_session_storage(
+            lambda: MyStorage(tmp_path / "db", session_id="one"),
+            reopen=lambda: MyStorage(tmp_path / "db", session_id="one"),
+        )
+
+The checks intentionally use only public session primitives and add no
+dependencies.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import NoReturn
+
+from tau_agent.messages import UserMessage
+from tau_agent.session.entries import (
+    BranchSummaryEntry,
+    CompactionEntry,
+    CustomEntry,
+    LabelEntry,
+    LeafEntry,
+    MessageEntry,
+    ModelChangeEntry,
+    SessionEntry,
+    SessionInfoEntry,
+    ThinkingLevelChangeEntry,
+)
+from tau_agent.session.jsonl import entry_to_json_line
+from tau_agent.session.storage import SessionStorage
+
+type StorageFactory = Callable[[], SessionStorage | Awaitable[SessionStorage]]
+"""A callable returning a SessionStorage (optionally awaitable for async setup)."""
+
+
+class SessionStorageConformanceError(AssertionError):
+    """Raised when a SessionStorage implementation violates the storage contract."""
+
+
+def conformance_entries() -> list[SessionEntry]:
+    """Return one linked session entry of every entry type.
+
+    The entries form a valid parent chain and exercise the shapes storage
+    backends most often get wrong: float timestamps, `None` fields, nested
+    JSON data, and list-valued fields.
+    """
+    info = SessionInfoEntry(cwd="/tmp/conformance", title="Conformance session")
+    model = ModelChangeEntry(parent_id=info.id, model="conformance-model")
+    thinking = ThinkingLevelChangeEntry(parent_id=model.id, thinking_level=None)
+    message = MessageEntry(
+        parent_id=thinking.id,
+        message=UserMessage(content='Hello, storage \u2014 with unicode and "quotes".'),
+    )
+    compaction = CompactionEntry(
+        parent_id=message.id,
+        summary="Summarized one greeting.",
+        replaces_entry_ids=[message.id],
+    )
+    branch = BranchSummaryEntry(
+        parent_id=compaction.id,
+        summary="Abandoned branch summary.",
+        branch_root_id=info.id,
+    )
+    label = LabelEntry(parent_id=branch.id, label="conformance")
+    custom = CustomEntry(
+        parent_id=label.id,
+        namespace="tau_agent.conformance",
+        data={"nested": {"ok": True, "values": [1, 2.5, None]}, "count": 3},
+    )
+    leaf = LeafEntry(parent_id=custom.id, entry_id=custom.id)
+    return [info, model, thinking, message, compaction, branch, label, custom, leaf]
+
+
+async def verify_session_storage(
+    make_storage: StorageFactory,
+    *,
+    reopen: StorageFactory | None = None,
+) -> None:
+    """Verify a SessionStorage implementation against the storage contract.
+
+    Checks, in order:
+
+    1. A fresh storage reads as an empty session (missing == empty, no error).
+    2. Appended entries come back from `read_all` in append order.
+    3. Every entry type round-trips with full serialization fidelity
+       (ids, parent ids, float timestamps, nested JSON, `None` fields).
+    4. `read_all` is repeatable and does not mutate stored entries.
+    5. With `reopen`: entries survive re-constructing the storage, and
+       appending after a reopen preserves earlier entries (append-only).
+
+    `make_storage` must return a **fresh, empty** storage. `reopen` should
+    return a new storage instance backed by the same underlying data; pass it
+    whenever the backend is durable so persistence is covered too.
+
+    Raises `SessionStorageConformanceError` on the first violation.
+    """
+    storage = await _built(make_storage)
+
+    initial = await storage.read_all()
+    if initial != []:
+        _fail(
+            "a fresh storage must read as an empty session (a missing session "
+            f"is an empty list, not an error); got {len(initial)} entries"
+        )
+
+    entries = conformance_entries()
+    for entry in entries:
+        await storage.append(entry)
+
+    _check_entries_match(await storage.read_all(), entries, context="after appending")
+    _check_entries_match(
+        await storage.read_all(),
+        entries,
+        context="on a repeated read_all (reads must not mutate storage)",
+    )
+
+    if reopen is None:
+        return
+
+    reopened = await _built(reopen)
+    _check_entries_match(
+        await reopened.read_all(),
+        entries,
+        context="after reopening the storage (entries must be durable)",
+    )
+
+    extra = LabelEntry(parent_id=entries[-1].id, label="appended-after-reopen")
+    await reopened.append(extra)
+    _check_entries_match(
+        await reopened.read_all(),
+        [*entries, extra],
+        context="after appending post-reopen (storage must stay append-only)",
+    )
+
+
+async def _built(factory: StorageFactory) -> SessionStorage:
+    storage = factory()
+    if inspect.isawaitable(storage):
+        return await storage  # ty:ignore[invalid-return-type]
+    return storage
+
+
+def _fail(message: str) -> NoReturn:
+    raise SessionStorageConformanceError(f"SessionStorage conformance failure: {message}")
+
+
+def _canonical_line(entry: SessionEntry, *, context: str) -> str:
+    try:
+        return entry_to_json_line(entry)
+    except Exception as error:  # noqa: BLE001 - surfaced as a conformance failure
+        _fail(f"{context}: read_all returned a value that is not a valid SessionEntry: {error!r}")
+
+
+def _check_entries_match(
+    got: list[SessionEntry],
+    expected: list[SessionEntry],
+    *,
+    context: str,
+) -> None:
+    if len(got) != len(expected):
+        _fail(f"{context}: expected {len(expected)} entries, read_all returned {len(got)}")
+    for index, (got_entry, expected_entry) in enumerate(zip(got, expected, strict=True)):
+        got_line = _canonical_line(got_entry, context=context)
+        expected_line = entry_to_json_line(expected_entry)
+        if got_line != expected_line:
+            _fail(
+                f"{context}: entry {index} does not round-trip.\n"
+                f"  expected: {expected_line.strip()}\n"
+                f"  got:      {got_line.strip()}"
+            )

src/tau_coding/session_storage.py

@@ -0,0 +1,110 @@
+"""Pluggable session storage resolution for coding sessions.
+
+`tau_agent` defines the `SessionStorage` protocol and `tau_coding` defaults to
+local JSONL files. This module is the seam between the two: it decides which
+storage backend a coding session uses, without adding any backend-specific
+code or dependencies to Tau itself.
+
+Resolution order for `create_session_storage(record)`:
+
+1. An explicit factory argument (embedders wiring things up in Python).
+2. The `TAU_SESSION_STORAGE` environment variable, a `package.module:attribute`
+   spec naming a callable that takes a `CodingSessionRecord` and returns a
+   `SessionStorage` (external packages plugging into the stock CLI/TUI).
+3. The default: `JsonlSessionStorage(record.path)` — existing behavior,
+   unchanged when nothing is configured.
+"""
+
+from __future__ import annotations
+
+import importlib
+import os
+from collections.abc import Callable, Mapping
+from typing import cast
+
+from tau_agent.session import JsonlSessionStorage, SessionStorage
+from tau_coding.session_manager import CodingSessionRecord
+
+SESSION_STORAGE_ENV_VAR = "TAU_SESSION_STORAGE"
+"""Environment variable holding a `package.module:attribute` factory spec."""
+
+type SessionStorageFactory = Callable[[CodingSessionRecord], SessionStorage]
+"""Builds a SessionStorage for one coding-session record.
+
+The record provides a stable `id` to key the session on, plus `cwd`, `path`,
+and metadata. Factories are free to ignore `path` — it only means something to
+the JSONL backend.
+"""
+
+
+class SessionStorageResolutionError(RuntimeError):
+    """Raised when a configured session storage factory cannot be loaded."""
+
+
+def load_session_storage_factory(spec: str) -> SessionStorageFactory:
+    """Load a session storage factory from a `package.module:attribute` spec.
+
+    The attribute part may be dotted (`module:Class.method`). Raises
+    `SessionStorageResolutionError` with a pointed message when the spec is
+    malformed, the module cannot be imported, the attribute is missing, or the
+    resolved object is not callable.
+    """
+    module_name, separator, attribute_path = spec.partition(":")
+    module_name = module_name.strip()
+    attribute_path = attribute_path.strip()
+    if not separator or not module_name or not attribute_path:
+        raise SessionStorageResolutionError(
+            f"Invalid session storage spec {spec!r}: expected the form 'package.module:attribute'."
+        )
+
+    try:
+        target: object = importlib.import_module(module_name)
+    except ImportError as error:
+        raise SessionStorageResolutionError(
+            f"Could not import session storage module {module_name!r} "
+            f"from {SESSION_STORAGE_ENV_VAR}: {error}"
+        ) from error
+
+    for part in attribute_path.split("."):
+        try:
+            target = getattr(target, part)
+        except AttributeError as error:
+            raise SessionStorageResolutionError(
+                f"Session storage module {module_name!r} has no attribute {attribute_path!r}."
+            ) from error
+
+    if not callable(target):
+        raise SessionStorageResolutionError(f"Session storage factory {spec!r} is not callable.")
+    return cast("SessionStorageFactory", target)
+
+
+def resolve_session_storage_factory(
+    env: Mapping[str, str] | None = None,
+) -> SessionStorageFactory | None:
+    """Return the factory configured via `TAU_SESSION_STORAGE`, if any.
+
+    Returns `None` when the variable is unset or blank, which means "use the
+    default JSONL storage". `env` exists for tests; production callers use the
+    process environment.
+    """
+    variables = os.environ if env is None else env
+    spec = variables.get(SESSION_STORAGE_ENV_VAR, "").strip()
+    if not spec:
+        return None
+    return load_session_storage_factory(spec)
+
+
+def create_session_storage(
+    record: CodingSessionRecord,
+    factory: SessionStorageFactory | None = None,
+) -> SessionStorage:
+    """Create the session storage for one coding-session record.
+
+    Uses `factory` when provided, otherwise the `TAU_SESSION_STORAGE`
+    environment factory, otherwise local JSONL at `record.path` (the existing
+    default behavior).
+    """
+    active_factory = factory if factory is not None else resolve_session_storage_factory()
+    if active_factory is None:
+        return JsonlSessionStorage(record.path)
+    return active_factory(record)