move agent-registry to agent-runtime

charleslien · charleslien · commit b231514128c8 · 2025-10-17T17:53:27.000-07:00
diff --git a/backend/src/__tests__/cost-aggregation.integration.test.ts b/backend/src/__tests__/cost-aggregation.integration.test.ts
@@ -1,3 +1,4 @@
+import * as agentRegistry from '@codebuff/agent-runtime/templates/agent-registry'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import {
   TEST_AGENT_RUNTIME_IMPL,
@@ -16,7 +17,6 @@ import {
 
 import * as messageCostTracker from '../llm-apis/message-cost-tracker'
 import { mainPrompt } from '../main-prompt'
-import * as agentRegistry from '../templates/agent-registry'
 import * as websocketAction from '../websockets/websocket-action'
 
 import type { AgentTemplate } from '@codebuff/agent-runtime/templates/types'
diff --git a/backend/src/__tests__/cost-aggregation.test.ts b/backend/src/__tests__/cost-aggregation.test.ts
@@ -1,3 +1,4 @@
+import * as agentRegistry from '@codebuff/agent-runtime/templates/agent-registry'
 import {
   TEST_AGENT_RUNTIME_IMPL,
   TEST_AGENT_RUNTIME_SCOPED_IMPL,
@@ -16,7 +17,6 @@ import {
   mock,
 } from 'bun:test'
 
-import * as agentRegistry from '../templates/agent-registry'
 import * as spawnAgentUtils from '../tools/handlers/tool/spawn-agent-utils'
 import { handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
 
diff --git a/backend/src/__tests__/loop-agent-steps.test.ts b/backend/src/__tests__/loop-agent-steps.test.ts
@@ -29,7 +29,7 @@ import { loopAgentSteps } from '../run-agent-step'
 import { clearAgentGeneratorCache } from '../run-programmatic-step'
 import { mockFileContext } from './test-utils'
 
-import type { getAgentTemplate } from '../templates/agent-registry'
+import type { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import type { AgentTemplate } from '@codebuff/agent-runtime/templates/types'
 import type { StepGenerator } from '@codebuff/common/types/agent-template'
 import type {
diff --git a/backend/src/__tests__/read-docs-tool.test.ts b/backend/src/__tests__/read-docs-tool.test.ts
@@ -1,3 +1,4 @@
+import { assembleLocalAgentTemplates } from '@codebuff/agent-runtime/templates/agent-registry'
 import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
@@ -24,7 +25,6 @@ import * as liveUserInputs from '../live-user-inputs'
 import { mockFileContext } from './test-utils'
 import * as context7Api from '../llm-apis/context7-api'
 import { runAgentStep } from '../run-agent-step'
-import { assembleLocalAgentTemplates } from '../templates/agent-registry'
 
 import type {
   AgentRuntimeDeps,
diff --git a/backend/src/__tests__/subagent-streaming.test.ts b/backend/src/__tests__/subagent-streaming.test.ts
@@ -1,3 +1,4 @@
+import { assembleLocalAgentTemplates } from '@codebuff/agent-runtime/templates/agent-registry'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import {
   TEST_AGENT_RUNTIME_IMPL,
@@ -17,7 +18,6 @@ import {
 
 import * as runAgentStep from '../run-agent-step'
 import { mockFileContext } from './test-utils'
-import { assembleLocalAgentTemplates } from '../templates/agent-registry'
 import { handleSpawnAgents } from '../tools/handlers/tool/spawn-agents'
 import * as loggerModule from '../util/logger'
 
diff --git a/backend/src/__tests__/web-search-tool.test.ts b/backend/src/__tests__/web-search-tool.test.ts
@@ -1,6 +1,7 @@
 // Set environment variables before any imports
 process.env.LINKUP_API_KEY = 'test-api-key'
 
+import { assembleLocalAgentTemplates } from '@codebuff/agent-runtime/templates/agent-registry'
 import * as bigquery from '@codebuff/bigquery'
 import * as analytics from '@codebuff/common/analytics'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
@@ -27,7 +28,6 @@ import * as liveUserInputs from '../live-user-inputs'
 import { mockFileContext } from './test-utils'
 import * as linkupApi from '../llm-apis/linkup-api'
 import { runAgentStep } from '../run-agent-step'
-import { assembleLocalAgentTemplates } from '../templates/agent-registry'
 
 import type {
   AgentRuntimeDeps,
diff --git a/backend/src/api/validate-agent-name.ts b/backend/src/api/validate-agent-name.ts
@@ -1,7 +1,7 @@
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import { AGENT_PERSONAS } from '@codebuff/common/constants/agents'
 import { z } from 'zod/v4'
 
-import { getAgentTemplate } from '../templates/agent-registry'
 import { extractAuthTokenFromHeader } from '../util/auth-helpers'
 import { logger } from '../util/logger'
 
diff --git a/backend/src/impl/agent-runtime.ts b/backend/src/impl/agent-runtime.ts
@@ -4,7 +4,7 @@ import {
   promptAiSdkStream,
   promptAiSdkStructured,
 } from '../llm-apis/vercel-ai-sdk/ai-sdk'
-import { fetchAgentFromDatabase } from '../templates/agent-registry'
+import { fetchAgentFromDatabase } from '../templates/agent-db'
 import { logger } from '../util/logger'
 import { getUserInfoFromApiKey } from '../websockets/auth'
 
diff --git a/backend/src/main-prompt.ts b/backend/src/main-prompt.ts
@@ -1,11 +1,11 @@
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import { expireMessages } from '@codebuff/agent-runtime/util/messages'
 import { AgentTemplateTypes } from '@codebuff/common/types/session-state'
 import { generateCompactId } from '@codebuff/common/util/string'
 import { uniq } from 'lodash'
 
 import { checkTerminalCommand } from './check-terminal-command'
 import { loopAgentSteps } from './run-agent-step'
-import { getAgentTemplate } from './templates/agent-registry'
 
 import type { AgentTemplate } from '@codebuff/agent-runtime/templates/types'
 import type { ClientAction } from '@codebuff/common/actions'
diff --git a/backend/src/run-agent-step.ts b/backend/src/run-agent-step.ts
@@ -1,5 +1,6 @@
 import { getMCPToolData } from '@codebuff/agent-runtime/mcp'
 import { getAgentStreamFromTemplate } from '@codebuff/agent-runtime/prompt-agent-stream'
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import {
   asSystemInstruction,
   asSystemMessage,
@@ -20,7 +21,6 @@ import { cloneDeep } from 'lodash'
 import { checkLiveUserInput } from './live-user-inputs'
 import { runProgrammaticStep } from './run-programmatic-step'
 import { additionalSystemPrompts } from './system-prompt/prompts'
-import { getAgentTemplate } from './templates/agent-registry'
 import { getAgentPrompt } from './templates/strings'
 import { processStreamWithTools } from './tools/stream-parser'
 import { getAgentOutput } from './util/agent-output'
@@ -401,7 +401,7 @@ export const runAgentStep = async (
   }
 }
 
-export const loopAgentSteps = async (
+export async function loopAgentSteps(
   params: {
     userInputId: string
     agentType: AgentTemplateType
@@ -448,7 +448,7 @@ export const loopAgentSteps = async (
 ): Promise<{
   agentState: AgentState
   output: AgentOutput
-}> => {
+}> {
   const {
     userInputId,
     agentType,
diff --git a/backend/src/templates/agent-db.ts b/backend/src/templates/agent-db.ts
@@ -0,0 +1,109 @@
+import db from '@codebuff/common/db'
+import * as schema from '@codebuff/common/db/schema'
+import { validateSingleAgent } from '@codebuff/common/templates/agent-validation'
+import { and, desc, eq } from 'drizzle-orm'
+
+import type { FetchAgentFromDatabaseFn } from '@codebuff/common/types/contracts/database'
+import type { DynamicAgentTemplate } from '@codebuff/common/types/dynamic-agent-template'
+import type { ParamsOf } from '@codebuff/common/types/function-params'
+
+/**
+ * Fetch and validate an agent from the database by publisher/agent-id[@version] format
+ */
+export async function fetchAgentFromDatabase(
+  params: ParamsOf<FetchAgentFromDatabaseFn>,
+): ReturnType<FetchAgentFromDatabaseFn> {
+  const { parsedAgentId, logger } = params
+  const { publisherId, agentId, version } = parsedAgentId
+
+  try {
+    let agentConfig
+
+    if (version && version !== 'latest') {
+      // Query for specific version
+      agentConfig = await db
+        .select()
+        .from(schema.agentConfig)
+        .where(
+          and(
+            eq(schema.agentConfig.id, agentId),
+            eq(schema.agentConfig.publisher_id, publisherId),
+            eq(schema.agentConfig.version, version),
+          ),
+        )
+        .then((rows) => rows[0])
+    } else {
+      // Query for latest version
+      agentConfig = await db
+        .select()
+        .from(schema.agentConfig)
+        .where(
+          and(
+            eq(schema.agentConfig.id, agentId),
+            eq(schema.agentConfig.publisher_id, publisherId),
+          ),
+        )
+        .orderBy(
+          desc(schema.agentConfig.major),
+          desc(schema.agentConfig.minor),
+          desc(schema.agentConfig.patch),
+        )
+        .limit(1)
+        .then((rows) => rows[0])
+    }
+
+    if (!agentConfig) {
+      logger.debug(
+        { publisherId, agentId, version },
+        'fetchAgentFromDatabase: Agent not found in database',
+      )
+      return null
+    }
+
+    const rawAgentData = agentConfig.data as DynamicAgentTemplate
+
+    // Validate the raw agent data with the original agentId (not full identifier)
+    const validationResult = validateSingleAgent({
+      template: { ...rawAgentData, id: agentId, version: agentConfig.version },
+      filePath: `${publisherId}/${agentId}@${agentConfig.version}`,
+    })
+
+    if (!validationResult.success) {
+      logger.error(
+        {
+          publisherId,
+          agentId,
+          version: agentConfig.version,
+          error: validationResult.error,
+        },
+        'fetchAgentFromDatabase: Agent validation failed',
+      )
+      return null
+    }
+
+    // Set the correct full agent ID for the final template
+    const agentTemplate = {
+      ...validationResult.agentTemplate!,
+      id: `${publisherId}/${agentId}@${agentConfig.version}`,
+    }
+
+    logger.debug(
+      {
+        publisherId,
+        agentId,
+        version: agentConfig.version,
+        fullAgentId: agentTemplate.id,
+        parsedAgentId,
+      },
+      'fetchAgentFromDatabase: Successfully loaded and validated agent from database',
+    )
+
+    return agentTemplate
+  } catch (error) {
+    logger.error(
+      { publisherId, agentId, version, error },
+      'fetchAgentFromDatabase: Error fetching agent from database',
+    )
+    return null
+  }
+}
diff --git a/backend/src/templates/prompts.ts b/backend/src/templates/prompts.ts
@@ -1,8 +1,7 @@
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import { buildArray } from '@codebuff/common/util/array'
 import { schemaToJsonStr } from '@codebuff/common/util/zod-schema'
 
-import { getAgentTemplate } from './agent-registry'
-
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type { ParamsExcluding } from '@codebuff/common/types/function-params'
diff --git a/backend/src/templates/strings.ts b/backend/src/templates/strings.ts
@@ -1,3 +1,4 @@
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import {
   PLACEHOLDER,
   placeholderValues,
@@ -8,7 +9,6 @@ import { escapeString } from '@codebuff/common/util/string'
 import { schemaToJsonStr } from '@codebuff/common/util/zod-schema'
 import { z } from 'zod/v4'
 
-import { getAgentTemplate } from './agent-registry'
 import { buildSpawnableAgentsDescription } from './prompts'
 import {
   getGitChangesPrompt,
diff --git a/backend/src/tools/handlers/tool/lookup-agent-info.ts b/backend/src/tools/handlers/tool/lookup-agent-info.ts
@@ -1,8 +1,7 @@
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import { removeUndefinedProps } from '@codebuff/common/util/object'
 import z from 'zod/v4'
 
-import { getAgentTemplate } from '../../../templates/agent-registry'
-
 import type { CodebuffToolHandlerFunction } from '@codebuff/agent-runtime/tools/handlers/handler-function-type'
 
 export const handleLookupAgentInfo: CodebuffToolHandlerFunction<
diff --git a/backend/src/tools/handlers/tool/set-output.ts b/backend/src/tools/handlers/tool/set-output.ts
@@ -1,4 +1,4 @@
-import { getAgentTemplate } from '../../../templates/agent-registry'
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 
 import type { CodebuffToolHandlerFunction } from '@codebuff/agent-runtime/tools/handlers/handler-function-type'
 
diff --git a/backend/src/tools/handlers/tool/spawn-agent-utils.ts b/backend/src/tools/handlers/tool/spawn-agent-utils.ts
@@ -1,10 +1,10 @@
+import { getAgentTemplate } from '@codebuff/agent-runtime/templates/agent-registry'
 import { MAX_AGENT_STEPS_DEFAULT } from '@codebuff/common/constants/agents'
 import { parseAgentId } from '@codebuff/common/util/agent-id-parsing'
 import { generateCompactId } from '@codebuff/common/util/string'
 
-import { getAgentTemplate } from '../../../templates/agent-registry'
+import { loopAgentSteps } from '../../../run-agent-step'
 
-import type { loopAgentSteps } from '../../../run-agent-step'
 import type { AgentTemplate } from '@codebuff/common/types/agent-template'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 import type {
@@ -333,9 +333,6 @@ export async function executeSubagent(
   }
   onResponseChunk(startEvent)
 
-  // Import loopAgentSteps dynamically to avoid circular dependency
-  const { loopAgentSteps } = await import('../../../run-agent-step')
-
   const result = await loopAgentSteps({
     ...withDefaults,
     agentType: agentTemplate.id,
diff --git a/backend/src/websockets/websocket-action.ts b/backend/src/websockets/websocket-action.ts
@@ -1,3 +1,4 @@
+import { assembleLocalAgentTemplates } from '@codebuff/agent-runtime/templates/agent-registry'
 import { calculateUsageAndBalance } from '@codebuff/billing'
 import { trackEvent } from '@codebuff/common/analytics'
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
@@ -14,7 +15,6 @@ import {
 import { mainPrompt } from '../main-prompt'
 import { protec } from './middleware'
 import { sendActionWs } from '../client-wrapper'
-import { assembleLocalAgentTemplates } from '../templates/agent-registry'
 import { withLoggerContext } from '../util/logger'
 
 import type { ClientAction, UsageResponse } from '@codebuff/common/actions'
diff --git a/evals/scaffolding.ts b/evals/scaffolding.ts
@@ -3,6 +3,7 @@ import { EventEmitter } from 'events'
 import fs from 'fs'
 import path from 'path'
 
+import { assembleLocalAgentTemplates } from '@codebuff/agent-runtime/templates/agent-registry'
 import {
   handleStepsLogChunkWs,
   requestFilesWs,
@@ -13,7 +14,6 @@ import {
   sendSubagentChunkWs,
 } from '@codebuff/backend/client-wrapper'
 import { runAgentStep } from '@codebuff/backend/run-agent-step'
-import { assembleLocalAgentTemplates } from '@codebuff/backend/templates/agent-registry'
 import { getFileTokenScores } from '@codebuff/code-map/parse'
 import { TEST_USER_ID } from '@codebuff/common/old-constants'
 import { mockModule } from '@codebuff/common/testing/mock-modules'
diff --git a/packages/agent-runtime/src/templates/__tests__/agent-registry.test.ts b/packages/agent-runtime/src/templates/__tests__/agent-registry.test.ts
diff --git a/packages/agent-runtime/src/templates/agent-registry.ts b/packages/agent-runtime/src/templates/agent-registry.ts
