a2a_agent_test.go

Overview

The a2a_agent_test.go file provides comprehensive unit and integration tests for the remote agent communication implementation using the Agent-To-Agent (A2A) protocol. It validates the interaction between local ADK agents and remote agents over the A2A protocol, focusing on:

• Correct translation and streaming of session events and A2A events.

• Proper handling of agent invocation contexts, including session event history.

• Verification of event conversion correctness, including complex cases with function calls, files, and streaming partial responses.

• Error handling in client-server communication.

• Agent card resolution from HTTP endpoints.

• Behavior under edge cases such as empty sessions or incompatible transports.

This test file directly exercises and verifies the functionality described in the [A2A Agent Implementation](80565) topic, ensuring that the client-side remote agent proxy correctly sends and receives A2A protocol messages and events, integrates with session management, and respects agent lifecycle semantics.

Key Types and Utilities

mockA2AExecutor

• Purpose: A mock implementation of the a2asrv.AgentExecutor interface that allows injecting custom execution logic for testing.

• Fields:

  • executeFn func(ctx context.Context, reqCtx *a2asrv.RequestContext, queue eventqueue.Queue) error: Custom function executed during Execute.

• Methods:

  • Execute: Calls executeFn if defined; otherwise returns "not implemented" error.

  • Cancel: Always returns "not implemented" error.

• Usage: Used to simulate server-side agent execution behavior in tests.

Helper Functions

• startA2AServer(t *testing.T, agentExecutor a2asrv.AgentExecutor, listener *bufconn.Listener):
  Starts a gRPC server serving the A2A protocol using the provided agent executor, listening on an in-memory buffer connection (bufconn) for testing without network overhead.

• newTestClientFactory(listener *bufconn.Listener) *a2aclient.Factory:
  Creates a new A2A client factory configured to dial the in-memory server listener with insecure credentials, suitable for testing.

• newA2ARemoteAgent(t *testing.T, name string, listener *bufconn.Listener) agent.Agent:
  Constructs a new remote A2A agent with a given name backed by the in-memory server listener for communication.

• newInvocationContext(t *testing.T, events []*session.Event) agent.InvocationContext:
  Creates a new invocation context with a fresh in-memory session pre-populated with the provided session events.

• runAndCollect(ic agent.InvocationContext, agnt agent.Agent) ([]*session.Event, error):
  Runs the agent with the given invocation context, collects all yielded session events, and returns them along with any error encountered.

• toLLMResponses(events []*session.Event) []model.LLMResponse:
  Extracts the model.LLMResponse payloads from a list of session events for comparison and verification.

• newADKEventReplay(t *testing.T, events []*session.Event) a2asrv.AgentExecutor:
  Returns an executor that replays a fixed sequence of ADK session events as the agent's output for testing the client side.

• newA2AEventReplay(t *testing.T, events []a2a.Event) a2asrv.AgentExecutor:
  Returns a mock executor that writes a predefined sequence of A2A events to the event queue during execution.

• newUserHello() *session.Event:
  Generates a simple user invocation event with text content "hello".

• newFinalStatusUpdate(task *a2a.Task, state a2a.TaskState, msgParts ...a2a.Part) *a2a.TaskStatusUpdateEvent:
  Creates a final task status update event for a given task with optional message parts, marking the task as complete.

Test Suites and Their Purposes

TestRemoteAgent_ADK2ADK

• Purpose:
  Tests the remote agent's ability to correctly handle ADK session events forwarded to a remote ADK agent over the A2A protocol and convert streamed A2A responses back into session events.

• Test Cases:
  Covers scenarios such as:

  • Streaming textual content.

  • Code execution and results.

  • Function call and response handling.

  • File content and URI references.

  • Escalation triggers.

  • Agent transfer requests.

• Validation:
  Compares received LLM responses against expected partial and complete responses, checks that escalation and transfer flags are correctly propagated, and verifies the presence of original A2A request/response metadata in events.

TestRemoteAgent_ADK2A2A

• Purpose:
  Verifies conversion correctness and streaming behavior when the remote agent receives A2A protocol events and translates them into ADK session events on the client side.

• Test Cases:
  Include:

  • Empty messages.

  • Messages with multiple text parts.

  • Tasks with various status updates.

  • Multipart artifacts handling.

  • Multiple task sequences.

  • Non-final and final status update messages treated as "thoughts" or user-visible content.

  • Ignoring empty non-final status updates.

• Checks:
  Compares the LLM response content, including partial results and turn completion flags, and verifies correct metadata in session events.

TestRemoteAgent_EmptyResultForEmptySession

