When connecting to a recent workspace, log into that deployment

code-asher · code-asher · commit d3c42ee2f06a · 2024-06-04T14:05:02.000-08:00
That way the sidebar displays the workspaces that belong to that same
deployment.

As a side effect, this will now avoid logging into a different
deployment first only to switch over to the one the remote belongs to.
diff --git a/src/extension.ts b/src/extension.ts
@@ -40,7 +40,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
   const storage = new Storage(output, ctx.globalState, ctx.secrets, ctx.globalStorageUri, ctx.logUri)
 
   // This client tracks the current login and will be used through the life of
-  // the plugin to poll workspaces for the current login.
+  // the plugin to poll workspaces for the current login, as well as being used
+  // in commands that operate on the current login.
   const url = storage.getUrl()
   const restClient = await makeCoderSdk(url || "", await storage.getSessionToken(), storage)
 
@@ -57,33 +58,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
     myWorkspacesProvider.setVisibility(event.visible)
   })
 
-  if (url) {
-    restClient
-      .getAuthenticatedUser()
-      .then(async (user) => {
-        if (user && user.roles) {
-          vscode.commands.executeCommand("setContext", "coder.authenticated", true)
-          if (user.roles.find((role) => role.name === "owner")) {
-            await vscode.commands.executeCommand("setContext", "coder.isOwner", true)
-          }
-
-          // Fetch and monitor workspaces, now that we know the client is good.
-          myWorkspacesProvider.fetchAndRefresh()
-          allWorkspacesProvider.fetchAndRefresh()
-        }
-      })
-      .catch((error) => {
-        // This should be a failure to make the request, like the header command
-        // errored.
-        vscodeProposed.window.showErrorMessage("Failed to check user authentication: " + error.message)
-      })
-      .finally(() => {
-        vscode.commands.executeCommand("setContext", "coder.loaded", true)
-      })
-  } else {
-    vscode.commands.executeCommand("setContext", "coder.loaded", true)
-  }
-
+  // Handle vscode:// URIs.
   vscode.window.registerUriHandler({
     handleUri: async (uri) => {
       const params = new URLSearchParams(uri.query)
@@ -127,12 +102,15 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
         await storage.configureCli(toSafeHost(url), url, token)
 
         vscode.commands.executeCommand("coder.open", owner, workspace, agent, folder, openRecent)
+      } else {
+        throw new Error(`Unknown path ${uri.path}`)
       }
     },
   })
 
+  // Register globally available commands.  Many of these have visibility
+  // controlled by contexts, see `when` in the package.json.
   const commands = new Commands(vscodeProposed, restClient, storage)
-
   vscode.commands.registerCommand("coder.login", commands.login.bind(commands))
   vscode.commands.registerCommand("coder.logout", commands.logout.bind(commands))
   vscode.commands.registerCommand("coder.open", commands.open.bind(commands))
@@ -153,38 +131,71 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
   // Since the "onResolveRemoteAuthority:ssh-remote" activation event exists
   // in package.json we're able to perform actions before the authority is
   // resolved by the remote SSH extension.
