Handle customer proxy re-auth response by retrying, not prompting use…

…r for different token

Fixes https://linear.app/sourcegraph/issue/CODY-4695/handle-customer-proxy-re-auth-response-by-retrying-not-prompting-user

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

@@ -1,4 +1,5 @@
 import { isDotCom } from '../sourcegraph-api/environments'
+import { NeedsAuthChallengeError } from '../sourcegraph-api/errors'
 import type { UserProductSubscription } from '../sourcegraph-api/userProductSubscription'
 
 /**
@@ -49,8 +50,19 @@ 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'
+
+    /**
+     * Whether the error is due to a proxy needing the user to complete an auth challenge. See
+     * {@link NeedsAuthChallengeError}.
+     */
+    needsAuthChallenge?: boolean
 }
 
 export interface InvalidAccessTokenError {
@@ -67,23 +79,27 @@ export interface AuthConfigError extends AuthenticationErrorMessage {
 }
 
 export type AuthenticationError =
-    | NetworkAuthError
+    | AvailabilityError
     | InvalidAccessTokenError
     | EnterpriseUserDotComError
     | AuthConfigError
 
 export interface AuthenticationErrorMessage {
     title?: string
     message: string
+    showTryAgain?: boolean
 }
 
 export function getAuthErrorMessage(error: AuthenticationError): AuthenticationErrorMessage {
     switch (error.type) {
-        case 'network-error':
-            return {
-                title: 'Network Error',
-                message: 'Cody is unreachable',
-            }
+        case 'availability-error':
+            return error.needsAuthChallenge
+                ? NeedsAuthChallengeError.TITLE_AND_MESSAGE
+                : {
+                      title: 'Network Error',
+                      message: 'Sourcegraph is unreachable',
+                      showTryAgain: true,
+                  }
         case 'invalid-access-token':
             return {
                 title: 'Invalid Access Token',

lib/shared/src/index.ts

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

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

@@ -107,7 +107,7 @@ export class ClientConfigSingleton {
                           // Don't update if the editor is in the background, to avoid network
                           // activity that can cause OS warnings or authorization flows when the
                           // user is not using Cody. See
-                          // linear.app/sourcegraph/issue/CODY-3745/codys-background-periodic-network-access-causes-2fa.
+                          // https://linear.app/sourcegraph/issue/CODY-3745/codys-background-periodic-network-access-causes-2fa.
                           filter((_value): _value is undefined => editorWindowIsFocused()),
                           startWith(undefined),
                           switchMap(() => promiseFactoryToObservable(signal => this.fetchConfig(signal)))

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

@@ -6,6 +6,7 @@ import { isError } from '../../utils'
 import { addClientInfoParams, addCodyClientIdentificationHeaders } from '../client-name-version'
 import { addAuthHeaders } from '../utils'
 
+import { verifyResponseCode } from '../graphql/client'
 import { CompletionsResponseBuilder } from './CompletionsResponseBuilder'
 import { type CompletionRequestParameters, SourcegraphCompletionsClient } from './client'
 import { parseCompletionJSON } from './parse'
@@ -60,18 +61,7 @@ export class SourcegraphBrowserCompletionsClient extends SourcegraphCompletionsC
             openWhenHidden: isRunningInWebWorker, // otherwise tries to call document.addEventListener
             async onopen(response) {
                 if (!response.ok && response.headers.get('content-type') !== 'text/event-stream') {
-                    let errorMessage: null | string = null
-                    try {
-                        errorMessage = await response.text()
-                    } catch (error) {
-                        // We show the generic error message in this case
-                        console.error(error)
-                    }
-                    const error = new Error(
-                        errorMessage === null || errorMessage.length === 0
-                            ? `Request failed with status code ${response.status}`
-                            : errorMessage
-                    )
+                    const error = await verifyResponseCode(response).catch(err => err)
                     cb.onError(error, response.status)
                     abort.abort()
                     return
@@ -128,6 +118,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,
         })
@@ -143,15 +134,7 @@ export class SourcegraphBrowserCompletionsClient extends SourcegraphCompletionsC
                 headers: headersInstance,
                 body: JSON.stringify(serializedParams),
                 signal,
-            })
-            if (!response.ok) {
-                const errorMessage = await response.text()
-                throw new Error(
-                    errorMessage.length === 0
-                        ? `Request failed with status code ${response.status}`
-                        : errorMessage
-                )
-            }
+            }).then(verifyResponseCode)
             const data = (await response.json()) as CompletionResponse
             if (data?.completion) {
                 cb.onChange(data.completion)

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

@@ -2,6 +2,7 @@ import { differenceInDays, format, formatDistanceStrict, formatRelative } from '
 
 import { isError } from '../utils'
 
+import type { AuthenticationErrorMessage } from '../auth/types'
 import type { BrowserOrNodeResponse } from './graphql/client'
 
 function formatRetryAfterDate(retryAfterDate: Date): string {
@@ -135,3 +136,30 @@ 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 =
+        `Tap your YubiKey to re-authenticate. (Your device's authentication expired and must be renewed to access Sourcegraph on your organization's network.)`
+
+    /** The same, but separated into title and message. */
+    static TITLE_AND_MESSAGE: AuthenticationErrorMessage = {
+        // See
+        // https://linear.app/sourcegraph/issue/CODY-4695/handle-customer-proxy-re-auth-response-by-retrying-not-prompting-user
+        // for an explanation of this message. If you need to change it to something more general,
+        // consult the customers mentioned in that issue.
+        title: 'Tap Your YubiKey to Authenticate...',
+        message: `Your device's authentication expired and must be renewed to access Sourcegraph on your organization's network.`,
+    }
+}
+
+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': 'true',
+                        }),
+                    })
+                )
+            )
+            const result = await client.getCurrentUserInfo()
+            expect(fetchMock).toHaveBeenCalled()
+            console.log('XX', result)
+            expect(result).toBeInstanceOf(NeedsAuthChallengeError)
+        })
+    })
+})

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

