From 34becea058964bfc07cce49dc578c5618a527e0a Mon Sep 17 00:00:00 2001 From: Junjie Wang Date: Wed, 3 Jun 2026 18:48:15 -0700 Subject: [PATCH 1/5] Add config specific to harness and controllerv2. --- Makefile | 4 +- cmd/ax/fork.go | 3 +- cmd/ax/internal/cliutil/cliutil.go | 15 +- cmd/ax/internal/cliutil/cliutil_harness.go | 52 +++- .../internal/cliutil/cliutil_harness_test.go | 70 ++++++ cmd/ax/main.go | 8 +- cmd/ax/trace.go | 6 +- internal/ax2.yaml | 31 +++ .../config/harnessconfig/harness_config.go | 124 ++++++++++ internal/controller2/controller_test.go | 17 +- internal/controller2/registry.go | 225 ++---------------- internal/controller2/registry_ate.go | 62 ----- internal/controller2/registry_test.go | 75 +++--- internal/controller2/validation.go | 14 +- internal/controller2/validation_test.go | 20 -- internal/manifests/ax-deployment2.yaml | 26 +- 16 files changed, 377 insertions(+), 375 deletions(-) create mode 100644 cmd/ax/internal/cliutil/cliutil_harness_test.go create mode 100644 internal/ax2.yaml create mode 100644 internal/config/harnessconfig/harness_config.go delete mode 100644 internal/controller2/registry_ate.go diff --git a/Makefile b/Makefile index 15425d4..dfcb6da 100644 --- a/Makefile +++ b/Makefile @@ -31,10 +31,12 @@ proto: proto/ax.proto proto/content.proto @echo "Protobuf generation complete!" -# Run tests +# Run tests. test: @echo "Running tests..." @go test -v ./... + @echo "Running tests (harness path)..." + @go test -v -tags harness ./... # Clean build artifacts clean: diff --git a/cmd/ax/fork.go b/cmd/ax/fork.go index 217df65..9210b4c 100644 --- a/cmd/ax/fork.go +++ b/cmd/ax/fork.go @@ -18,7 +18,6 @@ import ( "fmt" "github.com/google/ax/cmd/ax/internal/cliutil" - "github.com/google/ax/internal/config" "github.com/google/ax/proto" "github.com/spf13/cobra" ) @@ -54,7 +53,7 @@ func runFork(cmd *cobra.Command, args []string) error { if forkServerAddr == "" { // Headless mode - cfg, err := config.LoadFromFile(forkConfigFile) + cfg, err := cliutil.LoadFromFile(forkConfigFile) if err != nil { return fmt.Errorf("error loading config file '%s': %w", forkConfigFile, err) } diff --git a/cmd/ax/internal/cliutil/cliutil.go b/cmd/ax/internal/cliutil/cliutil.go index 3f4cb99..9b1c1f9 100644 --- a/cmd/ax/internal/cliutil/cliutil.go +++ b/cmd/ax/internal/cliutil/cliutil.go @@ -34,8 +34,21 @@ type Controller = *controller.Controller // ExecHandler is the handler type accepted by Controller.Exec. type ExecHandler = controller.ExecHandler +// Config is the configuration type for this build. +type Config = config.Config + +// LoadFromFile loads configuration from a YAML file. +func LoadFromFile(path string) (*Config, error) { + return config.LoadFromFile(path) +} + +// DefaultConfig returns a configuration with default values set. +func DefaultConfig() *Config { + return config.DefaultConfig() +} + // NewControllerFromConfig creates a new Controller instance based on the provided configuration. -func NewControllerFromConfig(ctx context.Context, cfg *config.Config) (*controller.Controller, error) { +func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Controller, error) { // Validate planner type early switch cfg.Planner.Type { case "gemini": diff --git a/cmd/ax/internal/cliutil/cliutil_harness.go b/cmd/ax/internal/cliutil/cliutil_harness.go index 11591b6..38eda48 100644 --- a/cmd/ax/internal/cliutil/cliutil_harness.go +++ b/cmd/ax/internal/cliutil/cliutil_harness.go @@ -18,8 +18,9 @@ package cliutil import ( "context" + "fmt" - "github.com/google/ax/internal/config" + "github.com/google/ax/internal/config/harnessconfig" "github.com/google/ax/internal/controller/executor" "github.com/google/ax/internal/controller2" "github.com/google/ax/internal/harness" @@ -31,10 +32,53 @@ type Controller = *controller2.Controller // ExecHandler is the handler type accepted by Controller.Exec. type ExecHandler = controller2.ExecHandler -// NewControllerFromConfig creates a controller2.Controller. -func NewControllerFromConfig(ctx context.Context, cfg *config.Config) (*controller2.Controller, error) { +// Config is the configuration type for this build. +type Config = harnessconfig.Config + +// LoadFromFile loads configuration from a YAML file. +func LoadFromFile(path string) (*Config, error) { + return harnessconfig.LoadFromFile(path) +} + +// DefaultConfig returns a configuration with default values set. +func DefaultConfig() *Config { + return harnessconfig.DefaultConfig() +} + +// NewControllerFromConfig creates a controller2.Controller instance based on the provided configuration. +func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller2.Controller, error) { reg := controller2.NewRegistry() - reg.RegisterHarness("antigravity", harness.NewAntigravityHarness("")) + + // Antigravity harnesses. + for _, hc := range cfg.Harnesses.Antigravity { + h := harness.NewAntigravityHarness(hc.Address) + if err := reg.RegisterHarness(hc.ID, h); err != nil { + return nil, fmt.Errorf("register antigravity harness %q: %w", hc.ID, err) + } + } + + // Substrate harnesses. + for _, sc := range cfg.Harnesses.Substrate { + sh, err := harness.NewSubstrateHarness(cfg.ATE.Endpoint, sc.Namespace, sc.Template, sc.Port) + if err != nil { + return nil, fmt.Errorf("substrate harness %q: %w", sc.ID, err) + } + if err := reg.RegisterHarness(sc.ID, sh); err != nil { + return nil, fmt.Errorf("register substrate harness %q: %w", sc.ID, err) + } + } + + // Register the configured default harness. + if id := cfg.Harnesses.Default; id != "" { + h, err := reg.Harness(id) + if err != nil { + return nil, fmt.Errorf("default harness %q not found", id) + } + if err := reg.RegisterHarness("", h); err != nil { + return nil, fmt.Errorf("register default harness %q: %w", id, err) + } + } + return controller2.New(ctx, controller2.Config{ Registry: reg, EventLogBuilder: func() (executor.EventLog, error) { diff --git a/cmd/ax/internal/cliutil/cliutil_harness_test.go b/cmd/ax/internal/cliutil/cliutil_harness_test.go new file mode 100644 index 0000000..a33d701 --- /dev/null +++ b/cmd/ax/internal/cliutil/cliutil_harness_test.go @@ -0,0 +1,70 @@ +//go:build harness + +// 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 cliutil + +import ( + "context" + "path/filepath" + "strings" + "testing" + + "github.com/google/ax/internal/config/harnessconfig" +) + +func TestNewControllerFromConfig_DefaultHarness(t *testing.T) { + cfg := &harnessconfig.Config{ + EventLog: harnessconfig.EventLogConfig{ + SQLiteConfig: harnessconfig.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: harnessconfig.HarnessesConfig{ + Default: "ag", + Antigravity: []harnessconfig.AntigravityHarnessConfig{ + {ID: "ag", Address: "localhost:50053"}, + }, + }, + } + + c, err := NewControllerFromConfig(context.Background(), cfg) + if err != nil { + t.Fatalf("NewControllerFromConfig: %v", err) + } + if c == nil { + t.Fatal("expected non-nil controller") + } + c.Close() +} + +func TestNewControllerFromConfig_UnknownDefaultHarness(t *testing.T) { + cfg := &harnessconfig.Config{ + Harnesses: harnessconfig.HarnessesConfig{ + Default: "missing", + Antigravity: []harnessconfig.AntigravityHarnessConfig{ + {ID: "ag", Address: "localhost:50053"}, + }, + }, + } + + _, err := NewControllerFromConfig(context.Background(), cfg) + if err == nil { + t.Fatal("expected error for unknown default harness, got nil") + } + if !strings.Contains(err.Error(), "missing") { + t.Errorf("expected error to mention %q, got: %v", "missing", err) + } +} diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 3d6348c..1114917 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -22,7 +22,7 @@ import ( "fmt" "os" - "github.com/google/ax/internal/config" + "github.com/google/ax/cmd/ax/internal/cliutil" "github.com/spf13/cobra" "google.golang.org/grpc" "google.golang.org/grpc/credentials/insecure" @@ -61,10 +61,10 @@ func connect(server string) (*grpc.ClientConn, error) { return conn, nil } -func newConfig(cmd *cobra.Command, configFile string) (*config.Config, error) { - cfg, err := config.LoadFromFile(configFile) +func newConfig(cmd *cobra.Command, configFile string) (*cliutil.Config, error) { + cfg, err := cliutil.LoadFromFile(configFile) if errors.Is(err, os.ErrNotExist) && !cmd.Flags().Changed("config") { - return config.DefaultConfig(), nil + return cliutil.DefaultConfig(), nil } if err != nil { return nil, fmt.Errorf("error loading config file '%s': %w", configFile, err) diff --git a/cmd/ax/trace.go b/cmd/ax/trace.go index 1e368ff..a95f2d5 100644 --- a/cmd/ax/trace.go +++ b/cmd/ax/trace.go @@ -24,7 +24,7 @@ import ( "runtime" "time" - "github.com/google/ax/internal/config" + "github.com/google/ax/cmd/ax/internal/cliutil" "github.com/google/ax/internal/controller/executor" "github.com/google/ax/proto" "github.com/spf13/cobra" @@ -115,7 +115,7 @@ type TraceData struct { Execs []ExecTrace `json:"execs"` } -func loadTraceData(ctx context.Context, cfg *config.Config, convID string) (*TraceData, error) { +func loadTraceData(ctx context.Context, cfg *cliutil.Config, convID string) (*TraceData, error) { events, rootExecID, execIDs, err := fetch(ctx, cfg, convID) if err != nil { return nil, err @@ -132,7 +132,7 @@ func loadTraceData(ctx context.Context, cfg *config.Config, convID string) (*Tra }, nil } -func fetch(ctx context.Context, cfg *config.Config, convID string) ([]*proto.ExecutionEvent, string, []string, error) { +func fetch(ctx context.Context, cfg *cliutil.Config, convID string) ([]*proto.ExecutionEvent, string, []string, error) { evLog, err := executor.OpenSQLiteEventLog(cfg.EventLog.SQLiteConfig.Filename) if err != nil { return nil, "", nil, fmt.Errorf("could not open sqlite eventlog: %w", err) diff --git a/internal/ax2.yaml b/internal/ax2.yaml new file mode 100644 index 0000000..fa76559 --- /dev/null +++ b/internal/ax2.yaml @@ -0,0 +1,31 @@ +# 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. + +# Sample LOCAL configuration for the AX *harness* build (compiled with -tags=harness). +# Usage: ax serve --config internal/ax2.yaml +# ax exec --config internal/ax2.yaml --input "hello" # uses harnesses.default + +server: + address: ":8494" + +eventlog: + sqlite: + filename: "eventlog/log.sqlite" + +harnesses: + default: antigravity-example + + antigravity: + - id: antigravity-example + address: "localhost:50053" diff --git a/internal/config/harnessconfig/harness_config.go b/internal/config/harnessconfig/harness_config.go new file mode 100644 index 0000000..1d98e84 --- /dev/null +++ b/internal/config/harnessconfig/harness_config.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 harnessconfig provides configuration structures for the AX harness build. +package harnessconfig + +import ( + "fmt" + "os" + + "gopkg.in/yaml.v3" +) + +// Config represents the main configuration for the AX harness server. +type Config struct { + Server ServerConfig `yaml:"server"` + EventLog EventLogConfig `yaml:"eventlog"` + Harnesses HarnessesConfig `yaml:"harnesses,omitempty"` + ATE ATEConfig `yaml:"ate,omitempty"` +} + +// ServerConfig configures the gRPC server. +type ServerConfig struct { + Address string `yaml:"address"` // Server address to listen on (e.g., ":8494") +} + +// SQLiteConfig configures the SQLite event log file. +type SQLiteConfig struct { + Filename string `yaml:"filename"` // SQLite file for event log storage +} + +// EventLogConfig configures the event log storage. +type EventLogConfig struct { + SQLiteConfig SQLiteConfig `yaml:"sqlite"` +} + +// ATEConfig configures the substrate control plane endpoint used by +// substrate harnesses. +type ATEConfig struct { + Endpoint string `yaml:"endpoint"` +} + +// HarnessesConfig groups harnesses to serve by type. Each type maps to a list +// of that type's configurations. +type HarnessesConfig struct { + // Default is the id of the harness to serve when a request specifies no harness. + Default string `yaml:"default,omitempty"` + // Antigravity harnesses connect to a gRPC server at a fixed address. + Antigravity []AntigravityHarnessConfig `yaml:"antigravity,omitempty"` + // Substrate harnesses are brought up as SubstrATE actors from an ActorTemplate. + Substrate []SubstrateHarnessConfig `yaml:"substrate,omitempty"` +} + +// AntigravityHarnessConfig registers an Antigravity harness, which connects to +// a gRPC server at a fixed address. +type AntigravityHarnessConfig struct { + ID string `yaml:"id"` // Unique harness identifier + Address string `yaml:"address"` // gRPC address of the harness server +} + +// SubstrateHarnessConfig registers a harness backed by a SubstrATE ActorTemplate. +type SubstrateHarnessConfig struct { + ID string `yaml:"id"` // Unique harness identifier + Namespace string `yaml:"namespace"` // ActorTemplate namespace + Template string `yaml:"template"` // ActorTemplate name + Port int `yaml:"port,omitempty"` // HarnessService port (default 50053) +} + +// LoadFromFile loads configuration from a YAML file. +func LoadFromFile(path string) (*Config, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("failed to read config file: %w", err) + } + + var cfg Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("failed to parse config file: %w", err) + } + + cfg.setDefaults() + + return &cfg, nil +} + +// DefaultConfig returns a configuration with default values set. +func DefaultConfig() *Config { + var cfg Config + cfg.setDefaults() + return &cfg +} + +// setDefaults sets default values for optional fields. +func (c *Config) setDefaults() { + if c.Server.Address == "" { + c.Server.Address = ":8494" + } + if c.EventLog.SQLiteConfig.Filename == "" { + c.EventLog.SQLiteConfig.Filename = "eventlog/log.sqlite" + } +} + +// Validate validates the configuration. +func (c *Config) Validate() error { + if c.Server.Address == "" { + return fmt.Errorf("server.address is required") + } + if c.EventLog.SQLiteConfig.Filename == "" { + return fmt.Errorf("eventlog.sqlite.filename is required") + } + + return nil +} diff --git a/internal/controller2/controller_test.go b/internal/controller2/controller_test.go index 259087f..26d0489 100644 --- a/internal/controller2/controller_test.go +++ b/internal/controller2/controller_test.go @@ -19,7 +19,6 @@ import ( "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" @@ -27,14 +26,6 @@ import ( "github.com/google/ax/proto" ) -type dummyAgent struct{} - -func (a *dummyAgent) Connect(ctx context.Context, conversationID string, execID string, start *proto.AgentStart, e agent.Executor, o agent.OutputHandler) error { - return nil -} - -func (a *dummyAgent) Close() error { return nil } - func TestController2_ExecHelloWorld(t *testing.T) { ctx := context.Background() cid := "test-conversation-id" @@ -93,7 +84,7 @@ func TestController2_ExecAntigravityFallback(t *testing.T) { 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" @@ -102,7 +93,9 @@ func TestController2_ExecAntigravityFallback(t *testing.T) { } else { badHarness = harness.NewAntigravityHarness(scriptPath) } - reg.RegisterHarness("antigravity", badHarness) + if err := reg.RegisterHarness("antigravity", badHarness); err != nil { + t.Fatal(err) + } c, err := New(ctx, Config{ Registry: reg, @@ -206,5 +199,3 @@ func TestController2_ExecRuntimeFallback(t *testing.T) { t.Errorf("expected 'Hello world' output text response due to runtime fallback, got %q", gotText) } } - - diff --git a/internal/controller2/registry.go b/internal/controller2/registry.go index 0296a12..30918b3 100644 --- a/internal/controller2/registry.go +++ b/internal/controller2/registry.go @@ -15,239 +15,45 @@ package controller2 import ( - "context" "fmt" - "strings" "sync" - "github.com/google/ax/internal/agent" - "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. -// It provides agent discovery, health monitoring, and load balancing. +// Registry manages a collection of harnesses. 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. +// NewRegistry creates a new harness registry. func NewRegistry() *Registry { return &Registry{ - agents: make(map[string]agent.Agent), - agentInfo: make(map[string]*agent.AgentInfo), harnesses: make(map[string]harness.Harness), } } -func (r *Registry) Map() map[string]agent.Agent { - r.mu.RLock() - defer r.mu.RUnlock() - - return r.agents -} - -// RegisterLocal registers a local (in-process) agent. -func (r *Registry) RegisterLocal(cfg config.LocalAgentConfig) error { - r.mu.Lock() - defer r.mu.Unlock() - - if err := validateID(cfg.ID); err != nil { - return err - } - - if _, ok := r.agents[cfg.ID]; ok { - return fmt.Errorf("agent %s already registered", cfg.ID) - } - - r.agents[cfg.ID] = cfg.Agent - r.agentInfo[cfg.ID] = &agent.AgentInfo{ - ID: cfg.ID, - Name: cfg.Name, - Description: cfg.Description, - Metadata: cfg.Metadata, - } - - return nil -} - -// RegisterRemote registers a remote agent by creating a remote agent client. -// The protocol field determines what kind of remote agent to register -// (matched case-insensitively): -// - "axp" (default): AX's proto.AgentService. -// - "a2a": A2A protocol. -func (r *Registry) RegisterRemote(ctx context.Context, cfg config.RemoteAgentConfig) error { - r.mu.Lock() - defer r.mu.Unlock() - - if err := validateID(cfg.ID); err != nil { - return err - } - if _, ok := r.agents[cfg.ID]; ok { - return fmt.Errorf("agent %s already registered", cfg.ID) - } - - switch strings.ToLower(cfg.Protocol) { - case "", "axp": - return r.registerRemote(cfg) - case "a2a": - return r.registerA2A(ctx, cfg) - default: - return fmt.Errorf("remote agent %s: invalid protocol %q (want \"axp\" or \"a2a\")", cfg.ID, cfg.Protocol) - } -} - -func (r *Registry) registerRemote(cfg config.RemoteAgentConfig) error { - remoteAgent, err := agent.NewRemoteAgent(agent.RemoteAgentConfig{ - Address: cfg.Address, - Reconnect: true, - MaxRetries: 3, - }) - if err != nil { - return fmt.Errorf("failed to create remote agent: %w", err) - } - r.agents[cfg.ID] = remoteAgent - r.agentInfo[cfg.ID] = &agent.AgentInfo{ - ID: cfg.ID, - Name: cfg.Name, - Description: cfg.Description, - Metadata: cfg.Metadata, - } - return nil -} - -// Creates an A2A-protocol agent client. The agent's AgentCard is resolved at -// registration time and used to populate the agent's information. -func (r *Registry) registerA2A(ctx context.Context, cfg config.RemoteAgentConfig) error { - a2aAgent, err := expagent.NewA2AAgent(ctx, expagent.A2AAgentConfig{ - ID: cfg.ID, - Address: cfg.Address, - Auth: cfg.Auth, - Headers: cfg.Headers, - Stateless: cfg.A2A.Stateless, - }) - if err != nil { - return fmt.Errorf("failed to create a2a agent: %w", err) - } - name, description := a2abridge.AgentMetadataFromCard(a2aAgent.Card(), cfg.Name, cfg.Description) - r.agents[cfg.ID] = a2aAgent - r.agentInfo[cfg.ID] = &agent.AgentInfo{ - ID: cfg.ID, - Name: name, - Description: description, - Metadata: cfg.Metadata, - } - return nil -} - -// RegisterColab registers a Colab agent that executes a local Python file -// on a remote Colab session via the colab CLI. -func (r *Registry) RegisterColab(cfg config.ColabAgentConfig) error { - r.mu.Lock() - defer r.mu.Unlock() - - if err := validateID(cfg.ID); err != nil { - return err - } - - if _, ok := r.agents[cfg.ID]; ok { - return fmt.Errorf("agent %s already registered", cfg.ID) - } - - colabAgent, err := expagent.NewColabAgent(expagent.ColabAgentConfig{ - ID: cfg.ID, - LocalFile: cfg.LocalFile, - DriveFile: cfg.DriveFile, - Accelerator: cfg.Accelerator, - DriveMountPath: cfg.DriveMountPath, - Requirements: cfg.Requirements, - InputFlag: cfg.InputFlag, - OutputImage: cfg.OutputImage, - OutputDrivePath: cfg.OutputDrivePath, - }) - if err != nil { - return fmt.Errorf("failed to create colab agent: %w", err) - } - - r.agents[cfg.ID] = colabAgent - r.agentInfo[cfg.ID] = &agent.AgentInfo{ - ID: cfg.ID, - Name: cfg.Name, - Description: cfg.Description, - Metadata: cfg.Metadata, - } - - return nil -} - -// Get retrieves an agent by ID. -func (r *Registry) Get(id string) (agent.Agent, error) { - r.mu.RLock() - defer r.mu.RUnlock() - - a, exists := r.agents[id] - if !exists { - return nil, fmt.Errorf("agent %s not found", id) - } - - return a, nil -} - -// AgentInfo retrieves agent metadata by ID. -func (r *Registry) AgentInfo(id string) (*agent.AgentInfo, error) { - r.mu.RLock() - defer r.mu.RUnlock() - - info, ok := r.agentInfo[id] - if !ok { - return nil, fmt.Errorf("agent %s not found", id) - } - - return info, nil -} - -// List returns all registered agent IDs. -func (r *Registry) List() []string { - r.mu.RLock() - defer r.mu.RUnlock() - - ids := make([]string, 0, len(r.agents)) - for id := range r.agents { - ids = append(ids, id) +// RegisterHarness registers a harness under the given id. An empty id registers +// the harness as the default, used when a request specifies no agent id. +func (r *Registry) RegisterHarness(id string, h harness.Harness) error { + if id != "" { + if err := validateID(id); err != nil { + return err + } } - return ids -} - -// Close stops the registry and closes all agents. -func (r *Registry) Close() error { - r.mu.Lock() defer r.mu.Unlock() - // Close all agents - var firstErr error - for id, a := range r.agents { - if err := a.Close(); err != nil && firstErr == nil { - firstErr = fmt.Errorf("failed to close agent %s: %w", id, err) - } + if _, ok := r.harnesses[id]; ok { + return fmt.Errorf("harness %q already registered", id) } - - return firstErr -} -// 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() r.harnesses[id] = h + return nil } -// Harness retrieves a harness by ID. +// Harness retrieves a harness by id. func (r *Registry) Harness(id string) (harness.Harness, error) { r.mu.RLock() defer r.mu.RUnlock() @@ -257,3 +63,8 @@ func (r *Registry) Harness(id string) (harness.Harness, error) { } return h, nil } + +// Close releases resources held by the registry. +func (r *Registry) Close() error { + return nil +} diff --git a/internal/controller2/registry_ate.go b/internal/controller2/registry_ate.go deleted file mode 100644 index b72181d..0000000 --- a/internal/controller2/registry_ate.go +++ /dev/null @@ -1,62 +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. - -package controller2 - -import ( - "context" - "fmt" - - "github.com/google/ax/internal/agent" - "github.com/google/ax/internal/config" - expagent "github.com/google/ax/internal/experimental/agent" -) - -// RegisterATE registers an ATE agent by creating an ATE agent client. -func (r *Registry) RegisterATE(ctx context.Context, ateTarget string, cfg config.SubstrateAgentConfig) error { - r.mu.Lock() - defer r.mu.Unlock() - - if err := validateID(cfg.ID); err != nil { - return err - } - - if _, ok := r.agents[cfg.ID]; ok { - return fmt.Errorf("agent %s already registered", cfg.ID) - } - - // Create ATE agent client - substrateAgent, err := expagent.NewSubstrateAgent(ateTarget, expagent.SubstrateAgentConfig{ - ID: cfg.ID, - Namespace: cfg.Namespace, - Template: cfg.Template, - Port: cfg.Port, - Protocol: cfg.Protocol, - Auth: cfg.Auth, - Headers: cfg.Headers, - }) - if err != nil { - return fmt.Errorf("failed to create ATE agent: %w", err) - } - - r.agents[cfg.ID] = substrateAgent - r.agentInfo[cfg.ID] = &agent.AgentInfo{ - ID: cfg.ID, - Name: cfg.Name, - Description: cfg.Description, - Metadata: cfg.Metadata, - } - - return nil -} diff --git a/internal/controller2/registry_test.go b/internal/controller2/registry_test.go index 9052f2a..07d43ad 100644 --- a/internal/controller2/registry_test.go +++ b/internal/controller2/registry_test.go @@ -14,62 +14,47 @@ package controller2 -// TODO(lhuan): Setup a better automated testing framework - import ( - "context" - "fmt" - "net" "testing" - "google.golang.org/grpc" - - "github.com/google/ax/internal/config" - "github.com/google/ax/proto" + "github.com/google/ax/internal/harness/harnesstest" ) -type mockAgentServer struct { - proto.UnimplementedAgentServiceServer - healthy bool -} +func TestRegistry_RegisterHarness(t *testing.T) { + r := NewRegistry() + h := harnesstest.New() -func startMockGRPCServer(t *testing.T, healthy bool) (string, func()) { - lis, err := net.Listen("tcp", "localhost:0") - if err != nil { - t.Fatalf("failed to listen: %v", err) + if err := r.RegisterHarness("antigravity", h); err != nil { + t.Fatalf("RegisterHarness(valid id): %v", err) } - s := grpc.NewServer() - proto.RegisterAgentServiceServer(s, &mockAgentServer{healthy: healthy}) - go func() { - if err := s.Serve(lis); err != nil { - // server might be closed - } - }() - return lis.Addr().String(), func() { - s.Stop() - lis.Close() + + // Duplicate id is rejected. + if err := r.RegisterHarness("antigravity", h); err == nil { + t.Error("expected error registering duplicate id, got nil") } -} -func TestRegistry_GracefulShutdown(t *testing.T) { - r := NewRegistry() + // Invalid id is rejected. + if err := r.RegisterHarness("bad id", h); err == nil { + t.Error("expected error registering invalid id, got nil") + } - address, cleanup := startMockGRPCServer(t, true) - defer cleanup() + // Empty id (the default) bypasses validation and is allowed. + if err := r.RegisterHarness("", h); err != nil { + t.Fatalf("RegisterHarness(default): %v", err) + } +} - // Register multiple agents to create workload - for i := range 50 { - err := r.RegisterRemote(context.Background(), config.RemoteAgentConfig{ - ID: fmt.Sprintf("remote-shutdown-test-%d", i), - Name: "Shutdown Test Remote", - Address: address, - }) - if err != nil { - t.Fatalf("Failed to register remote agent for shutdown test: %v", err) - } +func TestRegistry_FindHarness(t *testing.T) { + r := NewRegistry() + h := harnesstest.New() + if err := r.RegisterHarness("antigravity", h); err != nil { + t.Fatalf("RegisterHarness: %v", err) } - // Close should return specific errors for failed agents, but NOT panic or deadlock - // We are testing for absence of panic/deadlock here. - _ = r.Close() + if _, err := r.Harness("antigravity"); err != nil { + t.Errorf("Harness(antigravity): %v", err) + } + if _, err := r.Harness("missing"); err == nil { + t.Error("expected error looking up missing harness, got nil") + } } diff --git a/internal/controller2/validation.go b/internal/controller2/validation.go index 4f71a74..f0bb081 100644 --- a/internal/controller2/validation.go +++ b/internal/controller2/validation.go @@ -18,28 +18,18 @@ import ( "errors" "fmt" "regexp" - "strings" ) -var reservedAgentIDs = map[string]struct{}{ - "gemini": {}, - "__planner": {}, -} - var validIDRegex = regexp.MustCompile(`^[A-Za-z0-9\-_]+$`) -// validateID checks if an ID contains allowed characters and is not reserved. +// validateID checks if a harness ID contains allowed characters. func validateID(id string) error { if id == "" { return errors.New("empty ID") } if !validIDRegex.MatchString(id) { - return fmt.Errorf("invalid ID %q: must only contain A-Z, a-z, 0-9, -, and _", id) - } - - if _, isReserved := reservedAgentIDs[strings.ToLower(id)]; isReserved { - return fmt.Errorf("agent ID %q is reserved", id) + return fmt.Errorf("invalid harness ID %q: must only contain A-Z, a-z, 0-9, -, and _", id) } return nil diff --git a/internal/controller2/validation_test.go b/internal/controller2/validation_test.go index 16c15d0..92d1cb9 100644 --- a/internal/controller2/validation_test.go +++ b/internal/controller2/validation_test.go @@ -59,26 +59,6 @@ func TestValidateID(t *testing.T) { id: "", wantErr: true, }, - { - name: "reserved gemini", - id: "gemini", - wantErr: true, - }, - { - name: "reserved gemini mixed case", - id: "Gemini", - wantErr: true, - }, - { - name: "reserved __planner", - id: "__planner", - wantErr: true, - }, - { - name: "reserved __planner mixed case", - id: "__Planner", - wantErr: true, - }, } for _, tt := range tests { diff --git a/internal/manifests/ax-deployment2.yaml b/internal/manifests/ax-deployment2.yaml index e8bf391..c7e1e5d 100644 --- a/internal/manifests/ax-deployment2.yaml +++ b/internal/manifests/ax-deployment2.yaml @@ -74,7 +74,7 @@ spec: containers: - name: ax-server image: ko://github.com/google/ax/cmd/ax - command: ["/ko-app/ax", "serve"] + command: ["/ko-app/ax", "serve", "--config", "/etc/ax/ax.yaml"] ports: - containerPort: 8494 env: @@ -82,3 +82,27 @@ spec: value: "${GEMINI_API_KEY}" - name: AX_SUBSTRATE value: "1" + volumeMounts: + - name: ax-config + mountPath: /etc/ax + volumes: + - name: ax-config + configMap: + name: ax-server-config +--- +apiVersion: v1 +kind: ConfigMap +metadata: + name: ax-server-config + namespace: ax +data: + ax.yaml: | + ate: + endpoint: "api.ate-system.svc:443" + harnesses: + default: hello-world + substrate: + - id: hello-world + namespace: ax + template: ax-harness-template + port: 50053 From 3388cb2220bb4194d61ade1aae89fdf0390765b9 Mon Sep 17 00:00:00 2001 From: Junjie Wang Date: Wed, 3 Jun 2026 18:54:52 -0700 Subject: [PATCH 2/5] Small style change. --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index dfcb6da..eff9de2 100644 --- a/Makefile +++ b/Makefile @@ -31,7 +31,7 @@ proto: proto/ax.proto proto/content.proto @echo "Protobuf generation complete!" -# Run tests. +# Run tests test: @echo "Running tests..." @go test -v ./... From a15881e154c17a4274019479a955b4c358b096bd Mon Sep 17 00:00:00 2001 From: Junjie Wang Date: Thu, 4 Jun 2026 21:53:09 -0700 Subject: [PATCH 3/5] Refactor the harness config according to previous comments and discussions. --- cmd/ax/internal/cliutil/cliutil_harness.go | 32 +++- .../internal/cliutil/cliutil_harness_test.go | 96 +++++++++- .../config/harnessconfig/harness_config.go | 124 ------------ internal/config2/config.go | 181 ++++++++++++++++++ internal/config2/config_test.go | 113 +++++++++++ internal/manifests/README.md | 40 +++- internal/manifests/ax-deployment2.yaml | 77 +++++++- 7 files changed, 503 insertions(+), 160 deletions(-) delete mode 100644 internal/config/harnessconfig/harness_config.go create mode 100644 internal/config2/config.go create mode 100644 internal/config2/config_test.go diff --git a/cmd/ax/internal/cliutil/cliutil_harness.go b/cmd/ax/internal/cliutil/cliutil_harness.go index 38eda48..d936ff6 100644 --- a/cmd/ax/internal/cliutil/cliutil_harness.go +++ b/cmd/ax/internal/cliutil/cliutil_harness.go @@ -19,11 +19,11 @@ package cliutil import ( "context" "fmt" + "os" - "github.com/google/ax/internal/config/harnessconfig" + "github.com/google/ax/internal/config2" "github.com/google/ax/internal/controller/executor" "github.com/google/ax/internal/controller2" - "github.com/google/ax/internal/harness" ) // Controller is the active controller type for this build. @@ -33,37 +33,49 @@ type Controller = *controller2.Controller type ExecHandler = controller2.ExecHandler // Config is the configuration type for this build. -type Config = harnessconfig.Config +type Config = config2.Config // LoadFromFile loads configuration from a YAML file. func LoadFromFile(path string) (*Config, error) { - return harnessconfig.LoadFromFile(path) + return config2.LoadFromFile(path) } // DefaultConfig returns a configuration with default values set. func DefaultConfig() *Config { - return harnessconfig.DefaultConfig() + return config2.DefaultConfig() } // NewControllerFromConfig creates a controller2.Controller instance based on the provided configuration. func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller2.Controller, error) { reg := controller2.NewRegistry() - // Antigravity harnesses. + // AX_SUBSTRATE selects how built-in harnesses run: locally (unset) or as + // substrate actors ("1"). + substrateMode := os.Getenv("AX_SUBSTRATE") == "1" + // AX_SUBSTRATE_ENDPOINT is the control-plane endpoint for substrate server. + endpoint := os.Getenv("AX_SUBSTRATE_ENDPOINT") + + // Built-in harnesses. for _, hc := range cfg.Harnesses.Antigravity { - h := harness.NewAntigravityHarness(hc.Address) + h, err := hc.NewHarness(substrateMode, endpoint) + if err != nil { + return nil, fmt.Errorf("antigravity harness %q: %w", hc.ID, err) + } if err := reg.RegisterHarness(hc.ID, h); err != nil { return nil, fmt.Errorf("register antigravity harness %q: %w", hc.ID, err) } } - // Substrate harnesses. + // Custom substrate harnesses. + if len(cfg.Harnesses.Substrate) > 0 && !substrateMode { + return nil, fmt.Errorf("custom substrate harnesses require AX_SUBSTRATE=1") + } for _, sc := range cfg.Harnesses.Substrate { - sh, err := harness.NewSubstrateHarness(cfg.ATE.Endpoint, sc.Namespace, sc.Template, sc.Port) + h, err := sc.NewHarness(endpoint) if err != nil { return nil, fmt.Errorf("substrate harness %q: %w", sc.ID, err) } - if err := reg.RegisterHarness(sc.ID, sh); err != nil { + if err := reg.RegisterHarness(sc.ID, h); err != nil { return nil, fmt.Errorf("register substrate harness %q: %w", sc.ID, err) } } diff --git a/cmd/ax/internal/cliutil/cliutil_harness_test.go b/cmd/ax/internal/cliutil/cliutil_harness_test.go index a33d701..1655cfc 100644 --- a/cmd/ax/internal/cliutil/cliutil_harness_test.go +++ b/cmd/ax/internal/cliutil/cliutil_harness_test.go @@ -22,19 +22,19 @@ import ( "strings" "testing" - "github.com/google/ax/internal/config/harnessconfig" + "github.com/google/ax/internal/config2" ) func TestNewControllerFromConfig_DefaultHarness(t *testing.T) { - cfg := &harnessconfig.Config{ - EventLog: harnessconfig.EventLogConfig{ - SQLiteConfig: harnessconfig.SQLiteConfig{ + cfg := &config2.Config{ + EventLog: config2.EventLogConfig{ + SQLiteConfig: config2.SQLiteConfig{ Filename: filepath.Join(t.TempDir(), "log.sqlite"), }, }, - Harnesses: harnessconfig.HarnessesConfig{ + Harnesses: config2.HarnessesConfig{ Default: "ag", - Antigravity: []harnessconfig.AntigravityHarnessConfig{ + Antigravity: []config2.AntigravityHarnessConfig{ {ID: "ag", Address: "localhost:50053"}, }, }, @@ -51,10 +51,10 @@ func TestNewControllerFromConfig_DefaultHarness(t *testing.T) { } func TestNewControllerFromConfig_UnknownDefaultHarness(t *testing.T) { - cfg := &harnessconfig.Config{ - Harnesses: harnessconfig.HarnessesConfig{ + cfg := &config2.Config{ + Harnesses: config2.HarnessesConfig{ Default: "missing", - Antigravity: []harnessconfig.AntigravityHarnessConfig{ + Antigravity: []config2.AntigravityHarnessConfig{ {ID: "ag", Address: "localhost:50053"}, }, }, @@ -68,3 +68,81 @@ func TestNewControllerFromConfig_UnknownDefaultHarness(t *testing.T) { t.Errorf("expected error to mention %q, got: %v", "missing", err) } } + +func TestNewControllerFromConfig_BuiltinSubstrate(t *testing.T) { + t.Setenv("AX_SUBSTRATE", "1") + + cfg := &config2.Config{ + EventLog: config2.EventLogConfig{ + SQLiteConfig: config2.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: config2.HarnessesConfig{ + Default: "ag", + Antigravity: []config2.AntigravityHarnessConfig{ + {ID: "ag"}, + }, + }, + } + + c, err := NewControllerFromConfig(context.Background(), cfg) + if err != nil { + t.Fatalf("NewControllerFromConfig: %v", err) + } + if c == nil { + t.Fatal("expected non-nil controller") + } + c.Close() +} + +func TestNewControllerFromConfig_CustomHarnessRequiresSubstrateMode(t *testing.T) { + t.Setenv("AX_SUBSTRATE", "") + + cfg := &config2.Config{ + EventLog: config2.EventLogConfig{ + SQLiteConfig: config2.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: config2.HarnessesConfig{ + Substrate: []config2.SubstrateHarnessConfig{ + {ID: "custom", Namespace: "team-ns", Template: "custom-template"}, + }, + }, + } + + _, err := NewControllerFromConfig(context.Background(), cfg) + if err == nil { + t.Fatal("expected error for custom substrate harness without AX_SUBSTRATE=1, got nil") + } + if !strings.Contains(err.Error(), "AX_SUBSTRATE=1") { + t.Errorf("expected error to mention AX_SUBSTRATE=1, got: %v", err) + } +} + +func TestNewControllerFromConfig_CustomHarnessInSubstrateMode(t *testing.T) { + t.Setenv("AX_SUBSTRATE", "1") + + cfg := &config2.Config{ + EventLog: config2.EventLogConfig{ + SQLiteConfig: config2.SQLiteConfig{ + Filename: filepath.Join(t.TempDir(), "log.sqlite"), + }, + }, + Harnesses: config2.HarnessesConfig{ + Substrate: []config2.SubstrateHarnessConfig{ + {ID: "custom", Namespace: "team-ns", Template: "custom-template"}, + }, + }, + } + + c, err := NewControllerFromConfig(context.Background(), cfg) + if err != nil { + t.Fatalf("NewControllerFromConfig: %v", err) + } + if c == nil { + t.Fatal("expected non-nil controller") + } + c.Close() +} diff --git a/internal/config/harnessconfig/harness_config.go b/internal/config/harnessconfig/harness_config.go deleted file mode 100644 index 1d98e84..0000000 --- a/internal/config/harnessconfig/harness_config.go +++ /dev/null @@ -1,124 +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. - -// Package harnessconfig provides configuration structures for the AX harness build. -package harnessconfig - -import ( - "fmt" - "os" - - "gopkg.in/yaml.v3" -) - -// Config represents the main configuration for the AX harness server. -type Config struct { - Server ServerConfig `yaml:"server"` - EventLog EventLogConfig `yaml:"eventlog"` - Harnesses HarnessesConfig `yaml:"harnesses,omitempty"` - ATE ATEConfig `yaml:"ate,omitempty"` -} - -// ServerConfig configures the gRPC server. -type ServerConfig struct { - Address string `yaml:"address"` // Server address to listen on (e.g., ":8494") -} - -// SQLiteConfig configures the SQLite event log file. -type SQLiteConfig struct { - Filename string `yaml:"filename"` // SQLite file for event log storage -} - -// EventLogConfig configures the event log storage. -type EventLogConfig struct { - SQLiteConfig SQLiteConfig `yaml:"sqlite"` -} - -// ATEConfig configures the substrate control plane endpoint used by -// substrate harnesses. -type ATEConfig struct { - Endpoint string `yaml:"endpoint"` -} - -// HarnessesConfig groups harnesses to serve by type. Each type maps to a list -// of that type's configurations. -type HarnessesConfig struct { - // Default is the id of the harness to serve when a request specifies no harness. - Default string `yaml:"default,omitempty"` - // Antigravity harnesses connect to a gRPC server at a fixed address. - Antigravity []AntigravityHarnessConfig `yaml:"antigravity,omitempty"` - // Substrate harnesses are brought up as SubstrATE actors from an ActorTemplate. - Substrate []SubstrateHarnessConfig `yaml:"substrate,omitempty"` -} - -// AntigravityHarnessConfig registers an Antigravity harness, which connects to -// a gRPC server at a fixed address. -type AntigravityHarnessConfig struct { - ID string `yaml:"id"` // Unique harness identifier - Address string `yaml:"address"` // gRPC address of the harness server -} - -// SubstrateHarnessConfig registers a harness backed by a SubstrATE ActorTemplate. -type SubstrateHarnessConfig struct { - ID string `yaml:"id"` // Unique harness identifier - Namespace string `yaml:"namespace"` // ActorTemplate namespace - Template string `yaml:"template"` // ActorTemplate name - Port int `yaml:"port,omitempty"` // HarnessService port (default 50053) -} - -// LoadFromFile loads configuration from a YAML file. -func LoadFromFile(path string) (*Config, error) { - data, err := os.ReadFile(path) - if err != nil { - return nil, fmt.Errorf("failed to read config file: %w", err) - } - - var cfg Config - if err := yaml.Unmarshal(data, &cfg); err != nil { - return nil, fmt.Errorf("failed to parse config file: %w", err) - } - - cfg.setDefaults() - - return &cfg, nil -} - -// DefaultConfig returns a configuration with default values set. -func DefaultConfig() *Config { - var cfg Config - cfg.setDefaults() - return &cfg -} - -// setDefaults sets default values for optional fields. -func (c *Config) setDefaults() { - if c.Server.Address == "" { - c.Server.Address = ":8494" - } - if c.EventLog.SQLiteConfig.Filename == "" { - c.EventLog.SQLiteConfig.Filename = "eventlog/log.sqlite" - } -} - -// Validate validates the configuration. -func (c *Config) Validate() error { - if c.Server.Address == "" { - return fmt.Errorf("server.address is required") - } - if c.EventLog.SQLiteConfig.Filename == "" { - return fmt.Errorf("eventlog.sqlite.filename is required") - } - - return nil -} diff --git a/internal/config2/config.go b/internal/config2/config.go new file mode 100644 index 0000000..2aa8436 --- /dev/null +++ b/internal/config2/config.go @@ -0,0 +1,181 @@ +// 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 config2 provides configuration for the controller2 server path. +package config2 + +import ( + "fmt" + "os" + + "github.com/google/ax/internal/harness" + "gopkg.in/yaml.v3" +) + +const ( + // The substrate namespace reserved for AX's built-in harnesses. + defaultNamespace = "ax" + // The default port of HarnessService. + defaultPort = 50053 + // The Antigravity ActorTemplate name. + antigravityTemplate = "antigravity-template" +) + +// Config represents the main configuration for the AX harness server. +type Config struct { + Server ServerConfig `yaml:"server"` + EventLog EventLogConfig `yaml:"eventlog"` + Harnesses HarnessesConfig `yaml:"harnesses,omitempty"` +} + +// ServerConfig configures the gRPC server. +type ServerConfig struct { + Address string `yaml:"address"` // Server address to listen on (e.g., ":8494") +} + +// SQLiteConfig configures the SQLite event log file. +type SQLiteConfig struct { + Filename string `yaml:"filename"` // SQLite file for event log storage +} + +// EventLogConfig configures the event log storage. +type EventLogConfig struct { + SQLiteConfig SQLiteConfig `yaml:"sqlite"` +} + +// HarnessesConfig groups harnesses to serve by type. There are two categories: +// - Built-in harnesses (e.g. Antigravity) whose implementation and container +// image are provided by AX. +// - Custom harnesses on substrate whose implementation and container image are +// provided by the user via their own ActorTemplate. +type HarnessesConfig struct { + // Default is the id of the harness to serve when a request specifies no harness. + Default string `yaml:"default,omitempty"` + Antigravity []AntigravityHarnessConfig `yaml:"antigravity,omitempty"` + Substrate []SubstrateHarnessConfig `yaml:"substrate,omitempty"` +} + +// AntigravityHarnessConfig registers the built-in Antigravity harness. +type AntigravityHarnessConfig struct { + ID string `yaml:"id"` // Unique harness identifier + Address string `yaml:"address,omitempty"` // HarnessService address +} + +// SubstrateHarnessConfig registers a custom harness deployed on substrate +// from a user-provided container image. +type SubstrateHarnessConfig struct { + ID string `yaml:"id"` // Unique harness identifier + Namespace string `yaml:"namespace"` // ActorTemplate namespace (user-owned, not "ax") + Template string `yaml:"template"` // ActorTemplate name + Port int `yaml:"port,omitempty"` // HarnessService port +} + +// NewHarness builds the built-in Antigravity harness. In substrate mode it's deployed +// as a substrate actor; otherwise it runs locally. +func (c AntigravityHarnessConfig) NewHarness(substrate bool, endpoint string) (harness.Harness, error) { + if substrate { + return newSubstrateHarness(endpoint, defaultNamespace, antigravityTemplate, defaultPort) + } + address := c.Address + if address == "" { + address = fmt.Sprintf("localhost:%d", defaultPort) + } + return harness.NewAntigravityHarness(address), nil +} + +// NewHarness builds the custom harness. Custom harnesses always run as substrate +// actors from the user's own ActorTemplate. +func (c SubstrateHarnessConfig) NewHarness(endpoint string) (harness.Harness, error) { + port := c.Port + if port == 0 { + port = defaultPort + } + return newSubstrateHarness(endpoint, c.Namespace, c.Template, port) +} + +// newSubstrateHarness brings up a harness that is deployed as a substrate actor. +func newSubstrateHarness(endpoint, namespace, template string, port int) (harness.Harness, error) { + sh, err := harness.NewSubstrateHarness(endpoint, namespace, template, port) + if err != nil { + return nil, err + } + return sh, nil +} + +// LoadFromFile loads configuration from a YAML file. +func LoadFromFile(path string) (*Config, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("failed to read config file: %w", err) + } + + var cfg Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("failed to parse config file: %w", err) + } + + cfg.setDefaults() + + return &cfg, nil +} + +// DefaultConfig returns a configuration with default values set. +func DefaultConfig() *Config { + var cfg Config + cfg.setDefaults() + return &cfg +} + +// setDefaults sets default values for optional fields. +func (c *Config) setDefaults() { + if c.Server.Address == "" { + c.Server.Address = ":8494" + } + if c.EventLog.SQLiteConfig.Filename == "" { + c.EventLog.SQLiteConfig.Filename = "eventlog/log.sqlite" + } +} + +// Validate validates the configuration. +func (c *Config) Validate() error { + if c.Server.Address == "" { + return fmt.Errorf("server.address is required") + } + if c.EventLog.SQLiteConfig.Filename == "" { + return fmt.Errorf("eventlog.sqlite.filename is required") + } + + for _, hc := range c.Harnesses.Antigravity { + if hc.ID == "" { + return fmt.Errorf("antigravity harness id is required") + } + } + + for _, sc := range c.Harnesses.Substrate { + if sc.ID == "" { + return fmt.Errorf("substrate harness id is required") + } + if sc.Namespace == "" { + return fmt.Errorf("substrate harness %q: namespace is required", sc.ID) + } + if sc.Namespace == defaultNamespace { + return fmt.Errorf("substrate harness %q: namespace %q is reserved for built-in harnesses", sc.ID, defaultNamespace) + } + if sc.Template == "" { + return fmt.Errorf("substrate harness %q: template is required", sc.ID) + } + } + + return nil +} diff --git a/internal/config2/config_test.go b/internal/config2/config_test.go new file mode 100644 index 0000000..abdd6ed --- /dev/null +++ b/internal/config2/config_test.go @@ -0,0 +1,113 @@ +// 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 config2 + +import ( + "strings" + "testing" +) + +func TestAntigravityNewHarness_Local(t *testing.T) { + h, err := AntigravityHarnessConfig{ID: "ag"}.NewHarness(false, "") + if err != nil { + t.Fatalf("NewHarness: %v", err) + } + if h == nil { + t.Fatal("expected non-nil harness") + } +} + +func TestAntigravityNewHarness_Substrate(t *testing.T) { + h, err := AntigravityHarnessConfig{ID: "ag"}.NewHarness(true, "api.ate-system.svc:443") + if err != nil { + t.Fatalf("NewHarness: %v", err) + } + if h == nil { + t.Fatal("expected non-nil harness") + } +} + +func TestSubstrateNewHarness(t *testing.T) { + h, err := SubstrateHarnessConfig{ID: "c", Namespace: "team-ns", Template: "custom-template"}.NewHarness("api.ate-system.svc:443") + if err != nil { + t.Fatalf("NewHarness: %v", err) + } + if h == nil { + t.Fatal("expected non-nil harness") + } +} + +// validConfig returns a config that passes Validate, that tests can mutate. +func validConfig() *Config { + c := DefaultConfig() + c.Harnesses = HarnessesConfig{ + Antigravity: []AntigravityHarnessConfig{{ID: "ag"}}, + Substrate: []SubstrateHarnessConfig{ + {ID: "custom", Namespace: "team-ns", Template: "custom-template"}, + }, + } + return c +} + +func TestValidate_ValidConfig(t *testing.T) { + if err := validConfig().Validate(); err != nil { + t.Fatalf("Validate() = %v, want nil", err) + } +} + +func TestValidate_AntigravityIDRequired(t *testing.T) { + c := validConfig() + c.Harnesses.Antigravity[0].ID = "" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "antigravity harness id") { + t.Fatalf("Validate() = %v, want antigravity id error", err) + } +} + +func TestValidate_CustomIDRequired(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].ID = "" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "substrate harness id") { + t.Fatalf("Validate() = %v, want substrate id error", err) + } +} + +func TestValidate_CustomNamespaceRequired(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Namespace = "" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "namespace is required") { + t.Fatalf("Validate() = %v, want namespace-required error", err) + } +} + +func TestValidate_CustomNamespaceReserved(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Namespace = defaultNamespace + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "reserved") { + t.Fatalf("Validate() = %v, want reserved-namespace error", err) + } +} + +func TestValidate_CustomTemplateRequired(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Template = "" + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "template is required") { + t.Fatalf("Validate() = %v, want template-required error", err) + } +} diff --git a/internal/manifests/README.md b/internal/manifests/README.md index acb6cb2..ea410bf 100644 --- a/internal/manifests/README.md +++ b/internal/manifests/README.md @@ -15,9 +15,25 @@ The target Kubernetes cluster is assumed to have --- +## Harness types + +AX serves two kinds of harnesses: + +- **Built-in** (e.g. Antigravity): implementation + and container image are provided by AX. You configure only behavior; AX owns + deployment. A built-in runs **locally** or as a **SubstrATE actor** depending + on the `AX_SUBSTRATE` environment variable (`1` = substrate). Built-in actors + run in the reserved `ax` namespace. +- **Custom** (the `substrate` config key): implementation and container image are + provided by you via your own `ActorTemplate`. Custom harnesses always run on + SubstrATE, in **your own namespace** (the `ax` namespace is reserved for + built-ins), and require `AX_SUBSTRATE=1`. + +--- + ## 🚀 Deploying to Agent Substrate -This deploys the AX `harness` path: the AX harness `WorkerPool` and `ActorTemplate` — provisioned as isolated, warm-standby actors that are live-snapshotted on boot and instantly restored from GCS when a new conversation starts — together with an `ax-server` controller front-end (a `ReplicaSet`). +This deploys the AX `harness` path: a built-in harness `WorkerPool` and `ActorTemplate` (the `antigravity` example, in the reserved `ax` namespace), a custom harness `WorkerPool` and `ActorTemplate` (the `hello-world` example, in the `custom-harness` namespace) — provisioned as isolated, warm-standby actors that are live-snapshotted on boot and instantly restored from GCS when a new conversation starts — together with an `ax-server` controller front-end (a `ReplicaSet`) in the `ax` namespace. ### 1. Build and Deploy @@ -38,19 +54,27 @@ export KO_DEFAULTPLATFORMS="linux/amd64" This command will: - Build the AX images using `ko` with the `harness` build tag. -- Create the `ax` namespace. -- Create the `WorkerPool` and `ActorTemplate` for the AX harness. -- Create the `ax-server` `ReplicaSet` (the controller front-end). +- Create the `ax` namespace (AX control plane + built-in harnesses) and the + `custom-harness` namespace (the example custom harness). +- Create a shared `harness-workerpool` `WorkerPool` and the built-in + `antigravity-template` `ActorTemplate` in `ax` (all built-in harnesses share + this pool). +- Create a shared `harness-workerpool` `WorkerPool` and the `hello-world-template` + `ActorTemplate` in `custom-harness` (custom harnesses there share this pool). +- Create the `ax-server` `ReplicaSet` (the controller front-end) in `ax`. - Create the `ax-server-config` `ConfigMap` that tells the `ax-server` which harnesses to serve (mounted at `/etc/ax/ax.yaml`). -The harness registry lives in that `ConfigMap`. By default it registers a -substrate harness (`hello-world`) backed by the `ax-harness-template`, marked as +The harness registry lives in that `ConfigMap`. It registers a built-in +`antigravity` harness (AX-managed, in `ax`; currently a placeholder stub that +returns "hello world" until the real antigravity image lands) and a custom +substrate harness (`hello-world`, in `custom-harness`), with the latter marked as the default via `harnesses.default`. -Wait until the template is ready: +Wait until the templates are ready: ```bash -kubectl wait --for=condition=Ready actortemplate/ax-harness-template -n ax --timeout=5m +kubectl wait --for=condition=Ready actortemplate/antigravity-template -n ax --timeout=5m +kubectl wait --for=condition=Ready actortemplate/hello-world-template -n custom-harness --timeout=5m ``` ### 2. Port-Forward Services diff --git a/internal/manifests/ax-deployment2.yaml b/internal/manifests/ax-deployment2.yaml index c7e1e5d..458584b 100644 --- a/internal/manifests/ax-deployment2.yaml +++ b/internal/manifests/ax-deployment2.yaml @@ -14,15 +14,28 @@ # TODO(jbd): This yaml will eventually become the ax-deployment.yaml.tmpl. +# --------------------------------------------------------------------------- +# Namespaces +# --------------------------------------------------------------------------- +# Namespace "ax" is reserved for built-in harnesses. Users who bring their own +# harness images should add their own namespaces. apiVersion: v1 kind: Namespace metadata: name: ax --- +apiVersion: v1 +kind: Namespace +metadata: + name: custom-harness +--- +# --------------------------------------------------------------------------- +# Built-in harnesses +# --------------------------------------------------------------------------- apiVersion: ate.dev/v1alpha1 kind: WorkerPool metadata: - name: ax-harness-workerpool + name: harness-workerpool namespace: ax spec: replicas: 5 @@ -31,11 +44,11 @@ spec: apiVersion: ate.dev/v1alpha1 kind: ActorTemplate metadata: - name: ax-harness-template + name: antigravity-template namespace: ax spec: workerPoolRef: - name: ax-harness-workerpool + name: harness-workerpool namespace: ax runsc: amd64: @@ -45,6 +58,45 @@ spec: url: "gs://gvisor/releases/nightly/2026-05-19/aarch64/runsc" sha256Hash: "1ba2366ae2efceba166046f51a4104f9261c9cb72c6db8f5b3fe2dc57dea86b9" pauseImage: "gcr.io/gke-release/pause@sha256:bcbd57ba5653580ec647b16d8163cdd1112df3609129b01f912a8032e48265da" + containers: + # TODO(wjjclaud): update the hello-world demo container with real antigravity. + - name: "axharness" + image: ko://github.com/google/ax/cmd/ax + command: ["/ko-app/ax", "harness"] + ports: + - containerPort: 50053 + snapshotsConfig: + location: gs://${BUCKET_NAME}/antigravity/ +--- +# --------------------------------------------------------------------------- +# Custom harnesses +# --------------------------------------------------------------------------- +apiVersion: ate.dev/v1alpha1 +kind: WorkerPool +metadata: + name: harness-workerpool + namespace: custom-harness +spec: + replicas: 3 + ateomImage: ko://github.com/agent-substrate/substrate/cmd/servers/ateom-gvisor +--- +apiVersion: ate.dev/v1alpha1 +kind: ActorTemplate +metadata: + name: hello-world-template + namespace: custom-harness +spec: + workerPoolRef: + name: harness-workerpool + namespace: custom-harness + runsc: + amd64: + url: "gs://gvisor/releases/nightly/2026-05-19/x86_64/runsc" + sha256Hash: "a397be1abc2420d26bce6c70e6e2ff96c73aaaab929756c56f5e2089ea842b63" + arm64: + url: "gs://gvisor/releases/nightly/2026-05-19/aarch64/runsc" + sha256Hash: "1ba2366ae2efceba166046f51a4104f9261c9cb72c6db8f5b3fe2dc57dea86b9" + pauseImage: "gcr.io/gke-release/pause@sha256:bcbd57ba5653580ec647b16d8163cdd1112df3609129b01f912a8032e48265da" containers: - name: "axharness" image: ko://github.com/google/ax/cmd/ax @@ -52,8 +104,11 @@ spec: ports: - containerPort: 50053 snapshotsConfig: - location: gs://${BUCKET_NAME}/ax-harness/ + location: gs://${BUCKET_NAME}/hello-world/ --- +# --------------------------------------------------------------------------- +# AX Server +# --------------------------------------------------------------------------- apiVersion: apps/v1 kind: ReplicaSet metadata: @@ -82,6 +137,8 @@ spec: value: "${GEMINI_API_KEY}" - name: AX_SUBSTRATE value: "1" + - name: AX_SUBSTRATE_ENDPOINT + value: "api.ate-system.svc:443" volumeMounts: - name: ax-config mountPath: /etc/ax @@ -97,12 +154,14 @@ metadata: namespace: ax data: ax.yaml: | - ate: - endpoint: "api.ate-system.svc:443" harnesses: - default: hello-world + default: antigravity + # Built-in harness + antigravity: + - id: antigravity + # Custom harness substrate: - id: hello-world - namespace: ax - template: ax-harness-template + namespace: custom-harness + template: hello-world-template port: 50053 From ee16bcef87e247fe8874231ef5dfa6b68caf5354 Mon Sep 17 00:00:00 2001 From: Junjie Wang Date: Thu, 4 Jun 2026 22:04:34 -0700 Subject: [PATCH 4/5] Update wokerpool name. --- internal/manifests/README.md | 7 ++++--- internal/manifests/ax-deployment2.yaml | 8 ++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/internal/manifests/README.md b/internal/manifests/README.md index ea410bf..faa090f 100644 --- a/internal/manifests/README.md +++ b/internal/manifests/README.md @@ -56,11 +56,12 @@ This command will: - Build the AX images using `ko` with the `harness` build tag. - Create the `ax` namespace (AX control plane + built-in harnesses) and the `custom-harness` namespace (the example custom harness). -- Create a shared `harness-workerpool` `WorkerPool` and the built-in +- Create a shared `ax-harness-workerpool` `WorkerPool` and the built-in `antigravity-template` `ActorTemplate` in `ax` (all built-in harnesses share this pool). -- Create a shared `harness-workerpool` `WorkerPool` and the `hello-world-template` - `ActorTemplate` in `custom-harness` (custom harnesses there share this pool). +- Create a shared `custom-harness-workerpool` `WorkerPool` and the + `hello-world-template` `ActorTemplate` in `custom-harness` (custom harnesses + there share this pool). - Create the `ax-server` `ReplicaSet` (the controller front-end) in `ax`. - Create the `ax-server-config` `ConfigMap` that tells the `ax-server` which harnesses to serve (mounted at `/etc/ax/ax.yaml`). diff --git a/internal/manifests/ax-deployment2.yaml b/internal/manifests/ax-deployment2.yaml index 458584b..3657912 100644 --- a/internal/manifests/ax-deployment2.yaml +++ b/internal/manifests/ax-deployment2.yaml @@ -35,7 +35,7 @@ metadata: apiVersion: ate.dev/v1alpha1 kind: WorkerPool metadata: - name: harness-workerpool + name: ax-harness-workerpool namespace: ax spec: replicas: 5 @@ -48,7 +48,7 @@ metadata: namespace: ax spec: workerPoolRef: - name: harness-workerpool + name: ax-harness-workerpool namespace: ax runsc: amd64: @@ -74,7 +74,7 @@ spec: apiVersion: ate.dev/v1alpha1 kind: WorkerPool metadata: - name: harness-workerpool + name: custom-harness-workerpool namespace: custom-harness spec: replicas: 3 @@ -87,7 +87,7 @@ metadata: namespace: custom-harness spec: workerPoolRef: - name: harness-workerpool + name: custom-harness-workerpool namespace: custom-harness runsc: amd64: From f2f891c5cf3b2be2910cc339b37aaadee6e962d3 Mon Sep 17 00:00:00 2001 From: Junjie Wang Date: Fri, 5 Jun 2026 11:00:33 -0700 Subject: [PATCH 5/5] Remove the address override from the example. --- internal/ax2.yaml | 1 - 1 file changed, 1 deletion(-) diff --git a/internal/ax2.yaml b/internal/ax2.yaml index fa76559..9865b29 100644 --- a/internal/ax2.yaml +++ b/internal/ax2.yaml @@ -28,4 +28,3 @@ harnesses: antigravity: - id: antigravity-example - address: "localhost:50053"