statebus 0.1.1__py3-none-any.whl

statebus/__init__.py

@@ -0,0 +1,44 @@
+"""statebus — typed, versioned state for multi-agent workflows."""
+
+from .bus import SlotView, StateBus, TraceEvent, WriteReceipt
+from .errors import (
+    BreakingSchemaChangeError,
+    ConflictError,
+    InvalidArgumentError,
+    InvalidSchemaError,
+    NotJsonSerializableError,
+    SchemaVersionNotFoundError,
+    SlotArchivedError,
+    SlotExistsError,
+    SlotNameError,
+    SlotNotFoundError,
+    SlotVersionNotFoundError,
+    StateBusError,
+    ValidationError,
+    ValidationIssue,
+    ValueTooLargeError,
+)
+
+__version__ = "0.1.1"
+
+__all__ = [
+    "StateBus",
+    "WriteReceipt",
+    "SlotView",
+    "TraceEvent",
+    "StateBusError",
+    "SlotNotFoundError",
+    "SlotExistsError",
+    "SlotArchivedError",
+    "SlotNameError",
+    "SchemaVersionNotFoundError",
+    "SlotVersionNotFoundError",
+    "NotJsonSerializableError",
+    "ValueTooLargeError",
+    "ValidationError",
+    "ValidationIssue",
+    "ConflictError",
+    "BreakingSchemaChangeError",
+    "InvalidArgumentError",
+    "InvalidSchemaError",
+]

statebus/bus.py