• Purpose:
  Ensures that when the input session events are empty, the remote agent yields an empty result event properly annotated with invocation and branch metadata.

• Behavior Tested:
  Confirms that no invocation happens on the remote side, and a single event with correct context is returned.

TestRemoteAgent_ResolvesAgentCard

• Purpose:
  Tests the remote agent's ability to resolve agent cards from HTTP endpoints dynamically.

• Method:
  Sets up a test HTTP server serving a JSON agent card, then creates a remote agent that fetches this card.

• Validation:
  Confirms that running the remote agent produces the expected LLM response and that event metadata is correctly attached.

TestRemoteAgent_ErrorEventIfNoCompatibleTransport

• Purpose:
  Validates that the remote agent correctly handles the case when no compatible transports are found in the agent card, producing an error event.

• Verification:
  Checks that the error message in the yielded event contains the expected failure text.

TestRemoteAgent_ErrorEventOnServerError

• Purpose:
  Tests that server-side execution errors in the remote agent executor propagate back as error events to the client.

• Setup:
  Uses mockA2AExecutor to simulate a failure during execution.

• Validation:
  Checks that the error message in the session event matches the mock executor error.

Important Implementation Details

• In-Memory gRPC Server Using bufconn:
  Tests avoid network dependencies by using an in-memory buffered connection (bufconn.Listener) to run the gRPC server and client in-process, enabling fast, reliable integration tests.

• Event Conversion and Metadata Checking:
  Tests verify not only the content of session events but also the presence of custom metadata keys (adka2a.ToADKMetaKey("request") and adka2a.ToADKMetaKey("response")) to ensure that the conversion layer preserves traceability.

• Streaming Semantics:
  The interaction with the remote agent follows streaming RPC patterns, with events incrementally yielded, supporting partial responses and progressive streaming output.

• Use of cmp.Diff with Ignored Fields:
  To robustly compare expected and actual results, tests use google/go-cmp with options to ignore fields like timestamps and custom metadata that are not relevant to correctness.

• Invocation Context and Session Management:
  Tests utilize an in-memory session service to simulate realistic invocation contexts, appending events to sessions and ensuring stateful behavior consistent with actual use.

Interaction with Other Parts of the Application

• The tests exercise the Remote Agent Communication (A2A) module's client-side proxy (agent.Agent implementation) and the server-side executor (a2asrv.AgentExecutor), verifying the integration of:

  • Session Management (Session Management) to create and maintain invocation contexts and event histories.

  • Event Conversion and Processing (Event Conversion and Processing) for bi-directional translation between ADK session events and A2A events.

  • Agent Invocation Context (Agent Invocation Context) to manage invocation lifecycles during testing.

• The file indirectly tests ability to stream partial LLM responses, handle function calls, and escalate or transfer sessions, which are core to LLM interaction and agent lifecycle callbacks (LLM Integration and Agents, Agent Lifecycle and Callbacks).

• It uses the A2A client factory and mock executors to simulate remote agent behaviors for testing the entire remote invocation flow.

Visual Diagram: Flow of Core Test Setup and Execution

flowchart TD
Start[Test Initialization] --> SetupServer[Start bufconn gRPC Server]
SetupServer --> CreateExecutor["Create AgentExecutor (mock or replay)"]
CreateExecutor --> StartServer[Run gRPC Server with Executor]
StartServer --> CreateClient[Create A2A Client Factory]
CreateClient --> CreateAgent[Create Remote A2A Agent instance]
CreateAgent --> CreateInvocationCtx[Create InvocationContext with session events]
CreateInvocationCtx --> RunAgent[Run Remote Agent with InvocationContext]
RunAgent --> CollectEvents[Collect streamed Session Events]
CollectEvents --> Verify[Verify Events and Responses]
Verify --> End[Test Completion]

This flowchart outlines the common steps taken in each test case to start the in-memory server, create the remote agent, invoke it with a session context, collect streamed events, and verify results.

Summary of Functions and Their Parameters

Usage Examples in Tests

Example usage pattern in tests:

listener := bufconn.Listen(connBufSize)
executor := newADKEventReplay(t, remoteEvents) // or newA2AEventReplay for A2A events
go startA2AServer(t, executor, listener)

remoteAgent := newA2ARemoteAgent(t, "agent-name", listener)
ictx := newInvocationContext(t, []*session.Event{newUserHello()})
gotEvents, err := runAndCollect(ictx, remoteAgent)
if err != nil {
    t.Fatalf("agent.Run() error = %v", err)
}
// Assertions on gotEvents and converted LLM responses follow...

End of a2a_agent_test.go Documentation