sdk: implement fetchAgentFromDatabase

Author: charleslien

backend/src/__tests__/cost-aggregation.integration.test.ts

@@ -2,10 +2,7 @@ import { disableLiveUserInputCheck } from '@codebuff/agent-runtime/live-user-inp
 import { mainPrompt } from '@codebuff/agent-runtime/main-prompt'
 import * as agentRegistry from '@codebuff/agent-runtime/templates/agent-registry'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
-import {
-  TEST_AGENT_RUNTIME_IMPL,
-  TEST_AGENT_RUNTIME_SCOPED_IMPL,
-} from '@codebuff/common/testing/impl/agent-runtime'
+import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import { getInitialSessionState } from '@codebuff/common/types/session-state'
 import {
   spyOn,
@@ -90,17 +87,15 @@ const mockFileContext: ProjectFileContext = {
 
 describe('Cost Aggregation Integration Tests', () => {
   let mockLocalAgentTemplates: Record<string, any>
-  let agentRuntimeImpl: AgentRuntimeDeps
-  let agentRuntimeScopedImpl: AgentRuntimeScopedDeps
+  let agentRuntimeImpl: AgentRuntimeDeps & AgentRuntimeScopedDeps
 
   beforeAll(() => {
     disableLiveUserInputCheck()
   })
 
   beforeEach(async () => {
-    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
-    agentRuntimeScopedImpl = {
-      ...TEST_AGENT_RUNTIME_SCOPED_IMPL,
+    agentRuntimeImpl = {
+      ...TEST_AGENT_RUNTIME_IMPL,
       sendAction: mock(() => {}),
     }
 
@@ -178,7 +173,7 @@ describe('Cost Aggregation Integration Tests', () => {
     }
 
     // Mock tool call execution
-    agentRuntimeScopedImpl.requestToolCall = async ({ toolName, input }) => {
+    agentRuntimeImpl.requestToolCall = async ({ toolName, input }) => {
       if (toolName === 'write_file') {
         return {
           output: [
@@ -204,9 +199,7 @@ describe('Cost Aggregation Integration Tests', () => {
     }
 
     // Mock file reading
-    agentRuntimeScopedImpl.requestFiles = async (params: {
-      filePaths: string[]
-    }) => {
+    agentRuntimeImpl.requestFiles = async (params: { filePaths: string[] }) => {
       const results: Record<string, string | null> = {}
       params.filePaths.forEach((path) => {
         results[path] = path === 'hello.txt' ? 'Hello, World!' : null
@@ -244,7 +237,6 @@ describe('Cost Aggregation Integration Tests', () => {
 
     const result = await mainPrompt({
       ...agentRuntimeImpl,
-      ...agentRuntimeScopedImpl,
       repoId: undefined,
       repoUrl: undefined,
       action,
@@ -281,7 +273,6 @@ describe('Cost Aggregation Integration Tests', () => {
     // Call through websocket action handler to test full integration
     await websocketAction.callMainPrompt({
       ...agentRuntimeImpl,
-      ...agentRuntimeScopedImpl,
       repoId: undefined,
       repoUrl: undefined,
       action,
@@ -292,7 +283,7 @@ describe('Cost Aggregation Integration Tests', () => {
 
     // Verify final cost is included in prompt response
     const promptResponse = (
-      agentRuntimeScopedImpl.sendAction as Mock<SendActionFn>
+      agentRuntimeImpl.sendAction as Mock<SendActionFn>
     ).mock.calls
       .map((call) => call[0].action)
       .find((action: ServerAction) => action.type === 'prompt-response') as any
@@ -353,7 +344,6 @@ describe('Cost Aggregation Integration Tests', () => {
 
     const result = await mainPrompt({
       ...agentRuntimeImpl,
-      ...agentRuntimeScopedImpl,
       repoId: undefined,
       repoUrl: undefined,
       action,
@@ -411,7 +401,6 @@ describe('Cost Aggregation Integration Tests', () => {
     try {
       result = await mainPrompt({
         ...agentRuntimeImpl,
-        ...agentRuntimeScopedImpl,
         repoId: undefined,
         repoUrl: undefined,
         action,
@@ -462,7 +451,6 @@ describe('Cost Aggregation Integration Tests', () => {
 
     await mainPrompt({
       ...agentRuntimeImpl,
-      ...agentRuntimeScopedImpl,
       repoId: undefined,
       repoUrl: undefined,
       action,
@@ -504,7 +492,6 @@ describe('Cost Aggregation Integration Tests', () => {
     // Call through websocket action to test server-side reset
     await websocketAction.callMainPrompt({
       ...agentRuntimeImpl,
-      ...agentRuntimeScopedImpl,
       repoId: undefined,
       repoUrl: undefined,
       action,
@@ -515,7 +502,7 @@ describe('Cost Aggregation Integration Tests', () => {
 
     // Server should have reset the malicious value and calculated correct cost
     const promptResponse = (
-      agentRuntimeScopedImpl.sendAction as Mock<SendActionFn>
+      agentRuntimeImpl.sendAction as Mock<SendActionFn>
     ).mock.calls
       .map((call) => call[0].action)
       .find((action) => action.type === 'prompt-response') as any

backend/src/__tests__/main-prompt.integration.test.ts

@@ -4,15 +4,18 @@ import { describe, it, expect, beforeEach, afterEach, mock } from 'bun:test'
 
 import { validateAgentNameHandlerHelper } from '../validate-agent-name'
 
-import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
+import type {
+  AgentRuntimeDeps,
+  AgentRuntimeScopedDeps,
+} from '@codebuff/common/types/contracts/agent-runtime'
 import type { FetchAgentFromDatabaseFn } from '@codebuff/common/types/contracts/database'
 import type {
   Request as ExpressRequest,
   Response as ExpressResponse,
   NextFunction,
 } from 'express'
 
-let agentRuntimeImpl: AgentRuntimeDeps
+let agentRuntimeImpl: AgentRuntimeDeps & AgentRuntimeScopedDeps
 
 function createMockReq(query: Record<string, any>): Partial<ExpressRequest> {
   return {

backend/src/api/__tests__/validate-agent-name.test.ts

@@ -39,7 +39,7 @@ export async function validateAgentNameHandlerHelper(
     next: NextFunction
   } & ParamsExcluding<
     typeof getAgentTemplate,
-    'agentId' | 'localAgentTemplates'
+    'agentId' | 'localAgentTemplates' | 'apiKey'
   >,
 ): Promise<void | ExpressResponse> {
   const { req, res, next } = params
@@ -88,6 +88,7 @@ export async function validateAgentNameHandlerHelper(
       ...params,
       agentId,
       localAgentTemplates: {},
+      apiKey,
     })
     if (found) {
       const result = {

backend/src/api/validate-agent-name.ts

@@ -168,14 +168,14 @@ export class WebSocketMiddleware {
       clientSessionId: string,
       ws: WebSocket,
     ) => {
-      const userInfo =
-        'authToken' in action
-          ? await getUserInfoFromApiKey({
-              apiKey: action.authToken!,
-              fields: ['id', 'email', 'discord_id'],
-              logger,
-            })
-          : undefined
+      const authToken = 'authToken' in action ? action.authToken : undefined
+      const userInfo = authToken
+        ? await getUserInfoFromApiKey({
+            apiKey: authToken,
+            fields: ['id', 'email', 'discord_id'],
+            logger,
+          })
+        : undefined
 
       const scopedDeps: AgentRuntimeScopedDeps = {
         handleStepsLogChunk: (params) =>
@@ -187,6 +187,7 @@ export class WebSocketMiddleware {
           requestOptionalFileWs({ ...params, ws }),
         sendSubagentChunk: (params) => sendSubagentChunkWs({ ...params, ws }),
         sendAction: (params) => sendActionWs({ ...params, ws }),
+        apiKey: authToken ?? '',
       }
 
       // Use the new combined context - much cleaner!

backend/src/websockets/middleware.ts

@@ -19,7 +19,9 @@ testFetch.preconnect = async () => {
   throw new Error('fetch.preconnect not implemented in test runtime')
 }
 
-export const TEST_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
+export const TEST_AGENT_RUNTIME_IMPL = Object.freeze<
+  AgentRuntimeDeps & AgentRuntimeScopedDeps
+>({
   // Database
   getUserInfoFromApiKey: async () => ({
     id: 'test-user-id',
@@ -60,30 +62,31 @@ export const TEST_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
   // Other
   logger: testLogger,
   fetch: testFetch,
-})
 
-export const TEST_AGENT_RUNTIME_SCOPED_IMPL =
-  Object.freeze<AgentRuntimeScopedDeps>({
-    // Database
-    handleStepsLogChunk: () => {
-      throw new Error('handleStepsLogChunk not implemented in test runtime')
-    },
-    requestToolCall: () => {
-      throw new Error('requestToolCall not implemented in test runtime')
-    },
-    requestMcpToolData: () => {
-      throw new Error('requestMcpToolData not implemented in test runtime')
-    },
-    requestFiles: () => {
-      throw new Error('requestFiles not implemented in test runtime')
-    },
-    requestOptionalFile: () => {
-      throw new Error('requestOptionalFile not implemented in test runtime')
-    },
-    sendSubagentChunk: () => {
-      throw new Error('sendSubagentChunk not implemented in test runtime')
-    },
-    sendAction: () => {
-      throw new Error('sendAction not implemented in test runtime')
-    },
-  })
+  // Scoped deps
+
+  // Database
+  handleStepsLogChunk: () => {
+    throw new Error('handleStepsLogChunk not implemented in test runtime')
+  },
+  requestToolCall: () => {
+    throw new Error('requestToolCall not implemented in test runtime')
+  },
+  requestMcpToolData: () => {
+    throw new Error('requestMcpToolData not implemented in test runtime')
+  },
+  requestFiles: () => {
+    throw new Error('requestFiles not implemented in test runtime')
+  },
+  requestOptionalFile: () => {
+    throw new Error('requestOptionalFile not implemented in test runtime')
+  },
+  sendSubagentChunk: () => {
+    throw new Error('sendSubagentChunk not implemented in test runtime')
+  },
+  sendAction: () => {
+    throw new Error('sendAction not implemented in test runtime')
+  },
+
+  apiKey: 'test-api-key',
+})

common/src/testing/impl/agent-runtime.ts

@@ -25,6 +25,7 @@ import type {
 } from './llm'
 import type { Logger } from './logger'
 
+/** Shared dependencies */
 export type AgentRuntimeDeps = {
   // Database
   getUserInfoFromApiKey: GetUserInfoFromApiKeyFn
@@ -54,6 +55,7 @@ export type AgentRuntimeDeps = {
   fetch: typeof globalThis.fetch
 }
 
+/** Per-run dependencies */
 export type AgentRuntimeScopedDeps = {
   // Client (WebSocket)
   handleStepsLogChunk: HandleStepsLogChunkFn
@@ -63,4 +65,6 @@ export type AgentRuntimeScopedDeps = {
   requestOptionalFile: RequestOptionalFileFn
   sendAction: SendActionFn
   sendSubagentChunk: SendSubagentChunkFn
+
+  apiKey: string
 }

common/src/types/contracts/agent-runtime.ts

@@ -46,6 +46,7 @@ export type GetAgentRunFromIdFn = <T extends AgentRunColumn>(
  * Fetch and validate an agent from the database by `publisher/agent-id[@version]` format
  */
 export type FetchAgentFromDatabaseFn = (params: {
+  apiKey: string
   parsedAgentId: {
     publisherId: string
     agentId: string

common/src/types/contracts/database.ts

@@ -200,6 +200,7 @@ export async function runAgentStepScaffolding(
     sendSubagentChunk: (params) =>
       sendSubagentChunkWs({ ...params, ws: mockWs }),
     sendAction: (params) => sendActionWs({ ...params, ws: mockWs }),
+    apiKey: process.env.CODEBUFF_API_KEY ?? '',
   }
   const result = await runAgentStep({
     ...EVALS_AGENT_RUNTIME_IMPL,

evals/scaffolding.ts

@@ -1,7 +1,4 @@
-import {
-  TEST_AGENT_RUNTIME_IMPL,
-  TEST_AGENT_RUNTIME_SCOPED_IMPL,
-} from '@codebuff/common/testing/impl/agent-runtime'
+import { TEST_AGENT_RUNTIME_IMPL } from '@codebuff/common/testing/impl/agent-runtime'
 import {
   getInitialAgentState,
   getInitialSessionState,
@@ -182,7 +179,6 @@ describe('Cost Aggregation System', () => {
 
       const result = handleSpawnAgents({
         ...TEST_AGENT_RUNTIME_IMPL,
-        ...TEST_AGENT_RUNTIME_SCOPED_IMPL,
         repoId: undefined,
         repoUrl: undefined,
         previousToolCallFinished: Promise.resolve(),
@@ -264,7 +260,6 @@ describe('Cost Aggregation System', () => {
 
       const result = handleSpawnAgents({
         ...TEST_AGENT_RUNTIME_IMPL,
-        ...TEST_AGENT_RUNTIME_SCOPED_IMPL,
         repoId: undefined,
         repoUrl: undefined,
         previousToolCallFinished: Promise.resolve(),
@@ -423,7 +418,6 @@ describe('Cost Aggregation System', () => {
 
       const result = handleSpawnAgents({
         ...TEST_AGENT_RUNTIME_IMPL,
-        ...TEST_AGENT_RUNTIME_SCOPED_IMPL,
         repoId: undefined,
         repoUrl: undefined,
         previousToolCallFinished: Promise.resolve(),