mycelium-runtime 1.1.0__py3-none-any.whl

mycelium/__init__.py

@@ -0,0 +1,121 @@
+"""Mycelium runtime — failure prevention for AI agents."""
+
+from mycelium.action_ledger import (
+    ActionLedger,
+    FileLedgerStorage,
+    InMemoryLedgerStorage,
+    LedgerEntry,
+    LedgerError,
+    LedgerPendingError,
+    LedgerStorage,
+    get_ledger,
+    ledger,
+    ledger_sync,
+)
+from mycelium.audit_receipt import (
+    AuditReceiptEmitter,
+    AuditReceiptError,
+    AuditReceiptRecord,
+    FileAuditReceiptStorage,
+    InMemoryAuditReceiptStorage,
+    verify_receipt,
+)
+from mycelium.config import (
+    ConfigError,
+    MyceliumConfig,
+    load_config,
+    load_config_from_string,
+)
+from mycelium.history_guard import HistoryGuard, HistoryTruncatedError
+from mycelium.message_validator import MessageValidationError, MessageValidator
+from mycelium.protect import protect, protect_sync
+from mycelium.session import Session
+from mycelium.state_flush import (
+    FileStateFlushStorage,
+    InMemoryStateFlushStorage,
+    StateFlush,
+    StateFlushError,
+    StateSnapshot,
+)
+from mycelium.storage.postgres_ledger import PostgresLedgerStorage, PostgresTaskLedgerStorage
+from mycelium.storage.redis_ledger import RedisLedgerStorage, RedisTaskLedgerStorage
+from mycelium.task_ledger import (
+    TaskFileLedgerStorage,
+    TaskInMemoryLedgerStorage,
+    TaskLedger,
+    TaskLedgerEntry,
+    TaskLedgerError,
+    TaskLedgerPendingError,
+    TaskLedgerStorage,
+    get_task_ledger,
+    task_ledger,
+    task_ledger_sync,
+)
+from mycelium.tool_boundary import (
+    ToolBoundaryError,
+    ToolBoundaryExhaustedError,
+    bounded,
+    bounded_sync,
+    tool_error_message,
+)
+from mycelium.tool_registry import ToolRegistry
+from mycelium.tool_runner import ToolRunner
+
+__version__ = "1.1.0"
+
+__all__ = [
+    "ActionLedger",
+    "FileLedgerStorage",
+    "InMemoryLedgerStorage",
+    "LedgerEntry",
+    "LedgerError",
+    "LedgerPendingError",
+    "LedgerStorage",
+    "get_ledger",
+    "ledger",
+    "ledger_sync",
+    "AuditReceiptEmitter",
+    "AuditReceiptError",
+    "AuditReceiptRecord",
+    "FileAuditReceiptStorage",
+    "InMemoryAuditReceiptStorage",
+    "verify_receipt",
+    "TaskFileLedgerStorage",
+    "TaskInMemoryLedgerStorage",
+    "TaskLedger",
+    "TaskLedgerEntry",
+    "TaskLedgerError",
+    "TaskLedgerPendingError",
+    "TaskLedgerStorage",
+    "RedisLedgerStorage",
+    "RedisTaskLedgerStorage",
+    "PostgresLedgerStorage",
+    "PostgresTaskLedgerStorage",
+    "get_task_ledger",
+    "task_ledger",
+    "task_ledger_sync",
+    "ConfigError",
+    "MyceliumConfig",
+    "load_config",
+    "load_config_from_string",
+    "protect",
+    "protect_sync",
+    "Session",
+    "StateFlush",
+    "StateFlushError",
+    "StateSnapshot",
+    "FileStateFlushStorage",
+    "InMemoryStateFlushStorage",
+    "MessageValidator",
+    "MessageValidationError",
+    "HistoryGuard",
+    "HistoryTruncatedError",
+    "bounded",
+    "bounded_sync",
+    "ToolBoundaryError",
+    "ToolBoundaryExhaustedError",
+    "tool_error_message",
+    "ToolRegistry",
+    "ToolRunner",
+    "__version__",
+]

