diff --git a/cmd/ax/internal/cliutil/cliutil_harness.go b/cmd/ax/internal/cliutil/cliutil_harness.go index d936ff6..1187bf1 100644 --- a/cmd/ax/internal/cliutil/cliutil_harness.go +++ b/cmd/ax/internal/cliutil/cliutil_harness.go @@ -24,8 +24,11 @@ import ( "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" ) +const antigravityHarnessID = "antigravity" + // Controller is the active controller type for this build. type Controller = *controller2.Controller @@ -52,19 +55,29 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller2.Con // 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, err := hc.NewHarness(substrateMode, endpoint) - if err != nil { - return nil, fmt.Errorf("antigravity harness %q: %w", hc.ID, err) + var defaultHarnessID string + var antigravityHarness harness.Harness + var err error + if !substrateMode { + address := cfg.Harnesses.Antigravity.Endpoint + if address == "" { + address = "127.0.0.1:50053" } - if err := reg.RegisterHarness(hc.ID, h); err != nil { - return nil, fmt.Errorf("register antigravity harness %q: %w", hc.ID, err) + antigravityHarness = harness.NewAntigravityHarness(address) + } else { + antigravityHarness, err = harness.NewSubstrateHarness(antigravityHarnessID, "", "", "", 80) + if err != nil { + return nil, fmt.Errorf("antigravity harness: %w", err) } } + if err := reg.RegisterHarness(antigravityHarnessID, antigravityHarness); err != nil { + return nil, fmt.Errorf("register antigravity harness: %w", err) + } + if cfg.Harnesses.Antigravity.Default { + defaultHarnessID = antigravityHarnessID + } // Custom substrate harnesses. if len(cfg.Harnesses.Substrate) > 0 && !substrateMode { @@ -78,16 +91,19 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller2.Con if err := reg.RegisterHarness(sc.ID, h); err != nil { return nil, fmt.Errorf("register substrate harness %q: %w", sc.ID, err) } + if sc.Default { + defaultHarnessID = sc.ID + } } // Register the configured default harness. - if id := cfg.Harnesses.Default; id != "" { - h, err := reg.Harness(id) + if defaultHarnessID != "" { + h, err := reg.Harness(defaultHarnessID) if err != nil { - return nil, fmt.Errorf("default harness %q not found", id) + return nil, fmt.Errorf("default harness %q not found", defaultHarnessID) } if err := reg.RegisterHarness("", h); err != nil { - return nil, fmt.Errorf("register default harness %q: %w", id, err) + return nil, fmt.Errorf("register default harness %q: %w", defaultHarnessID, err) } } diff --git a/cmd/ax/internal/cliutil/cliutil_harness_test.go b/cmd/ax/internal/cliutil/cliutil_harness_test.go index 1655cfc..bc392fe 100644 --- a/cmd/ax/internal/cliutil/cliutil_harness_test.go +++ b/cmd/ax/internal/cliutil/cliutil_harness_test.go @@ -33,9 +33,9 @@ func TestNewControllerFromConfig_DefaultHarness(t *testing.T) { }, }, Harnesses: config2.HarnessesConfig{ - Default: "ag", - Antigravity: []config2.AntigravityHarnessConfig{ - {ID: "ag", Address: "localhost:50053"}, + Antigravity: config2.AntigravityHarnessConfig{ + Default: true, + Endpoint: "localhost:50053", }, }, } @@ -50,25 +50,6 @@ func TestNewControllerFromConfig_DefaultHarness(t *testing.T) { c.Close() } -func TestNewControllerFromConfig_UnknownDefaultHarness(t *testing.T) { - cfg := &config2.Config{ - Harnesses: config2.HarnessesConfig{ - Default: "missing", - Antigravity: []config2.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) - } -} - func TestNewControllerFromConfig_BuiltinSubstrate(t *testing.T) { t.Setenv("AX_SUBSTRATE", "1") @@ -79,9 +60,8 @@ func TestNewControllerFromConfig_BuiltinSubstrate(t *testing.T) { }, }, Harnesses: config2.HarnessesConfig{ - Default: "ag", - Antigravity: []config2.AntigravityHarnessConfig{ - {ID: "ag"}, + Antigravity: config2.AntigravityHarnessConfig{ + Default: true, }, }, } diff --git a/internal/ax2.yaml b/internal/ax2.yaml index 9865b29..5e3cd24 100644 --- a/internal/ax2.yaml +++ b/internal/ax2.yaml @@ -24,7 +24,5 @@ eventlog: filename: "eventlog/log.sqlite" harnesses: - default: antigravity-example - antigravity: - - id: antigravity-example + default: true diff --git a/internal/config2/config.go b/internal/config2/config.go index af5dfeb..12eb5ab 100644 --- a/internal/config2/config.go +++ b/internal/config2/config.go @@ -26,13 +26,9 @@ import ( const ( // The substrate namespace reserved for AX's built-in harnesses. defaultNamespace = "ax" - // The default HarnessService port for non-substrate harnesses. - defaultPort = 50053 // The port for harnesses running as substrate actors. Substrate's // actor networking DNATs inbound workerPodIP:80 to the actor. substrateDefaultPort = 80 - // The Antigravity ActorTemplate name. - antigravityTemplate = "antigravity-template" ) // Config represents the main configuration for the AX harness server. @@ -63,38 +59,24 @@ type EventLogConfig struct { // - 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"` + 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 + Default bool `yaml:"default,omitempty"` + Endpoint string `yaml:"endpoint,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(c.ID, endpoint, defaultNamespace, antigravityTemplate, substrateDefaultPort) - } - address := c.Address - if address == "" { - address = fmt.Sprintf("localhost:%d", defaultPort) - } - return harness.NewAntigravityHarness(address), nil + 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 + Default bool `yaml:"default,omitempty"` // Default harness or not } // NewHarness builds the custom harness. Custom harnesses always run as substrate @@ -159,16 +141,18 @@ func (c *Config) Validate() error { 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") - } + var defaultCount int + if c.Harnesses.Antigravity.Default { + defaultCount++ } for _, sc := range c.Harnesses.Substrate { if sc.ID == "" { return fmt.Errorf("substrate harness id is required") } + if sc.ID == "antigravity" { + return fmt.Errorf("substrate harness id %q is reserved for the built-in antigravity harness", sc.ID) + } if sc.Namespace == "" { return fmt.Errorf("substrate harness %q: namespace is required", sc.ID) } @@ -178,6 +162,13 @@ func (c *Config) Validate() error { if sc.Template == "" { return fmt.Errorf("substrate harness %q: template is required", sc.ID) } + if sc.Default { + defaultCount++ + } + } + + if defaultCount > 1 { + return fmt.Errorf("multiple harnesses marked as default") } return nil diff --git a/internal/config2/config_test.go b/internal/config2/config_test.go index abdd6ed..0423790 100644 --- a/internal/config2/config_test.go +++ b/internal/config2/config_test.go @@ -19,26 +19,6 @@ import ( "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 { @@ -53,7 +33,7 @@ func TestSubstrateNewHarness(t *testing.T) { func validConfig() *Config { c := DefaultConfig() c.Harnesses = HarnessesConfig{ - Antigravity: []AntigravityHarnessConfig{{ID: "ag"}}, + Antigravity: AntigravityHarnessConfig{Default: true}, Substrate: []SubstrateHarnessConfig{ {ID: "custom", Namespace: "team-ns", Template: "custom-template"}, }, @@ -67,21 +47,21 @@ func TestValidate_ValidConfig(t *testing.T) { } } -func TestValidate_AntigravityIDRequired(t *testing.T) { +func TestValidate_CustomIDRequired(t *testing.T) { c := validConfig() - c.Harnesses.Antigravity[0].ID = "" + c.Harnesses.Substrate[0].ID = "" err := c.Validate() - if err == nil || !strings.Contains(err.Error(), "antigravity harness id") { - t.Fatalf("Validate() = %v, want antigravity id error", err) + if err == nil || !strings.Contains(err.Error(), "substrate harness id") { + t.Fatalf("Validate() = %v, want substrate id error", err) } } -func TestValidate_CustomIDRequired(t *testing.T) { +func TestValidate_CustomIDReserved(t *testing.T) { c := validConfig() - c.Harnesses.Substrate[0].ID = "" + c.Harnesses.Substrate[0].ID = "antigravity" err := c.Validate() - if err == nil || !strings.Contains(err.Error(), "substrate harness id") { - t.Fatalf("Validate() = %v, want substrate id error", err) + if err == nil || !strings.Contains(err.Error(), "reserved") { + t.Fatalf("Validate() = %v, want reserved id error", err) } } @@ -111,3 +91,12 @@ func TestValidate_CustomTemplateRequired(t *testing.T) { t.Fatalf("Validate() = %v, want template-required error", err) } } + +func TestValidate_MultipleDefaults(t *testing.T) { + c := validConfig() + c.Harnesses.Substrate[0].Default = true + err := c.Validate() + if err == nil || !strings.Contains(err.Error(), "multiple harnesses marked as default") { + t.Fatalf("Validate() = %v, want multiple defaults error", err) + } +} diff --git a/internal/controller2/controller.go b/internal/controller2/controller.go index 91e20d6..f842049 100644 --- a/internal/controller2/controller.go +++ b/internal/controller2/controller.go @@ -88,7 +88,7 @@ func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler E return fmt.Errorf("failed to start harness session: %w", err) } defer exec.Close(ctx) - + if err := exec.Queue(ctx, req.Inputs...); err != nil { return fmt.Errorf("failed to queue inputs: %w", err) } diff --git a/internal/manifests/ax-deployment2.yaml b/internal/manifests/ax-deployment2.yaml index b0743de..5d12159 100644 --- a/internal/manifests/ax-deployment2.yaml +++ b/internal/manifests/ax-deployment2.yaml @@ -39,7 +39,7 @@ spec: apiVersion: ate.dev/v1alpha1 kind: ActorTemplate metadata: - name: antigravity-template + name: ax-harness-template namespace: ax spec: workerPoolRef: @@ -104,8 +104,6 @@ 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 @@ -122,7 +120,5 @@ metadata: data: ax.yaml: | harnesses: - default: antigravity - # Built-in harness antigravity: - - id: antigravity + default: true