# AISuite State Store Architecture: InMemory, File, and Postgres Implementation Guide

> Explore AISuite's flexible State Store architecture. Learn about InMemory, File, and Postgres implementations for efficient agent run persistence and easy backend switching.

- Repository: [Andrew Ng/aisuite](https://github.com/andrewyng/aisuite)
- Tags: architecture
- Published: 2026-07-27

---

**AISuite abstracts agent run persistence through a `StateStore` protocol implemented by three swappable backends—InMemory for transient testing, File for simple disk storage, and Postgres for production workloads—enabling seamless storage strategy changes without modifying consumer code.**

The `andrewyng/aisuite` repository provides a flexible persistence layer for agent conversation state, allowing developers to choose between lightweight in-process storage, portable JSON files, or a relational database depending on their operational requirements. This architecture centers on a common protocol that standardizes how `RunState` objects are saved, loaded, and versioned across disparate storage technologies.

## The StateStore Protocol

All storage implementations adhere to the `StateStore` protocol defined in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py) (lines 53-65). This interface mandates three core methods:

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

def load_state(self, thread_id: str) -> Optional[StoredRunState]: ...

def delete_state(self, thread_id: str) -> None: ...

```

By coding against this protocol, agent logic remains agnostic to whether the underlying storage is volatile memory, the local filesystem, or a remote PostgreSQL instance.

## Shared Data Model

The architecture relies on two primary data structures defined across the codebase.

### StoredRunState

The `StoredRunState` dataclass (defined in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py), lines 19-50) wraps the core `RunState` with metadata essential for concurrency control and auditing. It includes:

- **`revision`** – An integer implementing optimistic concurrency control
- **`created_at`/`updated_at`** – ISO timestamps tracking mutation history
- **`metadata`** – Optional user-defined dictionary for custom annotations

This class provides `to_dict()` and `from_dict()` methods for JSON serialization, ensuring consistent data representation across all backends.

### RunState

The actual agent execution state resides in `RunState` (located in [`aisuite/agents/types.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/types.py), lines 46-80). This dataclass contains the agent name, message history, execution status, step counters, and auxiliary fields. The state store system ensures all content is JSON-serializable via `ensure_json_serializable` before persistence.

## The Three Storage Backends

AISuite ships with three concrete implementations of the `StateStore` protocol, each optimized for specific deployment scenarios.

### InMemoryStateStore

Located in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py) (lines 68-90), **InMemoryStateStore** maintains a Python dictionary mapping thread IDs to `StoredRunState` objects. This backend provides the fastest access speeds but offers no durability across process restarts.

Key characteristics:
- Stores data in process memory using a private `_store` dict
- Returns **deep-copied** state objects to prevent accidental in-place mutations
- Validates revisions via `_assert_revision` before accepting writes

This implementation is ideal for **unit tests** and short-lived ephemeral agents where persistence is unnecessary.

### FileStateStore

The **FileStateStore** class (lines 101-131 in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py)) provides durable storage by serializing each thread's state to individual JSON files on disk.

Key characteristics:
- Default storage location is `.aisuite/state` (configurable via `root` parameter)
- Atomic write operations using temporary files followed by `os.replace()` to prevent corruption during write failures
- One JSON file per thread ID containing the serialized `StoredRunState`

This backend suits **local development** and single-node deployments requiring durability without external database dependencies.

### PostgresStateStore

For production workloads, **PostgresStateStore** (implemented in [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py), lines 95-115) offers scalable, relational storage with advanced features.

Key characteristics:
- Schema consists of three tables: `agent_thread_heads` (current state pointers), `agent_messages` (message history), and `agent_compactions` (summarized message archives)
- Uses `SELECT ... FOR UPDATE` for transactional consistency
- Supports **compaction**—replacing old messages with summaries to manage storage growth
- Can list compaction records for audit trails

Connect to the database using the `from_dsn()` class method with `create_schema=True` to initialize tables automatically.

## Concurrency Control and Data Integrity

