[PRO-1050] Fix solution saving for large problems

app/components/chat/LoadProblemButton.tsx

@@ -13,21 +13,48 @@ interface LoadProblemButtonProps {
 }
 
 export function setLastLoadedProblem(problem: BoltProblem) {
+  const problemSerialized = JSON.stringify(problem);
+
   try {
-    localStorage.setItem('loadedProblem', JSON.stringify(problem));
-  } catch (error) {
-    console.error('Failed to set last loaded problem:', error);
+    localStorage.setItem('loadedProblemId', problem.problemId);
+    localStorage.setItem('loadedProblem', problemSerialized);
+  } catch (error: any) {
+    // Remove loadedProblem, so we don't accidentally associate a solution with the wrong problem.
+    localStorage.removeItem('loadedProblem');
+    console.error(
+      `Failed to set last loaded problem (size=${(problemSerialized.length / 1024).toFixed(2)}kb):`,
+      error.stack || error,
+    );
   }
 }
 
-export function getLastLoadedProblem(): BoltProblem | undefined {
+export async function getOrFetchLastLoadedProblem(): Promise<BoltProblem | null> {
   const problemJSON = localStorage.getItem('loadedProblem');
+  let problem: BoltProblem | null = null;
 
-  if (!problemJSON) {
-    return undefined;
+  if (problemJSON) {
+    problem = JSON.parse(problemJSON);
+  } else {
+    /*
+     * Problem might not have fit into localStorage.
+     * Try to re-load it from server.
+     */
+    const problemId = localStorage.getItem('loadedProblemId');
+
+    if (!problemId) {
+      return null;
+    }
+
+    problem = await getProblem(problemId);
+
+    if (!problem) {
+      return null;
+    }
+
+    setLastLoadedProblem(problem);
   }
 
-  return JSON.parse(problemJSON);
+  return problem;
 }
 
 export async function loadProblem(

app/components/sidebar/SaveSolution.tsx

@@ -1,11 +1,15 @@
 import { toast } from 'react-toastify';
 import ReactModal from 'react-modal';
 import { useState } from 'react';
-import { workbenchStore } from '~/lib/stores/workbench';
 import { BoltProblemStatus, updateProblem } from '~/lib/replay/Problems';
-import type { BoltProblemInput } from '~/lib/replay/Problems';
-import { getLastLoadedProblem } from '../chat/LoadProblemButton';
-import { getLastUserSimulationData, getLastSimulationChatMessages } from '~/lib/replay/SimulationPrompt';
+import type { BoltProblem, BoltProblemInput } from '~/lib/replay/Problems';
+import { getOrFetchLastLoadedProblem } from '~/components/chat/LoadProblemButton';
+import {
+  getLastUserSimulationData,
+  getLastSimulationChatMessages,
+  getSimulationRecordingId,
+  isSimulatingOrHasFinished,
+} from '~/lib/replay/SimulationPrompt';
 
 ReactModal.setAppElement('#root');
 
@@ -20,13 +24,28 @@ export function SaveSolution() {
     evaluator: '',
   });
   const [savedSolution, setSavedSolution] = useState<boolean>(false);
+  const [problem, setProblem] = useState<BoltProblem | null>(null);
 
-  const handleSaveSolution = () => {
-    setIsModalOpen(true);
-    setFormData({
-      evaluator: '',
-    });
+  const handleSaveSolution = async () => {
+    const loadId = toast.loading('Loading problem...');
+
+    try {
+      const savedProblem = await getOrFetchLastLoadedProblem();
+
+      if (!savedProblem) {
+        toast.error('No problem loaded');
+        return;
+      }
+
+      setProblem(savedProblem);
+      setFormData({
+        evaluator: savedProblem.solution?.evaluator || '',
+      });
+    } finally {
+      toast.dismiss(loadId);
+    }
     setSavedSolution(false);
+    setIsModalOpen(true);
   };
 
   const handleInputChange = (e: React.ChangeEvent<HTMLInputElement>) => {
@@ -38,55 +57,82 @@ export function SaveSolution() {
   };
 
   const handleSubmitSolution = async () => {
-    const savedProblem = getLastLoadedProblem();
-
-    if (!savedProblem) {
+    if (!problem) {
       toast.error('No problem loaded');
       return;
     }
 
-    const simulationData = getLastUserSimulationData();
-
-    if (!simulationData) {
-      toast.error('No simulation data found');
+    if (!isSimulatingOrHasFinished()) {
+      toast.error('No simulation found (neither in progress nor finished)');
       return;
     }
 
-    const messages = getLastSimulationChatMessages();
+    try {
+      const loadId = toast.loading('Waiting for recording...');
 
-    if (!messages) {
-      toast.error('No user prompt found');
-      return;
+      try {
+        /*
+         * Wait for simulation to finish.
+         * const recordingId =
+         */
+        await getSimulationRecordingId();
+      } finally {
+        toast.dismiss(loadId);
+      }
+
+      toast.info('Submitting solution...');
+
+      console.log('SubmitSolution', formData);
+
+      const simulationData = getLastUserSimulationData();
+
+      if (!simulationData) {
+        toast.error('No simulation data found');
+        return;
+      }
+
+      const messages = getLastSimulationChatMessages();
+
+      if (!messages) {
+        toast.error('No user prompt found');
+        return;
+      }
+
+      /*
+       * The evaluator is only present when the problem has been solved.
+       * We still create a "solution" object even if it hasn't been
+       * solved quite yet, which is used for working on the problem.
+       *
+       * TODO: Split `solution` into `reproData` and `evaluator`.
+       */
+      const evaluator = formData.evaluator.length ? formData.evaluator : undefined;
+
+      const problemUpdatePacket: BoltProblemInput = {
+        version: 2,
+        title: problem.title,
+        description: problem.description,
+        username: problem.username,
+        repositoryContents: problem.repositoryContents,
+        status: evaluator ? BoltProblemStatus.Solved : BoltProblemStatus.Unsolved,
+        solution: {
+          simulationData,
+          messages,
+          evaluator,
+
+          /*
+           * TODO: Also store recordingId for easier debugging.
+           * recordingId,
+           */
+        },
+      };
+
+      await updateProblem(problem.problemId, problemUpdatePacket);
+
+      setSavedSolution(true);
+    } catch (error: any) {
+      console.error('Error saving solution', error?.stack || error);
+      toast.error(`Error saving solution: ${error?.message}`);
     }
-
-    toast.info('Submitting solution...');
-
-    console.log('SubmitSolution', formData);
-
-    /*
-     * The evaluator is only present when the problem has been solved.
-     * We still create a "solution" object even if it hasn't been
-     * solved quite yet, which is used for working on the problem.
-     */
-    const evaluator = formData.evaluator.length ? formData.evaluator : undefined;
-
-    const problem: BoltProblemInput = {
-      version: 2,
-      title: savedProblem.title,
-      description: savedProblem.description,
-      username: savedProblem.username,
-      repositoryContents: savedProblem.repositoryContents,
-      status: evaluator ? BoltProblemStatus.Solved : BoltProblemStatus.Unsolved,
-      solution: {
-        simulationData,
-        messages,
-        evaluator,
-      },
-    };
-
-    await updateProblem(savedProblem.problemId, problem);
-
-    setSavedSolution(true);
   };
 
   return (

app/lib/persistence/db.ts

@@ -208,6 +208,8 @@ export async function createChatFromMessages(
   const newId = await getNextId(db);
   const newUrlId = await getUrlId(db, newId); // Get a new urlId for the duplicated chat
 
+  // TODO: Call setLastLoadedProblem(null).
+
   await setMessages(
     db,
     newId,

app/lib/replay/SimulationPrompt.ts

@@ -8,9 +8,9 @@ import type { SimulationData, SimulationPacket } from './SimulationData';
 import { SimulationDataVersion } from './SimulationData';
 import { assert, generateRandomId, ProtocolClient } from './ReplayProtocolClient';
 import type { MouseData } from './Recording';
-import type { FileMap } from '../stores/files';
+import type { FileMap } from '~/lib/stores/files';
 import { shouldIncludeFile } from '~/utils/fileUtils';
-import { DeveloperSystemPrompt } from '../common/prompts/prompts';
+import { DeveloperSystemPrompt } from '~/lib/common/prompts/prompts';
 import { detectProjectCommands } from '~/utils/projectCommands';
 
 function createRepositoryContentsPacket(contents: string): SimulationPacket {
@@ -85,6 +85,10 @@ class ChatManager {
     })();
   }
 
+  isValid() {
+    return !!this.client;
+  }
+
   destroy() {
     this.client?.close();
     this.client = undefined;
@@ -294,6 +298,17 @@ export async function getSimulationRecording(): Promise<string> {
   return gChatManager.recordingIdPromise;
 }
 
+export function isSimulatingOrHasFinished(): boolean {
+  return gChatManager?.isValid() ?? false;
+}
+
+export async function getSimulationRecordingId(): Promise<string> {
+  assert(gChatManager, 'Chat not started');
+  assert(gChatManager.recordingIdPromise, 'Expected recording promise');
+
+  return gChatManager.recordingIdPromise;
+}
+
 let gLastSimulationChatMessages: ProtocolMessage[] | undefined;
 
 export function getLastSimulationChatMessages(): ProtocolMessage[] | undefined {