StateGuard Documentation

The transactional firewall for AI agents — a drop-in checkpointer for LangGraph, CrewAI, and vanilla Python workflows.

Python 3.9+MIT LicenseProduction Ready

What problem does StateGuard solve?

LLM agents are probabilistic. Even the best models occasionally output corrupted state — negative balances, invalid dates, bad JSON. Native checkpointers like MemorySaver blindly save this to your database. StateGuard intercepts the save, validates it against your business rules, and rolls back if anything fails.

How it works

Agent generates output

Your LangGraph / CrewAI / Python agent runs and produces a new state.

↓

StateGuard intercepts

Before the state is persisted, StateGuard runs your invariants against it.

↓

✓

Passes → State saved

If all invariants pass, the state is saved normally. Zero overhead.

↓

✗

Fails → Rollback + Saga

If an invariant fails, StateGuard blocks the save, rolls back memory, and fires your compensation functions (Sagas) to undo any external side-effects.

Installation

StateGuard is available on PyPI and is compatible with Python 3.9+.

terminal

pip install stateguard-core

To verify the installation:

terminal

python -c "import stateguard; print(stateguard.__version__)"

Optional dependencies

StateGuard works out of the box. If you use LangGraph or CrewAI, make sure those are installed too:

terminal

# For LangGraph users
pip install stateguard-core langgraph

# For CrewAI users
pip install stateguard-core crewai

# Install everything at once
pip install stateguard-core langgraph crewai

Requirements

Dependency	Version	Required?
Python	≥ 3.9	Required
LangGraph	≥ 0.1.0	Optional
CrewAI	≥ 0.28.0	Optional

Core Concepts

StateGuard introduces three foundational concepts borrowed from database engineering and applied to AI agent memory:

🔒

Invariants

Business rules your agent's state must always satisfy. An invariant is a plain Python function that receives the state and raises an exception if the rule is violated.

def no_negative_balance(state):
    if state["balance"] < 0:
        raise ValueError("Balance cannot go negative")

🔄

Transactions

All state changes within a checkpoint are atomic. If any invariant fails, the entire state is rolled back to the last valid checkpoint — just like a database ROLLBACK.

# Either ALL of this succeeds, or NONE of it persists
with checkpointer.transaction() as state:
    state["balance"] -= 500
    state["items"].append(order)

⚡

Sagas

Memory rollback is automatic. But what about external API calls your agent already made? Sagas are compensation functions that undo real-world side effects when a transaction fails.

def refund_user(state, exc):
    stripe_api.refund(state["charge_id"])

# Registered as a compensation
compensations={"stripe_charge": refund_user}

🦜LangGraph Integration

StateGuard is designed as a zero-refactoring drop-in replacement for LangGraph's native checkpointers (MemorySaver, SqliteSaver, etc.). You only change 2 lines.

Step 1 — Install

terminal

pip install stateguard-core langgraph

Step 2 — Define your invariants and replace your checkpointer

agent.py

from stateguard.langgraph import StateGuardCheckpointer
from stateguard.invariants import guard
from langgraph.graph import StateGraph
from typing import TypedDict

class AgentState(TypedDict):
    balance: float
    last_action: str

# 1. Define invariant rules for the state
@guard.invariant(name="positive_balance", mode="block")
def no_negative_balance(state):
    return state.get("balance", 0) >= 0

workflow = StateGraph(AgentState)
workflow.add_node("agent", call_agent)
workflow.add_node("tools", tool_executor)
workflow.set_entry_point("agent")

# 2. Drop in the StateGuard checkpointer
# It intercepts LangGraph's state saves and runs the invariants
checkpointer = StateGuardCheckpointer(backend="sqlite", db_path="memory.db")

# Compile exactly as before
app = workflow.compile(checkpointer=checkpointer)

Step 3 — Run your graph

agent.py

config = {"configurable": {"thread_id": "user-session-42"}}

# StateGuard silently protects every checkpoint
result = app.invoke(
    {"balance": 100.0, "last_action": ""},
    config=config
)

# If an agent hallucinates a negative balance:
# → StateGuard BLOCKS the save (InvariantViolation)
# → State rolls back to last valid checkpoint
# → Your database is never corrupted

