Use SubstrateHarness if AX_SUBSTRATE=1

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)
 		}
 	}
 

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,
 			},
 		},
 	}

internal/ax2.yaml

@@ -24,7 +24,5 @@ eventlog:
     filename: "eventlog/log.sqlite"
 
 harnesses:
-  default: antigravity-example
-
   antigravity:
-    - id: antigravity-example
+    default: true

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

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)
+	}
+}

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)
 	}

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