Agent persistence

Agent Persistence is a feature that provides checkpoint functionality for AI agents in the Koog framework. It lets you save and restore the state of an agent at specific points during execution, enabling capabilities such as:

• Resuming agent execution from a specific point
• Rolling back to previous states
• Persisting agent state across sessions

Key concepts

Checkpoints

A checkpoint captures the complete state of an agent at a specific point in its execution, including:

• Message history (all interactions between user, system, assistant, and tools)
• Current node being executed
• Input data for the current node
• Timestamp of creation

Checkpoints are identified by unique IDs and are associated with a specific agent.

Installation

To use the Agent Persistence feature, add it to your agent's configuration:

KotlinJava

val agent = AIAgent(
    promptExecutor = executor,
    llmModel = OllamaModels.Meta.LLAMA_3_2,
) {
    install(Persistence) {
        // Use in-memory storage for snapshots
        storage = InMemoryPersistenceStorageProvider()
    }
}

AIAgent<String, String> agent = AIAgent.<String, String>builder()
    .promptExecutor(SimplePromptExecutorsKt.simpleOllamaAIExecutor("http://localhost:11434"))
    .llmModel(OllamaModels.Meta.LLAMA_3_2)
    .install(Persistence.Feature, cfg -> {
        // Use in-memory storage for snapshots
        cfg.setStorage(new InMemoryPersistenceStorageProvider());
    })
.build();

Configuration options

The Agent Persistence feature has three main configuration options:

• Storage provider: the provider used to save and retrieve checkpoints.
• Continuous persistence: automatic creation of checkpoints after each node is run.
• Rollback strategy: determines which state will be restored when rolling back to a checkpoint.

Storage provider

Set the storage provider that will be used to save and retrieve checkpoints:

KotlinJava

install(Persistence) {
    storage = InMemoryPersistenceStorageProvider()
}

AIAgent<String, String> agent = AIAgent.<String, String>builder()
    .promptExecutor(SimplePromptExecutorsKt.simpleOllamaAIExecutor("http://localhost:11434"))
    .llmModel(OllamaModels.Meta.LLAMA_3_2)
    .install(Persistence.Feature, cfg -> {
        cfg.setStorage(new InMemoryPersistenceStorageProvider());
    })
    .build();

The framework includes the following built-in providers:

• InMemoryPersistenceStorageProvider: stores checkpoints in memory (lost when the application restarts).
• FilePersistenceStorageProvider: persists checkpoints to the file system.
• NoPersistenceStorageProvider: a no-op implementation that does not store checkpoints. This is the default provider.

You can also implement custom storage providers by implementing the PersistenceStorageProvider interface. For more information, see Custom storage providers.

Continuous persistence

Continuous persistence means that a checkpoint is automatically created after each node is run. To disable continuous persistence, use the code below:

KotlinJava

install(Persistence) {
    enableAutomaticPersistence = false
}

AIAgent<String, String> agent = AIAgent.<String, String>builder()
    .promptExecutor(SimplePromptExecutorsKt.simpleOllamaAIExecutor("http://localhost:11434"))
    .llmModel(OllamaModels.Meta.LLAMA_3_2)
    .install(Persistence.Feature, cfg -> {
        cfg.setEnableAutomaticPersistence(true);
    })
    .build();

If continuous persistence is disabled, you can still create checkpoints manually.

Basic usage

Creating a checkpoint

To learn how to create a checkpoint at a specific point in your agent's execution, see the code sample below:

KotlinJava

suspend fun example(context: AIAgentContext) {
    // Create a checkpoint with the current state
    val checkpoint = context.persistence().createCheckpointAfterNode(
        agentContext = context,
        nodePath = context.executionInfo.path(),
        lastOutput = outputData,
        lastOutputType = outputType,
        checkpointId = context.runId,
        version = 0L
    )

    // The checkpoint ID can be stored for later use
    val checkpointId = checkpoint?.checkpointId
}

Restoring from a checkpoint

