Agent: authentication capability - enable url handler and auth redire…

…ctions (#5325)

Allows agent to set up HTTP server to listen for authentication
redirections when agent capabilities for authentication are enabled.

TODO: 
- [x] authentication url handler
- [x] agent capability configuration 
- [x] work with secret storage (added by
#5348)

Follow-ups:
- we should close the opened browser to refocus the webview if possible.

## Test plan

<!-- Required. See
https://docs-legacy.sourcegraph.com/dev/background-information/testing_principles.
-->

Example of this feature working in Visual Studio:


https://github.com/user-attachments/assets/d21c8c2f-2667-426a-9c4d-991b7645d2e5

Updated onboarding view for non VS Code editors to support login with
browser for enterprise:


![image](https://github.com/user-attachments/assets/19a234a7-83a0-4d52-b918-b1680fa10e72)

## Changelog

<!-- OPTIONAL; info at
https://www.notion.so/sourcegraph/Writing-a-changelog-entry-dd997f411d524caabf0d8d38a24a878c
-->

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/ClientCapabilities.kt

@@ -4,6 +4,7 @@ package com.sourcegraph.cody.agent.protocol_generated;
 import com.google.gson.annotations.SerializedName;
 
 data class ClientCapabilities(
+  val authentication: AuthenticationEnum? = null, // Oneof: enabled, none
   val completions: CompletionsEnum? = null, // Oneof: none
   val chat: ChatEnum? = null, // Oneof: none, streaming
   val git: GitEnum? = null, // Oneof: none, enabled
@@ -23,6 +24,11 @@ data class ClientCapabilities(
   val webviewNativeConfig: WebviewNativeConfigParams? = null,
 ) {
 
+  enum class AuthenticationEnum {
+    @SerializedName("enabled") Enabled,
+    @SerializedName("none") None,
+  }
+
   enum class CompletionsEnum {
     @SerializedName("none") None,
   }

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/CodyAgentClient.kt

@@ -74,4 +74,6 @@ interface CodyAgentClient {
   fun webview_setHtml(params: Webview_SetHtmlParams)
   @JsonNotification("window/didChangeContext")
   fun window_didChangeContext(params: Window_DidChangeContextParams)
+  @JsonNotification("window/focusSidebar")
+  fun window_focusSidebar(params: Null?)
 }

agent/src/AgentAuthHandler.test.ts

@@ -0,0 +1,59 @@
+import http from 'node:http'
+import open from 'open'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { URI } from 'vscode-uri'
+import { AgentAuthHandler } from './AgentAuthHandler'
+
+vi.mock('open')
+vi.mock('node:http', () => ({
+    default: {
+        createServer: vi.fn().mockReturnValue({
+            listen: vi.fn().mockImplementation((port, uri, callback) => {
+                callback()
+            }),
+            on: vi.fn(),
+            address: vi.fn().mockReturnValue({ port: 123 }),
+        }),
+    },
+}))
+
+describe('AgentAuthHandler', () => {
+    let agentAuthHandler: AgentAuthHandler
+
+    beforeEach(() => {
+        agentAuthHandler = new AgentAuthHandler()
+        agentAuthHandler.setTokenCallbackHandler(uri => console.log(`Token received: ${uri}`))
+    })
+
+    afterEach(() => {
+        vi.clearAllMocks()
+    })
+
+    describe('handleCallback', () => {
+        it.each([
+            [
+                'valid endpointUri',
+                'https://sourcegraph.test/user/settings/tokens/new/callback?requestFrom=CODY_JETBRAINS',
+                'https://sourcegraph.test/user/settings/tokens/new/callback?requestFrom=CODY_JETBRAINS-123',
+            ],
+            [
+                'valid endpointUri with additional params appended',
+                'https://sourcegraph.com/user/settings/tokens/new/callback?requestFrom=VISUAL_STUDIO&tokenReceiverUrl=https%3A%2F%2Fexample.com',
+                'https://sourcegraph.com/user/settings/tokens/new/callback?requestFrom=VISUAL_STUDIO-123&tokenReceiverUrl=https%3A%2F%2Fexample.com',
+            ],
+            ['invalid IDE', 'https://sourcegraph.com', 'https://sourcegraph.com/'],
+            ['invalid endpointUri', 'invalid-url', undefined],
+        ])('%s', (_, endpointUri: string, expectedUrl?: string) => {
+            const uri = URI.parse(endpointUri)
+
+            agentAuthHandler.handleCallback(uri)
+
+            if (!expectedUrl) {
+                expect(http.createServer).not.toHaveBeenCalled()
+            } else {
+                expect(http.createServer).toHaveBeenCalled()
+                expect(open).toHaveBeenCalledWith(expect.stringContaining(expectedUrl))
+            }
+        })
+    })
+})

agent/src/AgentAuthHandler.ts

@@ -0,0 +1,161 @@
+import type { IncomingMessage, Server, ServerResponse } from 'node:http'
+import http from 'node:http'
+import type { AddressInfo } from 'node:net'
+import { logDebug } from '@sourcegraph/cody-shared'
+import open from 'open'
+import { URI } from 'vscode-uri'
+
+type CallbackHandler = (url: URI, token?: string) => void
+
+const SIX_MINUTES = 6 * 60 * 1000
+
+/**
+ * Handles the authentication flow for Agent clients.
+ * Manages the creation and lifecycle of an HTTP server to handle the token callback from the Sourcegraph login flow.
+ * Redirects the user to the Sourcegraph instance login page and handles the token callback.
+ */
+export class AgentAuthHandler {
+    private port = 0
+    private server: Server | null = null
+    private tokenCallbackHandlers: CallbackHandler[] = []
+
+    public setTokenCallbackHandler(handler: CallbackHandler): void {
+        this.tokenCallbackHandlers.push(handler)
+    }
+
+    public handleCallback(url: URI): void {
+        try {
+            const callbackUri = getValidCallbackUri(url.toString())
+            if (!callbackUri) {
+                throw new Error(url.toString() + ' is not a valid URL')
+            }
+            this.startServer(url.toString())
+        } catch (error) {
+            logDebug('AgentAuthHandler', `Invalid callback URL: ${error}`)
+        }
+    }
+
+    private startServer(callbackUri: string): void {
+        if (!this.tokenCallbackHandlers?.length) {
+            logDebug('AgentAuthHandler', 'Token callback handler is not set.')
+            return
+        }
+
+        if (this.server && this.port) {
+            logDebug('AgentAuthHandler', 'Server already running')
+            this.redirectToEndpointLoginPage(callbackUri)
+            return
+        }
+
+        // Create an HTTP server to handle the token callback
+        const server = http.createServer((req: IncomingMessage, res: ServerResponse) => {
+            if (req.url?.startsWith('/api/sourcegraph/token')) {
+                const url = new URL(req.url, `http://127.0.0.1:${this.port}`)
+                const token = url.searchParams.get('token')
+                if (token) {
+                    for (const handler of this.tokenCallbackHandlers) {
+                        handler(URI.parse(req.url), token)
+                    }
+                    res.writeHead(200, { 'Content-Type': 'text/html' })
+                    // Close the window once the token is received.
+                    res.end(`
+                        <html>
+                            <body>
+                                Token received. This window will close automatically.
+                                <script>
+                                    window.close();
+                                </script>
+                            </body>
+                        </html>
+                    `)
+                    this.closeServer()
+                } else {
+                    res.writeHead(400, { 'Content-Type': 'text/plain' })
+                    res.end('Token not found.')
+                }
+            } else {
+                res.writeHead(404, { 'Content-Type': 'text/plain' })
+                res.end('Not found')
+            }
+        })
+
+        // Bind the server to the loopback interface
+        server.listen(0, '127.0.0.1', () => {
+            // The server is now bound to the loopback interface (127.0.0.1).
+            // This ensures that only local processes can connect to it.
+            this.port = (server.address() as AddressInfo).port
+            this.server = server
+            logDebug('AgentAuthHandler', `Server listening on port ${this.port}`)
+            this.redirectToEndpointLoginPage(callbackUri)
+            // Automatically close the server after 6 minutes,
+            // as the startTokenReceiver in token-receiver.ts only listens for 5 minutes.
+            setTimeout(() => this.closeServer(), SIX_MINUTES)
+        })
+
+        // Handle server errors
+        server.on('error', error => {
+            logDebug('AgentAuthHandler', `Server error: ${error}`)
+        })
+    }
+
+    private closeServer(): void {
+        if (this.server) {
+            logDebug('AgentAuthHandler', 'Auth server closed')
+            this.server.close()
+        }
+        this.server = null
+        this.port = 0
+    }
+
+    /**
+     * Redirects the user to the endpoint login page with the updated callback URI.
+     *
+     * The callback URI is updated by finding the 'requestFrom' parameter in the query string,
+     * removing the old parameter, and adding a new parameter with the correct port number appended.
+     *
+     * @param callbackUri - The original callback URI to be updated.
+     */
+    private redirectToEndpointLoginPage(callbackUri: string): void {
+        const uri = new URL(callbackUri)
+        const params = new URLSearchParams(decodeURIComponent(uri.search))
+        const requestFrom = params.get('requestFrom')
+        if (requestFrom) {
+            // Add the new parameter with the correct port number appended.
+            const newRequestFrom = `${requestFrom}-${this.port}`
+            params.set('requestFrom', newRequestFrom)
+            const redirect = params.get('redirect')
+            if (redirect) {
+                params.set('redirect', redirect.replace(requestFrom, newRequestFrom))
+            }
+            uri.search = params.toString()
+        }
+        open(uri.toString())
+    }
+
+    public dispose(): void {
+        this.closeServer()
+    }
+}
+
+/**
+ * Validates and normalizes a given callback URI.
+ *
+ * @param uri - The callback URI to validate and normalize.
+ * @returns The validated and normalized URI, or `null` if the input URI is invalid.
+ * @throws {Error} If the input URI is empty or starts with `file:`.
+ */
+function getValidCallbackUri(uri: string): URI | null {
+    if (!uri || uri.startsWith('file:')) {
+        throw new Error('Empty URL')
+    }
+    try {
+        const endpointUri = new URL(uri)
+        if (!endpointUri.protocol.startsWith('http')) {
+            endpointUri.protocol = 'https:'
+        }
+        return URI.parse(endpointUri.href)
+    } catch (error) {
+        logDebug('Invalid URL: ', `${error}`)
+        return null
+    }
+}

agent/src/agent.ts

@@ -61,6 +61,7 @@ import { IndentationBasedFoldingRangeProvider } from '../../vscode/src/lsp/foldi
 import type { FixupActor, FixupFileCollection } from '../../vscode/src/non-stop/roles'
 import type { FixupControlApplicator } from '../../vscode/src/non-stop/strategies'
 import { AgentWorkspaceEdit } from '../../vscode/src/testutils/AgentWorkspaceEdit'
+import { AgentAuthHandler } from './AgentAuthHandler'
 import { AgentFixupControls } from './AgentFixupControls'
 import { AgentProviders } from './AgentProviders'
 import { AgentClientManagedSecretStorage, AgentStatelessSecretStorage } from './AgentSecretStorage'
@@ -340,6 +341,7 @@ export class Agent extends MessageHandler implements ExtensionClient {
     public webPanels = new AgentWebviewPanels()
     public webviewViewProviders = new Map<string, vscode.WebviewViewProvider>()
 
+    public authenticationHandler: AgentAuthHandler | null = null
     private authenticationPromise: Promise<AuthStatus | undefined> = Promise.resolve(undefined)
 
     private clientInfo: ClientInfo | null = null
@@ -410,6 +412,9 @@ export class Agent extends MessageHandler implements ExtensionClient {
                     this.notify('ignore/didChange', null)
                 })
             }
+            if (clientInfo.capabilities?.authentication === 'enabled') {
+                this.authenticationHandler = new AgentAuthHandler()
+            }
             if (process.env.CODY_DEBUG === 'true') {
                 console.error(
                     `Cody Agent: handshake with client '${clientInfo.name}' (version '${clientInfo.version}') at workspace root path '${clientInfo.workspaceRootUri}'\n`
@@ -434,11 +439,12 @@ export class Agent extends MessageHandler implements ExtensionClient {
                     secrets
                 )
 
-                this.authenticationPromise = clientInfo.extensionConfiguration
+                this.authenticationPromise = clientInfo.extensionConfiguration?.accessToken
                     ? this.handleConfigChanges(clientInfo.extensionConfiguration, {
                           forceAuthentication: true,
                       })
                     : this.authStatus()
+
                 const authStatus = await this.authenticationPromise
 
                 const webviewKind = clientInfo.capabilities?.webview || 'agentic'

agent/src/vscode-shim.ts

@@ -622,7 +622,18 @@ const _window: typeof vscode.window = {
     onDidChangeWindowState: emptyEvent(),
     onDidCloseTerminal: emptyEvent(),
     onDidOpenTerminal: emptyEvent(),
-    registerUriHandler: () => emptyDisposable,
+    registerUriHandler: (vsceHandler: vscode.UriHandler) => {
+        if (agent?.authenticationHandler && vsceHandler.handleUri) {
+            // Add the token callback handler implemented for VS Code to the agent authentication handler.
+            const handler = agent.authenticationHandler
+            handler.setTokenCallbackHandler(vsceHandler.handleUri)
+            // A callback to notify the client to focus the sidebar when the token callback is handled.
+            const callbackFocusNotifier = () => agent?.notify('window/focusSidebar', null)
+            handler.setTokenCallbackHandler(callbackFocusNotifier)
+            return new Disposable(() => handler.dispose())
+        }
+        return emptyDisposable
+    },
     registerWebviewViewProvider: (
         viewId: string,
         provider: vscode.WebviewViewProvider,
@@ -1039,6 +1050,12 @@ const _env: Partial<typeof vscode.env> = {
         writeText: () => Promise.resolve(),
     },
     openExternal: (uri: vscode.Uri): Thenable<boolean> => {
+        // Handle the case where the user is trying to authenticate with redirect URI.
+        if (uri.toString()?.includes('user/settings/tokens/new/callback?requestFrom')) {
+            agent?.authenticationHandler?.handleCallback(uri)
+            return Promise.resolve(true)
+        }
+
         try {
             open(uri.toString())
             return Promise.resolve(true)

lib/shared/src/auth/referral.ts

@@ -0,0 +1,29 @@
+import { CodyIDE } from '../configuration'
+
+/**
+ * Returns a known referral code to use based on the current VS Code environment.
+ */
+export function getCodyAuthReferralCode(ideName: CodyIDE, uriScheme?: string): string | undefined {
+    const referralCodes: Record<CodyIDE, string> = {
+        [CodyIDE.JetBrains]: 'CODY_JETBRAINS',
+        [CodyIDE.Neovim]: 'CODY_NEOVIM',
+        [CodyIDE.Emacs]: 'CODY_EMACS',
+        [CodyIDE.VisualStudio]: 'VISUAL_STUDIO',
+        [CodyIDE.Eclipse]: 'ECLIPSE',
+        [CodyIDE.VSCode]: 'CODY',
+        [CodyIDE.Web]: 'CODY',
+    }
+
+    if (ideName === CodyIDE.VSCode) {
+        switch (uriScheme) {
+            case 'vscode-insiders':
+                return 'CODY_INSIDERS'
+            case 'vscodium':
+                return 'CODY_VSCODIUM'
+            case 'cursor':
+                return 'CODY_CURSOR'
+        }
+    }
+
+    return referralCodes[ideName] || undefined
+}

lib/shared/src/index.ts

@@ -275,6 +275,7 @@ export {
 export type { CurrentUserCodySubscription } from './sourcegraph-api/graphql/client'
 export * from './auth/types'
 export * from './auth/tokens'
+export * from './auth/referral'
 export * from './chat/sse-iterator'
 export {
     parseMentionQuery,

vscode/src/auth/auth.ts

@@ -272,7 +272,7 @@ export async function tokenCallbackHandler(
     closeAuthProgressIndicator()
 
     const params = new URLSearchParams(uri.query)
-    const token = params.get('code')
+    const token = params.get('code') || params.get('token')
     const endpoint = authProvider.instance!.status.endpoint
     if (!token || !endpoint) {
         return

vscode/src/chat/chat-view/ChatController.ts

@@ -471,9 +471,11 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
 
                     const authProviderSimplified = new AuthProviderSimplified()
                     const authMethod = message.authMethod || 'dotcom'
+                    const config = getConfigWithEndpoint()
                     const successfullyOpenedUrl = await authProviderSimplified.openExternalAuthUrl(
                         authMethod,
-                        tokenReceiverUrl
+                        tokenReceiverUrl,
+                        config?.agentIDE
                     )
                     if (!successfullyOpenedUrl) {
                         closeAuthProgressIndicator()

vscode/src/jsonrpc/agent-protocol.ts

@@ -503,6 +503,8 @@ export type ServerNotifications = {
     // When the when-claude context has changed.
     // For example, 'cody.activated' is set based on user's latest authentication status.
     'window/didChangeContext': [{ key: string; value?: string | undefined | null }]
+    // Client should move the focus to the sidebar.
+    'window/focusSidebar': [null]
 }
 
 export interface WebviewCreateWebviewPanelOptions {
@@ -605,6 +607,7 @@ export interface ClientInfo {
 
 // The capability should match the name of the JSON-RPC methods.
 export interface ClientCapabilities {
+    authentication?: 'enabled' | 'none' | undefined | null
     completions?: 'none' | undefined | null
     //  When 'streaming', handles 'chat/updateMessageInProgress' streaming notifications.
     chat?: 'none' | 'streaming' | undefined | null