From 40aaf8c935af7f5c5fe5300d5998d078ebfebb90 Mon Sep 17 00:00:00 2001
From: Philippe Martin <phmartin@redhat.com>
Date: Tue, 16 Jan 2024 17:53:20 +0100
Subject: [PATCH] fix: display query when leaving and going back to the
 Playground page

---
 packages/backend/src/playground.ts            | 13 +++++---
 packages/backend/src/studio-api-impl.ts       |  4 +++
 .../frontend/src/pages/ModelPlayground.svelte | 30 +++++++++++++++----
 .../frontend/src/stores/playground-queries.ts |  2 ++
 packages/shared/StudioAPI.ts                  |  5 ++++
 5 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/packages/backend/src/playground.ts b/packages/backend/src/playground.ts
index 847de2efe..f9872ebe7 100644
--- a/packages/backend/src/playground.ts
+++ b/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()),
+    });
+  }
 }
diff --git a/packages/backend/src/studio-api-impl.ts b/packages/backend/src/studio-api-impl.ts
index 634d8e56f..bd780a9ce 100644
--- a/packages/backend/src/studio-api-impl.ts
+++ b/packages/backend/src/studio-api-impl.ts
@@ -108,4 +108,8 @@ export class StudioApiImpl implements StudioAPI {
     }
     return this.playgroundManager.askPlayground(localModelInfo[0], prompt);
   }
+
+  async askPlaygroundQueries(): Promise<void> {
+    this.playgroundManager.sendState();
+  }
 }
diff --git a/packages/frontend/src/pages/ModelPlayground.svelte b/packages/frontend/src/pages/ModelPlayground.svelte
index efc5afacc..b885dafd9 100644
--- a/packages/frontend/src/pages/ModelPlayground.svelte
+++ b/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>
diff --git a/packages/frontend/src/stores/playground-queries.ts b/packages/frontend/src/stores/playground-queries.ts
index d2533d08b..0a1d0fb7c 100644
--- a/packages/frontend/src/stores/playground-queries.ts
+++ b/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();
 });
diff --git a/packages/shared/StudioAPI.ts b/packages/shared/StudioAPI.ts
index ec23896ae..ae0a40a65 100644
--- a/packages/shared/StudioAPI.ts
+++ b/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>;
 }