pi_harness/packages/agent/learning/02-AGENT-LOOP-DETAILED.md

Agent Loop Deep Dive

Overview

The agent-loop.ts file contains the core async iteration logic that drives the agent. It's intentionally low-level and stateless - it takes a snapshot of context and drives it to completion.

---

Core Functions

1. runAgentLoop()

Purpose: Start a new agent run with initial prompt messages.

async function runAgentLoop(
  prompts: AgentMessage[],
  context: AgentContext,
  config: AgentLoopConfig,
  emit: AgentEventSink,
  signal: AbortSignal | undefined,
  streamFn: StreamFn,
): Promise<AgentMessage[]>

Flow:

1. Create newMessages = [...prompts]
2. Append prompts to context.messages
3. Emit: agent_start
4. Emit: turn_start
5. For each prompt:
   - Emit: message_start
   - Emit: message_end
6. Call: runLoop() - main iteration logic
7. Return: newMessages

2. runAgentLoopContinue()

Purpose: Continue from existing context (no new prompts).

async function runAgentLoopContinue(
  context: AgentContext,
  config: AgentLoopConfig,
  emit: AgentEventSink,
  signal: AbortSignal | undefined,
  streamFn: StreamFn,
): Promise<AgentMessage[]>

Constraints:

• Last message must convert to user or toolResult
• Throws if context is empty or last message is assistant

Flow:

1. Validate context (non-empty, last message is not assistant)
2. Create newMessages = [] (empty - we continue)
3. Emit: agent_start
4. Emit: turn_start
5. Call: runLoop()
6. Return: newMessages

3. runLoop() - The Heart of the Agent

Purpose: Main iteration loop that drives conversation.

async function runLoop(
  initialContext: AgentContext,
  newMessages: AgentMessage[],
  initialConfig: AgentLoopConfig,
  signal: AbortSignal | undefined,
  emit: AgentEventSink,
  streamFunction: StreamFn,
): Promise<void>

Structure:

async function runLoop(...) {
  let currentContext = initialContext;
  let config = initialConfig;
  let firstTurn = true;
  let pendingMessages: AgentMessage[] = [];
  
  // OUTER LOOP: Handles follow-up messages
  while (true) {
    let hasMoreToolCalls = true;
    
    // INNER LOOP: Handles tool calls and steering
    while (hasMoreToolCalls || pendingMessages.length > 0) {
      if (!firstTurn) {
        await emit({ type: "turn_start" });
      } else {
        firstTurn = false;
      }
      
      // 1. Process pending messages (steering/follow-up)
      if (pendingMessages.length > 0) {
        for (const message of pendingMessages) {
          await emit({ type: "message_start", message });
          await emit({ type: "message_end", message });
          currentContext.messages.push(message);
          newMessages.push(message);
        }
        pendingMessages = [];
      }
      
      // 2. Stream assistant response
      const message = await streamAssistantResponse(...);
      newMessages.push(message);
      
      // 3. Check for errors
      if (message.stopReason === "error" || message.stopReason === "aborted") {
        await emit({ type: "turn_end", message, toolResults: [] });
        await emit({ type: "agent_end", messages: newMessages });
        return;
      }
      
      // 4. Execute tool calls
      const toolCalls = message.content.filter(c => c.type === "toolCall");
      const toolResults: ToolResultMessage[] = [];
      hasMoreToolCalls = false;
      
      if (toolCalls.length > 0) {
        const executedBatch = await executeToolCalls(...);
        toolResults.push(...executedBatch.messages);
        hasMoreToolCalls = !executedBatch.terminate;
        
        for (const result of toolResults) {
          currentContext.messages.push(result);
          newMessages.push(result);
        }
      }
      
      // 5. Emit turn_end
      await emit({ type: "turn_end", message, toolResults });
      
      // 6. Prepare next turn
      const nextTurnContext = { message, toolResults, context, newMessages };
      const nextTurnSnapshot = await config.prepareNextTurn?.(nextTurnContext);
      
      if (nextTurnSnapshot) {
        currentContext = nextTurnSnapshot.context ?? currentContext;
        config = { ...config, model: nextTurnSnapshot.model };
      }
      
      // 7. Check termination
      if (await config.shouldStopAfterTurn?.(...)) {
        await emit({ type: "agent_end", messages: newMessages });
        return;
      }
      
      // 8. Drain steering queue
      pendingMessages = (await config.getSteeringMessages?.()) || [];
    }
    
    // Outer loop: Check for follow-up messages
    const followUpMessages = (await config.getFollowUpMessages?.()) || [];
    if (followUpMessages.length > 0) {
      pendingMessages = followUpMessages;
      continue; // Back to inner loop
    }
    
    // No more messages - exit
    break;
  }
  
  await emit({ type: "agent_end", messages: newMessages });
}

