From 7ca09e5c8e96fb84dd9de169a944b6c4347f5cda Mon Sep 17 00:00:00 2001 From: Anjali Sridhar Date: Tue, 2 Jun 2026 11:42:19 -0700 Subject: [PATCH] feat: strictly block legacy agent instantiation under controller V2 build tag --- cmd/ax/internal/cliutil/cliutil_harness.go | 9 +++++++++ internal/controller2/controller.go | 6 +----- internal/controller2/controller_test.go | 14 +++---------- internal/controller2/registry.go | 5 ++++- internal/controller2/registry_test.go | 23 ++++++++++++++++++++++ 5 files changed, 40 insertions(+), 17 deletions(-) diff --git a/cmd/ax/internal/cliutil/cliutil_harness.go b/cmd/ax/internal/cliutil/cliutil_harness.go index 11591b6..8813da2 100644 --- a/cmd/ax/internal/cliutil/cliutil_harness.go +++ b/cmd/ax/internal/cliutil/cliutil_harness.go @@ -35,6 +35,15 @@ type ExecHandler = controller2.ExecHandler func NewControllerFromConfig(ctx context.Context, cfg *config.Config) (*controller2.Controller, error) { reg := controller2.NewRegistry() reg.RegisterHarness("antigravity", harness.NewAntigravityHarness("")) + + // Register legacy agents so that the V2 registry can detect and block them with a clear validation error + for _, agentCfg := range cfg.Registry.RemoteAgents { + _ = reg.RegisterRemote(ctx, agentCfg) + } + for _, agentCfg := range cfg.Registry.ColabAgents { + _ = reg.RegisterColab(agentCfg) + } + return controller2.New(ctx, controller2.Config{ Registry: reg, EventLogBuilder: func() (executor.EventLog, error) { diff --git a/internal/controller2/controller.go b/internal/controller2/controller.go index c41edf1..a8fbcde 100644 --- a/internal/controller2/controller.go +++ b/internal/controller2/controller.go @@ -19,10 +19,8 @@ package controller2 import ( "context" "fmt" - "log" "github.com/google/ax/internal/controller/executor" - "github.com/google/ax/internal/harness/harnesstest" "github.com/google/ax/proto" "github.com/google/uuid" ) @@ -75,9 +73,7 @@ func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler E // Adding harness registration support temporarily. h, err := d.registry.Harness(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() + return fmt.Errorf("harness/agent %q not found in registry: %w", req.AgentId, err) } exec, err := h.Start(ctx, req.ConversationId) if err != nil { diff --git a/internal/controller2/controller_test.go b/internal/controller2/controller_test.go index 259087f..48b03a7 100644 --- a/internal/controller2/controller_test.go +++ b/internal/controller2/controller_test.go @@ -41,6 +41,7 @@ func TestController2_ExecHelloWorld(t *testing.T) { log := &executortest.MemoryEventLog{} reg := NewRegistry() + reg.RegisterHarness("", harnesstest.New()) c, err := New(ctx, Config{ Registry: reg, EventLogBuilder: func() (executor.EventLog, error) { @@ -193,17 +194,8 @@ func TestController2_ExecRuntimeFallback(t *testing.T) { Inputs: inputs, AgentId: "antigravity", }, handler) - if err != nil { - t.Fatalf("Controller2.Exec failed: %v", err) - } - - if len(outputs) != 1 { - t.Fatalf("expected exactly 1 output message, got %d", len(outputs)) - } - - gotText := outputs[0].GetContent().GetText().GetText() - if gotText != "Hello world" { - t.Errorf("expected 'Hello world' output text response due to runtime fallback, got %q", gotText) + if err == nil { + t.Fatal("expected error for unregistered agent lookup under V2, got nil") } } diff --git a/internal/controller2/registry.go b/internal/controller2/registry.go index 0296a12..ca619e0 100644 --- a/internal/controller2/registry.go +++ b/internal/controller2/registry.go @@ -253,7 +253,10 @@ func (r *Registry) Harness(id string) (harness.Harness, error) { defer r.mu.RUnlock() h, ok := r.harnesses[id] if !ok { - return nil, fmt.Errorf("harness %s not found", id) + if _, exists := r.agents[id]; exists { + return nil, fmt.Errorf("agent %q is registered under legacy agent config but has no V2 harness interface implementation", id) + } + return nil, fmt.Errorf("harness %q not found (local V2 execution only)", id) } return h, nil } diff --git a/internal/controller2/registry_test.go b/internal/controller2/registry_test.go index 9052f2a..9fcdaf4 100644 --- a/internal/controller2/registry_test.go +++ b/internal/controller2/registry_test.go @@ -73,3 +73,26 @@ func TestRegistry_GracefulShutdown(t *testing.T) { // We are testing for absence of panic/deadlock here. _ = r.Close() } + +func TestRegistry_LegacyAgentValidation(t *testing.T) { + r := NewRegistry() + + err := r.RegisterRemote(context.Background(), config.RemoteAgentConfig{ + ID: "legacy-remote-agent", + Name: "Legacy Agent", + Address: "localhost:8494", + }) + if err != nil { + t.Fatalf("failed to register remote agent: %v", err) + } + + _, err = r.Harness("legacy-remote-agent") + if err == nil { + t.Fatal("expected error for legacy agent lookup in V2 mode, got nil") + } + + expectedErr := `agent "legacy-remote-agent" is registered under legacy agent config but has no V2 harness interface implementation` + if err.Error() != expectedErr { + t.Errorf("expected error message %q, got %q", expectedErr, err.Error()) + } +}