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 — Replace your checkpointer

agent.py

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

# --- Your existing code stays EXACTLY the same ---
class AgentState(TypedDict):
    balance: float
    last_action: str
    charge_id: str

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

# --- ONLY THIS CHANGES ---
# Before: checkpointer = MemorySaver()
# After:
def no_negative_balance(state):
    if state["balance"] < 0:
        raise ValueError(f"Invariant failed: balance is {state['balance']}")

checkpointer = StateGuardCheckpointer(
    invariants=[no_negative_balance]
)

# --- 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": "", "charge_id": ""},
    config=config
)

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

With Saga compensation (for external APIs)

agent.py

import stripe

def refund_stripe_charge(state, exception):
    """Called automatically if an invariant fails after a charge."""
    if state.get("charge_id"):
        stripe.Refund.create(charge=state["charge_id"])
        print(f"Auto-refunded charge: {state['charge_id']}")

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

checkpointer = StateGuardCheckpointer(
    invariants=[no_negative_balance],
    compensations={
        # Key = the name of the action that caused the side-effect
        "stripe_charge": refund_stripe_charge
    }
)

app = workflow.compile(checkpointer=checkpointer)

✓ 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 — Wrap your crew's shared state

crew.py

from crewai import Agent, Task, Crew
from stateguard import GuardedCrewState

# --- Define your invariants ---
def order_total_is_positive(state):
    if state.get("order_total", 0) <= 0:
        raise ValueError("Order total must be positive before confirming.")

def customer_id_exists(state):
    if not state.get("customer_id"):
        raise ValueError("Customer ID cannot be empty.")

# --- Create a guarded shared state dict ---
shared_state = GuardedCrewState(
    initial_state={
        "order_total": 0,
        "customer_id": None,
        "items": []
    },
    invariants=[order_total_is_positive, customer_id_exists]
)

# --- 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 with the guarded state ---
research_task = Task(
    description="Research and populate order details in the shared state.",
    agent=researcher,
    # StateGuard validates state after each task completes
    callback=shared_state.validate_checkpoint
)

process_task = Task(
    description="Process the payment using the validated order state.",
    agent=processor,
    callback=shared_state.validate_checkpoint
)

# --- Run the crew ---
crew = Crew(
    agents=[researcher, processor],
    tasks=[research_task, process_task],
    verbose=True
)

result = crew.kickoff()

Handling validation failures

crew.py

from stateguard import GuardedCrewState, InvariantError

shared_state = GuardedCrewState(
    initial_state={"order_total": 0},
    invariants=[order_total_is_positive]
)

try:
    # This will be caught before it propagates to the next agent
    shared_state.update({"order_total": -50})
except InvariantError as e:
    print(f"Crew state violation caught: {e}")
    # State is automatically rolled back to last valid checkpoint
    print(f"Safe state restored: {shared_state.current}")

🐍Vanilla Python

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

Step 1 — Install

terminal

pip install stateguard-core

Step 2 — Wrap your state dict

agent.py

from stateguard import GuardedState

# --- Define invariants ---
def no_negative_balance(state):
    if state["balance"] < 0:
        raise ValueError(f"Balance cannot be negative: {state['balance']}")

def valid_status(state):
    allowed = {"idle", "processing", "complete", "failed"}
    if state["status"] not in allowed:
        raise ValueError(f"Invalid status: {state['status']}")

# --- Create guarded state ---
memory = GuardedState(
    initial_state={
        "balance": 1000.0,
        "status": "idle",
        "history": []
    },
    invariants=[no_negative_balance, valid_status]
)

Step 3 — Use inside your agent loop

agent.py

import openai

client = openai.OpenAI()

def run_agent_loop(memory: GuardedState):
    while memory.current["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(memory.current)}
            ]
        )
        
        action = parse_action(response.choices[0].message.content)
        
        # Use a transaction — either ALL changes persist or NONE do
        try:
            with memory.transaction() as state:
                if action["type"] == "charge":
                    state["balance"] -= action["amount"]
                    state["history"].append(action)
                elif action["type"] == "complete":
                    state["status"] = "complete"
                    
                # StateGuard auto-validates here before committing
                
        except ValueError as e:
            print(f"Agent action blocked: {e}")
            print(f"Memory safely rolled back to: {memory.current}")
            # Continue the loop — state is back to the last safe checkpoint
            continue

run_agent_loop(memory)
print("Final state:", memory.current)

Inspecting checkpoints

agent.py

# View the full checkpoint history
for i, checkpoint in enumerate(memory.history):
    print(f"Checkpoint {i}: {checkpoint}")

# Manually roll back to a specific checkpoint
memory.rollback(steps=2)
print("Rolled back state:", memory.current)

# Get the last N checkpoints
last_3 = memory.get_checkpoints(n=3)

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")

checkpointer = StateGuardCheckpointer(
    invariants=[no_negative_balance, valid_amount],
    compensations={"stripe_charge": refund_charge}
)

app = workflow.compile(checkpointer=checkpointer)

# If agent hallucinates amount=99999:
# 1. StateGuard detects balance would go negative
# 2. Memory is rolled back
# 3. refund_charge() 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
`invariants`	`list[callable]`	`[]`	List of invariant functions. Each receives the full state dict.
`compensations`	`dict[str, callable]`		Map of action name → compensation function for Saga rollbacks.
`storage`	`str`	`"memory"`	Storage backend. Options: `"memory"`, `"sqlite"`.

`GuardedState`

Vanilla Python transactional state container.

Method	Description
`.transaction()`	Context manager. Commits on success, rolls back on failure.
`.update(dict)`	Update state, running invariants before committing.
`.rollback(steps=1)`	Manually roll back to a previous checkpoint.
`.current`	Property. Returns the current validated state.
`.history`	Property. Returns list of all past checkpoints.
`.get_checkpoints(n)`	Returns the last N checkpoints.

`GuardedCrewState`

Shared state container for CrewAI crews.

Method	Description
`.validate_checkpoint(output)`	Task callback. Validates state after each CrewAI task.
`.update(dict)`	Update shared state with invariant validation.
`.current`	Property. Returns the current validated shared state.

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