diff --git a/src/tau_agent/session/__init__.py b/src/tau_agent/session/__init__.py index 08a3f28e8..cd9ff303a 100644 --- a/src/tau_agent/session/__init__.py +++ b/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", ] diff --git a/src/tau_agent/session/conformance.py b/src/tau_agent/session/conformance.py new file mode 100644 index 000000000..0c56c4b21 --- /dev/null +++ b/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()}" + ) diff --git a/src/tau_coding/session_storage.py b/src/tau_coding/session_storage.py new file mode 100644 index 000000000..54f783f3e --- /dev/null +++ b/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) diff --git a/tests/test_session_storage_conformance.py b/tests/test_session_storage_conformance.py new file mode 100644 index 000000000..70d79d86f --- /dev/null +++ b/tests/test_session_storage_conformance.py @@ -0,0 +1,144 @@ +"""Tests for the SessionStorage conformance checker.""" + +from pathlib import Path + +import pytest + +from tau_agent.session import ( + JsonlSessionStorage, + LabelEntry, + SessionEntry, +) +from tau_agent.session.conformance import ( + SessionStorageConformanceError, + conformance_entries, + verify_session_storage, +) + + +class MemorySessionStorage: + """Minimal conforming in-memory storage (the docs example).""" + + def __init__(self) -> None: + self.entries: list[SessionEntry] = [] + + async def append(self, entry: SessionEntry) -> None: + self.entries.append(entry) + + async def read_all(self) -> list[SessionEntry]: + return list(self.entries) + + +class _DropsEntriesStorage(MemorySessionStorage): + """Broken: silently drops every second append.""" + + def __init__(self) -> None: + super().__init__() + self._seen = 0 + + async def append(self, entry: SessionEntry) -> None: + self._seen += 1 + if self._seen % 2 == 1: + self.entries.append(entry) + + +class _ReordersStorage(MemorySessionStorage): + """Broken: returns entries out of append order.""" + + async def read_all(self) -> list[SessionEntry]: + return list(reversed(self.entries)) + + +class _TruncatesTimestampsStorage(MemorySessionStorage): + """Broken: loses float precision, like a careless column type.""" + + async def read_all(self) -> list[SessionEntry]: + return [ + entry.model_copy(update={"timestamp": float(int(entry.timestamp))}) + for entry in self.entries + ] + + +class _StaleStorage(MemorySessionStorage): + """Broken: a \"fresh\" storage that already contains entries.""" + + def __init__(self) -> None: + super().__init__() + self.entries.append(LabelEntry(label="stale")) + + +class _ForgetfulStorage(MemorySessionStorage): + """Broken across reopens: nothing is durable (state dies with the instance).""" + + +def test_conformance_entries_cover_every_entry_type() -> None: + types = {entry.type for entry in conformance_entries()} + assert types == { + "session_info", + "model_change", + "thinking_level_change", + "message", + "compaction", + "branch_summary", + "label", + "leaf", + "custom", + } + + +def test_conformance_entries_form_a_parent_chain() -> None: + entries = conformance_entries() + for previous, entry in zip(entries, entries[1:], strict=False): + assert entry.parent_id == previous.id + + +@pytest.mark.anyio +async def test_jsonl_storage_passes_conformance(tmp_path: Path) -> None: + path = tmp_path / "sessions" / "conformance.jsonl" + await verify_session_storage( + lambda: JsonlSessionStorage(path), + reopen=lambda: JsonlSessionStorage(path), + ) + + +@pytest.mark.anyio +async def test_memory_storage_passes_without_reopen() -> None: + await verify_session_storage(MemorySessionStorage) + + +@pytest.mark.anyio +async def test_async_factories_are_supported(tmp_path: Path) -> None: + async def make_storage() -> JsonlSessionStorage: + return JsonlSessionStorage(tmp_path / "async.jsonl") + + await verify_session_storage(make_storage) + + +@pytest.mark.anyio +async def test_dropping_entries_fails() -> None: + with pytest.raises(SessionStorageConformanceError, match="expected"): + await verify_session_storage(_DropsEntriesStorage) + + +@pytest.mark.anyio +async def test_reordering_entries_fails() -> None: + with pytest.raises(SessionStorageConformanceError, match="round-trip"): + await verify_session_storage(_ReordersStorage) + + +@pytest.mark.anyio +async def test_losing_timestamp_precision_fails() -> None: + with pytest.raises(SessionStorageConformanceError, match="round-trip"): + await verify_session_storage(_TruncatesTimestampsStorage) + + +@pytest.mark.anyio +async def test_non_empty_fresh_storage_fails() -> None: + with pytest.raises(SessionStorageConformanceError, match="empty session"): + await verify_session_storage(_StaleStorage) + + +@pytest.mark.anyio +async def test_non_durable_storage_fails_reopen_check() -> None: + with pytest.raises(SessionStorageConformanceError, match="durable"): + await verify_session_storage(_ForgetfulStorage, reopen=_ForgetfulStorage) diff --git a/tests/test_session_storage_resolution.py b/tests/test_session_storage_resolution.py new file mode 100644 index 000000000..ab59eef5e --- /dev/null +++ b/tests/test_session_storage_resolution.py @@ -0,0 +1,122 @@ +"""Tests for pluggable session storage resolution in tau_coding.""" + +from pathlib import Path + +import pytest + +from tau_agent.session import JsonlSessionStorage, SessionEntry +from tau_coding.session_manager import CodingSessionRecord +from tau_coding.session_storage import ( + SESSION_STORAGE_ENV_VAR, + SessionStorageResolutionError, + create_session_storage, + load_session_storage_factory, + resolve_session_storage_factory, +) + + +class MemorySessionStorage: + """Record-aware in-memory storage used as a custom-factory target.""" + + def __init__(self, record: CodingSessionRecord) -> None: + self.record = record + self.entries: list[SessionEntry] = [] + + async def append(self, entry: SessionEntry) -> None: + self.entries.append(entry) + + async def read_all(self) -> list[SessionEntry]: + return list(self.entries) + + +def memory_factory(record: CodingSessionRecord) -> MemorySessionStorage: + """Module-level factory target for `module:attribute` resolution tests.""" + return MemorySessionStorage(record) + + +not_callable = "this is not a factory" + + +def _record(tmp_path: Path) -> CodingSessionRecord: + return CodingSessionRecord( + id="session-one", + path=tmp_path / "session-one.jsonl", + cwd=tmp_path, + model="test-model", + title=None, + created_at=0.0, + updated_at=0.0, + ) + + +def test_default_storage_is_jsonl_at_record_path( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv(SESSION_STORAGE_ENV_VAR, raising=False) + record = _record(tmp_path) + + storage = create_session_storage(record) + + assert isinstance(storage, JsonlSessionStorage) + assert storage.path == record.path + + +def test_explicit_factory_wins(tmp_path: Path) -> None: + record = _record(tmp_path) + + storage = create_session_storage(record, memory_factory) + + assert isinstance(storage, MemorySessionStorage) + assert storage.record is record + + +def test_env_var_selects_custom_storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(SESSION_STORAGE_ENV_VAR, f"{__name__}:memory_factory") + record = _record(tmp_path) + + storage = create_session_storage(record) + + assert isinstance(storage, MemorySessionStorage) + assert storage.record is record + + +def test_blank_env_var_means_default(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(SESSION_STORAGE_ENV_VAR, " ") + + assert resolve_session_storage_factory() is None + + +def test_unset_env_var_means_default(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(SESSION_STORAGE_ENV_VAR, raising=False) + + assert resolve_session_storage_factory() is None + + +def test_env_mapping_override_is_supported() -> None: + factory = resolve_session_storage_factory( + {SESSION_STORAGE_ENV_VAR: f"{__name__}:memory_factory"} + ) + + assert factory is memory_factory + + +@pytest.mark.parametrize( + "spec", + [ + "missing-colon", + ":only_attribute", + "only.module:", + "definitely_missing_module_xyz:factory", + f"{__name__}:missing_attribute", + f"{__name__}:not_callable", + ], +) +def test_bad_specs_raise_pointed_errors(spec: str) -> None: + with pytest.raises(SessionStorageResolutionError): + load_session_storage_factory(spec) + + +def test_dotted_attribute_paths_resolve() -> None: + factory = load_session_storage_factory(f"{__name__}:MemorySessionStorage") + + assert factory is MemorySessionStorage