Ariawisp/postgres persistence provider (#705)

rebased and fixed version of #481
by @ariawisp

---------

Co-authored-by: Aria Wisp <[email protected]>
Co-authored-by: Claude <[email protected]>

Author: 3 people

agents/agents-features/agents-features-snapshot/build.gradle.kts

@@ -29,7 +29,7 @@ kotlin {
 
         jvmMain {
             dependencies {
-                api(libs.ktor.client.cio)
+                // SQL dependencies moved to agents-features-sql module
             }
         }
 
@@ -38,6 +38,8 @@ kotlin {
                 implementation(kotlin("test-junit5"))
                 implementation(project(":agents:agents-test"))
                 implementation(libs.mockk)
+                implementation(libs.testcontainers)
+                implementation(libs.testcontainers.postgresql)
             }
         }
     }

agents/agents-features/agents-features-sql/Module.md

@@ -0,0 +1,37 @@
+# SQL Features Module
+
+Provides SQL-based persistence providers for agent checkpoints using JetBrains Exposed ORM with support for multiple database engines.
+
+## Features
+
+- **Multi-Database Support**: PostgreSQL, MySQL, H2, SQLite
+- **Exposed ORM Integration**: Type-safe SQL operations with Kotlin DSL
+- **Connection Pooling**: HikariCP integration for production environments
+- **TTL Support**: Automatic cleanup of expired checkpoints with configurable intervals
+- **Database-Specific Optimizations**: Vendor-specific performance tuning
+
+## Dependencies
+
+- `org.jetbrains.exposed:*` - JetBrains Exposed ORM framework
+- `com.zaxxer:HikariCP` - Connection pooling
+- Database drivers: PostgreSQL, MySQL, H2, SQLite
+- `org.testcontainers:postgresql` - Testing infrastructure
+
+## Providers
+
+### ExposedPersistencyStorageProvider
+Base provider using Exposed ORM with configurable cleanup behavior.
+
+### PostgresPersistencyStorageProvider  
+Production-ready provider with JSONB support and HikariCP pooling.
+
+### MySQLPersistencyStorageProvider
+Enterprise provider with JSON column support (MySQL 5.7+).
+
+### H2PersistencyStorageProvider
+Perfect for testing and embedded applications.
+
+### SQLitePersistencyStorageProvider
+Zero-configuration provider for desktop and mobile applications.
+
+All providers implement `AutoCloseable` for proper resource management and support configurable TTL cleanup.

agents/agents-features/agents-features-sql/build.gradle.kts

@@ -0,0 +1,61 @@
+import ai.koog.gradle.publish.maven.Publishing.publishToMaven
+
+group = rootProject.group
+version = rootProject.version
+
+plugins {
+    id("ai.kotlin.multiplatform")
+    alias(libs.plugins.kotlin.serialization)
+}
+
+kotlin {
+    sourceSets {
+        commonMain {
+            dependencies {
+                api(project(":agents:agents-core"))
+                api(project(":agents:agents-features:agents-features-snapshot"))
+                api(project(":rag:rag-base"))
+
+                api(libs.kotlinx.serialization.json)
+                api(libs.ktor.serialization.kotlinx.json)
+            }
+        }
+
+        commonTest {
+            dependencies {
+                implementation(kotlin("test"))
+                implementation(libs.kotlinx.coroutines.test)
+            }
+        }
+
+        jvmMain {
+            dependencies {
+                api(libs.exposed.core)
+                api(libs.exposed.dao)
+                api(libs.exposed.jdbc)
+                api(libs.exposed.json)
+                api(libs.exposed.kotlin.datetime)
+                api(libs.postgresql)
+                api(libs.mysql)
+                api(libs.h2)
+                api(libs.sqlite)
+                implementation(libs.hikaricp)
+            }
+        }
+
+        jvmTest {
+            dependencies {
+                implementation(kotlin("test-junit5"))
+                implementation(project(":agents:agents-test"))
+                implementation(libs.mockk)
+                implementation(libs.testcontainers)
+                implementation(libs.testcontainers.postgresql)
+                implementation(libs.testcontainers.mysql)
+            }
+        }
+    }
+
+    explicitApi()
+}
+
+publishToMaven()

agents/agents-features/agents-features-sql/src/commonMain/kotlin/ai/koog/agents/features/sql/providers/SQLPersistenceSchemaMigrator.kt

