Initial commit of the Substrate refactor

Makefile

@@ -74,10 +74,14 @@ axepp-image:
 	@echo "Building axepp container image with ko..."
 	GOFLAGS="-tags=ate" ko build --base-import-paths ./cmd/axepp
 
+axharness-image:
+	@echo "Building axharness container image with ko..."
+	ko build --base-import-paths ./cmd/axharness
+
 ax-shell-image:
 	# Used to debug ax servers within a cluster.
 	@echo "Building ax shell container image with ko using busybox..."
 	KO_DOCKER_REPO=$(KO_DOCKER_REPO)/ax-shell KO_DEFAULTBASEIMAGE=busybox:1.36 GOFLAGS="-tags=ate" ko build --base-import-paths ./cmd/ax
 
 # Build all container images
-images: ax-image axepp-image ax-shell-image
+images: ax-image axepp-image axharness-image ax-shell-image

cmd/axharness/main.go

@@ -0,0 +1,61 @@
+// 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 main
+
+import (
+	"flag"
+	"fmt"
+	"log"
+	"net"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"google.golang.org/grpc"
+
+	"github.com/google/ax/proto"
+)
+
+var (
+	port = flag.Int("port", 50053, "The port for the gRPC HarnessService to listen on")
+)
+
+func main() {
+	flag.Parse()
+
+	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
+	if err != nil {
+		log.Fatalf("Failed to listen on port :%d: %v", *port, err)
+	}
+
+	// Start gRPC Server
+	grpcServer := grpc.NewServer()
+	harnessServer := NewHarnessServiceServer()
+	proto.RegisterHarnessServiceServer(grpcServer, harnessServer)
+
+	// Graceful shutdown handling
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-sigChan
+		log.Println("\nReceived shutdown signal, stopping gRPC HarnessService server gracefully...")
+		grpcServer.GracefulStop()
+	}()
+
+	log.Printf("gRPC HarnessService listening on port :%d...\n", *port)
+	if err := grpcServer.Serve(lis); err != nil {
+		log.Fatalf("Failed to serve gRPC: %v", err)
+	}
+}

cmd/axharness/service.go

