diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 9a7f95df87..44b8507906 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/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): Promise; + awaitFirstSave(threadId: string | undefined): Promise; +}; + +const pendingFirstThreadSaves = new Map>(); + +function wait(ms: number): Promise { + 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 { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - async function updateStoredChatThreadEventually( threadId: string, patch: Parameters[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) { diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 979d28a952..5e5780ec2a 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/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>()); + + const reportAutosaveError = useCallback((error: unknown): void => { + if (!isExpectedBackgroundChatStorageError(error)) { + console.error("Failed to autosave chat thread", error); + } + }, []); const saveThread = useCallback( async (threadId: string): Promise => { @@ -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); });