Use harness service in antigravity implementation.

cmd/ax/harness.go

@@ -16,11 +16,11 @@ package main
 
 import (
 	"fmt"
-	"io"
 	"log"
 	"net"
 	"os"
 	"os/signal"
+	"sync"
 	"syscall"
 
 	"github.com/google/ax/proto"
@@ -71,61 +71,116 @@ func runHarness(cmd *cobra.Command, args []string) error {
 	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{}
+	return &HarnessServiceServer{conversations: make(map[string]*conversationState)}
 }
 
-// Connect implements the bidirectional gRPC streaming capability.
-// It receives client inputs and responds with "hello world" unless the input message text is "go_away".
-// TODO(params): Update the implementation to be a proper one.
+// 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 {
-	for {
-		req, err := stream.Recv()
-		if err == io.EOF {
-			return nil
+	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)
 		}
-		if err != nil {
+	}
+
+	// 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
 		}
-
-		shouldGoAway := false
-		for _, m := range req.Messages {
-			if textBlock, ok := m.Content.Type.(*proto.Content_Text); ok {
-				// TODO(params): Replace this with a proper protocol for go away.
-				if textBlock.Text.Text == "go_away" {
-					shouldGoAway = true
-					break
-				}
-			}
+	}
+	if prior != "" {
+		if err := stream.Send(textOutput(convID, "previously you said: "+prior)); err != nil {
+			return err
 		}
+	}
 
-		if shouldGoAway {
-			log.Println("Received 'go_away' message, closing stream...")
-			return nil
-		}
+	return stream.Send(&proto.HarnessResponse{
+		ConversationId: convID,
+		Type: &proto.HarnessResponse_End{
+			End: &proto.HarnessEnd{State: proto.State_STATE_COMPLETED},
+		},
+	})
+}
 
-		err = stream.Send(&proto.HarnessMessage{
-			Messages: []*proto.Message{
-				{
-					Role: "assistant",
-					Content: &proto.Content{
-						Type: &proto.Content_Text{
-							Text: &proto.TextContent{
-								Text: "hello world",
+// 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},
 							},
 						},
 					},
 				},
 			},
-		})
-		if err != nil {
-			return err
-		}
+		},
 	}
 }

cmd/ax/harnessclient.go