---

Message Streaming

streamAssistantResponse()

Purpose: Stream assistant response from LLM provider.

async function streamAssistantResponse(
  context: AgentContext,
  config: AgentLoopConfig,
  signal: AbortSignal | undefined,
  emit: AgentEventSink,
  streamFunction: StreamFn,
): Promise<AssistantMessage>

Flow:

1. Apply transformContext() if configured
   ├─► messages = await config.transformContext(messages)
   └─► Returns new AgentMessage[]

2. Convert to LLM format
   ├─► llmMessages = await config.convertToLlm(messages)
   └─► Returns Message[] (filters custom messages)

3. Build LLM Context
   Context = {
     systemPrompt: context.systemPrompt,
     messages: llmMessages,
     tools: context.tools
   }

4. Resolve API key
   ├─► Get key from getApiKey() hook
   └─► Fallback to config.apiKey

5. Call streamFn()
   ├─► StreamFn(model, context, options)
   └─► Returns AssistantMessageEventStream

6. Process stream events
   for await (const event of response) {
     switch (event.type) {
       case "start":
         // Initialize partial message
         partialMessage = event.partial
         context.messages.push(partialMessage)
         emit({ type: "message_start", message })
         
       case "text_start" | "text_delta" | "text_end":
       case "thinking_start" | "thinking_delta" | "thinking_end":
       case "toolcall_start" | "toolcall_delta" | "toolcall_end":
         // Update partial message
         partialMessage = event.partial
         emit({ type: "message_update", ... })
         
       case "done" | "error":
         const finalMessage = await response.result()
         emit({ type: "message_end", message })
         return finalMessage
     }
   }

---

Tool Execution

Sequential vs Parallel

Sequential Mode:

• Each tool call prepared, executed, finalized before next
• Emit tool_execution_end immediately after each
• Tool results in source order

Parallel Mode:

• All tool calls prepared sequentially
• Allowed tools execute concurrently
• Emit tool_execution_end in completion order
• Tool results in source order

executeToolCalls()

async function executeToolCalls(...): Promise<ExecutedToolCallBatch> {
  const toolCalls = assistantMessage.content.filter(c => c.type === "toolCall");
  
  // Check if any tool requires sequential execution
  const hasSequentialToolCall = toolCalls.some(tc => {
    const tool = currentContext.tools?.find(t => t.name === tc.name);
    return tool?.executionMode === "sequential";
  });
  
  if (config.toolExecution === "sequential" || hasSequentialToolCall) {
    return executeToolCallsSequential(...);
  }
  
  return executeToolCallsParallel(...);
}

executeToolCallsSequential()

async function executeToolCallsSequential(...): Promise<ExecutedToolCallBatch> {
  const finalizedCalls: FinalizedToolCallOutcome[] = [];
  const messages: ToolResultMessage[] = [];
  
  for (const toolCall of toolCalls) {
    // 1. Prepare
    const preparation = await prepareToolCall(...);
    
    let finalized: FinalizedToolCallOutcome;
    if (preparation.kind === "immediate") {
      // Validation/permission hook blocked execution
      finalized = { toolCall, result: preparation.result, isError: preparation.isError };
    } else {
      // Execute
      const executed = await executePreparedToolCall(preparation, signal, emit);
      finalized = await finalizeExecutedToolCall(...);
    }
    
    // 2. Emit
    await emitToolExecutionEnd(finalized, emit);
    const toolResultMessage = createToolResultMessage(finalized);
    await emitToolResultMessage(toolResultMessage, emit);
    
    finalizedCalls.push(finalized);
    messages.push(toolResultMessage);
    
    if (signal?.aborted) break;
  }
  
  return {
    messages,
    terminate: shouldTerminateToolBatch(finalizedCalls)
  };
}

executeToolCallsParallel()

