genai.go

Overview

The genai.go file in the testutil package provides utilities for configuring HTTP client transport tailored for Google’s Gemini AI services, specifically enabling record and replay (RR) capabilities during testing. It wraps the HTTP transport with a custom round-tripper that records outbound requests and responses, and replays them during tests to ensure deterministic behavior. The file also implements request scrubbing logic to canonicalize Gemini API requests by removing or normalizing headers and JSON payloads that could vary between calls, such as API keys or version identifiers. This facilitates consistent replay and comparison of HTTP interactions.

Functions and Methods

NewGeminiTransport(rrfile string) (http.RoundTripper, error)

• Purpose:
  Creates and returns an HTTP RoundTripper configured for recording and replaying HTTP requests/responses from/to Gemini AI APIs using a specified record/replay file.

• Parameters:

  • rrfile (string): The path to the record/replay file storing serialized HTTP interactions.

• Returns:

  • http.RoundTripper: A RoundTripper configured to replay recorded Gemini HTTP requests or record new interactions.

  • error: Returns an error if opening the record/replay file fails.

• Usage:
  This function is typically called to create an HTTP client transport for test clients that interact with Gemini AI services, enabling deterministic testing by replaying prior recorded HTTP interactions.

• Implementation Details:
  The function calls httprr.Open, passing the file path and the default HTTP transport. It then registers a scrub function to clean the requests before replay or record. This scrub function removes headers that interfere with consistent replay and canonicalizes JSON bodies.

rr, err := httprr.Open(rrfile, http.DefaultTransport)
if err != nil {
    return nil, fmt.Errorf("httprr.Open(%q) failed: %w", rrfile, err)
}
rr.ScrubReq(scrubGeminiRequest)
return rr, nil

scrubGeminiRequest(req *http.Request) error

• Purpose:
  A helper function that sanitizes outgoing HTTP requests to Gemini AI APIs by removing or normalizing volatile headers and canonicalizing JSON request bodies.

• Parameters:

  • req (*http.Request): The HTTP request to be scrubbed.

• Returns:

  • error: Always returns nil currently; no error conditions are implemented.

• Implementation Details:
  The function deletes or removes multiple headers which contain sensitive or variable data that should not affect replay matching:

  • "x-goog-api-key" / "X-Goog-Api-Key"

  • "x-goog-api-client" / "X-Goog-Api-Client"

  • "user-agent" / "User-Agent"

  For requests with a JSON content type (e.g., "application/json"), the JSON body is canonicalized by compacting it to remove unnecessary whitespace. This normalization addresses the randomization introduced by protobuf JSON encoding (which may add or omit spaces after commas), ensuring consistent serialized JSON for replay.

• Key Algorithm:
  Uses json.Compact to minimize the JSON in the request body to a canonical form.

• Note:
  The request body is expected to be of type *httprr.Body, which exposes a Data field containing the raw request body bytes.

if ctype := req.Header.Get("Content-Type"); ctype == "application/json" || strings.HasPrefix(ctype, "application/json;") {
    b := req.Body.(*httprr.Body)
    var buf bytes.Buffer
    if err := json.Compact(&buf, b.Data); err == nil {
        b.Data = buf.Bytes()
    }
}

Important Implementation Details

• Record and Replay (RR) Integration:
  The file leverages the httprr package from the internal ADK library, which supports recording HTTP requests and responses to a file, and replaying them during tests to enable consistent, repeatable network interactions.

• Header Scrubbing:
  The scrubbing process removes headers that contain ephemeral or environment-specific information such as API keys and user-agent strings that include version numbers. This avoids spurious mismatches during replay tests.

• JSON Canonicalization:
  Protobuf JSON encoding tends to produce randomized whitespace patterns. The JSON compaction step ensures all JSON request bodies have a consistent compact representation, improving replay matching fidelity.

• Error Handling:
  Only errors arising from opening the RR file are propagated. The scrubbing function handles errors internally and does not propagate errors to avoid disrupting the RR flow.

Interaction With Other System Components

• Relies on the httprr package (google.golang.org/adk/internal/httprr) to provide the record and replay HTTP transport implementation. This package manages the persistence and matching of HTTP interactions.

• Intended to be used by Gemini AI client configurations (genai.ClientConfig) to inject a transport layer that supports deterministic testing through HTTP interaction recording/replaying.

• Removes or modifies HTTP headers and request bodies prior to sending requests to Gemini AI services, ensuring compatibility with Gemini’s API canonicalization requirements.

• Can be utilized in testing frameworks or integration tests that involve Gemini AI interactions, reducing the need for live network calls and enabling offline test execution.

Visual Diagram

flowchart TD
A["NewGeminiTransport(rrfile)"] -->|calls| B["httprr.Open(rrfile, DefaultTransport)"]
B -->|returns| C[httprr.RecordReplayTransport]
C -->|registers scrub function| D["scrubGeminiRequest(req)"]
D -->|removes headers| E[x-goog-api-key, x-goog-api-client, user-agent]
D -->|canonicalizes JSON body| F[json.Compact]

Summary of Key Elements

This file is a utility specifically crafted for Gemini AI HTTP client setup in test environments, ensuring reproducible and sanitized HTTP request/response cycles. It is a foundational piece in the test infrastructure surrounding Gemini AI integration, facilitating reliable and consistent test execution.

For further details on HTTP transport and adapter patterns, see Agent Invocation Context and REST API and Web Launchers. For integration with Gemini and other Google AI services, refer to Google Search Tool - Gemini-native tool.