With Saga compensation (for external APIs)

agent.py

from stateguard.saga import saga_step
import stripe

def refund_stripe_charge():
    print("Auto-refunding charge...")

# Wrap your LangGraph node to register compensating actions
@saga_step(compensate=refund_stripe_charge)
def tool_executor(state: AgentState):
    # Charge the user
    stripe.Charge.create(...)
    state["balance"] -= 500
    return state

# If an invariant fails later in the workflow, 
# refund_stripe_charge is automatically executed.

✓ Zero refactoring required.

StateGuard implements the same interface as LangGraph's BaseCheckpointSaver. Your graph logic, node functions, and state schema remain completely untouched.

🚣CrewAI Integration

StateGuard wraps CrewAI task outputs before they are passed to the next agent in a crew. Use GuardedCrewState to enforce invariants on the crew's shared memory.

Step 1 — Install

terminal

pip install stateguard-core crewai

Step 2 — Define your invariants

crew.py

from crewai import Agent, Task, Crew
from stateguard.invariants import guard
from stateguard.saga import Saga, InvariantViolation

# --- Define your invariants ---
@guard.invariant(name="order_total_is_positive", mode="block")
def check_order_total(state):
    return state.get("order_total", 0) > 0

@guard.invariant(name="customer_id_exists", mode="block")
def check_customer_id(state):
    return bool(state.get("customer_id"))

Step 3 — Run your crew within a Saga

crew.py

# --- Define your crew agents ---
researcher = Agent(
    role="Order Researcher",
    goal="Find and validate the order details",
    backstory="Expert at validating e-commerce order data.",
)

processor = Agent(
    role="Order Processor",
    goal="Process and confirm the order",
    backstory="Handles secure payment processing.",
)

# --- Define tasks ---
research_task = Task(description="Research order details.", agent=researcher)
process_task = Task(description="Process the payment.", agent=processor)

crew = Crew(agents=[researcher, processor], tasks=[research_task, process_task])

# --- Execute with StateGuard ---
initial_state = {"order_total": 0, "customer_id": None}

async def execute_crew():
    try:
        # Wrap the entire crew kickoff in a Saga transaction
        async with Saga(initial_state, saga_id="crew_run_1") as coord:
            
            # The crew mutates the state dictionary
            result = crew.kickoff(inputs={"state": coord.state})
            
            # If the crew hallucinates an invalid state, it is caught here 
            # before persisting to the database.
            
    except InvariantViolation as e:
        print(f"Crew hallucinated an invalid state: {e}")
        print("Safely blocked the transaction.")

import asyncio
asyncio.run(execute_crew())

🐍Vanilla Python

No LangGraph or CrewAI? No problem. Use Saga to add transactional protection to any plain Python dictionary in your agent loop.

Step 1 — Install

terminal

pip install stateguard-core

Step 2 — Define invariant rules

agent.py

from stateguard.invariants import guard

# --- Define invariants ---
@guard.invariant(name="no_negative_balance", mode="block")
def check_balance(state):
    return state.get("balance", 0) >= 0

@guard.invariant(name="valid_status", mode="block")
def check_status(state):
    allowed = {"idle", "processing", "complete", "failed"}
    return state.get("status") in allowed

Step 3 — Wrap your state dict in a Saga

agent.py

import openai
from stateguard.saga import Saga, InvariantViolation

client = openai.OpenAI()

initial_state = {
    "balance": 1000.0,
    "status": "idle",
    "history": []
}

async def run_agent_loop():
    # Use a Saga transaction — either ALL changes persist or NONE do
    try:
        async with Saga(initial_state, saga_id="tx_123") as coord:
            while coord.state["status"] != "complete":
                # Call your LLM
                response = client.chat.completions.create(
                    model="gpt-4o",
                    messages=[
                        {"role": "system", "content": "You are an order processing agent."},
                        {"role": "user", "content": str(coord.state.to_dict())}
                    ]
                )
                
                action = parse_action(response.choices[0].message.content)
                
                if action["type"] == "charge":
                    coord.state["balance"] -= action["amount"]
                    coord.state["history"].append(action)
                elif action["type"] == "complete":
                    coord.state["status"] = "complete"
                
                # Exiting the Saga context automatically validates invariants
                
    except InvariantViolation as e:
        print(f"Agent action blocked: {e}")
        print("Memory safely rolled back to the original dictionary state.")