@@ -0,0 +1,504 @@
+"""StateBus — typed, versioned state for multi-agent workflows.
+
+Agents read/write named *slots*. Every write is validated against the slot's
+JSON Schema, appended as a new version, and recorded in the handoff trace::
+
+    from statebus import StateBus
+
+    bus = StateBus("team.db")  # durable; StateBus() is in-memory
+    bus.register_slot("customer_record", schema={...}, schema_version="2.1")
+    receipt = bus.write("customer_record", {...}, agent="agent_a")
+    value = bus.read("customer_record")          # plain value
+    old = bus.replay("customer_record", version=2)  # exact value at v2
+"""
+
+from __future__ import annotations
+
+import copy
+import functools
+import json
+import re
+import threading
+from dataclasses import dataclass
+from typing import Any
+
+from . import store as store_module
+from .errors import (
+    BreakingSchemaChangeError,
+    ConflictError,
+    InvalidArgumentError,
+    NotJsonSerializableError,
+    SchemaVersionNotFoundError,
+    SlotArchivedError,
+    SlotExistsError,
+    SlotNameError,
+    SlotNotFoundError,
+    SlotVersionNotFoundError,
+    ValidationError,
+    ValueTooLargeError,
+)
+from .store import (
+    CasMismatch,
+    MemoryStore,
+    SlotRecord,
+    SQLiteStore,
+    Store,
+    TraceRecord,
+    utc_now,
+)
+from .validation import (
+    check_backward_compatible,
+    check_schema_valid,
+    validate_against_schema,
+)
+
+_SLOT_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.\-]{0,63}$")
+DEFAULT_MAX_VERSIONS = 100
+DEFAULT_MAX_VALUE_BYTES = 1_000_000  # 1 MiB per write
+DEFAULT_MAX_TRACE_EVENTS = 100_000
+
+
+def _locked(method):
+    """Serialize mutating bus operations in-process (stores handle the rest)."""
+
+    @functools.wraps(method)
+    def wrapper(self, *args, **kwargs):
+        with self._lock:
+            return method(self, *args, **kwargs)
+
+    return wrapper
+
+
+@dataclass(frozen=True)
+class WriteReceipt:
+    slot: str
+    version: int
+    schema_version: str
+    agent: str
+    ts: str
+
+
+@dataclass(frozen=True)
+class SlotView:
+    """Envelope for a slot's current (or historical) value."""
+
+    slot: str
+    version: int
+    value: Any
+    agent: str
+    schema_version: str
+    ts: str
+
+
+@dataclass(frozen=True)
+class TraceEvent:
+    seq: int
+    ts: str
+    kind: str
+    slot: str | None
+    agent: str | None
+    version: int | None
+    schema_version: str | None
+    detail: dict
+
+
+class StateBus:
+    """Typed, versioned, traced state for multi-agent workflows."""
+
+    def __init__(
+        self,
+        db_path: str | None = None,
+        *,
+        max_value_bytes: int = DEFAULT_MAX_VALUE_BYTES,
+        trace_reads: bool = False,
+        default_agent: str = "unknown",
+    ) -> None:
+        """
+        Args:
+            db_path: Path to a SQLite file for durable storage.
+                ``None`` selects the in-memory backend (nothing is persisted).
+            max_value_bytes: Per-write size cap; larger writes are rejected.
+            trace_reads: Record read events in the trace by default.
+            default_agent: Agent id used when a call does not name one.
+        """
+        self._store: Store = SQLiteStore(db_path) if db_path is not None else MemoryStore()
+        self._max_value_bytes = max_value_bytes
+        self._trace_reads = trace_reads
+        self._default_agent = default_agent
+        self._lock = threading.RLock()
+
+    # -- context manager -------------------------------------------------
+    def __enter__(self) -> StateBus:
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._store.close()
+
+    # -- helpers ----------------------------------------------------------
+    @staticmethod
+    def _check_slot_name(name: str) -> None:
+        if not isinstance(name, str) or not _SLOT_NAME_RE.match(name):
+            raise SlotNameError(name)
+
+    def _get_slot(self, name: str) -> SlotRecord:
+        self._check_slot_name(name)
+        record = self._store.get_slot(name)
+        if record is None:
+            raise SlotNotFoundError(name)
+        return record
+
+    def _schema_for(self, record: SlotRecord, schema_version: str | None) -> tuple[str, dict]:
+        wanted = schema_version or record.active_schema_version
+        for entry in record.schema_versions:
+            if entry["version"] == wanted:
+                return wanted, entry["schema"]
+        raise SchemaVersionNotFoundError(record.name, wanted)
+
+    def _serialize(self, value: Any) -> str:
+        try:
+            # allow_nan=False: NaN/Infinity are not valid JSON and would corrupt
+            # non-Python readers (Go, JS agents via REST/MCP).
+            text = json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False)
+        except (TypeError, ValueError) as exc:
+            raise NotJsonSerializableError(str(exc)) from exc
+        size = len(text.encode("utf-8"))
+        if size > self._max_value_bytes:
+            raise ValueTooLargeError(size, self._max_value_bytes)
+        return text
+
+    def _trace(
+        self,
+        kind: str,
+        slot: str | None,
+        agent: str | None,
+        version: int | None = None,
+        schema_version: str | None = None,
+        detail: dict | None = None,
+    ) -> None:
+        self._store.append_trace(
+            TraceRecord(
+                seq=0,
+                ts=utc_now(),
+                kind=kind,
+                slot=slot,
+                agent=agent,
+                version=version,
+                schema_version=schema_version,
+                detail=detail or {},
+            )
+        )
+
+    # -- slot lifecycle ----------------------------------------------------
+    @_locked
+    def register_slot(
+        self,
+        name: str,
+        schema: dict,
+        schema_version: str = "1.0",
+        *,
+        max_versions: int = DEFAULT_MAX_VERSIONS,
+        trace_reads: bool | None = None,
+        by: str | None = None,
+    ) -> None:
+        """Register a named slot with its first schema version.
+
+        Raises:
+            SlotExistsError: the slot is already registered.
+            SlotNameError: the name is invalid.
+        """
+        self._check_slot_name(name)
+        if not isinstance(schema, (dict, bool)):
+            raise TypeError("schema must be a JSON Schema (dict or boolean)")
+        check_schema_valid(schema)
+        if not schema_version or not isinstance(schema_version, str):
+            raise ValueError("schema_version must be a non-empty string")
+        record = SlotRecord(
+            name=name,
+            schema_versions=[
+                {
+                    "version": schema_version,
+                    "schema": copy.deepcopy(schema),
+                    "breaking": False,
+                    "registered_at": utc_now(),
+                    "registered_by": by or self._default_agent,
+                }
+            ],
+            active_schema_version=schema_version,
+            max_versions=max_versions,
+            trace_reads=self._trace_reads if trace_reads is None else trace_reads,
+            archived=False,
+            created_at=utc_now(),
+        )
+        try:
+            self._store.register_slot(record)
+        except KeyError:
+            raise SlotExistsError(name) from None
+        self._trace(
+            "slot.registered",
+            name,
+            by or self._default_agent,
+            schema_version=schema_version,
+            detail={"max_versions": max_versions},
+        )
+
+    @_locked
+    def register_schema_version(
+        self,
+        slot: str,
+        schema: dict,
+        schema_version: str,
+        *,
+        allow_breaking: bool = False,
+        by: str | None = None,
+    ) -> list[str]:
+        """Register a new schema version for a slot; it becomes the active one.
+
+        Returns the list of detected breakages (empty when compatible).
+        Raises BreakingSchemaChangeError unless ``allow_breaking=True``.
+        """
+        record = self._get_slot(slot)
+        if record.archived:
+            raise SlotArchivedError(slot)
+        if not isinstance(schema, (dict, bool)):
+            raise TypeError("schema must be a JSON Schema (dict or boolean)")
+        check_schema_valid(schema)
+        if any(e["version"] == schema_version for e in record.schema_versions):
+            raise ValueError(
+                f"schema version {schema_version!r} already registered for slot {slot!r}"
+            )
+        old_schema = record.schema_versions[-1]["schema"]
+        breakages = check_backward_compatible(old_schema, schema)
+        if breakages and not allow_breaking:
+            raise BreakingSchemaChangeError(slot, schema_version, breakages)
+        self._store.add_schema_version(
+            slot, copy.deepcopy(schema), schema_version, bool(breakages), by or self._default_agent
+        )
+        self._store.set_active_schema_version(slot, schema_version)
+        self._trace(
+            "slot.schema_version",
+            slot,
+            by or self._default_agent,
+            schema_version=schema_version,
+            detail={"breaking": bool(breakages), "breakages": breakages},
+        )
+        return breakages
+
+    @_locked
+    def archive_slot(self, slot: str, *, by: str | None = None) -> None:
+        """Archive a slot: writes are rejected, history is preserved forever."""
+        record = self._get_slot(slot)
+        if record.archived:
+            return
+        self._store.archive_slot(slot)
+        self._trace("slot.archived", slot, by or self._default_agent)
+
+    def slots(self) -> list[dict]:
+        """List all slots with their metadata (no values)."""
+        return [
+            {
+                "name": r.name,
+                "active_schema_version": r.active_schema_version,
+                "schema_versions": [e["version"] for e in r.schema_versions],
+                "max_versions": r.max_versions,
+                "archived": r.archived,
+                "created_at": r.created_at,
+            }
+            for r in self._store.list_slots()
+        ]
+
+    # -- writes -------------------------------------------------------------
+    @_locked
+    def write(
+        self,
+        slot: str,
+        value: Any,
+        *,
+        agent: str | None = None,
+        expected_version: int | None = None,
+        schema_version: str | None = None,
+    ) -> WriteReceipt:
+        """Validate ``value`` and append it as a new version of ``slot``.
+
+        Args:
+            expected_version: optimistic-concurrency guard — the write fails
+                with :class:`ConflictError` unless this matches the current
+                version. Omit for last-write-wins.
+            schema_version: validate against this registered schema version
+                instead of the active one.
+
+        Every write path funnels through here — there is no raw-write API.
+        """
+        record = self._get_slot(slot)
+        if record.archived:
+            raise SlotArchivedError(slot)
+        if expected_version is not None and (
+            isinstance(expected_version, bool)
+            or not isinstance(expected_version, int)
+            or expected_version < 0
+        ):
+            raise InvalidArgumentError(
+                f"expected_version must be a non-negative integer, got {expected_version!r}"
+            )
+        who = agent or self._default_agent
+        version_label, schema = self._schema_for(record, schema_version)
+        issues = validate_against_schema(schema, value)
+        if issues:
+            raise ValidationError(slot, version_label, issues)
+        payload = self._serialize(value)
+        try:
+            rec = self._store.append_version(slot, payload, who, version_label, expected_version)
+        except KeyError:
+            raise SlotNotFoundError(slot) from None
+        except CasMismatch as exc:
+            raise ConflictError(
+                slot=slot,
+                expected_version=exc.expected_version,
+                current_version=exc.current_version,
+                current_agent=exc.current_agent,
+            ) from None
+        self._trace("write", slot, who, version=rec.version, schema_version=version_label)
+        return WriteReceipt(
+            slot=slot,
+            version=rec.version,
+            schema_version=version_label,
+            agent=who,
+            ts=rec.ts,
+        )
+
+    # -- reads ----------------------------------------------------------------
+    def _view(self, slot: str, rec: store_module.VersionRecord) -> SlotView:
+        return SlotView(
+            slot=slot,
+            version=rec.version,
+            value=json.loads(rec.value_json),
+            agent=rec.agent,
+            schema_version=rec.schema_version,
+            ts=rec.ts,
+        )
+
+    def read(self, slot: str, *, agent: str | None = None, trace: bool | None = None) -> Any:
+        """Return the current value of ``slot`` (plain JSON data)."""
+        record = self._get_slot(slot)
+        rec = self._store.get_latest(slot)
+        if rec is None:
+            raise SlotVersionNotFoundError(slot, 0)
+        if trace if trace is not None else record.trace_reads:
+            self._trace(
+                "read",
+                slot,
+                agent or self._default_agent,
+                version=rec.version,
+                schema_version=rec.schema_version,
+            )
+        return json.loads(rec.value_json)
+
+    def inspect(self, slot: str) -> SlotView:
+        """Return the current value plus its envelope (version, writer, ...)."""
+        self._get_slot(slot)
+        rec = self._store.get_latest(slot)
+        if rec is None:
+            raise SlotVersionNotFoundError(slot, 0)
+        return self._view(slot, rec)
+
+    def history(self, slot: str, *, limit: int = 50, offset: int = 0) -> list[SlotView]:
+        """Newest-first version history for a slot."""
+        self._get_slot(slot)
+        return [self._view(slot, r) for r in self._store.list_versions(slot, limit, offset)]
+
+    def replay(self, slot: str, version: int) -> Any:
+        """Return the exact value ``slot`` held at ``version``."""
+        self._get_slot(slot)
+        rec = self._store.get_version(slot, version)
+        if rec is None:
+            raise SlotVersionNotFoundError(slot, version)
+        return json.loads(rec.value_json)
+
+    def version_count(self, slot: str) -> int:
+        self._get_slot(slot)
+        return self._store.count_versions(slot)
+
+    # -- trace -----------------------------------------------------------------
+    def trace(
+        self,
+        *,
+        slot: str | None = None,
+        agent: str | None = None,
+        since: int | None = None,
+        kinds: list[str] | None = None,
+        limit: int = 100,
+    ) -> list[TraceEvent]:
+        """Query the handoff trace, newest first.
+
+        ``since`` is a trace sequence number — poll ``trace(since=last_seq)``
+        to follow the bus live.
+        """
+        if slot is not None:
+            self._check_slot_name(slot)
+        records = self._store.list_trace(slot, agent, since, kinds, limit)
+        return [
+            TraceEvent(
+                seq=r.seq,
+                ts=r.ts,
+                kind=r.kind,
+                slot=r.slot,
+                agent=r.agent,
+                version=r.version,
+                schema_version=r.schema_version,
+                detail=r.detail,
+            )
+            for r in records
+        ]
+
+    # -- retention ---------------------------------------------------------------
+    @_locked
+    def prune(self, slot: str | None = None, *, by: str | None = None) -> dict[str, int]:
+        """Enforce each slot's ``max_versions`` retention. Returns {slot: pruned}."""
+        targets = [slot] if slot else [r.name for r in self._store.list_slots()]
+        result: dict[str, int] = {}
+        for name in targets:
+            record = self._get_slot(name)
+            pruned = self._store.prune(name, record.max_versions)
+            if pruned:
+                self._trace(
+                    "prune",
+                    name,
+                    by or self._default_agent,
+                    detail={"pruned": pruned, "kept_genesis": True},
+                )
+            result[name] = pruned
+        return result
+
+    @_locked
+    def prune_trace(self, keep_last: int = DEFAULT_MAX_TRACE_EVENTS) -> int:
+        """Bound trace storage: keep only the newest ``keep_last`` trace events.
+
+        Version retention (:meth:`prune`) does not touch the trace; call this
+        (or schedule it) on long-running buses. Returns the number removed.
+        """
+        if keep_last < 1:
+            raise ValueError("keep_last must be >= 1")
+        return self._store.prune_trace(keep_last)
+
+    def export_slot(self, slot: str) -> dict:
+        """Export a slot's full metadata + version history (backup before pruning)."""
+        record = self._get_slot(slot)
+        return {
+            "slot": record.name,
+            "active_schema_version": record.active_schema_version,
+            "schema_versions": copy.deepcopy(record.schema_versions),
+            "archived": record.archived,
+            "exported_at": utc_now(),
+            "versions": [
+                {
+                    "version": v.version,
+                    "value": v.value,
+                    "agent": v.agent,
+                    "schema_version": v.schema_version,
+                    "ts": v.ts,
+                }
+                for v in self.history(slot, limit=self._store.count_versions(slot))
+            ],
+        }