chore(aiagent): store memory in db (#7053)

* store to db

* no warning in context

* warning in frontend

* remove log

* 100kb + alert

* update sqlx

* update eeref

* Update ee-repo-ref.txt

* cleaning

---------

Co-authored-by: windmill-internal-app[bot] <windmill-internal-app[bot]@users.noreply.github.com>

Author: centdix

backend/.sqlx/query-38baa6cf7d1c2532d38486a01e6d27ac1a58540a51e97dec5613e9bc0791e890.json

@@ -1 +1 @@
-20f881d0c328cdb7ce888db8faf8989a4d8e5bad
+5b7afe50da442441747e7a8f6ef461c96faa9dc2

backend/.sqlx/query-6a1005b8ae5326c5d5955534e9228b6bd524253971ae09cbe100c0d6d42e6cdc.json

@@ -0,0 +1,4 @@
+-- Add down migration script here
+
+-- Drop table
+DROP TABLE IF EXISTS ai_agent_memory;

backend/.sqlx/query-d1d9cb0a7043760564171f481dd73f04eaa2c7d5c1ddc0501bc70c3fb34a07be.json

@@ -0,0 +1,16 @@
+-- Add up migration script here
+
+-- Create ai_agent_memory table for storing AI agent step memory when S3 is unavailable
+CREATE TABLE ai_agent_memory (
+    workspace_id VARCHAR(50) NOT NULL,
+    conversation_id UUID NOT NULL,
+    step_id VARCHAR(255) NOT NULL,
+    messages JSONB NOT NULL,
+    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
+    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
+    PRIMARY KEY (workspace_id, conversation_id, step_id)
+);
+
+-- Grant permissions
+GRANT ALL ON ai_agent_memory TO windmill_admin;
+GRANT ALL ON ai_agent_memory TO windmill_user;

backend/ee-repo-ref.txt

@@ -11,7 +11,7 @@ use uuid::Uuid;
 
 use crate::db::ApiAuthed;
 use windmill_common::{
-    db::UserDB,
+    db::{UserDB, DB},
     error::{JsonResult, Result},
     flow_conversations::MessageType,
     utils::{not_found_if_none, paginate, Pagination},
@@ -154,6 +154,7 @@ pub async fn get_or_create_conversation_with_id(
 async fn delete_conversation(
     authed: ApiAuthed,
     Extension(user_db): Extension<UserDB>,
+    Extension(db): Extension<DB>,
     Path((w_id, conversation_id)): Path<(String, Uuid)>,
 ) -> Result<String> {
     let mut tx = user_db.clone().begin(&authed).await?;
@@ -185,10 +186,14 @@ async fn delete_conversation(
 
     // Delete associated memory in background (non-blocking cleanup)
     let w_id_clone = w_id.clone();
+    let db_clone = db.clone();
     tokio::spawn(async move {
-        if let Err(e) =
-            windmill_worker::memory_oss::delete_conversation_memory(&w_id_clone, conversation_id)
-                .await
+        if let Err(e) = windmill_worker::memory_oss::delete_conversation_memory(
+            &db_clone,
+            &w_id_clone,
+            conversation_id,
+        )
+        .await
         {
             tracing::error!(
                 "Failed to delete memory for conversation {} in workspace {}: {:?}",

backend/migrations/20251104135208_ai_agent_memory.down.sql

@@ -498,7 +498,6 @@ pub async fn store_pull_query(wc: &WorkerConfig) {
 
 pub const TMP_DIR: &str = "/tmp/windmill";
 pub const TMP_LOGS_DIR: &str = concatcp!(TMP_DIR, "/logs");
-pub const TMP_MEMORY_DIR: &str = concatcp!(TMP_DIR, "/memory");
 
 pub const HUB_CACHE_DIR: &str = concatcp!(ROOT_CACHE_DIR, "hub");
 

backend/migrations/20251104135208_ai_agent_memory.up.sql

@@ -405,7 +405,7 @@ pub async fn run_agent(
                     .and_then(|fs| fs.memory_id)
                 {
                     // Read messages from memory
-                    match read_from_memory(&job.workspace_id, memory_id, step_id).await {
+                    match read_from_memory(db, &job.workspace_id, memory_id, step_id).await {
                         Ok(Some(loaded_messages)) => {
                             // Take the last n messages
                             let start_idx = loaded_messages.len().saturating_sub(context_length);
@@ -856,6 +856,7 @@ pub async fn run_agent(
 
                     if let Some(memory_id) = flow_context.flow_status.and_then(|fs| fs.memory_id) {
                         if let Err(e) = write_to_memory(
+                            db,
                             &job.workspace_id,
                             memory_id,
                             step_id,

backend/windmill-api/src/flow_conversations.rs

@@ -1,72 +1,123 @@
 use crate::ai::types::OpenAIMessage;
-use std::path::PathBuf;
-use tokio::{fs, io::AsyncWriteExt};
 use uuid::Uuid;
-use windmill_common::worker::TMP_MEMORY_DIR;
-
-/// Get the file path for storing memory for a specific AI agent step
-pub fn path_for(workspace_id: &str, conversation_id: Uuid, step_id: &str) -> PathBuf {
-    PathBuf::from(TMP_MEMORY_DIR)
-        .join(workspace_id)
-        .join(conversation_id.to_string())
-        .join(format!("{step_id}.json"))
-}
+use windmill_common::{db::DB, error::Error};
+
+pub const MAX_MEMORY_SIZE_BYTES: usize = 100_000; // 100KB per memory entry in database
 
-/// Read messages from disk storage
-pub async fn read_from_disk(
+/// Read AI agent memory from database
+pub async fn read_from_db(
+    db: &DB,
     workspace_id: &str,
     conversation_id: Uuid,
     step_id: &str,
-) -> anyhow::Result<Option<Vec<OpenAIMessage>>> {
-    let path = path_for(workspace_id, conversation_id, step_id);
-    if !fs::try_exists(&path).await? {
-        return Ok(None);
-    }
+) -> Result<Option<Vec<OpenAIMessage>>, Error> {
+    let result = sqlx::query!(
+        "SELECT messages FROM ai_agent_memory
+         WHERE workspace_id = $1 AND conversation_id = $2 AND step_id = $3",
+        workspace_id,
+        conversation_id,
+        step_id
+    )
+    .fetch_optional(db)
+    .await?;
 
-    let bytes = fs::read(&path).await?;
-    let messages: Vec<OpenAIMessage> = serde_json::from_slice(&bytes)?;
-    Ok(Some(messages))
+    match result {
+        Some(row) => {
+            let messages: Vec<OpenAIMessage> = serde_json::from_value(row.messages)?;
+            Ok(Some(messages))
+        }
+        None => Ok(None),
+    }
 }
 
-/// Write messages to disk storage
-pub async fn write_to_disk(
+/// Write AI agent memory to database with size checking and truncation
+pub async fn write_to_db(
+    db: &DB,
     workspace_id: &str,
     conversation_id: Uuid,
     step_id: &str,
     messages: &[OpenAIMessage],
-) -> anyhow::Result<()> {
-    let path = path_for(workspace_id, conversation_id, step_id);
-
-    // Ensure parent directories exist
-    if let Some(dir) = path.parent() {
-        fs::create_dir_all(dir).await?;
+) -> Result<(), Error> {
+    if messages.is_empty() {
+        return Ok(());
     }
 
-    // Write atomically using a temporary file
-    let tmp = path.with_extension("json.tmp");
-    let mut f = fs::File::create(&tmp).await?;
-    f.write_all(&serde_json::to_vec(messages)?).await?;
-    f.flush().await?;
-    drop(f);
+    // Serialize messages and check size
+    let mut messages_to_store = messages.to_vec();
+    let mut json_value = serde_json::to_value(&messages_to_store)?;
+    let size_bytes = json_value.to_string().len();
+
+    // Truncate if necessary
+    if size_bytes > MAX_MEMORY_SIZE_BYTES {
+        tracing::warn!(
+            "Memory size ({} bytes) exceeds limit ({} bytes) for workspace={} conversation={} step={}. Truncating messages. Use S3 storage in workspace settings to store full conversation history.",
+            size_bytes,
+            MAX_MEMORY_SIZE_BYTES,
+            workspace_id,
+            conversation_id,
+            step_id
+        );
+
+        messages_to_store = truncate_messages(messages, MAX_MEMORY_SIZE_BYTES)?;
+        json_value = serde_json::to_value(&messages_to_store)?;
+    }
 
-    // Atomic rename
-    fs::rename(tmp, &path).await?;
+    // Insert or update using UPSERT
+    sqlx::query!(
+        "INSERT INTO ai_agent_memory (workspace_id, conversation_id, step_id, messages, created_at, updated_at)
+         VALUES ($1, $2, $3, $4, NOW(), NOW())
+         ON CONFLICT (workspace_id, conversation_id, step_id)
+         DO UPDATE SET
+            messages = EXCLUDED.messages,
+            updated_at = NOW()",
+        workspace_id,
+        conversation_id,
+        step_id,
+        json_value
+    )
+    .execute(db)
+    .await?;
 
     Ok(())
 }
 
-/// Delete all memory for a conversation from disk storage
-pub async fn delete_conversation_from_disk(
+/// Delete all memory for a conversation from database
+pub async fn delete_conversation_from_db(
+    db: &DB,
     workspace_id: &str,
     conversation_id: Uuid,
-) -> anyhow::Result<()> {
-    let conversation_path = PathBuf::from(TMP_MEMORY_DIR)
-        .join(workspace_id)
-        .join(conversation_id.to_string());
+) -> Result<(), Error> {
+    sqlx::query!(
+        "DELETE FROM ai_agent_memory
+         WHERE workspace_id = $1 AND conversation_id = $2",
+        workspace_id,
+        conversation_id
+    )
+    .execute(db)
+    .await?;
 
-    if fs::try_exists(&conversation_path).await? {
-        fs::remove_dir_all(&conversation_path).await?;
+    Ok(())
+}
+
+/// Truncate messages to fit within the size limit
+fn truncate_messages(
+    messages: &[OpenAIMessage],
+    max_size_bytes: usize,
+) -> Result<Vec<OpenAIMessage>, Error> {
+    let mut result = messages.to_vec();
+
+    // Keep removing oldest messages until we're under the threshold
+    while !result.is_empty() {
+        let test_json = serde_json::to_value(&result)?;
+        let test_size = test_json.to_string().len();
+
+        if test_size <= max_size_bytes {
+            break;
+        }
+
+        // Remove the first (oldest) message
+        result.remove(0);
     }
 
-    Ok(())
+    Ok(result)
 }