test: Establish test safety net for controller2 refactoring

e2e.go

@@ -0,0 +1,124 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package main implements an end-to-end demonstration of the Antigravity harness
+// integration with AX Controller V2.
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"github.com/google/ax/internal/controller/executor"
+	"github.com/google/ax/internal/controller/executor/executortest"
+	"github.com/google/ax/internal/controller2"
+	"github.com/google/ax/internal/harness"
+	"github.com/google/ax/proto"
+)
+
+func main() {
+	ctx := context.Background()
+	fmt.Println("==================================================")
+	fmt.Println("AX Controller V2 - E2E Harness Demonstration")
+	fmt.Println("==================================================")
+
+	// -------------------------------------------------------------------------
+	// Demo 1: Runtime Fallback (No harness registered)
+	// -------------------------------------------------------------------------
+	fmt.Println("\n--- Demo 1: Runtime Fallback ---")
+	fmt.Println("Requesting 'unregistered-agent'. Should fallback to Test Harness (Hello World).")
+	runDemo(ctx, "unregistered-agent", func(reg *controller2.Registry) {
+		// Do not register any harness
+	})
+
+	// -------------------------------------------------------------------------
+	// Demo 2: Build-time Fallback (Antigravity with bad script path)
+	// -------------------------------------------------------------------------
+	fmt.Println("\n--- Demo 2: Build-time Fallback ---")
+	fmt.Println("Registering 'antigravity' with non-existent script. Should fallback to Test Harness.")
+	runDemo(ctx, "antigravity", func(reg *controller2.Registry) {
+		// Build harness with bad path
+		badHarness := controller2.BuildHarness(ctx, "antigravity", harness.HarnessConfig{
+			AntigravityScriptPath: "non-existent-script.py",
+		})
+		reg.RegisterHarness("antigravity", badHarness)
+	})
+
+	// -------------------------------------------------------------------------
+	// Demo 3: Antigravity Execution (Requires google-antigravity & GEMINI_API_KEY)
+	// -------------------------------------------------------------------------
+	fmt.Println("\n--- Demo 3: Antigravity Execution ---")
+	fmt.Println("Registering 'antigravity' with real script. Attempting execution.")
+	if os.Getenv("GEMINI_API_KEY") == "" {
+		fmt.Println("WARNING: GEMINI_API_KEY is not set. Execution will likely fail if dependencies are missing, but we will try anyway.")
+	}
+	runDemo(ctx, "antigravity", func(reg *controller2.Registry) {
+		// Build harness with real path (empty defaults to examples/antigravity_agent/agent.py)
+		realHarness := controller2.BuildHarness(ctx, "antigravity", harness.HarnessConfig{
+			AntigravityScriptPath: "examples/antigravity_agent/agent.py",
+		})
+		reg.RegisterHarness("antigravity", realHarness)
+	})
+}
+
+func runDemo(ctx context.Context, agentID string, setupRegistry func(reg *controller2.Registry)) {
+	reg := controller2.NewRegistry()
+	setupRegistry(reg)
+
+	log := &executortest.MemoryEventLog{}
+	c, err := controller2.New(ctx, controller2.Config{
+		Registry: reg,
+		EventLogBuilder: func() (executor.EventLog, error) {
+			return log, nil
+		},
+	})
+	if err != nil {
+		fmt.Printf("Error creating controller: %v\n", err)
+		return
+	}
+	defer c.Close()
+
+	handler := controller2.ExecHandler(func(resp *proto.ExecResponse) error {
+		for _, out := range resp.Outputs {
+			if textContent := out.GetContent().GetText().GetText(); textContent != "" {
+				fmt.Printf("Agent Output: %s\n", textContent)
+			}
+		}
+		return nil
+	})
+
+	inputs := []*proto.Message{
+		{
+			Role: "user",
+			Content: &proto.Content{
+				Type: &proto.Content_Text{
+					Text: &proto.TextContent{Text: "Who are you?"},
+				},
+			},
+		},
+	}
+
+	err = c.Exec(ctx, &proto.ExecRequest{
+		ConversationId: "e2e-conv",
+		Inputs:         inputs,
+		AgentId:        agentID,
+	}, handler)
+
+	if err != nil {
+		fmt.Printf("Execution Failed (as expected if environment is not ready): %v\n", err)
+	} else {
+		fmt.Println("Execution Succeeded!")
+	}
+}

examples/antigravity_agent/README.md

@@ -0,0 +1,27 @@
+# Antigravity Agent Example
+
+This directory contains a simple example of an agent built using the `google-antigravity` SDK.
+
+## Prerequisites
+
+Ensure you have Python 3.10+ installed.
+
+## Setup
+
+1.  Install the required dependencies:
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+2.  Set your Gemini API key in your environment:
+    ```bash
+    export GEMINI_API_KEY="your-api-key-here"
+    ```
+
+## Running the Agent
+
+Run the agent script directly:
+
+```bash
+python agent.py
+```

examples/antigravity_agent/agent.py

