From abf86937e974aee68785c95a805ed79cee14c0b0 Mon Sep 17 00:00:00 2001 From: Jaana Dogan Date: Thu, 18 Jun 2026 00:30:11 -0700 Subject: [PATCH] Remove ATE agents SubstrateHarness will replace this. SubstrateAgent is deadcode. --- cmd/ax/internal/cliutil/cliutil.go | 5 -- internal/config/config.go | 34 ++------ internal/controller/registry_ate.go | 62 -------------- internal/experimental/agent/ate.go | 126 ---------------------------- 4 files changed, 9 insertions(+), 218 deletions(-) delete mode 100644 internal/controller/registry_ate.go delete mode 100644 internal/experimental/agent/ate.go diff --git a/cmd/ax/internal/cliutil/cliutil.go b/cmd/ax/internal/cliutil/cliutil.go index 9b1c1f9..f2dc728 100644 --- a/cmd/ax/internal/cliutil/cliutil.go +++ b/cmd/ax/internal/cliutil/cliutil.go @@ -109,11 +109,6 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Cont } } - for _, agentCfg := range cfg.Registry.SubstrateAgents { - if err := c.Registry().RegisterATE(ctx, cfg.ATE.Endpoint, agentCfg); err != nil { - return nil, fmt.Errorf("failed to register ATE agent %s: %w", agentCfg.ID, err) - } - } return c, nil } diff --git a/internal/config/config.go b/internal/config/config.go index 49cfe0b..50479a4 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -27,22 +27,21 @@ import ( // Config represents the main configuration for AX server. type Config struct { - Server ServerConfig `yaml:"server"` - EventLog EventLogConfig `yaml:"eventlog"` - Planner PlannerConfig `yaml:"planner,omitempty"` - Registry RegistryConfig `yaml:"registry,omitempty"` - ATE ATEConfig `yaml:"ate,omitempty"` + Server ServerConfig `yaml:"server"` + EventLog EventLogConfig `yaml:"eventlog"` + Planner PlannerConfig `yaml:"planner,omitempty"` + Registry RegistryConfig `yaml:"registry,omitempty"` + Substrate SubstrateConfig `yaml:"substrate,omitempty"` } // RegistryConfig allows registring agents. type RegistryConfig struct { - RemoteAgents []RemoteAgentConfig `yaml:"remote_agents,omitempty"` - ColabAgents []ColabAgentConfig `yaml:"colab_agents,omitempty"` - SubstrateAgents []SubstrateAgentConfig `yaml:"substrate_agents,omitempty"` + RemoteAgents []RemoteAgentConfig `yaml:"remote_agents,omitempty"` + ColabAgents []ColabAgentConfig `yaml:"colab_agents,omitempty"` } -// ATEConfig configures the ATE integration. -type ATEConfig struct { +// SubstrateConfig configures the Substrate integration. +type SubstrateConfig struct { Endpoint string `yaml:"endpoint"` } @@ -97,21 +96,6 @@ type GeminiConfig struct { Tools []string `json:"tools,omitempty" yaml:"tools,omitempty"` } -// SubstrateAgentConfig allows registering a new agent with an ATE actor template. -type SubstrateAgentConfig struct { - ID string `yaml:"id"` // Unique agent identifier - Name string `yaml:"name"` // Human-readable name - Description string `yaml:"description"` // Description of agent capabilities - Port int `yaml:"port"` // Port the agent is listening on - Namespace string `yaml:"namespace"` // Namespace for actor - Template string `yaml:"template"` // Template for actor - Protocol string `yaml:"protocol,omitempty"` // "axp" (default) or "a2a" - Auth auth.Auth `yaml:"auth,omitempty"` // Optional auth - Headers auth.Headers `yaml:"headers,omitempty"` // Optional headers - Metadata map[string]string `yaml:"metadata,omitempty"` // Optional metadata - // TODO(jbd): Rename this struct before releasing. -} - // RemoteAgentConfig configures a remote agent to register on startup. type RemoteAgentConfig struct { ID string `yaml:"id"` // Unique agent identifier diff --git a/internal/controller/registry_ate.go b/internal/controller/registry_ate.go deleted file mode 100644 index 2443089..0000000 --- a/internal/controller/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 controller - -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/experimental/agent/ate.go b/internal/experimental/agent/ate.go deleted file mode 100644 index 8197b43..0000000 --- a/internal/experimental/agent/ate.go +++ /dev/null @@ -1,126 +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 agent - -import ( - "context" - "fmt" - "log" - "strings" - "time" - - "github.com/google/ax/internal/agent" - "github.com/google/ax/internal/auth" - "github.com/google/ax/internal/experimental/k8s/ate" - "github.com/google/ax/proto" -) - -// SubstrateAgent manages execution in a SubstrATE actor. -type SubstrateAgent struct { - ateClient *ate.Client - config SubstrateAgentConfig -} - -// SubstrateAgentConfig configures an ATE agent client. -type SubstrateAgentConfig struct { - ID string - Namespace string - Template string - Port int // Port where agent runs in the worker - Protocol string - Auth auth.Auth - Headers auth.Headers -} - -// NewSubstrateAgent creates a new ATE agent client. -func NewSubstrateAgent(endpoint string, config SubstrateAgentConfig) (*SubstrateAgent, error) { - if config.Port == 0 { - config.Port = 8494 // Default port - } - client, err := ate.NewClient(config.Namespace, config.Template, endpoint) - if err != nil { - return nil, fmt.Errorf("failed to create ATE client: %w", err) - } - return &SubstrateAgent{ - ateClient: client, - config: config, - }, nil -} - -// Connect handles processing of input content by creating an actor and delegating to RemoteAgent. -func (a *SubstrateAgent) Connect(ctx context.Context, conversationID string, execID string, start *proto.AgentStart, e agent.Executor, o agent.OutputHandler) error { - // 1. Create Actor - resp, err := a.ateClient.CreateActor(ctx, execID) - if err != nil { - return fmt.Errorf("failed to create actor: %w", err) - } - actor := resp.Actor - if actor == nil { - return fmt.Errorf("received nil actor in response") - } - if actor.AteomPodIp == "" { - return fmt.Errorf("actor has no active worker IP address (AteomPodIp is empty)") - } - - workerAddr := fmt.Sprintf("%s:%d", actor.AteomPodIp, a.config.Port) - // 2. Connect to the Actor. - var activeAgent agent.Agent - switch strings.ToLower(a.config.Protocol) { - case "", "axp": - activeAgent, err = agent.NewRemoteAgent(agent.RemoteAgentConfig{ - Address: workerAddr, - Reconnect: true, - MaxRetries: 3, - }) - case "a2a": - activeAgent, err = NewA2AAgent(ctx, A2AAgentConfig{ - ID: a.config.ID, - Address: workerAddr, - Auth: a.config.Auth, - Headers: a.config.Headers, - Stateless: true, - OverrideCardHosts: true, - }) - default: - return fmt.Errorf("ate agent %s: invalid protocol %q", a.config.ID, a.config.Protocol) - } - if err != nil { - return fmt.Errorf("failed to create remote agent connection: %w", err) - } - defer activeAgent.Close() - - // 3. Suspend Actor when done. - defer func() { - // TODO(jbd): Actors need to be deleted when they are done. - // DeleteActor once deletion is possible. For now, suspending allows - // us to return the worker back to the pool. - log.Printf("Suspending ATE actor for execution %s", execID) - suspendCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - if _, err := a.ateClient.SuspendActor(suspendCtx, execID); err != nil { - log.Printf("Failed to suspend actor %s: %v", execID, err) - } - }() - - return activeAgent.Connect(ctx, conversationID, execID, start, e, o) -} - -// Close gracefully shuts down the ATE agent connection. -func (a *SubstrateAgent) Close() error { - if a.ateClient != nil { - return a.ateClient.Close() - } - return nil -}