@@ -4,6 +4,7 @@ import { fetch } from '../../fetch'
 
 import type { TelemetryEventInput } from '@sourcegraph/telemetry'
 
+import type { IncomingMessage } from 'http'
 import escapeRegExp from 'lodash/escapeRegExp'
 import isEqual from 'lodash/isEqual'
 import omit from 'lodash/omit'
@@ -16,7 +17,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 +1613,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 +1622,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 +1639,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 +1670,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.
@@ -1686,13 +1696,55 @@ export const graphqlClient = SourcegraphGraphQLAPIClient.withGlobalConfig()
 export async function verifyResponseCode(
     response: BrowserOrNodeResponse
 ): Promise<BrowserOrNodeResponse> {
+    if (isCustomAuthChallengeResponse(response)) {
+        throw new NeedsAuthChallengeError()
+    }
+
     if (!response.ok) {
-        const body = await response.text()
+        let body: string | undefined
+        try {
+            body = await response.text()
+        } catch {}
         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: true` header. Detect these
+ * kinds of error responses.
+ */
+export function isCustomAuthChallengeResponse(
+    response:
+        | Pick<BrowserOrNodeResponse, 'status' | 'headers'>
+        | Pick<IncomingMessage, 'httpVersion' | 'statusCode' | 'headers'>
+): boolean {
+    ///////////////////////////////////
+    // TODO!(sqs): remove
+    ///////////////////////////////////
+    if (require('node:fs').existsSync('/tmp/is-custom-auth-challenge-error')) {
+        return true
+    }
+
+    function isIncomingMessageType(
+        v: typeof response
+    ): v is Pick<IncomingMessage, 'httpVersion' | 'statusCode' | 'headers'> {
+        return 'httpVersion' in response
+    }
+    const statusCode = isIncomingMessageType(response) ? response.statusCode : response.status
+    if (statusCode === 401) {
+        const headerEntries = isIncomingMessageType(response)
+            ? Object.entries(response.headers)
+            : Array.from(response.headers.entries())
+        return headerEntries.some(
+            ([name, value]) => /^x-.*-u2f-challenge$/i.test(name) && value === 'true'
+        )
+    }
+
+    return false
+}
+
 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