Skip to content

Commit c6d7a32

Browse files
charlesliencharleslien
committed
remove state from function handlers
1 parent a542e7f commit c6d7a32

35 files changed

+778
-1066
lines changed

‎packages/agent-runtime/src/__tests__/cost-aggregation.test.ts‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -183,14 +183,12 @@ describe('Cost Aggregation System', () => {
183183
},
184184
}
185185

186-
const result = handleSpawnAgents({
186+
await handleSpawnAgents({
187187
...params,
188188
agentState: parentAgentState,
189189
toolCall: mockToolCall,
190190
})
191191

192-
await result.result
193-
194192
// Parent should have aggregated costs: original 50 + subagent 75 + subagent 100 = 225
195193
expect(parentAgentState.creditsUsed).toBe(225)
196194
expect(mockExecuteAgent).toHaveBeenCalledTimes(2)
@@ -255,14 +253,12 @@ describe('Cost Aggregation System', () => {
255253
},
256254
}
257255

258-
const result = handleSpawnAgents({
256+
await handleSpawnAgents({
259257
...params,
260258
agentState: parentAgentState,
261259
toolCall: mockToolCall,
262260
})
263261

264-
await result.result
265-
266262
// Parent should aggregate costs: original 10 + successful subagent 50 + failed subagent 25 = 85
267263
expect(parentAgentState.creditsUsed).toBe(85)
268264
})
@@ -399,14 +395,12 @@ describe('Cost Aggregation System', () => {
399395
},
400396
}
401397

402-
const result = handleSpawnAgents({
398+
await handleSpawnAgents({
403399
...params,
404400
agentState: mainAgentState,
405401
toolCall: mockToolCall,
406402
})
407403

408-
await result.result
409-
410404
// Verify exact cost accounting
411405
expect(mainAgentState.creditsUsed).toBe(expectedTotal)
412406

‎packages/agent-runtime/src/__tests__/spawn-agents-message-history.test.ts‎

Lines changed: 4 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -128,16 +128,14 @@ describe('Spawn Agents Message History', () => {
128128
userMessage('How are you?'),
129129
]
130130

131-
const { result } = handleSpawnAgents({
131+
await handleSpawnAgents({
132132
...handleSpawnAgentsBaseParams,
133133
agentState: sessionState.mainAgentState,
134134
agentTemplate: parentAgent,
135135
localAgentTemplates: { 'child-agent': childAgent },
136136
toolCall,
137137
})
138138

139-
await result
140-
141139
// Verify that the spawned agent was called
142140
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
143141

@@ -188,16 +186,14 @@ describe('Spawn Agents Message History', () => {
188186
assistantMessage('Hi there!'),
189187
]
190188

191-
const { result } = handleSpawnAgents({
189+
await handleSpawnAgents({
192190
...handleSpawnAgentsBaseParams,
193191
agentState: sessionState.mainAgentState,
194192
agentTemplate: parentAgent,
195193
localAgentTemplates: { 'child-agent': childAgent },
196194
toolCall,
197195
})
198196

199-
await result
200-
201197
// Verify that the subagent's message history is empty when includeMessageHistory is false
202198
expect(capturedSubAgentState.messageHistory).toHaveLength(0)
203199
})
@@ -210,16 +206,14 @@ describe('Spawn Agents Message History', () => {
210206

211207
sessionState.mainAgentState.messageHistory = [] // Empty message history
212208

213-
const { result } = handleSpawnAgents({
209+
await handleSpawnAgents({
214210
...handleSpawnAgentsBaseParams,
215211
agentState: sessionState.mainAgentState,
216212
agentTemplate: parentAgent,
217213
localAgentTemplates: { 'child-agent': childAgent },
218214
toolCall,
219215
})
220216

221-
await result
222-
223217
// Verify that the subagent's message history is empty when there are no messages to pass
224218
expect(capturedSubAgentState.messageHistory).toHaveLength(0)
225219
})
@@ -235,16 +229,14 @@ describe('Spawn Agents Message History', () => {
235229
systemMessage('System prompt 2'),
236230
]
237231

238-
const { result } = handleSpawnAgents({
232+
await handleSpawnAgents({
239233
...handleSpawnAgentsBaseParams,
240234
agentState: sessionState.mainAgentState,
241235
agentTemplate: parentAgent,
242236
localAgentTemplates: { 'child-agent': childAgent },
243237
toolCall,
244238
})
245239

246-
await result
247-
248240
// Verify that system messages without timeToLive are included
249241
// expireMessages only filters messages with timeToLive='userPrompt'
250242
expect(capturedSubAgentState.messageHistory).toHaveLength(2)

‎packages/agent-runtime/src/__tests__/spawn-agents-permissions.test.ts‎

Lines changed: 19 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -253,15 +253,14 @@ describe('Spawn Agents Permissions', () => {
253253
const sessionState = getInitialSessionState(mockFileContext)
254254
const toolCall = createSpawnToolCall('thinker')
255255

256-
const { result } = handleSpawnAgents({
256+
const { output } = await handleSpawnAgents({
257257
...handleSpawnAgentsBaseParams,
258258
agentState: sessionState.mainAgentState,
259259
agentTemplate: parentAgent,
260260
localAgentTemplates: { thinker: childAgent },
261261
toolCall,
262262
})
263263

264-
const output = await result
265264
expect(JSON.stringify(output)).toContain('Mock agent response')
266265
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
267266
})
@@ -272,15 +271,14 @@ describe('Spawn Agents Permissions', () => {
272271
const sessionState = getInitialSessionState(mockFileContext)
273272
const toolCall = createSpawnToolCall('reviewer') // Try to spawn reviewer
274273

275-
const { result } = handleSpawnAgents({
274+
const { output } = await handleSpawnAgents({
276275
...handleSpawnAgentsBaseParams,
277276
agentState: sessionState.mainAgentState,
278277
agentTemplate: parentAgent,
279278
localAgentTemplates: { reviewer: childAgent },
280279
toolCall,
281280
})
282281

283-
const output = await result
284282
expect(JSON.stringify(output)).toContain('Error spawning agent')
285283
expect(JSON.stringify(output)).toContain(
286284
'is not allowed to spawn child agent type reviewer',
@@ -293,15 +291,14 @@ describe('Spawn Agents Permissions', () => {
293291
const sessionState = getInitialSessionState(mockFileContext)
294292
const toolCall = createSpawnToolCall('nonexistent')
295293

296-
const { result } = handleSpawnAgents({
294+
const { output } = await handleSpawnAgents({
297295
...handleSpawnAgentsBaseParams,
298296
agentState: sessionState.mainAgentState,
299297
agentTemplate: parentAgent,
300298
localAgentTemplates: {}, // Empty - agent not found
301299
toolCall,
302300
})
303301

304-
const output = await result
305302
console.log('output', output)
306303
expect(JSON.stringify(output)).toContain('Error spawning agent')
307304
expect(JSON.stringify(output)).toContain(
@@ -316,15 +313,14 @@ describe('Spawn Agents Permissions', () => {
316313
const sessionState = getInitialSessionState(mockFileContext)
317314
const toolCall = createSpawnToolCall('codebuff/thinker@1.0.0')
318315

319-
const { result } = handleSpawnAgents({
316+
const { output } = await handleSpawnAgents({
320317
...handleSpawnAgentsBaseParams,
321318
agentState: sessionState.mainAgentState,
322319
agentTemplate: parentAgent,
323320
localAgentTemplates: { 'codebuff/thinker@1.0.0': childAgent },
324321
toolCall,
325322
})
326323

327-
const output = await result
328324
expect(JSON.stringify(output)).toContain('Mock agent response')
329325
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
330326
})
@@ -335,7 +331,7 @@ describe('Spawn Agents Permissions', () => {
335331
const sessionState = getInitialSessionState(mockFileContext)
336332
const toolCall = createSpawnToolCall('thinker') // Simple name
337333

338-
const { result } = handleSpawnAgents({
334+
const { output } = await handleSpawnAgents({
339335
...handleSpawnAgentsBaseParams,
340336
agentState: sessionState.mainAgentState,
341337
agentTemplate: parentAgent,
@@ -346,7 +342,6 @@ describe('Spawn Agents Permissions', () => {
346342
toolCall,
347343
})
348344

349-
const output = await result
350345
expect(JSON.stringify(output)).toContain('Mock agent response')
351346
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
352347
})
@@ -357,15 +352,14 @@ describe('Spawn Agents Permissions', () => {
357352
const sessionState = getInitialSessionState(mockFileContext)
358353
const toolCall = createSpawnToolCall('codebuff/thinker@2.0.0')
359354

360-
const { result } = handleSpawnAgents({
355+
const { output } = await handleSpawnAgents({
361356
...handleSpawnAgentsBaseParams,
362357
agentState: sessionState.mainAgentState,
363358
agentTemplate: parentAgent,
364359
localAgentTemplates: { 'codebuff/thinker@2.0.0': childAgent },
365360
toolCall,
366361
})
367362

368-
const output = await result
369363
expect(JSON.stringify(output)).toContain('Error spawning agent')
370364
expect(JSON.stringify(output)).toContain(
371365
'is not allowed to spawn child agent type',
@@ -390,7 +384,7 @@ describe('Spawn Agents Permissions', () => {
390384
},
391385
}
392386

393-
const { result } = handleSpawnAgents({
387+
const { output } = await handleSpawnAgents({
394388
...handleSpawnAgentsBaseParams,
395389
agentState: sessionState.mainAgentState,
396390
agentTemplate: parentAgent,
@@ -401,7 +395,6 @@ describe('Spawn Agents Permissions', () => {
401395
toolCall,
402396
})
403397

404-
const output = await result
405398
expect(JSON.stringify(output)).toContain('Mock agent response') // Successful thinker spawn
406399
expect(JSON.stringify(output)).toContain('Error spawning agent') // Failed reviewer spawn
407400
expect(JSON.stringify(output)).toContain(
@@ -430,15 +423,15 @@ describe('Spawn Agents Permissions', () => {
430423
const sessionState = getInitialSessionState(mockFileContext)
431424
const toolCall = createInlineSpawnToolCall('thinker')
432425

433-
const { result } = handleSpawnAgentInline({
426+
// Should not throw
427+
await handleSpawnAgentInline({
434428
...handleSpawnAgentsBaseParams,
435429
agentState: sessionState.mainAgentState,
436430
agentTemplate: parentAgent,
437431
localAgentTemplates: { thinker: childAgent },
438432
toolCall,
439433
})
440434

441-
await result // Should not throw
442435
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
443436
})
444437

@@ -448,15 +441,15 @@ describe('Spawn Agents Permissions', () => {
448441
const sessionState = getInitialSessionState(mockFileContext)
449442
const toolCall = createInlineSpawnToolCall('reviewer') // Try to spawn reviewer
450443

451-
const { result } = handleSpawnAgentInline({
444+
const result = handleSpawnAgentInline({
452445
...handleSpawnAgentsBaseParams,
453446
agentState: sessionState.mainAgentState,
454447
agentTemplate: parentAgent,
455448
localAgentTemplates: { reviewer: childAgent },
456449
toolCall,
457450
})
458451

459-
await expect(result).rejects.toThrow(
452+
expect(result).rejects.toThrow(
460453
'is not allowed to spawn child agent type reviewer',
461454
)
462455
expect(mockLoopAgentSteps).not.toHaveBeenCalled()
@@ -467,15 +460,15 @@ describe('Spawn Agents Permissions', () => {
467460
const sessionState = getInitialSessionState(mockFileContext)
468461
const toolCall = createInlineSpawnToolCall('nonexistent')
469462

470-
const { result } = handleSpawnAgentInline({
463+
const result = handleSpawnAgentInline({
471464
...handleSpawnAgentsBaseParams,
472465
agentState: sessionState.mainAgentState,
473466
agentTemplate: parentAgent,
474467
localAgentTemplates: {}, // Empty - agent not found
475468
toolCall,
476469
})
477470

478-
await expect(result).rejects.toThrow('Agent type nonexistent not found')
471+
expect(result).rejects.toThrow('Agent type nonexistent not found')
479472
expect(mockLoopAgentSteps).not.toHaveBeenCalled()
480473
})
481474

@@ -485,15 +478,15 @@ describe('Spawn Agents Permissions', () => {
485478
const sessionState = getInitialSessionState(mockFileContext)
486479
const toolCall = createInlineSpawnToolCall('codebuff/thinker@1.0.0')
487480

488-
const { result } = handleSpawnAgentInline({
481+
// Should not throw
482+
await handleSpawnAgentInline({
489483
...handleSpawnAgentsBaseParams,
490484
agentState: sessionState.mainAgentState,
491485
agentTemplate: parentAgent,
492486
localAgentTemplates: { 'codebuff/thinker@1.0.0': childAgent },
493487
toolCall,
494488
})
495489

496-
await result // Should not throw
497490
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
498491
})
499492

@@ -503,7 +496,8 @@ describe('Spawn Agents Permissions', () => {
503496
const sessionState = getInitialSessionState(mockFileContext)
504497
const toolCall = createInlineSpawnToolCall('thinker') // Simple name
505498

506-
const { result } = handleSpawnAgentInline({
499+
// Should not throw
500+
await handleSpawnAgentInline({
507501
...handleSpawnAgentsBaseParams,
508502
agentState: sessionState.mainAgentState,
509503
agentTemplate: parentAgent,
@@ -514,7 +508,6 @@ describe('Spawn Agents Permissions', () => {
514508
toolCall,
515509
})
516510

517-
await result // Should not throw
518511
expect(mockLoopAgentSteps).toHaveBeenCalledTimes(1)
519512
})
520513

@@ -524,17 +517,15 @@ describe('Spawn Agents Permissions', () => {
524517
const sessionState = getInitialSessionState(mockFileContext)
525518
const toolCall = createInlineSpawnToolCall('codebuff/thinker@2.0.0')
526519

527-
const { result } = handleSpawnAgentInline({
520+
const result = handleSpawnAgentInline({
528521
...handleSpawnAgentsBaseParams,
529522
agentState: sessionState.mainAgentState,
530523
agentTemplate: parentAgent,
531524
localAgentTemplates: { 'codebuff/thinker@2.0.0': childAgent },
532525
toolCall,
533526
})
534527

535-
await expect(result).rejects.toThrow(
536-
'is not allowed to spawn child agent type',
537-
)
528+
expect(result).rejects.toThrow('is not allowed to spawn child agent type')
538529
expect(mockLoopAgentSteps).not.toHaveBeenCalled()
539530
})
540531
})

‎packages/agent-runtime/src/__tests__/subagent-streaming.test.ts‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ describe('Subagent Streaming', () => {
146146
},
147147
}
148148

149-
const { result } = handleSpawnAgents({
149+
await handleSpawnAgents({
150150
...handleSpawnAgentsBaseParams,
151151
agentState,
152152
agentTemplate: parentTemplate,
@@ -156,8 +156,6 @@ describe('Subagent Streaming', () => {
156156
toolCall,
157157
})
158158

159-
await result
160-
161159
// Verify that subagent streaming messages were sent
162160
expect(mockWriteToClient).toHaveBeenCalledTimes(2)
163161

@@ -197,7 +195,7 @@ describe('Subagent Streaming', () => {
197195
},
198196
}
199197

200-
const { result } = handleSpawnAgents({
198+
await handleSpawnAgents({
201199
...handleSpawnAgentsBaseParams,
202200
agentState,
203201
agentTemplate: parentTemplate,
@@ -206,7 +204,6 @@ describe('Subagent Streaming', () => {
206204
},
207205
toolCall,
208206
})
209-
await result
210207

211208
// Verify the streaming messages have consistent agentId and correct agentType
212209
expect(mockSendSubagentChunk.mock.calls.length).toBeGreaterThanOrEqual(2)

0 commit comments

Comments
 (0)