agent_test.go

Overview

The agent_test.go file contains unit tests and supporting helper functions to validate the behavior of the Parallel Agent, a workflow agent that executes multiple sub-agents concurrently. This test suite verifies the correct orchestration of sub-agents, proper event streaming, error propagation, and cancellation handling within the parallel execution context.

This file exercises the Parallel Agent functionality as part of the broader Agent Workflow Management framework, specifically focusing on concurrent execution patterns implemented in the parallelagent package. It uses sub-agents created with the loopagent package, which repeatedly emits events, to simulate realistic agent workloads. The tests ensure that the parallel agent correctly merges events from multiple sub-agents and handles error and cancellation scenarios gracefully.

The tests run agents inside an in-memory session service and utilize the agent execution runner (Agent Execution Runner) to simulate the actual runtime environment of agents.

Detailed Explanations

TestNewParallelAgent

func TestNewParallelAgent(t *testing.T)

• Purpose:
  Validates the behavior of the ParallelAgent under various conditions:

  • Normal completion of sub-agents emitting fixed numbers of events.

  • Handling of context cancellation to terminate infinite loops.

  • Proper error propagation when one sub-agent returns an error.

• Test Cases:

  • subagents complete run:
    Runs 3 sub-agents, each emitting 2 events. Verifies that all expected events are generated with correct authorship and content.

  • handle ctx cancel:
    Runs an infinite loop agent but cancels the context shortly after start. Verifies that execution stops and an error is returned.

  • agent returns error:
    Runs 100 sub-agents where one sub-agent returns a predefined error. Confirms that the error is propagated and execution halts.

• Parameters: None (standard testing.T parameter).

• Key Behavior:

  • Uses newParallelAgent helper to create a parallel agent configured per test case.

  • Runs the agent via runner.New and agentRunner.Run.

  • Collects events yielded from the agent and compares against expected events using cmp.Diff.

  • Handles timing and cancellation with goroutines and context manipulation.

  • Uses parallel subtests (t.Parallel()) for concurrency during testing.

• Usage Example:

t.Run("subagents complete run", func(t *testing.T) {
    // setup parallel agent with 3 sub-agents and max 2 iterations
    // run agent and collect events
    // compare received events to expected events
})

newParallelAgent

func newParallelAgent(t *testing.T, maxIterations uint, numSubAgents int, agentErr error) agent.Agent

• Purpose:
  Constructs a configured ParallelAgent instance with a given number of sub-agents, each implemented as a looping agent producing events up to maxIterations. Optionally adds a sub-agent that returns a specified error.

• Parameters:

  • t *testing.T: testing instance to report fatal errors.

  • maxIterations uint: number of iterations each sub-agent will run. Zero means infinite loop.

  • numSubAgents int: number of loop sub-agents to create.

  • agentErr error: if non-nil, a special error-returning sub-agent is added.

• Returns:

  • agent.Agent: the constructed parallel agent ready for execution.

• Implementation Details:

  • Creates multiple sub-agents using the loopagent.New function with MaxIterations.

  • Each loop agent contains a single sub-agent (via agent.New) that runs the customRun function generating events with author names like "sub1", "sub2", etc.

  • If agentErr is specified, appends an additional sub-agent that immediately returns the error.

  • Wraps all sub-agents into a parallelagent.New instance.

  • Uses the must helper to panic on any construction errors, ensuring test failures are explicit.

• Usage Example:

parallelAgent := newParallelAgent(t, 2, 3, nil)
// parallelAgent runs 3 sub-agents emitting 2 events each

must

func must[T agent.Agent](a T, err error) T

• Purpose:
  Utility helper to unwrap an agent creation call returning (agent.Agent, error). Panics if the error is non-nil.

• Parameters:

  • a T: agent instance returned.

  • err error: error returned.

• Returns:

  • T: agent instance if no error.

• Usage:
  Simplifies test setup by avoiding repetitive error checking.

customRun

func customRun(id int, agentErr error) func(agent.InvocationContext) iter.Seq2[*session.Event, error]

• Purpose:
  Generates a custom run function for an agent that produces a single event or an error.

• Parameters:

  • id int: identifier used to generate event content.

  • agentErr error: if non-nil, the run function immediately yields this error.

• Returns:

  • A function matching the agent run signature that yields either one event or an error.

