fix client disconnect (#134)

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>

Author: samuelcolvin

gateway/src/auth.ts

@@ -1,7 +1,7 @@
 import type { GatewayOptions } from '.'
 import type { RateLimiter } from './rateLimiter'
 import type { ApiKeyInfo } from './types'
-import { ResponseError, runAfter } from './utils'
+import { runAfter, textResponse } from './utils'
 
 const CACHE_TTL = 86400 * 30
 
@@ -10,12 +10,12 @@ export async function apiKeyAuth(
   ctx: ExecutionContext,
   options: GatewayOptions,
   rateLimiter: RateLimiter,
-): Promise<ApiKeyInfo> {
+): Promise<ApiKeyInfo | Response> {
   const authorization = request.headers.get('authorization')
   const xApiKey = request.headers.get('x-api-key')
 
   if (authorization && xApiKey) {
-    throw new ResponseError(401, 'Unauthorized - Both Authorization and X-API-Key headers are present, use only one')
+    return textResponse(401, 'Unauthorized - Both Authorization and X-API-Key headers are present, use only one')
   }
 
   const authHeader = authorization || xApiKey
@@ -28,11 +28,11 @@ export async function apiKeyAuth(
       key = authHeader
     }
   } else {
-    throw new ResponseError(401, 'Unauthorized - Missing Authorization Header')
+    return textResponse(401, 'Unauthorized - Missing Authorization Header')
   }
   // avoid very long queries to the DB
   if (key.length > 100) {
-    throw new ResponseError(401, 'Unauthorized - Key too long')
+    return textResponse(401, 'Unauthorized - Key too long')
   }
 
   const cacheKey = apiKeyCacheKey(key, options.kvVersion)
@@ -46,7 +46,10 @@ export async function apiKeyAuth(
       options.kv.get(projectStateCacheKey(apiKeyInfo.project, options.kvVersion)),
       rateLimiter.requestStart(apiKeyInfo),
     ])
