Skip to content

Commit 553b05f

Browse files
jahoomacodebuff-team
jahooma
and
codebuff-team
committed
Consolidate agent template types and simplify registry system
Replace AgentRegistry type with Record<string, AgentTemplate> across codebase. Update tests to use assembleLocalAgentTemplates instead of getAllAgentTemplates. Remove DynamicAgentLoadResult interface in favor of inline return type. Add missing mock agent templates to integration tests. 🤖 Generated with Codebuff Co-Authored-By: Codebuff <noreply@codebuff.com>
1 parent a8acd93 commit 553b05f

24 files changed

Lines changed: 613 additions & 345 deletions

‎backend/src/__tests__/agent-id-resolution.test.ts‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
import { resolveAgentId } from '@codebuff/common/util/agent-name-normalization'
22
import { describe, expect, it, beforeEach } from 'bun:test'
33

4-
import type { AgentRegistry } from '../templates/agent-registry'
4+
import type { AgentTemplate } from '../templates/types'
55

66
describe('Agent ID Resolution', () => {
7-
let mockRegistry: AgentRegistry
8-
7+
let mockRegistry: Record<string, AgentTemplate>
98
beforeEach(() => {
109
mockRegistry = {
1110
// Built-in agents

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,28 @@ const mockFileContext: ProjectFileContext = {
6262

6363
// --- Integration Test with Real LLM Call ---
6464
describe.skip('mainPrompt (Integration)', () => {
65+
let mockLocalAgentTemplates: Record<string, any>
66+
6567
beforeEach(() => {
68+
// Setup common mock agent templates
69+
mockLocalAgentTemplates = {
70+
base: {
71+
id: 'base',
72+
displayName: 'Base Agent',
73+
outputMode: 'last_message',
74+
inputSchema: {},
75+
parentPrompt: '',
76+
model: 'gpt-4o-mini',
77+
includeMessageHistory: true,
78+
toolNames: ['write_file', 'run_terminal_command'],
79+
subagents: [],
80+
systemPrompt: '',
81+
instructionsPrompt: '',
82+
stepPrompt: '',
83+
},
84+
}
85+
86+
6687
spyOn(websocketAction, 'requestToolCall').mockImplementation(
6788
async (
6889
ws: WebSocket,
@@ -362,6 +383,7 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
362383
} = await mainPrompt(new MockWebSocket() as unknown as WebSocket, action, {
363384
userId: TEST_USER_ID,
364385
clientSessionId: 'test-session-delete-function-integration',
386+
localAgentTemplates: mockLocalAgentTemplates,
365387
onResponseChunk: (chunk: string | PrintModeObject) => {
366388
if (typeof chunk !== 'string') {
367389
return
@@ -442,6 +464,22 @@ export function getMessagesSubset(messages: Message[], otherTokens: number) {
442464
await mainPrompt(new MockWebSocket() as unknown as WebSocket, action, {
443465
userId: TEST_USER_ID,
444466
clientSessionId: 'test-session-delete-function-integration',
467+
localAgentTemplates: {
468+
base: {
469+
id: 'base',
470+
displayName: 'Base Agent',
471+
outputMode: 'last_message',
472+
inputSchema: {},
473+
parentPrompt: '',
474+
model: 'gpt-4o-mini',
475+
includeMessageHistory: true,
476+
toolNames: ['write_file', 'run_terminal_command'],
477+
subagents: [],
478+
systemPrompt: '',
479+
instructionsPrompt: '',
480+
stepPrompt: '',
481+
},
482+
},
445483
onResponseChunk: (chunk: string | PrintModeObject) => {
446484
if (typeof chunk !== 'string') {
447485
return

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,42 @@ const mockAgentStream = (streamOutput: string) => {
4242
}
4343

4444
describe('mainPrompt', () => {
45+
let mockLocalAgentTemplates: Record<string, any>
46+
47+
beforeEach(() => {
48+
// Setup common mock agent templates
49+
mockLocalAgentTemplates = {
50+
base: {
51+
id: 'base',
52+
displayName: 'Base Agent',
53+
outputMode: 'last_message',
54+
inputSchema: {},
55+
parentPrompt: '',
56+
model: 'gpt-4o-mini',
57+
includeMessageHistory: true,
58+
toolNames: ['write_file', 'run_terminal_command'],
59+
subagents: [],
60+
systemPrompt: '',
61+
instructionsPrompt: '',
62+
stepPrompt: '',
63+
},
64+
base_max: {
65+
id: 'base_max',
66+
displayName: 'Base Max Agent',
67+
outputMode: 'last_message',
68+
inputSchema: {},
69+
parentPrompt: '',
70+
model: 'gpt-4o',
71+
includeMessageHistory: true,
72+
toolNames: ['write_file', 'run_terminal_command'],
73+
subagents: [],
74+
systemPrompt: '',
75+
instructionsPrompt: '',
76+
stepPrompt: '',
77+
},
78+
}
79+
})
80+
4581
beforeAll(() => {
4682
// Mock logger
4783
mockModule('@codebuff/backend/util/logger', () => ({
@@ -219,6 +255,36 @@ describe('mainPrompt', () => {
219255
userId: TEST_USER_ID,
220256
clientSessionId: 'test-session',
221257
onResponseChunk: () => {},
258+
localAgentTemplates: {
259+
base: {
260+
id: 'base',
261+
displayName: 'Base Agent',
262+
outputMode: 'last_message',
263+
inputSchema: {},
264+
parentPrompt: '',
265+
model: 'gpt-4o-mini',
266+
includeMessageHistory: true,
267+
toolNames: ['write_file', 'run_terminal_command'],
268+
subagents: [],
269+
systemPrompt: '',
270+
instructionsPrompt: '',
271+
stepPrompt: '',
272+
},
273+
base_max: {
274+
id: 'base_max',
275+
displayName: 'Base Max Agent',
276+
outputMode: 'last_message',
277+
inputSchema: {},
278+
parentPrompt: '',
279+
model: 'gpt-4o',
280+
includeMessageHistory: true,
281+
toolNames: ['write_file', 'run_terminal_command'],
282+
subagents: [],
283+
systemPrompt: '',
284+
instructionsPrompt: '',
285+
stepPrompt: '',
286+
},
287+
},
222288
},
223289
)
224290

@@ -271,6 +337,7 @@ describe('mainPrompt', () => {
271337
userId: TEST_USER_ID,
272338
clientSessionId: 'test-session',
273339
onResponseChunk: () => {},
340+
localAgentTemplates: mockLocalAgentTemplates,
274341
},
275342
)
276343

@@ -329,6 +396,36 @@ describe('mainPrompt', () => {
329396
userId: TEST_USER_ID,
330397
clientSessionId: 'test-session',
331398
onResponseChunk: () => {},
399+
localAgentTemplates: {
400+
base: {
401+
id: 'base',
402+
displayName: 'Base Agent',
403+
outputMode: 'last_message',
404+
inputSchema: {},
405+
parentPrompt: '',
406+
model: 'gpt-4o-mini',
407+
includeMessageHistory: true,
408+
toolNames: ['write_file', 'run_terminal_command'],
409+
subagents: [],
410+
systemPrompt: '',
411+
instructionsPrompt: '',
412+
stepPrompt: '',
413+
},
414+
base_max: {
415+
id: 'base_max',
416+
displayName: 'Base Max Agent',
417+
outputMode: 'last_message',
418+
inputSchema: {},
419+
parentPrompt: '',
420+
model: 'gpt-4o',
421+
includeMessageHistory: true,
422+
toolNames: ['write_file', 'run_terminal_command'],
423+
subagents: [],
424+
systemPrompt: '',
425+
instructionsPrompt: '',
426+
stepPrompt: '',
427+
},
428+
},
332429
})
333430

334431
// Assert that requestToolCall was called exactly once
@@ -374,6 +471,7 @@ describe('mainPrompt', () => {
374471
userId: TEST_USER_ID,
375472
clientSessionId: 'test-session',
376473
onResponseChunk: () => {},
474+
localAgentTemplates: mockLocalAgentTemplates,
377475
},
378476
)
379477

@@ -401,6 +499,7 @@ describe('mainPrompt', () => {
401499
userId: TEST_USER_ID,
402500
clientSessionId: 'test-session',
403501
onResponseChunk: () => {},
502+
localAgentTemplates: mockLocalAgentTemplates,
404503
},
405504
)
406505

@@ -432,6 +531,7 @@ describe('mainPrompt', () => {
432531
userId: TEST_USER_ID,
433532
clientSessionId: 'test-session',
434533
onResponseChunk: () => {},
534+
localAgentTemplates: mockLocalAgentTemplates,
435535
},
436536
)
437537

@@ -461,6 +561,7 @@ describe('mainPrompt', () => {
461561
userId: TEST_USER_ID,
462562
clientSessionId: 'test-session',
463563
onResponseChunk: () => {},
564+
localAgentTemplates: mockLocalAgentTemplates,
464565
},
465566
)
466567

@@ -498,6 +599,7 @@ describe('mainPrompt', () => {
498599
userId: TEST_USER_ID,
499600
clientSessionId: 'test-session',
500601
onResponseChunk: () => {},
602+
localAgentTemplates: mockLocalAgentTemplates,
501603
})
502604

503605
// Assert that requestToolCall was called exactly once

‎backend/src/__tests__/parent-instructions.test.ts‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { describe, expect, it } from 'bun:test'
22

3-
import { getAllAgentTemplates } from '../templates/agent-registry'
3+
import { assembleLocalAgentTemplates } from '../templates/agent-registry'
44
import { collectParentInstructions } from '../templates/strings'
55

66
import type { DynamicAgentTemplate } from '@codebuff/common/types/dynamic-agent-template'
@@ -87,10 +87,10 @@ describe('Parent Instructions Injection', () => {
8787
})
8888

8989
// Initialize the registry
90-
const { agentRegistry } = await getAllAgentTemplates({ fileContext })
90+
const { agentTemplates } = assembleLocalAgentTemplates(fileContext)
9191

9292
// Get the researcher template
93-
const researcherTemplate = agentRegistry['researcher']
93+
const researcherTemplate = agentTemplates['researcher']
9494
expect(researcherTemplate).toBeDefined()
9595

9696
// Create mock agent state
@@ -107,7 +107,7 @@ describe('Parent Instructions Injection', () => {
107107
// Test parent instructions collection directly
108108
const parentInstructions = await collectParentInstructions(
109109
'researcher',
110-
agentRegistry,
110+
agentTemplates,
111111
)
112112

113113
// Verify that parent instructions are collected
@@ -137,16 +137,16 @@ describe('Parent Instructions Injection', () => {
137137
})
138138

139139
// Initialize the registry
140-
const { agentRegistry } = await getAllAgentTemplates({ fileContext })
140+
const { agentTemplates } = assembleLocalAgentTemplates(fileContext)
141141

142142
// Get the researcher template
143-
const researcherTemplate = agentRegistry['researcher']
143+
const researcherTemplate = agentTemplates['researcher']
144144
expect(researcherTemplate).toBeDefined()
145145

146146
// Test parent instructions collection directly
147147
const parentInstructions = await collectParentInstructions(
148148
'researcher',
149-
agentRegistry,
149+
agentTemplates,
150150
)
151151

152152
// Verify that no parent instructions are collected
@@ -207,16 +207,16 @@ describe('Parent Instructions Injection', () => {
207207
})
208208

209209
// Initialize the registry
210-
const { agentRegistry } = await getAllAgentTemplates({ fileContext })
210+
const { agentTemplates } = assembleLocalAgentTemplates(fileContext)
211211

212212
// Get the researcher template
213-
const researcherTemplate = agentRegistry['researcher']
213+
const researcherTemplate = agentTemplates['researcher']
214214
expect(researcherTemplate).toBeDefined()
215215

216216
// Test parent instructions collection directly
217217
const parentInstructions = await collectParentInstructions(
218218
'researcher',
219-
agentRegistry,
219+
agentTemplates,
220220
)
221221

222222
// Verify that both parent instructions are collected

0 commit comments

Comments
 (0)