Skip to content

Commit 74b23a9

Browse files
charlesliencharleslien
committed
remove global agentRegistry
1 parent 671c1ee commit 74b23a9

17 files changed

+230
-448
lines changed

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

Lines changed: 8 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
import { AgentState } from '@codebuff/common/types/session-state'
22
import { FileTreeNode, ProjectFileContext } from '@codebuff/common/util/file'
3-
import { afterEach, beforeEach, describe, expect, it } from 'bun:test'
3+
import { describe, expect, it } from 'bun:test'
44

55
import { DynamicAgentTemplate } from '@codebuff/common/types/dynamic-agent-template'
6-
import { agentRegistry } from '../templates/agent-registry'
6+
import { getAllAgentTemplates } from '../templates/agent-registry'
77
import { collectParentInstructions } from '../templates/strings'
8-
import { AgentTemplate } from '../templates/types'
98

109
// Helper to create a mock ProjectFileContext
1110
const createMockFileContext = (
@@ -44,14 +43,6 @@ const createMockFileContext = (
4443
})
4544

4645
describe('Parent Instructions Injection', () => {
47-
beforeEach(() => {
48-
agentRegistry.reset()
49-
})
50-
51-
afterEach(() => {
52-
agentRegistry.reset()
53-
})
54-
5546
it('should inject parent instructions into userInputPrompt', async () => {
5647
// Mock file context with agent templates
5748
const fileContext = createMockFileContext({
@@ -101,12 +92,10 @@ describe('Parent Instructions Injection', () => {
10192
})
10293

10394
// Initialize the registry
104-
await agentRegistry.initialize(fileContext)
95+
const { agentRegistry } = await getAllAgentTemplates({ fileContext })
10596

10697
// Get the researcher template
107-
const researcherTemplate = agentRegistry.getTemplate(
108-
'researcher'
109-
) as AgentTemplate
98+
const researcherTemplate = agentRegistry['researcher']
11099
expect(researcherTemplate).toBeDefined()
111100

112101
// Create mock agent state
@@ -158,12 +147,10 @@ describe('Parent Instructions Injection', () => {
158147
})
159148

160149
// Initialize the registry
161-
await agentRegistry.initialize(fileContext)
150+
const { agentRegistry } = await getAllAgentTemplates({ fileContext })
162151

163152
// Get the researcher template
164-
const researcherTemplate = agentRegistry.getTemplate(
165-
'researcher'
166-
) as AgentTemplate
153+
const researcherTemplate = agentRegistry['researcher']
167154
expect(researcherTemplate).toBeDefined()
168155

169156
// Create mock agent state
@@ -256,12 +243,10 @@ describe('Parent Instructions Injection', () => {
256243
})
257244

258245
// Initialize the registry
259-
await agentRegistry.initialize(fileContext)
246+
const { agentRegistry } = await getAllAgentTemplates({ fileContext })
260247

261248
// Get the researcher template
262-
const researcherTemplate = agentRegistry.getTemplate(
263-
'researcher'
264-
) as AgentTemplate
249+
const researcherTemplate = agentRegistry['researcher']
265250
expect(researcherTemplate).toBeDefined()
266251

267252
// Create mock agent state

‎backend/src/__tests__/read-docs-tool.test.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import * as liveUserInputs from '../live-user-inputs'
2525
import * as context7Api from '../llm-apis/context7-api'
2626
import * as aisdk from '../llm-apis/vercel-ai-sdk/ai-sdk'
2727
import { runAgentStep } from '../run-agent-step'
28+
import { getAllAgentTemplates } from '../templates/agent-registry'
2829
import * as websocketAction from '../websockets/websocket-action'
2930
import { MockWebSocket, mockFileContext } from './test-utils'
3031

@@ -113,6 +114,9 @@ describe('read_docs tool with researcher agent', () => {
113114
...sessionState.mainAgentState,
114115
agentType: 'researcher' as const,
115116
}
117+
const { agentRegistry } = await getAllAgentTemplates({
118+
fileContext: mockFileContext,
119+
})
116120

117121
const { agentState: newAgentState } = await runAgentStep(
118122
new MockWebSocket() as unknown as WebSocket,
@@ -124,6 +128,7 @@ describe('read_docs tool with researcher agent', () => {
124128
onResponseChunk: () => {},
125129
agentType: 'researcher',
126130
fileContext: mockFileContext,
131+
agentRegistry,
127132
agentState,
128133
prompt: 'Get React documentation',
129134
params: undefined,
@@ -178,6 +183,9 @@ describe('read_docs tool with researcher agent', () => {
178183
...sessionState.mainAgentState,
179184
agentType: 'researcher' as const,
180185
}
186+
const { agentRegistry } = await getAllAgentTemplates({
187+
fileContext: mockFileContext,
188+
})
181189

182190
await runAgentStep(new MockWebSocket() as unknown as WebSocket, {
183191
userId: TEST_USER_ID,
@@ -187,6 +195,7 @@ describe('read_docs tool with researcher agent', () => {
187195
onResponseChunk: () => {},
188196
agentType: 'researcher',
189197
fileContext: mockFileContext,
198+
agentRegistry,
190199
agentState,
191200
prompt: 'Get React hooks documentation',
192201
params: undefined,
@@ -226,6 +235,9 @@ describe('read_docs tool with researcher agent', () => {
226235
...sessionState.mainAgentState,
227236
agentType: 'researcher' as const,
228237
}
238+
const { agentRegistry } = await getAllAgentTemplates({
239+
fileContext: mockFileContext,
240+
})
229241

230242
const { agentState: newAgentState } = await runAgentStep(
231243
new MockWebSocket() as unknown as WebSocket,
@@ -237,6 +249,7 @@ describe('read_docs tool with researcher agent', () => {
237249
onResponseChunk: () => {},
238250
agentType: 'researcher',
239251
fileContext: mockFileContext,
252+
agentRegistry,
240253
agentState,
241254
prompt: 'Get documentation for NonExistentLibrary',
242255
params: undefined,
@@ -285,6 +298,9 @@ describe('read_docs tool with researcher agent', () => {
285298
...sessionState.mainAgentState,
286299
agentType: 'researcher' as const,
287300
}
301+
const { agentRegistry } = await getAllAgentTemplates({
302+
fileContext: mockFileContext,
303+
})
288304

289305
const { agentState: newAgentState } = await runAgentStep(
290306
new MockWebSocket() as unknown as WebSocket,
@@ -296,6 +312,7 @@ describe('read_docs tool with researcher agent', () => {
296312
onResponseChunk: () => {},
297313
agentType: 'researcher',
298314
fileContext: mockFileContext,
315+
agentRegistry,
299316
agentState,
300317
prompt: 'Get React documentation',
301318
params: undefined,
@@ -344,6 +361,9 @@ describe('read_docs tool with researcher agent', () => {
344361
...sessionState.mainAgentState,
345362
agentType: 'researcher' as const,
346363
}
364+
const { agentRegistry } = await getAllAgentTemplates({
365+
fileContext: mockFileContext,
366+
})
347367

348368
const { agentState: newAgentState } = await runAgentStep(
349369
new MockWebSocket() as unknown as WebSocket,
@@ -355,6 +375,7 @@ describe('read_docs tool with researcher agent', () => {
355375
onResponseChunk: () => {},
356376
agentType: 'researcher',
357377
fileContext: mockFileContext,
378+
agentRegistry,
358379
agentState,
359380
prompt: 'Get React server components documentation',
360381
params: undefined,
@@ -401,6 +422,9 @@ describe('read_docs tool with researcher agent', () => {
401422
...sessionState.mainAgentState,
402423
agentType: 'researcher' as const,
403424
}
425+
const { agentRegistry } = await getAllAgentTemplates({
426+
fileContext: mockFileContext,
427+
})
404428

405429
const { agentState: newAgentState } = await runAgentStep(
406430
new MockWebSocket() as unknown as WebSocket,
@@ -412,6 +436,7 @@ describe('read_docs tool with researcher agent', () => {
412436
onResponseChunk: () => {},
413437
agentType: 'researcher',
414438
fileContext: mockFileContext,
439+
agentRegistry,
415440
agentState,
416441
prompt: 'Get React documentation',
417442
params: undefined,

‎backend/src/__tests__/run-agent-step-tools.test.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import {
2424
} from '@codebuff/common/testing/mock-modules'
2525
import * as aisdk from '../llm-apis/vercel-ai-sdk/ai-sdk'
2626
import { runAgentStep } from '../run-agent-step'
27+
import { getAllAgentTemplates } from '../templates/agent-registry'
2728
import * as tools from '../tools'
2829
import * as websocketAction from '../websockets/websocket-action'
2930

@@ -185,6 +186,9 @@ describe('runAgentStep - update_report tool', () => {
185186

186187
const sessionState = getInitialSessionState(mockFileContext)
187188
const agentState = sessionState.mainAgentState
189+
const { agentRegistry } = await getAllAgentTemplates({
190+
fileContext: mockFileContext,
191+
})
188192

189193
const result = await runAgentStep(
190194
new MockWebSocket() as unknown as WebSocket,
@@ -196,6 +200,7 @@ describe('runAgentStep - update_report tool', () => {
196200
onResponseChunk: () => {},
197201
agentType: 'base',
198202
fileContext: mockFileContext,
203+
agentRegistry,
199204
agentState,
200205
prompt: 'Analyze the codebase',
201206
params: undefined,
@@ -231,6 +236,9 @@ describe('runAgentStep - update_report tool', () => {
231236

232237
const sessionState = getInitialSessionState(mockFileContext)
233238
const agentState = sessionState.mainAgentState
239+
const { agentRegistry } = await getAllAgentTemplates({
240+
fileContext: mockFileContext,
241+
})
234242

235243
const result = await runAgentStep(
236244
new MockWebSocket() as unknown as WebSocket,
@@ -242,6 +250,7 @@ describe('runAgentStep - update_report tool', () => {
242250
onResponseChunk: () => {},
243251
agentType: 'base',
244252
fileContext: mockFileContext,
253+
agentRegistry,
245254
agentState,
246255
prompt: 'Analyze the codebase',
247256
params: undefined,
@@ -282,6 +291,9 @@ describe('runAgentStep - update_report tool', () => {
282291
existingField: 'original value',
283292
anotherField: 'unchanged',
284293
}
294+
const { agentRegistry } = await getAllAgentTemplates({
295+
fileContext: mockFileContext,
296+
})
285297

286298
const result = await runAgentStep(
287299
new MockWebSocket() as unknown as WebSocket,
@@ -293,6 +305,7 @@ describe('runAgentStep - update_report tool', () => {
293305
onResponseChunk: () => {},
294306
agentType: 'base',
295307
fileContext: mockFileContext,
308+
agentRegistry,
296309
agentState,
297310
prompt: 'Update the report',
298311
params: undefined,
@@ -325,6 +338,9 @@ describe('runAgentStep - update_report tool', () => {
325338
const sessionState = getInitialSessionState(mockFileContext)
326339
const agentState = sessionState.mainAgentState
327340
agentState.report = { existingField: 'value' }
341+
const { agentRegistry } = await getAllAgentTemplates({
342+
fileContext: mockFileContext,
343+
})
328344

329345
const result = await runAgentStep(
330346
new MockWebSocket() as unknown as WebSocket,
@@ -336,6 +352,7 @@ describe('runAgentStep - update_report tool', () => {
336352
onResponseChunk: () => {},
337353
agentType: 'base',
338354
fileContext: mockFileContext,
355+
agentRegistry,
339356
agentState,
340357
prompt: 'Update with empty object',
341358
params: undefined,

‎backend/src/__tests__/subagent-streaming.test.ts‎

Lines changed: 27 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -70,38 +70,34 @@ describe('Subagent Streaming', () => {
7070
})
7171

7272
// Mock agent registry
73-
spyOn(agentRegistryModule.agentRegistry, 'initialize').mockImplementation(
74-
async () => {}
75-
)
76-
spyOn(
77-
agentRegistryModule.agentRegistry,
78-
'getAllTemplates'
79-
).mockImplementation(() => ({
80-
thinker: {
81-
id: 'thinker',
82-
name: 'Thinker',
83-
outputMode: 'last_message',
84-
promptSchema: {
85-
prompt: {
86-
safeParse: () => ({ success: true }),
87-
} as any,
73+
spyOn(agentRegistryModule, 'getAllAgentTemplates').mockImplementation(
74+
async () => ({
75+
agentRegistry: {
76+
thinker: {
77+
id: 'thinker',
78+
name: 'Thinker',
79+
outputMode: 'last_message',
80+
promptSchema: {
81+
prompt: {
82+
safeParse: () => ({ success: true }),
83+
} as any,
84+
},
85+
purpose: '',
86+
model: '',
87+
includeMessageHistory: true,
88+
toolNames: [],
89+
spawnableAgents: [],
90+
initialAssistantMessage: '',
91+
initialAssistantPrefix: '',
92+
stepAssistantMessage: '',
93+
stepAssistantPrefix: '',
94+
systemPrompt: '',
95+
userInputPrompt: '',
96+
agentStepPrompt: '',
97+
},
8898
},
89-
purpose: '',
90-
model: '',
91-
includeMessageHistory: true,
92-
toolNames: [],
93-
spawnableAgents: [],
94-
initialAssistantMessage: '',
95-
initialAssistantPrefix: '',
96-
stepAssistantMessage: '',
97-
stepAssistantPrefix: '',
98-
systemPrompt: '',
99-
userInputPrompt: '',
100-
agentStepPrompt: '',
101-
},
102-
}))
103-
spyOn(agentRegistryModule.agentRegistry, 'getAgentName').mockImplementation(
104-
() => 'Thinker'
99+
validationErrors: [],
100+
})
105101
)
106102
})
107103

0 commit comments

Comments
 (0)