@@ -0,0 +1,63 @@
+// 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 main
+
+import (
+	"io"
+
+	"github.com/google/ax/proto"
+)
+
+// HarnessServiceServer implements the gRPC proto.HarnessServiceServer interface.
+type HarnessServiceServer struct {
+	proto.UnimplementedHarnessServiceServer
+}
+
+// NewHarnessServiceServer creates a new HarnessServiceServer.
+func NewHarnessServiceServer() *HarnessServiceServer {
+	return &HarnessServiceServer{}
+}
+
+// Connect implements the bidirectional gRPC streaming capability.
+// It receives client inputs and responds only with "Hello world".
+func (s *HarnessServiceServer) Connect(stream proto.HarnessService_ConnectServer) error {
+	for {
+		_, err := stream.Recv()
+		if err == io.EOF {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+
+		err = stream.Send(&proto.HarnessMessage{
+			Messages: []*proto.Message{
+				{
+					Role: "assistant",
+					Content: &proto.Content{
+						Type: &proto.Content_Text{
+							Text: &proto.TextContent{
+								Text: "Hello world",
+							},
+						},
+					},
+				},
+			},
+		})
+		if err != nil {
+			return err
+		}
+	}
+}

internal/SUBSTRATE_REFACTOR_CLEANUPS.md

@@ -0,0 +1,5 @@
+* Make internal/manifests/ax-deployment2.yaml the new templated ax-deployment.yaml.tmpl.
+* Remove harnesstest package.
+* Update HarnessService with the actual protocol.
+* Remove axepp.
+* Remove ate build tag.

internal/controller/registry_ate.go

@@ -1,5 +1,3 @@
-//go:build ate
-
 // Copyright 2026 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");

internal/controller2/controller.go

@@ -0,0 +1,218 @@
+// 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 implements the single-writer orchestrator that coordinates
+// agentic loops, manages executions, and communicates with local and remote agents.
+package controller2
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/google/ax/internal/agent"
+	"github.com/google/ax/internal/controller/executor"
+	"github.com/google/ax/internal/gemini"
+	"github.com/google/ax/internal/harness/harnesstest"
+	"github.com/google/ax/proto"
+	"github.com/google/uuid"
+)
+
+const (
+	plannerAgentID = "__planner"
+	geminiAgentID  = "gemini"
+)
+
+var reservedAgentIDs = map[string]struct{}{
+	plannerAgentID: {},
+	geminiAgentID:  {},
+}
+
+type ExecHandler func(resp *proto.ExecResponse) error
+
+// Controller is the main controller that coordinates all components.
+// It acts as a single-writer system for managing agentic loops.
+type Controller struct {
+	registry       *Registry
+	eventLog       executor.EventLog
+	plannerBuilder PlannerBuilder
+}
+
+// PlannerBuilder is a function that creates a PlanFunc given a Registry.
+type PlannerBuilder func(ctx context.Context, r *Registry) (agent.Agent, error)
+
+// Config configures the controller.
+type Config struct {
+	EventLogBuilder executor.EventLogBuilder
+	PlannerBuilder  PlannerBuilder
+}
+
+// New creates a new controller instance.
+func New(ctx context.Context, cfg Config) (*Controller, error) {
+	// Initialize agent registry
+	registry := NewRegistry()
+
+	// Determine plan function
+	// If no planner builder is provided, use the default Gemini planner.
+	if cfg.PlannerBuilder == nil {
+		cfg.PlannerBuilder = func(ctx context.Context, r *Registry) (agent.Agent, error) {
+			return gemini.NewGeminiPlannerAgent(ctx, r, gemini.GeminiPlannerConfig{})
+		}
+	}
+
+	if cfg.EventLogBuilder == nil {
+		return nil, fmt.Errorf("event log builder is required")
+	}
+	eventLog, err := cfg.EventLogBuilder()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create event log: %w", err)
+	}
+
+	return &Controller{
+		registry:       registry,
+		eventLog:       eventLog,
+		plannerBuilder: cfg.PlannerBuilder,
+	}, nil
+}
+
+// Exec executes a new agentic loop execution or resumes an existing one.
+// If id is empty, a UUID will be generated.
+// If the execution already exists, it will be resumed with optional new inputs.
+func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler ExecHandler) error {
+	if req.ConversationId == "" {
+		return fmt.Errorf("conversation_id is required")
+	}
+
+	// TODO(jbd): Resume an incomplete execution if there exists one.
+	// TODO(jbd): Enable bringing a remote harness that implements HarnessService.
+
+	h := harnesstest.New()
+	exec, err := h.Start(ctx, req.ConversationId)
+	if err != nil {
+		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)
+	}
+
+	hhandler := &harnessHandler{
+		execHandler: handler,
+	}
+	if err := exec.Run(ctx, hhandler); err != nil {
+		return fmt.Errorf("harness execution turn failed: %w", err)
+	}
+
+	return nil
+}
+
+type harnessHandler struct {
+	execHandler ExecHandler
+}
+
+func (a *harnessHandler) OnMessage(ctx context.Context, execID string, msg *proto.Message) error {
+	if a.execHandler == nil {
+		return nil
+	}
+	return a.execHandler(&proto.ExecResponse{
+		Outputs: []*proto.Message{msg},
+	})
+}
+
+func (a *harnessHandler) OnComplete(ctx context.Context, execID string) error {
+	return nil
+}
+
+// Delete deletes all events for a specific conversation ID.
+func (d *Controller) Delete(ctx context.Context, conversationID string) error {
+	if conversationID == "" {
+		return fmt.Errorf("conversation_id is required")
+	}
+
+	return d.eventLog.DeleteEvents(ctx, conversationID)
+}
+
+// Fork forks an event log from a specific conversation up to a checkpoint.
+func (d *Controller) Fork(ctx context.Context, srcConversationID string, srcSeq int32, destConversationID string) (string, error) {
+	if srcConversationID == "" {
+		return "", fmt.Errorf("src_conversation_id is required")
+	}
+	// TODO(anj-s): Check whether destination ID already exists and reject collisions.
+	if destConversationID == "" {
+		destConversationID = uuid.NewString()
+	}
+
+	events, err := d.eventLog.Events(ctx, srcConversationID)
+	if err != nil {
+		return "", fmt.Errorf("failed to retrieve source events: %w", err)
+	}
+	if len(events) == 0 {
+		return "", fmt.Errorf("source conversation %s not found or has no events", srcConversationID)
+	}
+
+	// When the caller specifies srcSeq, require that it actually exists in
+	// the source event log. Without this check a typo or stale checkpoint
+	// silently degrades to "fork all events", which is misleading. Walk
+	// the events once: stop as soon as we pass the requested seq, and
+	// truncate the slice on an exact match so the copy loop below doesn't
+	// need to re-check the bound.
+	if srcSeq > 0 {
+		found := false
+		for i, ev := range events {
+			if ev.Seq == srcSeq {
+				events = events[:i+1]
+				found = true
+				break
+			}
+			if ev.Seq > srcSeq {
+				break
+			}
+		}
+		if !found {
+			return "", fmt.Errorf("src_seq %d not found in conversation %s", srcSeq, srcConversationID)
+		}
+	}
+
+	for _, ev := range events {
+		// Clone the event to update the conversation ID.
+		newEvent := &proto.ConversationEvent{
+			ConversationId: destConversationID,
+			Seq:            ev.Seq,
+			ExecId:         ev.ExecId,
+			Messages:       ev.Messages,
+			State:          ev.State,
+		}
+		if _, err := d.eventLog.Append(ctx, newEvent); err != nil {
+			return "", fmt.Errorf("failed to append forked event: %w", err)
+		}
+	}
+
+	return destConversationID, nil
+}
+
+// Registry returns the agent registry.
+func (d *Controller) Registry() *Registry {
+	return d.registry
+}
+
+// Close gracefully shuts down the controller.
+func (d *Controller) Close() error {
+	if err := d.eventLog.Close(); err != nil {
+		return fmt.Errorf("failed to close event log: %w", err)
+	}
+	if err := d.registry.Close(); err != nil {
+		return fmt.Errorf("failed to close registry: %w", err)
+	}
+	return nil
+}