# How to Implement Custom State Stores in aisuite: Extending Beyond Built-In Options

> Learn to implement custom state stores in aisuite. Create a StateStore class with save load and delete methods for persistent agent conversations. Extend beyond built-in options.

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

---

**Implementing custom state stores in aisuite requires creating a class that conforms to the `StateStore` protocol defined in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py), implementing three required methods—`save_state`, `load_state`, and `delete_state`—and returning `StoredRunState` objects to enable persistent agent conversations across any backend.**

The aisuite library by Andrew Ng provides robust conversation persistence through built-in storage options, but production deployments often require specialized backends like Redis, DynamoDB, or proprietary databases. By implementing the `StateStore` protocol, you can integrate any storage mechanism while maintaining full compatibility with the `Runner` class and its conversation continuation features.

## Understanding the StateStore Protocol

The foundation of custom state storage lies in the `StateStore` protocol located in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py). This protocol defines the contract that any storage backend must satisfy to work with aisuite's agent runner.

The protocol requires three specific methods:

- **`save_state(self, thread_id: str, state: RunState, *, revision: int | None = None, metadata: Optional[dict[str, Any]] = None) -> StoredRunState`** – Persists the conversation state and returns a `StoredRunState` object containing the new revision number.
- **`load_state(self, thread_id: str) -> Optional[StoredRunState]`** – Retrieves the previously saved state for a given thread identifier, or `None` if no state exists.
- **`delete_state(self, thread_id: str) -> None`** – Removes the stored state for a specific thread.

The `RunState` class contains the complete conversation history, tool call results, and metadata, while `StoredRunState` wraps this data with revision tracking and timestamps essential for optimistic concurrency control.

## Handling Optimistic Concurrency with Revisions

Aisuite implements optimistic concurrency control through the `revision` parameter. When `Runner.continue_sync` resumes a conversation, it passes the current revision to `save_state`; if the stored revision has changed since loading, your implementation must raise `StateConflictError` (available from [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py)).

This pattern prevents lost updates when multiple processes interact with the same thread. The built-in `PostgresStateStore` and `FileStateStore` demonstrate this check via the `_assert_revision` helper logic, ensuring that stale writes are rejected before data corruption occurs.

## Complete Implementation: SQLite State Store

Below is a production-ready implementation using SQLite that demonstrates proper revision handling, JSON serialization, and the required protocol methods.

