# How to Create a Custom Extension Using the ExtensionState Trait in Goose (Rust)

> Learn to create custom extensions in Goose Rust by implementing the ExtensionState trait. Manage typed session data persisted automatically in Goose's ExtensionData store.

- Repository: [Block Open Source/goose](https://github.com/block/goose)
- Tags: how-to-guide
- Published: 2026-04-05

---

**Implement the `ExtensionState` trait on a serializable struct to read and write typed session data that persists automatically in Goose's `ExtensionData` store.**

The `block/goose` repository provides a powerful session management system that allows developers to create custom extensions with persistent state. By implementing the `ExtensionState` trait, you can create a custom extension in Rust that maintains typed data across session lifecycles without manual serialization boilerplate.

## Understanding the ExtensionState Architecture

### How Session Data is Stored

Goose stores per-extension data inside a **session** via the `ExtensionData` struct defined in [`crates/goose/src/session/extension_data.rs`](https://github.com/block/goose/blob/main/crates/goose/src/session/extension_data.rs). This struct maintains a flattened `HashMap` where each extension's state is stored as JSON values keyed by `"extension_name.version"` strings (for example, `"todo.v0"`).

```rust
// crates/goose/src/session/extension_data.rs#L14-L20
#[derive(Debug, Clone, Serialize, Deserialize, Default, ToSchema)]
pub struct ExtensionData {
    #[serde(flatten)]
    pub extension_states: HashMap<String, Value>,
}

```

When the session is persisted to a file or database, the entire `ExtensionData` object—containing every extension's JSON payload—is serialized automatically.

### The ExtensionState Trait Interface

The **`ExtensionState`** trait, located in the same file at lines 44-82, provides the conversion layer between your typed Rust struct and the JSON storage. It supplies four required constants and methods that handle the composite key generation and serialization:

```rust
// crates/goose/src/session/extension_data.rs#L44-L82
pub trait ExtensionState: Sized + Serialize + for<'de> Deserialize<'de> {
    const EXTENSION_NAME: &'static str;
    const VERSION: &'static str;

    fn from_value(value: &Value) -> Result<Self> { /* serde conversion */ }
    fn to_value(&self) -> Result<Value> { /* serde conversion */ }
    fn from_extension_data(extension_data: &ExtensionData) -> Option<Self> { /* key lookup + from_value */ }
    fn to_extension_data(&self, extension_data: &mut ExtensionData) -> Result<()> { /* key insertion + to_value */ }
}

```

Goose ships with a reference implementation (`TodoState`) that demonstrates this pattern exactly:

```rust
// crates/goose/src/session/extension_data.rs#L84-L93
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TodoState {
    pub content: String,
}
impl ExtensionState for TodoState {
    const EXTENSION_NAME: &'static str = "todo";
    const VERSION: &'static str = "v0";
}

```

## Implementing a Custom Extension

### Step 1: Define the State Struct

Create a struct deriving `Serialize` and `Deserialize`. This struct represents the data you want to persist across sessions.

```rust
use serde::{Deserialize, Serialize};
use goose::session::ExtensionState;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MyCounterState {
    /// Current count value persisted across sessions
    pub count: u64,
}

```

### Step 2: Implement the ExtensionState Trait

Set the `EXTENSION_NAME` and `VERSION` constants. The name must be unique across all extensions in the system, while the version string allows schema migrations by changing the storage key.

```rust
impl ExtensionState for MyCounterState {
    // Generates the storage key "my_counter.v0"
    const EXTENSION_NAME: &'static str = "my_counter";
    const VERSION: &'static str = "v0";
}

```

### Step 3: Persist State in Session Operations

Use `from_extension_data` to retrieve your typed state (returns `None` if not found) and `to_extension_data` to save changes back into the session's `ExtensionData` map.

```rust
use goose::session::{ExtensionData, SessionManager};

async fn increment_counter(
    session_manager: &SessionManager, 
    session_id: &str
) -> anyhow::Result<()> {
    // Load session with create_if_missing=true
    let mut session = session_manager.get_session(session_id, true).await?;
    
    // Retrieve existing state or default to zero
    let mut state = MyCounterState::from_extension_data(&session.extension_data)
        .unwrap_or(MyCounterState { count: 0 });
    
    // Mutate the state
    state.count += 1;
    
    // Write back to the session's ExtensionData under key "my_counter.v0"
    state.to_extension_data(&mut session.extension_data)?;
    
    // Persist the entire session (including all extension data)
    session_manager.save_session(session).await?;
    Ok(())
}

```

### Step 4: Retrieve State Later

Reading the state uses the same pattern without mutation:

```rust
async fn get_counter(
    session_manager: &SessionManager, 
    session_id: &str
) -> anyhow::Result<u64> {
    let session = session_manager.get_session(session_id, false).await?;
    let count = MyCounterState::from_extension_data(&session.extension_data)
        .map(|s| s.count)
        .unwrap_or(0);
    Ok(count)
}

```

## Registering Your Extension for Discovery

While state management works independently, you can register the extension in the configuration system defined in [`crates/goose/src/config/extensions.rs`](https://github.com/block/goose/blob/main/crates/goose/src/config/extensions.rs) to enable UI exposure and permission checks. This makes your extension selectable via the `ExtensionConfig::Builtin` variant:

```yaml
extensions:
  my_counter:
    enabled: true
    builtin:
      name: my_counter
      description: "Simple counter example"

```

The **`ExtensionManager`** in [`crates/goose/src/agents/extension_manager.rs`](https://github.com/block/goose/blob/main/crates/goose/src/agents/extension_manager.rs) handles runtime loading of resources for enabled extensions, connecting your state management with the broader Goose agent system.

## Summary

- **Implement `ExtensionState`** on a struct deriving `Serialize` and `Deserialize` to create persistent extension state in `block/goose`
- **Set `EXTENSION_NAME` and `VERSION`** constants to generate unique storage keys in the format `"name.version"` within `ExtensionData`
- **Use `from_extension_data`** to retrieve typed state and `to_extension_data` to persist changes back to the session
- **Session persistence is automatic** when calling `session_manager.save_session()` because `ExtensionData` is part of the core `Session` struct
- **Optional configuration** in [`crates/goose/src/config/extensions.rs`](https://github.com/block/goose/blob/main/crates/goose/src/config/extensions.rs) enables UI discovery and runtime permission management through the `ExtensionManager`

## Frequently Asked Questions

### What is the difference between ExtensionState and ExtensionConfig in Goose?

`ExtensionState` manages runtime data persistence within sessions, defined in [`crates/goose/src/session/extension_data.rs`](https://github.com/block/goose/blob/main/crates/goose/src/session/extension_data.rs), while `ExtensionConfig` in [`crates/goose/src/config/extensions.rs`](https://github.com/block/goose/blob/main/crates/goose/src/config/extensions.rs) handles static configuration like enabling or disabling extensions and setting permissions. You can use `ExtensionState` for programmatic extensions without registering in the config system, but `ExtensionConfig` is required for UI exposure and permission checks.

### How do I migrate data when changing my extension's schema?

Increment the `VERSION` constant in your `ExtensionState` implementation. Goose stores data under keys formatted as `"extension_name.version"`, so changing from `"v0"` to `"v1"` creates a new storage slot. Implement migration logic by manually reading from the old key using `from_value` and writing to the new key using `to_value` during your extension's initialization phase.

### Can multiple extensions access the same session data concurrently?

Yes, because `ExtensionData` stores each extension's state under a unique namespaced key in the `HashMap<String, Value>` defined in [`crates/goose/src/session/extension_data.rs`](https://github.com/block/goose/blob/main/crates/goose/src/session/extension_data.rs). Multiple extensions can read and write to the same `ExtensionData` instance without collisions, as each operates on distinct keys like `"todo.v0"` and `"my_counter.v0"`.

### Where is the extension state actually persisted?

The state is serialized as part of the `Session` struct when `session_manager.save_session()` is called. According to the `block/goose` source code, the `ExtensionData` object containing all extension JSON values is serialized alongside other session metadata, typically to a file or database depending on your `SessionManager` backend implementation.