refactor: remove test-only user special-casing

brandonkachen · brandonkachen · commit 2597c7aa8881 · 2025-12-17T01:02:29.000-08:00
diff --git a/common/src/old-constants.ts b/common/src/old-constants.ts
@@ -311,8 +311,6 @@ export function supportsCacheControl(model: Model): boolean {
   return !nonCacheableModels.includes(model)
 }
 
-export const TEST_USER_ID = 'test-user-id'
-
 export function getModelFromShortName(
   modelName: string | undefined,
 ): Model | undefined {
diff --git a/common/src/testing/constants.ts b/common/src/testing/constants.ts
@@ -0,0 +1 @@
+export const TEST_USER_ID = 'test-user-id'
diff --git a/evals/scaffolding.ts b/evals/scaffolding.ts
@@ -5,8 +5,9 @@ import path from 'path'
 import { runAgentStep } from '@codebuff/agent-runtime/run-agent-step'
 import { assembleLocalAgentTemplates } from '@codebuff/agent-runtime/templates/agent-registry'
 import { getFileTokenScores } from '@codebuff/code-map/parse'
-import { API_KEY_ENV_VAR, TEST_USER_ID } from '@codebuff/common/old-constants'
 import { clientToolCallSchema } from '@codebuff/common/tools/list'
+import { API_KEY_ENV_VAR } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { generateCompactId } from '@codebuff/common/util/string'
 import { getSystemInfo } from '@codebuff/common/util/system-info'
 import { ToolHelpers } from '@codebuff/sdk'
diff --git a/packages/agent-runtime/src/__tests__/fast-rewrite.test.ts b/packages/agent-runtime/src/__tests__/fast-rewrite.test.ts
@@ -1,6 +1,6 @@
 import path from 'path'
 
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { afterAll, beforeEach, describe, expect, it } from 'bun:test'
 import { createPatch } from 'diff'
diff --git a/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts b/packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts
@@ -1,5 +1,5 @@
 import * as analytics from '@codebuff/common/analytics'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import {
   mockAnalytics,
   mockBigQuery,
diff --git a/packages/agent-runtime/src/__tests__/main-prompt.test.ts b/packages/agent-runtime/src/__tests__/main-prompt.test.ts
@@ -1,6 +1,6 @@
 import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import {
   mockAnalytics,
   mockBigQuery,
diff --git a/packages/agent-runtime/src/__tests__/n-parameter.test.ts b/packages/agent-runtime/src/__tests__/n-parameter.test.ts
@@ -1,5 +1,5 @@
 import * as analytics from '@codebuff/common/analytics'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import {
   mockAnalytics,
   mockRandomUUID,
diff --git a/packages/agent-runtime/src/__tests__/process-file-block.test.ts b/packages/agent-runtime/src/__tests__/process-file-block.test.ts
@@ -1,4 +1,4 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { cleanMarkdownCodeBlock } from '@codebuff/common/util/file'
 import { beforeEach, describe, expect, it } from 'bun:test'
diff --git a/packages/agent-runtime/src/__tests__/prompt-caching-subagents.test.ts b/packages/agent-runtime/src/__tests__/prompt-caching-subagents.test.ts
@@ -1,4 +1,4 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import { assistantMessage, userMessage } from '@codebuff/common/util/messages'
diff --git a/packages/agent-runtime/src/__tests__/propose-tools.test.ts b/packages/agent-runtime/src/__tests__/propose-tools.test.ts
@@ -1,4 +1,4 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
diff --git a/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts b/packages/agent-runtime/src/__tests__/read-docs-tool.test.ts
@@ -1,6 +1,6 @@
 import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import {
   mockAnalytics,
   mockBigQuery,
diff --git a/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts b/packages/agent-runtime/src/__tests__/run-agent-step-tools.test.ts
@@ -1,6 +1,6 @@
 import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import {
   mockAnalytics,
   mockBigQuery,
diff --git a/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts b/packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts
@@ -1,5 +1,5 @@
 import * as analytics from '@codebuff/common/analytics'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import {
   mockAnalytics,
   mockRandomUUID,
diff --git a/packages/agent-runtime/src/__tests__/spawn-agents-image-content.test.ts b/packages/agent-runtime/src/__tests__/spawn-agents-image-content.test.ts
@@ -1,4 +1,4 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
diff --git a/packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts b/packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts
@@ -1,4 +1,4 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
diff --git a/packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts b/packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts
@@ -1,4 +1,4 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import { assistantMessage } from '@codebuff/common/util/messages'
diff --git a/packages/agent-runtime/src/__tests__/subagent-streaming.test.ts b/packages/agent-runtime/src/__tests__/subagent-streaming.test.ts
@@ -1,4 +1,4 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import { assistantMessage } from '@codebuff/common/util/messages'
diff --git a/packages/agent-runtime/src/__tests__/web-search-tool.test.ts b/packages/agent-runtime/src/__tests__/web-search-tool.test.ts
@@ -1,6 +1,6 @@
 import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
+import { TEST_USER_ID } from '@codebuff/common/testing/constants'
 import {
   mockAnalytics,
   mockBigQuery,
diff --git a/packages/billing/src/balance-calculator.ts b/packages/billing/src/balance-calculator.ts
@@ -1,4 +1,3 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { GrantTypeValues } from '@codebuff/common/types/grant'
 import { failure, getErrorObject, success } from '@codebuff/common/util/error'
 import db from '@codebuff/internal/db'
@@ -481,9 +480,6 @@ export async function consumeCreditsAndAddAgentStep(params: {
               tx,
             })
 
-            if (userId === TEST_USER_ID) {
-              return { ...result, agentStepId: 'test-step-id' }
-            }
           }
 
           try {
diff --git a/web/src/app/api/admin/relabel-for-user/route.ts b/web/src/app/api/admin/relabel-for-user/route.ts
@@ -14,7 +14,6 @@ import {
 import {
   finetunedVertexModels,
   models,
-  TEST_USER_ID,
 } from '@codebuff/common/old-constants'
 import { userMessage } from '@codebuff/common/util/messages'
 import { generateCompactId } from '@codebuff/common/util/string'
@@ -61,6 +60,7 @@ export async function GET(req: NextRequest) {
   if (authResult instanceof NextResponse) {
     return authResult
   }
+  const adminUser = authResult
 
   const userId = req.nextUrl.searchParams.get('userId')
   if (!userId) {
@@ -96,6 +96,7 @@ export async function POST(req: NextRequest) {
   if (authResult instanceof NextResponse) {
     return authResult
   }
+  const adminUser = authResult
 
   const userId = req.nextUrl.searchParams.get('userId')
   if (!userId) {
@@ -129,7 +130,7 @@ export async function POST(req: NextRequest) {
     const results = await relabelUserTraces({
       userId,
       limit,
-      promptContext: buildPromptContext(apiKey),
+      promptContext: buildPromptContext({ apiKey, adminUserId: adminUser.id }),
     })
 
     return NextResponse.json({
@@ -519,14 +520,15 @@ function extractQueryFromMessages(messages: unknown): string {
   return match?.[1] ?? 'Unknown query'
 }
 
-function buildPromptContext(apiKey: string) {
+function buildPromptContext(params: { apiKey: string; adminUserId: string }) {
+  const { apiKey, adminUserId } = params
   return {
     apiKey,
     runId: `admin-relabel-${Date.now()}`,
     clientSessionId: STATIC_SESSION_ID,
     fingerprintId: STATIC_SESSION_ID,
     userInputId: STATIC_SESSION_ID,
-    userId: TEST_USER_ID,
+    userId: adminUserId,
     sendAction: async () => {},
     trackEvent: async () => {},
     logger,
diff --git a/web/src/app/api/v1/agent-runs/[runId]/steps/__tests__/steps.test.ts b/web/src/app/api/v1/agent-runs/[runId]/steps/__tests__/steps.test.ts
@@ -1,4 +1,3 @@
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { beforeEach, describe, expect, mock, test } from 'bun:test'
 import { NextRequest } from 'next/server'
 
@@ -28,14 +27,6 @@ describe('agentRunsStepsPost', () => {
           ]),
         ) as any
       }
-      if (apiKey === 'test-key') {
-        return Object.fromEntries(
-          fields.map((field) => [
-            field,
-            field === 'id' ? TEST_USER_ID : undefined,
-          ]),
-        ) as any
-      }
       return null
     }
 
@@ -236,31 +227,6 @@ describe('agentRunsStepsPost', () => {
     expect(json.error).toBe('Unauthorized to add steps to this run')
   })
 
-  test('returns test step ID for test user', async () => {
-    const req = new NextRequest(
-      'http://localhost/api/v1/agent-runs/run-123/steps',
-      {
-        method: 'POST',
-        headers: { Authorization: 'Bearer test-key' },
-        body: JSON.stringify({ stepNumber: 1 }),
-      },
-    )
-
-    const response = await postAgentRunsSteps({
-      req,
-      runId: 'run-123',
-      getUserInfoFromApiKey: mockGetUserInfoFromApiKey,
-      logger: mockLogger,
-      loggerWithContext: mockLoggerWithContext,
-      trackEvent: mockTrackEvent,
-      db: mockDb,
-    })
-
-    expect(response.status).toBe(200)
-    const json = await response.json()
-    expect(json.stepId).toBe('test-step-id')
-  })
-
   test('successfully adds agent step', async () => {
     const req = new NextRequest(
       'http://localhost/api/v1/agent-runs/run-123/steps',
diff --git a/web/src/app/api/v1/agent-runs/[runId]/steps/_post.ts b/web/src/app/api/v1/agent-runs/[runId]/steps/_post.ts
@@ -1,5 +1,4 @@
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { getErrorObject } from '@codebuff/common/util/error'
 import * as schema from '@codebuff/internal/db/schema'
 import { eq } from 'drizzle-orm'
@@ -108,11 +107,6 @@ export async function postAgentRunsSteps(params: {
     startTime,
   } = data
 
-  // Skip database insert for test user
-  if (userInfo.id === TEST_USER_ID) {
-    return NextResponse.json({ stepId: 'test-step-id' })
-  }
-
   // Verify the run belongs to the authenticated user
   const agentRun = await db
     .select({ user_id: schema.agentRun.user_id })
diff --git a/web/src/app/api/v1/agent-runs/__tests__/agent-runs.test.ts b/web/src/app/api/v1/agent-runs/__tests__/agent-runs.test.ts
@@ -1,5 +1,4 @@
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test'
 import { NextRequest } from 'next/server'
 
@@ -26,9 +25,6 @@ describe('/api/v1/agent-runs POST endpoint', () => {
     'test-api-key-456': {
       id: 'user-456',
     },
-    'test-api-key-test': {
-      id: TEST_USER_ID,
-    },
   }
 
   const mockGetUserInfoFromApiKey: GetUserInfoFromApiKeyFn = async ({
@@ -744,34 +740,4 @@ describe('/api/v1/agent-runs POST endpoint', () => {
     })
   })
 
-  describe('Test user handling', () => {
-    test('skips database update for test user on FINISH action', async () => {
-      const req = new NextRequest('http://localhost:3000/api/v1/agent-runs', {
-        method: 'POST',
-        headers: { Authorization: 'Bearer test-api-key-test' },
-        body: JSON.stringify({
-          action: 'FINISH',
-          runId: 'run-test',
-          status: 'completed',
-          totalSteps: 5,
-          directCredits: 100,
-          totalCredits: 150,
-        }),
-      })
-
-      const response = await postAgentRuns({
-        req,
-        getUserInfoFromApiKey: mockGetUserInfoFromApiKey,
-        logger: mockLogger,
-        loggerWithContext: mockLoggerWithContext,
-        trackEvent: mockTrackEvent,
-        db: mockDb,
-      })
-
-      expect(response.status).toBe(200)
-      const body = await response.json()
-      expect(body).toEqual({ success: true })
-      expect(mockDb.update).not.toHaveBeenCalled()
-    })
-  })
 })
diff --git a/web/src/app/api/v1/agent-runs/_post.ts b/web/src/app/api/v1/agent-runs/_post.ts
@@ -1,5 +1,4 @@
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
-import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { getErrorObject } from '@codebuff/common/util/error'
 import * as schema from '@codebuff/internal/db/schema'
 import { eq } from 'drizzle-orm'
@@ -117,11 +116,6 @@ async function handleFinishAction(params: {
     errorMessage,
   } = data
 
-  // Skip database update for test user
-  if (userId === TEST_USER_ID) {
-    return NextResponse.json({ success: true })
-  }
-
   try {
     await db
       .update(schema.agentRun)

Original file line number	Diff line number	Diff line change
`@@ -311,8 +311,6 @@ export function supportsCacheControl(model: Model): boolean {`
`311`	`311`	`return !nonCacheableModels.includes(model)`
`312`	`312`	`}`
`313`	`313`
`314`		`-export const TEST_USER_ID = 'test-user-id'`
`315`		`-`
`316`	`314`	`export function getModelFromShortName(`
`317`	`315`	`modelName: string \| undefined,`
`318`	`316`	`): Model \| undefined {`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+export const TEST_USER_ID = 'test-user-id'`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-import { TEST_USER_ID } from '@codebuff/common/old-constants'`
	`1`	`+import { TEST_USER_ID } from '@codebuff/common/testing/constants'`
`2`	`2`	`import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'`
`3`	`3`	`import { cleanMarkdownCodeBlock } from '@codebuff/common/util/file'`
`4`	`4`	`import { beforeEach, describe, expect, it } from 'bun:test'`