Build one comprehensive ax image for the ax-server and harness

cmd/ax/Dockerfile

@@ -0,0 +1,76 @@
+# 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.
+
+# Comprehensive AX image: the Go `ax` binary plus the Antigravity Python sidecar
+# (SDK, localharness, agent), on a Debian-based image with shell tooling for
+# skills/code execution. Used by both the ax-server (`ax serve`) and the harness
+# actor (`ax harness`, which forks the Python sidecar).
+#
+# Build context: repository root. The Antigravity runtime deps -- including the
+# linux/amd64 google-antigravity wheel that bundles the linux `localharness`
+# binary -- are installed offline from a "wheels" build context, so the build
+# needs no access to the private package registry. Populate the wheel cache with
+# `internal/hack/install-ax.sh --fetch-wheels` (which also drives build/deploy).
+#
+# Requires an OCI builder that supports additional build contexts and
+# RUN --mount: docker with BuildKit (default since Docker 23; older Docker: use
+# `docker buildx build`), or podman/buildah >= 4.0. For example:
+#   docker build --platform linux/amd64 \
+#     --build-context wheels=$HOME/.cache/ax-antigravity-wheels \
+#     -f cmd/ax/Dockerfile -t <ref> . && docker push <ref>
+
+# --- Stage 1: build the ax Go binary (with the harness build tag) -------------
+FROM --platform=$BUILDPLATFORM golang:1.26 AS build
+ARG TARGETOS=linux
+ARG TARGETARCH=amd64
+WORKDIR /src
+COPY go.mod go.sum ./
+RUN go mod download
+COPY . .
+RUN CGO_ENABLED=0 GOOS=${TARGETOS} GOARCH=${TARGETARCH} \
+    go build -tags=harness -o /out/ax ./cmd/ax
+
+# --- Stage 2: comprehensive runtime (Debian + Python + Antigravity) -----------
+# Python image: a full Debian base (buildpack-deps) that
+# already ships bash, git, curl, wget, ca-certificates, build-essential, and the
+# common CLIs (find/ps/less/grep/sed/...).
+FROM python:3.13
+WORKDIR /app
+
+# Antigravity runtime deps, installed offline from the pre-downloaded wheels.
+COPY python/antigravity/requirements.txt /tmp/requirements.txt
+RUN --mount=type=bind,from=wheels,target=/tmp/wheels \
+    pip install --no-cache-dir --no-index --find-links=/tmp/wheels -r /tmp/requirements.txt
+
+# Python sidecar (with generated proto stubs), the agent it serves, and the ax
+# binary built in stage 1.
+COPY python/ /app/python
+COPY examples/antigravity_agent/ /app/examples/antigravity_agent
+COPY --from=build /out/ax /ax
+
+# `from python.proto import ...` needs /app on the path; the generated stubs do
+# `from proto import ...`, which needs /app/python.
+ENV PYTHONPATH=/app:/app/python
+
+# Flush stdout/stderr immediately so server logs surface in container logs.
+ENV PYTHONUNBUFFERED=1
+
+EXPOSE 80
+
+# Default command for local docker runs. Substrate ignores the image CMD and runs
+# the ActorTemplate `command` instead. `ax harness` forks the Python sidecar,
+# which serves the HarnessService on :80.
+CMD ["/ax", "harness", \
+     "--antigravity-agent-file", "/app/examples/antigravity_agent/agent.py", \
+     "--host", "0.0.0.0", "--port", "80"]

cmd/ax/harness.go

