From d93fe55ad15670d16dc92c3c2670d5d437697f48 Mon Sep 17 00:00:00 2001
From: Classic298 <27028174+Classic298@users.noreply.github.com>
Date: Fri, 3 Apr 2026 17:50:35 +0200
Subject: [PATCH] fix: guard non-dict JSON, incremental flushing, offline txn
 parity

---
 backend/open_webui/migrations/env.py          |  1 +
 .../8452d01d26d7_add_chat_message_table.py    | 26 ++++++++++++++-----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/backend/open_webui/migrations/env.py b/backend/open_webui/migrations/env.py
index f971001569..23e163de69 100644
--- a/backend/open_webui/migrations/env.py
+++ b/backend/open_webui/migrations/env.py
@@ -57,6 +57,7 @@ def run_migrations_offline() -> None:
         target_metadata=target_metadata,
         literal_binds=True,
         dialect_opts={'paramstyle': 'named'},
+        transaction_per_migration=True,
     )
 
     with context.begin_transaction():
diff --git a/backend/open_webui/migrations/versions/8452d01d26d7_add_chat_message_table.py b/backend/open_webui/migrations/versions/8452d01d26d7_add_chat_message_table.py
index 20001db092..ffe94c43a4 100644
--- a/backend/open_webui/migrations/versions/8452d01d26d7_add_chat_message_table.py
+++ b/backend/open_webui/migrations/versions/8452d01d26d7_add_chat_message_table.py
@@ -38,6 +38,9 @@ def _parse_chat_messages(chat_id, user_id, chat_data, now):
         except Exception:
             return []
 
+    if not isinstance(chat_data, dict):
+        return []
+
     history = chat_data.get('history', {})
     if not isinstance(history, dict):
         return []
@@ -367,23 +370,34 @@ def _upgrade_postgresql() -> None:
         last_id = rows[-1][0]
         total_chats += len(rows)
 
-        # Parse all messages from this page of chats
+        # Parse and flush incrementally to avoid memory spikes from
+        # message-dense chats (some chats can have thousands of messages).
         messages_batch = []
         for chat_row in rows:
             messages_batch.extend(
                 _parse_chat_messages(chat_row[0], chat_row[1], chat_row[2], now)
             )
+            # Flush when batch is full
+            while len(messages_batch) >= BATCH_SIZE:
+                batch = messages_batch[:BATCH_SIZE]
+                messages_batch = messages_batch[BATCH_SIZE:]
+                inserted, failed = _flush_batch_pg(
+                    conn, chat_message_table, batch
+                )
+                total_inserted += inserted
+                total_failed += failed
 
-        # Insert in sub-batches with ON CONFLICT DO NOTHING
-        for i in range(0, len(messages_batch), BATCH_SIZE):
-            batch = messages_batch[i : i + BATCH_SIZE]
-            inserted, failed = _flush_batch_pg(conn, chat_message_table, batch)
+        # Flush remaining messages from this page
+        if messages_batch:
+            inserted, failed = _flush_batch_pg(
+                conn, chat_message_table, messages_batch
+            )
             total_inserted += inserted
             total_failed += failed
 
         conn.commit()
 
-        if total_inserted % 50000 < max(len(messages_batch), 1):
+        if total_inserted % 50000 < BATCH_SIZE:
             log.info(
                 f'Migration progress: {total_chats} chats processed,'
                 f' {total_inserted} messages inserted...'