Skip to content

Commit 6f17f50

Browse files
charlesliencodebuff-team
charleslien
and
codebuff-team
committed
Refactor run-agent-step.ts to pass logger as parameter and consolidate ws into options object
🤖 Generated with Codebuff Co-Authored-By: Codebuff <noreply@codebuff.com>
1 parent be23932 commit 6f17f50

14 files changed

+611
-565
lines changed

‎backend/src/__tests__/loop-agent-steps.test.ts‎

Lines changed: 205 additions & 211 deletions
Large diffs are not rendered by default.

‎backend/src/__tests__/prompt-caching-subagents.test.ts‎

Lines changed: 30 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
176176
const ws = new MockWebSocket() as unknown as WebSocket
177177

178178
// Run parent agent first to establish system prompt
179-
const parentResult = await loopAgentSteps(ws, {
179+
const parentResult = await loopAgentSteps({
180+
ws,
180181
userInputId: 'test-parent',
181182
prompt: 'Parent task',
182183
params: undefined,
@@ -188,6 +189,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
188189
userId: TEST_USER_ID,
189190
clientSessionId: 'test-session',
190191
onResponseChunk: () => {},
192+
logger: console as any,
191193
})
192194

193195
// Capture parent's messages which include the system prompt
@@ -208,7 +210,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
208210
messageHistory: [],
209211
}
210212

211-
await loopAgentSteps(ws, {
213+
await loopAgentSteps({
214+
ws,
212215
userInputId: 'test-child',
213216
prompt: 'Child task',
214217
params: undefined,
@@ -221,6 +224,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
221224
clientSessionId: 'test-session',
222225
onResponseChunk: () => {},
223226
parentSystemPrompt: parentSystemPrompt,
227+
logger: console as any,
224228
})
225229

226230
// Verify child uses parent's system prompt
@@ -255,7 +259,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
255259
mockLocalAgentTemplates['standalone-child'] = standaloneChild
256260

257261
// Run parent agent first
258-
const parentResult = await loopAgentSteps(ws, {
262+
const parentResult = await loopAgentSteps({
263+
ws,
259264
userInputId: 'test-parent',
260265
prompt: 'Parent task',
261266
params: undefined,
@@ -267,6 +272,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
267272
userId: TEST_USER_ID,
268273
clientSessionId: 'test-session',
269274
onResponseChunk: () => {},
275+
logger: console as any,
270276
})
271277

272278
const parentMessages = capturedMessages
@@ -281,7 +287,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
281287
messageHistory: [],
282288
}
283289

284-
await loopAgentSteps(ws, {
290+
await loopAgentSteps({
291+
ws,
285292
userInputId: 'test-child',
286293
prompt: 'Child task',
287294
params: undefined,
@@ -294,6 +301,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
294301
clientSessionId: 'test-session',
295302
onResponseChunk: () => {},
296303
parentSystemPrompt: parentSystemPrompt,
304+
logger: console as any,
297305
})
298306

299307
const childMessages = capturedMessages
@@ -329,7 +337,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
329337
mockLocalAgentTemplates['message-history-child'] = messageHistoryChild
330338

331339
// Run parent agent first
332-
await loopAgentSteps(ws, {
340+
await loopAgentSteps({
341+
ws,
333342
userInputId: 'test-parent',
334343
prompt: 'Parent task',
335344
params: undefined,
@@ -341,6 +350,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
341350
userId: TEST_USER_ID,
342351
clientSessionId: 'test-session',
343352
onResponseChunk: () => {},
353+
logger: console as any,
344354
})
345355

346356
const parentMessages = capturedMessages
@@ -358,7 +368,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
358368
],
359369
}
360370

361-
await loopAgentSteps(ws, {
371+
await loopAgentSteps({
372+
ws,
362373
userInputId: 'test-child',
363374
prompt: 'Child task',
364375
params: undefined,
@@ -371,6 +382,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
371382
clientSessionId: 'test-session',
372383
onResponseChunk: () => {},
373384
parentSystemPrompt: parentSystemPrompt,
385+
logger: console as any,
374386
})
375387

376388
const childMessages = capturedMessages
@@ -432,7 +444,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
432444
const ws = new MockWebSocket() as unknown as WebSocket
433445

434446
// Run parent agent
435-
const parentResult = await loopAgentSteps(ws, {
447+
const parentResult = await loopAgentSteps({
448+
ws,
436449
userInputId: 'test-parent',
437450
prompt: 'Parent task',
438451
params: undefined,
@@ -444,6 +457,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
444457
userId: TEST_USER_ID,
445458
clientSessionId: 'test-session',
446459
onResponseChunk: () => {},
460+
logger: console as any,
447461
})
448462

449463
const parentMessages = capturedMessages
@@ -458,7 +472,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
458472
messageHistory: [],
459473
}
460474

461-
await loopAgentSteps(ws, {
475+
await loopAgentSteps({
476+
ws,
462477
userInputId: 'test-child',
463478
prompt: 'Child task',
464479
params: undefined,
@@ -471,6 +486,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
471486
clientSessionId: 'test-session',
472487
onResponseChunk: () => {},
473488
parentSystemPrompt: parentSystemPrompt,
489+
logger: console as any,
474490
})
475491

476492
const childMessages = capturedMessages
@@ -510,7 +526,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
510526
mockLocalAgentTemplates['full-inherit-child'] = fullInheritChild
511527

512528
// Run parent agent first with some message history
513-
const parentResult = await loopAgentSteps(ws, {
529+
const parentResult = await loopAgentSteps({
530+
ws,
514531
userInputId: 'test-parent',
515532
prompt: 'Parent task',
516533
params: undefined,
@@ -528,6 +545,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
528545
userId: TEST_USER_ID,
529546
clientSessionId: 'test-session',
530547
onResponseChunk: () => {},
548+
logger: console as any,
531549
})
532550

533551
const parentMessages = capturedMessages
@@ -545,7 +563,8 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
545563
],
546564
}
547565

548-
await loopAgentSteps(ws, {
566+
await loopAgentSteps({
567+
ws,
549568
userInputId: 'test-child',
550569
prompt: 'Child task',
551570
params: undefined,
@@ -558,6 +577,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
558577
clientSessionId: 'test-session',
559578
onResponseChunk: () => {},
560579
parentSystemPrompt: parentSystemPrompt,
580+
logger: console as any,
561581
})
562582

563583
const childMessages = capturedMessages

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

Lines changed: 83 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -297,23 +297,22 @@ describe('read_docs tool with researcher agent', () => {
297297
logger,
298298
})
299299

300-
const { agentState: newAgentState } = await runAgentStep(
301-
new MockWebSocket() as unknown as WebSocket,
302-
{
303-
system: 'Test system prompt',
304-
userId: TEST_USER_ID,
305-
userInputId: 'test-input',
306-
clientSessionId: 'test-session',
307-
fingerprintId: 'test-fingerprint',
308-
onResponseChunk: () => {},
309-
agentType: 'researcher',
310-
fileContext: mockFileContextWithAgents,
311-
localAgentTemplates: agentTemplates,
312-
agentState,
313-
prompt: 'Get React documentation',
314-
params: undefined,
315-
},
316-
)
300+
const { agentState: newAgentState } = await runAgentStep({
301+
ws: new MockWebSocket() as unknown as WebSocket,
302+
system: 'Test system prompt',
303+
userId: TEST_USER_ID,
304+
userInputId: 'test-input',
305+
clientSessionId: 'test-session',
306+
fingerprintId: 'test-fingerprint',
307+
onResponseChunk: () => {},
308+
agentType: 'researcher',
309+
fileContext: mockFileContextWithAgents,
310+
localAgentTemplates: agentTemplates,
311+
agentState,
312+
prompt: 'Get React documentation',
313+
params: undefined,
314+
logger,
315+
})
317316

318317
expect(context7Api.fetchContext7LibraryDocumentation).toHaveBeenCalledWith(
319318
'React',
@@ -372,7 +371,8 @@ describe('read_docs tool with researcher agent', () => {
372371
logger,
373372
})
374373

375-
await runAgentStep(new MockWebSocket() as unknown as WebSocket, {
374+
await runAgentStep({
375+
ws: new MockWebSocket() as unknown as WebSocket,
376376
system: 'Test system prompt',
377377
userId: TEST_USER_ID,
378378
userInputId: 'test-input',
@@ -385,6 +385,7 @@ describe('read_docs tool with researcher agent', () => {
385385
agentState,
386386
prompt: 'Get React hooks documentation',
387387
params: undefined,
388+
logger,
388389
})
389390

390391
expect(context7Api.fetchContext7LibraryDocumentation).toHaveBeenCalledWith(
@@ -421,23 +422,22 @@ describe('read_docs tool with researcher agent', () => {
421422
logger,
422423
})
423424

424-
const { agentState: newAgentState } = await runAgentStep(
425-
new MockWebSocket() as unknown as WebSocket,
426-
{
427-
system: 'Test system prompt',
428-
userId: TEST_USER_ID,
429-
userInputId: 'test-input',
430-
clientSessionId: 'test-session',
431-
fingerprintId: 'test-fingerprint',
432-
onResponseChunk: () => {},
433-
agentType: 'researcher',
434-
fileContext: mockFileContextWithAgents,
435-
localAgentTemplates: agentTemplates,
436-
agentState,
437-
prompt: 'Get documentation for NonExistentLibrary',
438-
params: undefined,
439-
},
440-
)
425+
const { agentState: newAgentState } = await runAgentStep({
426+
ws: new MockWebSocket() as unknown as WebSocket,
427+
system: 'Test system prompt',
428+
userId: TEST_USER_ID,
429+
userInputId: 'test-input',
430+
clientSessionId: 'test-session',
431+
fingerprintId: 'test-fingerprint',
432+
onResponseChunk: () => {},
433+
agentType: 'researcher',
434+
fileContext: mockFileContextWithAgents,
435+
localAgentTemplates: agentTemplates,
436+
agentState,
437+
prompt: 'Get documentation for NonExistentLibrary',
438+
params: undefined,
439+
logger,
440+
})
441441

442442
// Check that the "no documentation found" message was added
443443
const toolResultMessages = newAgentState.messageHistory.filter(
@@ -489,23 +489,22 @@ describe('read_docs tool with researcher agent', () => {
489489
logger,
490490
})
491491

492-
const { agentState: newAgentState } = await runAgentStep(
493-
new MockWebSocket() as unknown as WebSocket,
494-
{
495-
system: 'Test system prompt',
496-
userId: TEST_USER_ID,
497-
userInputId: 'test-input',
498-
clientSessionId: 'test-session',
499-
fingerprintId: 'test-fingerprint',
500-
onResponseChunk: () => {},
501-
agentType: 'researcher',
502-
fileContext: mockFileContextWithAgents,
503-
localAgentTemplates: agentTemplates,
504-
agentState,
505-
prompt: 'Get React documentation',
506-
params: undefined,
507-
},
508-
)
492+
const { agentState: newAgentState } = await runAgentStep({
493+
ws: new MockWebSocket() as unknown as WebSocket,
494+
system: 'Test system prompt',
495+
userId: TEST_USER_ID,
496+
userInputId: 'test-input',
497+
clientSessionId: 'test-session',
498+
fingerprintId: 'test-fingerprint',
499+
onResponseChunk: () => {},
500+
agentType: 'researcher',
501+
fileContext: mockFileContextWithAgents,
502+
localAgentTemplates: agentTemplates,
503+
agentState,
504+
prompt: 'Get React documentation',
505+
params: undefined,
506+
logger,
507+
})
509508

510509
// Check that the error message was added
511510
const toolResultMessages = newAgentState.messageHistory.filter(
@@ -556,23 +555,22 @@ describe('read_docs tool with researcher agent', () => {
556555
logger,
557556
})
558557

559-
const { agentState: newAgentState } = await runAgentStep(
560-
new MockWebSocket() as unknown as WebSocket,
561-
{
562-
system: 'Test system prompt',
563-
userId: TEST_USER_ID,
564-
userInputId: 'test-input',
565-
clientSessionId: 'test-session',
566-
fingerprintId: 'test-fingerprint',
567-
onResponseChunk: () => {},
568-
agentType: 'researcher',
569-
fileContext: mockFileContextWithAgents,
570-
localAgentTemplates: agentTemplates,
571-
agentState,
572-
prompt: 'Get React server components documentation',
573-
params: undefined,
574-
},
575-
)
558+
const { agentState: newAgentState } = await runAgentStep({
559+
ws: new MockWebSocket() as unknown as WebSocket,
560+
system: 'Test system prompt',
561+
userId: TEST_USER_ID,
562+
userInputId: 'test-input',
563+
clientSessionId: 'test-session',
564+
fingerprintId: 'test-fingerprint',
565+
onResponseChunk: () => {},
566+
agentType: 'researcher',
567+
fileContext: mockFileContextWithAgents,
568+
localAgentTemplates: agentTemplates,
569+
agentState,
570+
prompt: 'Get React server components documentation',
571+
params: undefined,
572+
logger,
573+
})
576574

577575
// Check that the topic is included in the error message
578576
const toolResultMessages = newAgentState.messageHistory.filter(
@@ -624,23 +622,22 @@ describe('read_docs tool with researcher agent', () => {
624622
logger,
625623
})
626624

627-
const { agentState: newAgentState } = await runAgentStep(
628-
new MockWebSocket() as unknown as WebSocket,
629-
{
630-
system: 'Test system prompt',
631-
userId: TEST_USER_ID,
632-
userInputId: 'test-input',
633-
clientSessionId: 'test-session',
634-
fingerprintId: 'test-fingerprint',
635-
onResponseChunk: () => {},
636-
agentType: 'researcher',
637-
fileContext: mockFileContextWithAgents,
638-
localAgentTemplates: agentTemplates,
639-
agentState,
640-
prompt: 'Get React documentation',
641-
params: undefined,
642-
},
643-
)
625+
const { agentState: newAgentState } = await runAgentStep({
626+
ws: new MockWebSocket() as unknown as WebSocket,
627+
system: 'Test system prompt',
628+
userId: TEST_USER_ID,
629+
userInputId: 'test-input',
630+
clientSessionId: 'test-session',
631+
fingerprintId: 'test-fingerprint',
632+
onResponseChunk: () => {},
633+
agentType: 'researcher',
634+
fileContext: mockFileContextWithAgents,
635+
localAgentTemplates: agentTemplates,
636+
agentState,
637+
prompt: 'Get React documentation',
638+
params: undefined,
639+
logger,
640+
})
644641

645642
// Check that the generic error message was added
646643
const toolResultMessages = newAgentState.messageHistory.filter(

0 commit comments

Comments
 (0)