All three backends implement **optimistic concurrency control** through the `_assert_revision` helper (lines 71-77 in [`state_store.py`](https://github.com/andrewyng/aisuite/blob/main/state_store.py)). Before writing, each store verifies that the provided `revision` parameter matches the current stored revision. If another process has modified the state in the interim, the method raises `StateConflictError`, preventing lost updates.

When saving state, implementations delegate to `_next_stored_state` (lines 151-167) to automatically increment the revision counter and update timestamps, ensuring monotonic versioning across storage media.

## Compaction Strategy

Only `PostgresStateStore` implements the `compact_state` method (lines 70-87 in [`postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/postgres_state_store.py)). This feature addresses long-running conversations by:

1. Accepting a list of source message IDs to archive
2. Inserting a summary message that captures the semantic content of the removed messages
3. Recording a `CompactionRecord` linking the summary to the original message IDs

Compaction reduces storage costs and context window pressure while maintaining conversation continuity through the summary message.

## Practical Usage Examples

The following demonstrations show identical API usage across all three backends.

### In-Memory Storage (Testing)

```python
from aisuite.agents.state_store import InMemoryStateStore

mem_store = InMemoryStateStore()
stored = mem_store.save_state("thread_1", run_state, metadata={"user": "test"})
loaded = mem_store.load_state("thread_1")
mem_store.delete_state("thread_1")

```

### File-Based Persistence (Local Development)

```python
from aisuite.agents.state_store import FileStateStore

file_store = FileStateStore(root=".aisuite/state")
file_store.save_state("thread_2", run_state, revision=1)
loaded = file_store.load_state("thread_2")
file_store.delete_state("thread_2")

```

### PostgreSQL Backend (Production)

```python
from aisuite.agents.postgres_state_store import PostgresStateStore

pg_store = PostgresStateStore.from_dsn(
    "postgresql://user:password@localhost/aisuite", 
    create_schema=True
)
pg_store.save_state("thread_3", run_state)

# Compact old messages to save space

pg_store.compact_state(
    thread_id="thread_3",
    source_message_ids=["msg1", "msg2"],
    summary_message={"role": "assistant", "content": "Summary of previous discussion"},
)
pg_store.delete_state("thread_3")

```

## Summary

- **Protocol-based design** allows swapping between InMemory, File, and Postgres storage without changing agent code, as implemented in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py).
- **Optimistic concurrency control** via revision checking prevents race conditions across all backends using the shared `_assert_revision` mechanism.
- **InMemoryStateStore** provides fast, transient storage ideal for testing scenarios requiring no disk I/O.
- **FileStateStore** offers atomic JSON file operations for simple durability needs without external dependencies.
- **PostgresStateStore** delivers production-grade persistence with SQL transactions, compaction for message archival, and scalable relational storage.
- **Shared data model** using `StoredRunState` and `RunState` ensures consistent serialization and metadata tracking across disparate storage technologies.

## Frequently Asked Questions

### When should I use each StateStore implementation?

**InMemoryStateStore** is optimal for unit tests and CI/CD pipelines where speed matters and data persistence is unnecessary. **FileStateStore** suits local development and single-instance deployments requiring restart durability without database administration overhead. **PostgresStateStore** is designed for production environments needing horizontal scalability, concurrent access from multiple processes, and long-term conversation archiving with compaction capabilities.

### How does optimistic concurrency control work in AISuite's state stores?

Before any write operation, the store compares the `revision` parameter provided to `save_state()` against the current revision stored for that thread ID. If they differ, `_assert_revision` raises `StateConflictError`, indicating another process modified the state since the last read. This mechanism prevents lost updates without requiring database locks in the File and InMemory implementations, while Postgres uses `SELECT ... FOR UPDATE` to handle concurrent modifications safely.

### What is message compaction and why is it only available in Postgres?

Compaction is the process of replacing a sequence of old messages with a single summary message to reduce storage costs and context window token usage. It is implemented only in `PostgresStateStore` because it requires relational schema support to maintain audit trails via the `agent_compactions` table and transactional integrity when replacing multiple message records. The InMemory and File stores are intended for simpler use cases where storage limitations are less critical, though developers could implement similar functionality manually if needed.

### Can I migrate conversation data from FileStateStore to PostgresStateStore?

Yes, migration is possible because both implementations use the same `StoredRunState` serialization format. You would iterate through thread directories in your `.aisuite/state` folder, load each state using `FileStateStore.load_state()`, and persist it using `PostgresStateStore.save_state()`. Since the `RunState` and revision metadata are identical across formats, the transition requires no data transformation—only copying the JSON-serializable content from files to database rows.