fix: persist Studio thread synchronously on first runStart so mid-stream refresh keeps the prompt (#5814)

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -134,6 +134,47 @@ function isServerSideBuiltinToolPart(
   return hasNativePart;
 }
 
+const FIRST_THREAD_SAVE_TIMEOUT_MS = 250;
+
+type ThreadAutosaveHandle = {
+  registerFirstSave(threadId: string, promise: Promise<void>): Promise<void>;
+  awaitFirstSave(threadId: string | undefined): Promise<void>;
+};
+
+const pendingFirstThreadSaves = new Map<string, Promise<void>>();
+
+function wait(ms: number): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+export const ThreadAutosaveHandle: ThreadAutosaveHandle = {
+  registerFirstSave(threadId, promise) {
+    const trackedPromise = promise.catch(() => {});
+    const cleanupPromise = trackedPromise.finally(() => {
+      if (pendingFirstThreadSaves.get(threadId) === cleanupPromise) {
+        pendingFirstThreadSaves.delete(threadId);
+      }
+    });
+    pendingFirstThreadSaves.set(threadId, cleanupPromise);
+    return cleanupPromise;
+  },
+
+  async awaitFirstSave(threadId) {
+    if (!threadId) {
+      return;
+    }
+    const pending = pendingFirstThreadSaves.get(threadId);
+    if (!pending) {
+      return;
+    }
+    await Promise.race([pending, wait(FIRST_THREAD_SAVE_TIMEOUT_MS)]);
+  },
+};
+
+export function useThreadAutosaveHandle(): ThreadAutosaveHandle {
+  return ThreadAutosaveHandle;
+}
+
 /**
  * Match error messages that indicate the request filled or would fill
  * the KV cache, so the UI can show a dedicated toast pointing at the
@@ -169,10 +210,6 @@ export function isContextLimitError(message: string): boolean {
   );
 }
 
-function wait(ms: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
-
 async function updateStoredChatThreadEventually(
   threadId: string,
   patch: Parameters<typeof updateStoredChatThread>[1],
@@ -2179,6 +2216,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
               throw error;
             }
             clearSelectedImageEditReference();
+            await ThreadAutosaveHandle.awaitFirstSave(resolvedThreadId);
             const stream = streamChatCompletions(requestPayload, abortSignal);
 
             for await (const chunk of stream) {

studio/frontend/src/features/chat/runtime-provider.tsx

@@ -34,7 +34,10 @@ import {
 import { extractText, getDocumentProxy } from "unpdf";
 import { toast } from "sonner";
 import { StudioWebSpeechDictationAdapter } from "./adapters/studio-web-speech-dictation-adapter";
-import { createOpenAIStreamAdapter } from "./api/chat-adapter";
+import {
+  ThreadAutosaveHandle,
+  createOpenAIStreamAdapter,
+} from "./api/chat-adapter";
 import {
   loadConnectionsEnabled,
   loadExternalProviders,
@@ -1132,6 +1135,13 @@ function ThreadBackendAutosave({
 }): ReactElement | null {
   const aui = useAui();
   const saveChainRef = useRef(Promise.resolve());
+  const pendingFirstSavesRef = useRef(new Map<string, Promise<void>>());
+
+  const reportAutosaveError = useCallback((error: unknown): void => {
+    if (!isExpectedBackgroundChatStorageError(error)) {
+      console.error("Failed to autosave chat thread", error);
+    }
+  }, []);
 
   const saveThread = useCallback(
     async (threadId: string): Promise<void> => {
@@ -1173,14 +1183,30 @@ function ThreadBackendAutosave({
     (threadId: string): void => {
       saveChainRef.current = saveChainRef.current
         .catch(() => {})
-        .then(() => saveThread(threadId))
-        .catch((error) => {
-          if (!isExpectedBackgroundChatStorageError(error)) {
-            console.error("Failed to autosave chat thread", error);
-          }
-        });
+        .then(async () => {
+          await pendingFirstSavesRef.current.get(threadId);
+          await saveThread(threadId);
+        })
+        .catch(reportAutosaveError);
     },
-    [saveThread],
+    [reportAutosaveError, saveThread],
+  );
+
+  const saveFirstThreadSnapshot = useCallback(
+    (threadId: string): void => {
+      if (pendingFirstSavesRef.current.has(threadId)) {
+        return;
+      }
+
+      const promise = saveThread(threadId)
+        .catch(reportAutosaveError)
+        .finally(() => {
+          pendingFirstSavesRef.current.delete(threadId);
+        });
+      pendingFirstSavesRef.current.set(threadId, promise);
+      ThreadAutosaveHandle.registerFirstSave(threadId, promise);
+    },
+    [reportAutosaveError, saveThread],
   );
 
   useAuiEvent("thread.runEnd", ({ threadId }) => {
@@ -1188,6 +1214,13 @@ function ThreadBackendAutosave({
   });
 
   useAuiEvent("thread.runStart", ({ threadId }) => {
+    const runtime = aui.threads().__internal_getAssistantRuntime?.();
+    const { remoteId } =
+      runtime?.threads.getItemById(threadId).getState() ?? {};
+    if (!remoteId) {
+      saveFirstThreadSnapshot(threadId);
+      return;
+    }
     queueSave(threadId);
   });