@@ -31,6 +31,7 @@ import (
 
 var (
 	harnessServerAddr string
+	harnessClientID   string
 )
 
 var harnessClientCmd = &cobra.Command{
@@ -42,6 +43,7 @@ var harnessClientCmd = &cobra.Command{
 
 func init() {
 	harnessClientCmd.Flags().StringVar(&harnessServerAddr, "server", "localhost:50053", "The server address for the gRPC HarnessService.")
+	harnessClientCmd.Flags().StringVar(&harnessClientID, "harness", "", "The harness id to send on the request envelope.")
 	rootCmd.AddCommand(harnessClientCmd)
 }
 
@@ -57,72 +59,63 @@ func runHarnessClient(cmd *cobra.Command, args []string) error {
 
 	client := proto.NewHarnessServiceClient(conn)
 
+	fmt.Print("Client > ")
+	scanner := bufio.NewScanner(os.Stdin)
+	scanner.Scan()
+	input := scanner.Text()
+
 	stream, err := client.Connect(ctx)
 	if err != nil {
-		return fmt.Errorf("Failed to open connection stream: %v", err)
+		return fmt.Errorf("failed to open connection stream: %v", err)
 	}
 
-	scanner := bufio.NewScanner(os.Stdin)
-	fmt.Println("Interactive client started. Type your messages below.")
-	fmt.Println("Type 'go_away' to close the stream and exit.")
-	for {
-		fmt.Print("\nClient > ")
-		if !scanner.Scan() {
-			break
-		}
-		text := scanner.Text()
-		if text == "" {
-			continue
-		}
-
-		msg := &proto.HarnessMessage{
-			Messages: []*proto.Message{
-				{
-					Role: "user",
-					Content: &proto.Content{
-						Type: &proto.Content_Text{
-							Text: &proto.TextContent{
-								Text: text,
-							},
+	// A single HarnessRequest{start} initiates the turn.
+	start := &proto.HarnessRequest{
+		ConversationId: "harnessclient",
+		HarnessId:      harnessClientID,
+		Type: &proto.HarnessRequest_Start{
+			Start: &proto.HarnessStart{
+				Messages: []*proto.Message{
+					{
+						Role: "user",
+						Content: &proto.Content{
+							Type: &proto.Content_Text{Text: &proto.TextContent{Text: input}},
 						},
 					},
 				},
 			},
-		}
+		},
+	}
+	if err := stream.Send(start); err != nil {
+		return fmt.Errorf("failed to send start: %v", err)
+	}
+	if err := stream.CloseSend(); err != nil {
+		return fmt.Errorf("failed to close send side: %v", err)
+	}
 
-		if err := stream.Send(msg); err != nil {
-			return fmt.Errorf("Failed to send message: %v", err)
-		}
-		// TODO(params): Replace this with a proper protocol for go away.
-		if text == "go_away" {
-			log.Println("Sending 'go_away' to close the stream...")
+	// Drain HarnessResponse frames until HarnessEnd / EOF.
+	for {
+		resp, err := stream.Recv()
+		if err == io.EOF {
 			break
 		}
-
-		resp, err := stream.Recv()
 		if err != nil {
-			return fmt.Errorf("Failed to receive response: %v", err)
+			return fmt.Errorf("failed to receive response: %v", err)
 		}
-
-		for i, m := range resp.Messages {
-			var textContent string
-			if textBlock, ok := m.Content.Type.(*proto.Content_Text); ok {
-				textContent = textBlock.Text.Text
+		switch payload := resp.Type.(type) {
+		case *proto.HarnessResponse_Outputs:
+			for i, m := range payload.Outputs.Messages {
+				var text string
+				if tb, ok := m.Content.Type.(*proto.Content_Text); ok {
+					text = tb.Text.Text
+				}
+				fmt.Printf("Server > message[%d] (%s): %s\n", i, m.Role, text)
 			}
-			fmt.Printf("Server > message[%d] (%s): %s\n", i, m.Role, textContent)
+		case *proto.HarnessResponse_End:
+			fmt.Printf("Server > [end] state=%s %s\n", payload.End.GetState(), payload.End.GetErrorMessage())
 		}
 	}
 
-	// Close send side to signal request completion
-	if err := stream.CloseSend(); err != nil {
-		return fmt.Errorf("Failed to close send side of stream: %v", err)
-	}
-
-	log.Println("Waiting for final stream EOF...")
-	_, err = stream.Recv()
-	if err != io.EOF {
-		return fmt.Errorf("Expected EOF from server, got: %v", err)
-	}
 	log.Println("Stream closed successfully by server.")
 	return nil
 }

internal/config2/config.go

@@ -85,7 +85,7 @@ type SubstrateHarnessConfig struct {
 // as a substrate actor; otherwise it runs locally.
 func (c AntigravityHarnessConfig) NewHarness(substrate bool, endpoint string) (harness.Harness, error) {
 	if substrate {
-		return newSubstrateHarness(endpoint, defaultNamespace, antigravityTemplate, defaultPort)
+		return newSubstrateHarness(c.ID, endpoint, defaultNamespace, antigravityTemplate, defaultPort)
 	}
 	address := c.Address
 	if address == "" {
@@ -101,12 +101,12 @@ func (c SubstrateHarnessConfig) NewHarness(endpoint string) (harness.Harness, er
 	if port == 0 {
 		port = defaultPort
 	}
-	return newSubstrateHarness(endpoint, c.Namespace, c.Template, port)
+	return newSubstrateHarness(c.ID, endpoint, c.Namespace, c.Template, port)
 }
 
 // newSubstrateHarness brings up a harness that is deployed as a substrate actor.
-func newSubstrateHarness(endpoint, namespace, template string, port int) (harness.Harness, error) {
-	sh, err := harness.NewSubstrateHarness(endpoint, namespace, template, port)
+func newSubstrateHarness(harnessID, endpoint, namespace, template string, port int) (harness.Harness, error) {
+	sh, err := harness.NewSubstrateHarness(harnessID, endpoint, namespace, template, port)
 	if err != nil {
 		return nil, err
 	}

internal/harness/antigravity.go

@@ -107,23 +107,30 @@ func (e *antigravityExecution) Run(ctx context.Context, handler Handler) error {
 	}
 	defer conn.Close()
 
-	// 2. Create AgentService client
-	client := proto.NewAgentServiceClient(conn)
+	// 2. Create HarnessService client.
+	client := proto.NewHarnessServiceClient(conn)
 
-	// 3. Build standard AgentRequest
-	req := &proto.AgentRequest{
+	// 3. Build standard HarnessRequest.
+	start := &proto.HarnessRequest{
 		ConversationId: e.conversationID,
-		ExecId:         e.id,
-		Start: &proto.AgentStart{
-			AgentId:  "antigravity",
-			Messages: inputs,
+		HarnessId:      "antigravity",
+		Type: &proto.HarnessRequest_Start{
+			Start: &proto.HarnessStart{
+				Messages: inputs,
+			},
 		},
 	}
 
 	// 4. Call Connect to start bidirectional streaming
-	stream, err := client.Connect(ctx, req)
+	stream, err := client.Connect(ctx)
 	if err != nil {
-		return fmt.Errorf("failed to call gRPC AgentService.Connect: %w", err)
+		return fmt.Errorf("failed to call gRPC HarnessService.Connect: %w", err)
+	}
+	if err := stream.Send(start); err != nil {
+		return fmt.Errorf("failed to send harness start: %w", err)
+	}
+	if err := stream.CloseSend(); err != nil {
+		return fmt.Errorf("failed to close stream send direction: %w", err)
 	}
 
 	// 5. Stream responses and trigger callbacks
@@ -137,14 +144,16 @@ func (e *antigravityExecution) Run(ctx context.Context, handler Handler) error {
 		}
 
 		switch payload := resp.Type.(type) {
-		case *proto.AgentResponse_Outputs:
+		case *proto.HarnessResponse_Outputs:
 			for _, outMsg := range payload.Outputs.Messages {
 				if err := handler.OnMessage(ctx, e.id, outMsg); err != nil {
 					return fmt.Errorf("failed to dispatch streamed output: %w", err)
 				}
 			}
-		case *proto.AgentResponse_End:
-			// Standard turn complete callback
+		case *proto.HarnessResponse_End:
+			if payload.End.GetState() == proto.State_STATE_FAILED {
+				return fmt.Errorf("harness failed: %s", payload.End.GetErrorMessage())
+			}
 			return handler.OnComplete(ctx, e.id)
 		}
 	}