async function executeToolCallsParallel(...): Promise<ExecutedToolCallBatch> {
  const finalizedCalls: FinalizedToolCallEntry[] = [];
  
  // Phase 1: Prepare all tool calls
  for (const toolCall of toolCalls) {
    const preparation = await prepareToolCall(...);
    
    if (preparation.kind === "immediate") {
      // Blocked or error - execute immediately
      const finalized = {
        toolCall,
        result: preparation.result,
        isError: preparation.isError
      };
      await emitToolExecutionEnd(finalized, emit);
      finalizedCalls.push(finalized);
    } else {
      // Schedule for concurrent execution
      finalizedCalls.push(async () => {
        const executed = await executePreparedToolCall(preparation, signal, emit);
        const finalized = await finalizeExecutedToolCall(...);
        await emitToolExecutionEnd(finalized, emit);
        return finalized;
      });
    }
    
    if (signal?.aborted) break;
  }
  
  // Phase 2: Execute concurrent tools and collect results
  const orderedFinalizedCalls = await Promise.all(
    finalizedCalls.map(entry => typeof entry === "function" ? entry() : Promise.resolve(entry))
  );
  
  // Phase 3: Emit tool result messages in source order
  const messages: ToolResultMessage[] = [];
  for (const finalized of orderedFinalizedCalls) {
    const toolResultMessage = createToolResultMessage(finalized);
    await emitToolResultMessage(toolResultMessage, emit);
    messages.push(toolResultMessage);
  }
  
  return {
    messages,
    terminate: shouldTerminateToolBatch(orderedFinalizedCalls)
  };
}

---

Tool Preparation Flow

prepareToolCall()

async function prepareToolCall(...): Promise<PreparedToolCall | ImmediateToolCallOutcome> {
  // 1. Find tool
  const tool = currentContext.tools?.find(t => t.name === toolCall.name);
  if (!tool) {
    return {
      kind: "immediate",
      result: createErrorToolResult(`Tool ${toolCall.name} not found`),
      isError: true
    };
  }
  
  try {
    // 2. Prepare arguments (optional shim)
    const preparedToolCall = prepareToolCallArguments(tool, toolCall);
    
    // 3. Validate arguments
    const validatedArgs = validateToolArguments(tool, preparedToolCall);
    
    // 4. beforeToolCall hook
    if (config.beforeToolCall) {
      const beforeResult = await config.beforeToolCall(
        { assistantMessage, toolCall, args: validatedArgs, context: currentContext },
        signal
      );
      
      if (signal?.aborted) {
        return immediateError("Operation aborted");
      }
      
      if (beforeResult?.block) {
        return {
          kind: "immediate",
          result: createErrorToolResult(beforeResult.reason || "Tool execution was blocked"),
          isError: true
        };
      }
    }
    
    if (signal?.aborted) {
      return immediateError("Operation aborted");
    }
    
    // 5. Return prepared call for execution
    return {
      kind: "prepared",
      toolCall,
      tool,
      args: validatedArgs
    };
  } catch (error) {
    return {
      kind: "immediate",
      result: createErrorToolResult(error.message),
      isError: true
    };
  }
}

---

Tool Execution Flow

executePreparedToolCall()

async function executePreparedToolCall(
  prepared: PreparedToolCall,
  signal: AbortSignal | undefined,
  emit: AgentEventSink,
): Promise<ExecutedToolCallOutcome> {
  const updateEvents: Promise<void>[] = [];
  let acceptingUpdates = true;
  
  try {
    // Call tool.execute() with onUpdate callback
    const result = await prepared.tool.execute(
      prepared.toolCall.id,
      prepared.args,
      signal,
      (partialResult) => {
        if (!acceptingUpdates) return;
        
        // Buffer update events to emit in order
        updateEvents.push(
          Promise.resolve(
            emit({
              type: "tool_execution_update",
              toolCallId: prepared.toolCall.id,
              toolName: prepared.toolCall.name,
              args: prepared.toolCall.arguments,
              partialResult
            })
          )
        );
      }
    );
    
    acceptingUpdates = false;
    await Promise.all(updateEvents); // Wait for all updates to flush
    return { result, isError: false };
  } catch (error) {
    acceptingUpdates = false;
    await Promise.all(updateEvents);
    return {
      result: createErrorToolResult(error.message),
      isError: true
    };
  } finally {
    acceptingUpdates = false;
  }
}

