agent_test.go

Overview

This file contains unit tests for validating the behavior and lifecycle of custom AI agents created using the core agent framework. These tests specifically focus on:

• The execution and effects of before and after agent lifecycle callbacks.

• Correct event generation and content modification during agent invocation.

• Proper handling of early termination of agent invocation via the EndInvocation mechanism.

• Verification of the number of LLM calls made during agent execution.

• Ensuring that callbacks can modify or short-circuit agent runs as intended.

The file utilizes a customAgent implementation as a test double to simulate agent behavior, counting calls and optionally ending invocations early. It integrates with session event objects and uses the cmp package to compare expected and actual event outputs.

This testing file validates core behaviors defined in the broader Agent Lifecycle and Callbacks and AI Agent Framework topics, ensuring that the callback mechanism and invocation flow operate correctly.

Detailed Explanations of Components

TestAgentCallbacks

This test verifies the lifecycle behavior of an agent with various configurations of before and after callbacks.

• Test Cases:

  • before agent callback runs, no llm calls:
    A before-agent callback returns new content, skipping the main agent run (so no LLM calls). It verifies that the event generated contains the callback content.

  • no callback effect if callbacks return nil:
    Both before and after callbacks return nil, so the main run executes normally, causing one LLM call and one event.

  • after agent callback create a new event with new content:
    After-agent callback returns new content, which generates an additional event after the main run.

• Parameters:

  • name (string): Describes the test scenario.

  • beforeAgent ([]BeforeAgentCallback): Slice of callbacks executed before the main agent run.

  • afterAgent ([]AfterAgentCallback): Slice of callbacks executed after the main agent run.

  • wantLLMCalls (int): Expected number of LLM calls made (tracked by customAgent).

  • wantEvents ([]*session.Event): Expected sequence of emitted session events.

• Process:

  • Creates a new agent with specified callbacks and the custom agent’s Run method.

  • Runs the agent and collects emitted events.

  • Verifies the number of LLM calls made.

  • Checks that the generated events match expectations, ignoring auto-generated fields like IDs and timestamps.

• Usage Example:

beforeCallback := func(ctx CallbackContext) (*genai.Content, error) {
    return genai.NewContentFromText("early response", genai.RoleModel), nil
}

agent, err := New(Config{
    Name: "test_agent",
    BeforeAgentCallbacks: []BeforeAgentCallback{beforeCallback},
    Run: custom.Run,
})

for event, err := range agent.Run(ctx) {
    // Process events here
}

TestEndInvocation_EndsBeforeMainCall

This test ensures that when the invocation context signals EndInvocation prior to the main agent run (even if before callbacks return nil), the main LLM call does not occur.

• Sets endInvocation: true in the invocation context.

• Verifies that no LLM call (callCounter) is made.

• Validates that no unexpected errors occur.

This tests the early termination mechanism described in Agent Lifecycle and Callbacks.

TestEndInvocation_EndsAfterMainCall

This test case verifies that when the EndInvocation flag is set during the main run (via the custom agent), subsequent after-agent callbacks are skipped.

• The custom agent sets endInvocation: true during its run.

• The agent is configured with an after-agent callback returning new content.

• The test confirms:

  • Only one LLM call occurs.

  • Only one event is emitted (the main run event).

  • The after-agent callback event is not emitted due to early termination.

customAgent Struct and Method

• Fields:

  • callCounter (int): Counts how many times the Run method is invoked (simulates LLM calls).

  • endInvocation (bool): If true, calls ctx.EndInvocation() during Run.

• Method:

  • Run(ctx InvocationContext) iter.Seq2[*session.Event, error]: Implements the agent's run logic as a generator function.

    • Increments callCounter.

    • If endInvocation is set, calls ctx.EndInvocation() to signal early termination.

    • Yields a single event with a fixed LLM response content "hello".

• Usage:
  The customAgent simulates an agent's run function for testing callbacks and invocation termination without making real LLM calls.

Important Implementation Details and Algorithms

• Callback Execution Order and Effect:
  The tests confirm that before-agent callbacks execute first and can short-circuit the main run by returning content. After-agent callbacks execute only if the invocation is not ended early.

• Invocation Early Termination:
  The EndInvocation flag in the invocation context is respected both before and during the main run, preventing further processing and event emission as needed.

• Event Comparison:
  The test uses go-cmp with options to ignore fields like ID, Timestamp, InvocationID, and StateDelta to focus on the substantive content of events.

• Event Generation:
  Events are generated either from callbacks (returning content) or from the main agent run. The test verifies that events have the correct author and content.

• Parallel Execution:
  The tests run in parallel (t.Parallel()) to allow concurrency without shared state conflicts.

Interaction with Other Parts of the System

• Agent Creation:
  Uses New(Config) from the core agent framework (AI Agent Framework) to create agents with lifecycle callbacks.

• Invocation Context:
  Passes a custom invocationContext to simulate real invocation context behavior (Agent Invocation Context).

• Session Events:
  Produces and verifies session.Event objects that represent outputs emitted by agents during invocation (Session Management).

• Callbacks:
  Implements and tests BeforeAgentCallback and AfterAgentCallback functions that modify or short-circuit agent runs (Agent Lifecycle and Callbacks).

• LLM Content:
  Uses genai.Content objects to represent LLM-generated text content and roles (LLM Integration and Agents).

Diagram: Agent Invocation Lifecycle in Tests

flowchart TD
Start["Test Start"]
CreateAgent["Create Agent with Callbacks"]
RunAgent["Run Agent.Run(ctx)"]
BeforeCB["Execute BeforeAgentCallbacks"]
CheckBefore["Callback returned content?"]
SkipRun["Skip Main Run, Yield Callback Event"]
MainRun["Execute Main Agent Run"]
AfterCB["Execute AfterAgentCallbacks"]
CheckAfter["Callback returned content?"]
YieldAfter["Yield After Callback Event"]
CollectEvents["Collect Events"]
VerifyCalls["Verify LLM calls count"]
VerifyEvents["Compare Produced Events"]
End["Test End"]
Start --> CreateAgent --> RunAgent --> BeforeCB --> CheckBefore
CheckBefore -- Yes --> SkipRun --> CollectEvents
CheckBefore -- No --> MainRun --> AfterCB --> CheckAfter
CheckAfter -- Yes --> YieldAfter --> CollectEvents
CheckAfter -- No --> CollectEvents
CollectEvents --> VerifyCalls --> VerifyEvents --> End

Summary of Functions and Methods

Usage Context

This test file is intended for developers working on the AI Agent Framework, particularly those extending or maintaining lifecycle callback functionality in agents. It demonstrates how to:

• Define and configure agents with lifecycle callbacks.

• Simulate agent runs with custom behavior for testing.

• Verify correct event emission and invocation flow.

• Utilize the invocation context and lifecycle control flags.

This file directly supports the robustness and correctness of agent invocation lifecycle management as described in Agent Lifecycle and Callbacks.

End of Documentation for agent_test.go