# How to Use PostgresStateStore for Scalable Agent State in aisuite

> Learn how to leverage PostgresStateStore for scalable agent state in aisuite. Achieve transactional persistence and optimistic concurrency control for high-throughput agent threads.

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

---

**PostgresStateStore implements aisuite's StateStore protocol on PostgreSQL, providing transactional persistence, optimistic concurrency control, and compaction support for high-throughput agent threads.**

The aisuite framework provides a state store abstraction that allows agents to persist conversation and run state across sessions. The `PostgresStateStore` class in [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py) delivers a production-ready implementation backed by PostgreSQL, enabling multiple workers to safely read and write shared thread state while supporting historical compaction for long-running conversations.

## Core Architecture of PostgresStateStore

The store separates message storage from thread metadata, allowing efficient updates without rewriting entire conversation histories.

### StateStore Protocol Implementation

As defined in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py), the `StateStore` protocol requires three core methods: `save_state`, `load_state`, and `delete_state`. `PostgresStateStore` implements this interface, making it interchangeable with in-memory or file-based stores while adding database-backed durability. This protocol compliance ensures your agent code remains agnostic to the underlying storage mechanism.

### Database Schema Design

When instantiated with `create_schema=True`, the store executes `SCHEMA_STATEMENTS` to create three tables in your PostgreSQL database:

- **`agent_thread_heads`** – Stores one row per thread containing the current context, full-history message IDs, step IDs, JSON-encoded state, revision counter, and timestamps.
- **`agent_messages`** – Contains individual message JSON objects with optional artifact references.
- **`agent_compactions`** – Records summarization events for audit trails.

This schema design keeps the state JSON small by storing messages as separate rows, minimizing write amplification during updates.

## Installing and Configuring PostgresStateStore

### Prerequisites

Install the PostgreSQL driver required by the store:

```bash
pip install psycopg[binary]

```

Create a PostgreSQL database (e.g., `aisuite`) and a user with read/write privileges.

### Instantiation Methods

You can initialize the store using either an existing connection or a DSN string:

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

# Option 1: Pass a live psycopg connection

import psycopg
conn = psycopg.connect("host=localhost dbname=aisuite user=aisuite password=secret")
store = PostgresStateStore(conn, create_schema=True)

# Option 2: Use a DSN string (recommended for most applications)

store = PostgresStateStore.from_dsn(
    "postgresql://aisuite:secret@localhost/aisuite",
    create_schema=True,
)

```

Setting `create_schema=True` ensures the three core tables exist before operations begin.

## Saving and Loading Agent State

### Persisting RunState

The `save_state` method in [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py) handles inserting new messages and updating thread heads within a single transaction. It accepts a `thread_id` and a `RunState` object from [`aisuite/framework/run_state.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/framework/run_state.py):

```python
from aisuite.framework.run_state import RunState
from aisuite.agents.state_store import StoredRunState

# Construct a run state

run_state = RunState(
    messages=[{"role": "user", "content": "What is the weather?"}],
    steps=[]
)

# Persist with optimistic concurrency disabled (revision=None)

stored: StoredRunState = store.save_state(
    thread_id="weather-thread",
    state=run_state,
)
print(f"Saved revision: {stored.revision}")

```

### Retrieving State

Use `load_state` to reconstruct a `StoredRunState` from the database:

```python
restored = store.load_state("weather-thread")
if restored:
    print(f"Loaded {len(restored.state.messages)} messages")

```

### Handling Concurrent Updates

`PostgresStateStore` implements **optimistic concurrency control** via the `_assert_revision` helper. Each write operation increments a `revision` column in `agent_thread_heads`. If you pass a specific `revision` to `save_state` and the database row has changed since that revision, the store raises a `StateConflictError`:

```python
try:
    # Attempt to update only if revision matches

    store.save_state(
        thread_id="weather-thread",
        state=new_state,
        revision=expected_revision,
    )
except StateConflictError:
    # Handle concurrent modification (reload and retry)

    pass

```

This mechanism ensures atomic updates even when multiple workers access the same thread simultaneously.

## Compacting Large Conversation Histories

