feat: add device flow to vscode extension (#4189)

benfdking · web-flow · commit beb24f015612 · 2025-04-20T09:45:12.000+01:00
diff --git a/vscode/extension/package.json b/vscode/extension/package.json
@@ -41,6 +41,11 @@
         "title": "Sign in to Tobiko Cloud",
         "description": "SQLMesh"
       },
+      {
+          "command": "sqlmesh.signinSpecifyFlow",
+          "title": "Sign in to Tobiko Cloud (Specify Auth Flow)",
+          "description": "SQLMesh"
+      },
       {
         "command": "sqlmesh.signout",
         "title": "Sign out from Tobiko Cloud",
diff --git a/vscode/extension/src/auth/auth.ts b/vscode/extension/src/auth/auth.ts
@@ -19,13 +19,13 @@ export const AUTH_TYPE = "tobikodata"
 export const AUTH_NAME = "Tobiko"
 
 const tokenSchema = z.object({
-    iss: z.string(),
-    aud: z.string(),
-    sub: z.string(),
-    scope: z.string(),
-    iat: z.number(),
-    exp: z.number(),
-    email: z.string(),
+  iss: z.string(),
+  aud: z.string(),
+  sub: z.string(),
+  scope: z.string(),
+  iat: z.number(),
+  exp: z.number(),
+  email: z.string(),
 })
 const statusResponseSchema = z.object({
   is_logged_in: z.boolean(),
@@ -39,6 +39,14 @@ const loginUrlResponseSchema = z.object({
   verifier_code: z.string(),
 })
 
+const deviceCodeResponseSchema = z.object({
+  device_code: z.string(),
+  user_code: z.string(),
+  verification_uri: z.string(),
+  verification_uri_complete: z.string(),
+  expires_in: z.number(),
+})
+
 export class AuthenticationProviderTobikoCloud
   implements AuthenticationProvider
 {
@@ -91,8 +99,8 @@ export class AuthenticationProviderTobikoCloud
     const session = {
       id: token.email,
       account: {
-        id: token.email,
-        label: "Tobiko",
+        id: token.sub,
+        label: token.email,
       },
       scopes: token.scope.split(" "),
       accessToken: "",
@@ -101,6 +109,60 @@ export class AuthenticationProviderTobikoCloud
   }
 
   async createSession(): Promise<AuthenticationSession> {
+    await this.sign_in_oauth_flow()
+    const status = await this.get_status()
+    if (isErr(status)) {
+      throw new Error("Failed to get tcloud auth status")
+    }
+    const statusResponse = status.value
+    if (!statusResponse.is_logged_in) {
+      throw new Error("Failed to login to tcloud")
+    }
+    const token = statusResponse.id_token
+    const session: AuthenticationSession = {
+      id: token.email,
+      account: {
+        id: token.email,
+        label: "Tobiko",
+      },
+      scopes: token.scope.split(" "),
+      accessToken: "",
+    }
+    this._sessionChangeEmitter.fire({
+      added: [session],
+      removed: [],
+      changed: [],
+    })
+    return session
+  }
+
+  async removeSession(): Promise<void> {
+    // Get current sessions before logging out
+    const currentSessions = await this.getSessions()
+    const tcloudBin = await get_tcloud_bin()
+    const workspacePath = await getProjectRoot()
+    if (isErr(tcloudBin)) {
+      throw new Error("Failed to get tcloud bin")
+    }
+    const tcloudBinPath = tcloudBin.value
+    const result = await execAsync(tcloudBinPath, ["auth", "logout"], {
+      cwd: workspacePath.uri.fsPath,
+    })
+    if (result.exitCode !== 0) {
+      throw new Error("Failed to logout from tcloud")
+    }
+
+    // Emit event with the actual sessions that were removed
+    if (currentSessions.length > 0) {
+      this._sessionChangeEmitter.fire({
+        added: [],
+        removed: currentSessions,
+        changed: [],
+      })
+    }
+  }
+
+  async sign_in_oauth_flow(): Promise<void> {
     const workspacePath = await getProjectRoot()
     const tcloudBin = await get_tcloud_bin()
     if (isErr(tcloudBin)) {
@@ -117,83 +179,156 @@ export class AuthenticationProviderTobikoCloud
     if (result.exitCode !== 0) {
       throw new Error("Failed to get tcloud login url")
     }
-    const resultToJson = JSON.parse(result.stdout)
-    const urlCode = loginUrlResponseSchema.parse(resultToJson)
-    const url = urlCode.url
 
-    const ac = new AbortController()
-    const timeout = setTimeout(() => ac.abort(), 1000 * 60 * 5)
-    const backgroundServerForLogin = execAsync(
-      tcloudBinPath,
-      ["auth", "vscode", "start-server", urlCode.verifier_code],
-      {
-        cwd: workspacePath.uri.fsPath,
-        signal: ac.signal,
+    try {
+      const resultToJson = JSON.parse(result.stdout)
+      const urlCode = loginUrlResponseSchema.parse(resultToJson)
+      const url = urlCode.url
+
+      if (!url) {
+        throw new Error("Invalid login URL received")
       }
-    )
 
-    const messageResult = await window.showInformationMessage(
-      "Please login to Tobiko Cloud",
-      {
-        modal: true,
-      },
-      "Sign in with browser",
-      "Cancel"
-    )
+      const ac = new AbortController()
+      const timeout = setTimeout(() => ac.abort(), 1000 * 60 * 5)
+      const backgroundServerForLogin = execAsync(
+        tcloudBinPath,
+        ["auth", "vscode", "start-server", urlCode.verifier_code],
+        {
+          cwd: workspacePath.uri.fsPath,
+          signal: ac.signal,
+        }
+      )
 
-    if (messageResult === "Sign in with browser") {
-      await env.openExternal(Uri.parse(url))
-    }
-    if (messageResult === "Cancel") {
-      ac.abort()
-      throw new Error("Login cancelled")
-    }
+      const messageResult = await window.showInformationMessage(
+        "Please login to Tobiko Cloud",
+        {
+          modal: true,
+        },
+        "Sign in with browser",
+        "Cancel"
+      )
 
-    try {
-      const output = await backgroundServerForLogin
-      if (output.exitCode !== 0) {
-        throw new Error(`Failed to start server: ${output.stderr}`)
+      if (messageResult === "Sign in with browser") {
+        await env.openExternal(Uri.parse(url))
+      } else {
+        // Always abort the server if not proceeding with sign in
+        ac.abort()
+        clearTimeout(timeout)
+        if (messageResult === "Cancel") {
+          throw new Error("Login cancelled")
+        }
+        return
       }
-    } catch (error) {
-      traceError(`Server error: ${error}`)
-      throw error
-    }
 
-    clearTimeout(timeout)
-
-    const status = await this.get_status()
-    if (isErr(status)) {
-      throw new Error("Failed to get tcloud auth status")
-    }
-    const statusResponse = status.value
-    if (!statusResponse.is_logged_in) {
-      throw new Error("Failed to login to tcloud")
-    }
-    const scopes = statusResponse.id_token.scope.split(" ")
-    const session: AuthenticationSession = {
-      id: AuthenticationProviderTobikoCloud.id,
-      account: {
-        id: AuthenticationProviderTobikoCloud.id,
-        label: "Tobiko",
-      },
-      scopes: scopes,
-      accessToken: ""
+      try {
+        const output = await backgroundServerForLogin
+        if (output.exitCode !== 0) {
+          throw new Error(
+            `Failed to complete authentication: ${output.stderr}`
+          )
+        }
+        // Get updated session and notify about the change
+        const sessions = await this.getSessions()
+        if (sessions.length > 0) {
+          this._sessionChangeEmitter.fire({
+            added: sessions,
+            removed: [],
+            changed: [],
+          })
+        }
+      } catch (error) {
+        if (error instanceof Error && error.name === "AbortError") {
+          throw new Error("Authentication timeout or aborted")
+        }
+        traceError(`Server error: ${error}`)
+        throw error
+      } finally {
+        clearTimeout(timeout)
+      }
+    } catch (error) {
+      if (error instanceof Error && error.message === "Login cancelled") {
+        throw error
+      }
+      traceError(`Authentication flow error: ${error}`)
+      throw new Error("Failed to complete authentication flow")
     }
-    return session
   }
 
-  async removeSession(): Promise<void> {
-    const tcloudBin = await get_tcloud_bin()
+  async sign_in_device_flow(): Promise<void> {
     const workspacePath = await getProjectRoot()
+    const tcloudBin = await get_tcloud_bin()
     if (isErr(tcloudBin)) {
       throw new Error("Failed to get tcloud bin")
     }
     const tcloudBinPath = tcloudBin.value
-    const result = await execAsync(tcloudBinPath, ["auth", "logout"], {
-      cwd: workspacePath.uri.fsPath,
-    })
+    const result = await execAsync(
+      tcloudBinPath,
+      ["auth", "vscode", "device"],
+      {
+        cwd: workspacePath.uri.fsPath,
+      }
+    )
     if (result.exitCode !== 0) {
-      throw new Error("Failed to logout from tcloud")
+      throw new Error("Failed to get device code")
+    }
+
+    try {
+      const resultToJson = JSON.parse(result.stdout)
+      const deviceCodeResponse = deviceCodeResponseSchema.parse(resultToJson)
+
+      const ac = new AbortController()
+      const timeout = setTimeout(() => ac.abort(), 1000 * 60 * 5)
+      const waiting = execAsync(
+        tcloudBinPath,
+        ["auth", "vscode", "poll_device", deviceCodeResponse.device_code],
+        {
+          cwd: workspacePath.uri.fsPath,
+          signal: ac.signal,
+        }
+      )
+
+      const messageResult = await window.showInformationMessage(
+        `Confirm the code ${deviceCodeResponse.user_code} at ${deviceCodeResponse.verification_uri}`,
+        {
+          modal: true,
+        },
+        "Open browser",
+        "Cancel"
+      )
+
+      if (messageResult === "Open browser") {
+        await env.openExternal(Uri.parse(deviceCodeResponse.verification_uri_complete))
+      }
+      if (messageResult === "Cancel") {
+        ac.abort()
+        throw new Error("Login cancelled")
+      }
+
+      try {
+        const output = await waiting
+        if (output.exitCode !== 0) {
+          throw new Error(`Failed to authenticate: ${output.stderr}`)
+        }
+
+        // Get updated session and notify about the change
+        const sessions = await this.getSessions()
+        if (sessions.length > 0) {
+          this._sessionChangeEmitter.fire({
+            added: sessions,
+            removed: [],
+            changed: [],
+          })
+        }
+      } catch (error) {
+        traceError(`Authentication error: ${error}`)
+        throw error
+      } finally {
+        clearTimeout(timeout)
+      }
+    } catch (error) {
+      traceError(`JSON parsing error: ${error}`)
+      throw new Error("Failed to parse device code response")
     }
   }
 }
diff --git a/vscode/extension/src/commands/signin.ts b/vscode/extension/src/commands/signin.ts
@@ -1,8 +1,13 @@
 import { AuthenticationProviderTobikoCloud } from "../auth/auth"
 import * as vscode from "vscode"
+import { isCodespaces } from "../utilities/isCodespaces"
 
-
-export const signIn = (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => {
-    await authenticationProvider.createSession()
+export const signIn =
+  (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => {
+    if (isCodespaces()) {
+      await authenticationProvider.sign_in_device_flow()
+    } else {
+      await authenticationProvider.createSession()
+    }
     await vscode.window.showInformationMessage("Signed in successfully")
-}
+  }
diff --git a/vscode/extension/src/commands/signinSpecifyFlow.ts b/vscode/extension/src/commands/signinSpecifyFlow.ts
@@ -0,0 +1,33 @@
+import { AuthenticationProviderTobikoCloud } from "../auth/auth"
+import { traceInfo } from "../utilities/common/log"
+import { window } from "vscode"
+
+export const signInSpecifyFlow = (authenticationProvider: AuthenticationProviderTobikoCloud) => async () => {
+    traceInfo("Sign in specify flow")
+    const flowOptions = [
+        { label: "OAuth Flow", description: "Sign in using OAuth flow in your browser" },
+        { label: "Device Flow", description: "Sign in using a device code" }
+    ]
+    const selectedFlow = await window.showQuickPick(flowOptions, {
+        placeHolder: "Select authentication flow method",
+        ignoreFocusOut: true
+    })
+    if (!selectedFlow) {
+        traceInfo("Sign in cancelled by user")
+        return
+    }
+    if (selectedFlow.label === "OAuth Flow") {
+        await authenticationProvider.sign_in_oauth_flow()
+        await authenticationProvider.getSessions()
+        await window.showInformationMessage("Sign in success")
+        return
+    } else if (selectedFlow.label === "Device Flow") {
+        await authenticationProvider.sign_in_device_flow()
+        await authenticationProvider.getSessions()
+        await window.showInformationMessage("Sign in success")
+        return
+    } else {
+        traceInfo("Invalid flow selected")
+        return
+    }
+}
diff --git a/vscode/extension/src/extension.ts b/vscode/extension/src/extension.ts
diff --git a/vscode/extension/src/utilities/isCodespaces.ts b/vscode/extension/src/utilities/isCodespaces.ts