@@ -0,0 +1,34 @@
+package ai.koog.agents.features.sql.providers
+
+/**
+ * Represents an interface to handle database schema migrations using Exposed SQL library.
+ *
+ * This interface is designed to be implemented by classes that need to handle schema migrations
+ * for databases, ensuring compatibility and flexibility in schema evolution.
+ *
+ * ExposedSQLMigrator provides an abstraction for executing migrations asynchronously, allowing
+ * for better control and management of database schema changes as part of the application's lifecycle.
+ *
+ * The primary function to be implemented is [migrate], which encapsulates the details of performing
+ * the required schema update or adjustments based on the application's requirements.
+ */
+public interface SQLPersistenceSchemaMigrator {
+    /**
+     * Performs a database schema migration asynchronously.
+     */
+    public suspend fun migrate()
+}
+
+/**
+ * A no-operation implementation of the [SQLPersistenceSchemaMigrator] interface.
+ *
+ * This class is designed to be used in scenarios where schema migrations are not required
+ * or are managed externally. It provides an empty implementation of the [migrate] method,
+ * effectively acting as a placeholder.
+ *
+ * Use this class if you need to satisfy the dependency on [SQLPersistenceSchemaMigrator] but do not
+ * want to perform any migration actions.
+ */
+public object NoOpSQLPersistenceSchemaMigrator : SQLPersistenceSchemaMigrator {
+    override suspend fun migrate() { }
+}

agents/agents-features/agents-features-sql/src/commonMain/kotlin/ai/koog/agents/features/sql/providers/SQLPersistencyStorageProvider.kt

@@ -0,0 +1,91 @@
+package ai.koog.agents.features.sql.providers
+
+import ai.koog.agents.snapshot.providers.PersistencyStorageProvider
+import kotlinx.datetime.Instant
+import kotlinx.serialization.json.Json
+
+/**
+ * Abstract base class for SQL-based implementations of [PersistencyStorageProvider].
+ *
+ * This provider offers a generic SQL abstraction for persisting agent checkpoints
+ * to relational databases. Concrete implementations should handle specific SQL
+ * dialects and connection management.
+ *
+ * ## Storage Schema:
+ * Implementations should create a table with the following structure:
+ * - persistence_id: String (part of primary key)
+ * - checkpoint_id: String (part of primary key)
+ * - created_at: Long (epoch milliseconds)
+ * - checkpoint_json: String (JSON-serialized checkpoint data)
+ * - ttl_timestamp: Long? (optional expiration timestamp)
+ *
+ * ## Design Decisions:
+ * - Uses JSON serialization for checkpoint storage (leveraging database JSON support where available)
+ * - Composite key on (persistence_id, checkpoint_id) ensures uniqueness
+ * - Timestamp stored as epoch milliseconds for cross-database compatibility
+ * - TTL is implemented via a nullable ttl_timestamp column for query-based cleanup
+ *
+ * ## Thread Safety:
+ * Implementations must ensure thread-safe database access, typically through connection pooling.
+ *
+ * @constructor Initializes the SQL persistence provider.
+ * @param persistenceId Unique identifier for this agent's persistence data
+ * @param tableName Name of the table to store checkpoints (default: "agent_checkpoints")
+ * @param ttlSeconds Optional TTL for checkpoint entries in seconds (null = no expiration)
+ */
+public abstract class SQLPersistencyStorageProvider(
+    protected val persistenceId: String,
+    protected val tableName: String = "agent_checkpoints",
+    protected val ttlSeconds: Long? = null,
+    protected val migrator: SQLPersistenceSchemaMigrator
+) : PersistencyStorageProvider {
+
+    protected val json: Json = Json {
+        prettyPrint = true
+        ignoreUnknownKeys = true
+    }
+
+    /**
+     * Initializes the database schema if it doesn't exist.
+     * This should be called once during provider initialization.
+     */
+    public open suspend fun migrate() {
+        migrator.migrate()
+    }
+
+    /**
+     * Executes a database transaction with the given operations.
+     * Implementations should ensure proper transaction isolation and rollback on failure.
+     */
+    protected abstract suspend fun <T> transaction(block: suspend () -> T): T
+
+    /**
+     * Cleans up expired checkpoints based on TTL.
+     * This should be called periodically or before operations to maintain database hygiene.
+     */
+    public abstract suspend fun cleanupExpired()
+
+    /**
+     * Calculates the TTL timestamp for a checkpoint if TTL is configured.
+     */
+    public fun calculateTtlTimestamp(timestamp: Instant): Long? {
+        return ttlSeconds?.let {
+            timestamp.toEpochMilliseconds() + (it * 1000)
+        }
+    }
+
+    /**
+     * Deletes a specific checkpoint by ID
+     */
+    public abstract suspend fun deleteCheckpoint(checkpointId: String)
+
+    /**
+     * Deletes all checkpoints for this persistence ID
+     */
+    public abstract suspend fun deleteAllCheckpoints()
+
+    /**
+     * Gets the total number of checkpoints stored
+     */
+    public abstract suspend fun getCheckpointCount(): Long
+}