```python
import json
import sqlite3
from pathlib import Path
from aisuite.agents.state_store import (
    StoredRunState, 
    RunState, 
    StateConflictError,
    _next_stored_state
)

class SQLiteStateStore:
    """SQLite-backed implementation of the StateStore protocol."""
    
    def __init__(self, db_path: str | Path = "aisuite_state.db"):
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self._init_db()
    
    def _init_db(self):
        """Initialize the database schema."""
        self.conn.execute(
            """CREATE TABLE IF NOT EXISTS thread_state (
                thread_id TEXT PRIMARY KEY,
                revision INTEGER NOT NULL,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                state_json TEXT NOT NULL,
                metadata TEXT
            )"""
        )
        self.conn.commit()

    def save_state(
        self,
        thread_id: str,
        state: RunState,
        *,
        revision: int | None = None,
        metadata: dict | None = None,
    ) -> StoredRunState:
        cur = self.conn.cursor()
        cur.execute(
            "SELECT revision FROM thread_state WHERE thread_id = ?", 
            (thread_id,)
        )
        row = cur.fetchone()
        current_rev = row[0] if row else None
        
        # Enforce optimistic concurrency control

        if revision is not None and current_rev != revision:
            raise StateConflictError(
                f"State revision conflict for {thread_id!r}: "
                f"expected {revision}, found {current_rev}"
            )
        
        # Generate next stored state using the helper

        current_stored = None
        if row:
            current_stored = self.load_state(thread_id)
            
        stored = _next_stored_state(
            thread_id,
            state,
            current=current_stored,
            metadata=metadata,
        )
        
        # Upsert the state

        cur.execute(
            """INSERT INTO thread_state
               (thread_id, revision, created_at, updated_at, state_json, metadata)
               VALUES (?, ?, ?, ?, ?, ?)
               ON CONFLICT(thread_id) DO UPDATE SET
                 revision = excluded.revision,
                 updated_at = excluded.updated_at,
                 state_json = excluded.state_json,
                 metadata = excluded.metadata
            """,
            (
                thread_id,
                stored.revision,
                stored.created_at,
                stored.updated_at,
                json.dumps(state.to_dict(), sort_keys=True),
                json.dumps(metadata) if metadata else None,
            ),
        )
        self.conn.commit()
        return stored

    def load_state(self, thread_id: str) -> StoredRunState | None:
        cur = self.conn.cursor()
        cur.execute(
            """SELECT revision, created_at, updated_at, state_json, metadata 
               FROM thread_state WHERE thread_id = ?""",
            (thread_id,),
        )
        row = cur.fetchone()
        if not row:
            return None
            
        revision, created_at, updated_at, state_json, meta_json = row
        state = RunState.from_dict(json.loads(state_json))
        metadata = json.loads(meta_json) if meta_json else {}
        
        return StoredRunState(
            thread_id=thread_id,
            state=state,
            revision=revision,
            created_at=created_at,
            updated_at=updated_at,
            metadata=metadata,
        )

    def delete_state(self, thread_id: str) -> None:
        self.conn.execute(
            "DELETE FROM thread_state WHERE thread_id = ?", 
            (thread_id,)
        )
        self.conn.commit()

```

## Complete Implementation: Redis State Store

For distributed systems requiring high availability, Redis provides an excellent backend. This implementation stores state as JSON strings with proper revision checking.

```python
import json
import redis
from aisuite.agents.state_store import (
    RunState, 
    StoredRunState, 
    StateConflictError,
    _next_stored_state
)

class RedisStateStore:
    """Redis-backed implementation of the StateStore protocol."""
    
    def __init__(self, url: str = "redis://localhost:6379/0"):
        self.client = redis.from_url(url)
        self.key_prefix = "aisuite:state:"

    def save_state(
        self, 
        thread_id: str, 
        state: RunState, 
        *, 
        revision: int | None = None,
        metadata: dict | None = None,
    ) -> StoredRunState:
        key = f"{self.key_prefix}{thread_id}"
        existing = self.client.get(key)
        
        current_rev = None
        current_stored = None
        if existing:
            data = json.loads(existing)
            current_rev = data.get("revision")
            current_stored = StoredRunState.from_dict(data)
        
        # Check for concurrent modifications

        if revision is not None and current_rev != revision:
            raise StateConflictError(
                f"Revision mismatch for {thread_id}: expected {revision}, got {current_rev}"
            )
        
        # Generate new stored state

        stored = _next_stored_state(
            thread_id, 
            state, 
            current=current_stored, 
            metadata=metadata
        )
        
        # Store with JSON serialization

        self.client.set(
            key, 
            json.dumps(stored.to_dict(), default=str),
            # Optional: set expiration for cleanup

            # ex=86400  # 24 hours

        )
        return stored

    def load_state(self, thread_id: str) -> StoredRunState | None:
        key = f"{self.key_prefix}{thread_id}"
        data = self.client.get(key)
        if not data:
            return None
        return StoredRunState.from_dict(json.loads(data))

    def delete_state(self, thread_id: str) -> None:
        key = f"{self.key_prefix}{thread_id}"
        self.client.delete(key)

```

## Integrating Custom Stores with the Runner