• Behavior:

  • Simulates variable processing delay (1-5 ms).

  • If agentErr is present, yields error immediately.

  • Otherwise, yields one session.Event with LLMResponse content "hello <id>".

• Usage Example:

runFn := customRun(1, nil) // yields one event with "hello 1"
runFnWithError := customRun(-1, fmt.Errorf("error")) // yields error immediately

Important Implementation Details and Algorithms

• Parallel Execution:
  The tested ParallelAgent runs each sub-agent concurrently in isolated invocation contexts, leveraging Go concurrency primitives. It collects events via channels and merges streams into a unified output.

• Event Collection and Sorting:
  Tests use cmp.Diff and sorting by author name to compare unordered event streams, reflecting that parallel execution may yield events in non-deterministic order.

• Context Cancellation Handling:
  Tests simulate context cancellation to verify that the agent properly terminates ongoing operations, which is critical for resource management and responsiveness.

• Error Propagation:
  Propagation of errors from any sub-agent causes the entire parallel agent run to terminate with an error, validating fail-fast behavior.

Interaction with Other Parts of the System

• parallelagent Package:
  This test file exercises the parallelagent.New parallel agent implementation, which is part of the Agent Workflow Management system.

• loopagent Package:
  Sub-agents are implemented as loop agents (loopagent.New) to simulate iterative event emission, enabling testing of both finite and infinite iteration scenarios.

• Agent Execution Runner:
  Uses runner.New and agentRunner.Run from the Agent Execution Runner to execute agents within session contexts, managing lifecycle and event streaming.

• Session Service:
  Utilizes in-memory session service (session.InMemoryService) from Session Management to manage session lifecycles and event histories.

• Invocation Context:
  Each sub-agent runs in a forked invocation context (agent.InvocationContext) isolating branches while sharing common session and artifact layers.

• Event Model:
  Uses session.Event and model.LLMResponse to represent agent output, consistent with the overall agent event streaming architecture.

Visual Diagram: Structure and Workflow of the Test Setup and Parallel Agent Execution

flowchart TD
TestNewParallelAgent[TestNewParallelAgent]
newParallelAgent[newParallelAgent]
parallelAgent[ParallelAgent Instance]
SubAgentLoop1[LoopAgent Sub-Agent 1]
SubAgentLoop2[LoopAgent Sub-Agent 2]
...
SubAgentLoopN[LoopAgent Sub-Agent N]
customRunFn[customRun Function]
runner[Agent Runner]
sessionService[InMemory Session Service]
events[Collected Events]
ctx[Context & Cancellation]
TestNewParallelAgent -->|calls| newParallelAgent
newParallelAgent --> parallelAgent
parallelAgent --> SubAgentLoop1
parallelAgent --> SubAgentLoop2
parallelAgent --> SubAgentLoopN
SubAgentLoop1 --> customRunFn
SubAgentLoop2 --> customRunFn
SubAgentLoopN --> customRunFn
TestNewParallelAgent --> runner
runner --> parallelAgent
runner --> sessionService
runner --> ctx
parallelAgent --> events
runner --> events
TestNewParallelAgent --> events

• Explanation:
  The diagram shows the high-level relationships during test execution:

  • TestNewParallelAgent creates a ParallelAgent via newParallelAgent.

  • The ParallelAgent manages multiple LoopAgent sub-agents.

  • Each LoopAgent invokes a customRun function that generates events.

  • The agent execution is managed by the runner, which interacts with an in-memory sessionService and a context for cancellation.

  • Finally, the events produced by sub-agents are collected and asserted in the test.

References to Related Topics

• The ParallelAgent tested here is a key part of the Agent Workflow Management system enabling concurrent sub-agent orchestration.

• The sub-agents are implemented using the Loop Agent pattern to simulate iterative event emission.

• The execution environment is managed by the Agent Execution Runner coordinating agent runs and session state.

• Event streaming and session management interfaces are defined in Session Management.

• The agent run functions utilize the agent.InvocationContext interface from Agent Invocation Context.

• The produced events conform to the session.Event model, part of the overall agent event streaming mechanism.

This documentation provides a comprehensive technical overview of the agent_test.go file, covering its role in verifying the parallel agent's execution correctness, the structure of helper functions used, and its integration within the agent workflow and execution framework.