Skip to content

Commit e423dcf

Browse files
jahoomajahooma
committed
Merge branch 'main' into set_output
2 parents c6f05b6 + 74b23a9 commit e423dcf

18 files changed

+231
-449
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
// Test parent instructions collection directly
@@ -245,12 +232,10 @@ describe('Parent Instructions Injection', () => {
245232
})
246233

247234
// Initialize the registry
248-
await agentRegistry.initialize(fileContext)
235+
const { agentRegistry } = await getAllAgentTemplates({ fileContext })
249236

250237
// Get the researcher template
251-
const researcherTemplate = agentRegistry.getTemplate(
252-
'researcher'
253-
) as AgentTemplate
238+
const researcherTemplate = agentRegistry['researcher']
254239
expect(researcherTemplate).toBeDefined()
255240

256241
// Test parent instructions collection directly

‎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

@@ -183,6 +184,9 @@ describe('runAgentStep - update_report tool', () => {
183184

184185
const sessionState = getInitialSessionState(mockFileContext)
185186
const agentState = sessionState.mainAgentState
187+
const { agentRegistry } = await getAllAgentTemplates({
188+
fileContext: mockFileContext,
189+
})
186190

187191
const result = await runAgentStep(
188192
new MockWebSocket() as unknown as WebSocket,
@@ -194,6 +198,7 @@ describe('runAgentStep - update_report tool', () => {
194198
onResponseChunk: () => {},
195199
agentType: 'base',
196200
fileContext: mockFileContext,
201+
agentRegistry,
197202
agentState,
198203
prompt: 'Analyze the codebase',
199204
params: undefined,
@@ -227,6 +232,9 @@ describe('runAgentStep - update_report tool', () => {
227232

228233
const sessionState = getInitialSessionState(mockFileContext)
229234
const agentState = sessionState.mainAgentState
235+
const { agentRegistry } = await getAllAgentTemplates({
236+
fileContext: mockFileContext,
237+
})
230238

231239
const result = await runAgentStep(
232240
new MockWebSocket() as unknown as WebSocket,
@@ -238,6 +246,7 @@ describe('runAgentStep - update_report tool', () => {
238246
onResponseChunk: () => {},
239247
agentType: 'base',
240248
fileContext: mockFileContext,
249+
agentRegistry,
241250
agentState,
242251
prompt: 'Analyze the codebase',
243252
params: undefined,
@@ -276,6 +285,9 @@ describe('runAgentStep - update_report tool', () => {
276285
existingField: 'original value',
277286
anotherField: 'unchanged',
278287
}
288+
const { agentRegistry } = await getAllAgentTemplates({
289+
fileContext: mockFileContext,
290+
})
279291

280292
const result = await runAgentStep(
281293
new MockWebSocket() as unknown as WebSocket,
@@ -287,6 +299,7 @@ describe('runAgentStep - update_report tool', () => {
287299
onResponseChunk: () => {},
288300
agentType: 'base',
289301
fileContext: mockFileContext,
302+
agentRegistry,
290303
agentState,
291304
prompt: 'Update the output',
292305
params: undefined,
@@ -313,6 +326,9 @@ describe('runAgentStep - update_report tool', () => {
313326
const sessionState = getInitialSessionState(mockFileContext)
314327
const agentState = sessionState.mainAgentState
315328
agentState.output = { existingField: 'value' }
329+
const { agentRegistry } = await getAllAgentTemplates({
330+
fileContext: mockFileContext,
331+
})
316332

317333
const result = await runAgentStep(
318334
new MockWebSocket() as unknown as WebSocket,
@@ -324,6 +340,7 @@ describe('runAgentStep - update_report tool', () => {
324340
onResponseChunk: () => {},
325341
agentType: 'base',
326342
fileContext: mockFileContext,
343+
agentRegistry,
327344
agentState,
328345
prompt: 'Update with empty object',
329346
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)