Handle default harness correctly when resume a conversation.

cmd/ax/internal/cliutil/cliutil.go

@@ -138,14 +138,10 @@ func NewControllerFromConfig(ctx context.Context, cfg *Config) (*controller.Cont
 		}
 	}
 
-	// Register the configured default harness.
+	// Set the default harness.
 	if defaultHarnessID != "" {
-		h, err := reg.Harness(defaultHarnessID)
-		if err != nil {
-			return nil, fmt.Errorf("default harness %q not found", defaultHarnessID)
-		}
-		if err := reg.RegisterHarness("", h); err != nil {
-			return nil, fmt.Errorf("register default harness %q: %w", defaultHarnessID, err)
+		if err := reg.SetDefaultHarness(defaultHarnessID); err != nil {
+			return nil, fmt.Errorf("set default harness %q: %w", defaultHarnessID, err)
 		}
 	}
 

internal/controller/controller.go

@@ -73,15 +73,31 @@ func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler E
 	// TODO(jbd): Enable bringing a remote harness that implements HarnessService.
 	// TODO(anj): We need to consolidate agents and harness registration.
 	// Adding harness registration support temporarily.
-	h, err := d.registry.Harness(req.HarnessId)
+	l := newLogger(d.eventLog, req.ConversationId, req.HarnessId)
+	state, storedHarnessID, err := l.ResumptionState(ctx)
 	if err != nil {
-		return fmt.Errorf("failed to get harness %q: %w", req.HarnessId, err)
+		return fmt.Errorf("failed to check resumption state: %w", err)
 	}
 
-	l := newLogger(d.eventLog, req.ConversationId, req.HarnessId)
-	state, err := l.ResumptionState(ctx)
+	// On resume, use the conversation's recorded harness. Using a different harness
+	// for the same conversation is not allowed.
+	if req.HarnessId != "" && storedHarnessID != "" && req.HarnessId != storedHarnessID {
+		return fmt.Errorf("resumption not allowed: harness ID changed from %s to %s", storedHarnessID, req.HarnessId)
+	}
+	harnessID := req.HarnessId
+	// Use the conversations's stored harness if no harness is specified.
+	if harnessID == "" {
+		harnessID = storedHarnessID
+	}
+	// For new conversations, use the default harness if no harness is specified.
+	if harnessID == "" {
+		harnessID = d.registry.defaultHarness
+	}
+	l.harnessID = harnessID
+
+	h, err := d.registry.Harness(harnessID)
 	if err != nil {
-		return fmt.Errorf("failed to check resumption state: %w", err)
+		return fmt.Errorf("failed to get harness %q: %w", harnessID, err)
 	}
 
 	hhandler := &harnessHandler{
@@ -207,24 +223,26 @@ type logger struct {
 	harnessID      string
 }
 
-func (l *logger) ResumptionState(ctx context.Context) (proto.State, error) {
+// ResumptionState returns the conversation's current state and the harness it used.
+func (l *logger) ResumptionState(ctx context.Context) (proto.State, string, error) {
 	events, err := l.el.Events(ctx, l.conversationID)
 	if err != nil {
-		return proto.State_STATE_UNSPECIFIED, err
+		return proto.State_STATE_UNSPECIFIED, "", err
 	}
 
 	var state proto.State
+	var harnessID string
 	for _, ev := range events {
-		if ev.HarnessId != "" && ev.HarnessId != l.harnessID {
-			return proto.State_STATE_UNSPECIFIED, fmt.Errorf("resumption not allowed: harness ID changed from %s to %s", ev.HarnessId, l.harnessID)
+		if harnessID == "" && ev.HarnessId != "" {
+			harnessID = ev.HarnessId
 		}
 		if l.execID == "" || ev.ExecId == l.execID {
 			if ev.State != proto.State_STATE_UNSPECIFIED {
 				state = ev.State
 			}
 		}
 	}
-	return state, nil
+	return state, harnessID, nil
 }
 
 func (l *logger) LogInputs(ctx context.Context, inputs []*proto.Message, harnessConfig []byte) (int32, error) {

internal/controller/controller_test.go

@@ -17,6 +17,7 @@ package controller
 import (
 	"context"
 	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/google/ax/internal/controller/eventlog"
@@ -72,7 +73,10 @@ func TestController2_ExecHelloWorld(t *testing.T) {
 
 	log := &eventlogtest.MemoryEventLog{}
 	reg := NewRegistry()
-	if err := reg.RegisterHarness("", &fakeHarness{}); err != nil {
+	if err := reg.RegisterHarness("default-harness", &fakeHarness{}); err != nil {
+		t.Fatal(err)
+	}
+	if err := reg.SetDefaultHarness("default-harness"); err != nil {
 		t.Fatal(err)
 	}
 
@@ -514,3 +518,148 @@ func TestController2_ExecResumptionFlow(t *testing.T) {
 		}
 	})
 }
+
+// A conversation started on one harness must resume on that harness when the
+// harness id is empty, not fall back to the registry default.
+func TestExec_ResumeEmptyHarnessUsesStored(t *testing.T) {
+	ctx := context.Background()
+	cid := "resume-empty"
+	log := &eventlogtest.MemoryEventLog{}
+	reg := NewRegistry()
+
+	def := &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) {
+		return &testExecution{}, nil
+	}}
+	stored := &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) {
+		return &testExecution{}, nil
+	}}
+	if err := reg.RegisterHarness("harness-a", def); err != nil {
+		t.Fatal(err)
+	}
+	if err := reg.RegisterHarness("harness-b", stored); err != nil {
+		t.Fatal(err)
+	}
+	if err := reg.SetDefaultHarness("harness-a"); err != nil {
+		t.Fatal(err)
+	}
+
+	c, err := New(ctx, Config{
+		Registry:        reg,
+		EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil },
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+
+	noop := ExecHandler(func(*proto.ExecResponse) error { return nil })
+
+	// Turn 1: explicitly run the NON-default harness.
+	if err := c.Exec(ctx, &proto.ExecRequest{
+		ConversationId: cid,
+		HarnessId:      "harness-b",
+		Inputs:         []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "hi"}}}}},
+	}, noop); err != nil {
+		t.Fatalf("turn 1: %v", err)
+	}
+
+	// Turn 2: resume WITHOUT a harness id. Must reuse harness-b, not the default.
+	before := stored.startCalls
+	if err := c.Exec(ctx, &proto.ExecRequest{
+		ConversationId: cid,
+		Inputs:         []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "more"}}}}},
+	}, noop); err != nil {
+		t.Fatalf("turn 2 (resume, empty harness): %v", err)
+	}
+	if stored.startCalls <= before {
+		t.Errorf("stored harness not used on resume (startCalls stayed %d)", stored.startCalls)
+	}
+	if def.startCalls != 0 {
+		t.Errorf("default harness used on resume (startCalls=%d), want the stored harness", def.startCalls)
+	}
+}
+
+// Verifies that an explicit harness change during resume is rejected.
+func TestExec_ResumeExplicitDifferentHarnessRejected(t *testing.T) {
+	ctx := context.Background()
+	cid := "resume-mismatch"
+	log := &eventlogtest.MemoryEventLog{}
+	reg := NewRegistry()
+	if err := reg.RegisterHarness("harness-a", &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) {
+		return &testExecution{}, nil
+	}}); err != nil {
+		t.Fatal(err)
+	}
+	if err := reg.RegisterHarness("harness-b", &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) {
+		return &testExecution{}, nil
+	}}); err != nil {
+		t.Fatal(err)
+	}
+
+	c, err := New(ctx, Config{
+		Registry:        reg,
+		EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil },
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+
+	noop := ExecHandler(func(*proto.ExecResponse) error { return nil })
+
+	if err := c.Exec(ctx, &proto.ExecRequest{
+		ConversationId: cid,
+		HarnessId:      "harness-a",
+		Inputs:         []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "hi"}}}}},
+	}, noop); err != nil {
+		t.Fatalf("turn 1: %v", err)
+	}
+	err = c.Exec(ctx, &proto.ExecRequest{
+		ConversationId: cid,
+		HarnessId:      "harness-b",
+		Inputs:         []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "more"}}}}},
+	}, noop)
+	if err == nil || !strings.Contains(err.Error(), "harness ID changed from harness-a to harness-b") {
+		t.Fatalf("resume with a different harness: got %v, want error 'harness ID changed from harness-a to harness-b'", err)
+	}
+}
+
+// Verifies that a new conversation started without a harness id records the
+// default harness's canonical id, not "".
+func TestExec_NewConversationLogsCanonicalDefault(t *testing.T) {
+	ctx := context.Background()
+	cid := "canonical-log"
+	log := &eventlogtest.MemoryEventLog{}
+	reg := NewRegistry()
+	if err := reg.RegisterHarness("harness-a", &testHarness{startFunc: func(context.Context, string) (harness.Execution, error) {
+		return &testExecution{}, nil
+	}}); err != nil {
+		t.Fatal(err)
+	}
+	if err := reg.SetDefaultHarness("harness-a"); err != nil {
+		t.Fatal(err)
+	}
+
+	c, err := New(ctx, Config{
+		Registry:        reg,
+		EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil },
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+
+	if err := c.Exec(ctx, &proto.ExecRequest{
+		ConversationId: cid,
+		Inputs:         []*proto.Message{{Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "hi"}}}}},
+	}, ExecHandler(func(*proto.ExecResponse) error { return nil })); err != nil {
+		t.Fatalf("exec: %v", err)
+	}
+	_, stored, err := newLogger(log, cid, "").ResumptionState(ctx)
+	if err != nil {
+		t.Fatalf("ResumptionState: %v", err)
+	}
+	if stored != "harness-a" {
+		t.Errorf("logged harness id = %q, want canonical %q (not empty)", stored, "harness-a")
+	}
+}

