diff --git a/e2e.go b/e2e.go new file mode 100644 index 0000000..2a24f97 --- /dev/null +++ b/e2e.go @@ -0,0 +1,124 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package main implements an end-to-end demonstration of the Antigravity harness +// integration with AX Controller V2. +package main + +import ( + "context" + "fmt" + "os" + + "github.com/google/ax/internal/controller/executor" + "github.com/google/ax/internal/controller/executor/executortest" + "github.com/google/ax/internal/controller2" + "github.com/google/ax/internal/harness" + "github.com/google/ax/proto" +) + +func main() { + ctx := context.Background() + fmt.Println("==================================================") + fmt.Println("AX Controller V2 - E2E Harness Demonstration") + fmt.Println("==================================================") + + // ------------------------------------------------------------------------- + // Demo 1: Runtime Fallback (No harness registered) + // ------------------------------------------------------------------------- + fmt.Println("\n--- Demo 1: Runtime Fallback ---") + fmt.Println("Requesting 'unregistered-agent'. Should fallback to Test Harness (Hello World).") + runDemo(ctx, "unregistered-agent", func(reg *controller2.Registry) { + // Do not register any harness + }) + + // ------------------------------------------------------------------------- + // Demo 2: Build-time Fallback (Antigravity with bad script path) + // ------------------------------------------------------------------------- + fmt.Println("\n--- Demo 2: Build-time Fallback ---") + fmt.Println("Registering 'antigravity' with non-existent script. Should fallback to Test Harness.") + runDemo(ctx, "antigravity", func(reg *controller2.Registry) { + // Build harness with bad path + badHarness := controller2.BuildHarness(ctx, "antigravity", harness.HarnessConfig{ + AntigravityScriptPath: "non-existent-script.py", + }) + reg.RegisterHarness("antigravity", badHarness) + }) + + // ------------------------------------------------------------------------- + // Demo 3: Antigravity Execution (Requires google-antigravity & GEMINI_API_KEY) + // ------------------------------------------------------------------------- + fmt.Println("\n--- Demo 3: Antigravity Execution ---") + fmt.Println("Registering 'antigravity' with real script. Attempting execution.") + if os.Getenv("GEMINI_API_KEY") == "" { + fmt.Println("WARNING: GEMINI_API_KEY is not set. Execution will likely fail if dependencies are missing, but we will try anyway.") + } + runDemo(ctx, "antigravity", func(reg *controller2.Registry) { + // Build harness with real path (empty defaults to examples/antigravity_agent/agent.py) + realHarness := controller2.BuildHarness(ctx, "antigravity", harness.HarnessConfig{ + AntigravityScriptPath: "examples/antigravity_agent/agent.py", + }) + reg.RegisterHarness("antigravity", realHarness) + }) +} + +func runDemo(ctx context.Context, agentID string, setupRegistry func(reg *controller2.Registry)) { + reg := controller2.NewRegistry() + setupRegistry(reg) + + log := &executortest.MemoryEventLog{} + c, err := controller2.New(ctx, controller2.Config{ + Registry: reg, + EventLogBuilder: func() (executor.EventLog, error) { + return log, nil + }, + }) + if err != nil { + fmt.Printf("Error creating controller: %v\n", err) + return + } + defer c.Close() + + handler := controller2.ExecHandler(func(resp *proto.ExecResponse) error { + for _, out := range resp.Outputs { + if textContent := out.GetContent().GetText().GetText(); textContent != "" { + fmt.Printf("Agent Output: %s\n", textContent) + } + } + return nil + }) + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "Who are you?"}, + }, + }, + }, + } + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: "e2e-conv", + Inputs: inputs, + AgentId: agentID, + }, handler) + + if err != nil { + fmt.Printf("Execution Failed (as expected if environment is not ready): %v\n", err) + } else { + fmt.Println("Execution Succeeded!") + } +} 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/config/config.go b/internal/config/config.go index 49cfe0b..e15e21c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -39,6 +39,17 @@ type RegistryConfig struct { RemoteAgents []RemoteAgentConfig `yaml:"remote_agents,omitempty"` ColabAgents []ColabAgentConfig `yaml:"colab_agents,omitempty"` SubstrateAgents []SubstrateAgentConfig `yaml:"substrate_agents,omitempty"` + Harnesses []HarnessConfig `yaml:"harnesses,omitempty"` +} + +type HarnessConfig struct { + ID string `yaml:"id"` + Type string `yaml:"type"` // "antigravity" + Antigravity AntigravityConfig `yaml:"antigravity,omitempty"` +} + +type AntigravityConfig struct { + ScriptPath string `yaml:"script_path"` } // ATEConfig configures the ATE integration. diff --git a/internal/controller2/builder.go b/internal/controller2/builder.go new file mode 100644 index 0000000..1666ecc --- /dev/null +++ b/internal/controller2/builder.go @@ -0,0 +1,57 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controller2 + +import ( + "context" + "log" + "os" + "os/exec" + + "github.com/google/ax/internal/harness" + "github.com/google/ax/internal/harness/harnesstest" +) + +// BuildHarness builds a harness based on the requested type, with fallback to test harness. +func BuildHarness(ctx context.Context, harnessType string, cfg harness.HarnessConfig) harness.Harness { + switch harnessType { + case "antigravity": + // Check if python3 is available + if _, err := exec.LookPath("python3"); err != nil { + log.Printf("WARNING: python3 not found in PATH, falling back to test harness: %v", err) + return harnesstest.New() + } + // Check if script exists + scriptPath := cfg.AntigravityScriptPath + if scriptPath == "" { + scriptPath = "examples/antigravity_agent/agent.py" + } + if _, err := os.Stat(scriptPath); err != nil { + log.Printf("WARNING: Antigravity agent script not found at %s, falling back to test harness: %v", scriptPath, err) + return harnesstest.New() + } + log.Printf("Using Antigravity harness with script: %s", scriptPath) + + builder := &harness.AntigravityHarnessBuilder{ + Config: harness.HarnessConfig{ + AntigravityScriptPath: scriptPath, + }, + } + return builder.Build() + default: + log.Printf("Using default test harness") + return harnesstest.New() + } +} 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..fa73157 100644 --- a/internal/controller2/controller_test.go +++ b/internal/controller2/controller_test.go @@ -18,26 +18,135 @@ import ( "context" "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" ) -type dummyAgent struct{} +func TestController2_ExecHelloWorld(t *testing.T) { + ctx := context.Background() + cid := "test-conversation-id" -func (a *dummyAgent) Connect(ctx context.Context, conversationID string, execID string, start *proto.AgentStart, e agent.Executor, o agent.OutputHandler) error { - return nil + log := &executortest.MemoryEventLog{} + reg := NewRegistry() + 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"}, + }, + }, + }, + } + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: inputs, + }, 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, got %q", gotText) + } } -func (a *dummyAgent) Close() error { return nil } +func TestController2_ExecAntigravityFallback(t *testing.T) { + ctx := context.Background() + cid := "test-conversation-id" -func TestController2_ExecHelloWorld(t *testing.T) { + log := &executortest.MemoryEventLog{} + reg := NewRegistry() + + // Build and register harness with bad path to trigger build-time fallback + badHarness := BuildHarness(ctx, "antigravity", harness.HarnessConfig{ + AntigravityScriptPath: "non-existent-script.py", + }) + 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 }, @@ -64,9 +173,11 @@ func TestController2_ExecHelloWorld(t *testing.T) { }, } + // 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) @@ -78,6 +189,132 @@ func TestController2_ExecHelloWorld(t *testing.T) { gotText := outputs[0].GetContent().GetText().GetText() if gotText != "Hello world" { - t.Errorf("expected 'Hello world' output text response, got %q", gotText) + t.Errorf("expected 'Hello world' output text response due to runtime fallback, got %q", gotText) + } +} + +func TestController2_Exec_ResumptionAndIDGeneration(t *testing.T) { + t.Skip("Feature Gap: Resumption and Event Logging are not yet implemented in controller2") + + // This test is sketched out for when the features are implemented. + ctx := context.Background() + cid := "test-conv" + + inputs := []*proto.Message{ + { + Role: "user", + Content: &proto.Content{ + Type: &proto.Content_Text{ + Text: &proto.TextContent{Text: "hello"}, + }, + }, + }, + } + + log := &executortest.MemoryEventLog{} + reg := NewRegistry() + mockHarness := harnesstest.New() + reg.RegisterHarness("mock-agent", mockHarness) + + c, err := New(ctx, Config{ + Registry: reg, + EventLogBuilder: func() (executor.EventLog, error) { + return log, nil + }, + }) + if err != nil { + t.Fatal(err) + } + defer c.Close() + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + Inputs: inputs, + AgentId: "mock-agent", + }, nil) + if err != nil { + t.Fatal(err) } + + // Verify that events were logged and exec ID was generated. + // ... +} + +func TestController2_Exec_LastSeq_Empty(t *testing.T) { + t.Skip("Feature Gap: History playback and LastSeq are not yet implemented in controller2") +} + +func TestController2_Exec_LastSeq(t *testing.T) { + t.Skip("Feature Gap: History playback and LastSeq are not yet implemented in controller2") +} + +func TestController2_Exec_LastSeq_NotFound(t *testing.T) { + t.Skip("Feature Gap: History playback and LastSeq are not yet implemented in controller2") +} + +func TestController2_Exec_WaitsForConfirmation(t *testing.T) { + t.Skip("Feature Gap: Resumption and Confirmation handling are not yet implemented in controller2") + + // Sketch for when implemented: + ctx := context.Background() + cid := "test-conv-conf" + + log := &executortest.MemoryEventLog{} + reg := NewRegistry() + mockHarness := harnesstest.New() + + // Configure mock harness to return a confirmation question on first Run + mockHarness.DefaultRunFunc = func(ctx context.Context, execID string, handler harness.Handler) error { + questionMsg := &proto.Message{ + Role: "assistant", + Content: &proto.Content{ + Type: &proto.Content_Confirmation{ + Confirmation: &proto.ConfirmationContent{ + Question: "Are you sure?", + }, + }, + }, + } + if err := handler.OnMessage(ctx, execID, questionMsg); err != nil { + return err + } + return handler.OnComplete(ctx, execID) + } + reg.RegisterHarness("mock-agent", mockHarness) + + 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 msgs []*proto.Message + handler := ExecHandler(func(resp *proto.ExecResponse) error { + msgs = append(msgs, resp.Outputs...) + return nil + }) + + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + AgentId: "mock-agent", + }, handler) + if err != nil { + t.Fatal(err) + } + + if len(msgs) != 1 { + t.Fatalf("expected 1 message, got %d", len(msgs)) + } + if msgs[0].GetContent().GetConfirmation().GetQuestion() != "Are you sure?" { + t.Fatalf("expected 'Are you sure?', got %v", msgs[0].GetContent().GetConfirmation().GetQuestion()) + } +} + +func TestController2_Exec_InternalOnly(t *testing.T) { + t.Skip("Feature Gap: Event Logging and InternalOnly filtering are not yet implemented in controller2") } 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..274178e --- /dev/null +++ b/internal/harness/antigravity.go @@ -0,0 +1,158 @@ +// 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, + } +} + +// AntigravityHarnessBuilder builds an AntigravityHarness using HarnessConfig. +type AntigravityHarnessBuilder struct { + Config HarnessConfig +} + +// Build constructs the AntigravityHarness. +func (b *AntigravityHarnessBuilder) Build() Harness { + return NewAntigravityHarness(b.Config.AntigravityScriptPath) +} + + +// 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 +} diff --git a/internal/harness/config.go b/internal/harness/config.go new file mode 100644 index 0000000..df244c1 --- /dev/null +++ b/internal/harness/config.go @@ -0,0 +1,21 @@ +// 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 + +// HarnessConfig holds configuration options for various harnesses. +type HarnessConfig struct { + AntigravityScriptPath string + // Future configuration options for other harnesses can be added here. +} diff --git a/internal/harness/harnesstest/harnesstest.go b/internal/harness/harnesstest/harnesstest.go index 77990b9..1eb7ef6 100644 --- a/internal/harness/harnesstest/harnesstest.go +++ b/internal/harness/harnesstest/harnesstest.go @@ -27,8 +27,9 @@ import ( // Harness implements the harness.Harness interface for testing purposes. type Harness struct { - mu sync.Mutex - Active map[string]*Execution + mu sync.Mutex + Active map[string]*Execution + DefaultRunFunc func(ctx context.Context, execID string, handler harness.Handler) error } // New creates a new Harness instance. @@ -48,6 +49,7 @@ func (h *Harness) Start(ctx context.Context, conversationID string) (harness.Exe harness: h, conversationID: conversationID, id: execID, + RunFunc: h.DefaultRunFunc, } h.Active[execID] = exec return exec, nil @@ -59,9 +61,10 @@ type Execution struct { conversationID string id string - mu sync.Mutex - queued []*proto.Message - closed bool + mu sync.Mutex + queued []*proto.Message + closed bool + RunFunc func(ctx context.Context, execID string, handler harness.Handler) error } // ID implements harness.Execution. @@ -78,8 +81,12 @@ func (e *Execution) Queue(ctx context.Context, msg ...*proto.Message) error { } // Run implements harness.Execution. -// It generates a "Hello world" message and completes the turn. +// It uses the configured RunFunc, or falls back to generating a "Hello world" message. func (e *Execution) Run(ctx context.Context, handler harness.Handler) error { + if e.RunFunc != nil { + return e.RunFunc(ctx, e.id, handler) + } + msg := &proto.Message{ Role: "assistant", Content: &proto.Content{