savant-code 0.0.16 → 0.0.19

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

@@ -0,0 +1,1512 @@
+import * as analytics from '@savant-code/common/analytics'
+import { TEST_USER_ID } from '@savant-code/common/old-constants'
+import {
+  createTestAgentRuntimeParams,
+  emptyMcpServers,
+} from '@savant-code/common/testing/fixtures/agent-runtime'
+import { clearMockedModules } from '@savant-code/common/testing/mock-modules'
+import {
+  createMockDbOperations,
+  setupDbSpies,
+} from '@savant-code/common/testing/mocks/database'
+import { getInitialSessionState } from '@savant-code/common/types/session-state'
+import { AbortError, promptSuccess } from '@savant-code/common/util/error'
+import {
+  assistantMessage,
+  userMessage,
+} from '@savant-code/common/util/messages'
+import { APICallError, RetryError } from 'ai'
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  mock,
+  spyOn,
+} from 'bun:test'
+import { z } from 'zod/v4'
+
+import { createToolCallChunk, mockFileContext } from './test-utils'
+import { loopAgentSteps } from '../run-agent-step'
+import { clearAgentGeneratorCache } from '../run-programmatic-step'
+import { clearThinkerConvergenceStateForTests } from '../tools/thinker-convergence-gate'
+import { clearAllThoughtSessionsForTests } from '../tools/thought-session-store'
+
+import type { AgentTemplate } from '../templates/types'
+import type { DbSpies } from '@savant-code/common/testing/mocks/database'
+import type { StepGenerator } from '@savant-code/common/types/agent-template'
+import type { AgentState } from '@savant-code/common/types/session-state'
+
+describe('loopAgentSteps - runAgentStep vs runProgrammaticStep behavior', () => {
+  let mockTemplate: AgentTemplate
+  let mockAgentState: AgentState
+  let llmCallCount: number
+  let agentRuntimeImpl: Omit<
+    ReturnType<typeof createTestAgentRuntimeParams>,
+    'agentTemplate' | 'localAgentTemplates'
+  > & {
+    promptAiSdkStream?: ReturnType<typeof mock>
+  }
+  let loopAgentStepsBaseParams: Parameters<typeof loopAgentSteps>[0]
+  let dbSpies: DbSpies
+
+  beforeAll(async () => {
+    // Set up mocks.
+  })
+
+  beforeEach(() => {
+    const {
+      agentTemplate: _,
+      localAgentTemplates: __,
+      ...baseRuntimeParams
+    } = createTestAgentRuntimeParams()
+
+    agentRuntimeImpl = {
+      ...baseRuntimeParams,
+    }
+
+    llmCallCount = 0
+
+    // Setup spies for database operations using typed helper
+    dbSpies = setupDbSpies(createMockDbOperations())
+
+    agentRuntimeImpl.promptAiSdkStream = mock(async function* ({}) {
+      llmCallCount++
+      yield { type: 'text' as const, text: 'LLM response\n\n' }
+      yield createToolCallChunk('end_turn', {})
+      return promptSuccess('mock-message-id')
+    })
+
+    // Mock analytics
+    spyOn(analytics, 'trackEvent').mockImplementation(() => {})
+
+    // Mock crypto.randomUUID
+    spyOn(crypto, 'randomUUID').mockImplementation(
+      () => 'mock-uuid-0000-0000-0000-000000000000' as const,
+    )
+
+    // Create mock template with programmatic agent
+    mockTemplate = {
+      id: 'test-agent',
+      displayName: 'Test Agent',
+      spawnerPrompt: 'Testing',
+      model: 'claude-3-5-sonnet-20241022',
+      inputSchema: {},
+      outputMode: 'structured_output',
+      includeMessageHistory: true,
+      inheritParentSystemPrompt: false,
+      mcpServers: emptyMcpServers,
+      toolNames: ['read_files', 'write_file', 'end_turn'],
+      spawnableAgents: [],
+      systemPrompt: 'Test system prompt',
+      instructionsPrompt: 'Test user prompt',
+      stepPrompt: 'Test agent step prompt',
+      handleSteps: undefined, // Will be set in individual tests
+    } satisfies AgentTemplate as AgentTemplate
+
+    // Create mock agent state
+    const sessionState = getInitialSessionState(mockFileContext)
+    mockAgentState = {
+      ...sessionState.mainAgentState,
+      agentId: 'test-agent-id',
+      messageHistory: [
+        userMessage('Initial message'),
+        assistantMessage('Initial response'),
+      ],
+      output: undefined,
+      stepsRemaining: 10, // Ensure we don't hit the limit
+    }
+
+    loopAgentStepsBaseParams = {
+      ...agentRuntimeImpl,
+      agentType: 'test-agent',
+      localAgentTemplates: { 'test-agent': mockTemplate },
+      repoId: undefined,
+      repoUrl: undefined,
+      userInputId: 'test-user-input',
+      agentState: mockAgentState,
+      prompt: 'Test prompt',
+      spawnParams: undefined,
+      fingerprintId: 'test-fingerprint',
+      fileContext: mockFileContext,
+      userId: TEST_USER_ID,
+      clientSessionId: 'test-session',
+      ancestorRunIds: [],
+      onResponseChunk: () => {},
+      signal: new AbortController().signal,
+    }
+  })
+
+  afterEach(() => {
+    clearAgentGeneratorCache(agentRuntimeImpl)
+    dbSpies.restore()
+    mock.restore()
+    const {
+      agentTemplate: _,
+      localAgentTemplates: __,
+      ...baseRuntimeParams
+    } = createTestAgentRuntimeParams()
+    agentRuntimeImpl = {
+      ...baseRuntimeParams,
+    }
+  })
+
+  afterAll(() => {
+    clearMockedModules()
+  })
+
+  it('should verify correct STEP behavior - LLM called once after STEP', async () => {
+    // This test verifies that when a programmatic agent yields STEP,
+    // the LLM should be called once in the next iteration
+
+    let stepCount = 0
+    const mockGeneratorFunction = function* () {
+      stepCount++
+      // Execute a tool, then STEP
+      yield { toolName: 'read_files', input: { paths: ['file1.txt'] } }
+      yield 'STEP' // Should pause here and let LLM run
+      // Continue after LLM runs (this won't be reached in this test since LLM ends turn)
+      yield {
+        toolName: 'write_file',
+        input: { path: 'output.txt', content: 'test' },
+      }
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const _result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    console.log(`LLM calls made: ${llmCallCount}`)
+    console.log(`Step count: ${stepCount}`)
+
+    // CORRECT BEHAVIOR: After STEP, LLM should be called once
+    // The programmatic agent yields STEP, then LLM runs once and ends turn
+    expect(llmCallCount).toBe(1) // LLM called once after STEP
+
+    // The programmatic agent should have been called once (yielded STEP)
+    expect(stepCount).toBe(1)
+  })
+
+  it('should demonstrate correct behavior when programmatic agent completes without STEP', async () => {
+    // This test shows that when a programmatic agent doesn't yield STEP,
+    // it should complete without calling the LLM at all (since it ends with end_turn)
+
+    const mockGeneratorFunction = function* () {
+      yield { toolName: 'read_files', input: { paths: ['file1.txt'] } }
+      yield {
+        toolName: 'write_file',
+        input: { path: 'output.txt', content: 'test' },
+      }
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Should NOT call LLM since the programmatic agent ended with end_turn
+    expect(llmCallCount).toBe(0)
+    // The result should have agentState
+    expect(result.agentState).toBeDefined()
+  })
+
+  it('should run programmatic step first, then LLM step, then continue', async () => {
+    // This test verifies the correct execution order in loopAgentSteps:
+    // 1. Programmatic step runs first and yields STEP
+    // 2. LLM step runs once
+    // 3. Loop continues but generator is complete after first STEP
+
+    let stepCount = 0
+    const mockGeneratorFunction = function* () {
+      stepCount++
+      // First execution: do some work, then STEP
+      yield { toolName: 'read_files', input: { paths: ['file1.txt'] } }
+      yield 'STEP' // Hand control to LLM
+      // After LLM runs, continue (this happens in the same generator instance)
+      yield {
+        toolName: 'write_file',
+        input: { path: 'output.txt', content: 'updated by LLM' },
+      }
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Verify execution order:
+    // 1. Programmatic step function was called once (creates generator)
+    // 2. LLM was called once after STEP
+    // 3. Generator continued after LLM step
+    expect(stepCount).toBe(1) // Generator function called once
+    expect(llmCallCount).toBe(1) // LLM called once after first STEP
+    expect(result.agentState).toBeDefined()
+  })
+
+  it('should handle programmatic agent that yields STEP_ALL', async () => {
+    // Test STEP_ALL behavior - should run LLM then continue with programmatic step
+
+    let stepCount = 0
+    const mockGeneratorFunction = function* () {
+      stepCount++
+      yield { toolName: 'read_files', input: { paths: ['file1.txt'] } }
+      yield 'STEP_ALL' // Hand all remaining control to LLM
+      // Should continue after LLM completes all its steps
+      yield {
+        toolName: 'write_file',
+        input: { path: 'final.txt', content: 'done' },
+      }
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    expect(stepCount).toBe(1) // Generator function called once
+    expect(llmCallCount).toBe(1) // LLM should be called once
+    expect(result.agentState).toBeDefined()
+  })
+
+  it('should not call LLM when programmatic agent returns without STEP', async () => {
+    // Test that programmatic agents that don't yield STEP don't trigger LLM
+
+    const mockGeneratorFunction = function* () {
+      yield { toolName: 'read_files', input: { paths: ['test.txt'] } }
+      yield {
+        toolName: 'write_file',
+        input: { path: 'result.txt', content: 'processed' },
+      }
+      // No STEP - agent completes without LLM involvement
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    expect(llmCallCount).toBe(0) // No LLM calls should be made
+    expect(result.agentState).toBeDefined()
+  })
+
+  it('should handle LLM-only agent (no handleSteps)', async () => {
+    // Test traditional LLM-based agents that don't have handleSteps
+
+    const llmOnlyTemplate = {
+      ...mockTemplate,
+      handleSteps: undefined, // No programmatic step function
+    }
+
+    const localAgentTemplates = {
+      'test-agent': llmOnlyTemplate,
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    expect(llmCallCount).toBe(1) // LLM should be called once
+    expect(result.agentState).toBeDefined()
+  })
+
+  it('should pass the full message history to the traceWriter when provided', async () => {
+    const recordedSteps: Array<{ agentId: string; messages: unknown[] }> = []
+    const traceWriter = {
+      recordStep: (params: { agentId: string; messages: unknown[] }) => {
+        recordedSteps.push(params)
+      },
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      traceWriter,
+      agentType: 'test-agent',
+      localAgentTemplates: {
+        'test-agent': { ...mockTemplate, handleSteps: undefined },
+      },
+    })
+
+    expect(result.agentState).toBeDefined()
+    // Called at least at the start and end of the step
+    expect(recordedSteps.length).toBeGreaterThanOrEqual(2)
+    expect(recordedSteps[0]!.agentId).toBe('test-agent-id')
+    // End-of-step call sees the assistant response appended to the history
+    const lastMessages = recordedSteps[recordedSteps.length - 1]!.messages
+    expect(lastMessages.length).toBeGreaterThan(
+      recordedSteps[0]!.messages.length,
+    )
+  })
+
+  it('should handle programmatic agent error and still call LLM', async () => {
+    // Test error handling in programmatic step - should still allow LLM to run
+
+    const mockGeneratorFunction = function* () {
+      yield { toolName: 'read_files', input: { paths: ['file1.txt'] } }
+      throw new Error('Programmatic step failed')
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // After programmatic step error, should end turn and not call LLM
+    expect(llmCallCount).toBe(0)
+    expect(result.agentState).toBeDefined()
+    expect(result.agentState.output?.error).toContain(
+      'Error executing handleSteps for agent test-agent',
+    )
+  })
+
+  it('should handle mixed execution with multiple STEP yields', async () => {
+    // Test complex scenario with multiple STEP yields and LLM interactions
+    // Note: In current implementation, LLM typically ends turn after running,
+    // so this tests the first STEP interaction
+
+    let stepCount = 0
+    const mockGeneratorFunction = function* () {
+      stepCount++
+      yield { toolName: 'read_files', input: { paths: ['input.txt'] } }
+      yield 'STEP' // First LLM interaction
+      yield {
+        toolName: 'write_file',
+        input: { path: 'temp.txt', content: 'intermediate' },
+      }
+      yield {
+        toolName: 'write_file',
+        input: { path: 'final.txt', content: 'complete' },
+      }
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    expect(stepCount).toBe(1) // Generator function called once
+    expect(llmCallCount).toBe(1) // LLM called once after STEP
+    expect(result.agentState).toBeDefined()
+  })
+
+  it('should pass shouldEndTurn: true as stepsComplete when end_turn tool is called', async () => {
+    // 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* () {
+      // 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
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // 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 () => {
+    // Test that handleSteps endTurn: false takes precedence over LLM end_turn tool call
+
+    let programmaticStepCount = 0
+    let llmStepCount = 0
+
+    const mockGeneratorFunction = function* () {
+      // First iteration: return endTurn: false
+      programmaticStepCount++
+      yield 'STEP'
+
+      // Second iteration: also return endTurn: false
+      programmaticStepCount++
+      yield 'STEP'
+
+      // Third iteration: finally return endTurn: true to end the loop
+      programmaticStepCount++
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    // Mock LLM to always call end_turn, but handleSteps should override it
+    let promptCallCount = 0
+    loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+      promptCallCount++
+      llmStepCount++
+
+      // LLM always tries to end turn
+      yield { type: 'text' as const, text: 'LLM response\n\n' }
+      yield createToolCallChunk('end_turn', {})
+      return promptSuccess(`mock-message-id-${promptCallCount}`)
+    }
+
+    await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Verify handleSteps ran 3 times (yielded STEP twice, then end_turn)
+    expect(programmaticStepCount).toBe(3)
+
+    // Verify LLM was called 2 times (once per STEP yield)
+    expect(llmStepCount).toBe(2)
+
+    // This confirms that even though LLM called end_turn every time,
+    // the loop continued because handleSteps kept yielding STEP before finally ending
+  })
+
+  it('should restart loop when agent finishes without setting required output', async () => {
+    // Test that when an agent has outputSchema but finishes without calling set_output,
+    // the loop restarts with a system message
+
+    const outputSchema = z.object({
+      result: z.string(),
+      status: z.string(),
+    })
+
+    const templateWithOutputSchema = {
+      ...mockTemplate,
+      outputSchema,
+      toolNames: ['set_output', 'end_turn'], // Add set_output to available tools
+      handleSteps: undefined, // LLM-only agent
+    }
+
+    const localAgentTemplates = {
+      'test-agent': templateWithOutputSchema,
+    }
+
+    let llmCallNumber = 0
+    let capturedAgentState: AgentState | null = null
+
+    loopAgentStepsBaseParams.promptAiSdkStream = async function* ({}) {
+      llmCallNumber++
+      if (llmCallNumber === 1) {
+        // First call: agent tries to end turn without setting output
+        yield {
+          type: 'text' as const,
+          text: 'First response without output\n\n',
+        }
+        yield createToolCallChunk('end_turn', {})
+      } else if (llmCallNumber === 2) {
+        // Second call: agent sets output after being reminded
+        // Manually set the output to simulate the set_output tool execution
+        if (capturedAgentState) {
+          capturedAgentState.output = {
+            result: 'test result',
+            status: 'success',
+          }
+        }
+        yield { type: 'text' as const, text: 'Setting output now\n\n' }
+        yield createToolCallChunk('set_output', {
+          result: 'test result',
+          status: 'success',
+        })
+        yield { type: 'text' as const, text: '\n\n' }
+        yield createToolCallChunk('end_turn', {})
+      } else {
+        // Safety: if called more than twice, just end
+        yield { type: 'text' as const, text: 'Ending\n\n' }
+        yield createToolCallChunk('end_turn', {})
+      }
+      return promptSuccess('mock-message-id')
+    }
+
+    mockAgentState.output = undefined
+    capturedAgentState = mockAgentState
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Should call LLM twice: once to try ending without output, once after reminder
+    expect(llmCallNumber).toBe(2)
+
+    // Should have output set after the second attempt
+    expect(result.agentState.output).toEqual({
+      result: 'test result',
+      status: 'success',
+    })
+
+    // Check that a system message was added to message history
+    const systemMessages = result.agentState.messageHistory.filter(
+      (msg) =>
+        msg.role === 'user' &&
+        msg.content[0].type === 'text' &&
+        msg.content[0].text.includes('set_output'),
+    )
+    expect(systemMessages.length).toBeGreaterThan(0)
+  })
+
+  it('should not restart loop if output is set correctly', async () => {
+    // Test that when an agent has outputSchema and sets output correctly,
+    // the loop ends normally without restarting
+
+    const outputSchema = z.object({
+      result: z.string(),
+    })
+
+    const templateWithOutputSchema = {
+      ...mockTemplate,
+      outputSchema,
+      toolNames: ['set_output', 'end_turn'],
+      handleSteps: undefined,
+    }
+
+    const localAgentTemplates = {
+      'test-agent': templateWithOutputSchema,
+    }
+
+    let llmCallNumber = 0
+    let capturedAgentState: AgentState | null = null
+
+    loopAgentStepsBaseParams.promptAiSdkStream = async function* ({}) {
+      llmCallNumber++
+      // Agent sets output correctly on first call
+      if (capturedAgentState) {
+        capturedAgentState.output = { result: 'success' }
+      }
+      yield { type: 'text' as const, text: 'Setting output\n\n' }
+      yield createToolCallChunk('set_output', { result: 'success' })
+      yield { type: 'text' as const, text: '\n\n' }
+      yield createToolCallChunk('end_turn', {})
+      return promptSuccess('mock-message-id')
+    }
+
+    mockAgentState.output = undefined
+    capturedAgentState = mockAgentState
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Should only call LLM once since output was set correctly
+    expect(llmCallNumber).toBe(1)
+
+    // Should have output set
+    expect(result.agentState.output).toEqual({ result: 'success' })
+  })
+
+  it('should pass generateN from programmatic step to runAgentStep as n parameter', async () => {
+    // Test that when programmatic step returns generateN, it's passed to runAgentStep
+
+    let agentStepN: number | undefined
+
+    const mockGeneratorFunction = function* () {
+      // Yield GENERATE_N to trigger n parameter
+      yield { type: 'GENERATE_N', n: 5 }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    // Mock promptAiSdk to capture the n parameter
+    loopAgentStepsBaseParams.promptAiSdk = async (params: any) => {
+      agentStepN = params.n
+      return promptSuccess(
+        JSON.stringify([
+          'Response 1',
+          'Response 2',
+          'Response 3',
+          'Response 4',
+          'Response 5',
+        ]),
+      )
+    }
+
+    await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Verify generateN was passed to runAgentStep as n
+    expect(agentStepN).toBe(5)
+  })
+
+  it('should pass nResponses from runAgentStep back to programmatic step', async () => {
+    // Test that nResponses returned by runAgentStep are passed to next programmatic step
+
+    let receivedNResponses: string[] | undefined
+
+    const mockGeneratorFunction = function* () {
+      const { nResponses } = yield { type: 'GENERATE_N', n: 3 }
+      receivedNResponses = nResponses
+      const _step = yield {
+        toolName: 'read_files',
+        input: { paths: ['test.txt'] },
+      }
+      yield { toolName: 'end_turn', input: {} }
+    } as () => StepGenerator
+
+    mockTemplate.handleSteps = mockGeneratorFunction
+
+    const localAgentTemplates = {
+      'test-agent': mockTemplate,
+    }
+
+    const expectedResponses = [
+      'Implementation A',
+      'Implementation B',
+      'Implementation C',
+    ]
+    loopAgentStepsBaseParams.promptAiSdk = async () => {
+      return promptSuccess(JSON.stringify(expectedResponses))
+    }
+
+    await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    expect(receivedNResponses).toEqual(expectedResponses)
+  })
+
+  it('should allow agents without outputSchema to end normally', async () => {
+    // Test that agents without outputSchema can end without setting output
+
+    const templateWithoutOutputSchema = {
+      ...mockTemplate,
+      outputSchema: undefined,
+      handleSteps: undefined,
+    }
+
+    const localAgentTemplates = {
+      'test-agent': templateWithoutOutputSchema,
+    }
+
+    let llmCallNumber = 0
+    loopAgentStepsBaseParams.promptAiSdkStream = async function* ({}) {
+      llmCallNumber++
+      yield { type: 'text' as const, text: 'Response without output\n\n' }
+      yield createToolCallChunk('end_turn', {})
+      return promptSuccess('mock-message-id')
+    }
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Should only call LLM once and end normally
+    expect(llmCallNumber).toBe(1)
+
+    // Output should be undefined since no outputSchema required
+    expect(result.agentState.output).toBeUndefined()
+  })
+
+  it('should continue loop if agent does not end turn (has more work)', async () => {
+    // Test that validation only triggers when shouldEndTurn is true
+
+    const outputSchema = z.object({
+      result: z.string(),
+    })
+
+    const templateWithOutputSchema = {
+      ...mockTemplate,
+      outputSchema,
+      toolNames: ['read_files', 'set_output', 'end_turn'],
+      handleSteps: undefined,
+    }
+
+    const localAgentTemplates = {
+      'test-agent': templateWithOutputSchema,
+    }
+
+    let llmCallNumber = 0
+    let capturedAgentState: AgentState | null = null
+
+    loopAgentStepsBaseParams.promptAiSdkStream = async function* ({}) {
+      llmCallNumber++
+      if (llmCallNumber === 1) {
+        // First call: agent does some work but doesn't end turn
+        yield { type: 'text' as const, text: 'Doing work\n\n' }
+        yield createToolCallChunk('read_files', { paths: ['test.txt'] })
+      } else {
+        // Second call: agent sets output and ends
+        if (capturedAgentState) {
+          capturedAgentState.output = { result: 'done' }
+        }
+        yield { type: 'text' as const, text: 'Finishing\n\n' }
+        yield createToolCallChunk('set_output', { result: 'done' })
+        yield { type: 'text' as const, text: '\n\n' }
+        yield createToolCallChunk('end_turn', {})
+      }
+      return promptSuccess('mock-message-id')
+    }
+
+    mockAgentState.output = undefined
+    capturedAgentState = mockAgentState
+
+    const result = await loopAgentSteps({
+      ...loopAgentStepsBaseParams,
+      agentType: 'test-agent',
+      localAgentTemplates,
+    })
+
+    // Should call LLM twice: once for work, once to set output and end
+    expect(llmCallNumber).toBe(2)
+
+    // Should have output set
+    expect(result.agentState.output).toEqual({ result: 'done' })
+  })
+
+  describe('abort handling', () => {
+    it('should handle AbortError and finish with cancelled status', async () => {
+      // Test that when an AbortError is thrown (e.g., from a tool handler),
+      // loopAgentSteps catches it, finishes with 'cancelled' status, and returns
+      // an error output indicating the run was cancelled.
+
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      // Track finishAgentRun calls
+      let finishAgentRunStatus: string | undefined
+      const mockFinishAgentRun = mock(async (params: { status: string }) => {
+        finishAgentRunStatus = params.status
+      })
+
+      // Mock promptAiSdkStream to throw an AbortError (simulating user cancellation mid-stream)
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        // Yield some content first
+        yield { type: 'text' as const, text: 'Starting work...\n' }
+        // Then throw AbortError to simulate user cancellation
+        throw new AbortError('User pressed Ctrl+C')
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+        finishAgentRun: mockFinishAgentRun,
+      })
+
+      // Verify the output indicates cancellation
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        expect(result.output.message).toBe('Run cancelled by user')
+      }
+
+      // Verify finishAgentRun was called with 'cancelled' status
+      expect(mockFinishAgentRun).toHaveBeenCalled()
+      expect(finishAgentRunStatus).toBe('cancelled')
+    })
+
+    it('should distinguish AbortError from other errors', async () => {
+      // Test that non-abort errors are NOT treated as cancellations
+
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      // Track finishAgentRun calls
+      let finishAgentRunStatus: string | undefined
+      const mockFinishAgentRun = mock(async (params: { status: string }) => {
+        finishAgentRunStatus = params.status
+      })
+
+      // Mock promptAiSdkStream to throw a regular error (not AbortError)
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        yield { type: 'text' as const, text: 'Starting...\n' }
+        throw new Error('Network connection failed')
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+        finishAgentRun: mockFinishAgentRun,
+      })
+
+      // Verify the output indicates an error (not cancellation)
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        expect(result.output.message).toContain('Network connection failed')
+        expect(result.output.message).not.toBe('Run cancelled by user')
+      }
+
+      // Verify finishAgentRun was called with 'failed' status (not 'cancelled')
+      expect(mockFinishAgentRun).toHaveBeenCalled()
+      expect(finishAgentRunStatus).toBe('failed')
+    })
+
+    it('should handle signal.aborted before loop starts', async () => {
+      // Test that if signal is already aborted when loopAgentSteps is called,
+      // it returns immediately with a cancelled message
+
+      const abortController = new AbortController()
+      abortController.abort() // Abort immediately
+
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+        signal: abortController.signal,
+      })
+
+      // Verify the output indicates cancellation
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        expect(result.output.message).toBe('Run cancelled by user')
+      }
+
+      // LLM should not have been called since we aborted before starting
+      expect(llmCallCount).toBe(0)
+    })
+  })
+
+  describe('API error handling', () => {
+    it('should propagate error code and server message from 403 APICallError responseBody', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      // Mock promptAiSdkStream to throw an APICallError with a 403 status
+      // and a responseBody containing the server's structured error
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        throw new APICallError({
+          statusCode: 403,
+          message: 'Forbidden',
+          url: 'https://api.savant-code.com/v1/chat/completions',
+          requestBodyValues: {},
+          responseBody: JSON.stringify({
+            error: 'free_mode_unavailable',
+            message: 'Free mode is not available in your country.',
+            countryCode: 'US',
+            countryBlockReason: 'anonymous_network',
+            ipPrivacySignals: ['vpn', 'hosting'],
+          }),
+          isRetryable: false,
+        })
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+      })
+
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        // Should use the server's message, NOT the generic "Forbidden"
+        expect(result.output.message).toBe(
+          'Free mode is not available in your country.',
+        )
+        // Should NOT have the 'Agent run error: ' prefix since message came from responseBody
+        expect(result.output.message).not.toContain('Agent run error:')
+        // Should propagate the error code so the CLI can match on it
+        expect(result.output.error).toBe('free_mode_unavailable')
+        // Should propagate the status code
+        expect(result.output.statusCode).toBe(403)
+        expect(result.output.countryCode).toBe('US')
+        expect(result.output.countryBlockReason).toBe('anonymous_network')
+        expect(result.output.ipPrivacySignals).toEqual(['vpn', 'hosting'])
+      }
+    })
+
+    it('should prefix with "Agent run error:" when responseBody has no parseable message', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      // APICallError with no responseBody
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        throw new APICallError({
+          statusCode: 500,
+          message: 'Internal Server Error',
+          url: 'https://api.savant-code.com/v1/chat/completions',
+          requestBodyValues: {},
+          responseBody: undefined,
+          isRetryable: true,
+        })
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+      })
+
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        // Should have the prefix since there's no server message
+        expect(result.output.message).toContain('Agent run error:')
+        expect(result.output.message).toContain('Internal Server Error')
+        // No error code since responseBody wasn't parseable
+        expect(result.output.error).toBeUndefined()
+      }
+    })
+
+    it('should unwrap retry errors to propagate underlying 409 gate errors', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      const apiError = new APICallError({
+        statusCode: 409,
+        message: 'Conflict',
+        url: 'https://api.savant-code.com/v1/chat/completions',
+        requestBodyValues: {},
+        responseBody: JSON.stringify({
+          error: 'session_superseded',
+          message:
+            'Another instance of savant-free has taken over this session. Only one instance per account is allowed.',
+        }),
+        isRetryable: true,
+      })
+
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        throw new RetryError({
+          message: 'Failed after 4 attempts. Last error: Conflict',
+          reason: 'maxRetriesExceeded',
+          errors: [apiError],
+        })
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+      })
+
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        expect(result.output.message).toBe(
+          'Another instance of savant-free has taken over this session. Only one instance per account is allowed.',
+        )
+        expect(result.output.message).not.toContain('Agent run error:')
+        expect(result.output.error).toBe('session_superseded')
+        expect(result.output.statusCode).toBe(409)
+      }
+    })
+
+    it('should explain fetch idle timeouts instead of showing the raw runtime message', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      // Bun aborts a fetch after 5 minutes without receiving bytes, throwing a
+      // DOMException named TimeoutError with this exact message.
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        const timeoutError = new Error('The operation timed out.')
+        timeoutError.name = 'TimeoutError'
+        throw timeoutError
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+      })
+
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        expect(result.output.message).toContain(
+          'no data was received from the server for 5 minutes',
+        )
+        expect(result.output.message).not.toContain('Agent run error:')
+        expect(result.output.message).not.toBe('The operation timed out.')
+      }
+    })
+
+    it('should explain dropped socket connections instead of showing the raw runtime message', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+
+      const localAgentTemplates = {
+        'test-agent': llmOnlyTemplate,
+      }
+
+      // Bun's fetch throws a plain Error with this message (and code
+      // ECONNRESET/ConnectionClosed) when the TCP connection is dropped.
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        const socketError = new Error(
+          'The socket connection was closed unexpectedly. For more information, pass `verbose: true` in the second argument to fetch()',
+        ) as Error & { code: string }
+        socketError.code = 'ECONNRESET'
+        throw socketError
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        localAgentTemplates,
+      })
+
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        expect(result.output.message).toContain('Connection interrupted')
+        expect(result.output.message).not.toContain('Agent run error:')
+        expect(result.output.message).not.toContain(
+          'pass `verbose: true` in the second argument to fetch()',
+        )
+      }
+    })
+  })
+
+  describe('native tool-call recovery (FID-2026-0801-010)', () => {
+    it('retries once, then fails visibly without a third model call', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+      let finishStatus: string | undefined
+      let finishError: string | undefined
+
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        llmCallCount++
+        yield {
+          type: 'error' as const,
+          message: 'Incomplete arguments for tool sequentialthinking',
+          errorClass: 'native-incomplete' as const,
+          toolName: 'sequentialthinking',
+        }
+        return promptSuccess(`native-incomplete-${llmCallCount}`)
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentTemplate: llmOnlyTemplate,
+        localAgentTemplates: { 'test-agent': llmOnlyTemplate },
+        finishAgentRun: mock(
+          async (params: { status: string; errorMessage?: string }) => {
+            finishStatus = params.status
+            finishError = params.errorMessage
+          },
+        ),
+      })
+
+      expect(llmCallCount).toBe(2)
+      expect(result.output.type).toBe('error')
+      if (result.output.type === 'error') {
+        expect(result.output.message).toContain(
+          'Native tool-call recovery failed twice consecutively',
+        )
+      }
+      expect(finishStatus).toBe('failed')
+      expect(finishError).toContain(
+        'Native tool-call recovery failed twice consecutively',
+      )
+
+      const history = result.agentState.messageHistory
+      expect(
+        history.some(
+          (message) =>
+            message.role === 'assistant' &&
+            message.content.some((part) => part.type === 'tool-call'),
+        ),
+      ).toBe(false)
+      expect(history.some((message) => message.role === 'tool')).toBe(false)
+      expect(
+        history.some(
+          (message) =>
+            message.role === 'user' &&
+            message.tags?.includes('TOOL_CALL_ERROR'),
+        ),
+      ).toBe(true)
+    })
+
+    it('recovers on the next step with one valid sequentialthinking result', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        id: 'thinker-test-agent',
+        handleSteps: undefined,
+        toolNames: ['sequentialthinking', 'end_turn'],
+      } satisfies AgentTemplate
+
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        llmCallCount++
+        if (llmCallCount === 1) {
+          yield {
+            type: 'error' as const,
+            message: 'Incomplete arguments for tool sequentialthinking',
+            errorClass: 'native-incomplete' as const,
+            toolName: 'sequentialthinking',
+          }
+        } else {
+          yield createToolCallChunk('sequentialthinking', {
+            thought: 'The continuation is executing the complete native call.',
+            thoughtNumber: 1,
+            totalThoughts: 1,
+            nextThoughtNeeded: false,
+          })
+          yield createToolCallChunk('end_turn', {})
+        }
+        return promptSuccess(`recovery-${llmCallCount}`)
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentTemplate: llmOnlyTemplate,
+        localAgentTemplates: { 'test-agent': llmOnlyTemplate },
+      })
+
+      expect(llmCallCount).toBe(2)
+      expect(result.output.type).not.toBe('error')
+
+      const assistantToolCalls = result.agentState.messageHistory.filter(
+        (message) =>
+          message.role === 'assistant' &&
+          message.content.some(
+            (part) =>
+              part.type === 'tool-call' &&
+              part.toolName === 'sequentialthinking',
+          ),
+      )
+      const sequentialThinkingResults = result.agentState.messageHistory.filter(
+        (message) =>
+          message.role === 'tool' && message.toolName === 'sequentialthinking',
+      )
+
+      expect(assistantToolCalls).toHaveLength(1)
+      expect(sequentialThinkingResults).toHaveLength(1)
+      expect(
+        result.agentState.messageHistory.some(
+          (message) =>
+            message.role === 'user' &&
+            message.tags?.includes('TOOL_CALL_ERROR'),
+        ),
+      ).toBe(true)
+    })
+
+    it('resets the native-incomplete streak after an unrelated tool error', async () => {
+      const llmOnlyTemplate = {
+        ...mockTemplate,
+        handleSteps: undefined,
+      }
+      let finishStatus: string | undefined
+
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        llmCallCount++
+        if (llmCallCount === 2) {
+          yield {
+            type: 'error' as const,
+            message: 'An unrelated tool validation error',
+          }
+        } else {
+          yield {
+            type: 'error' as const,
+            message: 'Incomplete arguments for tool sequentialthinking',
+            errorClass: 'native-incomplete' as const,
+            toolName: 'sequentialthinking',
+          }
+        }
+        return promptSuccess(`recovery-${llmCallCount}`)
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentTemplate: llmOnlyTemplate,
+        localAgentTemplates: { 'test-agent': llmOnlyTemplate },
+        finishAgentRun: mock(async (params: { status: string }) => {
+          finishStatus = params.status
+        }),
+      })
+
+      expect(llmCallCount).toBe(4)
+      expect(result.output.type).toBe('error')
+      expect(finishStatus).toBe('failed')
+    })
+  })
+
+  describe('Thinker convergence gate integration (FID-2026-0801-012)', () => {
+    const thinkerTemplate = (): AgentTemplate => ({
+      ...mockTemplate,
+      id: 'thinker-test-agent',
+      outputMode: 'structured_output',
+      outputSchema: z.object({
+        status: z.string(),
+        payload: z.object({ message: z.string() }).nullable(),
+      }),
+      toolNames: ['sequentialthinking', 'end_turn'],
+      handleSteps: undefined,
+    })
+
+    afterEach(() => {
+      clearAllThoughtSessionsForTests()
+      clearThinkerConvergenceStateForTests()
+    })
+
+    it('sets output from the session snapshot and breaks without the set_output restart', async () => {
+      // The Thinker converges with a single nextThoughtNeeded=false thought and
+      // ends its turn. The gate must build the FinalArtifact from the session
+      // snapshot and set agentState.output BEFORE the loop-top
+      // `output === undefined && shouldEndTurn` restart check can inject the
+      // "You must use set_output" message (which would reintroduce the null).
+      const template = thinkerTemplate()
+      let llmCallNumber = 0
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        llmCallNumber++
+        yield createToolCallChunk('sequentialthinking', {
+          thought: 'Conclusion: use the hybrid approach.',
+          thoughtNumber: 1,
+          totalThoughts: 1,
+          nextThoughtNeeded: false,
+        })
+        yield { type: 'text' as const, text: '\n\n' }
+        yield createToolCallChunk('end_turn', {})
+        return promptSuccess('mock-message-id')
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'thinker-test-agent',
+        localAgentTemplates: { 'thinker-test-agent': template },
+      })
+
+      // Exactly one LLM call: the loop must NOT have restarted.
+      expect(llmCallNumber).toBe(1)
+      expect(result.agentState.output).toBeDefined()
+      expect((result.agentState.output as { status?: string }).status).toBe(
+        'success',
+      )
+      const restartMessages = result.agentState.messageHistory.filter(
+        (m) =>
+          m.role === 'user' &&
+          m.content[0].type === 'text' &&
+          m.content[0].text.includes('set_output'),
+      )
+      expect(restartMessages.length).toBe(0)
+    })
+
+    it('retries on non-convergence instead of restarting, then converges', async () => {
+      // First turn: the model thinks (nextThoughtNeeded=true) and ends the turn
+      // unconverged. The gate appends a typed retry message and keeps the loop
+      // going — it must NOT hit the set_output restart path. Second turn
+      // converges and the gate sets the artifact.
+      const template = thinkerTemplate()
+      let llmCallNumber = 0
+      loopAgentStepsBaseParams.promptAiSdkStream = async function* () {
+        llmCallNumber++
+        if (llmCallNumber === 1) {
+          yield createToolCallChunk('sequentialthinking', {
+            thought: 'Partial analysis, not done yet.',
+            thoughtNumber: 1,
+            totalThoughts: 2,
+            nextThoughtNeeded: true,
+          })
+        } else {
+          yield createToolCallChunk('sequentialthinking', {
+            thought: 'Final conclusion: hybrid wins.',
+            thoughtNumber: 2,
+            totalThoughts: 2,
+            nextThoughtNeeded: false,
+          })
+        }
+        yield { type: 'text' as const, text: '\n\n' }
+        yield createToolCallChunk('end_turn', {})
+        return promptSuccess('mock-message-id')
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'thinker-test-agent',
+        localAgentTemplates: { 'thinker-test-agent': template },
+      })
+
+      // Two LLM calls: first unconverged turn + retry, second converged turn.
+      expect(llmCallNumber).toBe(2)
+      expect(result.agentState.output).toBeDefined()
+      expect((result.agentState.output as { status?: string }).status).toBe(
+        'success',
+      )
+      // The typed retry message was appended after the first turn...
+      const retryMessages = result.agentState.messageHistory.filter(
+        (m) =>
+          m.role === 'user' &&
+          m.content[0].type === 'text' &&
+          m.content[0].text.includes('nextThoughtNeeded=false'),
+      )
+      expect(retryMessages.length).toBeGreaterThan(0)
+      // ...and the set_output restart message was never injected.
+      const restartMessages = result.agentState.messageHistory.filter(
+        (m) =>
+          m.role === 'user' &&
+          m.content[0].type === 'text' &&
+          m.content[0].text.includes('set_output'),
+      )
+      expect(restartMessages.length).toBe(0)
+    })
+  })
+
+  describe('steering (drainSteeringMessages)', () => {
+    it('appends a steering message at the step boundary and continues the turn', async () => {
+      // The mock LLM ends the turn after one step. A steering message that arrives
+      // during that step should be appended to history and keep the turn going, so
+      // the agent runs a second step that can see (and act on) the new message.
+      const steerText = 'Also rename the variable to fooBar'
+      let drainCalls = 0
+      const drainSteeringMessages = () => {
+        drainCalls++
+        return drainCalls === 1 ? [steerText] : []
+      }
+
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        drainSteeringMessages,
+      })
+
+      // Step 1 wanted to end the turn, but the steer kept it going → a 2nd step ran.
+      expect(llmCallCount).toBe(2)
+
+      // The steered text landed in history as a user message.
+      const steered = result.agentState.messageHistory.find(
+        (m) =>
+          m.role === 'user' && JSON.stringify(m.content).includes(steerText),
+      )
+      expect(steered).toBeDefined()
+      expect((steered as { tags?: string[] }).tags).toContain('USER_PROMPT')
+    })
+
+    it('does not extend the turn when no steering messages arrive', async () => {
+      const result = await loopAgentSteps({
+        ...loopAgentStepsBaseParams,
+        agentType: 'test-agent',
+        drainSteeringMessages: () => [],
+      })
+
+      // No steer → the agent ends the turn after its single step, as usual.
+      expect(llmCallCount).toBe(1)
+      expect(result.agentState).toBeDefined()
+    })
+  })
+})