To restore the state of an agent from a specific checkpoint, follow the code sample below:

KotlinJava

suspend fun example(context: AIAgentContext, checkpointId: String) {
    // Roll back to a specific checkpoint
    context.persistence().rollbackToCheckpoint(checkpointId, context)

    // Or roll back to the latest checkpoint
    context.persistence().rollbackToLatestCheckpoint(context)
}

Rolling back all side-effects produced by tools

It's quite common for some tools to produce side-effects. Specifically, when you are running your agents on the backend, some of the tools would likely perform some database transactions. This makes it much harder for your agent to travel back in time.

Imagine you have a tool createUser that creates a new user in your database. And your agent has populated multiple tool calls overtime:

tool call: createUser "Alex"

->>>> checkpoint-1 <<<<-

tool call: createUser "Daniel"
tool call: createUser "Maria"

And now you would like to roll back to a checkpoint. Restoring the agent's state (including message history, and strategy graph node) alone would not be sufficient to achieve the exact state of the world before the checkpoint. You should also restore the side-effects produced by your tool calls. In our example, this would mean removing Maria and Daniel from the database.

With Koog Persistence you can achieve that by providing a RollbackToolRegistry to Persistence feature config:

KotlinJava

install(Persistence) {
    enableAutomaticPersistence = true
    rollbackToolRegistry = RollbackToolRegistry {
        // For every `createUser` tool call there will be a `removeUser` invocation in the reverse order 
        // when rolling back to the desired execution point.
        // Note: `removeUser` tool should take the same exact arguments as `createUser`. 
        // It's the developer's responsibility to make sure that `removeUser` invocation rolls back all side-effects of `createUser`:
        registerRollback(::createUser, ::removeUser)
    }
}

Using extension functions

The Agent Persistence feature provides convenient extension functions for working with checkpoints:

KotlinJava

suspend fun example(context: AIAgentContext) {
    // Access the checkpoint feature
    val checkpointFeature = context.persistence()

    // Or perform an action with the checkpoint feature
    context.withPersistence { ctx ->
        // 'this' is the checkpoint feature
        createCheckpointAfterNode(
            agentContext = ctx,
            nodePath = ctx.executionInfo.path(),
            lastOutput = outputData,
            lastOutputType = outputType,
            checkpointId = ctx.runId,
            version = 0L
        )
    }
}

Advanced usage

Custom storage providers

You can implement custom storage providers by implementing the PersistenceStorageProvider interface:

KotlinJava

class MyCustomStorageProvider<MyFilterType> : PersistenceStorageProvider<MyFilterType> {
    override suspend fun getCheckpoints(sessionId: String, filter: MyFilterType?): List<AgentCheckpointData> {
        TODO("Not yet implemented")
    }

    override suspend fun saveCheckpoint(sessionId: String, agentCheckpointData: AgentCheckpointData) {
        TODO("Not yet implemented")
    }

    override suspend fun getLatestCheckpoint(sessionId: String, filter: MyFilterType?): AgentCheckpointData? {
        TODO("Not yet implemented")
    }
}

To use your custom provider in the feature configuration, set it as the storage when configuring the Agent Persistence feature in your agent.

KotlinJava

install(Persistence) {
    storage = MyCustomStorageProvider<Any>()
}

Setting execution points

For advanced control, you can directly set the execution point of an agent:

KotlinJava

fun example(context: AIAgentContext) {
    // You can set the execution point before some node and provide an input for it:
    context.persistence().setExecutionPoint(
        agentContext = context,
        nodePath = context.executionInfo.path(),
        messageHistory = customMessageHistory,
        input = customInput
    )

    // Or after some node and provide an output from the node:
    context.persistence().setExecutionPointAfterNode(
        agentContext = context,
        nodePath = context.executionInfo.path(),
        messageHistory = customMessageHistory,
        output = customOutput
    )
}

This allows for more fine-grained control over the agent's state beyond just restoring from checkpoints.