From 7b7d8bdbee17094ed96204f951e8b1482a9febcf Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Fri, 22 May 2026 11:46:26 -0700 Subject: [PATCH 01/16] 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 +} From cf7b3539c522a32578583bb01da9363a2c7c3058 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 10:08:11 -0700 Subject: [PATCH 02/16] AX V2 Harness Redesign: Implement Antigravity WebSocket Streaming Move the Google Antigravity SDK integration from a one-off subprocess to an always-on, stateful WebSocket streaming harness protocol. - Python WebSocket Server (harness_server.py): Exposes a /ws endpoint wrapping google-antigravity SDK, hydrates event log history, and streams thought/text delta frames over WebSockets. - Go Harness Client (antigravity.go): Refactored client to connect over gorilla/websocket and stream output callbacks (OnMessage/OnComplete). - Unit Tests: Implemented robust unit test suites in both Go and Python. - Verification: Created hack/run-antigravity-streaming.sh for E2E verification. TAG=agy CONV=ae7297af-6274-4477-81b3-a5eb53abb0e1 --- e2e.go | 27 ++- examples/antigravity_agent/agent.py | 11 +- hack/run-antigravity-streaming.sh | 74 ++++++ internal/harness/antigravity.go | 148 +++++++----- internal/harness/antigravity_test.go | 176 +++++++++++++++ python/antigravity/__init__.py | 13 ++ python/antigravity/harness_server.py | 184 +++++++++++++++ python/antigravity/harness_server_test.py | 131 +++++++++++ python/proto/ax_pb2.py | 102 ++++----- python/proto/ax_pb2_grpc.py | 264 ++++++++++------------ python/proto/content_pb2.py | 14 -- python/proto/content_pb2_grpc.py | 14 -- 12 files changed, 852 insertions(+), 306 deletions(-) create mode 100755 hack/run-antigravity-streaming.sh create mode 100644 internal/harness/antigravity_test.go create mode 100644 python/antigravity/__init__.py create mode 100644 python/antigravity/harness_server.py create mode 100644 python/antigravity/harness_server_test.py diff --git a/e2e.go b/e2e.go index 68727b6..36f85f3 100644 --- a/e2e.go +++ b/e2e.go @@ -20,7 +20,8 @@ import ( "context" "fmt" "os" - "os/exec" + + "github.com/gorilla/websocket" "github.com/google/ax/internal/controller/executor" "github.com/google/ax/internal/controller/executor/executortest" @@ -54,14 +55,11 @@ func main() { // 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 { + 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) + badHarness = harness.NewAntigravityHarness("ws://localhost:50054/ws") } reg.RegisterHarness("antigravity", badHarness) }) @@ -75,17 +73,18 @@ func main() { 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 + // Check if Python WebSocket server is active, otherwise fallback 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) + address := "ws://localhost:50053/ws" + dialer := websocket.DefaultDialer + conn, _, err := dialer.Dial(address, nil) + if err != nil { + fmt.Printf("WARNING: Antigravity harness server not active at %s, falling back to test harness: %v\n", address, err) realHarness = harnesstest.New() } else { - realHarness = harness.NewAntigravityHarness(scriptPath) + conn.Close() + fmt.Printf("Connected to Antigravity harness server at %s\n", address) + realHarness = harness.NewAntigravityHarness(address) } reg.RegisterHarness("antigravity", realHarness) }) diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py index fa8603a..b9bbf6b 100644 --- a/examples/antigravity_agent/agent.py +++ b/examples/antigravity_agent/agent.py @@ -16,10 +16,15 @@ import sys from google.antigravity import Agent, LocalAgentConfig +# Expose agent_config globally for harness_server.py loading +agent_config = LocalAgentConfig( + system_instructions="You are a helpful assistant powered by Google Antigravity." +) + 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: + # Initialize the agent session using the global config. + # It automatically picks up GEMINI_API_KEY from the environment. + async with Agent(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()) diff --git a/hack/run-antigravity-streaming.sh b/hack/run-antigravity-streaming.sh new file mode 100755 index 0000000..fc0da11 --- /dev/null +++ b/hack/run-antigravity-streaming.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# 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. + +set -e + +# Check if GEMINI_API_KEY is set +if [ -z "$GEMINI_API_KEY" ]; then + echo "ERROR: GEMINI_API_KEY environment variable is not set." + echo "Please set it using: export GEMINI_API_KEY=\"your-key\"" + exit 1 +fi + +PORT=50053 +ADDRESS="localhost:$PORT" +AGENT_FILE="examples/antigravity_agent/agent.py" + +# 1. Start Python WebSocket server in the background +echo "Starting Python WebSocket Harness Server on port $PORT..." +PYTHONPATH=python:. .venv/bin/python -m python.antigravity.harness_server --agent_file "$AGENT_FILE" --port "$PORT" > /tmp/antigravity_harness.log 2>&1 & +SERVER_PID=$! + +# Register trap to ensure server is killed on script exit +cleanup() { + echo "Cleaning up: killing Python server (PID: $SERVER_PID)..." + kill "$SERVER_PID" || true + wait "$SERVER_PID" 2>/dev/null || true + echo "Cleanup complete!" +} +trap cleanup EXIT + +# 2. Wait for the Python server to be healthy +echo "Waiting for Python server to become healthy..." +MAX_ATTEMPTS=30 +ATTEMPT=1 +HEALTHY=false + +while [ $ATTEMPT -le $MAX_ATTEMPTS ]; do + # We can check if the port is open using nc (netcat) + if nc -z localhost "$PORT"; then + HEALTHY=true + break + fi + sleep 0.2 + ATTEMPT=$((ATTEMPT + 1)) +done + +if [ "$HEALTHY" = false ]; then + echo "ERROR: Python server failed to start within 6 seconds." + echo "Server logs (/tmp/antigravity_harness.log):" + cat /tmp/antigravity_harness.log + exit 1 +fi +echo "Python server is active!" + +# 3. Build and run the Go E2E V2 demonstration +echo "Building e2e..." +/opt/homebrew/bin/go build -o bin/e2e e2e.go + +echo "Executing E2E Demo with Antigravity WebSocket Harness..." +bin/e2e + +echo "Success!" diff --git a/internal/harness/antigravity.go b/internal/harness/antigravity.go index 5bb2802..e6e8d5a 100644 --- a/internal/harness/antigravity.go +++ b/internal/harness/antigravity.go @@ -16,32 +16,33 @@ package harness import ( "context" + "encoding/json" "fmt" - "os/exec" - "strings" "sync" + "github.com/gorilla/websocket" + "google.golang.org/protobuf/encoding/protojson" + "github.com/google/ax/proto" "github.com/google/uuid" ) -// AntigravityHarness implements the Harness interface by running the -// Antigravity Python agent as a subprocess. +// AntigravityHarness implements the Harness interface by connecting to the +// Antigravity Python agent server over WebSockets. type AntigravityHarness struct { - scriptPath string + address string } -// NewAntigravityHarness creates a new AntigravityHarness with a configurable script path. -func NewAntigravityHarness(scriptPath string) *AntigravityHarness { - if scriptPath == "" { - scriptPath = "examples/antigravity_agent/agent.py" +// NewAntigravityHarness creates a new AntigravityHarness with a configurable address. +func NewAntigravityHarness(address string) *AntigravityHarness { + if address == "" { + address = "ws://localhost:50053/ws" } return &AntigravityHarness{ - scriptPath: scriptPath, + address: address, } } - // Start implements Harness.Start. func (h *AntigravityHarness) Start(ctx context.Context, conversationID string) (Execution, error) { return &antigravityExecution{ @@ -79,64 +80,107 @@ func (e *antigravityExecution) Queue(ctx context.Context, msg ...*proto.Message) } // Run implements Execution.Run. -// It executes the Python agent as a subprocess, passing the last user message as an argument. +// It connects to the Python server over WebSockets, sends the start payload containing history, +// and streams responses back to the handler. 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 - } + // 1. Establish WebSocket connection + dialer := websocket.DefaultDialer + conn, _, err := dialer.DialContext(ctx, e.harness.address, nil) + if err != nil { + return fmt.Errorf("failed to dial antigravity harness websocket at %s: %w", e.harness.address, err) + } + defer conn.Close() + + // 2. Serialize inputs using protojson to match Python Parse() requirements + var serializedMessages []json.RawMessage + for _, msg := range inputs { + bytes, err := protojson.Marshal(msg) + if err != nil { + return fmt.Errorf("failed to marshal message to JSON: %w", err) } + serializedMessages = append(serializedMessages, json.RawMessage(bytes)) } - // 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) + // 3. Construct and send start payload + startPayload := map[string]any{ + "conversation_id": e.conversationID, + "exec_id": e.id, + "messages": serializedMessages, + } + payloadBytes, err := json.Marshal(startPayload) + if err != nil { + return fmt.Errorf("failed to marshal start payload: %w", err) } - cmd := exec.CommandContext(ctx, "python3", args...) - - // Capture stdout and stderr - var stdout, stderr strings.Builder - cmd.Stdout = &stdout - cmd.Stderr = &stderr + err = conn.WriteMessage(websocket.TextMessage, payloadBytes) + if err != nil { + return fmt.Errorf("failed to send start payload over WebSocket: %w", err) + } - // Run the command - if err := cmd.Run(); err != nil { - return fmt.Errorf("failed to run antigravity agent (stderr: %s): %w", stderr.String(), err) + // 4. Stream responses from WebSocket + type WSResponse struct { + Type string `json:"type"` + Content string `json:"content"` + Error string `json:"error"` } - output := strings.TrimSpace(stdout.String()) + for { + _, message, err := conn.ReadMessage() + if err != nil { + return fmt.Errorf("failed to read message from WebSocket: %w", err) + } - // Send the output back to the handler - msg := &proto.Message{ - Role: "assistant", - Content: &proto.Content{ - Type: &proto.Content_Text{ - Text: &proto.TextContent{ - Text: output, - }, - }, - }, - } + var resp WSResponse + if err := json.Unmarshal(message, &resp); err != nil { + return fmt.Errorf("failed to unmarshal WebSocket response: %w", err) + } - if err := handler.OnMessage(ctx, e.id, msg); err != nil { - return fmt.Errorf("failed to send message to handler: %w", err) + switch resp.Type { + case "text": + msg := &proto.Message{ + Role: "assistant", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: resp.Content}, + }, + }, + } + if err := handler.OnMessage(ctx, e.id, msg); err != nil { + return fmt.Errorf("failed to send message to handler: %w", err) + } + case "thought": + msg := &proto.Message{ + Role: "model", + Content: &proto.Content{ + Type: &proto.Content_Thought{ + Thought: &proto.ThoughtContent{ + Summary: []*proto.ThoughtSummaryContent{ + { + Type: &proto.ThoughtSummaryContent_Text{ + Text: &proto.TextContent{Text: resp.Content}, + }, + }, + }, + }, + }, + }, + } + if err := handler.OnMessage(ctx, e.id, msg); err != nil { + return fmt.Errorf("failed to send thought to handler: %w", err) + } + case "complete": + return handler.OnComplete(ctx, e.id) + case "error": + return fmt.Errorf("antigravity harness server error: %s", resp.Error) + default: + return fmt.Errorf("unknown response type from WebSocket: %q", resp.Type) + } } - - return handler.OnComplete(ctx, e.id) } // Close implements Execution.Close. diff --git a/internal/harness/antigravity_test.go b/internal/harness/antigravity_test.go new file mode 100644 index 0000000..9c6243b --- /dev/null +++ b/internal/harness/antigravity_test.go @@ -0,0 +1,176 @@ +// 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" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/gorilla/websocket" + + "github.com/google/ax/proto" +) + +type mockHandler struct { + mu sync.Mutex + messages []*proto.Message + complete bool + err error +} + +func (h *mockHandler) OnMessage(ctx context.Context, execID string, msg *proto.Message) error { + h.mu.Lock() + defer h.mu.Unlock() + h.messages = append(h.messages, msg) + return h.err +} + +func (h *mockHandler) OnComplete(ctx context.Context, execID string) error { + h.mu.Lock() + defer h.mu.Unlock() + h.complete = true + return nil +} + +func TestAntigravityHarness_Run_Success(t *testing.T) { + upgrader := websocket.Upgrader{} + + // Spin up a local mock WebSocket server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("failed to upgrade connection: %v", err) + return + } + defer conn.Close() + + // 1. Read initial start message + _, msg, err := conn.ReadMessage() + if err != nil { + t.Errorf("failed to read start message: %v", err) + return + } + + var payload map[string]any + if err := json.Unmarshal(msg, &payload); err != nil { + t.Errorf("failed to unmarshal start payload: %v", err) + return + } + + if payload["conversation_id"] != "conv-test" { + t.Errorf("expected conversation ID 'conv-test', got %v", payload["conversation_id"]) + } + + // 2. Stream response chunks + chunks := []map[string]string{ + {"type": "text", "content": "Hello "}, + {"type": "text", "content": "world!"}, + {"type": "complete"}, + } + + for _, chunk := range chunks { + bytes, _ := json.Marshal(chunk) + if err := conn.WriteMessage(websocket.TextMessage, bytes); err != nil { + t.Errorf("failed to write chunk: %v", err) + return + } + } + })) + defer server.Close() + + // Convert http:// to ws:// + wsURL := strings.Replace(server.URL, "http://", "ws://", 1) + "/ws" + + harnessClient := NewAntigravityHarness(wsURL) + exec, err := harnessClient.Start(context.Background(), "conv-test") + if err != nil { + t.Fatalf("failed to start execution: %v", err) + } + defer exec.Close(context.Background()) + + msg := &proto.Message{ + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{Text: &proto.TextContent{Text: "Hi"}}, + }, + } + if err := exec.Queue(context.Background(), msg); err != nil { + t.Fatalf("failed to queue message: %v", err) + } + + handler := &mockHandler{} + err = exec.Run(context.Background(), handler) + if err != nil { + t.Fatalf("Run failed: %v", err) + } + + handler.mu.Lock() + defer handler.mu.Unlock() + + if !handler.complete { + t.Error("expected OnComplete to be called") + } + if len(handler.messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(handler.messages)) + } + if handler.messages[0].GetContent().GetText().GetText() != "Hello " { + t.Errorf("expected 'Hello ', got %q", handler.messages[0].GetContent().GetText().GetText()) + } + if handler.messages[1].GetContent().GetText().GetText() != "world!" { + t.Errorf("expected 'world!', got %q", handler.messages[1].GetContent().GetText().GetText()) + } +} + +func TestAntigravityHarness_Run_ErrorFrame(t *testing.T) { + upgrader := websocket.Upgrader{} + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + + _, _, _ = conn.ReadMessage() + + errFrame := map[string]string{ + "type": "error", + "error": "internal model crash", + } + bytes, _ := json.Marshal(errFrame) + conn.WriteMessage(websocket.TextMessage, bytes) + })) + defer server.Close() + + wsURL := strings.Replace(server.URL, "http://", "ws://", 1) + "/ws" + + harnessClient := NewAntigravityHarness(wsURL) + exec, _ := harnessClient.Start(context.Background(), "conv-test") + defer exec.Close(context.Background()) + + handler := &mockHandler{} + err := exec.Run(context.Background(), handler) + if err == nil { + t.Fatal("expected error from Run(), got nil") + } + if !strings.Contains(err.Error(), "antigravity harness server error: internal model crash") { + t.Errorf("unexpected error message: %v", err) + } +} diff --git a/python/antigravity/__init__.py b/python/antigravity/__init__.py new file mode 100644 index 0000000..58d482e --- /dev/null +++ b/python/antigravity/__init__.py @@ -0,0 +1,13 @@ +# 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. diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py new file mode 100644 index 0000000..e4eae33 --- /dev/null +++ b/python/antigravity/harness_server.py @@ -0,0 +1,184 @@ +# 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 +# +# https://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 argparse +import asyncio +import importlib.util +import json +import logging +import os +import sys +import uuid +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +import uvicorn + +from google.protobuf.json_format import Parse +from proto import ax_pb2 +from google.antigravity import Agent, AgentConfig +from google.antigravity.types import Step, StepType, StepSource, StepTarget, StepStatus, Text, Thought + +app = FastAPI() + +# Global placeholder for loaded agent config +loaded_config: AgentConfig | None = None + +def load_agent_config(agent_file: str) -> AgentConfig: + print(f"Loading agent config from {agent_file}...") + spec = importlib.util.spec_from_file_location("agent_module", agent_file) + if spec is None or spec.loader is None: + raise FileNotFoundError(f"Could not find or load agent file: {agent_file}") + agent_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(agent_module) + + config = getattr(agent_module, "agent_config", None) + if not config: + raise ValueError(f"No 'agent_config' found in {agent_file}") + print("Agent config loaded successfully.") + return config + +def hydrate_ax_history_to_steps(historical_messages) -> list[Step]: + steps = [] + for i, msg in enumerate(historical_messages): + source = StepSource.UNKNOWN + target = StepTarget.UNSPECIFIED + step_type = StepType.TEXT_RESPONSE + content = "" + thinking = "" + + # Determine source and target based on role + if msg.role == "user": + source = StepSource.USER + target = StepTarget.ENVIRONMENT + elif msg.role in ("assistant", "model"): + source = StepSource.MODEL + target = StepTarget.USER + + # Extract content/thinking + active_type = msg.content.WhichOneof('type') + if active_type == 'text': + content = msg.content.text.text + elif active_type == 'thought': + step_type = StepType.TEXT_RESPONSE + if msg.content.thought.summary: + texts = [] + for s in msg.content.thought.summary: + if s.WhichOneof('type') == 'text': + texts.append(s.text.text) + thinking = "".join(texts) + + step = Step( + id=f"hist-{i}", + step_index=i, + type=step_type, + source=source, + target=target, + status=StepStatus.DONE, + content=content, + thinking=thinking, + is_complete_response=True + ) + steps.append(step) + return steps + +@app.websocket("/ws") +async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + print("[WS] Connection accepted.") + try: + # 1. Receive the start message + data = await websocket.receive_text() + payload = json.loads(data) + + conversation_id = payload.get("conversation_id") + exec_id = payload.get("exec_id") + raw_messages = payload.get("messages", []) + + print(f"[WS] Starting turn. conv_id={conversation_id}, exec_id={exec_id}, messages_count={len(raw_messages)}") + + # Deserialize AX protobuf messages + ax_messages = [] + for raw_msg in raw_messages: + msg_str = json.dumps(raw_msg) + ax_msg = Parse(msg_str, ax_pb2.Message()) + ax_messages.append(ax_msg) + + if not ax_messages: + raise ValueError("No messages found in start payload") + + historical_messages = ax_messages[:-1] + latest_message = ax_messages[-1] + + # Only support text queries for now in latest_message + if latest_message.content.WhichOneof('type') != 'text': + raise ValueError("Latest message must contain text content") + latest_query_text = latest_message.content.text.text + + # 2. Initialize the Antigravity Agent session + global loaded_config + if not loaded_config: + raise RuntimeError("Agent config is not loaded on the server") + + async with Agent(loaded_config) as agent: + conversation = agent.conversation + + # Hydrate history + print(f"[WS] Hydrating {len(historical_messages)} historical messages...") + history_steps = hydrate_ax_history_to_steps(historical_messages) + conversation._steps.extend(history_steps) + + # Run the turn with streaming + print(f"[WS] Running chat query: {latest_query_text}") + response = await conversation.chat(latest_query_text) + + async for chunk in response.chunks: + if isinstance(chunk, Text): + await websocket.send_json({"type": "text", "content": chunk.text}) + elif isinstance(chunk, Thought): + await websocket.send_json({"type": "thought", "content": chunk.text}) + + # Send complete frame + await websocket.send_json({"type": "complete"}) + print("[WS] Turn completed successfully.") + + except WebSocketDisconnect: + print("[WS] Client disconnected.") + except Exception as e: + logging.exception("Error in WebSocket turn handler") + try: + await websocket.send_json({"type": "error", "error": str(e)}) + except Exception: + pass + finally: + await websocket.close() + +def main(): + parser = argparse.ArgumentParser(description="Antigravity WebSocket Harness Server") + parser.add_argument("--agent_file", default="examples/antigravity_agent/agent.py", help="Path to the agent config file") + parser.add_argument("--port", type=int, default=50053, help="Port to bind the server to") + parser.add_argument("--host", default="localhost", help="Host to bind the server to") + args = parser.parse_args() + + # Load the agent config globally + global loaded_config + try: + loaded_config = load_agent_config(args.agent_file) + except Exception as e: + print(f"ERROR: Failed to load agent config: {e}", file=sys.stderr) + sys.exit(1) + + print(f"Starting WebSocket server on {args.host}:{args.port}...") + uvicorn.run(app, host=args.host, port=args.port) + +if __name__ == "__main__": + main() diff --git a/python/antigravity/harness_server_test.py b/python/antigravity/harness_server_test.py new file mode 100644 index 0000000..290be33 --- /dev/null +++ b/python/antigravity/harness_server_test.py @@ -0,0 +1,131 @@ +# 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 +# +# https://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 pytest +from fastapi.testclient import TestClient +from unittest.mock import AsyncMock, MagicMock, patch +import json + +from python.antigravity.harness_server import app, hydrate_ax_history_to_steps +from google.antigravity.types import Step, StepType, StepSource, StepTarget, StepStatus, Text, Thought +from proto import ax_pb2 + +client = TestClient(app) + +def test_hydrate_ax_history_to_steps(): + # Create mock AX Message protobuf objects + msg = ax_pb2.Message() + msg.role = "user" + msg.content.text.text = "Hi" + + steps = hydrate_ax_history_to_steps([msg]) + + assert len(steps) == 1 + assert steps[0].source == StepSource.USER + assert steps[0].content == "Hi" + assert steps[0].is_complete_response is True + +@patch("python.antigravity.harness_server.Agent") +def test_websocket_endpoint_success(mock_agent_class): + # 1. Setup mocks for Agent, Conversation, and ChatResponse + mock_agent = MagicMock() + mock_agent_class.return_value = mock_agent + + # Mock context manager methods + mock_agent.__aenter__ = AsyncMock(return_value=mock_agent) + mock_agent.__aexit__ = AsyncMock(return_value=None) + + mock_conversation = MagicMock() + mock_agent.conversation = mock_conversation + mock_conversation._steps = [] + + # Mock response stream chunks + async def mock_chunks(): + yield Text(step_index=0, text="Hello ") + yield Text(step_index=1, text="world!") + + mock_chat_response = MagicMock() + mock_chat_response.chunks = mock_chunks() + mock_conversation.chat = AsyncMock(return_value=mock_chat_response) + + # Load a dummy config globally to pass server validation + import python.antigravity.harness_server as server + server.loaded_config = MagicMock() + + # 2. Build start payload + start_payload = { + "conversation_id": "conv-123", + "exec_id": "exec-456", + "messages": [ + # Raw protobuf JSON message + { + "role": "user", + "content": { + "text": {"text": "Hi"} + } + } + ] + } + + # 3. Run WebSocket test client + with client.websocket_connect("/ws") as websocket: + # Send start payload + websocket.send_text(json.dumps(start_payload)) + + # Receive streamed text chunks + resp1 = websocket.receive_json() + assert resp1["type"] == "text" + assert resp1["content"] == "Hello " + + resp2 = websocket.receive_json() + assert resp2["type"] == "text" + assert resp2["content"] == "world!" + + # Receive complete + resp3 = websocket.receive_json() + assert resp3["type"] == "complete" + + # Verify mocks + mock_conversation.chat.assert_called_once_with("Hi") + +@patch("python.antigravity.harness_server.Agent") +def test_websocket_endpoint_error(mock_agent_class): + mock_agent = MagicMock() + mock_agent_class.return_value = mock_agent + mock_agent.__aenter__ = AsyncMock(return_value=mock_agent) + mock_agent.__aexit__ = AsyncMock(return_value=None) + + mock_conversation = MagicMock() + mock_agent.conversation = mock_conversation + mock_conversation._steps = [] + + # Mock chat to throw an exception + mock_conversation.chat = AsyncMock(side_effect=RuntimeError("Gemini connection timeout")) + + import python.antigravity.harness_server as server + server.loaded_config = MagicMock() + + start_payload = { + "conversation_id": "conv-123", + "exec_id": "exec-456", + "messages": [{"role": "user", "content": {"text": {"text": "Hi"}}}] + } + + with client.websocket_connect("/ws") as websocket: + websocket.send_text(json.dumps(start_payload)) + + # Expect error frame + resp = websocket.receive_json() + assert resp["type"] == "error" + assert "Gemini connection timeout" in resp["error"] diff --git a/python/proto/ax_pb2.py b/python/proto/ax_pb2.py index 988e4ae..37a01a1 100644 --- a/python/proto/ax_pb2.py +++ b/python/proto/ax_pb2.py @@ -1,17 +1,3 @@ -# 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. - # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE @@ -37,10 +23,10 @@ from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 -import content_pb2 as proto_dot_content__pb2 +from proto import content_pb2 as proto_dot_content__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0eproto/ax.proto\x12\x02\x61x\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x13proto/content.proto\"S\n\nAgentStart\x12\x10\n\x08\x61gent_id\x18\x01 \x01(\t\x12\x14\n\x0c\x61gent_config\x18\x02 \x01(\x0c\x12\x1d\n\x08messages\x18\x03 \x03(\x0b\x32\x0b.ax.Message\"-\n\x0c\x41gentOutputs\x12\x1d\n\x08messages\x18\x01 \x03(\x0b\x32\x0b.ax.Message\"\n\n\x08\x41gentEnd\"\xa3\x01\n\x0c\x41gentMessage\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0f\n\x07\x65xec_id\x18\x02 \x01(\t\x12\x1f\n\x05start\x18\x03 \x01(\x0b\x32\x0e.ax.AgentStartH\x00\x12#\n\x07outputs\x18\x04 \x01(\x0b\x32\x10.ax.AgentOutputsH\x00\x12\x1b\n\x03\x65nd\x18\x05 \x01(\x0b\x32\x0c.ax.AgentEndH\x00\x42\x06\n\x04type\"L\n\x07Message\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x1c\n\x07\x63ontent\x18\x02 \x01(\x0b\x32\x0b.ax.Content\x12\x15\n\rinternal_only\x18\x03 \x01(\x08\"\x83\x01\n\x11\x43onversationEvent\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0b\n\x03seq\x18\x02 \x01(\x05\x12\x0f\n\x07\x65xec_id\x18\x03 \x01(\t\x12\x1d\n\x08messages\x18\x04 \x03(\x0b\x32\x0b.ax.Message\x12\x18\n\x05state\x18\x05 \x01(\x0e\x32\t.ax.State\"\xcd\x01\n\x0e\x45xecutionEvent\x12\x0f\n\x07\x65xec_id\x18\x01 \x01(\t\x12\x10\n\x08\x61gent_id\x18\x02 \x01(\t\x12\x14\n\x0c\x61gent_config\x18\x03 \x01(\x0c\x12\x1b\n\x06inputs\x18\x04 \x03(\x0b\x32\x0b.ax.Message\x12\x1c\n\x07outputs\x18\x05 \x03(\x0b\x32\x0b.ax.Message\x12\x18\n\x05state\x18\x06 \x01(\x0e\x32\t.ax.State\x12-\n\ttimestamp\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"\x14\n\x12HealthCheckRequest\"7\n\x13HealthCheckResponse\x12\x0f\n\x07healthy\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"}\n\x0b\x45xecRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x1b\n\x06inputs\x18\x02 \x03(\x0b\x32\x0b.ax.Message\x12\x10\n\x08last_seq\x18\x03 \x01(\x05\x12\x10\n\x08\x61gent_id\x18\x04 \x01(\t\x12\x14\n\x0c\x61gent_config\x18\x05 \x01(\x0c\"9\n\x0c\x45xecResponse\x12\x1c\n\x07outputs\x18\x01 \x03(\x0b\x32\x0b.ax.Message\x12\x0b\n\x03seq\x18\x02 \x01(\x05\"$\n\x11RemoteAgentConfig\x12\x0f\n\x07\x61\x64\x64ress\x18\x01 \x01(\t\"\xe9\x01\n\x14RegisterAgentRequest\x12\x10\n\x08\x61gent_id\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x03 \x01(\t\x12\x38\n\x08metadata\x18\x04 \x03(\x0b\x32&.ax.RegisterAgentRequest.MetadataEntry\x12\'\n\x06remote\x18\x05 \x01(\x0b\x32\x15.ax.RemoteAgentConfigH\x00\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x42\x08\n\x06\x63onfig\"\x17\n\x15RegisterAgentResponse\"&\n\x0bListRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\"5\n\x0cListResponse\x12%\n\x06\x65vents\x18\x01 \x03(\x0b\x32\x15.ax.ConversationEvent\"(\n\rDeleteRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\"\x10\n\x0e\x44\x65leteResponse\"Y\n\x0b\x46orkRequest\x12\x1b\n\x13src_conversation_id\x18\x01 \x01(\t\x12\x0f\n\x07src_seq\x18\x02 \x01(\x05\x12\x1c\n\x14\x64\x65st_conversation_id\x18\x03 \x01(\t\"\'\n\x0c\x46orkResponse\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t*X\n\x05State\x12\x15\n\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n\rSTATE_PENDING\x10\x01\x12\x10\n\x0cSTATE_FAILED\x10\x02\x12\x13\n\x0fSTATE_COMPLETED\x10\x03\x32\x81\x01\n\x0c\x41gentService\x12\x31\n\x07\x43onnect\x12\x10.ax.AgentMessage\x1a\x10.ax.AgentMessage(\x01\x30\x01\x12>\n\x0bHealthCheck\x12\x16.ax.HealthCheckRequest\x1a\x17.ax.HealthCheckResponse2\x86\x01\n\x11\x43ontrollerService\x12+\n\x04\x45xec\x12\x0f.ax.ExecRequest\x1a\x10.ax.ExecResponse0\x01\x12\x44\n\rRegisterAgent\x12\x18.ax.RegisterAgentRequest\x1a\x19.ax.RegisterAgentResponse2\x98\x01\n\x0f\x45ventLogService\x12)\n\x04List\x12\x0f.ax.ListRequest\x1a\x10.ax.ListResponse\x12/\n\x06\x44\x65lete\x12\x11.ax.DeleteRequest\x1a\x12.ax.DeleteResponse\x12)\n\x04\x46ork\x12\x0f.ax.ForkRequest\x1a\x10.ax.ForkResponseB\x1cZ\x1agithub.com/google/ax/protob\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0eproto/ax.proto\x12\x02\x61x\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x13proto/content.proto\"S\n\nAgentStart\x12\x10\n\x08\x61gent_id\x18\x01 \x01(\t\x12\x14\n\x0c\x61gent_config\x18\x02 \x01(\x0c\x12\x1d\n\x08messages\x18\x03 \x03(\x0b\x32\x0b.ax.Message\"-\n\x0c\x41gentOutputs\x12\x1d\n\x08messages\x18\x01 \x03(\x0b\x32\x0b.ax.Message\"\n\n\x08\x41gentEnd\"W\n\x0c\x41gentRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0f\n\x07\x65xec_id\x18\x02 \x01(\t\x12\x1d\n\x05start\x18\x03 \x01(\x0b\x32\x0e.ax.AgentStart\"\x83\x01\n\rAgentResponse\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0f\n\x07\x65xec_id\x18\x02 \x01(\t\x12#\n\x07outputs\x18\x03 \x01(\x0b\x32\x10.ax.AgentOutputsH\x00\x12\x1b\n\x03\x65nd\x18\x04 \x01(\x0b\x32\x0c.ax.AgentEndH\x00\x42\x06\n\x04type\"L\n\x07Message\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x1c\n\x07\x63ontent\x18\x02 \x01(\x0b\x32\x0b.ax.Content\x12\x15\n\rinternal_only\x18\x03 \x01(\x08\"\x83\x01\n\x11\x43onversationEvent\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0b\n\x03seq\x18\x02 \x01(\x05\x12\x0f\n\x07\x65xec_id\x18\x03 \x01(\t\x12\x1d\n\x08messages\x18\x04 \x03(\x0b\x32\x0b.ax.Message\x12\x18\n\x05state\x18\x05 \x01(\x0e\x32\t.ax.State\"\xcd\x01\n\x0e\x45xecutionEvent\x12\x0f\n\x07\x65xec_id\x18\x01 \x01(\t\x12\x10\n\x08\x61gent_id\x18\x02 \x01(\t\x12\x14\n\x0c\x61gent_config\x18\x03 \x01(\x0c\x12\x1b\n\x06inputs\x18\x04 \x03(\x0b\x32\x0b.ax.Message\x12\x1c\n\x07outputs\x18\x05 \x03(\x0b\x32\x0b.ax.Message\x12\x18\n\x05state\x18\x06 \x01(\x0e\x32\t.ax.State\x12-\n\ttimestamp\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"\x14\n\x12HealthCheckRequest\"7\n\x13HealthCheckResponse\x12\x0f\n\x07healthy\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"/\n\x0eHarnessMessage\x12\x1d\n\x08messages\x18\x01 \x03(\x0b\x32\x0b.ax.Message\"}\n\x0b\x45xecRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x1b\n\x06inputs\x18\x02 \x03(\x0b\x32\x0b.ax.Message\x12\x10\n\x08last_seq\x18\x03 \x01(\x05\x12\x10\n\x08\x61gent_id\x18\x04 \x01(\t\x12\x14\n\x0c\x61gent_config\x18\x05 \x01(\x0c\"9\n\x0c\x45xecResponse\x12\x1c\n\x07outputs\x18\x01 \x03(\x0b\x32\x0b.ax.Message\x12\x0b\n\x03seq\x18\x02 \x01(\x05\"4\n\x19\x44\x65leteConversationRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\"\x1c\n\x1a\x44\x65leteConversationResponse\"e\n\x17\x46orkConversationRequest\x12\x1b\n\x13src_conversation_id\x18\x01 \x01(\t\x12\x0f\n\x07src_seq\x18\x02 \x01(\x05\x12\x1c\n\x14\x64\x65st_conversation_id\x18\x03 \x01(\t\"3\n\x18\x46orkConversationResponse\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t*X\n\x05State\x12\x15\n\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n\rSTATE_PENDING\x10\x01\x12\x10\n\x0cSTATE_FAILED\x10\x02\x12\x13\n\x0fSTATE_COMPLETED\x10\x03\x32\x80\x01\n\x0c\x41gentService\x12\x30\n\x07\x43onnect\x12\x10.ax.AgentRequest\x1a\x11.ax.AgentResponse0\x01\x12>\n\x0bHealthCheck\x12\x16.ax.HealthCheckRequest\x1a\x17.ax.HealthCheckResponse2G\n\x0eHarnessService\x12\x35\n\x07\x43onnect\x12\x12.ax.HarnessMessage\x1a\x12.ax.HarnessMessage(\x01\x30\x01\x32@\n\x11\x43ontrollerService\x12+\n\x04\x45xec\x12\x0f.ax.ExecRequest\x1a\x10.ax.ExecResponse0\x01\x32\xb9\x01\n\x13\x43onversationService\x12S\n\x12\x44\x65leteConversation\x12\x1d.ax.DeleteConversationRequest\x1a\x1e.ax.DeleteConversationResponse\x12M\n\x10\x46orkConversation\x12\x1b.ax.ForkConversationRequest\x1a\x1c.ax.ForkConversationResponseB\x1cZ\x1agithub.com/google/ax/protob\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -48,56 +34,48 @@ if not _descriptor._USE_C_DESCRIPTORS: _globals['DESCRIPTOR']._loaded_options = None _globals['DESCRIPTOR']._serialized_options = b'Z\032github.com/google/ax/proto' - _globals['_REGISTERAGENTREQUEST_METADATAENTRY']._loaded_options = None - _globals['_REGISTERAGENTREQUEST_METADATAENTRY']._serialized_options = b'8\001' - _globals['_STATE']._serialized_start=1657 - _globals['_STATE']._serialized_end=1745 + _globals['_STATE']._serialized_start=1417 + _globals['_STATE']._serialized_end=1505 _globals['_AGENTSTART']._serialized_start=76 _globals['_AGENTSTART']._serialized_end=159 _globals['_AGENTOUTPUTS']._serialized_start=161 _globals['_AGENTOUTPUTS']._serialized_end=206 _globals['_AGENTEND']._serialized_start=208 _globals['_AGENTEND']._serialized_end=218 - _globals['_AGENTMESSAGE']._serialized_start=221 - _globals['_AGENTMESSAGE']._serialized_end=384 - _globals['_MESSAGE']._serialized_start=386 - _globals['_MESSAGE']._serialized_end=462 - _globals['_CONVERSATIONEVENT']._serialized_start=465 - _globals['_CONVERSATIONEVENT']._serialized_end=596 - _globals['_EXECUTIONEVENT']._serialized_start=599 - _globals['_EXECUTIONEVENT']._serialized_end=804 - _globals['_HEALTHCHECKREQUEST']._serialized_start=806 - _globals['_HEALTHCHECKREQUEST']._serialized_end=826 - _globals['_HEALTHCHECKRESPONSE']._serialized_start=828 - _globals['_HEALTHCHECKRESPONSE']._serialized_end=883 - _globals['_EXECREQUEST']._serialized_start=885 - _globals['_EXECREQUEST']._serialized_end=1010 - _globals['_EXECRESPONSE']._serialized_start=1012 - _globals['_EXECRESPONSE']._serialized_end=1069 - _globals['_REMOTEAGENTCONFIG']._serialized_start=1071 - _globals['_REMOTEAGENTCONFIG']._serialized_end=1107 - _globals['_REGISTERAGENTREQUEST']._serialized_start=1110 - _globals['_REGISTERAGENTREQUEST']._serialized_end=1343 - _globals['_REGISTERAGENTREQUEST_METADATAENTRY']._serialized_start=1286 - _globals['_REGISTERAGENTREQUEST_METADATAENTRY']._serialized_end=1333 - _globals['_REGISTERAGENTRESPONSE']._serialized_start=1345 - _globals['_REGISTERAGENTRESPONSE']._serialized_end=1368 - _globals['_LISTREQUEST']._serialized_start=1370 - _globals['_LISTREQUEST']._serialized_end=1408 - _globals['_LISTRESPONSE']._serialized_start=1410 - _globals['_LISTRESPONSE']._serialized_end=1463 - _globals['_DELETEREQUEST']._serialized_start=1465 - _globals['_DELETEREQUEST']._serialized_end=1505 - _globals['_DELETERESPONSE']._serialized_start=1507 - _globals['_DELETERESPONSE']._serialized_end=1523 - _globals['_FORKREQUEST']._serialized_start=1525 - _globals['_FORKREQUEST']._serialized_end=1614 - _globals['_FORKRESPONSE']._serialized_start=1616 - _globals['_FORKRESPONSE']._serialized_end=1655 - _globals['_AGENTSERVICE']._serialized_start=1748 - _globals['_AGENTSERVICE']._serialized_end=1877 - _globals['_CONTROLLERSERVICE']._serialized_start=1880 - _globals['_CONTROLLERSERVICE']._serialized_end=2014 - _globals['_EVENTLOGSERVICE']._serialized_start=2017 - _globals['_EVENTLOGSERVICE']._serialized_end=2169 + _globals['_AGENTREQUEST']._serialized_start=220 + _globals['_AGENTREQUEST']._serialized_end=307 + _globals['_AGENTRESPONSE']._serialized_start=310 + _globals['_AGENTRESPONSE']._serialized_end=441 + _globals['_MESSAGE']._serialized_start=443 + _globals['_MESSAGE']._serialized_end=519 + _globals['_CONVERSATIONEVENT']._serialized_start=522 + _globals['_CONVERSATIONEVENT']._serialized_end=653 + _globals['_EXECUTIONEVENT']._serialized_start=656 + _globals['_EXECUTIONEVENT']._serialized_end=861 + _globals['_HEALTHCHECKREQUEST']._serialized_start=863 + _globals['_HEALTHCHECKREQUEST']._serialized_end=883 + _globals['_HEALTHCHECKRESPONSE']._serialized_start=885 + _globals['_HEALTHCHECKRESPONSE']._serialized_end=940 + _globals['_HARNESSMESSAGE']._serialized_start=942 + _globals['_HARNESSMESSAGE']._serialized_end=989 + _globals['_EXECREQUEST']._serialized_start=991 + _globals['_EXECREQUEST']._serialized_end=1116 + _globals['_EXECRESPONSE']._serialized_start=1118 + _globals['_EXECRESPONSE']._serialized_end=1175 + _globals['_DELETECONVERSATIONREQUEST']._serialized_start=1177 + _globals['_DELETECONVERSATIONREQUEST']._serialized_end=1229 + _globals['_DELETECONVERSATIONRESPONSE']._serialized_start=1231 + _globals['_DELETECONVERSATIONRESPONSE']._serialized_end=1259 + _globals['_FORKCONVERSATIONREQUEST']._serialized_start=1261 + _globals['_FORKCONVERSATIONREQUEST']._serialized_end=1362 + _globals['_FORKCONVERSATIONRESPONSE']._serialized_start=1364 + _globals['_FORKCONVERSATIONRESPONSE']._serialized_end=1415 + _globals['_AGENTSERVICE']._serialized_start=1508 + _globals['_AGENTSERVICE']._serialized_end=1636 + _globals['_HARNESSSERVICE']._serialized_start=1638 + _globals['_HARNESSSERVICE']._serialized_end=1709 + _globals['_CONTROLLERSERVICE']._serialized_start=1711 + _globals['_CONTROLLERSERVICE']._serialized_end=1775 + _globals['_CONVERSATIONSERVICE']._serialized_start=1778 + _globals['_CONVERSATIONSERVICE']._serialized_end=1963 # @@protoc_insertion_point(module_scope) diff --git a/python/proto/ax_pb2_grpc.py b/python/proto/ax_pb2_grpc.py index 5463f20..e042b66 100644 --- a/python/proto/ax_pb2_grpc.py +++ b/python/proto/ax_pb2_grpc.py @@ -1,23 +1,9 @@ -# 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. - # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc import warnings -import ax_pb2 as proto_dot_ax__pb2 +from proto import ax_pb2 as proto_dot_ax__pb2 GRPC_GENERATED_VERSION = '1.80.0' GRPC_VERSION = grpc.__version__ @@ -55,10 +41,10 @@ def __init__(self, channel): Args: channel: A grpc.Channel. """ - self.Connect = channel.stream_stream( + self.Connect = channel.unary_stream( '/ax.AgentService/Connect', - request_serializer=proto_dot_ax__pb2.AgentMessage.SerializeToString, - response_deserializer=proto_dot_ax__pb2.AgentMessage.FromString, + request_serializer=proto_dot_ax__pb2.AgentRequest.SerializeToString, + response_deserializer=proto_dot_ax__pb2.AgentResponse.FromString, _registered_method=True) self.HealthCheck = channel.unary_unary( '/ax.AgentService/HealthCheck', @@ -77,7 +63,7 @@ class AgentServiceServicer(object): we are solidifying resumption on the wire. """ - def Connect(self, request_iterator, context): + def Connect(self, request, context): """Connect is used by agents to connect to the controller. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) @@ -94,10 +80,10 @@ def HealthCheck(self, request, context): def add_AgentServiceServicer_to_server(servicer, server): rpc_method_handlers = { - 'Connect': grpc.stream_stream_rpc_method_handler( + 'Connect': grpc.unary_stream_rpc_method_handler( servicer.Connect, - request_deserializer=proto_dot_ax__pb2.AgentMessage.FromString, - response_serializer=proto_dot_ax__pb2.AgentMessage.SerializeToString, + request_deserializer=proto_dot_ax__pb2.AgentRequest.FromString, + response_serializer=proto_dot_ax__pb2.AgentResponse.SerializeToString, ), 'HealthCheck': grpc.unary_unary_rpc_method_handler( servicer.HealthCheck, @@ -123,7 +109,7 @@ class AgentService(object): """ @staticmethod - def Connect(request_iterator, + def Connect(request, target, options=(), channel_credentials=None, @@ -133,12 +119,12 @@ def Connect(request_iterator, wait_for_ready=None, timeout=None, metadata=None): - return grpc.experimental.stream_stream( - request_iterator, + return grpc.experimental.unary_stream( + request, target, '/ax.AgentService/Connect', - proto_dot_ax__pb2.AgentMessage.SerializeToString, - proto_dot_ax__pb2.AgentMessage.FromString, + proto_dot_ax__pb2.AgentRequest.SerializeToString, + proto_dot_ax__pb2.AgentResponse.FromString, options, channel_credentials, insecure, @@ -177,6 +163,78 @@ def HealthCheck(request, _registered_method=True) +class HarnessServiceStub(object): + """Missing associated documentation comment in .proto file.""" + + def __init__(self, channel): + """Constructor. + + Args: + channel: A grpc.Channel. + """ + self.Connect = channel.stream_stream( + '/ax.HarnessService/Connect', + request_serializer=proto_dot_ax__pb2.HarnessMessage.SerializeToString, + response_deserializer=proto_dot_ax__pb2.HarnessMessage.FromString, + _registered_method=True) + + +class HarnessServiceServicer(object): + """Missing associated documentation comment in .proto file.""" + + def Connect(self, request_iterator, context): + """Missing associated documentation comment in .proto file.""" + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + + +def add_HarnessServiceServicer_to_server(servicer, server): + rpc_method_handlers = { + 'Connect': grpc.stream_stream_rpc_method_handler( + servicer.Connect, + request_deserializer=proto_dot_ax__pb2.HarnessMessage.FromString, + response_serializer=proto_dot_ax__pb2.HarnessMessage.SerializeToString, + ), + } + generic_handler = grpc.method_handlers_generic_handler( + 'ax.HarnessService', rpc_method_handlers) + server.add_generic_rpc_handlers((generic_handler,)) + server.add_registered_method_handlers('ax.HarnessService', rpc_method_handlers) + + + # This class is part of an EXPERIMENTAL API. +class HarnessService(object): + """Missing associated documentation comment in .proto file.""" + + @staticmethod + def Connect(request_iterator, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.stream_stream( + request_iterator, + target, + '/ax.HarnessService/Connect', + proto_dot_ax__pb2.HarnessMessage.SerializeToString, + proto_dot_ax__pb2.HarnessMessage.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + + class ControllerServiceStub(object): """Missing associated documentation comment in .proto file.""" @@ -191,11 +249,6 @@ def __init__(self, channel): request_serializer=proto_dot_ax__pb2.ExecRequest.SerializeToString, response_deserializer=proto_dot_ax__pb2.ExecResponse.FromString, _registered_method=True) - self.RegisterAgent = channel.unary_unary( - '/ax.ControllerService/RegisterAgent', - request_serializer=proto_dot_ax__pb2.RegisterAgentRequest.SerializeToString, - response_deserializer=proto_dot_ax__pb2.RegisterAgentResponse.FromString, - _registered_method=True) class ControllerServiceServicer(object): @@ -209,13 +262,6 @@ def Exec(self, request, context): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') - def RegisterAgent(self, request, context): - """RegisterAgent registers a new agent with the controller - """ - context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details('Method not implemented!') - raise NotImplementedError('Method not implemented!') - def add_ControllerServiceServicer_to_server(servicer, server): rpc_method_handlers = { @@ -224,11 +270,6 @@ def add_ControllerServiceServicer_to_server(servicer, server): request_deserializer=proto_dot_ax__pb2.ExecRequest.FromString, response_serializer=proto_dot_ax__pb2.ExecResponse.SerializeToString, ), - 'RegisterAgent': grpc.unary_unary_rpc_method_handler( - servicer.RegisterAgent, - request_deserializer=proto_dot_ax__pb2.RegisterAgentRequest.FromString, - response_serializer=proto_dot_ax__pb2.RegisterAgentResponse.SerializeToString, - ), } generic_handler = grpc.method_handlers_generic_handler( 'ax.ControllerService', rpc_method_handlers) @@ -267,35 +308,8 @@ def Exec(request, metadata, _registered_method=True) - @staticmethod - def RegisterAgent(request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None): - return grpc.experimental.unary_unary( - request, - target, - '/ax.ControllerService/RegisterAgent', - proto_dot_ax__pb2.RegisterAgentRequest.SerializeToString, - proto_dot_ax__pb2.RegisterAgentResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - _registered_method=True) - -class EventLogServiceStub(object): +class ConversationServiceStub(object): """Missing associated documentation comment in .proto file.""" def __init__(self, channel): @@ -304,34 +318,22 @@ def __init__(self, channel): Args: channel: A grpc.Channel. """ - self.List = channel.unary_unary( - '/ax.EventLogService/List', - request_serializer=proto_dot_ax__pb2.ListRequest.SerializeToString, - response_deserializer=proto_dot_ax__pb2.ListResponse.FromString, - _registered_method=True) - self.Delete = channel.unary_unary( - '/ax.EventLogService/Delete', - request_serializer=proto_dot_ax__pb2.DeleteRequest.SerializeToString, - response_deserializer=proto_dot_ax__pb2.DeleteResponse.FromString, + self.DeleteConversation = channel.unary_unary( + '/ax.ConversationService/DeleteConversation', + request_serializer=proto_dot_ax__pb2.DeleteConversationRequest.SerializeToString, + response_deserializer=proto_dot_ax__pb2.DeleteConversationResponse.FromString, _registered_method=True) - self.Fork = channel.unary_unary( - '/ax.EventLogService/Fork', - request_serializer=proto_dot_ax__pb2.ForkRequest.SerializeToString, - response_deserializer=proto_dot_ax__pb2.ForkResponse.FromString, + self.ForkConversation = channel.unary_unary( + '/ax.ConversationService/ForkConversation', + request_serializer=proto_dot_ax__pb2.ForkConversationRequest.SerializeToString, + response_deserializer=proto_dot_ax__pb2.ForkConversationResponse.FromString, _registered_method=True) -class EventLogServiceServicer(object): +class ConversationServiceServicer(object): """Missing associated documentation comment in .proto file.""" - def List(self, request, context): - """List conversational events. - """ - context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details('Method not implemented!') - raise NotImplementedError('Method not implemented!') - - def Delete(self, request, context): + def DeleteConversation(self, request, context): """Deletes conversational events and all event log resources for its children executions. """ @@ -339,7 +341,7 @@ def Delete(self, request, context): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') - def Fork(self, request, context): + def ForkConversation(self, request, context): """Fork forks an event log from a specific conversation. """ context.set_code(grpc.StatusCode.UNIMPLEMENTED) @@ -347,63 +349,31 @@ def Fork(self, request, context): raise NotImplementedError('Method not implemented!') -def add_EventLogServiceServicer_to_server(servicer, server): +def add_ConversationServiceServicer_to_server(servicer, server): rpc_method_handlers = { - 'List': grpc.unary_unary_rpc_method_handler( - servicer.List, - request_deserializer=proto_dot_ax__pb2.ListRequest.FromString, - response_serializer=proto_dot_ax__pb2.ListResponse.SerializeToString, - ), - 'Delete': grpc.unary_unary_rpc_method_handler( - servicer.Delete, - request_deserializer=proto_dot_ax__pb2.DeleteRequest.FromString, - response_serializer=proto_dot_ax__pb2.DeleteResponse.SerializeToString, + 'DeleteConversation': grpc.unary_unary_rpc_method_handler( + servicer.DeleteConversation, + request_deserializer=proto_dot_ax__pb2.DeleteConversationRequest.FromString, + response_serializer=proto_dot_ax__pb2.DeleteConversationResponse.SerializeToString, ), - 'Fork': grpc.unary_unary_rpc_method_handler( - servicer.Fork, - request_deserializer=proto_dot_ax__pb2.ForkRequest.FromString, - response_serializer=proto_dot_ax__pb2.ForkResponse.SerializeToString, + 'ForkConversation': grpc.unary_unary_rpc_method_handler( + servicer.ForkConversation, + request_deserializer=proto_dot_ax__pb2.ForkConversationRequest.FromString, + response_serializer=proto_dot_ax__pb2.ForkConversationResponse.SerializeToString, ), } generic_handler = grpc.method_handlers_generic_handler( - 'ax.EventLogService', rpc_method_handlers) + 'ax.ConversationService', rpc_method_handlers) server.add_generic_rpc_handlers((generic_handler,)) - server.add_registered_method_handlers('ax.EventLogService', rpc_method_handlers) + server.add_registered_method_handlers('ax.ConversationService', rpc_method_handlers) # This class is part of an EXPERIMENTAL API. -class EventLogService(object): +class ConversationService(object): """Missing associated documentation comment in .proto file.""" @staticmethod - def List(request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None): - return grpc.experimental.unary_unary( - request, - target, - '/ax.EventLogService/List', - proto_dot_ax__pb2.ListRequest.SerializeToString, - proto_dot_ax__pb2.ListResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - _registered_method=True) - - @staticmethod - def Delete(request, + def DeleteConversation(request, target, options=(), channel_credentials=None, @@ -416,9 +386,9 @@ def Delete(request, return grpc.experimental.unary_unary( request, target, - '/ax.EventLogService/Delete', - proto_dot_ax__pb2.DeleteRequest.SerializeToString, - proto_dot_ax__pb2.DeleteResponse.FromString, + '/ax.ConversationService/DeleteConversation', + proto_dot_ax__pb2.DeleteConversationRequest.SerializeToString, + proto_dot_ax__pb2.DeleteConversationResponse.FromString, options, channel_credentials, insecure, @@ -430,7 +400,7 @@ def Delete(request, _registered_method=True) @staticmethod - def Fork(request, + def ForkConversation(request, target, options=(), channel_credentials=None, @@ -443,9 +413,9 @@ def Fork(request, return grpc.experimental.unary_unary( request, target, - '/ax.EventLogService/Fork', - proto_dot_ax__pb2.ForkRequest.SerializeToString, - proto_dot_ax__pb2.ForkResponse.FromString, + '/ax.ConversationService/ForkConversation', + proto_dot_ax__pb2.ForkConversationRequest.SerializeToString, + proto_dot_ax__pb2.ForkConversationResponse.FromString, options, channel_credentials, insecure, diff --git a/python/proto/content_pb2.py b/python/proto/content_pb2.py index dff1333..e5c138d 100644 --- a/python/proto/content_pb2.py +++ b/python/proto/content_pb2.py @@ -1,17 +1,3 @@ -# 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. - # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE diff --git a/python/proto/content_pb2_grpc.py b/python/proto/content_pb2_grpc.py index 8ead1e8..08a981a 100644 --- a/python/proto/content_pb2_grpc.py +++ b/python/proto/content_pb2_grpc.py @@ -1,17 +1,3 @@ -# 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. - # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc From e4a2e147a387367315c3b007995ffb6af0b44db7 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 11:46:55 -0700 Subject: [PATCH 03/16] examples: Refactor antigravity agent.py to perform real-time token streaming to console --- examples/antigravity_agent/agent.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py index b9bbf6b..cd782e3 100644 --- a/examples/antigravity_agent/agent.py +++ b/examples/antigravity_agent/agent.py @@ -27,7 +27,12 @@ async def main(): async with Agent(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()) + + # Stream token deltas to console in real-time + async for token in response: + sys.stdout.write(token) + sys.stdout.flush() + print() if __name__ == "__main__": asyncio.run(main()) From 9767aa4c637cf8b7dfa76c00ace8796ccc6c7130 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 11:52:57 -0700 Subject: [PATCH 04/16] examples: Refactor agent.py to use low-level L2 Conversation and LocalConnectionStrategy API --- examples/antigravity_agent/agent.py | 40 +++++++++++++++++++++-------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py index cd782e3..a9211cf 100644 --- a/examples/antigravity_agent/agent.py +++ b/examples/antigravity_agent/agent.py @@ -14,25 +14,45 @@ import asyncio import sys -from google.antigravity import Agent, LocalAgentConfig +from google.antigravity import LocalAgentConfig +from google.antigravity.connections.local import LocalConnectionStrategy +from google.antigravity.conversation.conversation import Conversation +from google.antigravity.tools.tool_runner import ToolRunner +from google.antigravity.types import Text, Thought -# Expose agent_config globally for harness_server.py loading +# Expose agent_config globally for harness_server.py config loading agent_config = LocalAgentConfig( system_instructions="You are a helpful assistant powered by Google Antigravity." ) +# Expose the L2 configuration strategy for custom loaders if needed +strategy_factory = lambda: LocalConnectionStrategy(tool_runner=ToolRunner()) + async def main(): - # Initialize the agent session using the global config. - # It automatically picks up GEMINI_API_KEY from the environment. - async with Agent(agent_config) as agent: + # 1. Initialize the local connection strategy + strategy = strategy_factory() + + # 2. Create the stateful conversation session + print("Starting stateful Antigravity conversation (L2 API)...") + async with Conversation.create(strategy) as conversation: prompt = sys.argv[1] if len(sys.argv) > 1 else "Explain quantum computing in one sentence." - response = await agent.chat(prompt) - # Stream token deltas to console in real-time - async for token in response: - sys.stdout.write(token) - sys.stdout.flush() + # 3. Send query and receive streaming ChatResponse + response = await conversation.chat(prompt) + + # 4. Stream semantic chunks (Thoughts and Text) in real-time + async for chunk in response.chunks: + if isinstance(chunk, Text): + sys.stdout.write(chunk.text) + sys.stdout.flush() + elif isinstance(chunk, Thought): + # Display thought process in comment style + sys.stdout.write(f"\n[Thinking]: {chunk.text}") + sys.stdout.flush() print() if __name__ == "__main__": asyncio.run(main()) + +if __name__ == "__main__": + asyncio.run(main()) From 1a65a2805e9cdd6ba47afc6753eac56cd32d9eca Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 12:06:06 -0700 Subject: [PATCH 05/16] harness: Add full E2E ToolCall streaming and local execution support --- e2e.go | 4 +- examples/antigravity_agent/weather_agent.py | 77 +++++++++++++++++++++ hack/run-antigravity-streaming.sh | 2 +- internal/harness/antigravity.go | 41 ++++++++++- internal/harness/antigravity_test.go | 34 ++++++--- python/antigravity/harness_server.py | 9 ++- 6 files changed, 152 insertions(+), 15 deletions(-) create mode 100644 examples/antigravity_agent/weather_agent.py diff --git a/e2e.go b/e2e.go index 36f85f3..aafad0f 100644 --- a/e2e.go +++ b/e2e.go @@ -111,6 +111,8 @@ func runDemo(ctx context.Context, agentID string, setupRegistry func(reg *contro for _, out := range resp.Outputs { if textContent := out.GetContent().GetText().GetText(); textContent != "" { fmt.Printf("Agent Output: %s\n", textContent) + } else if toolCall := out.GetContent().GetToolCall(); toolCall != nil { + fmt.Printf("Agent Triggered Tool Call: %s (ID: %s)\n", toolCall.GetFunctionCall().Name, toolCall.Id) } } return nil @@ -121,7 +123,7 @@ func runDemo(ctx context.Context, agentID string, setupRegistry func(reg *contro Role: "user", Content: &proto.Content{ Type: &proto.Content_Text{ - Text: &proto.TextContent{Text: "Who are you?"}, + Text: &proto.TextContent{Text: "What is the weather in New York?"}, }, }, }, diff --git a/examples/antigravity_agent/weather_agent.py b/examples/antigravity_agent/weather_agent.py new file mode 100644 index 0000000..9352e44 --- /dev/null +++ b/examples/antigravity_agent/weather_agent.py @@ -0,0 +1,77 @@ +# 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 LocalAgentConfig +from google.antigravity.connections.local import LocalConnectionStrategy +from google.antigravity.conversation.conversation import Conversation +from google.antigravity.tools.tool_runner import ToolRunner +from google.antigravity.types import Text, Thought, ToolCall + +# Define the local python tool +def get_weather(city: str) -> str: + """Retrieves the current weather report for a specified city. + + Args: + city (str): The name of the city for which to retrieve the weather report. + + Returns: + str: Weather report status and details. + """ + # Output directly to stderr to not pollute the clean stream capture of stdout + sys.stderr.write(f"\n[PYTHON TOOL get_weather executed for city: {city}]\n") + sys.stderr.flush() + c = city.lower() + if "new york" in c or "nyc" in c: + return "The weather in New York is sunny with a temperature of 25 degrees Celsius (77 degrees Fahrenheit)." + elif "san francisco" in c or "sf" in c: + return "The weather in San Francisco is foggy with a temperature of 16 degrees Celsius (60.8 degrees Fahrenheit)." + else: + return f"Weather information for '{city}' is not available." + +# Expose agent_config globally for harness_server.py config loading +agent_config = LocalAgentConfig( + system_instructions="You are a helpful weather agent. Use the get_weather tool to answer weather questions.", + tools=[get_weather] +) + +async def main(): + # 1. Initialize local connection strategy with local tool + tool_runner = ToolRunner(tools=[get_weather]) + strategy = LocalConnectionStrategy(tool_runner=tool_runner) + + # 2. Create stateful conversation + print("Starting stateful Antigravity Weather Agent (L2 API)...") + async with Conversation.create(strategy) as conversation: + prompt = sys.argv[1] if len(sys.argv) > 1 else "What is the weather in New York?" + + # 3. Execute chat query + response = await conversation.chat(prompt) + + # 4. Stream chunks (Thought, Text, and ToolCall) in real-time + async for chunk in response.chunks: + if isinstance(chunk, Text): + sys.stdout.write(chunk.text) + sys.stdout.flush() + elif isinstance(chunk, Thought): + sys.stdout.write(f"\n[Thinking]: {chunk.text}") + sys.stdout.flush() + elif isinstance(chunk, ToolCall): + sys.stdout.write(f"\n[Tool Call]: {chunk.name} with args {chunk.args}\n") + sys.stdout.flush() + print() + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/hack/run-antigravity-streaming.sh b/hack/run-antigravity-streaming.sh index fc0da11..eda1877 100755 --- a/hack/run-antigravity-streaming.sh +++ b/hack/run-antigravity-streaming.sh @@ -24,7 +24,7 @@ fi PORT=50053 ADDRESS="localhost:$PORT" -AGENT_FILE="examples/antigravity_agent/agent.py" +AGENT_FILE="examples/antigravity_agent/weather_agent.py" # 1. Start Python WebSocket server in the background echo "Starting Python WebSocket Harness Server on port $PORT..." diff --git a/internal/harness/antigravity.go b/internal/harness/antigravity.go index e6e8d5a..57ad10a 100644 --- a/internal/harness/antigravity.go +++ b/internal/harness/antigravity.go @@ -22,6 +22,7 @@ import ( "github.com/gorilla/websocket" "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/types/known/structpb" "github.com/google/ax/proto" "github.com/google/uuid" @@ -124,9 +125,12 @@ func (e *antigravityExecution) Run(ctx context.Context, handler Handler) error { // 4. Stream responses from WebSocket type WSResponse struct { - Type string `json:"type"` - Content string `json:"content"` - Error string `json:"error"` + Type string `json:"type"` + Content string `json:"content"` + Error string `json:"error"` + ID string `json:"id"` + Name string `json:"name"` + Args json.RawMessage `json:"args"` } for { @@ -173,6 +177,37 @@ func (e *antigravityExecution) Run(ctx context.Context, handler Handler) error { if err := handler.OnMessage(ctx, e.id, msg); err != nil { return fmt.Errorf("failed to send thought to handler: %w", err) } + case "tool_call": + var argsMap map[string]any + if len(resp.Args) > 0 { + if err := json.Unmarshal(resp.Args, &argsMap); err != nil { + return fmt.Errorf("failed to unmarshal tool call args: %w", err) + } + } + structArgs, err := structpb.NewStruct(argsMap) + if err != nil { + return fmt.Errorf("failed to create structpb from tool call args: %w", err) + } + + msg := &proto.Message{ + Role: "model", + Content: &proto.Content{ + Type: &proto.Content_ToolCall{ + ToolCall: &proto.ToolCallContent{ + Id: resp.ID, + Type: &proto.ToolCallContent_FunctionCall{ + FunctionCall: &proto.FunctionCallContent{ + Name: resp.Name, + Arguments: structArgs, + }, + }, + }, + }, + }, + } + if err := handler.OnMessage(ctx, e.id, msg); err != nil { + return fmt.Errorf("failed to send tool call to handler: %w", err) + } case "complete": return handler.OnComplete(ctx, e.id) case "error": diff --git a/internal/harness/antigravity_test.go b/internal/harness/antigravity_test.go index 9c6243b..cd06f73 100644 --- a/internal/harness/antigravity_test.go +++ b/internal/harness/antigravity_test.go @@ -79,9 +79,10 @@ func TestAntigravityHarness_Run_Success(t *testing.T) { } // 2. Stream response chunks - chunks := []map[string]string{ - {"type": "text", "content": "Hello "}, - {"type": "text", "content": "world!"}, + chunks := []map[string]any{ + {"type": "thought", "content": "Analyzing request"}, + {"type": "tool_call", "id": "call-123", "name": "get_weather", "args": map[string]any{"city": "Paris"}}, + {"type": "text", "content": "The weather in Paris is rainy."}, {"type": "complete"}, } @@ -127,14 +128,29 @@ func TestAntigravityHarness_Run_Success(t *testing.T) { if !handler.complete { t.Error("expected OnComplete to be called") } - if len(handler.messages) != 2 { - t.Fatalf("expected 2 messages, got %d", len(handler.messages)) + if len(handler.messages) != 3 { + t.Fatalf("expected 3 messages, got %d", len(handler.messages)) } - if handler.messages[0].GetContent().GetText().GetText() != "Hello " { - t.Errorf("expected 'Hello ', got %q", handler.messages[0].GetContent().GetText().GetText()) + if handler.messages[0].GetContent().GetThought().GetSummary()[0].GetText().GetText() != "Analyzing request" { + t.Errorf("expected 'Analyzing request', got %q", handler.messages[0].GetContent().GetThought().GetSummary()[0].GetText().GetText()) } - if handler.messages[1].GetContent().GetText().GetText() != "world!" { - t.Errorf("expected 'world!', got %q", handler.messages[1].GetContent().GetText().GetText()) + + toolCall := handler.messages[1].GetContent().GetToolCall() + if toolCall == nil { + t.Fatal("expected tool call message, got nil") + } + if toolCall.Id != "call-123" { + t.Errorf("expected ID 'call-123', got %q", toolCall.Id) + } + if toolCall.GetFunctionCall().Name != "get_weather" { + t.Errorf("expected name 'get_weather', got %q", toolCall.GetFunctionCall().Name) + } + if toolCall.GetFunctionCall().Arguments.GetFields()["city"].GetStringValue() != "Paris" { + t.Errorf("expected arg city='Paris', got %q", toolCall.GetFunctionCall().Arguments.GetFields()["city"].GetStringValue()) + } + + if handler.messages[2].GetContent().GetText().GetText() != "The weather in Paris is rainy." { + t.Errorf("expected 'The weather in Paris is rainy.', got %q", handler.messages[2].GetContent().GetText().GetText()) } } diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index e4eae33..88e7099 100644 --- a/python/antigravity/harness_server.py +++ b/python/antigravity/harness_server.py @@ -26,7 +26,7 @@ from google.protobuf.json_format import Parse from proto import ax_pb2 from google.antigravity import Agent, AgentConfig -from google.antigravity.types import Step, StepType, StepSource, StepTarget, StepStatus, Text, Thought +from google.antigravity.types import Step, StepType, StepSource, StepTarget, StepStatus, Text, Thought, ToolCall app = FastAPI() @@ -146,6 +146,13 @@ async def websocket_endpoint(websocket: WebSocket): await websocket.send_json({"type": "text", "content": chunk.text}) elif isinstance(chunk, Thought): await websocket.send_json({"type": "thought", "content": chunk.text}) + elif isinstance(chunk, ToolCall): + await websocket.send_json({ + "type": "tool_call", + "id": chunk.id or "", + "name": str(chunk.name), + "args": chunk.args + }) # Send complete frame await websocket.send_json({"type": "complete"}) From 25f8f84d0e0692af5db3af9c37b3a1561eb92a97 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 12:14:06 -0700 Subject: [PATCH 06/16] examples: Consolidate agent.py and weather_agent.py into a single ultimate agent.py example --- examples/antigravity_agent/agent.py | 48 +++++++++---- examples/antigravity_agent/weather_agent.py | 77 --------------------- hack/run-antigravity-streaming.sh | 2 +- 3 files changed, 36 insertions(+), 91 deletions(-) delete mode 100644 examples/antigravity_agent/weather_agent.py diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py index a9211cf..bb61e40 100644 --- a/examples/antigravity_agent/agent.py +++ b/examples/antigravity_agent/agent.py @@ -18,29 +18,51 @@ from google.antigravity.connections.local import LocalConnectionStrategy from google.antigravity.conversation.conversation import Conversation from google.antigravity.tools.tool_runner import ToolRunner -from google.antigravity.types import Text, Thought +from google.antigravity.types import Text, Thought, ToolCall -# Expose agent_config globally for harness_server.py config loading +# 1. Define a custom local python tool +def get_weather(city: str) -> str: + """Retrieves the current weather report for a specified city. + + Args: + city (str): The name of the city for which to retrieve the weather report. + + Returns: + str: Weather report status and details. + """ + # Output to stderr so it does not pollute the stdout stream capture + sys.stderr.write(f"\n[PYTHON TOOL get_weather executed for city: {city}]\n") + sys.stderr.flush() + c = city.lower() + if "new york" in c or "nyc" in c: + return "The weather in New York is sunny with a temperature of 25 degrees Celsius (77 degrees Fahrenheit)." + elif "san francisco" in c or "sf" in c: + return "The weather in San Francisco is foggy with a temperature of 16 degrees Celsius (60.8 degrees Fahrenheit)." + else: + return f"Weather information for '{city}' is not available." + +# 2. Expose agent_config globally for harness_server.py config loading agent_config = LocalAgentConfig( - system_instructions="You are a helpful assistant powered by Google Antigravity." + system_instructions="You are a helpful agent. Use the get_weather tool to answer weather questions.", + tools=[get_weather] ) -# Expose the L2 configuration strategy for custom loaders if needed -strategy_factory = lambda: LocalConnectionStrategy(tool_runner=ToolRunner()) +# Expose the L2 configuration strategy factory for custom loaders +strategy_factory = lambda: LocalConnectionStrategy(tool_runner=ToolRunner(tools=[get_weather])) async def main(): - # 1. Initialize the local connection strategy + # 3. Initialize the local connection strategy strategy = strategy_factory() - # 2. Create the stateful conversation session + # 4. Create the stateful conversation session print("Starting stateful Antigravity conversation (L2 API)...") async with Conversation.create(strategy) as conversation: - prompt = sys.argv[1] if len(sys.argv) > 1 else "Explain quantum computing in one sentence." + prompt = sys.argv[1] if len(sys.argv) > 1 else "What is the weather in New York?" - # 3. Send query and receive streaming ChatResponse + # 5. Send query and receive streaming ChatResponse response = await conversation.chat(prompt) - # 4. Stream semantic chunks (Thoughts and Text) in real-time + # 6. Stream semantic chunks (Thoughts, Text, and ToolCalls) in real-time async for chunk in response.chunks: if isinstance(chunk, Text): sys.stdout.write(chunk.text) @@ -49,10 +71,10 @@ async def main(): # Display thought process in comment style sys.stdout.write(f"\n[Thinking]: {chunk.text}") sys.stdout.flush() + elif isinstance(chunk, ToolCall): + sys.stdout.write(f"\n[Tool Call]: {chunk.name} with args {chunk.args}\n") + sys.stdout.flush() print() if __name__ == "__main__": asyncio.run(main()) - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/examples/antigravity_agent/weather_agent.py b/examples/antigravity_agent/weather_agent.py deleted file mode 100644 index 9352e44..0000000 --- a/examples/antigravity_agent/weather_agent.py +++ /dev/null @@ -1,77 +0,0 @@ -# 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 LocalAgentConfig -from google.antigravity.connections.local import LocalConnectionStrategy -from google.antigravity.conversation.conversation import Conversation -from google.antigravity.tools.tool_runner import ToolRunner -from google.antigravity.types import Text, Thought, ToolCall - -# Define the local python tool -def get_weather(city: str) -> str: - """Retrieves the current weather report for a specified city. - - Args: - city (str): The name of the city for which to retrieve the weather report. - - Returns: - str: Weather report status and details. - """ - # Output directly to stderr to not pollute the clean stream capture of stdout - sys.stderr.write(f"\n[PYTHON TOOL get_weather executed for city: {city}]\n") - sys.stderr.flush() - c = city.lower() - if "new york" in c or "nyc" in c: - return "The weather in New York is sunny with a temperature of 25 degrees Celsius (77 degrees Fahrenheit)." - elif "san francisco" in c or "sf" in c: - return "The weather in San Francisco is foggy with a temperature of 16 degrees Celsius (60.8 degrees Fahrenheit)." - else: - return f"Weather information for '{city}' is not available." - -# Expose agent_config globally for harness_server.py config loading -agent_config = LocalAgentConfig( - system_instructions="You are a helpful weather agent. Use the get_weather tool to answer weather questions.", - tools=[get_weather] -) - -async def main(): - # 1. Initialize local connection strategy with local tool - tool_runner = ToolRunner(tools=[get_weather]) - strategy = LocalConnectionStrategy(tool_runner=tool_runner) - - # 2. Create stateful conversation - print("Starting stateful Antigravity Weather Agent (L2 API)...") - async with Conversation.create(strategy) as conversation: - prompt = sys.argv[1] if len(sys.argv) > 1 else "What is the weather in New York?" - - # 3. Execute chat query - response = await conversation.chat(prompt) - - # 4. Stream chunks (Thought, Text, and ToolCall) in real-time - async for chunk in response.chunks: - if isinstance(chunk, Text): - sys.stdout.write(chunk.text) - sys.stdout.flush() - elif isinstance(chunk, Thought): - sys.stdout.write(f"\n[Thinking]: {chunk.text}") - sys.stdout.flush() - elif isinstance(chunk, ToolCall): - sys.stdout.write(f"\n[Tool Call]: {chunk.name} with args {chunk.args}\n") - sys.stdout.flush() - print() - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/hack/run-antigravity-streaming.sh b/hack/run-antigravity-streaming.sh index eda1877..fc0da11 100755 --- a/hack/run-antigravity-streaming.sh +++ b/hack/run-antigravity-streaming.sh @@ -24,7 +24,7 @@ fi PORT=50053 ADDRESS="localhost:$PORT" -AGENT_FILE="examples/antigravity_agent/weather_agent.py" +AGENT_FILE="examples/antigravity_agent/agent.py" # 1. Start Python WebSocket server in the background echo "Starting Python WebSocket Harness Server on port $PORT..." From 4c857c661d46c4c2c79e8ba4af0362dce354b350 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 14:32:46 -0700 Subject: [PATCH 07/16] harness: Cleanly relocate Go E2E program to cmd/e2e/main.go --- e2e.go => cmd/e2e/main.go | 0 hack/run-antigravity-streaming.sh | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename e2e.go => cmd/e2e/main.go (100%) diff --git a/e2e.go b/cmd/e2e/main.go similarity index 100% rename from e2e.go rename to cmd/e2e/main.go diff --git a/hack/run-antigravity-streaming.sh b/hack/run-antigravity-streaming.sh index fc0da11..9449b07 100755 --- a/hack/run-antigravity-streaming.sh +++ b/hack/run-antigravity-streaming.sh @@ -66,7 +66,7 @@ echo "Python server is active!" # 3. Build and run the Go E2E V2 demonstration echo "Building e2e..." -/opt/homebrew/bin/go build -o bin/e2e e2e.go +/opt/homebrew/bin/go build -o bin/e2e ./cmd/e2e echo "Executing E2E Demo with Antigravity WebSocket Harness..." bin/e2e From ef0342e5811aeadf755ba09e91abacf1b3320d59 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 14:40:32 -0700 Subject: [PATCH 08/16] harness: Add explicit compile-time interface assertions in antigravity.go --- internal/harness/antigravity.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/internal/harness/antigravity.go b/internal/harness/antigravity.go index 57ad10a..30a8d1c 100644 --- a/internal/harness/antigravity.go +++ b/internal/harness/antigravity.go @@ -28,6 +28,10 @@ import ( "github.com/google/uuid" ) +// Compile-time interface assertions. +var _ Harness = (*AntigravityHarness)(nil) +var _ Execution = (*antigravityExecution)(nil) + // AntigravityHarness implements the Harness interface by connecting to the // Antigravity Python agent server over WebSockets. type AntigravityHarness struct { From bccd63fc65bf56fcb759b9d416be24f827d4ca6b Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 14:49:43 -0700 Subject: [PATCH 09/16] python: Add mandatory Google license headers to regenerated proto modules --- python/proto/ax_pb2.py | 14 ++++++++++++++ python/proto/ax_pb2_grpc.py | 14 ++++++++++++++ python/proto/content_pb2.py | 14 ++++++++++++++ python/proto/content_pb2_grpc.py | 14 ++++++++++++++ 4 files changed, 56 insertions(+) diff --git a/python/proto/ax_pb2.py b/python/proto/ax_pb2.py index 37a01a1..9c62f51 100644 --- a/python/proto/ax_pb2.py +++ b/python/proto/ax_pb2.py @@ -1,3 +1,17 @@ +# 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. + # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE diff --git a/python/proto/ax_pb2_grpc.py b/python/proto/ax_pb2_grpc.py index e042b66..48b9ebd 100644 --- a/python/proto/ax_pb2_grpc.py +++ b/python/proto/ax_pb2_grpc.py @@ -1,3 +1,17 @@ +# 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. + # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc diff --git a/python/proto/content_pb2.py b/python/proto/content_pb2.py index e5c138d..dff1333 100644 --- a/python/proto/content_pb2.py +++ b/python/proto/content_pb2.py @@ -1,3 +1,17 @@ +# 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. + # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # NO CHECKED-IN PROTOBUF GENCODE diff --git a/python/proto/content_pb2_grpc.py b/python/proto/content_pb2_grpc.py index 08a981a..8ead1e8 100644 --- a/python/proto/content_pb2_grpc.py +++ b/python/proto/content_pb2_grpc.py @@ -1,3 +1,17 @@ +# 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. + # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! """Client and server classes corresponding to protobuf-defined services.""" import grpc From b950b896f79b1864b150a361a5236ddf1f7444d5 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Thu, 28 May 2026 16:22:32 -0700 Subject: [PATCH 10/16] harness: Re-architect always-on streaming harness to use standard gRPC AgentService protocol --- cmd/e2e/main.go | 13 +- hack/run-antigravity-streaming.sh | 6 +- internal/harness/antigravity.go | 168 +++++------------ internal/harness/antigravity_test.go | 212 ++++++++++++---------- python/antigravity/harness_server.py | 182 +++++++++++-------- python/antigravity/harness_server_test.py | 185 ++++++++----------- 6 files changed, 359 insertions(+), 407 deletions(-) diff --git a/cmd/e2e/main.go b/cmd/e2e/main.go index aafad0f..0e98a1c 100644 --- a/cmd/e2e/main.go +++ b/cmd/e2e/main.go @@ -19,9 +19,9 @@ package main import ( "context" "fmt" + "net" "os" - - "github.com/gorilla/websocket" + "time" "github.com/google/ax/internal/controller/executor" "github.com/google/ax/internal/controller/executor/executortest" @@ -73,17 +73,16 @@ func main() { 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) { - // Check if Python WebSocket server is active, otherwise fallback + // Check if Python gRPC server is active, otherwise fallback var realHarness harness.Harness - address := "ws://localhost:50053/ws" - dialer := websocket.DefaultDialer - conn, _, err := dialer.Dial(address, nil) + address := "localhost:50053" + conn, err := net.DialTimeout("tcp", address, 1*time.Second) if err != nil { fmt.Printf("WARNING: Antigravity harness server not active at %s, falling back to test harness: %v\n", address, err) realHarness = harnesstest.New() } else { conn.Close() - fmt.Printf("Connected to Antigravity harness server at %s\n", address) + fmt.Printf("Connected to Antigravity gRPC harness server at %s\n", address) realHarness = harness.NewAntigravityHarness(address) } reg.RegisterHarness("antigravity", realHarness) diff --git a/hack/run-antigravity-streaming.sh b/hack/run-antigravity-streaming.sh index 9449b07..7e80428 100755 --- a/hack/run-antigravity-streaming.sh +++ b/hack/run-antigravity-streaming.sh @@ -26,8 +26,8 @@ PORT=50053 ADDRESS="localhost:$PORT" AGENT_FILE="examples/antigravity_agent/agent.py" -# 1. Start Python WebSocket server in the background -echo "Starting Python WebSocket Harness Server on port $PORT..." +# 1. Start Python gRPC server in the background +echo "Starting Python gRPC Harness Server on port $PORT..." PYTHONPATH=python:. .venv/bin/python -m python.antigravity.harness_server --agent_file "$AGENT_FILE" --port "$PORT" > /tmp/antigravity_harness.log 2>&1 & SERVER_PID=$! @@ -68,7 +68,7 @@ echo "Python server is active!" echo "Building e2e..." /opt/homebrew/bin/go build -o bin/e2e ./cmd/e2e -echo "Executing E2E Demo with Antigravity WebSocket Harness..." +echo "Executing E2E Demo with Antigravity gRPC Harness..." bin/e2e echo "Success!" diff --git a/internal/harness/antigravity.go b/internal/harness/antigravity.go index 30a8d1c..15158de 100644 --- a/internal/harness/antigravity.go +++ b/internal/harness/antigravity.go @@ -16,13 +16,12 @@ package harness import ( "context" - "encoding/json" "fmt" + "io" "sync" - "github.com/gorilla/websocket" - "google.golang.org/protobuf/encoding/protojson" - "google.golang.org/protobuf/types/known/structpb" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" "github.com/google/ax/proto" "github.com/google/uuid" @@ -33,15 +32,16 @@ var _ Harness = (*AntigravityHarness)(nil) var _ Execution = (*antigravityExecution)(nil) // AntigravityHarness implements the Harness interface by connecting to the -// Antigravity Python agent server over WebSockets. +// Antigravity Python agent server over gRPC. type AntigravityHarness struct { address string } // NewAntigravityHarness creates a new AntigravityHarness with a configurable address. +// Address defaults to "localhost:50053" (gRPC TCP connection). func NewAntigravityHarness(address string) *AntigravityHarness { if address == "" { - address = "ws://localhost:50053/ws" + address = "localhost:50053" } return &AntigravityHarness{ address: address, @@ -78,148 +78,78 @@ func (e *antigravityExecution) Queue(ctx context.Context, msg ...*proto.Message) e.mu.Lock() defer e.mu.Unlock() if e.closed { - return fmt.Errorf("execution is closed") + return fmt.Errorf("execution session already closed") } e.queued = append(e.queued, msg...) return nil } -// Run implements Execution.Run. -// It connects to the Python server over WebSockets, sends the start payload containing history, -// and streams responses back to the handler. +// Run executes the turn over gRPC bidirectional streaming and forwards events to the handler. func (e *antigravityExecution) Run(ctx context.Context, handler Handler) error { e.mu.Lock() + if e.closed { + e.mu.Unlock() + return fmt.Errorf("execution session already closed") + } + // Retrieve queued inputs inputs := e.queued e.queued = nil e.mu.Unlock() - // 1. Establish WebSocket connection - dialer := websocket.DefaultDialer - conn, _, err := dialer.DialContext(ctx, e.harness.address, nil) - if err != nil { - return fmt.Errorf("failed to dial antigravity harness websocket at %s: %w", e.harness.address, err) + if len(inputs) == 0 { + return fmt.Errorf("no input messages queued for execution turn") } - defer conn.Close() - // 2. Serialize inputs using protojson to match Python Parse() requirements - var serializedMessages []json.RawMessage - for _, msg := range inputs { - bytes, err := protojson.Marshal(msg) - if err != nil { - return fmt.Errorf("failed to marshal message to JSON: %w", err) - } - serializedMessages = append(serializedMessages, json.RawMessage(bytes)) - } - - // 3. Construct and send start payload - startPayload := map[string]any{ - "conversation_id": e.conversationID, - "exec_id": e.id, - "messages": serializedMessages, - } - payloadBytes, err := json.Marshal(startPayload) + // 1. Connect to the gRPC server + conn, err := grpc.DialContext(ctx, e.harness.address, grpc.WithTransportCredentials(insecure.NewCredentials())) if err != nil { - return fmt.Errorf("failed to marshal start payload: %w", err) + return fmt.Errorf("failed to connect to gRPC harness server at %s: %w", e.harness.address, err) } + defer conn.Close() - err = conn.WriteMessage(websocket.TextMessage, payloadBytes) - if err != nil { - return fmt.Errorf("failed to send start payload over WebSocket: %w", err) + // 2. Create AgentService client + client := proto.NewAgentServiceClient(conn) + + // 3. Build standard AgentRequest + req := &proto.AgentRequest{ + ConversationId: e.conversationID, + ExecId: e.id, + Start: &proto.AgentStart{ + AgentId: "antigravity", + Messages: inputs, + }, } - // 4. Stream responses from WebSocket - type WSResponse struct { - Type string `json:"type"` - Content string `json:"content"` - Error string `json:"error"` - ID string `json:"id"` - Name string `json:"name"` - Args json.RawMessage `json:"args"` + // 4. Call Connect to start bidirectional streaming + stream, err := client.Connect(ctx, req) + if err != nil { + return fmt.Errorf("failed to call gRPC AgentService.Connect: %w", err) } + // 5. Stream responses and trigger callbacks for { - _, message, err := conn.ReadMessage() - if err != nil { - return fmt.Errorf("failed to read message from WebSocket: %w", err) + resp, err := stream.Recv() + if err == io.EOF { + break } - - var resp WSResponse - if err := json.Unmarshal(message, &resp); err != nil { - return fmt.Errorf("failed to unmarshal WebSocket response: %w", err) + if err != nil { + return fmt.Errorf("gRPC harness streaming failure: %w", err) } - switch resp.Type { - case "text": - msg := &proto.Message{ - Role: "assistant", - Content: &proto.Content{ - Type: &proto.Content_Text{ - Text: &proto.TextContent{Text: resp.Content}, - }, - }, - } - if err := handler.OnMessage(ctx, e.id, msg); err != nil { - return fmt.Errorf("failed to send message to handler: %w", err) - } - case "thought": - msg := &proto.Message{ - Role: "model", - Content: &proto.Content{ - Type: &proto.Content_Thought{ - Thought: &proto.ThoughtContent{ - Summary: []*proto.ThoughtSummaryContent{ - { - Type: &proto.ThoughtSummaryContent_Text{ - Text: &proto.TextContent{Text: resp.Content}, - }, - }, - }, - }, - }, - }, - } - if err := handler.OnMessage(ctx, e.id, msg); err != nil { - return fmt.Errorf("failed to send thought to handler: %w", err) - } - case "tool_call": - var argsMap map[string]any - if len(resp.Args) > 0 { - if err := json.Unmarshal(resp.Args, &argsMap); err != nil { - return fmt.Errorf("failed to unmarshal tool call args: %w", err) + switch payload := resp.Type.(type) { + case *proto.AgentResponse_Outputs: + for _, outMsg := range payload.Outputs.Messages { + if err := handler.OnMessage(ctx, e.id, outMsg); err != nil { + return fmt.Errorf("failed to dispatch streamed output: %w", err) } } - structArgs, err := structpb.NewStruct(argsMap) - if err != nil { - return fmt.Errorf("failed to create structpb from tool call args: %w", err) - } - - msg := &proto.Message{ - Role: "model", - Content: &proto.Content{ - Type: &proto.Content_ToolCall{ - ToolCall: &proto.ToolCallContent{ - Id: resp.ID, - Type: &proto.ToolCallContent_FunctionCall{ - FunctionCall: &proto.FunctionCallContent{ - Name: resp.Name, - Arguments: structArgs, - }, - }, - }, - }, - }, - } - if err := handler.OnMessage(ctx, e.id, msg); err != nil { - return fmt.Errorf("failed to send tool call to handler: %w", err) - } - case "complete": + case *proto.AgentResponse_End: + // Standard turn complete callback return handler.OnComplete(ctx, e.id) - case "error": - return fmt.Errorf("antigravity harness server error: %s", resp.Error) - default: - return fmt.Errorf("unknown response type from WebSocket: %q", resp.Type) } } + + return nil } // Close implements Execution.Close. diff --git a/internal/harness/antigravity_test.go b/internal/harness/antigravity_test.go index cd06f73..18d23b9 100644 --- a/internal/harness/antigravity_test.go +++ b/internal/harness/antigravity_test.go @@ -16,14 +16,14 @@ package harness import ( "context" - "encoding/json" - "net/http" - "net/http/httptest" + "net" "strings" "sync" "testing" - "github.com/gorilla/websocket" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "github.com/google/ax/proto" ) @@ -49,57 +49,101 @@ func (h *mockHandler) OnComplete(ctx context.Context, execID string) error { return nil } -func TestAntigravityHarness_Run_Success(t *testing.T) { - upgrader := websocket.Upgrader{} - - // Spin up a local mock WebSocket server - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) - if err != nil { - t.Errorf("failed to upgrade connection: %v", err) - return - } - defer conn.Close() +// mockAgentServer implements proto.AgentServiceServer for testing. +type mockAgentServer struct { + proto.UnimplementedAgentServiceServer + failConnect bool +} - // 1. Read initial start message - _, msg, err := conn.ReadMessage() - if err != nil { - t.Errorf("failed to read start message: %v", err) - return - } +func (s *mockAgentServer) Connect(req *proto.AgentRequest, stream proto.AgentService_ConnectServer) error { + if s.failConnect { + return status.Error(codes.Internal, "internal mock server crash") + } - var payload map[string]any - if err := json.Unmarshal(msg, &payload); err != nil { - t.Errorf("failed to unmarshal start payload: %v", err) - return - } + // 1. Verify conversation details + if req.ConversationId != "conv-test" { + return status.Error(codes.InvalidArgument, "invalid conversation_id") + } - if payload["conversation_id"] != "conv-test" { - t.Errorf("expected conversation ID 'conv-test', got %v", payload["conversation_id"]) - } + // 2. Stream thought frame + tMsg := &proto.Message{ + Role: "model", + Content: &proto.Content{ + Type: &proto.Content_Thought{ + Thought: &proto.ThoughtContent{ + Summary: []*proto.ThoughtSummaryContent{ + { + Type: &proto.ThoughtSummaryContent_Text{ + Text: &proto.TextContent{Text: "Analyzing"}, + }, + }, + }, + }, + }, + }, + } + err := stream.Send(&proto.AgentResponse{ + ConversationId: req.ConversationId, + ExecId: req.ExecId, + Type: &proto.AgentResponse_Outputs{ + Outputs: &proto.AgentOutputs{Messages: []*proto.Message{tMsg}}, + }, + }) + if err != nil { + return err + } - // 2. Stream response chunks - chunks := []map[string]any{ - {"type": "thought", "content": "Analyzing request"}, - {"type": "tool_call", "id": "call-123", "name": "get_weather", "args": map[string]any{"city": "Paris"}}, - {"type": "text", "content": "The weather in Paris is rainy."}, - {"type": "complete"}, - } + // 3. Stream text frame + txtMsg := &proto.Message{ + Role: "assistant", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "Hello world"}, + }, + }, + } + err = stream.Send(&proto.AgentResponse{ + ConversationId: req.ConversationId, + ExecId: req.ExecId, + Type: &proto.AgentResponse_Outputs{ + Outputs: &proto.AgentOutputs{Messages: []*proto.Message{txtMsg}}, + }, + }) + if err != nil { + return err + } - for _, chunk := range chunks { - bytes, _ := json.Marshal(chunk) - if err := conn.WriteMessage(websocket.TextMessage, bytes); err != nil { - t.Errorf("failed to write chunk: %v", err) - return - } - } - })) - defer server.Close() + // 4. Stream end frame + return stream.Send(&proto.AgentResponse{ + ConversationId: req.ConversationId, + ExecId: req.ExecId, + Type: &proto.AgentResponse_End{ + End: &proto.AgentEnd{}, + }, + }) +} - // Convert http:// to ws:// - wsURL := strings.Replace(server.URL, "http://", "ws://", 1) + "/ws" +func TestAntigravityHarness_Run_Success(t *testing.T) { + // Spin up a local TCP listener + lis, err := net.Listen("tcp", "localhost:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + defer lis.Close() + + // Initialize and start local gRPC server + grpcServer := grpc.NewServer() + mockServer := &mockAgentServer{} + proto.RegisterAgentServiceServer(grpcServer, mockServer) - harnessClient := NewAntigravityHarness(wsURL) + go func() { + if err := grpcServer.Serve(lis); err != nil && err != grpc.ErrServerStopped { + t.Errorf("Serve failed: %v", err) + } + }() + defer grpcServer.Stop() + + harnessClient := NewAntigravityHarness(lis.Addr().String()) exec, err := harnessClient.Start(context.Background(), "conv-test") if err != nil { t.Fatalf("failed to start execution: %v", err) @@ -128,65 +172,51 @@ func TestAntigravityHarness_Run_Success(t *testing.T) { if !handler.complete { t.Error("expected OnComplete to be called") } - if len(handler.messages) != 3 { - t.Fatalf("expected 3 messages, got %d", len(handler.messages)) - } - if handler.messages[0].GetContent().GetThought().GetSummary()[0].GetText().GetText() != "Analyzing request" { - t.Errorf("expected 'Analyzing request', got %q", handler.messages[0].GetContent().GetThought().GetSummary()[0].GetText().GetText()) + if len(handler.messages) != 2 { + t.Fatalf("expected 2 messages, got %d", len(handler.messages)) } - - toolCall := handler.messages[1].GetContent().GetToolCall() - if toolCall == nil { - t.Fatal("expected tool call message, got nil") + if handler.messages[0].GetContent().GetThought().GetSummary()[0].GetText().GetText() != "Analyzing" { + t.Errorf("expected 'Analyzing', got %q", handler.messages[0].GetContent().GetThought().GetSummary()[0].GetText().GetText()) } - if toolCall.Id != "call-123" { - t.Errorf("expected ID 'call-123', got %q", toolCall.Id) - } - if toolCall.GetFunctionCall().Name != "get_weather" { - t.Errorf("expected name 'get_weather', got %q", toolCall.GetFunctionCall().Name) - } - if toolCall.GetFunctionCall().Arguments.GetFields()["city"].GetStringValue() != "Paris" { - t.Errorf("expected arg city='Paris', got %q", toolCall.GetFunctionCall().Arguments.GetFields()["city"].GetStringValue()) - } - - if handler.messages[2].GetContent().GetText().GetText() != "The weather in Paris is rainy." { - t.Errorf("expected 'The weather in Paris is rainy.', got %q", handler.messages[2].GetContent().GetText().GetText()) + if handler.messages[1].GetContent().GetText().GetText() != "Hello world" { + t.Errorf("expected 'Hello world', got %q", handler.messages[1].GetContent().GetText().GetText()) } } func TestAntigravityHarness_Run_ErrorFrame(t *testing.T) { - upgrader := websocket.Upgrader{} - - server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) - if err != nil { - return - } - defer conn.Close() - - _, _, _ = conn.ReadMessage() + lis, err := net.Listen("tcp", "localhost:0") + if err != nil { + t.Fatalf("failed to listen: %v", err) + } + defer lis.Close() - errFrame := map[string]string{ - "type": "error", - "error": "internal model crash", - } - bytes, _ := json.Marshal(errFrame) - conn.WriteMessage(websocket.TextMessage, bytes) - })) - defer server.Close() + grpcServer := grpc.NewServer() + mockServer := &mockAgentServer{failConnect: true} + proto.RegisterAgentServiceServer(grpcServer, mockServer) - wsURL := strings.Replace(server.URL, "http://", "ws://", 1) + "/ws" + go func() { + _ = grpcServer.Serve(lis) + }() + defer grpcServer.Stop() - harnessClient := NewAntigravityHarness(wsURL) + harnessClient := NewAntigravityHarness(lis.Addr().String()) exec, _ := harnessClient.Start(context.Background(), "conv-test") defer exec.Close(context.Background()) + msg := &proto.Message{ + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{Text: &proto.TextContent{Text: "Hi"}}, + }, + } + _ = exec.Queue(context.Background(), msg) + handler := &mockHandler{} - err := exec.Run(context.Background(), handler) + err = exec.Run(context.Background(), handler) if err == nil { t.Fatal("expected error from Run(), got nil") } - if !strings.Contains(err.Error(), "antigravity harness server error: internal model crash") { + if !strings.Contains(err.Error(), "internal mock server crash") { t.Errorf("unexpected error message: %v", err) } } diff --git a/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index 88e7099..03b4af9 100644 --- a/python/antigravity/harness_server.py +++ b/python/antigravity/harness_server.py @@ -15,21 +15,17 @@ import argparse import asyncio import importlib.util -import json import logging -import os import sys -import uuid -from fastapi import FastAPI, WebSocket, WebSocketDisconnect -import uvicorn +import grpc +from google.protobuf.struct_pb2 import Struct -from google.protobuf.json_format import Parse -from proto import ax_pb2 +from python.proto import ax_pb2 +from python.proto import ax_pb2_grpc +from python.proto import content_pb2 from google.antigravity import Agent, AgentConfig from google.antigravity.types import Step, StepType, StepSource, StepTarget, StepStatus, Text, Thought, ToolCall -app = FastAPI() - # Global placeholder for loaded agent config loaded_config: AgentConfig | None = None @@ -91,86 +87,123 @@ def hydrate_ax_history_to_steps(historical_messages) -> list[Step]: steps.append(step) return steps -@app.websocket("/ws") -async def websocket_endpoint(websocket: WebSocket): - await websocket.accept() - print("[WS] Connection accepted.") - try: - # 1. Receive the start message - data = await websocket.receive_text() - payload = json.loads(data) - - conversation_id = payload.get("conversation_id") - exec_id = payload.get("exec_id") - raw_messages = payload.get("messages", []) - - print(f"[WS] Starting turn. conv_id={conversation_id}, exec_id={exec_id}, messages_count={len(raw_messages)}") +class AntigravityAgentServiceServicer(ax_pb2_grpc.AgentServiceServicer): + """Implements the standard ax.AgentService protocol over gRPC.""" + + async def Connect(self, request: ax_pb2.AgentRequest, context): + print(f"[gRPC] Connect turn requested. conv_id={request.conversation_id}, exec_id={request.exec_id}") - # Deserialize AX protobuf messages - ax_messages = [] - for raw_msg in raw_messages: - msg_str = json.dumps(raw_msg) - ax_msg = Parse(msg_str, ax_pb2.Message()) - ax_messages.append(ax_msg) - + # 1. Retrieve and check messages + ax_messages = request.start.messages if not ax_messages: - raise ValueError("No messages found in start payload") + context.set_code(grpc.StatusCode.INVALID_ARGUMENT) + context.set_details("No messages found in start payload") + return historical_messages = ax_messages[:-1] latest_message = ax_messages[-1] - # Only support text queries for now in latest_message if latest_message.content.WhichOneof('type') != 'text': - raise ValueError("Latest message must contain text content") + context.set_code(grpc.StatusCode.INVALID_ARGUMENT) + context.set_details("Latest message must contain text content") + return latest_query_text = latest_message.content.text.text # 2. Initialize the Antigravity Agent session global loaded_config if not loaded_config: - raise RuntimeError("Agent config is not loaded on the server") - - async with Agent(loaded_config) as agent: - conversation = agent.conversation - - # Hydrate history - print(f"[WS] Hydrating {len(historical_messages)} historical messages...") - history_steps = hydrate_ax_history_to_steps(historical_messages) - conversation._steps.extend(history_steps) - - # Run the turn with streaming - print(f"[WS] Running chat query: {latest_query_text}") - response = await conversation.chat(latest_query_text) + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details("Agent config is not loaded on the server") + return - async for chunk in response.chunks: - if isinstance(chunk, Text): - await websocket.send_json({"type": "text", "content": chunk.text}) - elif isinstance(chunk, Thought): - await websocket.send_json({"type": "thought", "content": chunk.text}) - elif isinstance(chunk, ToolCall): - await websocket.send_json({ - "type": "tool_call", - "id": chunk.id or "", - "name": str(chunk.name), - "args": chunk.args - }) - - # Send complete frame - await websocket.send_json({"type": "complete"}) - print("[WS] Turn completed successfully.") - - except WebSocketDisconnect: - print("[WS] Client disconnected.") - except Exception as e: - logging.exception("Error in WebSocket turn handler") try: - await websocket.send_json({"type": "error", "error": str(e)}) - except Exception: - pass - finally: - await websocket.close() + async with Agent(loaded_config) as agent: + conversation = agent.conversation + + # Hydrate history + print(f"[gRPC] Hydrating {len(historical_messages)} historical messages...") + history_steps = hydrate_ax_history_to_steps(historical_messages) + conversation._steps.extend(history_steps) + + # Run the turn with streaming + print(f"[gRPC] Running chat query: {latest_query_text}") + response = await conversation.chat(latest_query_text) + + async for chunk in response.chunks: + if isinstance(chunk, Text): + msg = ax_pb2.Message( + role="assistant", + content=content_pb2.Content(text=content_pb2.TextContent(text=chunk.text)) + ) + yield ax_pb2.AgentResponse( + conversation_id=request.conversation_id, + exec_id=request.exec_id, + outputs=ax_pb2.AgentOutputs(messages=[msg]) + ) + elif isinstance(chunk, Thought): + summary = [ + content_pb2.ThoughtSummaryContent(text=content_pb2.TextContent(text=chunk.text)) + ] + msg = ax_pb2.Message( + role="model", + content=content_pb2.Content(thought=content_pb2.ThoughtContent(summary=summary)) + ) + yield ax_pb2.AgentResponse( + conversation_id=request.conversation_id, + exec_id=request.exec_id, + outputs=ax_pb2.AgentOutputs(messages=[msg]) + ) + elif isinstance(chunk, ToolCall): + struct_args = Struct() + struct_args.update(chunk.args) + + func_call = content_pb2.FunctionCallContent( + name=str(chunk.name), + arguments=struct_args + ) + msg = ax_pb2.Message( + role="model", + content=content_pb2.Content(tool_call=content_pb2.ToolCallContent( + id=chunk.id or "", + function_call=func_call + )) + ) + yield ax_pb2.AgentResponse( + conversation_id=request.conversation_id, + exec_id=request.exec_id, + outputs=ax_pb2.AgentOutputs(messages=[msg]) + ) + + # Yield completion end frame + yield ax_pb2.AgentResponse( + conversation_id=request.conversation_id, + exec_id=request.exec_id, + end=ax_pb2.AgentEnd() + ) + print("[gRPC] Turn completed successfully.") + + except Exception as e: + logging.exception("Error inside Connect servicer execution") + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details(f"Agent execution terminated due to error. ({str(e)})") + return + + async def HealthCheck(self, request: ax_pb2.HealthCheckRequest, context): + """Simple health-probe responder.""" + return ax_pb2.HealthCheckResponse(healthy=True, message="Antigravity gRPC harness active") + +async def serve(host: str, port: int): + server = grpc.aio.server() + ax_pb2_grpc.add_AgentServiceServicer_to_server(AntigravityAgentServiceServicer(), server) + + listen_addr = f"{host}:{port}" + server.add_insecure_port(listen_addr) + print(f"Starting gRPC harness server on {listen_addr}...") + await server.start() + await server.wait_for_termination() def main(): - parser = argparse.ArgumentParser(description="Antigravity WebSocket Harness Server") + parser = argparse.ArgumentParser(description="Antigravity gRPC Harness Server") parser.add_argument("--agent_file", default="examples/antigravity_agent/agent.py", help="Path to the agent config file") parser.add_argument("--port", type=int, default=50053, help="Port to bind the server to") parser.add_argument("--host", default="localhost", help="Host to bind the server to") @@ -184,8 +217,7 @@ def main(): print(f"ERROR: Failed to load agent config: {e}", file=sys.stderr) sys.exit(1) - print(f"Starting WebSocket server on {args.host}:{args.port}...") - uvicorn.run(app, host=args.host, port=args.port) + asyncio.run(serve(args.host, args.port)) if __name__ == "__main__": main() diff --git a/python/antigravity/harness_server_test.py b/python/antigravity/harness_server_test.py index 290be33..fbff32a 100644 --- a/python/antigravity/harness_server_test.py +++ b/python/antigravity/harness_server_test.py @@ -12,120 +12,81 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import pytest -from fastapi.testclient import TestClient -from unittest.mock import AsyncMock, MagicMock, patch -import json +import grpc +from python.proto import ax_pb2, ax_pb2_grpc, content_pb2 +from python.antigravity.harness_server import AntigravityAgentServiceServicer, loaded_config +from google.antigravity import LocalAgentConfig -from python.antigravity.harness_server import app, hydrate_ax_history_to_steps -from google.antigravity.types import Step, StepType, StepSource, StepTarget, StepStatus, Text, Thought -from proto import ax_pb2 +@pytest.fixture +def mock_config(monkeypatch): + cfg = LocalAgentConfig(system_instructions="Test instructions") + import python.antigravity.harness_server as hs + hs.loaded_config = cfg + return cfg -client = TestClient(app) - -def test_hydrate_ax_history_to_steps(): - # Create mock AX Message protobuf objects - msg = ax_pb2.Message() - msg.role = "user" - msg.content.text.text = "Hi" - - steps = hydrate_ax_history_to_steps([msg]) - - assert len(steps) == 1 - assert steps[0].source == StepSource.USER - assert steps[0].content == "Hi" - assert steps[0].is_complete_response is True - -@patch("python.antigravity.harness_server.Agent") -def test_websocket_endpoint_success(mock_agent_class): - # 1. Setup mocks for Agent, Conversation, and ChatResponse - mock_agent = MagicMock() - mock_agent_class.return_value = mock_agent - - # Mock context manager methods - mock_agent.__aenter__ = AsyncMock(return_value=mock_agent) - mock_agent.__aexit__ = AsyncMock(return_value=None) - - mock_conversation = MagicMock() - mock_agent.conversation = mock_conversation - mock_conversation._steps = [] - - # Mock response stream chunks - async def mock_chunks(): - yield Text(step_index=0, text="Hello ") - yield Text(step_index=1, text="world!") - - mock_chat_response = MagicMock() - mock_chat_response.chunks = mock_chunks() - mock_conversation.chat = AsyncMock(return_value=mock_chat_response) - - # Load a dummy config globally to pass server validation - import python.antigravity.harness_server as server - server.loaded_config = MagicMock() - - # 2. Build start payload - start_payload = { - "conversation_id": "conv-123", - "exec_id": "exec-456", - "messages": [ - # Raw protobuf JSON message - { - "role": "user", - "content": { - "text": {"text": "Hi"} - } - } - ] - } - - # 3. Run WebSocket test client - with client.websocket_connect("/ws") as websocket: - # Send start payload - websocket.send_text(json.dumps(start_payload)) - - # Receive streamed text chunks - resp1 = websocket.receive_json() - assert resp1["type"] == "text" - assert resp1["content"] == "Hello " +def test_grpc_connect_success(mock_config, monkeypatch): + async def _run(): + # 1. Start temporary local gRPC server on random open port + server = grpc.aio.server() + servicer = AntigravityAgentServiceServicer() + ax_pb2_grpc.add_AgentServiceServicer_to_server(servicer, server) + port = server.add_insecure_port("localhost:0") + await server.start() - resp2 = websocket.receive_json() - assert resp2["type"] == "text" - assert resp2["content"] == "world!" - - # Receive complete - resp3 = websocket.receive_json() - assert resp3["type"] == "complete" - - # Verify mocks - mock_conversation.chat.assert_called_once_with("Hi") + # 2. Connect async stub channel + addr = f"localhost:{port}" + async with grpc.aio.insecure_channel(addr) as channel: + stub = ax_pb2_grpc.AgentServiceStub(channel) + + # Mock the underlying Antigravity SDK class calls + class MockConversation: + def __init__(self): + self._steps = [] + async def chat(self, text): + class MockResponse: + def __init__(self): + self.chunks = self._chunk_generator() + async def _chunk_generator(self): + from google.antigravity.types import Text, Thought + yield Thought(text="Thinking details", step_index=0) + yield Text(text="Hello human", step_index=0) + return MockResponse() + + class MockAgent: + def __init__(self, config): + self.conversation = MockConversation() + async def __aenter__(self): + return self + async def __aexit__(self, exc_type, exc, tb): + pass + + monkeypatch.setattr("python.antigravity.harness_server.Agent", MockAgent) + + # 3. Construct and fire standard AgentRequest + start_payload = ax_pb2.AgentStart( + agent_id="test", + messages=[ + ax_pb2.Message(role="user", content=content_pb2.Content(text=content_pb2.TextContent(text="Hi"))) + ] + ) + req = ax_pb2.AgentRequest( + conversation_id="conv-test", + exec_id="exec-test", + start=start_payload + ) + + responses = [] + async for resp in stub.Connect(req): + responses.append(resp) + + # 4. Assert outputs are correctly mapped and completed + assert len(responses) == 3 # Thought + Text + End + assert responses[0].outputs.messages[0].content.thought.summary[0].text.text == "Thinking details" + assert responses[1].outputs.messages[0].content.text.text == "Hello human" + assert responses[2].WhichOneof('type') == 'end' + + await server.stop(0) -@patch("python.antigravity.harness_server.Agent") -def test_websocket_endpoint_error(mock_agent_class): - mock_agent = MagicMock() - mock_agent_class.return_value = mock_agent - mock_agent.__aenter__ = AsyncMock(return_value=mock_agent) - mock_agent.__aexit__ = AsyncMock(return_value=None) - - mock_conversation = MagicMock() - mock_agent.conversation = mock_conversation - mock_conversation._steps = [] - - # Mock chat to throw an exception - mock_conversation.chat = AsyncMock(side_effect=RuntimeError("Gemini connection timeout")) - - import python.antigravity.harness_server as server - server.loaded_config = MagicMock() - - start_payload = { - "conversation_id": "conv-123", - "exec_id": "exec-456", - "messages": [{"role": "user", "content": {"text": {"text": "Hi"}}}] - } - - with client.websocket_connect("/ws") as websocket: - websocket.send_text(json.dumps(start_payload)) - - # Expect error frame - resp = websocket.receive_json() - assert resp["type"] == "error" - assert "Gemini connection timeout" in resp["error"] + asyncio.run(_run()) From b216469c09b12d250fe6ccd1a0d9722403e847fe Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Mon, 1 Jun 2026 11:17:02 -0700 Subject: [PATCH 11/16] docs: Add architectural note comments to agent.py and harness_server.py --- examples/antigravity_agent/agent.py | 11 ++++++++++- internal/controller2/controller.go | 24 ++++++++---------------- python/antigravity/harness_server.py | 6 ++++++ 3 files changed, 24 insertions(+), 17 deletions(-) diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py index bb61e40..f0c7742 100644 --- a/examples/antigravity_agent/agent.py +++ b/examples/antigravity_agent/agent.py @@ -12,7 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +# NOTE ON ARCHITECTURE: +# This file plays a dual-purpose role: +# 1. Standalone Sandbox: Can be run directly via CLI (python agent.py "prompt") for local L2 debugging. +# 2. Declarative Config Module: Exposes 'agent_config' globally, which python/antigravity/harness_server.py +# dynamically imports to serve this agent over production gRPC. + import asyncio + import sys from google.antigravity import LocalAgentConfig from google.antigravity.connections.local import LocalConnectionStrategy @@ -57,7 +64,9 @@ async def main(): # 4. Create the stateful conversation session print("Starting stateful Antigravity conversation (L2 API)...") async with Conversation.create(strategy) as conversation: - prompt = sys.argv[1] if len(sys.argv) > 1 else "What is the weather in New York?" + prompt = sys.argv[1] if len(sys.argv) > 1 else None + if not prompt: + raise ValueError("Please provide a prompt for your agent. Usage: python agent.py ") # 5. Send query and receive streaming ChatResponse response = await conversation.chat(prompt) diff --git a/internal/controller2/controller.go b/internal/controller2/controller.go index 086b41d..6095515 100644 --- a/internal/controller2/controller.go +++ b/internal/controller2/controller.go @@ -19,7 +19,6 @@ package controller2 import ( "context" "fmt" - "log" "github.com/google/ax/internal/controller/executor" "github.com/google/ax/internal/harness/harnesstest" @@ -32,21 +31,20 @@ 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) { - if cfg.Registry == nil { - return nil, fmt.Errorf("registry is required") - } + // Initialize agent registry + registry := NewRegistry() + if cfg.EventLogBuilder == nil { return nil, fmt.Errorf("event log builder is required") } @@ -56,8 +54,8 @@ func New(ctx context.Context, cfg Config) (*Controller, error) { } return &Controller{ - registry: cfg.Registry, - eventLog: eventLog, + registry: registry, + eventLog: eventLog, }, nil } @@ -72,13 +70,7 @@ 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. - // 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() - } + 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/python/antigravity/harness_server.py b/python/antigravity/harness_server.py index 03b4af9..04f65f6 100644 --- a/python/antigravity/harness_server.py +++ b/python/antigravity/harness_server.py @@ -12,7 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +# NOTE ON ARCHITECTURE: +# This is a generic, reusable gRPC server that does not define tools or personas. +# Instead, it dynamically imports any agent configuration file (defaulting to examples/antigravity_agent/agent.py) +# passed via the --agent_file CLI argument, then hosts it over the AX AgentService protocol. + import argparse + import asyncio import importlib.util import logging From 300cb354d7cd4d71e4827c7e96d7665e5772aa9e Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Mon, 1 Jun 2026 11:27:49 -0700 Subject: [PATCH 12/16] docs: Add E2E setup and running instructions header to cmd/e2e/main.go --- cmd/e2e/main.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cmd/e2e/main.go b/cmd/e2e/main.go index 0e98a1c..374c894 100644 --- a/cmd/e2e/main.go +++ b/cmd/e2e/main.go @@ -14,8 +14,17 @@ // Package main implements an end-to-end demonstration of the Antigravity harness // integration with AX Controller V2. +// +// TO RUN THIS E2E DEMONSTRATION: +// +// Step 1: Start the Python gRPC Harness Server (in a separate terminal or background): +// PYTHONPATH=python:. /Users/anjalisridhar/.gemini/jetski/worktrees/harness-interface-3/implement-agy-sdk-streaming-20260528/.venv/bin/python python/antigravity/harness_server.py --port 50053 +// +// Step 2: Run this Go E2E client: +// go run cmd/e2e/main.go package main + import ( "context" "fmt" From c1e2b6959b3651ab79b46071e8d7aa8887e69f84 Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Mon, 1 Jun 2026 11:28:32 -0700 Subject: [PATCH 13/16] feat: Upgrade agent.py to support an interactive multi-turn console chat loop --- examples/antigravity_agent/agent.py | 69 +++++++++++++++++++++-------- 1 file changed, 50 insertions(+), 19 deletions(-) diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py index f0c7742..9967618 100644 --- a/examples/antigravity_agent/agent.py +++ b/examples/antigravity_agent/agent.py @@ -62,28 +62,59 @@ async def main(): strategy = strategy_factory() # 4. Create the stateful conversation session - print("Starting stateful Antigravity conversation (L2 API)...") async with Conversation.create(strategy) as conversation: - prompt = sys.argv[1] if len(sys.argv) > 1 else None - if not prompt: - raise ValueError("Please provide a prompt for your agent. Usage: python agent.py ") + # Check if a prompt was passed via CLI arguments (single-turn compatibility mode) + cli_prompt = sys.argv[1] if len(sys.argv) > 1 else None - # 5. Send query and receive streaming ChatResponse - response = await conversation.chat(prompt) + if cli_prompt: + await run_turn(conversation, cli_prompt) + print() + return + + # Multi-turn interactive chat mode + print("=" * 50) + print("Antigravity Interactive Console Chat (L2 API)") + print("Type 'exit' or 'quit' to end the conversation.") + print("=" * 50) - # 6. Stream semantic chunks (Thoughts, Text, and ToolCalls) in real-time - async for chunk in response.chunks: - if isinstance(chunk, Text): - sys.stdout.write(chunk.text) - sys.stdout.flush() - elif isinstance(chunk, Thought): - # Display thought process in comment style - sys.stdout.write(f"\n[Thinking]: {chunk.text}") - sys.stdout.flush() - elif isinstance(chunk, ToolCall): - sys.stdout.write(f"\n[Tool Call]: {chunk.name} with args {chunk.args}\n") - sys.stdout.flush() - print() + while True: + try: + # Standard Python input prompt + user_input = input("\nYou: ") + except (EOFError, KeyboardInterrupt): + print("\nGoodbye!") + break + + clean_input = user_input.strip() + if not clean_input: + continue + + if clean_input.lower() in ("exit", "quit"): + print("Goodbye!") + break + + sys.stdout.write("Agent: ") + sys.stdout.flush() + await run_turn(conversation, clean_input) + sys.stdout.write("\n") + +async def run_turn(conversation, prompt): + # Send query and receive streaming ChatResponse + response = await conversation.chat(prompt) + + # Stream semantic chunks (Thoughts, Text, and ToolCalls) in real-time + async for chunk in response.chunks: + if isinstance(chunk, Text): + sys.stdout.write(chunk.text) + sys.stdout.flush() + elif isinstance(chunk, Thought): + # Display thought process in comment style + sys.stdout.write(f"\n[Thinking]: {chunk.text}") + sys.stdout.flush() + elif isinstance(chunk, ToolCall): + sys.stdout.write(f"\n[Tool Call]: {chunk.name} with args {chunk.args}\n") + sys.stdout.flush() if __name__ == "__main__": asyncio.run(main()) + From 598976af7492fb85481d65ad4b7ffd94ef93b89d Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Mon, 1 Jun 2026 11:33:00 -0700 Subject: [PATCH 14/16] revert: Revert interactive multi-turn console chat loop to preserve single-turn subprocess execution turn compatibility --- examples/antigravity_agent/agent.py | 69 +++++++++-------------------- 1 file changed, 20 insertions(+), 49 deletions(-) diff --git a/examples/antigravity_agent/agent.py b/examples/antigravity_agent/agent.py index 9967618..b5bf538 100644 --- a/examples/antigravity_agent/agent.py +++ b/examples/antigravity_agent/agent.py @@ -62,59 +62,30 @@ async def main(): strategy = strategy_factory() # 4. Create the stateful conversation session + print("Starting stateful Antigravity conversation (L2 API)...") async with Conversation.create(strategy) as conversation: - # Check if a prompt was passed via CLI arguments (single-turn compatibility mode) - cli_prompt = sys.argv[1] if len(sys.argv) > 1 else None + prompt = sys.argv[1] if len(sys.argv) > 1 else None + if not prompt: + raise ValueError("Please provide a prompt for your agent. Usage: python agent.py ") - if cli_prompt: - await run_turn(conversation, cli_prompt) - print() - return - - # Multi-turn interactive chat mode - print("=" * 50) - print("Antigravity Interactive Console Chat (L2 API)") - print("Type 'exit' or 'quit' to end the conversation.") - print("=" * 50) + # 5. Send query and receive streaming ChatResponse + response = await conversation.chat(prompt) - while True: - try: - # Standard Python input prompt - user_input = input("\nYou: ") - except (EOFError, KeyboardInterrupt): - print("\nGoodbye!") - break - - clean_input = user_input.strip() - if not clean_input: - continue - - if clean_input.lower() in ("exit", "quit"): - print("Goodbye!") - break - - sys.stdout.write("Agent: ") - sys.stdout.flush() - await run_turn(conversation, clean_input) - sys.stdout.write("\n") - -async def run_turn(conversation, prompt): - # Send query and receive streaming ChatResponse - response = await conversation.chat(prompt) - - # Stream semantic chunks (Thoughts, Text, and ToolCalls) in real-time - async for chunk in response.chunks: - if isinstance(chunk, Text): - sys.stdout.write(chunk.text) - sys.stdout.flush() - elif isinstance(chunk, Thought): - # Display thought process in comment style - sys.stdout.write(f"\n[Thinking]: {chunk.text}") - sys.stdout.flush() - elif isinstance(chunk, ToolCall): - sys.stdout.write(f"\n[Tool Call]: {chunk.name} with args {chunk.args}\n") - sys.stdout.flush() + # 6. Stream semantic chunks (Thoughts, Text, and ToolCalls) in real-time + async for chunk in response.chunks: + if isinstance(chunk, Text): + sys.stdout.write(chunk.text) + sys.stdout.flush() + elif isinstance(chunk, Thought): + # Display thought process in comment style + sys.stdout.write(f"\n[Thinking]: {chunk.text}") + sys.stdout.flush() + elif isinstance(chunk, ToolCall): + sys.stdout.write(f"\n[Tool Call]: {chunk.name} with args {chunk.args}\n") + sys.stdout.flush() + print() if __name__ == "__main__": asyncio.run(main()) + From a0e96f590ae56cbc2d86313f4d0a628c5516768f Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Mon, 1 Jun 2026 12:31:24 -0700 Subject: [PATCH 15/16] revert: Revert comment above RegisterHarness to keep the TODO for registry consolidation --- internal/controller2/registry.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/controller2/registry.go b/internal/controller2/registry.go index c77ba82..0296a12 100644 --- a/internal/controller2/registry.go +++ b/internal/controller2/registry.go @@ -240,7 +240,7 @@ func (r *Registry) Close() error { return firstErr } -// RegisterHarness registers a harness. +// TODO(anj): Remove this registration once we have harness and agent registration consolidated. func (r *Registry) RegisterHarness(id string, h harness.Harness) { r.mu.Lock() defer r.mu.Unlock() From 60830d7bd24e65540a635b4772300ecdbf9a887d Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Mon, 1 Jun 2026 14:51:49 -0700 Subject: [PATCH 16/16] refactor: remove obsolete local script path checks in E2E main --- cmd/e2e/main.go | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/cmd/e2e/main.go b/cmd/e2e/main.go index 374c894..676cf32 100644 --- a/cmd/e2e/main.go +++ b/cmd/e2e/main.go @@ -56,33 +56,16 @@ func main() { }) // ------------------------------------------------------------------------- - // Demo 2: Build-time Fallback (Antigravity with bad script path) + // Demo 2: Antigravity Execution (Requires google-antigravity & GEMINI_API_KEY) // ------------------------------------------------------------------------- - 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 := 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("ws://localhost:50054/ws") - } - reg.RegisterHarness("antigravity", badHarness) - }) - - // ------------------------------------------------------------------------- - // Demo 3: Antigravity Execution (Requires google-antigravity & GEMINI_API_KEY) - // ------------------------------------------------------------------------- - fmt.Println("\n--- Demo 3: Antigravity Execution ---") + fmt.Println("\n--- Demo 2: 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) { - // Check if Python gRPC server is active, otherwise fallback + // With the new stateful gRPC-based streaming harness, connectivity checks on the + // server address replace the build-time checks for local script file presence. var realHarness harness.Harness address := "localhost:50053" conn, err := net.DialTimeout("tcp", address, 1*time.Second)