feat(auth): Allow workspace to pre-populate URL for quick sign-in (#6653

) (#6812)

This is cherry-picking @mmanela commit from
#6653

It was recently reverted here
#6794 due to problems with
reverting other PRs and merge conflicts.

This PR contains a required change for a successful Workspace Auth flow
so we should include this in the upcoming 1.64 release

## Original PR description
VS Code side part of

https://linear.app/sourcegraph/issue/SRC-908/automate-cody-authflow-vs-code

After a user creates a workspace we want to make it easy for them to
sign-in to the instance in their editor. For enterprise instances today,
they need to copy the instance url, go into VS Code and then choose to
signin with enterprise instance and then paste the url.

This change will help stream-line that flow by allowing the UI in
sourcegraph to link to VS Code with a url like

`vscode://sourcegraph.cody-ai?instance=someinstance.sourcegraphdev.app`

VS Code will handle and either switch to that account if you already are
auth'd OR pop-open the signin menu directly

![image](https://github.com/user-attachments/assets/cdab4135-2dae-48d3-945c-2e70187c37f0)

Note: For initial version I limited just to workspace domains. But we
can generalize that to allow other instance domains.

## Test plan

1. Validate when client doesn't have instance already
2. Validate when client already has instance connected

(cherry-picked from commit 9719dc4)

---------

Co-authored-by: Matthew Manela <mmanela@users.noreply.github.com>
(cherry picked from commit ef65034)

lib/shared/src/index.ts

@@ -405,7 +405,7 @@ export {
     checkVersion,
 } from './sourcegraph-api/siteVersion'
 export { configOverwrites } from './models/configOverwrites'
-export { isS2 } from './sourcegraph-api/environments'
+export { isS2, isWorkspaceInstance } from './sourcegraph-api/environments'
 export { createGitDiff } from './editor/create-git-diff'
 
 export { serialize, deserialize } from './lexicalEditor/atMentionsSerializer'

lib/shared/src/sourcegraph-api/environments.ts

@@ -41,3 +41,25 @@ export function isS2(arg: Pick<AuthStatus, 'endpoint'> | undefined | string): bo
 // TODO: Update to live link https://linear.app/sourcegraph/issue/CORE-535/cody-clients-migrate-ctas-to-live-links
 export const DOTCOM_WORKSPACE_UPGRADE_URL = new URL('https://sourcegraph.com/cody/manage')
 export const SG_WORKSPACES_URL = new URL('https://workspaces.sourcegraph.com')
+
+export const Workspaces_Host_Prod = '.sourcegraph.app'
+export const Workspaces_Host_Dev = '.sourcegraphapp.test:3443'
+
+// 🚨 SECURITY: This is used to validate a set of URLs we will allow to be passed in
+//              to the editor in the URL handler.
+export function isWorkspaceInstance(authStatus: Pick<AuthStatus, 'endpoint'> | undefined): boolean
+export function isWorkspaceInstance(url: string): boolean
+export function isWorkspaceInstance(arg: Pick<AuthStatus, 'endpoint'> | undefined | string): boolean {
+    const url = typeof arg === 'string' ? arg : arg?.endpoint
+    if (url === undefined) {
+        return false
+    }
+    try {
+        return (
+            new URL(url).host.endsWith(Workspaces_Host_Prod) ||
+            new URL(url).host.endsWith(Workspaces_Host_Dev)
+        )
+    } catch {
+        return false
+    }
+}

vscode/src/auth/auth.ts

@@ -18,6 +18,7 @@ import {
     isDotCom,
     isError,
     isNetworkLikeError,
+    isWorkspaceInstance,
     telemetryRecorder,
 } from '@sourcegraph/cody-shared'
 import { isSourcegraphToken } from '../chat/protocol'
@@ -38,6 +39,32 @@ interface LoginMenuItem {
 
 type AuthMenuType = 'signin' | 'switch'
 
+/**
+ *  Handles trying to directly sign-in or add to an enterprise instance.
+ * First tries to sign in with the current token, if it's valid. Otherwise,
+ * opens the sign-in flow and has user confirm.
+ */
+async function showEnterpriseInstanceUrlFlow(endpoint: string): Promise<void> {
+    const token = await secretStorage.getToken(endpoint)
+    const tokenSource = await secretStorage.getTokenSource(endpoint)
+    const authStatus = token
+        ? await authProvider.validateAndStoreCredentials(
+              { serverEndpoint: endpoint, accessToken: token, tokenSource },
+              'store-if-valid'
+          )
+        : undefined
+
+    if (!authStatus?.authenticated) {
+        const instanceUrl = await showInstanceURLInputBox(endpoint)
+        if (!instanceUrl) {
+            return
+        }
+        authProvider.setAuthPendingToEndpoint(instanceUrl)
+        redirectToEndpointLogin(instanceUrl)
+    } else {
+        await showAuthResultMessage(endpoint, authStatus)
+    }
+}
 /**
  * Show a quickpick to select the endpoint to sign into.
  */
@@ -142,12 +169,12 @@ async function showAuthMenu(type: AuthMenuType): Promise<LoginMenuItem | null> {
 /**
  * Show a VS Code input box to ask the user to enter a Sourcegraph instance URL.
  */
-async function showInstanceURLInputBox(title: string): Promise<string | undefined> {
+async function showInstanceURLInputBox(url?: string): Promise<string | undefined> {
     const result = await vscode.window.showInputBox({
-        title,
+        title: 'Connect to a Sourcegraph instance',
         prompt: 'Enter the URL of the Sourcegraph instance. For example, https://sourcegraph.example.com.',
         placeHolder: 'https://sourcegraph.example.com',
-        value: 'https://',
+        value: url ?? 'https://',
         password: false,
         ignoreFocusOut: true,
         // valide input to ensure the user is not entering a token as URL
@@ -305,8 +332,22 @@ export async function tokenCallbackHandler(uri: vscode.Uri): Promise<void> {
     closeAuthProgressIndicator()
 
     const params = new URLSearchParams(uri.query)
+
     const token = params.get('code') || params.get('token')
     const endpoint = currentAuthStatus().endpoint
+
+    // If we were provided an instance URL then it means we are
+    // request the user setup auth with a different sourcegraph instance
+    // We want to prompt them to switch to this instance and if needed
+    // start the auth flow
+    const instanceHost = params.get('instance')
+    const instanceUrl = instanceHost ? new URL(instanceHost).origin : undefined
+    if (instanceUrl && isWorkspaceInstance(instanceUrl)) {
+        // Prompt the user to switch/setup with the new instance
+        await showEnterpriseInstanceUrlFlow(instanceUrl)
+        return
+    }
+
     if (!token || !endpoint) {
         return
     }

vscode/src/main.ts

@@ -624,7 +624,7 @@ function registerUpgradeHandlers(disposables: vscode.Disposable[]): void {
                 if (uri.path === '/app-done') {
                     // This is an old re-entrypoint from App that is a no-op now.
                 } else {
-                    tokenCallbackHandler(uri)
+                    void tokenCallbackHandler(uri)
                 }
             },
         }),