From 9042bd22fb12eb64fa50ae101870514895cdaec6 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Wed, 31 Jan 2024 19:53:53 -0900
Subject: [PATCH 1/3] Prevent refreshing in parallel

While adding an automatic refresh I discovered it was easy to
accidentally cause two refreshes to go off at the same time which causes
some races.
---
 src/workspacesProvider.ts | 61 +++++++++++++++++++++++++++++----------
 1 file changed, 45 insertions(+), 16 deletions(-)

diff --git a/src/workspacesProvider.ts b/src/workspacesProvider.ts
index f485b9d5..b066dd8c 100644
--- a/src/workspacesProvider.ts
+++ b/src/workspacesProvider.ts
@@ -14,6 +14,7 @@ export enum WorkspaceQuery {
 export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeItem> {
   private workspaces: WorkspaceTreeItem[] = []
   private agentWatchers: Record<WorkspaceAgent["id"], { dispose: () => void; metadata?: AgentMetadataEvent[] }> = {}
+  private fetching = false
 
   constructor(
     private readonly getWorkspacesQuery: WorkspaceQuery,
@@ -22,26 +23,54 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
     this.fetchAndRefresh()
   }
 
-  // fetchAndRefrehsh fetches new workspaces then re-renders the entire tree.
+  // fetchAndRefresh fetches new workspaces then re-renders the entire tree.
+  // Trying to call this while already refreshing is a no-op and will return
+  // immediately.
   async fetchAndRefresh() {
-    const token = await this.storage.getSessionToken()
-    const workspacesTreeItem: WorkspaceTreeItem[] = []
+    if (this.fetching) {
+      return
+    }
+    this.fetching = true
+
     Object.values(this.agentWatchers).forEach((watcher) => watcher.dispose())
-    // If the URL is set then we are logged in.
-    if (this.storage.getURL()) {
-      const resp = await getWorkspaces({ q: this.getWorkspacesQuery })
-      resp.workspaces.forEach((workspace) => {
-        const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine
-        if (showMetadata && token) {
-          const agents = extractAgents(workspace)
-          agents.forEach((agent) => this.monitorMetadata(agent.id, token)) // monitor metadata for all agents
-        }
-        const treeItem = new WorkspaceTreeItem(workspace, this.getWorkspacesQuery === WorkspaceQuery.All, showMetadata)
-        workspacesTreeItem.push(treeItem)
-      })
+
+    try {
+      this.workspaces = await this.fetch()
+    } catch (error) {
+      this.workspaces = []
     }
-    this.workspaces = workspacesTreeItem
+
     this.refresh()
+    this.fetching = false
+  }
+
+  /**
+   * Fetch workspaces and turn them into tree items.  Throw an error if not
+   * logged in or the query fails.
+   */
+  async fetch(): Promise<WorkspaceTreeItem[]> {
+    // Assume that no URL or no token means we are not logged in.
+    const url = this.storage.getURL()
+    const token = await this.storage.getSessionToken()
+    if (!url || !token) {
+      throw new Error("not logged in")
+    }
+
+    const resp = await getWorkspaces({ q: this.getWorkspacesQuery })
+
+    // We could have logged out while waiting for the query.
+    if (!url || !token) {
+      throw new Error("not logged in")
+    }
+
+    return resp.workspaces.map((workspace) => {
+      const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine
+      if (showMetadata) {
+        const agents = extractAgents(workspace)
+        agents.forEach((agent) => this.monitorMetadata(agent.id, token)) // monitor metadata for all agents
+      }
+      return new WorkspaceTreeItem(workspace, this.getWorkspacesQuery === WorkspaceQuery.All, showMetadata)
+    })
   }
 
   private _onDidChangeTreeData: vscode.EventEmitter<vscode.TreeItem | undefined | null | void> =

From b4f046b74db57c35a01ceb98f6cb0ba4b7d628d2 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Wed, 31 Jan 2024 19:55:51 -0900
Subject: [PATCH 2/3] Pass URL

We already pass the token and since we check the URL before calling
monitorMetadata I think it makes more sense to pass in the URL we
checked over getting it off the global again.
---
 src/workspacesProvider.ts | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/workspacesProvider.ts b/src/workspacesProvider.ts
index b066dd8c..ec3c1171 100644
--- a/src/workspacesProvider.ts
+++ b/src/workspacesProvider.ts
@@ -67,7 +67,7 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
       const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine
       if (showMetadata) {
         const agents = extractAgents(workspace)
-        agents.forEach((agent) => this.monitorMetadata(agent.id, token)) // monitor metadata for all agents
+        agents.forEach((agent) => this.monitorMetadata(agent.id, url, token)) // monitor metadata for all agents
       }
       return new WorkspaceTreeItem(workspace, this.getWorkspacesQuery === WorkspaceQuery.All, showMetadata)
     })
@@ -107,8 +107,8 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
 
   // monitorMetadata opens an SSE endpoint to monitor metadata on the specified
   // agent and registers a disposer that can be used to stop the watch.
-  monitorMetadata(agentId: WorkspaceAgent["id"], token: string): void {
-    const agentMetadataURL = new URL(`${this.storage.getURL()}/api/v2/workspaceagents/${agentId}/watch-metadata`)
+  monitorMetadata(agentId: WorkspaceAgent["id"], url: string, token: string): void {
+    const agentMetadataURL = new URL(`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`)
     const agentMetadataEventSource = new EventSource(agentMetadataURL.toString(), {
       headers: {
         "Coder-Session-Token": token,

From 63642280d8f8a4d0a107f27c88b1e307d3f31533 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Thu, 1 Feb 2024 13:07:15 -0900
Subject: [PATCH 3/3] Fix logout check

We have to check storage again, otherwise we are just checking the
values we already checked.
---
 src/workspacesProvider.ts | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/workspacesProvider.ts b/src/workspacesProvider.ts
index ec3c1171..fba0f65a 100644
--- a/src/workspacesProvider.ts
+++ b/src/workspacesProvider.ts
@@ -58,16 +58,25 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
 
     const resp = await getWorkspaces({ q: this.getWorkspacesQuery })
 
-    // We could have logged out while waiting for the query.
-    if (!url || !token) {
+    // We could have logged out while waiting for the query, or logged into a
+    // different deployment.
+    const url2 = this.storage.getURL()
+    const token2 = await this.storage.getSessionToken()
+    if (!url2 || !token2) {
       throw new Error("not logged in")
+    } else if (url !== url2) {
+      // In this case we need to fetch from the new deployment instead.
+      // TODO: It would be better to cancel this fetch when that happens,
+      // because this means we have to wait for the old fetch to finish before
+      // finally getting workspaces for the new one.
+      return this.fetch()
     }
 
     return resp.workspaces.map((workspace) => {
       const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine
       if (showMetadata) {
         const agents = extractAgents(workspace)
-        agents.forEach((agent) => this.monitorMetadata(agent.id, url, token)) // monitor metadata for all agents
+        agents.forEach((agent) => this.monitorMetadata(agent.id, url, token2)) // monitor metadata for all agents
       }
       return new WorkspaceTreeItem(workspace, this.getWorkspacesQuery === WorkspaceQuery.All, showMetadata)
     })