Handle customer proxy re-auth response by retrying, not prompting user for different token

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

@@ -13,7 +13,7 @@ sealed class AuthenticationError {
     val deserializer: JsonDeserializer<AuthenticationError> =
       JsonDeserializer { element: JsonElement, _: Type, context: JsonDeserializationContext ->
         when (element.getAsJsonObject().get("type").getAsString()) {
-          "network-error" -> context.deserialize<NetworkAuthError>(element, NetworkAuthError::class.java)
+          "availability-error" -> context.deserialize<AvailabilityError>(element, AvailabilityError::class.java)
           "invalid-access-token" -> context.deserialize<InvalidAccessTokenError>(element, InvalidAccessTokenError::class.java)
           "enterprise-user-logged-into-dotcom" -> context.deserialize<EnterpriseUserDotComError>(element, EnterpriseUserDotComError::class.java)
           "auth-config-error" -> context.deserialize<AuthConfigError>(element, AuthConfigError::class.java)
@@ -23,12 +23,12 @@ sealed class AuthenticationError {
   }
 }
 
-data class NetworkAuthError(
-  val type: TypeEnum, // Oneof: network-error
+data class AvailabilityError(
+  val type: TypeEnum, // Oneof: Availability-error
 ) : AuthenticationError() {
 
   enum class TypeEnum {
-    @SerializedName("network-error") `Network-error`,
+    @SerializedName("availability-error") `Availability-error`,
   }
 }
 

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

@@ -60,7 +60,7 @@ object Constants {
   const val method = "method"
   const val multiple = "multiple"
   const val native = "native"
-  const val `network-error` = "network-error"
+  const val `availability-error` = "availability-error"
   const val none = "none"
   const val notification = "notification"
   const val `object-encoded` = "object-encoded"

lib/shared/src/auth/types.ts

@@ -49,8 +49,14 @@ export interface UnauthenticatedAuthStatus {
     pendingValidation: boolean
 }
 
-export interface NetworkAuthError {
-    type: 'network-error'
+/**
+ * An error representing the condition where the endpoint is not available due to lack of network
+ * connectivity, server downtime, or other configuration issues *unrelated to* the validity of the
+ * credentials.
+ */
+export interface AvailabilityError {
+    type: 'availability-error'
+    reason: string
 }
 
 export interface InvalidAccessTokenError {
@@ -67,7 +73,7 @@ export interface AuthConfigError extends AuthenticationErrorMessage {
 }
 
 export type AuthenticationError =
-    | NetworkAuthError
+    | AvailabilityError
     | InvalidAccessTokenError
     | EnterpriseUserDotComError
     | AuthConfigError
@@ -79,10 +85,10 @@ export interface AuthenticationErrorMessage {
 
 export function getAuthErrorMessage(error: AuthenticationError): AuthenticationErrorMessage {
     switch (error.type) {
-        case 'network-error':
+        case 'availability-error':
             return {
                 title: 'Network Error',
-                message: 'Cody is unreachable',
+                message: `Sourcegraph is unreachable: ${error.reason}`,
             }
         case 'invalid-access-token':
             return {

lib/shared/src/index.ts

@@ -252,6 +252,8 @@ export {
     isNetworkError,
     isNetworkLikeError,
     isRateLimitError,
+    NeedsAuthChallengeError,
+    isNeedsAuthChallengeError,
 } from './sourcegraph-api/errors'
 export {
     SourcegraphGraphQLAPIClient,
@@ -280,6 +282,7 @@ export {
     type NLSSearchDynamicFilter,
     type NLSSearchDynamicFilterKind,
     type GraphQLAPIClientConfig,
+    setJSONAcceptContentTypeHeaders,
 } from './sourcegraph-api/graphql/client'
 export type {
     CodyLLMSiteConfiguration,

lib/shared/src/sourcegraph-api/completions/browserClient.ts

@@ -128,6 +128,7 @@ export class SourcegraphBrowserCompletionsClient extends SourcegraphCompletionsC
         const { url, serializedParams } = await this.prepareRequest(params, requestParams)
         const headersInstance = new Headers({
             'Content-Type': 'application/json; charset=utf-8',
+            Accept: 'text/event-stream',
             ...configuration.customHeaders,
             ...requestParams.customHeaders,
         })

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

@@ -135,3 +135,19 @@ export function isNetworkLikeError(error: Error): boolean {
         message.includes('SELF_SIGNED_CERT_IN_CHAIN')
     )
 }
+
+/**
+ * An error indicating that the user needs to complete an authentication challenge.
+ */
+export class NeedsAuthChallengeError extends Error {
+    constructor() {
+        super(NeedsAuthChallengeError.DESCRIPTION)
+    }
+
+    /** A human-readable description of this error. */
+    static DESCRIPTION = 'user must complete authentication challenge'
+}
+
+export function isNeedsAuthChallengeError(error: Error): error is NeedsAuthChallengeError {
+    return error instanceof NeedsAuthChallengeError
+}

lib/shared/src/sourcegraph-api/graphql/client.test.ts

@@ -0,0 +1,55 @@
+import { describe, expect, it, vi } from 'vitest'
+import { SourcegraphGraphQLAPIClient } from '../..'
+import * as fetchModule from '../../fetch'
+import { NeedsAuthChallengeError } from '../errors'
+
+vi.mocked(fetchModule)
+
+describe('SourcegraphGraphQLClient', () => {
+    const client = SourcegraphGraphQLAPIClient.withStaticConfig({
+        auth: {
+            credentials: { token: 'test-token' },
+            serverEndpoint: 'https://test.sourcegraph.com',
+        },
+        clientState: { anonymousUserID: 'a' },
+        configuration: {
+            telemetryLevel: 'off',
+        },
+    })
+
+    describe('fetchHTTP', () => {
+        it('sets Accept header', async () => {
+            const fetchMock = vi
+                .spyOn(fetchModule, 'fetch')
+                .mockImplementation(async () =>
+                    Promise.resolve(new Response(JSON.stringify({ data: {} }), { status: 200 }))
+                )
+            await client.fetchHTTP('TestQuery', 'POST', '/graphql', '{}')
+
+            expect(fetchMock).toHaveBeenCalled()
+            const headers = fetchMock.mock.calls[0][1]?.headers as Headers
+            expect(headers.get('Accept')).toBe('application/json')
+
+            fetchMock.mockRestore()
+        })
+    })
+
+    describe('getCurrentUserInfo', () => {
+        it('returns NeedsAuthChallengeError when response requires auth challenge', async () => {
+            const fetchMock = vi.spyOn(fetchModule, 'fetch').mockImplementation(async () =>
+                Promise.resolve(
+                    new Response(null, {
+                        status: 401,
+                        headers: new Headers({
+                            'X-CustomerName-U2f-Challenge': 'test-challenge',
+                        }),
+                    })
+                )
+            )
+            const result = await client.getCurrentUserInfo()
+            expect(fetchMock).toHaveBeenCalled()
+            console.log('XX', result)
+            expect(result).toBeInstanceOf(NeedsAuthChallengeError)
+        })
+    })
+})

lib/shared/src/sourcegraph-api/graphql/client.ts

@@ -16,7 +16,7 @@ import { distinctUntilChanged, firstValueFrom } from '../../misc/observable'
 import { addTraceparent, wrapInActiveSpan } from '../../tracing'
 import { isError } from '../../utils'
 import { addCodyClientIdentificationHeaders } from '../client-name-version'
-import { isAbortError } from '../errors'
+import { NeedsAuthChallengeError, isAbortError, isNeedsAuthChallengeError } from '../errors'
 import { addAuthHeaders } from '../utils'
 import { type GraphQLResultCache, ObservableInvalidatedGraphQLResultCacheFactory } from './cache'
 import {
@@ -1612,7 +1612,6 @@ export class SourcegraphGraphQLAPIClient {
             })())
 
         const headers = new Headers(config.configuration?.customHeaders as HeadersInit | undefined)
-        headers.set('Content-Type', 'application/json; charset=utf-8')
         if (config.clientState.anonymousUserID && !process.env.CODY_WEB_DONT_SET_SOME_HEADERS) {
             headers.set('X-Sourcegraph-Actor-Anonymous-UID', config.clientState.anonymousUserID)
         }
@@ -1622,6 +1621,7 @@ export class SourcegraphGraphQLAPIClient {
         addTraceparent(headers)
         addCodyClientIdentificationHeaders(headers)
         addAuthHeaders(config.auth, headers, url)
+        setJSONAcceptContentTypeHeaders(headers)
 
         const { abortController, timeoutSignal } = dependentAbortControllerWithTimeout(signal)
         return wrapInActiveSpan(`httpapi.fetch${queryName ? `.${queryName}` : ''}`, () =>
@@ -1638,6 +1638,11 @@ export class SourcegraphGraphQLAPIClient {
     }
 }
 
+export function setJSONAcceptContentTypeHeaders(headers: Headers): void {
+    headers.set('Accept', 'application/json')
+    headers.set('Content-Type', 'application/json; charset=utf-8')
+}
+
 const DEFAULT_TIMEOUT_MSEC = 20000
 
 /**
@@ -1664,6 +1669,10 @@ function catchHTTPError(
     timeoutSignal: Pick<AbortSignal, 'aborted'>
 ): (error: any) => Error {
     return (error: any) => {
+        if (isNeedsAuthChallengeError(error)) {
+            return error
+        }
+
         // Throw the plain AbortError for intentional aborts so we handle them with call
         // flow, but treat timeout aborts as a network error (below) and include the
         // URL.
@@ -1687,12 +1696,27 @@ export async function verifyResponseCode(
     response: BrowserOrNodeResponse
 ): Promise<BrowserOrNodeResponse> {
     if (!response.ok) {
+        if (isCustomAuthChallengeResponse(response)) {
+            throw new NeedsAuthChallengeError()
+        }
         const body = await response.text()
         throw new Error(`HTTP status code ${response.status}${body ? `: ${body}` : ''}`)
     }
     return response
 }
 
+/**
+ * Some customers use an HTTP proxy in front of Sourcegraph that, when the user needs to complete an
+ * auth challenge, returns HTTP 401 with a `X-${CUSTOMER}-U2f-Challenge` header. Detect these kinds
+ * of error responses.
+ */
+function isCustomAuthChallengeResponse(response: BrowserOrNodeResponse): boolean {
+    return (
+        response.status === 401 &&
+        Array.from(response.headers.keys()).some(header => /^x-.*-u2f-challenge$/i.test(header))
+    )
+}
+
 function hasOutdatedAPIErrorMessages(error: Error): boolean {
     // Sourcegraph 5.2.3 returns an empty string ("") instead of an error message
     // when querying non-existent codyContextFilters; this produces

vscode/src/auth/auth.test.ts

@@ -0,0 +1,124 @@
+import {
+    CLIENT_CAPABILITIES_FIXTURE,
+    NeedsAuthChallengeError,
+    SourcegraphGraphQLAPIClient,
+    mockAuthStatus,
+    mockClientCapabilities,
+    mockResolvedConfig,
+} from '@sourcegraph/cody-shared'
+import { afterEach, beforeAll, beforeEach, describe, expect, it } from 'vitest'
+import { type MockInstance, vi } from 'vitest'
+import type * as vscode from 'vscode'
+import * as AuthProviderModule from '../services/AuthProvider'
+import { mockLocalStorage } from '../services/LocalStorageProvider'
+import { vsCodeMocks } from '../testutils/mocks'
+import { showSignInMenu, validateCredentials } from './auth'
+
+describe('showSignInMenu', () => {
+    beforeAll(() => {
+        mockResolvedConfig({
+            configuration: { serverEndpoint: 'https://example.com', overrideAuthToken: 'x' },
+            auth: { serverEndpoint: 'https://example.com' },
+        })
+        mockClientCapabilities(CLIENT_CAPABILITIES_FIXTURE)
+        mockAuthStatus({ authenticated: false })
+        mockLocalStorage('noop')
+        vi.spyOn<typeof vscode.window, 'showQuickPick'>(
+            vsCodeMocks.window as unknown as typeof vscode.window,
+            'showQuickPick'
+        ).mockImplementation(async () => ({ id: '1', label: 'l', uri: 'https://example.com' }))
+    })
+
+    it('does not show access token input box when authentication fails with availability error', async () => {
+        vi.spyOn(AuthProviderModule, 'authProvider', 'get').mockReturnValue({
+            validateAndStoreCredentials: vi
+                .fn<(typeof AuthProviderModule.authProvider)['validateAndStoreCredentials']>()
+                .mockResolvedValue({
+                    authenticated: false,
+                    error: { type: 'availability-error', reason: 'foo' },
+                    endpoint: 'https://example.com',
+                    pendingValidation: false,
+                }),
+        } satisfies Pick<
+            typeof AuthProviderModule.authProvider,
+            'validateAndStoreCredentials'
+        > as unknown as typeof AuthProviderModule.authProvider)
+        const mockShowInputBox = vi.spyOn<typeof vscode.window, 'showInputBox'>(
+            vsCodeMocks.window as unknown as typeof vscode.window,
+            'showInputBox'
+        )
+        const mockShowErrorMessage = vi.spyOn<typeof vscode.window, 'showErrorMessage'>(
+            vsCodeMocks.window as unknown as typeof vscode.window,
+            'showErrorMessage'
+        )
+
+        await showSignInMenu()
+        expect(mockShowInputBox).not.toHaveBeenCalled()
+        expect(mockShowErrorMessage).toHaveBeenCalledWith('Sourcegraph is unreachable: foo')
+    })
+
+    it('shows access token input box when authentication fails with invalid access token error', async () => {
+        vi.spyOn(AuthProviderModule, 'authProvider', 'get').mockReturnValue({
+            validateAndStoreCredentials: vi
+                .fn<(typeof AuthProviderModule.authProvider)['validateAndStoreCredentials']>()
+                .mockResolvedValue({
+                    authenticated: false,
+                    error: { type: 'invalid-access-token' },
+                    endpoint: 'https://example.com',
+                    pendingValidation: false,
+                }),
+        } satisfies Pick<
+            typeof AuthProviderModule.authProvider,
+            'validateAndStoreCredentials'
+        > as unknown as typeof AuthProviderModule.authProvider)
+        const mockShowInputBox = vi
+            .spyOn<typeof vscode.window, 'showInputBox'>(
+                vsCodeMocks.window as unknown as typeof vscode.window,
+                'showInputBox'
+            )
+            .mockResolvedValue('my-token')
+        const mockShowErrorMessage = vi.spyOn<typeof vscode.window, 'showErrorMessage'>(
+            vsCodeMocks.window as unknown as typeof vscode.window,
+            'showErrorMessage'
+        )
+
+        await showSignInMenu()
+        expect(mockShowInputBox).toHaveBeenCalled()
+        expect(mockShowErrorMessage).toHaveBeenCalledWith('The access token is invalid or has expired')
+    })
+})
+
+describe('validateCredentials', () => {
+    let getCurrentUserInfoSpy: MockInstance
+    beforeEach(() => {
+        getCurrentUserInfoSpy = vi.spyOn(SourcegraphGraphQLAPIClient.prototype, 'getCurrentUserInfo')
+    })
+
+    afterEach(() => {
+        vi.resetAllMocks()
+    })
+
+    it('returns unauthenticated status when NeedsAuthChallengeError occurs', async () => {
+        getCurrentUserInfoSpy.mockResolvedValue(new NeedsAuthChallengeError())
+
+        const result = await validateCredentials({
+            auth: {
+                serverEndpoint: 'https://sourcegraph.test',
+                credentials: { token: 'test-token' },
+            },
+            configuration: {
+                customHeaders: {},
+            },
+            clientState: {
+                anonymousUserID: 'test-user',
+            },
+        })
+
+        expect(result).toEqual({
+            authenticated: false,
+            error: { type: 'availability-error', reason: 'user must complete authentication challenge' },
+            endpoint: 'https://sourcegraph.test',
+            pendingValidation: false,
+        })
+    })
+})