From 6538055a9455b13f07568ad726a19d7b58249e0c Mon Sep 17 00:00:00 2001 From: Junjie Wang Date: Tue, 14 Jul 2026 20:05:55 -0700 Subject: [PATCH] Handle default harness correctly. --- cmd/ax/internal/cliutil/cliutil.go | 10 +- internal/controller/controller.go | 38 +++++-- internal/controller/controller_test.go | 151 ++++++++++++++++++++++++- internal/controller/registry.go | 25 ++-- internal/controller/registry_test.go | 26 ++++- 5 files changed, 221 insertions(+), 29 deletions(-) diff --git a/cmd/ax/internal/cliutil/cliutil.go b/cmd/ax/internal/cliutil/cliutil.go index f8d0ebb1..4cef29e6 100644 --- a/cmd/ax/internal/cliutil/cliutil.go +++ b/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) } } diff --git a/internal/controller/controller.go b/internal/controller/controller.go index 1537ce4b..a3775284 100644 --- a/internal/controller/controller.go +++ b/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,16 +223,18 @@ 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 { @@ -224,7 +242,7 @@ func (l *logger) ResumptionState(ctx context.Context) (proto.State, error) { } } } - return state, nil + return state, harnessID, nil } func (l *logger) LogInputs(ctx context.Context, inputs []*proto.Message, harnessConfig []byte) (int32, error) { diff --git a/internal/controller/controller_test.go b/internal/controller/controller_test.go index d70382f0..24ab7559 100644 --- a/internal/controller/controller_test.go +++ b/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") + } +} diff --git a/internal/controller/registry.go b/internal/controller/registry.go index 2a6ef482..7c1ab57b 100644 --- a/internal/controller/registry.go +++ b/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 diff --git a/internal/controller/registry_test.go b/internal/controller/registry_test.go index 5fda9373..c3eb7516 100644 --- a/internal/controller/registry_test.go +++ b/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") + } +}