Fix types/tests

Author: jahooma

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

@@ -148,6 +148,8 @@ describe('Cost Aggregation System', () => {
         stepsRemaining: 10,
         creditsUsed: 50, // Parent starts with some cost
         directCreditsUsed: 50,
+        systemPrompt: 'Test system prompt',
+        toolDefinitions: {},
       }
 
       // Mock executeAgent to return results with different credit costs

packages/agent-runtime/src/__tests__/loop-agent-steps.test.ts

@@ -141,6 +141,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
       ancestorRunIds: [],
       onResponseChunk: () => {},
       signal: new AbortController().signal,
+      tools: {},
     }
   })
 
@@ -418,32 +419,24 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
   })
 
   it('should pass shouldEndTurn: true as stepsComplete when end_turn tool is called', async () => {
-    // Test that when LLM calls end_turn, shouldEndTurn is correctly passed to runProgrammaticStep
-
-    let runProgrammaticStepCalls: any[] = []
-
-    // Mock runProgrammaticStep module to capture calls and verify stepsComplete parameter
-    const mockedRunProgrammaticStep = await mockModule(
-      '@codebuff/agent-runtime/run-programmatic-step',
-      () => ({
-        runProgrammaticStep: async (params: any) => {
-          runProgrammaticStepCalls.push(params)
-          // First call: return endTurn false to continue
-          // Second call: return endTurn true to end the loop
-          const shouldEnd = runProgrammaticStepCalls.length >= 2
-          return {
-            agentState: params.agentState,
-            endTurn: shouldEnd,
-            stepNumber: params.stepNumber,
-          }
-        },
-        clearAgentGeneratorCache: () => {},
-        runIdToStepAll: new Set(),
-      }),
-    )
+    // Test that when LLM calls end_turn, shouldEndTurn (stepsComplete) is correctly passed
+    // to the handleSteps generator via the step result.
+    //
+    // Flow:
+    // 1. Generator yields 'STEP', runProgrammaticStep returns
+    // 2. loopAgentSteps calls runAgentStep (LLM), which calls end_turn -> shouldEndTurn = true
+    // 3. loopAgentSteps calls runProgrammaticStep again with stepsComplete: true
+    // 4. Generator resumes from yield 'STEP' and receives { stepsComplete: true }
+
+    let stepsCompleteValues: boolean[] = []
 
     const mockGeneratorFunction = function* () {
-      yield 'STEP' // Hand control to LLM
+      // First STEP - after LLM runs and calls end_turn, we receive stepsComplete: true
+      const result1 = yield 'STEP'
+      stepsCompleteValues.push(result1.stepsComplete)
+
+      // Since stepsComplete was true, we should end gracefully
+      yield { toolName: 'end_turn', input: {} }
     } as () => StepGenerator
 
     mockTemplate.handleSteps = mockGeneratorFunction
@@ -458,18 +451,11 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
       localAgentTemplates,
     })
 
-    mockedRunProgrammaticStep.clear()
-
-    // Verify that runProgrammaticStep was called twice:
-    // 1. First with stepsComplete: false (initial call)
-    // 2. Second with stepsComplete: true (after LLM called end_turn)
-    expect(runProgrammaticStepCalls).toHaveLength(2)
-
-    // First call should have stepsComplete: false
-    expect(runProgrammaticStepCalls[0].stepsComplete).toBe(false)
-
-    // Second call should have stepsComplete: true (after end_turn tool was called)
-    expect(runProgrammaticStepCalls[1].stepsComplete).toBe(true)
+    // Verify that stepsComplete was passed correctly:
+    // After yielding STEP and LLM running (which calls end_turn), 
+    // the generator receives stepsComplete: true
+    expect(stepsCompleteValues).toHaveLength(1)
+    expect(stepsCompleteValues[0]).toBe(true)
   })
 
   it('should continue loop when handleSteps returns endTurn: false even if LLM calls end_turn', async () => {

packages/agent-runtime/src/__tests__/main-prompt.test.ts

@@ -98,6 +98,7 @@ describe('mainPrompt', () => {
       onResponseChunk: () => {},
       localAgentTemplates: mockLocalAgentTemplates,
       signal: new AbortController().signal,
+      tools: {},
     }
 
     // Mock analytics and tracing

packages/agent-runtime/src/__tests__/malformed-tool-call.test.ts

@@ -84,6 +84,7 @@ describe('malformed tool call error handling', () => {
       fullResponse: '',
       prompt: '',
       signal: new AbortController().signal,
+      tools: {},
     }
 
     // Mock analytics and tracing

packages/agent-runtime/src/__tests__/n-parameter.test.ts

@@ -259,6 +259,7 @@ describe('n parameter and GENERATE_N functionality', () => {
           stepNumber: 1,
           logger,
           signal: new AbortController().signal,
+          tools: {},
         })
 
         expect(result.generateN).toBe(nValue)
@@ -298,6 +299,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       })
 
       expect(result.generateN).toBeUndefined()
@@ -362,6 +364,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       }
 
       // First call: programmatic step yields GENERATE_N
@@ -444,6 +447,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       }
 
       // First call: execute read_files and yield GENERATE_N
@@ -539,6 +543,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       }
 
       // First GENERATE_N
@@ -603,6 +608,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       })
 
       expect(result.generateN).toBe(1)
@@ -641,6 +647,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       }
 
       await runProgrammaticStep(mockParams)
@@ -687,6 +694,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       }
 
       await runProgrammaticStep(mockParams)
@@ -730,6 +738,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       }
 
       const result1 = await runProgrammaticStep(mockParams)
@@ -792,6 +801,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       }
 
       // First call yields GENERATE_N
@@ -839,6 +849,7 @@ describe('n parameter and GENERATE_N functionality', () => {
         stepNumber: 1,
         logger,
         signal: new AbortController().signal,
+        tools: {},
       })
 
       // Should still set generateN even though endTurn will be true

packages/agent-runtime/src/__tests__/prompt-caching-subagents.test.ts

@@ -136,6 +136,7 @@ describe('Prompt Caching for Subagents with inheritParentSystemPrompt', () => {
       ancestorRunIds: [],
       onResponseChunk: () => {},
       signal: new AbortController().signal,
+      tools: {},
     }
   })
 

packages/agent-runtime/src/__tests__/run-programmatic-step.test.ts

@@ -135,6 +135,7 @@ describe('runProgrammaticStep', () => {
       system: 'Test system prompt',
       stepsComplete: false,
       stepNumber: 1,
+      tools: {},
 
       logger,
       signal: new AbortController().signal,

packages/agent-runtime/src/__tests__/sandbox-generator.test.ts

@@ -89,6 +89,7 @@ describe('QuickJS Sandbox Generator', () => {
       stepsComplete: false,
       stepNumber: 1,
       signal: new AbortController().signal,
+      tools: {},
     }
   })