mycelium/__main__.py

@@ -0,0 +1,64 @@
+"""CLI entrypoint: ``mycelium init`` scaffolds a config file in the user's project."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from importlib import resources
+from pathlib import Path
+
+_TEMPLATE_FULL = "mycelium.template.yaml"
+_TEMPLATE_MINIMAL = "mycelium.minimal.yaml"
+
+
+def _load_template(minimal: bool) -> str:
+    filename = _TEMPLATE_MINIMAL if minimal else _TEMPLATE_FULL
+    path = resources.files("mycelium") / "templates" / filename
+    return path.read_text(encoding="utf-8")
+
+
+def cmd_init(output: Path, *, minimal: bool, force: bool) -> int:
+    if output.exists() and not force:
+        print(f"error: {output} already exists (use --force to overwrite)", file=sys.stderr)
+        return 1
+    output.write_text(_load_template(minimal), encoding="utf-8")
+    variant = "minimal" if minimal else "full"
+    print(f"Wrote {output} ({variant} template)")
+    print("Next: edit tool/task names, then load_config(...) in your agent code.")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="mycelium",
+        description="Mycelium runtime — scaffold config and utilities",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    init_parser = sub.add_parser("init", help="Create mycelium.yaml in the current project")
+    init_parser.add_argument(
+        "-o",
+        "--output",
+        type=Path,
+        default=Path("mycelium.yaml"),
+        help="Output path (default: ./mycelium.yaml)",
+    )
+    init_parser.add_argument(
+        "--minimal",
+        action="store_true",
+        help="Use the smaller template (fewer commented examples)",
+    )
+    init_parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite an existing file",
+    )
+
+    args = parser.parse_args(argv)
+    if args.command == "init":
+        return cmd_init(args.output, minimal=args.minimal, force=args.force)
+    return 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mycelium/action_ledger.py

@@ -0,0 +1,431 @@
+"""ActionLedger — AF-002 durable action records and idempotency guard."""
+
+from __future__ import annotations
+
+import functools
+import hashlib
+import json
+import time
+import uuid
+import warnings
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
+
+from mycelium.session import Session, _session_var
+from mycelium.storage.json_file import LockedJsonDictFile
+
+if TYPE_CHECKING:
+    from mycelium.audit_receipt import AuditReceiptEmitter
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class LedgerError(Exception):
+    """Raised when the action ledger cannot record or verify an action."""
+
+
+class LedgerPendingError(Exception):
+    """Raised when the same request is already in-flight."""
+
+
+@dataclass(frozen=True)
+class LedgerEntry:
+    """Immutable record of a single tool invocation."""
+
+    request_id: str
+    tool: str
+    args: list[Any]
+    kwargs: dict[str, Any]
+    status: str  # "in-flight" | "completed" | "failed"
+    result: Any = None
+    error: str | None = None
+    started_at: float = field(default_factory=time.time)
+    finished_at: float | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "request_id": self.request_id,
+            "tool": self.tool,
+            "args": self.args,
+            "kwargs": self.kwargs,
+            "status": self.status,
+            "result": self.result,
+            "error": self.error,
+            "started_at": self.started_at,
+            "finished_at": self.finished_at,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> LedgerEntry:
+        return cls(**data)
+
+
+class LedgerStorage:
+    """Backend interface for durable action ledger records."""
+
+    def get(self, request_id: str) -> LedgerEntry | None:
+        """Return the entry for request_id, or None if not found."""
+        raise NotImplementedError
+
+    def set(self, entry: LedgerEntry) -> None:
+        """Persist entry, replacing any existing entry with the same request_id."""
+        raise NotImplementedError
+
+    def try_claim_inflight(
+        self,
+        entry: LedgerEntry,
+    ) -> tuple[str, LedgerEntry | None]:
+        """Atomically claim an in-flight entry.
+
+        Returns ``("claimed", None)``, ``("completed", entry)``, or
+        ``("in_flight", entry)``. Redis/Postgres backends override with
+        atomic primitives; file storage uses an exclusive lock.
+        """
+        from mycelium.storage._helpers import default_try_claim_inflight
+
+        return default_try_claim_inflight(self, entry)
+
+    def list_all(self) -> list[LedgerEntry]:
+        """Return all entries. Intended for debugging/auditing only."""
+        raise NotImplementedError
+
+
+class InMemoryLedgerStorage(LedgerStorage):
+    """Default in-memory storage. Survives within the process only."""
+
+    def __init__(self) -> None:
+        self._entries: dict[str, LedgerEntry] = {}
+
+    def get(self, request_id: str) -> LedgerEntry | None:
+        return self._entries.get(request_id)
+
+    def set(self, entry: LedgerEntry) -> None:
+        self._entries[entry.request_id] = entry
+
+    def list_all(self) -> list[LedgerEntry]:
+        return list(self._entries.values())
+
+
+class FileLedgerStorage(LedgerStorage):
+    """JSON-file-backed storage with ``fcntl`` locking for multi-process safety."""
+
+    def __init__(self, path: str | Path) -> None:
+        self._file = LockedJsonDictFile(path)
+
+    def get(self, request_id: str) -> LedgerEntry | None:
+        def read(data: dict[str, dict[str, Any]]) -> LedgerEntry | None:
+            raw = data.get(request_id)
+            if raw is None:
+                return None
+            return LedgerEntry.from_dict(raw)
+
+        return self._file.read_modify_write_no_save(read)
+
+    def set(self, entry: LedgerEntry) -> None:
+        def mutate(data: dict[str, dict[str, Any]]) -> None:
+            data[entry.request_id] = entry.to_dict()
+
+        self._file.read_modify_write(mutate)
+
+    def try_claim_inflight(
+        self,
+        entry: LedgerEntry,
+    ) -> tuple[str, LedgerEntry | None]:
+        outcome: list[tuple[str, LedgerEntry | None]] = []
+
+        def mutate(data: dict[str, dict[str, Any]]) -> None:
+            raw = data.get(entry.request_id)
+            if raw is not None:
+                existing = LedgerEntry.from_dict(raw)
+                if existing.status == "completed":
+                    outcome.append(("completed", existing))
+                    return
+                if existing.status == "in-flight":
+                    outcome.append(("in_flight", existing))
+                    return
+            data[entry.request_id] = entry.to_dict()
+            outcome.append(("claimed", None))
+
+        self._file.read_modify_write(mutate)
+        return outcome[0]
+
+    def list_all(self) -> list[LedgerEntry]:
+        data = self._file.load()
+        return [LedgerEntry.from_dict(raw) for raw in data.values()]
+
+
+class ActionLedger:
+    """Durable ledger of tool invocations for idempotency and audit."""
+
+    def __init__(self, storage: LedgerStorage | None = None) -> None:
+        self._storage = storage if storage is not None else InMemoryLedgerStorage()
+
+    # --- public API ---
+
+    def get(self, request_id: str) -> LedgerEntry | None:
+        return self._storage.get(request_id)
+
+    def claim(
+        self,
+        request_id: str,
+        tool: str,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> LedgerEntry:
+        """Claim a request idempotency key before execution.
+
+        Returns the existing completed entry if the request already succeeded.
+        Raises LedgerPendingError if the request is currently in-flight.
+        """
+        bound = _bind_args(args, kwargs)
+        entry = LedgerEntry(
+            request_id=request_id,
+            tool=tool,
+            args=bound["args"],
+            kwargs=bound["kwargs"],
+            status="in-flight",
+        )
+        outcome, existing = self._storage.try_claim_inflight(entry)
+        if outcome == "completed" and existing is not None:
+            return existing
+        if outcome == "in_flight":
+            raise LedgerPendingError(
+                f"Tool {tool!r} request {request_id!r} is already in-flight"
+            )
+        return entry
+
+    def complete(self, request_id: str, result: Any) -> LedgerEntry:
+        existing = self._storage.get(request_id)
+        if existing is None:
+            raise LedgerError(f"Cannot complete unknown request {request_id!r}")
+        entry = LedgerEntry(
+            request_id=existing.request_id,
+            tool=existing.tool,
+            args=existing.args,
+            kwargs=existing.kwargs,
+            status="completed",
+            result=result,
+            started_at=existing.started_at,
+            finished_at=time.time(),
+        )
+        self._storage.set(entry)
+        return entry
+
+    def fail(self, request_id: str, error: BaseException) -> LedgerEntry:
+        existing = self._storage.get(request_id)
+        if existing is None:
+            raise LedgerError(f"Cannot fail unknown request {request_id!r}")
+        entry = LedgerEntry(
+            request_id=existing.request_id,
+            tool=existing.tool,
+            args=existing.args,
+            kwargs=existing.kwargs,
+            status="failed",
+            error=f"{type(error).__name__}: {error}",
+            started_at=existing.started_at,
+            finished_at=time.time(),
+        )
+        self._storage.set(entry)
+        return entry
+
+    # --- request id derivation ---
+
+    def derive_request_id(
+        self,
+        tool: str,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> str:
+        """Determine the request id for a tool invocation.
+
+        Priority:
+        1. kwargs["request_id"]
+        2. kwargs["tool_call_id"]
+        3. Session-derived id (run + tool + args hash)
+        4. Random UUID (no idempotency, still audited)
+
+        Note: valid repeats within the same Session with identical args will be
+        deduplicated unless an explicit request_id is supplied.
+        """
+        if "request_id" in kwargs:
+            return str(kwargs["request_id"])
+        if "tool_call_id" in kwargs:
+            return str(kwargs["tool_call_id"])
+
+        session = _session_var.get()
+        if session is not None:
+            return self._session_request_id(session, tool, args, kwargs)
+
+        warnings.warn(
+            f"Tool {tool!r} has no request_id, tool_call_id, or Session; "
+            "ActionLedger cannot deduplicate this call. A random UUID will be used.",
+            stacklevel=4,
+        )
+        return f"no-session:{tool}:{uuid.uuid4()}"
+
+    def _session_request_id(
+        self, session: Session, tool: str, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> str:
+        # Stable within the process for the lifetime of the Session object.
+        run_key = f"run-{id(session)}"
+        args_hash = self._hash_args(args, kwargs)
+        return f"{run_key}:{tool}:{args_hash}"
+
+    @staticmethod
+    def _hash_args(args: tuple[Any, ...], kwargs: dict[str, Any]) -> str:
+        payload = json.dumps(
+            {"args": args, "kwargs": kwargs},
+            sort_keys=True,
+            default=str,
+        )
+        return hashlib.sha256(payload.encode()).hexdigest()[:16]
+
+
+def _bind_args(args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Store a serializable snapshot of the call arguments."""
+    return {
+        "args": list(args),
+        "kwargs": dict(kwargs),
+    }
+
+
+def _drop_ledger_keys(kwargs: dict[str, Any]) -> dict[str, Any]:
+    """Remove Mycelium bookkeeping keys before calling the actual tool."""
+    return {k: v for k, v in kwargs.items() if k not in ("request_id", "tool_call_id")}
+
+
+def _emit_tool_receipt(
+    audit_emitter: AuditReceiptEmitter | None,
+    ledger: ActionLedger,
+    request_id: str,
+) -> None:
+    if audit_emitter is None:
+        return
+    entry = ledger.get(request_id)
+    if entry is not None and entry.status in ("completed", "failed"):
+        audit_emitter.emit_from_tool_entry(entry)
+
+
+def _run_ledgered(
+    func: Callable[P, R],
+    tool_name: str,
+    ledger: ActionLedger,
+    args: P.args,
+    kwargs: P.kwargs,
+    audit_emitter: AuditReceiptEmitter | None = None,
+) -> R:
+    request_id = ledger.derive_request_id(tool_name, args, kwargs)
+    clean_kwargs = _drop_ledger_keys(kwargs)
+    existing = ledger.claim(request_id, tool_name, args, clean_kwargs)
+    if existing.status == "completed":
+        return existing.result
+
+    try:
+        result = func(*args, **clean_kwargs)
+    except Exception as exc:
+        ledger.fail(request_id, exc)
+        _emit_tool_receipt(audit_emitter, ledger, request_id)
+        raise
+
+    ledger.complete(request_id, result)
+    _emit_tool_receipt(audit_emitter, ledger, request_id)
+    return result
+
+
+async def _run_ledgered_async(
+    func: Callable[P, Awaitable[R]],
+    tool_name: str,
+    ledger: ActionLedger,
+    args: P.args,
+    kwargs: P.kwargs,
+    audit_emitter: AuditReceiptEmitter | None = None,
+) -> R:
+    request_id = ledger.derive_request_id(tool_name, args, kwargs)
+    clean_kwargs = _drop_ledger_keys(kwargs)
+    existing = ledger.claim(request_id, tool_name, args, clean_kwargs)
+    if existing.status == "completed":
+        return existing.result
+
+    try:
+        result = await func(*args, **clean_kwargs)
+    except Exception as exc:
+        ledger.fail(request_id, exc)
+        _emit_tool_receipt(audit_emitter, ledger, request_id)
+        raise
+
+    ledger.complete(request_id, result)
+    _emit_tool_receipt(audit_emitter, ledger, request_id)
+    return result
+
+
+def _mark_ledgered(wrapper: Callable[..., Any], ledger: ActionLedger) -> None:
+    wrapper._mycelium_ledger = True  # type: ignore[attr-defined]
+    wrapper._mycelium_ledger_instance = ledger  # type: ignore[attr-defined]
+
+
+def ledger(
+    storage: LedgerStorage | None = None,
+    audit_emitter: AuditReceiptEmitter | None = None,
+) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+    """Decorator that records async tool invocations in an ActionLedger."""
+
+    action_ledger = ActionLedger(storage=storage)
+
+    def decorator(func: Callable[P, Awaitable[R]]) -> Callable[P, Awaitable[R]]:
+        tool_name = func.__name__
+
+        @functools.wraps(func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            return await _run_ledgered_async(
+                func, tool_name, action_ledger, args, kwargs, audit_emitter
+            )
+
+        _mark_ledgered(wrapper, action_ledger)
+        return wrapper
+
+    return decorator
+
+
+def ledger_sync(
+    storage: LedgerStorage | None = None,
+    audit_emitter: AuditReceiptEmitter | None = None,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Decorator that records sync tool invocations in an ActionLedger."""
+
+    action_ledger = ActionLedger(storage=storage)
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        tool_name = func.__name__
+
+        @functools.wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            return _run_ledgered(
+                func, tool_name, action_ledger, args, kwargs, audit_emitter
+            )
+
+        _mark_ledgered(wrapper, action_ledger)
+        return wrapper
+
+    return decorator
+
+
+def get_ledger(func: Callable[..., Any]) -> ActionLedger | None:
+    """Return the ActionLedger attached to a wrapped function, if any."""
+    return getattr(func, "_mycelium_ledger_instance", None)
+
+
+__all__ = [
+    "ActionLedger",
+    "FileLedgerStorage",
+    "InMemoryLedgerStorage",
+    "LedgerEntry",
+    "LedgerError",
+    "LedgerPendingError",
+    "LedgerStorage",
+    "get_ledger",
+    "ledger",
+    "ledger_sync",
+]