Long-running agents accumulate message histories that may exceed token limits. The `compact_state` method replaces a span of messages with a summary while maintaining auditability:

```python
summary_msg = {
    "role": "assistant", 
    "content": "User asked about weather; assistant replied sunny."
}

compacted = store.compact_state(
    thread_id="weather-thread",
    source_message_ids=["msg-1", "msg-2"],  # IDs from get_thread_head()

    summary_message=summary_msg,
    reason="reduce token usage",
    model="gpt-4",
)
print(f"New revision after compaction: {compacted.revision}")

```

The method inserts the summary into `agent_messages`, records the event in `agent_compactions`, and updates the thread head's model-context ID list. Query compaction history using `list_compactions`:

```python
for record in store.list_compactions("weather-thread"):
    print(f"{record.created_at}: {record.summary_text[:50]}...")

```

## Complete Integration Example

This end-to-end example demonstrates schema creation, state persistence, compaction, and cleanup:

```python
from aisuite.agents.postgres_state_store import PostgresStateStore
from aisuite.framework.run_state import RunState

# 1. Initialize store

store = PostgresStateStore.from_dsn(
    "postgresql://aisuite:secret@localhost/aisuite",
    create_schema=True,
)

# 2. Create and save initial state

state = RunState(
    messages=[
        {"role": "user", "content": "Explain quantum computing"},
        {"role": "assistant", "content": "Quantum computing uses qubits..."},
    ],
    steps=[]
)
saved = store.save_state(thread_id="quantum-demo", state=state)

# 3. Load state (simulating a separate worker process)

loaded = store.load_state("quantum-demo")
print(f"Message count: {len(loaded.state.messages)}")

# 4. Compact the conversation

summary = {"role": "assistant", "content": "Summary: Quantum computing explanation provided."}
store.compact_state(
    thread_id="quantum-demo",
    source_message_ids=["msg-1", "msg-2"],
    summary_message=summary,
    reason="token optimization",
    model="gpt-4",
)

# 5. Verify compaction records

compactions = store.list_compactions("quantum-demo")
print(f"Total compactions: {len(compactions)}")

# 6. Delete thread when finished

store.delete_state("quantum-demo")

```

## Summary

- **PostgresStateStore** in [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py) provides a PostgreSQL-backed implementation of the `StateStore` protocol for production agent deployments.
- The store creates three tables (`agent_thread_heads`, `agent_messages`, `agent_compactions`) automatically when `create_schema=True` is passed during initialization.
- **Optimistic concurrency control** via revision checking prevents lost updates when multiple workers access the same thread.
- The **compaction** feature allows you to replace message spans with summaries, storing the reduction events in `agent_compactions` for audit purposes.
- Use `PostgresStateStore.from_dsn()` for simple configuration or pass an existing `psycopg.Connection` for connection pool integration.

## Frequently Asked Questions

### How does PostgresStateStore handle concurrent writes from multiple workers?

The store implements optimistic concurrency control using a `revision` column in the `agent_thread_heads` table. When calling `save_state` with a specific `revision` argument, the internal `_assert_revision` method verifies the row matches before updating. If another process has modified the thread, a `StateConflictError` is raised, allowing your application to reload the current state and retry.

### Can I use PostgresStateStore with an existing PostgreSQL connection pool?

Yes. Instead of using `from_dsn()`, pass an existing `psycopg.Connection` directly to the `PostgresStateStore` constructor. This allows integration with connection pools or custom connection configurations: `PostgresStateStore(conn, create_schema=True)`.

### What is the difference between `agent_messages` and `agent_thread_heads`?

The `agent_messages` table stores individual message JSON objects with unique IDs, while `agent_thread_heads` maintains the current state of each conversation thread, including references to which messages form the current context. This separation allows `save_state` to insert only new messages and update a single row in `agent_thread_heads`, rather than rewriting the entire conversation history on every save.

### How do I obtain message IDs for the `compact_state` method?

Call `get_thread_head(thread_id)` to retrieve the current thread metadata, which includes `model_context_message_ids` and `full_history_message_ids`. These lists contain the message IDs needed for the `source_message_ids` parameter in `compact_state`.