-    processLimiterResult(limiterResult)
+    const limiterResponse = processLimiterResult(limiterResult)
+    if (limiterResponse) {
+      return limiterResponse
+    }
     // we only return a cache match if the project state is the same, so updating the project state invalidates the cache
     // projectState is null if we have never invalidated the cache which will only be true for the first request after a deployment
     if (projectState === null || projectState === cacheResult.metadata) {
@@ -58,12 +61,16 @@ export async function apiKeyAuth(
   const apiKeyInfo = await options.keysDb.getApiKey(key)
   if (apiKeyInfo) {
     if (!rateLimiterStarted) {
-      processLimiterResult(await rateLimiter.requestStart(apiKeyInfo))
+      const limiterResult = await rateLimiter.requestStart(apiKeyInfo)
+      const limiterResponse = processLimiterResult(limiterResult)
+      if (limiterResponse) {
+        return limiterResponse
+      }
     }
     runAfter(ctx, 'setApiKeyCache', setApiKeyCache(apiKeyInfo, options))
     return apiKeyInfo
   }
-  throw new ResponseError(401, 'Unauthorized - Key not found')
+  return textResponse(401, 'Unauthorized - Key not found')
 }
 
 export async function setApiKeyCache(
@@ -99,6 +106,6 @@ const projectStateCacheKey = (project: number, kvVersion: string) => `projectSta
 
 function processLimiterResult(limiterResult: string | null) {
   if (typeof limiterResult === 'string') {
-    throw new ResponseError(429, limiterResult)
+    return textResponse(429, limiterResult)
   }
 }

gateway/src/gateway.ts

@@ -35,7 +35,11 @@ export async function gateway(
   }
 
   const rateLimiter = options.rateLimiter ?? noopLimiter
-  const apiKeyInfo = await apiKeyAuth(request, ctx, options, rateLimiter)
+  const authResult = await apiKeyAuth(request, ctx, options, rateLimiter)
+  if (authResult instanceof Response) {
+    return authResult
+  }
+  const apiKeyInfo = authResult
   try {
     return await gatewayWithLimiter(request, restOfPath, apiType, apiKeyInfo, ctx, options)
   } finally {

gateway/src/index.ts

@@ -20,7 +20,7 @@ import { gateway } from './gateway'
 import type { DefaultProviderProxy, Middleware, Next } from './providers/default'
 import type { RateLimiter } from './rateLimiter'
 import type { SubFetch } from './types'
-import { ctHeader, ResponseError, response405, textResponse } from './utils'
+import { ctHeader, response405, runAfter, textResponse } from './utils'
 
 export { changeProjectState as setProjectState, deleteApiKeyCache, setApiKeyCache } from './auth'
 export type { DefaultProviderProxy, Middleware, Next }
@@ -56,15 +56,13 @@ export async function gatewayFetch(
     if (proxyPath === '/') {
       return index(request, options)
     } else {
-      return await gateway(request, `${proxyPath}${queryString}`, ctx, options)
+      const gatewayPromise = gateway(request, `${proxyPath}${queryString}`, ctx, options)
+      runAfter(ctx, 'gatewayPromise', gatewayPromise)
+      return await gatewayPromise
     }
   } catch (error) {
-    if (error instanceof ResponseError) {
-      logfire.reportError('ResponseError', error)
-      return error.response()
-    } else {
-      throw error
-    }
+    logfire.reportError('ResponseError', error as Error)
+    throw error
   }
 }
 

gateway/src/providers/anthropic.ts

@@ -1,6 +1,6 @@
 import type { ModelAPI } from '../api'
 import { AnthropicAPI } from '../api/anthropic'
-import { DefaultProviderProxy } from './default'
+import { DefaultProviderProxy, type ProxyInvalidRequest } from './default'
 
 export class AnthropicProvider extends DefaultProviderProxy {
   protected isWhitelistedEndpoint(): boolean {
@@ -15,8 +15,9 @@ export class AnthropicProvider extends DefaultProviderProxy {
   }
 
   // biome-ignore lint/suspicious/useAwait: required by google auth
-  protected async requestHeaders(headers: Headers): Promise<void> {
+  protected async requestHeaders(headers: Headers): Promise<ProxyInvalidRequest | null> {
     headers.set('x-api-key', this.providerProxy.credentials)
+    return null
   }
 
   protected responseHeaders(headers: Headers): Headers {

gateway/src/providers/default.ts

@@ -186,8 +186,9 @@ export class DefaultProviderProxy {
   }
 
   // biome-ignore lint/suspicious/useAwait: required by google auth
-  protected async requestHeaders(headers: Headers): Promise<void> {
+  protected async requestHeaders(headers: Headers): Promise<ProxyInvalidRequest | null> {
     headers.set('Authorization', `Bearer ${this.providerProxy.credentials}`)
+    return null
   }
 
   protected async prepRequest(): Promise<Prepare | ProxyInvalidRequest> {
@@ -289,7 +290,10 @@ export class DefaultProviderProxy {
     requestHeaders.set('user-agent', this.userAgent())
     // authorization header was used by the gateway auth, it definitely should not be forwarded to the target api
     requestHeaders.delete('authorization')
-    await this.requestHeaders(requestHeaders)
+    const requestHeadersError = await this.requestHeaders(requestHeaders)
+    if (requestHeadersError) {
+      return requestHeadersError
+    }
 
     const prepResult = await this.prepRequest()
     if ('error' in prepResult) {

gateway/src/providers/google/auth.ts

@@ -1,43 +1,48 @@
-import { ResponseError } from '../../utils'
+import type { ProxyInvalidRequest } from '../default'
 
-export async function authToken(credentials: string, kv: KVNamespace, subFetch: typeof fetch): Promise<string> {
+export async function authToken(
+  credentials: string,
+  kv: KVNamespace,
+  subFetch: typeof fetch,
+): Promise<{ token: string } | ProxyInvalidRequest> {
   const serviceAccountHash = await hash(credentials)
   const cacheKey = `gcp-auth:${serviceAccountHash}`
   const cachedToken = await kv.get(cacheKey, { cacheTtl: 300 })
   if (cachedToken) {
-    return cachedToken
+    return { token: cachedToken }
   }
-  const serviceAccount = getServiceAccount(credentials)
-  const jwt = await jwtSign(serviceAccount)
-  const token = await getAccessToken(jwt, subFetch)
-  await kv.put(cacheKey, token, { expirationTtl: 3000 })
-  return token
+  const serviceAccountResult = getServiceAccount(credentials)
+  if ('error' in serviceAccountResult) {
+    return serviceAccountResult
+  }
+  const jwt = await jwtSign(serviceAccountResult)
+  const tokenResult = await getAccessToken(jwt, subFetch)
+  if ('error' in tokenResult) {
+    return tokenResult
+  }
+  await kv.put(cacheKey, tokenResult.token, { expirationTtl: 3000 })
+  return tokenResult
 }
 
-function getServiceAccount(credentials: string): ServiceAccount {
+export function getServiceAccount(credentials: string): ServiceAccount | ProxyInvalidRequest {
   let sa: ServiceAccount
   try {
     sa = JSON.parse(credentials)
   } catch (error) {
-    throw new ResponseError(400, `provider credentials are not valid JSON: ${error as Error}`)
+    return { error: `provider credentials are not valid JSON: ${error as Error}` }
   }
   if (typeof sa.client_email !== 'string') {
-    throw new ResponseError(400, `"client_email" should be a string, not ${typeof sa.client_email}`)
+    return { error: `"client_email" should be a string, not ${typeof sa.client_email}` }
   }
   if (typeof sa.private_key !== 'string') {
-    throw new ResponseError(400, `"private_key" should be a string, not ${typeof sa.private_key}`)
+    return { error: `"private_key" should be a string, not ${typeof sa.private_key}` }
   }
   if (typeof sa.project_id !== 'string') {
-    throw new ResponseError(400, `"project_id" should be a string, not ${typeof sa.project_id}`)
+    return { error: `"project_id" should be a string, not ${typeof sa.project_id}` }
   }
   return { client_email: sa.client_email, private_key: sa.private_key, project_id: sa.project_id }
 }
 
-export function getProjectId(credentials: string): string {
-  const sa = getServiceAccount(credentials)
-  return sa.project_id
-}
-
 interface ServiceAccount {
   client_email: string
   private_key: string
@@ -80,7 +85,7 @@ async function jwtSign(serviceAccount: ServiceAccount): Promise<string> {
   return `${signingInput}.${b64UrlEncodeArray(signature)}`
 }
 
-async function getAccessToken(jwt: string, subFetch: typeof fetch): Promise<string> {
+async function getAccessToken(jwt: string, subFetch: typeof fetch): Promise<{ token: string } | ProxyInvalidRequest> {
   const body = new URLSearchParams({ grant_type: 'urn:ietf:params:oauth:grant-type:jwt-bearer', assertion: jwt })
 
   const response = await subFetch(tokenUrl, {
@@ -92,10 +97,10 @@ async function getAccessToken(jwt: string, subFetch: typeof fetch): Promise<stri
 
   if (response.ok) {
     const responseData: TokenResponse = await response.json()
-    return responseData.access_token
+    return { token: responseData.access_token }
   } else {
     const text = await response.text()
-    throw new ResponseError(400, `Failed to get GCP access token, response:\n${response.status}: ${text}`)
+    return { error: `Failed to get GCP access token, response:\n${response.status}: ${text}` }
   }
 }
 

gateway/src/providers/google/index.ts

@@ -1,16 +1,20 @@
 import type { ModelAPI } from '../../api'
 import { AnthropicAPI } from '../../api/anthropic'
 import { GoogleAPI, type GoogleRequest } from '../../api/google'
-import { DefaultProviderProxy } from '../default'
-import { authToken, getProjectId } from './auth'
+import { DefaultProviderProxy, type ProxyInvalidRequest } from '../default'
+import { authToken, getServiceAccount } from './auth'
 
 export class GoogleVertexProvider extends DefaultProviderProxy {
   protected usageField = 'usageMetadata'
   flavor: 'default' | 'anthropic' = 'default'
 
   url() {
     if (this.providerProxy.baseUrl) {
-      const projectId = getProjectId(this.providerProxy.credentials)
+      const serviceAccountResult = getServiceAccount(this.providerProxy.credentials)
+      if ('error' in serviceAccountResult) {
+        return serviceAccountResult
+      }
+      const projectId = serviceAccountResult.project_id
       const region = regionFromUrl(this.providerProxy.baseUrl)
       if (!region) {
         return { error: 'Unable to extract region from URL' }
@@ -92,9 +96,14 @@ export class GoogleVertexProvider extends DefaultProviderProxy {
     }
   }
 
-  async requestHeaders(headers: Headers): Promise<void> {
-    const token = await authToken(this.providerProxy.credentials, this.options.kv, this.options.subFetch)
-    headers.set('Authorization', `Bearer ${token}`)
+  async requestHeaders(headers: Headers): Promise<ProxyInvalidRequest | null> {
+    const tokenResult = await authToken(this.providerProxy.credentials, this.options.kv, this.options.subFetch)
+    if ('error' in tokenResult) {
+      return tokenResult
+    } else {
+      headers.set('Authorization', `Bearer ${tokenResult.token}`)
+      return null
+    }
   }
 }
 

gateway/src/utils.ts

@@ -30,21 +30,6 @@ export function getIP(request: Request): string {
   }
 }
 
-export class ResponseError extends Error {
-  status: number
-  message: string
-
-  constructor(status: number, message: string) {
-    super(message)
-    this.status = status
-    this.message = message
-  }
-
-  response(): Response {
-    return textResponse(this.status, this.message)
-  }
-}
-
 export function runAfter(ctx: ExecutionContext, name: string, promise: Promise<unknown>) {
   ctx.waitUntil(wrapLogfire(name, promise))
 }

gateway/test/auth.spec.ts

@@ -25,6 +25,24 @@ class CountingKeysDb implements KeysDb {
   }
 }
 
+describe('apiKeyAuth fails', () => {
+  test('no header', async () => {
+    const ctx = createExecutionContext()
+    const baseOptions = buildGatewayEnv(env, [], fetch)
+    const countingDb = new CountingKeysDb(baseOptions.keysDb)
+    const options = { ...baseOptions, keysDb: countingDb }
+
+    const request = new Request('https://example.com')
+
+    // First call should fetch from DB
+    const result = await apiKeyAuth(request, ctx, options, noopLimiter)
+    expect(result).instanceOf(Response)
+    const response = result as Response
+    expect(response.status).toBe(401)
+    expect(await response.text()).toEqual('Unauthorized - Missing Authorization Header')
+  })
+})
+
 describe('apiKeyAuth cache invalidation', () => {
   test('caches api key and returns cached value', async () => {
     const ctx = createExecutionContext()
@@ -36,7 +54,7 @@ describe('apiKeyAuth cache invalidation', () => {
 
     // First call should fetch from DB
     const apiKey1 = await apiKeyAuth(request, ctx, options, noopLimiter)
-    expect(apiKey1.key).toBe('healthy')
+    expect((apiKey1 as ApiKeyInfo).key).toBe('healthy')
     // Wait for cache to be set (it's set asynchronously via runAfter)
     await waitOnExecutionContext(ctx)
     expect(countingDb.callCount).toBe(1)
@@ -48,7 +66,7 @@ describe('apiKeyAuth cache invalidation', () => {
     // Second call should use cache, not hit DB
     const ctx2 = createExecutionContext()
     const apiKey2 = await apiKeyAuth(request, ctx2, options, noopLimiter)
-    expect(apiKey2.key).toBe('healthy')
+    expect((apiKey2 as ApiKeyInfo).key).toBe('healthy')
 
     expect(countingDb.callCount).toBe(1)
   })
@@ -85,7 +103,7 @@ describe('apiKeyAuth cache invalidation', () => {
     // Third call - cache is invalidated, should hit DB again
     const ctx3 = createExecutionContext()
     const apiKey3 = await apiKeyAuth(request, ctx3, options, noopLimiter)
-    expect(apiKey3.key).toBe('healthy')
+    expect((apiKey3 as ApiKeyInfo).key).toBe('healthy')
     await waitOnExecutionContext(ctx3)
 
     expect(countingDb.callCount).toBe(2)