agent-saga 0.1.0__py3-none-any.whl

agent_saga/__init__.py

@@ -0,0 +1,127 @@
+"""agent-saga -- transactional boundaries for non-deterministic AI agents.
+
+The core of AgentRollback. Three ideas, in order of commercial importance:
+
+  1. Compensation is typed (REVERSIBLE / COMPENSABLE / IRREVERSIBLE). "Undo"
+     is not one thing.
+  2. The pre-flight gate refuses uncompensable actions *before* they happen.
+  3. Compensations are derived at runtime from the forward call's result,
+     because the agent -- not a developer at authoring time -- chose the action.
+"""
+
+from .context import RollbackReport, SagaAborted, SagaContext
+from .decorator import current_saga, saga, saga_scope, tool
+from .gate import (
+    Decision,
+    GateContext,
+    PreFlightGate,
+    PreFlightViolation,
+    Rule,
+    Verdict,
+    arg_exceeds,
+    semantics_is,
+)
+from .recovery import (
+    DanglingSaga,
+    DanglingStep,
+    RecoveryDaemon,
+    RecoveryOutcome,
+    Resolution,
+    parse_wal,
+    recovery_token,
+)
+from .registry import compensator, registered, resolve
+from .semantics import ActionSemantics, Compensation, SagaStep, StepState
+from .durable import (
+    FileSnapshotStore,
+    SnapshotStore,
+    StaleFile,
+    get_snapshot_store,
+    restore_file,
+    set_snapshot_store,
+    snapshot_file,
+)
+from .encryption import (
+    EncryptedRecordError,
+    FernetEncryptor,
+    WALEncryptor,
+    generate_key,
+    get_wal_encryptor,
+    set_wal_encryptor,
+)
+from .gc import GCReport, SnapshotGC
+from .locks import FileLock, InProcessLock, RecoveryLock
+from .snapshot import (
+    AttributeSnapshot,
+    MappingSnapshot,
+    SequenceSnapshot,
+    SetSnapshot,
+    SnapshotStrategy,
+    auto_strategy,
+    reversible,
+)
+from .wal import AsyncWAL, BackpressurePolicy, WALBackpressure
+
+__version__ = "0.1.0"
+__author__ = "SagaOps"
+
+__all__ = [
+    "ActionSemantics",
+    "AsyncWAL",
+    "BackpressurePolicy",
+    "WALBackpressure",
+    "Compensation",
+    "DanglingSaga",
+    "DanglingStep",
+    "Decision",
+    "RecoveryDaemon",
+    "RecoveryOutcome",
+    "Resolution",
+    "compensator",
+    "parse_wal",
+    "recovery_token",
+    "registered",
+    "resolve",
+    "GateContext",
+    "PreFlightGate",
+    "PreFlightViolation",
+    "RollbackReport",
+    "Rule",
+    "SagaAborted",
+    "SagaContext",
+    "SagaStep",
+    "StepState",
+    "Verdict",
+    "arg_exceeds",
+    "current_saga",
+    "saga",
+    "saga_scope",
+    "semantics_is",
+    "tool",
+    "AttributeSnapshot",
+    "MappingSnapshot",
+    "SequenceSnapshot",
+    "SetSnapshot",
+    "SnapshotStrategy",
+    "auto_strategy",
+    "reversible",
+    "FileSnapshotStore",
+    "SnapshotStore",
+    "StaleFile",
+    "get_snapshot_store",
+    "restore_file",
+    "set_snapshot_store",
+    "snapshot_file",
+    "GCReport",
+    "SnapshotGC",
+    "EncryptedRecordError",
+    "FernetEncryptor",
+    "WALEncryptor",
+    "generate_key",
+    "get_wal_encryptor",
+    "set_wal_encryptor",
+    "FileLock",
+    "InProcessLock",
+    "RecoveryLock",
+    "__version__",
+]

agent_saga/adapters/__init__.py

