fix: display query when leaving and going back to the Playground page

packages/backend/src/playground.ts

@@ -141,10 +141,7 @@ export class PlayGroundManager {
             }
             q.response = result as ModelResponse;
             this.queries.set(query.id, q);
-            this.webview.postMessage({ 
-              id: MSG_NEW_PLAYGROUND_QUERIES_STATE, 
-              body: Array.from(this.queries.values()),
-            });
+            this.sendState();
             break;
         }
       });
@@ -154,10 +151,18 @@ export class PlayGroundManager {
     post_req.end();
 
     this.queries.set(query.id, query);
+    this.sendState();
     return query.id;
   }
 
   getNextQueryId() {
     return ++this.queryIdCounter;
   }
+
+  sendState() {
+    this.webview.postMessage({ 
+      id: MSG_NEW_PLAYGROUND_QUERIES_STATE, 
+      body: Array.from(this.queries.values()),
+    });
+  }
 }

packages/backend/src/studio-api-impl.ts

@@ -110,4 +110,8 @@ export class StudioApiImpl implements StudioAPI {
     }
     return this.playgroundManager.askPlayground(localModelInfo[0], prompt);
   }
+
+  async askPlaygroundQueries(): Promise<void> {
+    this.playgroundManager.sendState();
+  }
 }

packages/frontend/src/pages/ModelPlayground.svelte

@@ -20,23 +20,43 @@
     studioClient.startPlayground(model.id);
 
     playgroundQueries.subscribe((queries: QueryState[]) => {
-      const myQuery = queries.find(q => q.id === queryId);
+      if (queryId === -1) {
+        return;
+      }
+      let myQuery = queries.find(q => q.id === queryId);
       if (!myQuery) {
-        return;        
+        myQuery = queries.findLast(q => q.modelId === model?.id);
       }
-      inProgress = false;
-      if (myQuery.response?.choices.length) {
-        result = myQuery.response?.choices[0];
+      if (!myQuery) {
+        return;
       }
+      displayQuery(myQuery);
     });
   });
 
+  function displayQuery(query: QueryState) {
+    if (query.response) {
+      inProgress = false;
+      prompt = query.prompt;
+      if (query.response?.choices.length) {
+        result = query.response?.choices[0];
+      }
+    } else {
+      inProgress = true;
+      prompt = query.prompt;
+      queryId = query.id;
+    }
+  }
+
   async function askPlayground() {
     if (!model) {
       return;
     }
     inProgress = true;
     result = undefined;
+    // do not display anything before we get a response from askPlayground
+    // (we can receive a new queryState before the new QueryId)
+    queryId = -1; 
     queryId = await studioClient.askPlayground(model.id, prompt);
   }
 </script>

packages/frontend/src/stores/playground-queries.ts

@@ -3,9 +3,11 @@ import { readable } from 'svelte/store';
 import type { QueryState } from '@shared/models/IPlaygroundQueryState';
 import { rpcBrowser } from '../utils/client';
 import { MSG_NEW_PLAYGROUND_QUERIES_STATE } from '@shared/Messages';
+import { studioClient } from '/@/utils/client';
 
 export const playgroundQueries: Readable<QueryState[]> = readable<QueryState[]>([], (set) => {  
   rpcBrowser.subscribe(MSG_NEW_PLAYGROUND_QUERIES_STATE, (msg) => {
     set(msg);
   });
+  studioClient.askPlaygroundQueries();
 });

packages/shared/StudioAPI.ts

@@ -29,5 +29,10 @@ export abstract class StudioAPI {
    * @param label
    */
   abstract getTasksByLabel(label: string): Promise<Task[]>;
+
+  /**
+   * Ask to send a message MSG_NEW_PLAYGROUND_QUERIES_STATE with the current Playground queries
+   */
+  abstract askPlaygroundQueries(): Promise<void>;
 }