Skip to content

Commit 7c7315e

Browse files
committed
Put back new behavior of not ending until handleSteps returns endTurn: true.
1 parent 032cd5f commit 7c7315e

File tree

2 files changed

+77
-6
lines changed

2 files changed

+77
-6
lines changed

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

Lines changed: 76 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,7 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
527527
expect(llmCallCount).toBe(1) // LLM called once after STEP
528528
expect(result.agentState).toBeDefined()
529529
})
530+
530531
it('should pass shouldEndTurn: true as stepsComplete when end_turn tool is called', async () => {
531532
// Test that when LLM calls end_turn, shouldEndTurn is correctly passed to runProgrammaticStep
532533

@@ -538,11 +539,17 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
538539
() => ({
539540
runProgrammaticStep: async (params: any) => {
540541
runProgrammaticStepCalls.push(params)
541-
// Return default behavior
542-
return { agentState: params.agentState, endTurn: false }
542+
// First call: return endTurn false to continue
543+
// Second call: return endTurn true to end the loop
544+
const shouldEnd = runProgrammaticStepCalls.length >= 2
545+
return {
546+
agentState: params.agentState,
547+
endTurn: shouldEnd,
548+
stepNumber: params.stepNumber,
549+
}
543550
},
544551
clearAgentGeneratorCache: () => {},
545-
agentIdToStepAll: new Set(),
552+
runIdToStepAll: new Set(),
546553
}),
547554
)
548555

@@ -586,6 +593,72 @@ describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () =>
586593
expect(runProgrammaticStepCalls[1].stepsComplete).toBe(true)
587594
})
588595

596+
it('should continue loop when handleSteps returns endTurn: false even if LLM calls end_turn', async () => {
597+
// Test that handleSteps endTurn: false takes precedence over LLM end_turn tool call
598+
599+
let programmaticStepCount = 0
600+
let llmStepCount = 0
601+
602+
const mockGeneratorFunction = function* () {
603+
// First iteration: return endTurn: false
604+
programmaticStepCount++
605+
yield 'STEP'
606+
607+
// Second iteration: also return endTurn: false
608+
programmaticStepCount++
609+
yield 'STEP'
610+
611+
// Third iteration: finally return endTurn: true to end the loop
612+
programmaticStepCount++
613+
yield { toolName: 'end_turn', input: {} }
614+
} as () => StepGenerator
615+
616+
mockTemplate.handleSteps = mockGeneratorFunction
617+
618+
const localAgentTemplates = {
619+
'test-agent': mockTemplate,
620+
}
621+
622+
// Mock LLM to always call end_turn, but handleSteps should override it
623+
let promptCallCount = 0
624+
agentRuntimeImpl.promptAiSdkStream = async function* () {
625+
promptCallCount++
626+
llmStepCount++
627+
628+
// LLM always tries to end turn
629+
yield {
630+
type: 'text' as const,
631+
text: `LLM response\n\n${getToolCallString('end_turn', {})}`,
632+
}
633+
return `mock-message-id-${promptCallCount}`
634+
}
635+
636+
await runLoopAgentStepsWithContext({
637+
...agentRuntimeImpl,
638+
...agentRuntimeScopedImpl,
639+
userInputId: 'test-user-input',
640+
agentType: 'test-agent',
641+
agentState: mockAgentState,
642+
prompt: 'Test handleSteps endTurn override',
643+
spawnParams: undefined,
644+
fingerprintId: 'test-fingerprint',
645+
fileContext: mockFileContext,
646+
localAgentTemplates,
647+
userId: TEST_USER_ID,
648+
clientSessionId: 'test-session',
649+
onResponseChunk: () => {},
650+
})
651+
652+
// Verify handleSteps ran 3 times (yielded STEP twice, then end_turn)
653+
expect(programmaticStepCount).toBe(3)
654+
655+
// Verify LLM was called 2 times (once per STEP yield)
656+
expect(llmStepCount).toBe(2)
657+
658+
// This confirms that even though LLM called end_turn every time,
659+
// the loop continued because handleSteps kept yielding STEP before finally ending
660+
})
661+
589662
it('should restart loop when agent finishes without setting required output', async () => {
590663
// Test that when an agent has outputSchema but finishes without calling set_output,
591664
// the loop restarts with a system message

backend/src/run-agent-step.ts

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -620,9 +620,7 @@ export const loopAgentSteps = async (
620620
currentAgentState = programmaticAgentState
621621
totalSteps = stepNumber
622622

623-
if (endTurn) {
624-
shouldEndTurn = true
625-
}
623+
shouldEndTurn = endTurn
626624
}
627625

628626
// Check if output is required but missing

0 commit comments

Comments
 (0)