@@ -0,0 +1,11 @@
+"""Framework adapters.
+
+Each adapter makes agent-saga transparent to an existing agent framework: you
+keep writing tools and graphs the framework's way, and wrapping them routes
+every tool execution through a saga boundary.
+
+Adapters import their framework lazily, inside functions -- importing this
+package pulls in nothing. `from agent_saga.adapters.langgraph import wrap_tool`.
+"""
+
+__all__: list[str] = []

agent_saga/adapters/_common.py

@@ -0,0 +1,53 @@
+"""Framework-agnostic saga routing shared by the connector adapters.
+
+Every adapter reduces to the same thing: given some way to actually call the
+underlying tool, produce a coroutine that -- inside a saga -- routes through
+`SagaContext.execute` (gate, WAL, compensation) and -- outside one -- calls the
+tool untouched. Only the packaging around this differs per framework, and that
+packaging is where the lazy framework import lives.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Awaitable, Callable, Optional
+
+from ..decorator import current_saga
+from ..semantics import ActionSemantics, CompensationFactory
+
+
+def build_runner(
+    call: Callable[..., Awaitable[Any]],
+    *,
+    name: str,
+    semantics: ActionSemantics,
+    compensate: Optional[CompensationFactory] = None,
+    timeout: Optional[float] = None,
+) -> Callable[..., Awaitable[Any]]:
+    """Wrap an async `call(**kwargs)` so it records on the active saga.
+
+    Passes the tool arguments as `policy_args` so pre-flight threshold rules can
+    see them -- an argument captured only in the forward closure is invisible to
+    the gate, the exact bug that once let a connector bypass a limit.
+    """
+
+    async def _run(**kwargs: Any) -> Any:
+        async def _forward() -> Any:
+            return await call(**kwargs)
+
+        ctx = current_saga()
+        if ctx is None:
+            return await _forward()
+        return await ctx.execute(
+            tool=name,
+            semantics=semantics,
+            forward=_forward,
+            compensate=compensate,
+            policy_args=dict(kwargs),
+            timeout=timeout,
+        )
+
+    _run.__name__ = f"saga_{name}".replace(".", "_").replace("-", "_")
+    return _run
+
+
+__all__ = ["build_runner"]

agent_saga/adapters/autogen.py

@@ -0,0 +1,96 @@
+"""AutoGen adapter.
+
+AutoGen has moved through several tool APIs (pyautogen's `register_function`,
+autogen-core's `FunctionTool`, AG2), but they share one root: a tool is a plain
+Python callable whose name and docstring become the schema the model sees. So
+this adapter wraps at that stable layer -- a callable in, a saga-aware callable
+out -- which works across those variants without pinning to one.
+
+    saga_charge = wrap_tool(charge, semantics=COMPENSABLE, compensate=...)
+    # register saga_charge with your agent exactly as you would `charge`
+
+The returned callable keeps `__name__`, `__doc__`, and signature metadata, and
+routes through the active saga (passing through untouched outside one). The
+routing core is the shared `build_runner`.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+from typing import Any, Callable, Optional
+
+from ..context import SagaAborted
+from ..decorator import saga_scope
+from ..semantics import ActionSemantics, CompensationFactory
+from ._common import build_runner
+
+
+def wrap_tool(
+    fn: Callable[..., Any],
+    *,
+    semantics: ActionSemantics,
+    compensate: Optional[CompensationFactory] = None,
+    timeout: Optional[float] = None,
+    name: Optional[str] = None,
+) -> Callable[..., Any]:
+    """Wrap a tool callable so its execution records on the active saga.
+
+    `fn` may be sync or async. The wrapper is always async (AutoGen awaits async
+    tools), preserves the original name/docstring/signature so AutoGen builds the
+    same schema, and derives the compensation from the call's result at runtime.
+    """
+    tool_name = name or getattr(fn, "__name__", "tool")
+
+    if asyncio.iscoroutinefunction(fn):
+        async def _call(**kwargs: Any) -> Any:
+            return await fn(**kwargs)
+    else:
+        async def _call(**kwargs: Any) -> Any:
+            return await asyncio.to_thread(lambda: fn(**kwargs))
+
+    runner = build_runner(
+        _call, name=tool_name, semantics=semantics,
+        compensate=compensate, timeout=timeout,
+    )
+
+    @functools.wraps(fn)
+    async def wrapper(**kwargs: Any) -> Any:
+        return await runner(**kwargs)
+
+    # functools.wraps copies __wrapped__/__doc__/__name__; make sure the name
+    # override (if any) wins for schema generation.
+    wrapper.__name__ = tool_name
+    return wrapper
+
+
+async def saga_run(
+    run: Callable[[Any], Any],
+    *,
+    gate: Any = None,
+    wal: Any = None,
+    halt_on_compensation_failure: bool = True,
+    reraise: bool = True,
+) -> Any:
+    """Run an AutoGen conversation inside a saga boundary.
+
+    AutoGen's entry points vary (`a_initiate_chat`, `run`, a team's `run`), so
+    this takes a coroutine `run(saga_context) -> result` rather than guessing.
+    On any exception, wrapped tools that executed are compensated LIFO and
+    `SagaAborted` is raised unless `reraise=False`.
+
+        await saga_run(lambda saga: agent.a_initiate_chat(other, message="..."))
+    """
+    try:
+        async with saga_scope(
+            gate=gate, wal=wal,
+            halt_on_compensation_failure=halt_on_compensation_failure,
+        ) as ctx:
+            return await run(ctx)
+    except SagaAborted as aborted:
+        if reraise:
+            raise
+        return aborted.report
+
+
+__all__ = ["wrap_tool", "saga_run"]

agent_saga/adapters/crewai.py

@@ -0,0 +1,131 @@
+"""CrewAI adapter.
+
+Drop a saga boundary around a CrewAI crew, and wrap the tools whose effects
+must unwind. A wrapped tool keeps its name, description, and args schema, so the
+agent and the crew treat it exactly like the original.
+
+CrewAI tools are typically *synchronous* (`BaseTool._run`), unlike LangChain's
+async-first tools. The wrapper bridges that: the saga runs the sync `_run` on a
+worker thread (so it never blocks the event loop), while presenting the tool to
+CrewAI unchanged.
+
+The routing core lives in `.._common.build_runner` and is tested without CrewAI
+installed; CrewAI is imported lazily, only inside `wrap_tool`.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable, Optional
+
+from ..context import SagaAborted
+from ..decorator import saga_scope
+from ..semantics import ActionSemantics, CompensationFactory
+from ._common import build_runner
+
+
+def _sync_caller(run_sync: Callable[..., Any]) -> Callable[..., Any]:
+    """Adapt a CrewAI tool's synchronous `_run(**kwargs)` into the async `call`
+    that build_runner expects, off the event loop."""
+
+    async def _call(**kwargs: Any) -> Any:
+        return await asyncio.to_thread(lambda: run_sync(**kwargs))
+
+    return _call
+
+
+def wrap_tool(
+    tool: Any,
+    *,
+    semantics: ActionSemantics,
+    compensate: Optional[CompensationFactory] = None,
+    timeout: Optional[float] = None,
+    name: Optional[str] = None,
+) -> Any:
+    """Return a saga-aware drop-in for a CrewAI `BaseTool`.
+
+    The returned object is a CrewAI `BaseTool` subclass instance whose `_run`
+    routes through the active saga. Its `name`, `description`, and `args_schema`
+    are copied from the original so the crew sees no difference.
+    """
+    from crewai.tools import BaseTool  # lazy
+
+    if not isinstance(tool, BaseTool):
+        raise TypeError(
+            f"expected a crewai.tools.BaseTool, got {type(tool).__name__}. Define "
+            f"the tool as a BaseTool (or via crewai's @tool) and wrap that."
+        )
+
+    tool_name = name or tool.name
+    runner = build_runner(
+        _sync_caller(tool._run), name=tool_name, semantics=semantics,
+        compensate=compensate, timeout=timeout,
+    )
+
+    def _run(self, **kwargs: Any) -> Any:
+        # CrewAI calls _run synchronously. Usually there is no running loop, so
+        # asyncio.run is fine. But if _run is invoked from a thread that already
+        # drives a loop, asyncio.run would raise -- so fall back to a dedicated
+        # worker thread with its own loop, carrying the current context (and thus
+        # the active saga) across with copy_context().
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            return asyncio.run(runner(**kwargs))
+
+        import concurrent.futures
+        import contextvars
+
+        ctx = contextvars.copy_context()
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
+            return ex.submit(lambda: ctx.run(asyncio.run, runner(**kwargs))).result()
+
+    wrapped_cls = type(
+        f"Saga{type(tool).__name__}",
+        (BaseTool,),
+        {"_run": _run},
+    )
+    return wrapped_cls(
+        name=tool_name,
+        description=tool.description,
+        args_schema=tool.args_schema,
+    )
+
+
+async def saga_kickoff(
+    crew: Any,
+    inputs: Optional[dict] = None,
+    *,
+    gate: Any = None,
+    wal: Any = None,
+    halt_on_compensation_failure: bool = True,
+    reraise: bool = True,
+    kickoff: Optional[Callable[[Any], Any]] = None,
+) -> Any:
+    """Run a CrewAI crew inside a saga boundary.
+
+    On success returns the crew's output. On any exception, every wrapped tool
+    that already executed is compensated LIFO, then `SagaAborted` is raised
+    (carrying the `RollbackReport`) unless `reraise=False`.
+
+    CrewAI's `kickoff` is synchronous, so it is driven on a worker thread to keep
+    the saga's event loop responsive. `crew.kickoff_async` is used automatically
+    when present; override the whole call with `kickoff=<coroutine (ctx)->result>`.
+    """
+    try:
+        async with saga_scope(
+            gate=gate, wal=wal,
+            halt_on_compensation_failure=halt_on_compensation_failure,
+        ) as ctx:
+            if kickoff is not None:
+                return await kickoff(ctx)
+            if hasattr(crew, "kickoff_async"):
+                return await crew.kickoff_async(inputs=inputs)
+            return await asyncio.to_thread(lambda: crew.kickoff(inputs=inputs))
+    except SagaAborted as aborted:
+        if reraise:
+            raise
+        return aborted.report
+
+
+__all__ = ["wrap_tool", "saga_kickoff"]

agent_saga/adapters/langgraph.py

@@ -0,0 +1,172 @@
+"""LangGraph / LangChain adapter.
+
+The promise: drop this into an existing LangGraph agent and its tool calls
+become transactional, without rewriting the graph.
+
+Two entry points:
+
+  * `wrap_tool(tool, semantics=..., compensate=...)` -- takes a LangChain tool
+    and returns a drop-in replacement whose execution is recorded on the active
+    saga, with a runtime-derived compensation. Same name, description, and args
+    schema, so the model and the graph see no difference.
+
+  * `saga_run(graph, input)` -- runs a compiled graph inside a saga boundary.
+    If the graph raises, every tool that already ran is compensated LIFO.
+
+The routing logic (`build_runner`) is deliberately separable from the LangChain
+packaging so it can be tested without LangChain installed -- the framework is
+imported lazily, only inside `wrap_tool`.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Awaitable, Callable, Optional
+
+from ..decorator import current_saga, saga_scope
+from ..context import SagaAborted
+from ..semantics import ActionSemantics, CompensationFactory
+
+# A LangChain tool's `ainvoke` takes a single input dict and returns the result.
+ToolAInvoke = Callable[[dict], Awaitable[Any]]
+
+
+def build_runner(
+    ainvoke: ToolAInvoke,
+    *,
+    name: str,
+    semantics: ActionSemantics,
+    compensate: Optional[CompensationFactory] = None,
+    timeout: Optional[float] = None,
+) -> Callable[..., Awaitable[Any]]:
+    """Wrap a tool's `ainvoke` so it records on the active saga.
+
+    The returned coroutine takes keyword arguments (LangChain calls a structured
+    tool's function with the schema fields as kwargs) and:
+
+      * outside a saga, calls the tool untouched -- so the same wrapped tool is
+        usable in a plain script or a unit test with no ceremony;
+      * inside a saga, routes through `SagaContext.execute`, passing the tool's
+        arguments as `policy_args` so pre-flight rules can actually see them
+        (an argument hidden in a closure is invisible to the gate -- the exact
+        bug that let a connector bypass a threshold rule).
+    """
+
+    async def _run(**kwargs: Any) -> Any:
+        # The tool's forward call. An async def (not a lambda returning a
+        # coroutine) so SagaContext._invoke recognizes it as awaitable rather
+        # than shipping it to a worker thread.
+        async def _forward() -> Any:
+            return await ainvoke(dict(kwargs))
+
+        ctx = current_saga()
+        if ctx is None:
+            return await _forward()
+
+        return await ctx.execute(
+            tool=name,
+            semantics=semantics,
+            forward=_forward,
+            compensate=compensate,
+            policy_args=dict(kwargs),
+            timeout=timeout,
+        )
+
+    _run.__name__ = f"saga_{name}".replace(".", "_")
+    return _run
+
+
+def wrap_tool(
+    tool: Any,
+    *,
+    semantics: ActionSemantics,
+    compensate: Optional[CompensationFactory] = None,
+    timeout: Optional[float] = None,
+    name: Optional[str] = None,
+) -> Any:
+    """Return a saga-aware drop-in for a LangChain tool.
+
+    `tool` may be a `BaseTool` or a plain function (which is promoted to a
+    `StructuredTool` first). The result is a `StructuredTool` with the original
+    name, description, and args schema, so `model.bind_tools([...])` and
+    `ToolNode` treat it exactly like the original.
+
+    `compensate` is the runtime factory `(tool_result) -> Compensation | None`.
+    For an `IRREVERSIBLE` tool it may be omitted -- the gate stops it before it
+    runs, so no inverse is needed.
+    """
+    from langchain_core.tools import BaseTool, StructuredTool
+
+    if not isinstance(tool, BaseTool):
+        if not (getattr(tool, "__doc__", None) or "").strip():
+            raise ValueError(
+                f"{getattr(tool, '__name__', tool)!r} needs a docstring (or pass a "
+                f"BaseTool): LangChain uses it as the tool description the model "
+                f"sees. This is LangChain's requirement, surfaced early."
+            )
+        tool = StructuredTool.from_function(tool)
+
+    tool_name = name or tool.name
+    runner = build_runner(
+        tool.ainvoke, name=tool_name, semantics=semantics,
+        compensate=compensate, timeout=timeout,
+    )
+
+    if tool.args_schema is not None:
+        # Reuse the original schema verbatim; do not let StructuredTool try to
+        # infer one from the runner's **kwargs signature.
+        return StructuredTool.from_function(
+            coroutine=runner,
+            name=tool_name,
+            description=tool.description,
+            args_schema=tool.args_schema,
+            infer_schema=False,
+        )
+    return StructuredTool.from_function(
+        coroutine=runner, name=tool_name, description=tool.description,
+    )
+
+
+async def saga_run(
+    graph: Any,
+    input: Any = None,
+    *,
+    config: Any = None,
+    gate: Any = None,
+    wal: Any = None,
+    halt_on_compensation_failure: bool = True,
+    reraise: bool = True,
+    invoke: Optional[Callable[[Any], Awaitable[Any]]] = None,
+) -> Any:
+    """Run a compiled LangGraph graph inside a saga boundary.
+
+    On success returns the graph's output. On any exception, every wrapped tool
+    that already executed is compensated LIFO; then, if `reraise` (default),
+    `SagaAborted` is raised carrying the `RollbackReport`, otherwise the report
+    is returned.
+
+    `invoke` overrides how the graph is driven -- pass a coroutine
+    `(saga_context) -> result` for `astream`, custom configs, or to interleave
+    your own steps. When omitted, `graph.ainvoke(input, config)` is used.
+
+    Note on propagation: the active saga is stored in a contextvar. LangGraph
+    runs tools within the same context (async tasks copy it; sync tools run via
+    a context-preserving executor), so wrapped tools see the saga without any
+    threading of arguments through graph state.
+    """
+    try:
+        async with saga_scope(
+            gate=gate, wal=wal,
+            halt_on_compensation_failure=halt_on_compensation_failure,
+        ) as ctx:
+            if invoke is not None:
+                return await invoke(ctx)
+            if config is not None:
+                return await graph.ainvoke(input, config)
+            return await graph.ainvoke(input)
+    except SagaAborted as aborted:
+        if reraise:
+            raise
+        return aborted.report
+
+
+__all__ = ["wrap_tool", "saga_run", "build_runner"]

agent_saga/adapters/llamaindex.py

@@ -0,0 +1,104 @@
+"""LlamaIndex adapter.
+
+Wrap a LlamaIndex `FunctionTool` so its execution records on the active saga,
+and run a LlamaIndex agent inside a saga boundary. The wrapped tool keeps the
+original name and description, so the agent and the LLM see no difference.
+
+LlamaIndex tools expose the underlying callable as `.fn` / `.async_fn` and their
+metadata as `.metadata.name` / `.metadata.description`. We route the callable
+through the saga and rebuild a `FunctionTool` around it. The framework is
+imported lazily, only inside `wrap_tool`; the routing core is the shared
+`build_runner`, tested without LlamaIndex installed.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable, Optional
+
+from ..context import SagaAborted
+from ..decorator import saga_scope
+from ..semantics import ActionSemantics, CompensationFactory
+from ._common import build_runner
+
+
+def _async_callable(tool: Any) -> Callable[..., Any]:
+    """An async `call(**kwargs)` over the tool's underlying function, preferring
+    the native async implementation and dropping a sync one onto a worker
+    thread so it never blocks the loop."""
+    async_fn = getattr(tool, "async_fn", None)
+    sync_fn = getattr(tool, "fn", None)
+
+    if async_fn is not None:
+        async def _call(**kwargs: Any) -> Any:
+            return await async_fn(**kwargs)
+        return _call
+    if sync_fn is not None:
+        async def _call(**kwargs: Any) -> Any:
+            return await asyncio.to_thread(lambda: sync_fn(**kwargs))
+        return _call
+    raise TypeError("tool exposes neither .async_fn nor .fn; is it a FunctionTool?")
+
+
+def wrap_tool(
+    tool: Any,
+    *,
+    semantics: ActionSemantics,
+    compensate: Optional[CompensationFactory] = None,
+    timeout: Optional[float] = None,
+    name: Optional[str] = None,
+) -> Any:
+    """Return a saga-aware `FunctionTool` drop-in.
+
+    `tool` is a LlamaIndex `FunctionTool`. The result is a new `FunctionTool`
+    with the same name and description whose call routes through the active saga
+    (and passes through untouched outside one).
+    """
+    from llama_index.core.tools import FunctionTool  # lazy
+
+    meta = getattr(tool, "metadata", None)
+    tool_name = name or (getattr(meta, "name", None) or getattr(tool, "name", "tool"))
+    description = getattr(meta, "description", "") or ""
+
+    runner = build_runner(
+        _async_callable(tool), name=tool_name, semantics=semantics,
+        compensate=compensate, timeout=timeout,
+    )
+
+    return FunctionTool.from_defaults(
+        async_fn=runner, name=tool_name, description=description,
+    )
+
+
+async def saga_run(
+    agent: Any,
+    message: Any = None,
+    *,
+    gate: Any = None,
+    wal: Any = None,
+    halt_on_compensation_failure: bool = True,
+    reraise: bool = True,
+    run: Optional[Callable[[Any], Any]] = None,
+) -> Any:
+    """Run a LlamaIndex agent inside a saga boundary.
+
+    By default drives `agent.achat(message)`. Override with `run=<coroutine
+    (saga_context) -> result>` for a workflow, `AgentRunner`, or streaming. On
+    any exception every wrapped tool that executed is compensated LIFO, then
+    `SagaAborted` is raised unless `reraise=False`.
+    """
+    try:
+        async with saga_scope(
+            gate=gate, wal=wal,
+            halt_on_compensation_failure=halt_on_compensation_failure,
+        ) as ctx:
+            if run is not None:
+                return await run(ctx)
+            return await agent.achat(message)
+    except SagaAborted as aborted:
+        if reraise:
+            raise
+        return aborted.report
+
+
+__all__ = ["wrap_tool", "saga_run"]