import asyncio
asyncio.run(run_agent_loop())

Writing Invariants

Invariants are the core of StateGuard. They are plain Python functions — nothing special is required.

Signature

# An invariant takes the full state dict and raises an Exception if invalid.
# If no exception is raised, the invariant is considered PASSED.
def my_invariant(state: dict) -> None:
    if not some_condition(state):
        raise ValueError("Descriptive error message")

Examples: Common invariant patterns

invariants.py

# 1. Numeric bounds
def no_negative_balance(state):
    if state["balance"] < 0:
        raise ValueError(f"Balance {state['balance']} is negative")

# 2. Type validation
def valid_user_id(state):
    if not isinstance(state.get("user_id"), str):
        raise TypeError("user_id must be a string")

# 3. Cross-field validation
def start_before_end(state):
    if state["start_date"] >= state["end_date"]:
        raise ValueError("Start date must be before end date")

# 4. Required fields
def required_fields(state):
    for field in ["user_id", "order_id", "amount"]:
        if not state.get(field):
            raise ValueError(f"Required field missing: {field}")

# 5. Enum validation
def valid_status(state):
    allowed = {"pending", "processing", "complete", "refunded"}
    if state.get("status") not in allowed:
        raise ValueError(f"Invalid status: {state['status']}")

# 6. Using Pydantic for schema validation
from pydantic import BaseModel, validator

class OrderState(BaseModel):
    balance: float
    items: list
    status: str

    @validator("balance")
    def balance_must_be_positive(cls, v):
        assert v >= 0, "Balance cannot be negative"
        return v

def pydantic_invariant(state):
    OrderState(**state)  # Raises ValidationError if invalid

Registering multiple invariants

You can register as many invariants as needed. They run in order — the first failure stops execution and triggers a rollback.

checkpointer = StateGuardCheckpointer(
    invariants=[
        no_negative_balance,
        valid_status,
        required_fields,
        start_before_end,
    ]
)

The Saga Pattern

State rollback reverts your in-memory checkpoint. But what about the real-world side effects your agent already triggered — a Stripe charge, a SendGrid email, an AWS S3 upload?

The Saga Pattern solves this by pairing each action with a compensation function that undoes it.

⚠️ When do you need Sagas?

Use Sagas whenever your agent calls external APIs or services that have real-world consequences. If your agent only reads data or updates an in-memory state, Sagas are not needed.

Defining a Saga compensation

sagas.py

import stripe
import boto3

# Compensation functions take (state, exception) as arguments
def refund_stripe_charge(state, exception):
    """Called if an invariant fails after a Stripe charge was made."""
    charge_id = state.get("last_charge_id")
    if charge_id:
        refund = stripe.Refund.create(charge=charge_id)
        print(f"Auto-refunded: {refund.id}")

def delete_s3_upload(state, exception):
    """Called if an invariant fails after a file was uploaded."""
    key = state.get("uploaded_s3_key")
    if key:
        s3 = boto3.client("s3")
        s3.delete_object(Bucket="my-bucket", Key=key)
        print(f"Auto-deleted S3 object: {key}")

def cancel_calendar_event(state, exception):
    """Called if an invariant fails after a calendar event was created."""
    event_id = state.get("calendar_event_id")
    if event_id:
        calendar_api.events().delete(calendarId="primary", eventId=event_id).execute()

Registering compensations

agent.py

checkpointer = StateGuardCheckpointer(
    invariants=[
        no_negative_balance,
        valid_order_status
    ],
    compensations={
        # Map action name → compensation function
        "stripe_charge": refund_stripe_charge,
        "s3_upload": delete_s3_upload,
        "calendar_event": cancel_calendar_event,
    }
)

Full end-to-end Saga example

payment_agent.py

from stateguard import StateGuardCheckpointer
from langgraph.graph import StateGraph
from typing import TypedDict
import stripe