-  if (!vscodeProposed.env.remoteAuthority) {
-    return
-  }
-  const remote = new Remote(vscodeProposed, storage, commands, ctx.extensionMode)
-  try {
-    await remote.setup(vscodeProposed.env.remoteAuthority)
-  } catch (ex) {
-    if (ex instanceof CertificateError) {
-      await ex.showModal("Failed to open workspace")
-    } else if (isAxiosError(ex)) {
-      const msg = getErrorMessage(ex, "")
-      const detail = getErrorDetail(ex)
-      const urlString = axios.getUri(ex.response?.config)
-      let path = urlString
-      try {
-        path = new URL(urlString).pathname
-      } catch (e) {
-        // ignore, default to full url
+  if (vscodeProposed.env.remoteAuthority) {
+    const remote = new Remote(vscodeProposed, storage, commands, ctx.extensionMode)
+    try {
+      const details = await remote.setup(vscodeProposed.env.remoteAuthority)
+      if (details) {
+        // Authenticate the plugin client which is used in the sidebar to display
+        // workspaces belonging to this deployment.
+        restClient.setHost(details.url)
+        restClient.setSessionToken(details.token)
       }
-      await vscodeProposed.window.showErrorMessage("Failed to open workspace", {
-        detail: `API ${ex.response?.config.method?.toUpperCase()} to '${path}' failed with code ${ex.response?.status}.\nMessage: ${msg}\nDetail: ${detail}`,
-        modal: true,
-        useCustom: true,
+    } catch (ex) {
+      if (ex instanceof CertificateError) {
+        await ex.showModal("Failed to open workspace")
+      } else if (isAxiosError(ex)) {
+        const msg = getErrorMessage(ex, "")
+        const detail = getErrorDetail(ex)
+        const urlString = axios.getUri(ex.response?.config)
+        let path = urlString
+        try {
+          path = new URL(urlString).pathname
+        } catch (e) {
+          // ignore, default to full url
+        }
+        await vscodeProposed.window.showErrorMessage("Failed to open workspace", {
+          detail: `API ${ex.response?.config.method?.toUpperCase()} to '${path}' failed with code ${ex.response?.status}.\nMessage: ${msg}\nDetail: ${detail}`,
+          modal: true,
+          useCustom: true,
+        })
+      } else {
+        await vscodeProposed.window.showErrorMessage("Failed to open workspace", {
+          detail: (ex as string).toString(),
+          modal: true,
+          useCustom: true,
+        })
+      }
+      // Always close remote session when we fail to open a workspace.
+      await remote.closeRemote()
+    }
+  }
+
+  // See if the plugin client is authenticated.
+  if (restClient.getAxiosInstance().defaults.baseURL) {
+    restClient
+      .getAuthenticatedUser()
+      .then(async (user) => {
+        if (user && user.roles) {
+          vscode.commands.executeCommand("setContext", "coder.authenticated", true)
+          if (user.roles.find((role) => role.name === "owner")) {
+            await vscode.commands.executeCommand("setContext", "coder.isOwner", true)
+          }
+
+          // Fetch and monitor workspaces, now that we know the client is good.
+          myWorkspacesProvider.fetchAndRefresh()
+          allWorkspacesProvider.fetchAndRefresh()
+        }
       })
-    } else {
-      await vscodeProposed.window.showErrorMessage("Failed to open workspace", {
-        detail: (ex as string).toString(),
-        modal: true,
-        useCustom: true,
+      .catch((error) => {
+        // This should be a failure to make the request, like the header command
+        // errored.
+        vscode.window.showErrorMessage("Failed to check user authentication: " + error.message)
       })
-    }
-    // Always close remote session when we fail to open a workspace.
-    await remote.closeRemote()
+      .finally(() => {
+        vscode.commands.executeCommand("setContext", "coder.loaded", true)
+      })
+  } else {
+    vscode.commands.executeCommand("setContext", "coder.loaded", true)
   }
 }
diff --git a/src/remote.ts b/src/remote.ts
@@ -21,6 +21,11 @@ import { AuthorityPrefix, parseRemoteAuthority } from "./util"
 import { supportsCoderAgentLogDirFlag } from "./version"
 import { WorkspaceAction } from "./workspaceAction"
 
+export interface RemoteDetails extends vscode.Disposable {
+  url: string
+  token: string
+}
+
 export class Remote {
   public constructor(
     private readonly vscodeProposed: typeof vscode,
@@ -29,7 +34,12 @@ export class Remote {
     private readonly mode: vscode.ExtensionMode,
   ) {}
 
-  public async setup(remoteAuthority: string): Promise<vscode.Disposable | undefined> {
+  /**
+   * Ensure the workspace specified by the remote authority is ready to receive
+   * SSH connections.  Return undefined if the authority is not for a Coder
+   * workspace or when explicitly closing the remote.
+   */
+  public async setup(remoteAuthority: string): Promise<RemoteDetails | undefined> {
     const parts = parseRemoteAuthority(remoteAuthority)
     if (!parts) {
       // Not a Coder host.
@@ -63,16 +73,17 @@ export class Remote {
       return
     }
 
-    // It is possible to connect to any previously connected workspace, which
-    // might not belong to the deployment the plugin is currently logged into.
-    // For that reason, create a separate REST client instead of using the
-    // global one generally used by the plugin.
-    const restClient = await makeCoderSdk(baseUrlRaw, token, this.storage)
+    // We could use the plugin client, but it is possible for the user to log
+    // out or log into a different deployment while still connected, which would
+    // break this connection.  We could force close the remote session or
+    // disallow logging out/in altogether, but for now just use a separate
+    // client to remain unaffected by whatever the plugin is doing.
+    const workspaceRestClient = await makeCoderSdk(baseUrlRaw, token, this.storage)
     // Store for use in commands.
-    this.commands.workspaceRestClient = restClient
+    this.commands.workspaceRestClient = workspaceRestClient
 
     // First thing is to check the version.
-    const buildInfo = await restClient.getBuildInfo()
+    const buildInfo = await workspaceRestClient.getBuildInfo()
     const parsedVersion = semver.parse(buildInfo.version)
     // Server versions before v0.14.1 don't support the vscodessh command!
     if (
@@ -98,7 +109,7 @@ export class Remote {
     // Next is to find the workspace from the URI scheme provided.
     let workspace: Workspace
     try {
-      workspace = await restClient.getWorkspaceByOwnerAndName(parts.username, parts.workspace)
+      workspace = await workspaceRestClient.getWorkspaceByOwnerAndName(parts.username, parts.workspace)
       this.commands.workspace = workspace
     } catch (error) {
       if (!isAxiosError(error)) {
@@ -149,7 +160,7 @@ export class Remote {
     disposables.push(this.registerLabelFormatter(remoteAuthority, workspace.owner_name, workspace.name))
 
     // Initialize any WorkspaceAction notifications (auto-off, upcoming deletion)
-    const action = await WorkspaceAction.init(this.vscodeProposed, restClient, this.storage)
+    const action = await WorkspaceAction.init(this.vscodeProposed, workspaceRestClient, this.storage)
 
     // Make sure the workspace has started.
     let buildComplete: undefined | (() => void)
@@ -175,7 +186,7 @@ export class Remote {
           }),
       )
 
-      const latestBuild = await restClient.startWorkspace(workspace.id, versionID)
+      const latestBuild = await workspaceRestClient.startWorkspace(workspace.id, versionID)
       workspace = {
         ...workspace,
         latest_build: latestBuild,
@@ -206,7 +217,7 @@ export class Remote {
         } as Partial<vscode.Pseudoterminal> as any,
       })
       // This fetches the initial bunch of logs.
-      const logs = await restClient.getWorkspaceBuildLogs(workspace.latest_build.id, new Date())
+      const logs = await workspaceRestClient.getWorkspaceBuildLogs(workspace.latest_build.id, new Date())
       logs.forEach((log) => writeEmitter.fire(log.output + "\r\n"))
       terminal.show(true)
       // This follows the logs for new activity!
@@ -243,7 +254,7 @@ export class Remote {
         }
       })
       writeEmitter.fire("Build complete")
-      workspace = await restClient.getWorkspace(workspace.id)
+      workspace = await workspaceRestClient.getWorkspace(workspace.id)
       this.commands.workspace = workspace
       terminal.dispose()
 
@@ -372,10 +383,10 @@ export class Remote {
       if (newWorkspace.outdated) {
         if (!workspace.outdated || !hasShownOutdatedNotification) {
           hasShownOutdatedNotification = true
-          restClient
+          workspaceRestClient
             .getTemplate(newWorkspace.template_id)
             .then((template) => {
-              return restClient.getTemplateVersion(template.active_version_id)
+              return workspaceRestClient.getTemplateVersion(template.active_version_id)
             })
             .then((version) => {
               let infoMessage = `A new version of your workspace is available.`
@@ -384,7 +395,7 @@ export class Remote {
               }
               vscode.window.showInformationMessage(infoMessage, "Update").then((action) => {
                 if (action === "Update") {
-                  vscode.commands.executeCommand("coder.workspace.update", newWorkspace, restClient)
+                  vscode.commands.executeCommand("coder.workspace.update", newWorkspace, workspaceRestClient)
                 }
               })
             })
@@ -498,7 +509,7 @@ export class Remote {
     // If we didn't write to the SSH config file, connecting would fail with
     // "Host not found".
     try {
-      await this.updateSSHConfig(restClient, parts.label, parts.host, hasCoderLogs)
+      await this.updateSSHConfig(workspaceRestClient, parts.label, parts.host, hasCoderLogs)
     } catch (error) {
       this.storage.writeToCoderOutputChannel(`Failed to configure SSH: ${error}`)
       throw error
@@ -522,7 +533,13 @@ export class Remote {
       }),
     )
 
+    // Returning the URL and token allows the plugin to authenticate its own
+    // client, for example to display the list of workspaces belonging to this
+    // deployment in the sidebar.  We use our own client in here for reasons
+    // explained above.
     return {
+      url: baseUrlRaw,
+      token,
       dispose: () => {
         eventSource.close()
         action.cleanupWorkspaceActions()