The `Runner` class in [`aisuite/agents/runner.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/runner.py) accepts your custom store through the `state_store` parameter. At lines 52 and 58-59, the runner validates that `state_store` and `thread_id` are provided together, ensuring that persistent conversations always have unique identifiers.

When using custom stores, initialize your implementation and pass it to either `run_sync` for new conversations or `continue_sync` for existing threads:

```python
from aisuite import Agent, Runner

# Initialize your custom store

store = SQLiteStateStore(db_path="conversations.db")

# Create an agent

agent = Agent(
    name="assistant",
    model="openai:gpt-4o",
    instructions="You are a helpful assistant."
)

# Start a new conversation

result = Runner.run_sync(
    agent,
    "Hello, how are you?",
    thread_id="user_123_session_456",
    state_store=store
)

# Later, continue the same conversation

result = Runner.continue_sync(
    agent,
    "Can you summarize our discussion?",
    thread_id="user_123_session_456",
    state_store=store
)

```

At lines 74-75 and 112 in [`runner.py`](https://github.com/andrewyng/aisuite/blob/main/runner.py), the runner automatically calls `load_state` before execution and `save_state` after completion, handling the revision logic internally so your store receives the correct concurrency tokens.

## Key Files for Implementation Reference

| File | Purpose |
|------|---------|
| [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py) | Defines the `StateStore` protocol, `RunState`, `StoredRunState`, and helper functions like `_next_stored_state`. Also contains `InMemoryStateStore` and `FileStateStore` reference implementations. |
| [`aisuite/agents/postgres_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/postgres_state_store.py) | Production-grade PostgreSQL implementation demonstrating connection pooling, table schemas, and robust error handling. |
| [`aisuite/agents/runner.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/runner.py) | Contains the `Runner` class logic at lines 44-120 showing how stores are integrated with `run_sync` and `continue_sync` methods. |
| [`tests/agents/test_state_store.py`](https://github.com/andrewyng/aisuite/blob/main/tests/agents/test_state_store.py) | Unit tests validating the protocol contract; useful for verifying your custom implementation. |

## Summary

- **Protocol Compliance**: Implement `save_state`, `load_state`, and `delete_state` from the `StateStore` protocol in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py) to create compatible custom stores.
- **Concurrency Control**: Handle the `revision` parameter in `save_state` to support optimistic concurrency, raising `StateConflictError` when revisions mismatch.
- **Return Types**: Always return `StoredRunState` objects from `save_state` and `load_state` to maintain compatibility with the runner's internal logic.
- **Runner Integration**: Pass custom store instances to `Runner.run_sync` or `Runner.continue_sync` alongside a unique `thread_id` to enable persistent conversations.
- **Serialization**: Use `RunState.to_dict()` and `RunState.from_dict()` for consistent JSON serialization across different storage backends.

## Frequently Asked Questions

### What is the StateStore protocol in aisuite?

The **StateStore protocol** is a Python interface defined in [`aisuite/agents/state_store.py`](https://github.com/andrewyng/aisuite/blob/main/aisuite/agents/state_store.py) that specifies three methods: `save_state`, `load_state`, and `delete_state`. Any class implementing these methods can serve as a persistence layer for agent conversations, allowing aisuite to store and resume chat history across different sessions.

### How does revision handling work in custom state stores?

Revision handling implements **optimistic concurrency control**. The `save_state` method receives a `revision` parameter representing the last known state version. If the stored revision differs from this value, indicating another process modified the state, your implementation must raise `StateConflictError`. This prevents data loss when multiple clients interact with the same conversation thread simultaneously.

### Can I use custom state stores with both run_sync and continue_sync?

Yes, custom state stores work with both methods. Use `Runner.run_sync` with a new `thread_id` to initialize conversations, and `Runner.continue_sync` with the same `thread_id` to resume them. The runner automatically manages the revision lifecycle, calling `load_state` before execution and `save_state` after completion in both code paths.

### What metadata can I store in a custom state store?

The `metadata` parameter in `save_state` accepts any JSON-serializable dictionary. This allows you to store auxiliary information such as user identifiers, timestamps, conversation tags, or application-specific flags alongside the conversation state. The metadata is preserved in the returned `StoredRunState` object and passed through to subsequent operations.