class PaymentState(TypedDict):
    user_id: str
    amount: float
    balance: float
    last_charge_id: str
    status: str

# --- Invariants ---
def no_negative_balance(state):
    if state["balance"] < 0:
        raise ValueError(f"Cannot charge more than balance: {state['balance']}")

def valid_amount(state):
    if state["amount"] <= 0:
        raise ValueError("Charge amount must be positive")

# --- Saga compensations ---
def refund_charge(state, exception):
    print(f"Rolling back Stripe charge due to: {exception}")
    if state.get("last_charge_id"):
        stripe.Refund.create(charge=state["last_charge_id"])

# --- Agent node ---
def charge_node(state: PaymentState):
    charge = stripe.Charge.create(
        amount=int(state["amount"] * 100),
        currency="usd",
        customer=state["user_id"]
    )
    return {
        **state,
        "last_charge_id": charge.id,
        "balance": state["balance"] - state["amount"],
        "status": "charged"
    }

# --- Assemble graph ---
workflow = StateGraph(PaymentState)
workflow.add_node("charge", charge_node)
workflow.set_entry_point("charge")

# StateGuard intercepts LangGraph checkpoints
checkpointer = StateGuardCheckpointer(backend="sqlite", db_path="memory.db")

app = workflow.compile(checkpointer=checkpointer)

# If agent hallucinates amount=99999:
# 1. StateGuard detects valid_amount invariant failure
# 2. InvariantViolation is raised and memory is rolled back
# 3. refund_charge() compensation fires automatically
# 4. Your user is protected

API Reference

`StateGuardCheckpointer`

The main class for LangGraph integration. Drop-in replacement for MemorySaver.

Parameter	Type	Default	Description
`backend`	`str`	`"memory"`	Storage backend. Options: `"memory"`, `"sqlite"`.
`registry`	`Registry`	`None`	Optional explicit invariant registry. Defaults to global registry.
`db_path`	`str`	`None`	Path to SQLite database if using sqlite backend.

`Saga`

Transactional state context manager for Vanilla Python.

Parameter	Description
`initial_state`	The Python dictionary to wrap in the transaction.
`saga_id`	Unique string identifier for tracking the transaction.
`.state`	Property on the yielded context coordinator to read/write state.

`@guard.invariant`

Decorator to define invariants.

Parameter	Description
`name`	Unique name for the invariant rule.
`mode`	`"block"` (raises InvariantViolation) or `"warn"` (logs warning).

`@saga_step`

Decorator to define compensating actions for graph nodes or functions.

Parameter	Description
`compensate`	Function to execute if a subsequent invariant fails.

FAQ

Does StateGuard add latency to my LLM calls?

No. StateGuard runs entirely in-process on your Python dictionary before persisting to disk or memory. Invariant evaluation is typically sub-millisecond. There is no network hop, no external service, and no added latency to your LLM API calls.

Do I need to rewrite my LangGraph graph?

Absolutely not. StateGuard implements the same BaseCheckpointSaver interface as LangGraph's native checkpointers. You change 2 lines: your import and the checkpointer constructor. Everything else stays identical.

What state formats are supported?

StateGuard supports plain Python dicts, TypedDicts, Pydantic v1 and v2 models, and dataclasses. If LangGraph can use it as a state schema, StateGuard can protect it.

What happens if a Saga compensation itself fails?

StateGuard will log the compensation failure and raise a CompensationError. We recommend wrapping compensation functions in try/except and implementing idempotent operations (e.g., using Stripe's idempotency keys) to prevent double-refunds.

Is there persistent storage (not just in-memory)?

Yes. Pass storage='sqlite' and a db_path to StateGuardCheckpointer to persist checkpoints across restarts. Cloud-based storage (Redis, PostgreSQL) is on the roadmap for StateGuard Cloud.

Is this production-ready?

StateGuard itself is stable and MIT-licensed. We use it internally in production agent pipelines. That said, always test your invariants thoroughly before deploying, as poorly written invariants can incorrectly block valid state.

Ready to secure your agents?

Stop trusting LLMs blindly with your production state. Add StateGuard in 2 lines.

Back to Home ⭐ Star on GitHub