@@ -12,30 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package main implements a demo HarnessService.
-// It is intended for testing purposes only and should be replaced with a real
-// implementation in production.
-// TODO(wjjclaud): Replace this file with a real harness implementation.
+// Package main: the `ax harness` command. It supervises the Antigravity Python
+// sidecar server (which serves the HarnessService and gRPC health), forking it
+// as a child process and forwarding termination signals.
 package main
 
 import (
 	"fmt"
 	"log"
-	"net"
 	"os"
+	"os/exec"
 	"os/signal"
-	"sync"
+	"strconv"
 	"syscall"
 
-	"github.com/google/ax/proto"
 	"github.com/spf13/cobra"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/health"
-	"google.golang.org/grpc/health/grpc_health_v1"
 )
 
 var (
-	harnessPort int
+	harnessPort                 int
+	harnessHost                 string
+	harnessAntigravityAgentFile string
 )
 
 var harnessCmd = &cobra.Command{
@@ -46,152 +43,42 @@ var harnessCmd = &cobra.Command{
 }
 
 func init() {
-	harnessCmd.Flags().IntVar(&harnessPort, "port", 50053, "The port for the gRPC HarnessService to listen on")
+	harnessCmd.Flags().IntVar(&harnessPort, "port", 50053, "Port for the HarnessService to listen on")
+	harnessCmd.Flags().StringVar(&harnessHost, "host", "127.0.0.1", "Host interface for the HarnessService to bind")
+	harnessCmd.Flags().StringVar(&harnessAntigravityAgentFile, "antigravity-agent-file", "examples/antigravity_agent/agent.py", "Path to the agent config file the Python sidecar serves")
 	rootCmd.AddCommand(harnessCmd)
 }
 
+// runHarness forks the Antigravity Python sidecar server, which serves the
+// HarnessService (and gRPC health) on the configured port. ax harness supervises
+// the child: it forwards termination signals and exits with the child's status.
 func runHarness(cmd *cobra.Command, args []string) error {
-	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", harnessPort))
-	if err != nil {
-		return fmt.Errorf("failed to listen on port :%d: %w", harnessPort, err)
+	py := exec.Command("python3", "-m", "python.antigravity.harness_server",
+		"--host", harnessHost,
+		"--port", strconv.Itoa(harnessPort),
+		"--agent_file", harnessAntigravityAgentFile,
+	)
+	py.Stdin = os.Stdin
+	py.Stdout = os.Stdout
+	py.Stderr = os.Stderr
+	py.Env = os.Environ()
+
+	if err := py.Start(); err != nil {
+		return fmt.Errorf("failed to start antigravity harness server: %w", err)
 	}
+	log.Printf("forked antigravity harness server (pid %d) on %s:%d", py.Process.Pid, harnessHost, harnessPort)
 
-	// Start gRPC Server
-	grpcServer := grpc.NewServer()
-	harnessServer := NewHarnessServiceServer()
-	proto.RegisterHarnessServiceServer(grpcServer, harnessServer)
-
-	// Serve the standard gRPC health protocol.
-	healthServer := health.NewServer()
-	healthServer.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING)
-	grpc_health_v1.RegisterHealthServer(grpcServer, healthServer)
-
-	// Graceful shutdown handling
+	// Forward termination signals to the child so substrate can stop the actor.
 	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()
+		for sig := range sigChan {
+			_ = py.Process.Signal(sig)
+		}
 	}()
 
-	log.Printf("gRPC HarnessService listening on port :%d...\n", harnessPort)
-	if err := grpcServer.Serve(lis); err != nil {
-		return fmt.Errorf("failed to serve gRPC: %w", err)
+	if err := py.Wait(); err != nil {
+		return fmt.Errorf("antigravity harness server exited: %w", err)
 	}
 	return nil
 }
-
-// conversationState is the per-conversation state the stub keeps in process memory.
-// On substrate this state is preserved across turns by snapshot/suspend/resume.
-type conversationState struct {
-	turns   int
-	history []string
-}
-
-// HarnessServiceServer implements the gRPC proto.HarnessServiceServer interface.
-type HarnessServiceServer struct {
-	proto.UnimplementedHarnessServiceServer
-
-	mu            sync.Mutex
-	conversations map[string]*conversationState
-}
-
-// NewHarnessServiceServer creates a new HarnessServiceServer.
-func NewHarnessServiceServer() *HarnessServiceServer {
-	return &HarnessServiceServer{conversations: make(map[string]*conversationState)}
-}
-
-// Connect implements one HarnessService turn. It reads the initial
-// HarnessRequest{start} frame, then replies with a "hello world (turn N)"
-// frame, an echo of each input, and a recap of the inputs from prior turns,
-// terminating with HarnessEnd{STATE_COMPLETED}.
-//
-// The per-conversation turn count and input history are kept in process
-// memory and persist across turns.
-func (s *HarnessServiceServer) Connect(stream proto.HarnessService_ConnectServer) error {
-	req, err := stream.Recv()
-	if err != nil {
-		return err
-	}
-
-	convID := req.GetConversationId()
-	if req.GetStart() == nil {
-		return stream.Send(&proto.HarnessResponse{
-			ConversationId: convID,
-			Type: &proto.HarnessResponse_End{
-				End: &proto.HarnessEnd{
-					State:        proto.State_STATE_FAILED,
-					ErrorMessage: "expected HarnessRequest{start} as the first frame",
-				},
-			},
-		})
-	}
-
-	// Collect this turn's input text(s).
-	var inputs []string
-	for _, m := range req.GetStart().GetMessages() {
-		if text := m.GetContent().GetText().GetText(); text != "" {
-			inputs = append(inputs, text)
-		}
-	}
-
-	// Update per-conversation state held in process memory.
-	s.mu.Lock()
-	st := s.conversations[convID]
-	if st == nil {
-		st = &conversationState{}
-		s.conversations[convID] = st
-	}
-	st.turns++
-	turn := st.turns
-	prior := ""
-	if len(st.history) > 0 {
-		prior = st.history[len(st.history)-1]
-	}
-	st.history = append(st.history, inputs...)
-	s.mu.Unlock()
-
-	// Reply: turn number, this turn's inputs, and the remembered prior inputs.
-	if err := stream.Send(textOutput(convID, fmt.Sprintf("hello world (turn %d)", turn))); err != nil {
-		return err
-	}
-	for _, in := range inputs {
-		if err := stream.Send(textOutput(convID, "received: "+in)); err != nil {
-			return err
-		}
-	}
-	if prior != "" {
-		if err := stream.Send(textOutput(convID, "previously you said: "+prior)); err != nil {
-			return err
-		}
-	}
-
-	return stream.Send(&proto.HarnessResponse{
-		ConversationId: convID,
-		Type: &proto.HarnessResponse_End{
-			End: &proto.HarnessEnd{State: proto.State_STATE_COMPLETED},
-		},
-	})
-}
-
-// textOutput builds a HarnessResponse carrying a single assistant text Message.
-func textOutput(convID, text string) *proto.HarnessResponse {
-	return &proto.HarnessResponse{
-		ConversationId: convID,
-		Type: &proto.HarnessResponse_Outputs{
-			Outputs: &proto.HarnessOutputs{
-				Messages: []*proto.Message{
-					{
-						Role: "assistant",
-						Content: &proto.Content{
-							Type: &proto.Content_Text{
-								Text: &proto.TextContent{Text: text},
-							},
-						},
-					},
-				},
-			},
-		},
-	}
-}

go.mod

@@ -6,7 +6,7 @@ require (
 	charm.land/huh/v2 v2.0.3
 	charm.land/lipgloss/v2 v2.0.3
 	github.com/a2aproject/a2a-go/v2 v2.2.0
-	github.com/agent-substrate/substrate v0.0.0-20260616211440-0ec9b650a0ab
+	github.com/agent-substrate/substrate v0.0.0-20260617164646-00a90a8eda76
 	github.com/envoyproxy/go-control-plane/envoy v1.37.0
 	github.com/google/uuid v1.6.0
 	github.com/spf13/cobra v1.10.2

go.sum

@@ -16,8 +16,8 @@ github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ
 github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE=
 github.com/a2aproject/a2a-go/v2 v2.2.0 h1:eayiNXYpyTOLVhhQrGmIHlcy8GnOdnwaNdYQPvS84Ik=
 github.com/a2aproject/a2a-go/v2 v2.2.0/go.mod h1:htTxMwicNXXXEwwfjuB/Pd1g7UHDrswhSievncmTVcE=
-github.com/agent-substrate/substrate v0.0.0-20260616211440-0ec9b650a0ab h1:ESQ0AWzhw3GROJ5ZIT6tzqHg92mIEqdp/esSmK3s6zQ=
-github.com/agent-substrate/substrate v0.0.0-20260616211440-0ec9b650a0ab/go.mod h1:TgdtEUV6iaflJTwmS8ONiGsyyJD+5okPZj2H6mM8WlA=
+github.com/agent-substrate/substrate v0.0.0-20260617164646-00a90a8eda76 h1:GO7ajMplVYHUXAqT4yWTWKGMWsWehoL3TI5GXOwJSvo=
+github.com/agent-substrate/substrate v0.0.0-20260617164646-00a90a8eda76/go.mod h1:TgdtEUV6iaflJTwmS8ONiGsyyJD+5okPZj2H6mM8WlA=
 github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
 github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
 github.com/aymanbagabas/go-udiff v0.4.1 h1:OEIrQ8maEeDBXQDoGCbbTTXYJMYRCRO1fnodZ12Gv5o=