From 7b7d8bdbee17094ed96204f951e8b1482a9febcf Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Fri, 22 May 2026 11:46:26 -0700 Subject: [PATCH] feat: Integrate Antigravity harness in Controller V2 with Registry-centric design This PR implements the integration of the Antigravity agent as the first built-in harness in AX Controller V2, satisfying the first item on the AX roadmap. It also refactors the architecture to make `Registry` the Single Source of Truth (SSOT) and introduces robust fallback mechanisms. ### Goals 1. **Built-in Harness Integration**: Enable AX Controller V2 to execute Python-based Antigravity agents. 2. **Registry-Centric Architecture (SSOT)**: Centralize the management of both agents and harnesses in the `Registry`. 3. **Inverted Control (Dependency Injection)**: Decouple the `Controller` from harness creation by passing `Registry` as an input config. 4. **Resilient Fallbacks**: Implement build-time (script check) and runtime (registration check) fallbacks to a Test Harness to prevent failures in unconfigured environments. 5. **End-to-End Verification**: Provide a comprehensive E2E demonstration script. ### Key Changes - **Registry Updates (`internal/controller2/registry.go`)**: Added support for registering and retrieving Go `Harness` instances in `Registry`. - **Controller V2 Refactoring (`internal/controller2/controller.go`)**: - Received `Registry` as an input config in `New`. - Updated `Exec` to retrieve the harness from the registry using `AgentId` at runtime, falling back to the test harness if not found. - **Antigravity Go Harness (`internal/harness/antigravity.go`)**: - Implemented `AntigravityHarness` which executes the Python agent as a subprocess, passing the prompt as an argument. - Captured stdout and returned it as a streamed message. - Added a TODO to migrate to a gRPC server in the next step to avoid subprocess overhead. - **Python Agent Modification (`examples/antigravity_agent/agent.py`)**: Modified the script to accept dynamic prompt inputs from command line arguments. - **Verification (`internal/controller2/controller_test.go`, `fork_test.go`, `e2e.go`)**: - Updated unit tests to adapt to the new `Registry` injection. - Added unit tests for both build-time (implemented manually in test setup) and runtime fallbacks. - Created `e2e.go` in the root to demonstrate all 3 execution paths (Runtime Fallback, Build-time Fallback, and actual Antigravity Happy Path). --- e2e.go | 142 +++++++++++++++++++ examples/antigravity_agent/README.md | 27 ++++ examples/antigravity_agent/agent.py | 28 ++++ examples/antigravity_agent/requirements.txt | 1 + internal/controller2/controller.go | 24 ++-- internal/controller2/controller_test.go | 127 +++++++++++++++++ internal/controller2/fork_test.go | 4 + internal/controller2/registry.go | 20 +++ internal/harness/antigravity.go | 148 ++++++++++++++++++++ 9 files changed, 513 insertions(+), 8 deletions(-) create mode 100644 e2e.go create mode 100644 examples/antigravity_agent/README.md create mode 100644 examples/antigravity_agent/agent.py create mode 100644 examples/antigravity_agent/requirements.txt create mode 100644 internal/harness/antigravity.go diff --git a/e2e.go b/e2e.go new file mode 100644 index 0000000..68727b6 --- /dev/null +++ b/e2e.go @@ -0,0 +1,142 @@ +// 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" + "os/exec" + + "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/internal/harness/harnesstest" + "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, manually implementing fallback check + var badHarness harness.Harness + scriptPath := "non-existent-script.py" + if _, err := exec.LookPath("python3"); err != nil { + fmt.Printf("WARNING: python3 not found, falling back to test harness: %v\n", err) + badHarness = harnesstest.New() + } else if _, err := os.Stat(scriptPath); err != nil { + fmt.Printf("WARNING: Antigravity agent script not found at %s, falling back to test harness: %v\n", scriptPath, err) + badHarness = harnesstest.New() + } else { + badHarness = harness.NewAntigravityHarness(scriptPath) + } + 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, manually implementing fallback check + var realHarness harness.Harness + scriptPath := "examples/antigravity_agent/agent.py" + if _, err := exec.LookPath("python3"); err != nil { + fmt.Printf("WARNING: python3 not found, falling back to test harness: %v\n", err) + realHarness = harnesstest.New() + } else if _, err := os.Stat(scriptPath); err != nil { + fmt.Printf("WARNING: Antigravity agent script not found at %s, falling back to test harness: %v\n", scriptPath, err) + realHarness = harnesstest.New() + } else { + realHarness = harness.NewAntigravityHarness(scriptPath) + } + 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!") + } +} diff --git a/examples/antigravity_agent/README.md b/examples/antigravity_agent/README.md new file mode 100644 index 0000000..400de30 --- /dev/null +++ b/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 +``` diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py new file mode 100644 index 0000000..fa8603a --- /dev/null +++ b/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()) diff --git a/examples/antigravity_agent/requirements.txt b/examples/antigravity_agent/requirements.txt new file mode 100644 index 0000000..48da6dd --- /dev/null +++ b/examples/antigravity_agent/requirements.txt @@ -0,0 +1 @@ +google-antigravity diff --git a/internal/controller2/controller.go b/internal/controller2/controller.go index 6095515..086b41d 100644 --- a/internal/controller2/controller.go +++ b/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) diff --git a/internal/controller2/controller_test.go b/internal/controller2/controller_test.go index 3dae3a4..259087f 100644 --- a/internal/controller2/controller_test.go +++ b/internal/controller2/controller_test.go @@ -16,11 +16,14 @@ package controller2 import ( "context" + "os" "testing" "github.com/google/ax/internal/agent" "github.com/google/ax/internal/controller/executor" "github.com/google/ax/internal/controller/executor/executortest" + "github.com/google/ax/internal/harness" + "github.com/google/ax/internal/harness/harnesstest" "github.com/google/ax/proto" ) @@ -37,7 +40,9 @@ func TestController2_ExecHelloWorld(t *testing.T) { cid := "test-conversation-id" log := &executortest.MemoryEventLog{} + reg := NewRegistry() c, err := New(ctx, Config{ + Registry: reg, EventLogBuilder: func() (executor.EventLog, error) { return log, nil }, @@ -81,3 +86,125 @@ func TestController2_ExecHelloWorld(t *testing.T) { t.Errorf("expected 'Hello world' output text response, got %q", gotText) } } + +func TestController2_ExecAntigravityFallback(t *testing.T) { + ctx := context.Background() + cid := "test-conversation-id" + + log := &executortest.MemoryEventLog{} + reg := NewRegistry() + + // Build and register harness with bad path to trigger build-time fallback + var badHarness harness.Harness + scriptPath := "non-existent-script.py" + if _, err := os.Stat(scriptPath); err != nil { + badHarness = harnesstest.New() // Fallback + } else { + badHarness = harness.NewAntigravityHarness(scriptPath) + } + reg.RegisterHarness("antigravity", badHarness) + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (executor.EventLog, error) { + return log, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + var outputs []*proto.Message + handler := ExecHandler(func(resp *proto.ExecResponse) error { + outputs = append(outputs, resp.Outputs...) + return nil + }) + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "Trigger prompt"}, + }, + }, + }, + } + + // Request "antigravity" agent + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: inputs, + AgentId: "antigravity", + }, handler) + if err != nil { + t.Fatalf("Controller2.Exec failed: %v", err) + } + + if len(outputs) != 1 { + t.Fatalf("expected exactly 1 output message, got %d", len(outputs)) + } + + gotText := outputs[0].GetContent().GetText().GetText() + if gotText != "Hello world" { + t.Errorf("expected 'Hello world' output text response due to fallback, got %q", gotText) + } +} + +func TestController2_ExecRuntimeFallback(t *testing.T) { + ctx := context.Background() + cid := "test-conversation-id" + + log := &executortest.MemoryEventLog{} + reg := NewRegistry() // Empty registry, will force runtime fallback for any requested agent + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (executor.EventLog, error) { + return log, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + var outputs []*proto.Message + handler := ExecHandler(func(resp *proto.ExecResponse) error { + outputs = append(outputs, resp.Outputs...) + return nil + }) + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "Trigger prompt"}, + }, + }, + }, + } + + // Request "antigravity" agent, which is NOT registered + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: inputs, + AgentId: "antigravity", + }, handler) + if err != nil { + t.Fatalf("Controller2.Exec failed: %v", err) + } + + if len(outputs) != 1 { + t.Fatalf("expected exactly 1 output message, got %d", len(outputs)) + } + + gotText := outputs[0].GetContent().GetText().GetText() + if gotText != "Hello world" { + t.Errorf("expected 'Hello world' output text response due to runtime fallback, got %q", gotText) + } +} + + diff --git a/internal/controller2/fork_test.go b/internal/controller2/fork_test.go index 3f5fff4..2946d4a 100644 --- a/internal/controller2/fork_test.go +++ b/internal/controller2/fork_test.go @@ -49,7 +49,9 @@ func TestController_Fork(t *testing.T) { }, } + reg := NewRegistry() c, err := New(ctx, Config{ + Registry: reg, EventLogBuilder: func() (executor.EventLog, error) { return log, nil }, @@ -125,7 +127,9 @@ func TestController_Fork_SrcSeqNotFound(t *testing.T) { }, } + reg := NewRegistry() c, err := New(ctx, Config{ + Registry: reg, EventLogBuilder: func() (executor.EventLog, error) { return log, nil }, diff --git a/internal/controller2/registry.go b/internal/controller2/registry.go index 373c2af..1e64fa4 100644 --- a/internal/controller2/registry.go +++ b/internal/controller2/registry.go @@ -24,6 +24,7 @@ import ( "github.com/google/ax/internal/config" "github.com/google/ax/internal/experimental/a2abridge" expagent "github.com/google/ax/internal/experimental/agent" + "github.com/google/ax/internal/harness" ) // Registry manages a collection of local and remote agents. @@ -32,6 +33,7 @@ type Registry struct { mu sync.RWMutex agents map[string]agent.Agent agentInfo map[string]*agent.AgentInfo + harnesses map[string]harness.Harness } // NewRegistry creates a new agent registry. @@ -39,6 +41,7 @@ func NewRegistry() *Registry { return &Registry{ agents: make(map[string]agent.Agent), agentInfo: make(map[string]*agent.AgentInfo), + harnesses: make(map[string]harness.Harness), } } @@ -237,3 +240,20 @@ func (r *Registry) Close() error { return firstErr } +// RegisterHarness registers a harness. +func (r *Registry) RegisterHarness(id string, h harness.Harness) { + r.mu.Lock() + defer r.mu.Unlock() + r.harnesses[id] = h +} + +// GetHarness retrieves a harness by ID. +func (r *Registry) GetHarness(id string) (harness.Harness, error) { + r.mu.RLock() + defer r.mu.RUnlock() + h, ok := r.harnesses[id] + if !ok { + return nil, fmt.Errorf("harness %s not found", id) + } + return h, nil +} diff --git a/internal/harness/antigravity.go b/internal/harness/antigravity.go new file mode 100644 index 0000000..5bb2802 --- /dev/null +++ b/internal/harness/antigravity.go @@ -0,0 +1,148 @@ +// 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 harness + +import ( + "context" + "fmt" + "os/exec" + "strings" + "sync" + + "github.com/google/ax/proto" + "github.com/google/uuid" +) + +// AntigravityHarness implements the Harness interface by running the +// Antigravity Python agent as a subprocess. +type AntigravityHarness struct { + scriptPath string +} + +// NewAntigravityHarness creates a new AntigravityHarness with a configurable script path. +func NewAntigravityHarness(scriptPath string) *AntigravityHarness { + if scriptPath == "" { + scriptPath = "examples/antigravity_agent/agent.py" + } + return &AntigravityHarness{ + scriptPath: scriptPath, + } +} + + +// Start implements Harness.Start. +func (h *AntigravityHarness) Start(ctx context.Context, conversationID string) (Execution, error) { + return &antigravityExecution{ + harness: h, + conversationID: conversationID, + id: uuid.NewString(), + }, nil +} + +// antigravityExecution implements the Execution interface. +type antigravityExecution struct { + harness *AntigravityHarness + conversationID string + id string + + mu sync.Mutex + queued []*proto.Message + closed bool +} + +// ID implements Execution.ID. +func (e *antigravityExecution) ID() string { + return e.id +} + +// Queue implements Execution.Queue. +func (e *antigravityExecution) Queue(ctx context.Context, msg ...*proto.Message) error { + e.mu.Lock() + defer e.mu.Unlock() + if e.closed { + return fmt.Errorf("execution is closed") + } + e.queued = append(e.queued, msg...) + return nil +} + +// Run implements Execution.Run. +// It executes the Python agent as a subprocess, passing the last user message as an argument. +func (e *antigravityExecution) Run(ctx context.Context, handler Handler) error { + e.mu.Lock() + inputs := e.queued + e.queued = nil + e.mu.Unlock() + + // Find the last user message to pass to the agent + var prompt string + for i := len(inputs) - 1; i >= 0; i-- { + msg := inputs[i] + if msg.Role == "user" { + if textContent := msg.GetContent().GetText().GetText(); textContent != "" { + prompt = textContent + break + } + } + } + + // TODO: As a next step, we should implement this as a gRPC server to avoid subprocess overhead. + + // Prepare the command + args := []string{e.harness.scriptPath} + if prompt != "" { + args = append(args, prompt) + } + + cmd := exec.CommandContext(ctx, "python3", args...) + + // Capture stdout and stderr + var stdout, stderr strings.Builder + cmd.Stdout = &stdout + cmd.Stderr = &stderr + + // Run the command + if err := cmd.Run(); err != nil { + return fmt.Errorf("failed to run antigravity agent (stderr: %s): %w", stderr.String(), err) + } + + output := strings.TrimSpace(stdout.String()) + + // Send the output back to the handler + msg := &proto.Message{ + Role: "assistant", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{ + Text: output, + }, + }, + }, + } + + if err := handler.OnMessage(ctx, e.id, msg); err != nil { + return fmt.Errorf("failed to send message to handler: %w", err) + } + + return handler.OnComplete(ctx, e.id) +} + +// Close implements Execution.Close. +func (e *antigravityExecution) Close(ctx context.Context) error { + e.mu.Lock() + defer e.mu.Unlock() + e.closed = true + return nil +}