fix: stream data export to prevent 504 timeouts (BasedHardware#5117)

mdmohsin7 · web-flow · commit 7742764113d2 · 2026-02-25T15:49:19.000+05:30
## Summary - Converts `/v1/users/export` endpoint from sync to async `StreamingResponse`, fetching conversations in batches of 200 using `get_conversations_without_photos` (no photo loading) and streaming JSON incrementally — bypasses the 120s TimeoutMiddleware - Updates the Next.js API proxy to pass `Content-Disposition` responses through as a `ReadableStream` instead of buffering with `response.json()` - Changes frontend `exportAllData()` to use `response.blob()` instead of parsing the full JSON into memory - Adds an export-in-progress banner ("Exporting your data... Please don't close this tab") visible at the top of the settings page during export ## Test plan - [x] Export data for a user with 1400+ conversations — should complete without 504 - [x] Verify the downloaded `omi-export.json` is valid JSON with profile, conversations, memories, people, action_items, chat_messages - [ ] Verify export for a user with 0 conversations returns valid JSON with empty arrays - [x] Verify the banner appears during export and disappears on completion - [ ] Verify error toast appears if export fails 🤖 Generated with [Claude Code](https://claude.com/claude-code)
diff --git a/backend/database/chat.py b/backend/database/chat.py
@@ -240,6 +240,25 @@ def get_messages(
     return messages
 
 
+def iter_all_messages(uid: str, batch_size: int = 1000):
+    """Yield all chat messages for a user, decrypted, in batches. Used for streaming data export."""
+    user_ref = db.collection('users').document(uid)
+    msgs_ref = user_ref.collection('messages').order_by('created_at', direction=firestore.Query.DESCENDING)
+    offset = 0
+    while True:
+        batch_ref = msgs_ref.limit(batch_size).offset(offset)
+        batch = []
+        for doc in batch_ref.stream():
+            msg = doc.to_dict()
+            msg['id'] = doc.id
+            msg = _prepare_message_for_read(msg, uid) or msg
+            batch.append(msg)
+        yield from batch
+        if len(batch) < batch_size:
+            break
+        offset += batch_size
+
+
 def get_message(uid: str, message_id: str) -> tuple[Message, str] | None:
     user_ref = db.collection('users').document(uid)
     message_ref = user_ref.collection('messages').where('id', '==', message_id).limit(1).stream()
diff --git a/backend/database/conversations.py b/backend/database/conversations.py
@@ -259,6 +259,26 @@ def get_conversations_without_photos(
     return conversations
 
 
+def iter_all_conversations(uid: str, batch_size: int = 400, include_discarded: bool = True):
+    """Yield all conversations for a user, decrypted, in batches. Used for streaming data export."""
+    conversations_ref = db.collection('users').document(uid).collection(conversations_collection)
+    if not include_discarded:
+        conversations_ref = conversations_ref.where(filter=FieldFilter('discarded', '==', False))
+    conversations_ref = conversations_ref.order_by('created_at', direction=firestore.Query.DESCENDING)
+    offset = 0
+    while True:
+        batch_ref = conversations_ref.limit(batch_size).offset(offset)
+        batch = []
+        for doc in batch_ref.stream():
+            conv = doc.to_dict()
+            conv = _prepare_conversation_for_read(conv, uid) or conv
+            batch.append(conv)
+        yield from batch
+        if len(batch) < batch_size:
+            break
+        offset += batch_size
+
+
 def update_conversation(uid: str, conversation_id: str, update_data: dict):
     doc_ref = db.collection('users').document(uid).collection(conversations_collection).document(conversation_id)
     doc_snapshot = doc_ref.get()
diff --git a/backend/routers/users.py b/backend/routers/users.py
@@ -1,3 +1,4 @@
+import json
 import threading
 import uuid
 from typing import List, Dict, Any, Union, Optional
@@ -6,6 +7,7 @@
 
 import pytz
 from fastapi import APIRouter, Depends, HTTPException, Query
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 
 from database import (
@@ -1137,38 +1139,54 @@ def get_llm_top_features(
     return llm_usage_db.get_top_features(uid, days=days, limit=limit)
 
 
-@router.get('/v1/users/export', tags=['v1'])
-def export_all_user_data(uid: str = Depends(auth.get_current_user_uid)):
-    """Export all user data for GDPR/CCPA compliance."""
-    profile = get_user_profile(uid)
+def _json_default(obj):
+    if isinstance(obj, datetime):
+        return obj.isoformat()
+    raise TypeError(f"Type {type(obj)} not serializable")
 
-    conversations = conversations_db.get_conversations(uid, limit=10000, offset=0, include_discarded=True)
 
-    memories_list = memories_db.get_memories(uid, limit=10000, offset=0)
-
-    people = get_people(uid)
-
-    action_items = get_standalone_action_items(uid, limit=10000, offset=0)
-
-    # Get chat messages (all, no app_id filter)
-    chat_messages = []
-    try:
-        user_ref = chat_db.db.collection('users').document(uid)
-        msgs_ref = user_ref.collection('messages').order_by(
-            'created_at', direction=cloud_firestore.Query.DESCENDING
-        ).limit(10000)
-        for doc in msgs_ref.stream():
-            msg = doc.to_dict()
-            msg['id'] = doc.id
-            chat_messages.append(msg)
-    except Exception as e:
-        logger.warning(f'export_all_user_data: failed to fetch chat messages: {e}')
-
-    return {
-        'profile': profile if profile else {},
-        'conversations': conversations,
-        'memories': memories_list,
-        'people': people,
-        'action_items': action_items,
-        'chat_messages': chat_messages,
-    }
+@router.get('/v1/users/export', tags=['v1'])
+async def export_all_user_data(uid: str = Depends(auth.get_current_user_uid)):
+    """Export all user data for GDPR/CCPA compliance. Streams response to avoid timeouts."""
+
+    def generate():
+        profile = get_user_profile(uid)
+        memories_list = memories_db.get_memories(uid, limit=10000, offset=0)
+        people = get_people(uid)
+        action_items = get_standalone_action_items(uid, limit=10000, offset=0)
+
+        # Stream pretty-printed JSON, yielding conversations and messages one at a time
+        yield '{\n'
+        yield '  "profile": ' + json.dumps(profile if profile else {}, default=_json_default, indent=2) + ',\n'
+
+        # Stream conversations via generator (batched internally, never all in memory)
+        yield '  "conversations": [\n'
+        first = True
+        for conv in conversations_db.iter_all_conversations(uid, include_discarded=True):
+            if not first:
+                yield ',\n'
+            first = False
+            yield '    ' + json.dumps(conv, default=_json_default, indent=4)
+        yield '\n  ],\n'
+
+        yield '  "memories": ' + json.dumps(memories_list, default=_json_default, indent=2) + ',\n'
+        yield '  "people": ' + json.dumps(people, default=_json_default, indent=2) + ',\n'
+        yield '  "action_items": ' + json.dumps(action_items, default=_json_default, indent=2) + ',\n'
+
+        # Stream chat messages via generator (batched internally, never all in memory)
+        yield '  "chat_messages": [\n'
+        first = True
+        for msg in chat_db.iter_all_messages(uid):
+            if not first:
+                yield ',\n'
+            first = False
+            yield '    ' + json.dumps(msg, default=_json_default, indent=4)
+        yield '\n  ]\n'
+
+        yield '}\n'
+
+    return StreamingResponse(
+        generate(),
+        media_type='application/json',
+        headers={'Content-Disposition': 'attachment; filename="omi-export.json"'},
+    )
diff --git a/web/app/src/app/api/proxy/[...path]/route.ts b/web/app/src/app/api/proxy/[...path]/route.ts
@@ -117,6 +117,18 @@ async function handleRequest(
       });
     }
 
+    // Handle download/streaming responses (e.g., data export) — pass body through without buffering
+    const contentDisposition = response.headers.get('content-disposition');
+    if (contentDisposition) {
+      return new NextResponse(response.body, {
+        status: response.status,
+        headers: {
+          'Content-Type': responseContentType || 'application/octet-stream',
+          'Content-Disposition': contentDisposition,
+        },
+      });
+    }
+
     // Handle JSON responses
     if (responseContentType?.includes('application/json')) {
       const data = await response.json();
diff --git a/web/app/src/components/settings/SettingsPage.tsx b/web/app/src/components/settings/SettingsPage.tsx
@@ -2848,9 +2848,7 @@ export function SettingsPage() {
     if (isExporting) return;
     setIsExporting(true);
     try {
-      const data = await exportAllData();
-      const json = JSON.stringify(data, null, 2);
-      const blob = new Blob([json], { type: 'application/json' });
+      const blob = await exportAllData();
       const url = URL.createObjectURL(blob);
       const a = document.createElement('a');
       a.href = url;
@@ -3029,6 +3027,30 @@ export function SettingsPage() {
 
   return (
     <div className="h-full flex flex-col">
+      {/* Export in-progress dialog */}
+      {isExporting && (
+        <div className="fixed inset-0 z-50 flex items-center justify-center">
+          <div className="absolute inset-0 bg-black/60" />
+          <div className="relative bg-bg-secondary rounded-2xl p-6 max-w-sm w-full mx-4 shadow-2xl border border-white/[0.06]">
+            <div className="flex flex-col items-center text-center gap-4">
+              <div className="p-3 rounded-full bg-purple-500/10">
+                <Loader2 className="w-8 h-8 text-purple-400 animate-spin" />
+              </div>
+              <div>
+                <h3 className="text-lg font-semibold text-text-primary">Exporting Your Data</h3>
+                <p className="text-text-secondary mt-2 text-sm">
+                  This may take a moment depending on the amount of data in your account.
+                </p>
+              </div>
+              <div className="flex items-center gap-2 bg-yellow-500/10 rounded-xl px-4 py-2">
+                <AlertTriangle className="w-4 h-4 text-yellow-400 flex-shrink-0" />
+                <span className="text-xs text-yellow-400">Please don&apos;t close this tab</span>
+              </div>
+            </div>
+          </div>
+        </div>
+      )}
+
       {/* Page Header */}
       <PageHeader title={sectionInfo.title} icon={Settings} showBackButton />
 
diff --git a/web/app/src/lib/api.ts b/web/app/src/lib/api.ts
@@ -1741,10 +1741,20 @@ export async function deleteMcpApiKey(keyId: string): Promise<void> {
 // ============================================================================
 
 /**
- * Export all conversations as JSON
+ * Export all user data as a downloadable JSON blob (streamed from backend).
  */
-export async function exportAllData(): Promise<Record<string, unknown>> {
-  return fetchWithAuth<Record<string, unknown>>('/v1/users/export');
+export async function exportAllData(): Promise<Blob> {
+  const token = await getIdToken();
+  if (!token) {
+    throw new Error('Not authenticated');
+  }
+  const response = await fetch(`${API_BASE_URL}/v1/users/export`, {
+    headers: { Authorization: `Bearer ${token}` },
+  });
+  if (!response.ok) {
+    throw new Error(`Export failed: ${response.status} ${response.statusText}`);
+  }
+  return response.blob();
 }
 
 /**
