Aura-State

Prove your AI agent can't be prompt-injected — before you ship it.

A type-checker for LLM agents. Static analysis + formal proofs — Z3, CTL model checking, taint dataflow, conformal risk — over your agent's design. Catches injection paths, unsafe tool calls, hallucinated outputs, and leaked secrets. Runs locally, in CI, no API key.

_{Verify a design → import a real agent from its code → the lethal trifecta lights up with the exact path → save a baseline and gate regressions.  ·  ▶ watch the full walkthrough}

pip install aura-state          # or:  uv pip install aura-state  ·  uv add aura-state

Catch a prompt-injection in one command

Point it at your agent design. It fails the build if an untrusted input can reach a real tool call unsanitized:

$ aura-state check my_agent.json

  aura-state check · sql-agent (naive) · 3 nodes

  ✗ taint [Execute]: untrusted data from 'Ask' can reach sink 'Execute' with no sanitizer — injection path
  ✗ obligation [GenSQL]: 'read_only' not proven — the generated SQL may not be read-only

  ✗ NOT PROVEN — 2 blocking findings          # exit code 1 → CI fails

Zero install — run it straight from PyPI with uv:

uvx aura-state check agents/*.json          # no venv, no install — perfect for CI

Drop it into CI and every PR is checked:

# .github/workflows/aura.yml
- uses: your-org/aura-state@v0
  with: { paths: "agents/*.json" }

Catch the lethal trifecta — before an attacker does

The lethal trifecta is the scariest failure mode for a tool-using agent: the moment it can (1) read private data, (2) ingest untrusted content, and (3) communicate externally — all on one reachable path — a prompt injection hidden in that content can read your data and ship it out. Everyone warns about it. Aura proves your agent can't close it, statically, before you ship:

$ aura-state check dev_assistant.tools.json   # ← your MCP tools/list export

  ✗ trifecta [slack_post_message]: lethal trifecta closed — prompt injection at 'fetch'
    (untrusted) can reach external sink 'slack_post_message' unsanitized while 'read_file'
    brings private data into scope. Path: fetch → Agent → slack_post_message.
    Break it with a sanitizer between 'fetch' and 'slack_post_message'.

Point it straight at your MCP setup. Export your agent's tools/list (or hand it your client config) and Aura models the worst case — the LLM can call any tool in any order — then decides whether fetch + filesystem + slack together form an exfiltration channel. It never connects to or runs a server.

Or point it at your agent's code. Aura statically imports the tool surface from a LangGraph / CrewAI / LangChain source file or repo — no install, and it never runs your code:

aura-state check your_agent.py          # extracts @tool functions & framework tools
aura-state check ./my-agent-repo/       # a whole directory

✗ trifecta [post_to_slack]: prompt injection at 'fetch_url' (untrusted) can reach
  'post_to_slack' unsanitized while 'read_customer_file' brings private data into scope.

We audited real MCP-server compositions and agent frameworks (GitHub, filesystem, fetch, Slack, Postgres · CrewAI, AutoGPT) modeled from their published tool surfaces. 5 of 6 can close the lethal trifecta — including the exact shape of the Invariant Labs GitHub-MCP exploit. Full writeup + sources: docs/AUDIT.md. Reproduce it: python examples/audit/run.py.

(Also python examples/audit.py: 9 hand-built agent patterns, 5 with an unguarded injection path, 4 closing the trifecta — each fixable with one sanitizer via Aura's auto-repair.)

Gate regressions, not existing debt

Absolute checks are noisy on a legacy agent. Give check a baseline and it fails only on paths this change newly opened — the line that gets Aura into every agent repo's CI:

$ aura-state check agents/*.json --baseline main-baseline.json

  ✗ taint [send_email] [NEW]: untrusted data from 'read_web' can reach sink 'send_email'…
  ✗ REGRESSION — 1 new blocking finding this change      # exit 1 → PR fails

aura-state check agents/*.json --json > main-baseline.json   # on main
aura-state check agents/*.json --baseline main-baseline.json # on the PR

See it in 10 seconds (no API key)

Paste this after installing — Z3 rejects a hallucinated value and accepts a correct one, in the loop:

from aura_state import prove_extraction

print(prove_extraction({"area": 100, "rate": 3, "total": 999}, ["total == area * rate"]).verified)  # False
print(prove_extraction({"area": 100, "rate": 3, "total": 300}, ["total == area * rate"]).verified)  # True

For the full runnable proofs against the real solvers — each ~10s, no API key — clone the repo:

git clone https://github.com/munshi007/Aura-State && cd Aura-State
pip install -e .

Aura Studio — the whole platform, running locally (no cloud, no key)

Prefer a console to code? Launch a local web app that runs the real verifiers on your own machine. One command, nothing leaves your laptop:

pip install "aura-state[ui]" && aura-state ui     # opens http://127.0.0.1:8155
# or zero-install:  uvx --from "aura-state[ui]" aura-state ui

A full agent IDE — design an agent, prove it, run it, watch it execute, and hand it off to code — all in one local console. Not a sandbox you type into: every check is the actual framework verifier. Zero setup: it opens on a built-in demo provider that simulates the LLM steps, so you can hit Run and watch the agent execute end-to-end with no key and no Ollama. Paste an OpenAI / Gemini / DeepSeek key (or point at a local Ollama) — right by the Run button — for a real run.

An agent is a typed graph of four node kinds — you name them anything:

• Extract — an LLM step producing structured output (Z3 obligations + conformal calibration prove it)
• Decision — a verified branching rule (compiled to a whitelisted evaluator, never eval)
• Tool — a declared external call (db.write, http.get, payment.refund, …). Aura proves its preconditions; it does not execute it — you bind the real tool in your code / aura-runtime
• Sanitizer — a taint boundary that makes a downstream tool safe

_{A real OSS agent — the LangGraph SQL agent — modeled in Aura and proven read-only + injection-safe.}

Bring any model. A local Ollama model needs no key. For OpenAI / Gemini / DeepSeek, paste your key in Settings → Providers — it's held in memory for the session (never written to disk) and works immediately, no restart. Keys set in your environment before launch are picked up too. Nothing leaves your laptop.

Connect your agent (any code · CrewAI · LangGraph)

You don't rebuild your agent in the studio — you verify it where it already runs, and stream the results in. A tiny SDK does both:

from aura_state.hooks import Monitor, verified

mon = Monitor()   # points at a local `aura-state ui`

# verify + stream any output:
mon.record({"refund_amount": 180, "order_total": 200},
           ["refund_amount <= order_total"], node="RefundAgent")

# or decorate your extraction / CrewAI task / LangGraph node (fail-closed):
@verified(["total == area * rate"], monitor=mon, strict=True)
def extract_quote(text) -> dict:
    ...        # your LLM call
    return {"area": 100, "rate": 3, "total": 300}

Verification runs in your code even if the studio is closed; when it's open, the Monitor shows every output proven or rejected, live.

Build — click a node, configure everything (prompt, model, schema, obligations, transitions)

Run — the real engine executes the whole agent; each step extracted, Z3-verified, and a contract emitted

Monitor — your real agent (via the SDK hook) streams outputs in, proven or rejected live

What this is

Most LLM frameworks let you chain API calls and hope for the best. Aura-State takes a different approach: you define your workflow as a typed graph of nodes, and verification runs inside the loop — every extraction must satisfy its formal contract before the workflow moves on.

The key difference is what happens between nodes:

• Extractions are checked against Z3 proof obligations, in the extract→verify→retry loop — a value that can't be proven is not accepted (fail-closed)
• Math runs in a no-exec sandboxed interpreter, never hallucinated
• Uncertainty uses split-conformal / jackknife+ calibration (with the honest worst-case coverage); repeated same-input runs are reported as dispersion, not a coverage guarantee
• Workflows are model-checked (CTL) for reachability/completion/ordering before they run
• Routing (when a node returns an ambiguous edge) is a Thompson-sampling bandit, not an LLM guess

Scope & honest limits

Aura-State is a design-time verifier, and worth being precise about:

• It reasons over your agent's structure (a typed graph / tool surface), not its runtime behaviour — it does not execute your agent or tools.
• The lethal-trifecta / taint checks are sound but heuristic: tool roles are inferred from names + declared side-effects, so the analysis over-flags rather than misses (fail-closed), and you can correct any call with a data_class / exfil override. It is not a semantic proof that a specific injection string succeeds.
• The Z3 / CTL guarantees are exact for the obligations and properties you write; they don't invent the spec for you.
• It's early (v0.x) — verified on real OSS agent shapes, but treat it as a strong linter you can trust to fail closed, not a certified security product.

Quick example

from aura_state import AuraEngine, Node, CompiledTransition
from pydantic import BaseModel, Field
from openai import OpenAI

# Define what you want to extract
class LeadData(BaseModel):
    name: str = Field(description="Full name")
    budget: int = Field(description="Budget in USD")
    timeline: str = Field(description="Buying timeline")

# Define a node that extracts it — with a Z3 obligation the value must satisfy
class ExtractLead(Node):
    system_prompt = "Extract lead info from a sales call transcript."
    extracts = LeadData
    obligations = ["budget > 0"]   # proven in the loop; unprovable -> not accepted

    def handle(self, user_text, extracted_data=None, memory=None):
        return "QualifyBudget", extracted_data.model_dump()

# Define a decision node that does deterministic math (no LLM). Its rule runs
# even though the node does no extraction — it reads prior state from memory.
class QualifyBudget(Node):
    system_prompt = "Score the lead."
    sandbox_rule = "result = budget > 100000"  # runs in the no-exec sandbox

    def handle(self, user_text, extracted_data=None, memory=None):
        return "END", memory

# Wire it up
engine = AuraEngine(llm_client=OpenAI())
engine.register(ExtractLead, QualifyBudget)
engine.connect([
    CompiledTransition(from_node=ExtractLead, to_node=QualifyBudget),
])

# Run
next_state, data = engine.process("ExtractLead", user_text="Hi, I'm Sarah. Budget is $450k.")

What happens under the hood

When you call engine.process(), it runs through these steps in order:

1. Few-shot injection      →  optional: inject similar past successes as examples.
2. Verification loop       →  extract → check (sandbox rule + Z3 obligations) → retry.
                              A value that fails its contract is not accepted (fail-closed).
3. Consensus dispersion    →  with consensus > 1, measure agreement across the runs
                              (dispersion, not a calibrated coverage guarantee).
4. Your handle() method    →  your routing / business logic runs here.
5. Bandit router           →  if handle() returns an invalid edge, Thompson-sample a feasible one.
6. State serialization     →  save state (JSON, tamper-evident) for time-travel debugging.

Graph-level properties (reachability, completion, ordering) are checked separately at design time with engine.verify([...]) — CTL model checking over the whole graph, which per-transition checks can't do. Everything above happens in the loop; engine.verification_reports() returns the proof results and intervals per step.

Formal verification (the interesting part)

This is what actually makes Aura-State different from other frameworks.

Verify your workflow graph before it runs

Your node graph gets compiled into a Kripke structure and checked against temporal logic properties:

from aura_state import verify_engine, reachability, mutual_exclusion, eventual_completion

results = verify_engine(engine, [
    {"description": "QualifyBudget is reachable", "formula": reachability("QualifyBudget")},
    {"description": "All paths terminate", "formula": eventual_completion("QualifyBudget")},
])
# Result: PROVEN or VIOLATED, with the exact states that satisfy/violate

This is the same technique used to verify hardware circuits and flight control systems (CTL model checking, Clarke et al. 1986).

Prove that extracted data is correct

After the LLM extracts values, Z3 (a theorem prover from Microsoft Research) can formally prove they satisfy your constraints:

from aura_state import prove_extraction

result = prove_extraction(
    {"budget": 450000, "cost_per_sqft": 3, "total": 1350000},
    obligations=["budget > 0", "total == budget * cost_per_sqft"],
)
# result.verified = True
# If False, Z3 gives you a counterexample showing exactly what broke

It also proves your obligations aren't self-contradictory. ["x > 5", "x < 3"] can never hold — Z3 catches that symbolically (variables ranging freely over the declared field bounds, not pinned to one value), and the design→contract compiler flags it per node before you ship.

Confidence intervals on extractions

Run the extraction multiple times and get distribution-free confidence intervals:

from aura_state import conformal_interval

budgets = [450000, 452000, 448000, 450000, 451000]
ci = conformal_interval(budgets, confidence=0.95)
# ci.lower = 447800, ci.upper = 452200

This uses conformal prediction (Vovk et al., 2005) — no distributional assumptions required.

Pipeline-aware (PASC): a 95% guarantee at each node is not 95% end-to-end — errors compound. PipelineConformal calibrates on the composed output so the guarantee holds for the final answer. In the demo, per-step conformal covers the end-to-end result only ~48%; PASC hits the nominal 90%. See python examples/pasc_demo.py.

Compile a runtime contract from the design

The obligations, CTL verdicts, and confidence a workflow was proven against compile into a single portable, versioned contract. Because it's derived from the same typed design the engine runs, the specification is faithful by construction — spec and implementation are one artifact and can't drift.

contract = engine.compile_contract(properties=[
    {"description": "RouteLead is reachable", "formula": reachability("RouteLead")},
])
contract.to_json()                        # portable, content-addressable
check_faithfulness(contract, "QualifyLead", extracted)   # contract agrees with the loop
diff_contracts(old, contract)             # design-time regression gate

Every other assurance system consumes a behavioral contract it can't author — and hand-written policy drifts from the code (and is only 24–35% faithful when translated from prose). Here the contract is emitted from the design that was proven. See python examples/emit_contract_demo.py.

Prove untrusted data can't reach a dangerous tool (injection-proof)

Label nodes with capability types and the compiler statically proves — over the typed graph — that no untrusted source can reach a dangerous sink without passing a sanitizer. It tracks provenance, not content, so it can't be fooled by the encodings that defeat runtime scanners. The verdict compiles into the contract, so a runtime can refuse to deploy a VIOLATED graph.

class Ingest(Node):    untrusted_source = True     # LLM / external tool output
class Review(Node):    sanitizer = True            # clears taint
class SendEmail(Node): dangerous_sink = True       # irreversible action

analyze_taint(engine)   # -> VIOLATED (Ingest -> SendEmail) unless Review is in the path

It's field-level: label individual fields, and a clean field passes a sink untouched while only a tainted field reaching it is a violation — with the exact field and its origin named. A field-specific sanitizer clears just its field.

class Ingest(Node): untrusted_fields = ["note"]   # free text is untrusted
class Send(Node):   sink_fields = ["account_id"]  # the action consumes account_id

analyze_field_taint(engine)   # PROVEN — the tainted `note` never reaches the sink arg

Everyone else sells injection detection (probabilistic). This is impossibility over the design. See python examples/taint_proof_demo.py.

Act only if calibrated risk ≤ ε, otherwise escalate

The "knows when it doesn't know" story made into an actual gate. Calibrate a controller on a labeled set and the agent auto-acts only when its false-action rate is provably within budget — everything below the threshold escalates to a human, never a silent guess.

ctrl = RiskController(epsilon=0.05).calibrate(scores, correct)   # false-action rate ≤ 5%

class Decide(Node):
    risk_controller = ctrl
    escalation_node = "HumanReview"
    def risk_score(self, extracted_data=None, conformal=None, memory=None):
        return confidence   # in [0,1]

Uses Conformal Risk Control (arXiv:2208.02814); Learn-Then-Test (arXiv:2110.01052) for tuning several thresholds. Abstention is a first-class engine outcome. See python examples/risk_abstention_demo.py.

Let the verifier repair the plan (counterexample-guided replanning)

Verification is usually a gate that says VIOLATED and stops. Here the counterexample — a tainted path, a CTL violating state, a Z3 assignment — is fed back to a replanner, which edits the plan and re-verifies, until it's proven or a budget is hit. The plan is provably correct because the verifier drove it there.

result = engine.repair()          # verify → counterexample → repair → re-verify
result.verified                   # True: driven to PROVEN (e.g. a sanitizer inserted)
result.unresolved                 # if it aborts: the explicit remaining violations

Never a silent pass — an unrepairable design aborts with the violation named. Refs: PAT-Agent (arXiv:2509.23675), VERIMAP (arXiv:2510.17109). See python examples/replan_demo.py.

Results — real data, reproducible, no API key

Every number below comes from running Aura-State on real, public data (or a real local model). Nothing here is simulated.

# real data through the verifier — no LLM, no key
python examples/real_data/verify_real_dataset.py
pip install scikit-learn && python examples/real_data/conformal_on_real_data.py

# a real LangGraph agent on a local model — no key, no cloud
ollama pull qwen2.5:0.5b && pip install langgraph
OLLAMA_MODEL=qwen2.5:0.5b python examples/integrations/langgraph_verified.py

# verification rejects a hallucination in the loop — no key
python examples/verified_loop_demo.py

See examples/real_data/ and examples/integrations/ for the full runs.

Project structure

aura_state/
├── core/
│   ├── engine.py              # Main engine — verified process() loop + bandit router
│   ├── adaptive_graph.py      # Node health metrics + per-edge Beta-Bernoulli posteriors
│   ├── verification_loop.py   # Extract → verify (sandbox + Z3) → retry loop
│   └── providers.py           # Multi-model routing + cost tracking
├── verification/             # ← the core: correct, adversarially-tested primitives
│   ├── proof_engine.py        # Z3 proofs (fail-closed AST→Z3 compiler, no eval)
│   ├── conformal.py           # jackknife+ prediction intervals (order statistic)
│   └── temporal_verifier.py   # Kripke + CTL model checking (init-state, structural deadlocks)
├── execution/
│   ├── sandbox.py             # No-exec allowlist AST evaluator (deny-by-default)
│   └── tracer.py              # State serialization, tamper-evident JSON (time-travel debug)
├── compiler/
│   ├── schema_compiler.py     # JSON Schema → Node classes
│   └── dspy_tuner.py          # KNN few-shot selection (real embedder required)
├── memory/
│   └── pruner.py              # Context window optimization
└── consensus/
    └── auto_vote.py           # Multi-run extraction with voting

Installation

pip install aura-state

Or the latest from source:

pip install git+https://github.com/munshi007/Aura-State.git

Python 3.10+ required. Dependencies: pydantic, instructor, openai, networkx, pyModelChecking, z3-solver, pyyaml.

Tests

python -m pytest tests/ -v
# 225 tests passing

Works with any LLM provider

Aura-State's verification is independent of the model — only extraction calls one, and nearly every provider speaks the OpenAI-compatible API. Point the client anywhere; the Z3/CTL/taint/conformal code is identical:

from openai import OpenAI
# Gemini, DeepSeek, Together, or a local model via Ollama/vLLM — just the base_url
engine = AuraEngine(llm_client=OpenAI(api_key=key, base_url="https://api.deepseek.com"))

See examples/cookbook/ for OpenAI / Gemini / DeepSeek / local recipes. Your API key is read from your environment — never hard-coded.

Docs

• Cookbook — realistic agents, verified end to end; any provider
• LangGraph integration — a real LangGraph agent on a local model (Ollama), verified by Aura-State — no API key
• Real-data verification — Z3 verifies 1,000 real sales records (3,000 obligations, ~1,200/sec); conformal hits 91.3% coverage on real diabetes records — no LLM, no key
• Comparison — Aura-State vs LangGraph / CrewAI / Guardrails (they orchestrate; we verify)
• Usage Guide — code examples for every feature
• Algorithm Reference — deep-dive into CTL, Z3, Thompson sampling, conformal prediction
• Contributing — architecture overview and how to contribute
• Pipeline walkthrough — an illustrative end-to-end pipeline (mocked LLM); for measured results see real-data

License

MIT