---

Tool Finalization Flow

finalizeExecutedToolCall()

async function finalizeExecutedToolCall(
  currentContext: AgentContext,
  assistantMessage: AssistantMessage,
  prepared: PreparedToolCall,
  executed: ExecutedToolCallOutcome,
  config: AgentLoopConfig,
  signal: AbortSignal | undefined,
): Promise<FinalizedToolCallOutcome> {
  let result = executed.result;
  let isError = executed.isError;
  
  // afterToolCall hook - can override result
  if (config.afterToolCall) {
    try {
      const afterResult = await config.afterToolCall(
        {
          assistantMessage,
          toolCall: prepared.toolCall,
          args: prepared.args,
          result,
          isError,
          context: currentContext
        },
        signal
      );
      
      if (afterResult) {
        // Field-by-field override (no deep merge)
        result = {
          ...result,
          content: afterResult.content ?? result.content,
          details: afterResult.details ?? result.details,
          usage: afterResult.usage ?? result.usage,
          terminate: afterResult.terminate ?? result.terminate,
        };
        isError = afterResult.isError ?? isError;
      }
    } catch (error) {
      result = createErrorToolResult(error.message);
      isError = true;
    }
  }
  
  return {
    toolCall: prepared.toolCall,
    result,
    isError
  };
}

---

Termination Logic

shouldTerminateToolBatch()

function shouldTerminateToolBatch(finalizedCalls: FinalizedToolCallOutcome[]): boolean {
  return finalizedCalls.length > 0 && 
         finalizedCalls.every(f => f.result.terminate === true);
}

Key Points:

• Only terminates if ALL tool calls set terminate: true
• Allows partial tool execution while signaling early termination

shouldStopAfterTurn()

Called after turn_end, before checking steering/follow-up queues:

if (await config.shouldStopAfterTurn?.({
  message,
  toolResults,
  context: currentContext,
  newMessages
})) {
  await emit({ type: "agent_end", messages: newMessages });
  return;
}

Common use cases:

• Stop before context gets too large
• Stop after completing a specific goal
• Stop on error

---

Queue Management

Steering Queue

Purpose: Interrupt agent while it's working.

When drained: After each turn ends, before next LLM call.

Mode: "all" or "one-at-a-time"

// Example: Steer agent mid-execution
agent.steer("Wait, let me check something else first");
agent.steer("Also, use a different approach");

Follow-up Queue

Purpose: Queue messages for after agent would naturally stop.

When drained: When agent has no more tool calls and no steering messages.

Mode: "all" or "one-at-a-time"

// Example: Follow up after agent finishes
agent.followUp("Now summarize what you did");
agent.followUp("What's next?");

---

Error Handling

Truncated Tool Calls

async function failToolCallsFromTruncatedMessage(
  toolCalls: AgentToolCall[],
  emit: AgentEventSink
): Promise<ExecutedToolCallBatch> {
  // All tool calls from truncated assistant message fail
  // Reason: tool call arguments may be incomplete
  
  for (const toolCall of toolCalls) {
    await emit({ type: "tool_execution_start", ... });
    await emit({
      type: "tool_execution_end",
      toolCallId: toolCall.id,
      toolName: toolCall.name,
      result: createErrorToolResult(
        `Tool call was not executed: response hit output token limit, arguments may be truncated.`
      ),
      isError: true
    });
  }
  
  return { messages: [], terminate: false };
}

---

Abort Handling

All async operations respect the abort signal:

// In prepareToolCall
if (signal?.aborted) {
  return immediateError("Operation aborted");
}

// In executePreparedToolCall
const result = await tool.execute(id, args, signal, onUpdate);
// Tool can check signal.aborted and cancel long-running operations

// In streamAssistantResponse
for await (const event of response) {
  if (signal?.aborted) {
    throw new Error("Aborted");
  }
  // Process event
}

---

Summary

The agent loop is a two-level iterator:

1. Outer loop: Handles follow-up messages after agent would stop
2. Inner loop: Handles tool calls and steering messages

Each iteration:

• Streams assistant response (LLM)
• Executes tool calls (sequential or parallel)
• Emits events for UI updates
• Updates context with new messages

The loop terminates when:

• shouldStopAfterTurn() returns true
• Error or abort occurs
• No more steering/follow-up messages