@@ -0,0 +1,28 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import sys
+from google.antigravity import Agent, LocalAgentConfig
+
+async def main():
+    # Initialize the agent configuration. It automatically picks up GEMINI_API_KEY from the environment.
+    config = LocalAgentConfig()
+    async with Agent(config) as agent:
+        prompt = sys.argv[1] if len(sys.argv) > 1 else "Explain quantum computing in one sentence."
+        response = await agent.chat(prompt)
+        print(await response.text())
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/antigravity_agent/requirements.txt

@@ -0,0 +1 @@
+google-antigravity

internal/config/config.go

@@ -39,6 +39,17 @@ type RegistryConfig struct {
 	RemoteAgents    []RemoteAgentConfig    `yaml:"remote_agents,omitempty"`
 	ColabAgents     []ColabAgentConfig     `yaml:"colab_agents,omitempty"`
 	SubstrateAgents []SubstrateAgentConfig `yaml:"substrate_agents,omitempty"`
+	Harnesses       []HarnessConfig        `yaml:"harnesses,omitempty"`
+}
+
+type HarnessConfig struct {
+	ID          string            `yaml:"id"`
+	Type        string            `yaml:"type"` // "antigravity"
+	Antigravity AntigravityConfig `yaml:"antigravity,omitempty"`
+}
+
+type AntigravityConfig struct {
+	ScriptPath string `yaml:"script_path"`
 }
 
 // ATEConfig configures the ATE integration.

internal/controller2/builder.go

@@ -0,0 +1,57 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package controller2
+
+import (
+	"context"
+	"log"
+	"os"
+	"os/exec"
+
+	"github.com/google/ax/internal/harness"
+	"github.com/google/ax/internal/harness/harnesstest"
+)
+
+// BuildHarness builds a harness based on the requested type, with fallback to test harness.
+func BuildHarness(ctx context.Context, harnessType string, cfg harness.HarnessConfig) harness.Harness {
+	switch harnessType {
+	case "antigravity":
+		// Check if python3 is available
+		if _, err := exec.LookPath("python3"); err != nil {
+			log.Printf("WARNING: python3 not found in PATH, falling back to test harness: %v", err)
+			return harnesstest.New()
+		}
+		// Check if script exists
+		scriptPath := cfg.AntigravityScriptPath
+		if scriptPath == "" {
+			scriptPath = "examples/antigravity_agent/agent.py"
+		}
+		if _, err := os.Stat(scriptPath); err != nil {
+			log.Printf("WARNING: Antigravity agent script not found at %s, falling back to test harness: %v", scriptPath, err)
+			return harnesstest.New()
+		}
+		log.Printf("Using Antigravity harness with script: %s", scriptPath)
+
+		builder := &harness.AntigravityHarnessBuilder{
+			Config: harness.HarnessConfig{
+				AntigravityScriptPath: scriptPath,
+			},
+		}
+		return builder.Build()
+	default:
+		log.Printf("Using default test harness")
+		return harnesstest.New()
+	}
+}

internal/controller2/controller.go

@@ -19,6 +19,7 @@ package controller2
 import (
 	"context"
 	"fmt"
+	"log"
 
 	"github.com/google/ax/internal/controller/executor"
 	"github.com/google/ax/internal/harness/harnesstest"
@@ -31,20 +32,21 @@ type ExecHandler func(resp *proto.ExecResponse) error
 // Controller is the main controller that coordinates all components.
 // It acts as a single-writer system for managing agentic loops.
 type Controller struct {
-	registry       *Registry
-	eventLog       executor.EventLog
+	registry      *Registry
+	eventLog      executor.EventLog
 }
 
 // Config configures the controller.
 type Config struct {
+	Registry        *Registry
 	EventLogBuilder executor.EventLogBuilder
 }
 
 // New creates a new controller instance.
 func New(ctx context.Context, cfg Config) (*Controller, error) {
-	// Initialize agent registry
-	registry := NewRegistry()
-
+	if cfg.Registry == nil {
+		return nil, fmt.Errorf("registry is required")
+	}
 	if cfg.EventLogBuilder == nil {
 		return nil, fmt.Errorf("event log builder is required")
 	}
@@ -54,8 +56,8 @@ func New(ctx context.Context, cfg Config) (*Controller, error) {
 	}
 
 	return &Controller{
-		registry:       registry,
-		eventLog:       eventLog,
+		registry:      cfg.Registry,
+		eventLog:      eventLog,
 	}, nil
 }
 
@@ -70,7 +72,13 @@ func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler E
 	// TODO(jbd): Resume an incomplete execution if there exists one.
 	// TODO(jbd): Enable bringing a remote harness that implements HarnessService.
 
-	h := harnesstest.New()
+	// Retrieve harness from registry
+	h, err := d.registry.GetHarness(req.AgentId)
+	if err != nil {
+		// Fallback to test harness
+		log.Printf("WARNING: harness %s not found in registry, falling back to test harness: %v", req.AgentId, err)
+		h = harnesstest.New()
+	}
 	exec, err := h.Start(ctx, req.ConversationId)
 	if err != nil {
 		return fmt.Errorf("failed to start harness session: %w", err)