add getUserInfoFromApiKey to AgentRuntimeDeps

Author: charleslien

backend/src/__tests__/request-files-prompt.test.ts

@@ -9,7 +9,6 @@ import {
   beforeAll,
   beforeEach,
   mock as bunMockFn,
-  spyOn as bunSpyOn,
   describe,
   expect,
   it,
@@ -63,8 +62,6 @@ describe('requestRelevantFiles', () => {
   const mockCostMode: CostMode = 'normal'
   const mockRepoId = 'owner/repo'
 
-  let getCustomFilePickerConfigForOrgSpy: any // Explicitly typed as any
-
   beforeAll(() => {
     mockModule('@codebuff/backend/llm-apis/gemini-with-fallbacks', () => ({
       promptFlashWithFallbacks: bunMockFn(() =>
@@ -100,21 +97,6 @@ describe('requestRelevantFiles', () => {
   beforeEach(() => {
     agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
 
-    // If the spy was created in a previous test, restore it
-    if (
-      getCustomFilePickerConfigForOrgSpy &&
-      typeof getCustomFilePickerConfigForOrgSpy.mockRestore === 'function'
-    ) {
-      getCustomFilePickerConfigForOrgSpy.mockRestore()
-      getCustomFilePickerConfigForOrgSpy = undefined
-    }
-
-    // Use the directly imported bunSpyOn
-    getCustomFilePickerConfigForOrgSpy = bunSpyOn(
-      OriginalRequestFilesPromptModule,
-      'getCustomFilePickerConfigForOrg',
-    ).mockResolvedValue(null)
-
     const promptFlashWithFallbacksMock =
       geminiWithFallbacksModule.promptFlashWithFallbacks as Mock<
         typeof geminiWithFallbacksModule.promptFlashWithFallbacks
@@ -140,7 +122,6 @@ describe('requestRelevantFiles', () => {
     expect(
       geminiWithFallbacksModule.promptFlashWithFallbacks,
     ).toHaveBeenCalled()
-    expect(getCustomFilePickerConfigForOrgSpy).toHaveBeenCalled()
   })
 
   it('should use custom file counts from config', async () => {
@@ -149,7 +130,6 @@ describe('requestRelevantFiles', () => {
       customFileCounts: { normal: 5 },
       maxFilesPerRequest: 10,
     }
-    getCustomFilePickerConfigForOrgSpy!.mockResolvedValue(customConfig as any)
 
     await OriginalRequestFilesPromptModule.requestRelevantFiles({
       ...agentRuntimeImpl,
@@ -167,15 +147,13 @@ describe('requestRelevantFiles', () => {
     expect(
       geminiWithFallbacksModule.promptFlashWithFallbacks,
     ).toHaveBeenCalled()
-    expect(getCustomFilePickerConfigForOrgSpy).toHaveBeenCalled()
   })
 
   it('should use custom maxFilesPerRequest from config', async () => {
     const customConfig = {
       modelName: 'ft_filepicker_005',
       maxFilesPerRequest: 3,
     }
-    getCustomFilePickerConfigForOrgSpy!.mockResolvedValue(customConfig as any)
 
     const result = await OriginalRequestFilesPromptModule.requestRelevantFiles({
       ...agentRuntimeImpl,
@@ -194,14 +172,12 @@ describe('requestRelevantFiles', () => {
     if (result) {
       expect(result.length).toBeLessThanOrEqual(3)
     }
-    expect(getCustomFilePickerConfigForOrgSpy).toHaveBeenCalled()
   })
 
   it('should use custom modelName from config', async () => {
     const customConfig = {
       modelName: 'ft_filepicker_010',
     }
-    getCustomFilePickerConfigForOrgSpy!.mockResolvedValue(customConfig as any)
 
     await OriginalRequestFilesPromptModule.requestRelevantFiles({
       ...agentRuntimeImpl,
@@ -223,14 +199,12 @@ describe('requestRelevantFiles', () => {
         useFinetunedModel: finetunedVertexModels.ft_filepicker_010,
       }),
     )
-    expect(getCustomFilePickerConfigForOrgSpy).toHaveBeenCalled()
   })
 
   it('should use default model if custom modelName is invalid', async () => {
     const customConfig = {
       modelName: 'invalid-model-name',
     }
-    getCustomFilePickerConfigForOrgSpy!.mockResolvedValue(customConfig as any)
 
     await OriginalRequestFilesPromptModule.requestRelevantFiles({
       ...agentRuntimeImpl,
@@ -253,6 +227,5 @@ describe('requestRelevantFiles', () => {
         useFinetunedModel: expectedModel,
       }),
     )
-    expect(getCustomFilePickerConfigForOrgSpy).toHaveBeenCalled()
   })
 })

backend/src/api/org.ts

@@ -4,7 +4,7 @@ import { z } from 'zod/v4'
 
 import { extractAuthTokenFromHeader } from '../util/auth-helpers'
 import { logger } from '../util/logger'
-import { getUserIdFromAuthToken } from '../websockets/websocket-action'
+import { getUserInfoFromApiKey } from '../websockets/auth'
 
 import type {
   Request as ExpressRequest,
@@ -35,7 +35,9 @@ async function isRepoCoveredHandler(
         .status(401)
         .json({ error: 'Missing x-codebuff-api-key header' })
     }
-    const userId = await getUserIdFromAuthToken({ authToken })
+    const userId = (
+      await getUserInfoFromApiKey({ apiKey: authToken, fields: ['id'] })
+    )?.id
 
     if (!userId) {
       return res.status(401).json({ error: INVALID_AUTH_TOKEN_MESSAGE })

backend/src/api/usage.ts

@@ -1,12 +1,10 @@
 import { getOrganizationUsageResponse } from '@codebuff/billing'
-import db from '@codebuff/common/db'
-import * as schema from '@codebuff/common/db/schema'
 import { INVALID_AUTH_TOKEN_MESSAGE } from '@codebuff/common/old-constants'
-import { eq } from 'drizzle-orm'
 import { z } from 'zod/v4'
 
 import { checkAuth } from '../util/check-auth'
 import { logger } from '../util/logger'
+import { getUserInfoFromApiKey } from '../websockets/auth'
 import { genUsageResponse } from '../websockets/websocket-action'
 
 import type {
@@ -21,18 +19,6 @@ const usageRequestSchema = z.object({
   orgId: z.string().optional(),
 })
 
-async function getUserIdFromAuthToken(
-  token: string,
-): Promise<string | undefined> {
-  const user = await db
-    .select({ userId: schema.user.id })
-    .from(schema.user)
-    .innerJoin(schema.session, eq(schema.user.id, schema.session.userId))
-    .where(eq(schema.session.sessionToken, token))
-    .then((users) => users[0]?.userId)
-  return user
-}
-
 async function usageHandler(
   req: ExpressRequest,
   res: ExpressResponse,
@@ -59,7 +45,7 @@ async function usageHandler(
     }
 
     const userId = authToken
-      ? await getUserIdFromAuthToken(authToken)
+      ? (await getUserInfoFromApiKey({ apiKey: authToken, fields: ['id'] }))?.id
       : undefined
 
     if (!userId) {

backend/src/impl/agent-runtime.ts

@@ -5,11 +5,13 @@ import {
   promptAiSdkStructured,
 } from '../llm-apis/vercel-ai-sdk/ai-sdk'
 import { logger } from '../util/logger'
+import { getUserInfoFromApiKey } from '../websockets/auth'
 
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 
 export const BACKEND_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
   // Database
+  getUserInfoFromApiKey,
   startAgentRun,
   finishAgentRun,
   addAgentStep,

backend/src/websockets/auth.ts

@@ -2,40 +2,29 @@ import db from '@codebuff/common/db'
 import * as schema from '@codebuff/common/db/schema'
 import { eq } from 'drizzle-orm'
 
-export interface UserInfo {
-  id: string
-  email: string
-  discord_id: string | null
-}
+import type {
+  GetUserInfoFromApiKeyInput,
+  GetUserInfoFromApiKeyOutput,
+  UserColumn,
+} from '@codebuff/common/types/contracts/database'
 
-export async function getUserIdFromAuthToken(
-  authToken: string,
-): Promise<string | undefined> {
-  const user = await db
-    .select({ id: schema.user.id })
-    .from(schema.user)
-    .leftJoin(schema.session, eq(schema.user.id, schema.session.userId))
-    .where(eq(schema.session.sessionToken, authToken))
-    .limit(1)
-    .then((rows) => rows[0])
+export async function getUserInfoFromApiKey<T extends UserColumn>(
+  params: GetUserInfoFromApiKeyInput<T>,
+): GetUserInfoFromApiKeyOutput<T> {
+  const { apiKey, fields } = params
 
-  return user?.id
-}
+  // Build a typed selection object for user columns
+  const userSelection = Object.fromEntries(
+    fields.map((field) => [field, schema.user[field]]),
+  ) as { [K in T]: (typeof schema.user)[K] }
 
-export async function getUserInfoFromAuthToken(
-  authToken: string,
-): Promise<UserInfo | undefined> {
-  const user = await db
-    .select({
-      id: schema.user.id,
-      email: schema.user.email,
-      discord_id: schema.user.discord_id,
-    })
+  const rows = await db
+    .select({ user: userSelection }) // <-- important: nest under 'user'
     .from(schema.user)
     .leftJoin(schema.session, eq(schema.user.id, schema.session.userId))
-    .where(eq(schema.session.sessionToken, authToken))
+    .where(eq(schema.session.sessionToken, apiKey))
     .limit(1)
-    .then((rows) => rows[0])
 
-  return user
+  // Drizzle returns { user: ..., session: ... }, we return only the user part
+  return rows[0]?.user ?? null
 }

backend/src/websockets/middleware.ts

@@ -12,24 +12,24 @@ import * as schema from '@codebuff/common/db/schema'
 import { pluralize } from '@codebuff/common/util/string'
 import { eq } from 'drizzle-orm'
 
-import { getUserInfoFromAuthToken } from './auth'
+import { getUserInfoFromApiKey } from './auth'
 import { updateRequestContext } from './request-context'
 import { sendAction } from './websocket-action'
 import { withAppContext } from '../context/app-context'
 import { BACKEND_AGENT_RUNTIME_IMPL } from '../impl/agent-runtime'
 import { checkAuth } from '../util/check-auth'
 
-import type { UserInfo } from './auth'
 import type { ClientAction, ServerAction } from '@codebuff/common/actions'
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
+import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { WebSocket } from 'ws'
 
 type MiddlewareCallback = (params: {
   action: ClientAction
   clientSessionId: string
   ws: WebSocket
-  userInfo: UserInfo | undefined
+  userInfo: { id: string } | null
   logger: Logger
 }) => Promise<void | ServerAction>
 
@@ -64,7 +64,7 @@ export class WebSocketMiddleware {
       action: ClientAction<T>
       clientSessionId: string
       ws: WebSocket
-      userInfo: UserInfo | undefined
+      userInfo: { id: string } | null
       logger: Logger
     }) => Promise<void | ServerAction>,
   ) {
@@ -76,14 +76,18 @@ export class WebSocketMiddleware {
     clientSessionId: string
     ws: WebSocket
     silent?: boolean
+    getUserInfoFromApiKey: GetUserInfoFromApiKeyFn
     logger: Logger
   }): Promise<boolean> {
     const { action, clientSessionId, ws, silent, logger } = params
 
     const userInfo =
       'authToken' in action && action.authToken
-        ? await getUserInfoFromAuthToken(action.authToken)
-        : undefined
+        ? await getUserInfoFromApiKey({
+            apiKey: action.authToken,
+            fields: ['id'],
+          })
+        : null
 
     for (const middleware of this.middlewares) {
       const actionOrContinue = await middleware({
@@ -130,7 +134,10 @@ export class WebSocketMiddleware {
     ) => {
       const userInfo =
         'authToken' in action
-          ? await getUserInfoFromAuthToken(action.authToken!)
+          ? await getUserInfoFromApiKey({
+              apiKey: action.authToken!,
+              fields: ['id', 'email', 'discord_id'],
+            })
           : undefined
 
       // Use the new combined context - much cleaner!