pass in databaseAgentCache

Author: charleslien

backend/src/__tests__/agent-registry.test.ts

@@ -19,7 +19,6 @@ import {
 import {
   getAgentTemplate,
   assembleLocalAgentTemplates,
-  clearDatabaseCache,
 } from '../templates/agent-registry'
 
 import type { AgentTemplate } from '@codebuff/agent-runtime/templates/types'
@@ -147,10 +146,12 @@ describe('Agent Registry', () => {
   })
 
   beforeEach(async () => {
-    agentRuntimeImpl = { ...TEST_AGENT_RUNTIME_IMPL }
+    agentRuntimeImpl = {
+      ...TEST_AGENT_RUNTIME_IMPL,
+    }
+
+    agentRuntimeImpl.databaseAgentCache.clear()
 
-    // Clear cache before each test
-    clearDatabaseCache()
     mockFileContext = getStubProjectFileContext()
 
     // Spy on validation functions
@@ -520,8 +521,7 @@ describe('Agent Registry', () => {
       })
       expect(selectSpy).toHaveBeenCalledTimes(1)
 
-      // Clear cache
-      clearDatabaseCache()
+      agentRuntimeImpl.databaseAgentCache.clear()
 
       // Third call - should hit database again after cache clear
       await getAgentTemplate({

backend/src/impl/agent-runtime.ts

@@ -8,6 +8,7 @@ import { fetchAgentFromDatabase } from '../templates/agent-registry'
 import { logger } from '../util/logger'
 import { getUserInfoFromApiKey } from '../websockets/auth'
 
+import type { AgentTemplate } from '@codebuff/agent-runtime/templates/types'
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 
 export const BACKEND_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
@@ -17,6 +18,7 @@ export const BACKEND_AGENT_RUNTIME_IMPL: AgentRuntimeDeps = Object.freeze({
   startAgentRun,
   finishAgentRun,
   addAgentStep,
+  databaseAgentCache: new Map<string, AgentTemplate | null>(),
 
   // LLM
   promptAiSdkStream,

backend/src/templates/agent-registry.ts

@@ -18,9 +18,6 @@ import type { ProjectFileContext } from '@codebuff/common/util/file'
 
 export type AgentRegistry = Record<string, AgentTemplate>
 
-// Global database cache - only state in the system
-const databaseAgentCache = new Map<string, AgentTemplate | null>()
-
 /**
  * Fetch and validate an agent from the database by publisher/agent-id[@version] format
  */
@@ -132,10 +129,16 @@ export async function getAgentTemplate(params: {
   agentId: string
   localAgentTemplates: Record<string, AgentTemplate>
   fetchAgentFromDatabase: FetchAgentFromDatabaseFn
+  databaseAgentCache: Map<string, AgentTemplate | null>
   logger: Logger
 }): Promise<AgentTemplate | null> {
-  const { agentId, localAgentTemplates, fetchAgentFromDatabase, logger } =
-    params
+  const {
+    agentId,
+    localAgentTemplates,
+    fetchAgentFromDatabase,
+    databaseAgentCache,
+    logger,
+  } = params
   // 1. Check localAgentTemplates first (dynamic agents + static templates)
   if (localAgentTemplates[agentId]) {
     return localAgentTemplates[agentId]
@@ -203,6 +206,10 @@ export function assembleLocalAgentTemplates(params: {
 /**
  * Clear the database agent cache (useful for testing)
  */
-export function clearDatabaseCache(): void {
+export function clearDatabaseCache(params: {
+  databaseAgentCache: Map<string, AgentTemplate | null>
+}): void {
+  const { databaseAgentCache } = params
+
   databaseAgentCache.clear()
 }

backend/src/websockets/middleware.ts

@@ -32,7 +32,6 @@ import type {
   AgentRuntimeDeps,
   AgentRuntimeScopedDeps,
 } 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'
 

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

@@ -1,8 +1,9 @@
+import type { AgentTemplate } from '../../types/agent-template'
 import type {
   AgentRuntimeDeps,
   AgentRuntimeScopedDeps,
-} from '@codebuff/common/types/contracts/agent-runtime'
-import type { Logger } from '@codebuff/common/types/contracts/logger'
+} from '../../types/contracts/agent-runtime'
+import type { Logger } from '../../types/contracts/logger'
 
 export const testLogger: Logger = {
   debug: () => {},
@@ -15,13 +16,14 @@ export const TEST_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
   // Database
   getUserInfoFromApiKey: async () => ({
     id: 'test-user-id',
-   email: 'test-email',
+    email: 'test-email',
     discord_id: 'test-discord-id',
   }),
   fetchAgentFromDatabase: async () => null,
   startAgentRun: async () => 'test-agent-run-id',
   finishAgentRun: async () => {},
   addAgentStep: async () => 'test-agent-step-id',
+  databaseAgentCache: new Map<string, AgentTemplate | null>(),
 
   // LLM
   promptAiSdkStream: async function* () {

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

@@ -1,3 +1,4 @@
+import type { AgentTemplate } from '../agent-template'
 import type {
   HandleStepsLogChunkFn,
   RequestFilesFn,
@@ -28,6 +29,7 @@ export type AgentRuntimeDeps = {
   startAgentRun: StartAgentRunFn
   finishAgentRun: FinishAgentRunFn
   addAgentStep: AddAgentStepFn
+  databaseAgentCache: Map<string, AgentTemplate | null>
 
   // LLM
   promptAiSdkStream: PromptAiSdkStreamFn

evals/impl/agent-runtime.ts

@@ -1,3 +1,4 @@
+import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { AgentRuntimeDeps } from '@codebuff/common/types/contracts/agent-runtime'
 
 export const EVALS_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
@@ -11,6 +12,7 @@ export const EVALS_AGENT_RUNTIME_IMPL = Object.freeze<AgentRuntimeDeps>({
   startAgentRun: async () => 'test-agent-run-id',
   finishAgentRun: async () => {},
   addAgentStep: async () => 'test-agent-step-id',
+  databaseAgentCache: new Map<string, AgentTemplate | null>(),
 
   // LLM
   promptAiSdkStream: async function* () {