# How to Use State Stores in AISuite: File, In-Memory, and Postgres Implementations

> Learn to use AISuite's FileStateStore, InMemoryStateStore, and PostgresStateStore for robust agent state persistence. Ensure seamless agent restarts and concurrent access.

- Repository: [Andrew Ng/aisuite](https://github.com/andrewyng/aisuite)
- Tags: how-to-guide
- Published: 2026-08-04

---

**AISuite provides three built-in state store implementations—`InMemoryStateStore`, `FileStateStore`, and `PostgresStateStore`—that all implement the `StateStore` protocol for persisting agent run-state across process restarts and concurrent access.**

The **state stores in aisuite** abstract conversation persistence behind a unified interface defined in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py). Whether you are unit testing with ephemeral storage, prototyping with local files, or deploying to production with PostgreSQL, the same API surface handles thread lifecycle management, optimistic concurrency control, and metadata storage.

## The StateStore Protocol Interface

All persistence backends implement the **`StateStore`** protocol, which requires three methods with the following signatures:

```python
save_state(thread_id: str, state: RunState, *, revision: int | None = None,
           metadata: dict[str, Any] | None = None) -> StoredRunState
load_state(thread_id: str) -> Optional[StoredRunState]
delete_state(thread_id: str) -> None

```

The **`Runner`** class consumes this interface when you supply both `state_store` and `thread_id` arguments to `run_sync()` or `continue_sync()`. During execution, the runner dehydrates the final `RunState` and delegates persistence to the store's `save_state` method. On continuation, `load_state` rehydrates the conversation history before appending new user messages.

### Optimistic Concurrency with Revisions

Each store implements optimistic locking via the optional **`revision`** parameter. When you supply a revision integer, the store validates it against the persisted value before writing. If the revision has changed, the store raises **`StateConflictError`**, preventing stale overwrites. Successful increments automatically bump the revision counter—Postgres uses `revision = revision + 1` within a transaction, while the file store atomically replaces the JSON file only after verification.

## InMemoryStateStore for Unit Testing

**`InMemoryStateStore`** provides a pure-Python dictionary backend that lives only for the process lifetime. Located in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py), this store is ideal for unit tests or ephemeral conversational state where persistence across restarts is unnecessary.

```python
import aisuite as ai

# Initialize the ephemeral store

store = ai.InMemoryStateStore()

# Run an agent with persistence enabled

result = ai.Runner.run_sync(
    agent=my_agent,
    input="Hello",
    state_store=store,
    thread_id="test_thread_001",
)

# Verify state was captured

saved = store.load_state("test_thread_001")
print(saved.revision)  # → 1

print(saved.state.messages[-1]["content"])  # → agent response

# Continue the conversation

next_result = ai.Runner.continue_sync(
    target=my_agent,
    input="Follow-up question",
    state_store=store,
    thread_id="test_thread_001",
)

```

The store maintains all data in memory and resets when the process exits. See [`tests/agents/test_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/tests/agents/test_state_store.py) for comprehensive round-trip and revision handling tests.

## FileStateStore for Local Development

**`FileStateStore`** persists state as JSON files under a configurable directory, defaulting to `.aisuite/state`. Also defined in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py), this implementation writes to a temporary file first, then atomically moves it into place using `os.replace()` to prevent corruption during crashes.

```python
import pathlib
import aisuite as ai

# Configure a custom state directory

state_dir = pathlib.Path("./local_agent_state")
store = ai.FileStateStore(state_dir)

# First run creates `local_agent_state/thread%2Fuser%3A1.json`

result = ai.Runner.run_sync(
    agent=my_agent,
    input="What is the weather?",
    state_store=store,
    thread_id="thread/user:1",
)

# Inspect the persisted JSON

state_file = state_dir / "thread%2Fuser%3A1.json"
print(state_file.read_text())  # Contains schema_version, messages, revision

# Load and continue after process restart

saved = store.load_state("thread/user:1")
assert saved.revision == 1

```

File names are URL-encoded to safely handle special characters in thread IDs. The JSON payload includes a `schema_version` field for future migration compatibility.

## PostgresStateStore for Production

**`PostgresStateStore`** implements durable persistence using three relational tables: `agent_thread_heads`, `agent_messages`, and `agent_compactions`. Implemented in [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py), this backend supports high-concurrency workloads and long-term conversation history via a compaction API.

```python
import aisuite as ai

# Initialize with automatic schema creation

store = ai.PostgresStateStore.from_dsn(
    dsn="postgresql://user:password@localhost/aisuite",
    create_schema=True,  # Creates tables if missing

)

# Run with durable persistence

result = ai.Runner.run_sync(
    agent=my_agent,
    input="Summarize the quarterly report",
    state_store=store,
    thread_id="proj-q3-2024",
)

# Compact old messages to manage table growth

store.compact_state(
    thread_id="proj-q3-2024",
    source_message_ids=["msg-001", "msg-002", "msg-003"],
    summary_message={"role": "assistant", "content": "Summary of Q3..."},
    reason="token limit management",
)

```

The compaction feature archives older messages into summary entries, reducing storage overhead while preserving conversation context. All operations occur within database transactions that enforce revision consistency.

## Summary

- **StateStore protocol** defines a uniform interface (`save_state`, `load_state`, `delete_state`) implemented by all persistence backends in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py).
- **InMemoryStateStore** offers zero-configuration, ephemeral storage ideal for testing scenarios where state should not survive process restarts.
- **FileStateStore** provides atomic JSON file persistence suitable for single-node applications, with automatic directory creation and URL-safe thread ID encoding.
- **PostgresStateStore** delivers production-grade durability with optimistic concurrency, schema management, and conversation compaction via [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py).
- **Revision handling** across all stores prevents lost updates through integer-based optimistic locking, raising `StateConflictError` on mismatched versions.

## Frequently Asked Questions

### What is the primary difference between InMemoryStateStore and FileStateStore?

**InMemoryStateStore** maintains conversation state exclusively in a Python dictionary, losing all data when the process terminates, whereas **FileStateStore** serializes state to JSON files on disk, enabling persistence across process restarts. Both implement identical method signatures and support revision-based concurrency control, making them interchangeable during testing and local development.

### How does PostgresStateStore handle database schema initialization?

When initialized with `create_schema=True`, `PostgresStateStore.from_dsn()` automatically executes DDL to create the required tables (`agent_thread_heads`, `agent_messages`, `agent_compactions`) if they do not already exist. For production deployments, you may disable this flag and manage schema migrations externally using the definitions found in [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py).

### What happens if two processes attempt to update the same thread simultaneously?

All three stores implement optimistic concurrency control. If both processes read revision `N` and attempt to save, the first writer succeeds and increments the revision to `N+1`. The second writer's `save_state` call supplies the stale revision `N`, triggering a `StateConflictError`. Your application should catch this exception, reload the current state via `load_state`, and retry the operation with the updated revision.

### Can I implement a custom StateStore for Redis or MongoDB?

Yes. The `StateStore` protocol in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py) uses Python's `typing.Protocol`, allowing any class that implements the three required methods to be passed to `Runner`. Your custom implementation must handle `thread_id` uniqueness, atomic updates for the `revision` parameter, and round-trip preservation of the `RunState` and `metadata` fields to be fully compatible with the runner's persistence lifecycle.