internal/controller/registry.go

@@ -23,8 +23,9 @@ import (
 
 // Registry manages a collection of harnesses.
 type Registry struct {
-	mu        sync.RWMutex
-	harnesses map[string]harness.Harness
+	mu             sync.RWMutex
+	harnesses      map[string]harness.Harness
+	defaultHarness string
 }
 
 // NewRegistry creates a new harness registry.
@@ -34,13 +35,10 @@ func NewRegistry() *Registry {
 	}
 }
 
-// RegisterHarness registers a harness under the given id. An empty id registers
-// the harness as the default, used when a request specifies no agent id.
+// RegisterHarness registers a harness under the given id.
 func (r *Registry) RegisterHarness(id string, h harness.Harness) error {
-	if id != "" {
-		if err := validateID(id); err != nil {
-			return err
-		}
+	if err := validateID(id); err != nil {
+		return err
 	}
 
 	r.mu.Lock()
@@ -64,6 +62,17 @@ func (r *Registry) Harness(id string) (harness.Harness, error) {
 	return h, nil
 }
 
+// SetDefaultHarness marks a registered harness id as the default.
+func (r *Registry) SetDefaultHarness(id string) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if _, ok := r.harnesses[id]; !ok {
+		return fmt.Errorf("harness %q not found", id)
+	}
+	r.defaultHarness = id
+	return nil
+}
+
 // Close releases resources held by the registry.
 func (r *Registry) Close() error {
 	return nil

internal/controller/registry_test.go

@@ -45,9 +45,9 @@ func TestRegistry_RegisterHarness(t *testing.T) {
 		t.Error("expected error registering invalid id, got nil")
 	}
 
-	// Empty id (the default) bypasses validation and is allowed.
-	if err := r.RegisterHarness("", h); err != nil {
-		t.Fatalf("RegisterHarness(default): %v", err)
+	// Empty id is reserved for the default harness.
+	if err := r.RegisterHarness("", h); err == nil {
+		t.Error("expected error registering empty id, got nil")
 	}
 }
 
@@ -65,3 +65,23 @@ func TestRegistry_FindHarness(t *testing.T) {
 		t.Error("expected error looking up missing harness, got nil")
 	}
 }
+
+func TestRegistry_SetDefaultHarness(t *testing.T) {
+	r := NewRegistry()
+	if err := r.RegisterHarness("antigravity", &dummyHarness{}); err != nil {
+		t.Fatalf("RegisterHarness: %v", err)
+	}
+
+	// An unregistered id is rejected.
+	if err := r.SetDefaultHarness("missing"); err == nil {
+		t.Error("SetDefaultHarness(missing): expected error, got nil")
+	}
+
+	// A registered id becomes the default.
+	if err := r.SetDefaultHarness("antigravity"); err != nil {
+		t.Fatalf("SetDefaultHarness(antigravity): %v", err)
+	}
+	if r.defaultHarness != "antigravity" {
+		t.Errorf("defaultHarness = %q, want %q", r.defaultHarness, "antigravity")
+	}
+}