From 01dd9f41b28736b67d3ad12a2e9f20da2de3e8d0 Mon Sep 17 00:00:00 2001 From: PengFei Bai Date: Sun, 26 Jul 2026 10:16:28 +0800 Subject: [PATCH] agentfleet: warm actor lifecycle, terminal harness metadata, substrate aa1d14a7 Single squashed patch carrying every agentfleet-owned change on top of google/ax main. Rebuilt directly on upstream so `git diff origin/main` is exactly this patch and nothing else -- the previous long-lived merge branch had accumulated drift because upstream rewrote its history, which made 3-way merge bases unreliable. To follow upstream: branch from the new upstream main, cherry-pick this commit, resolve conflicts, done. Contents: - proto: `harness_metadata` on ConversationEvent, HarnessEnd and ExecResponse so a harness can attach an opaque terminal blob (token usage) to the event that ends a turn. Generated artifacts regenerated with upstream's toolchain (protoc 28.2, protoc-gen-go-grpc 1.5.1, grpcio-tools 1.62.2) so only ax.pb.go / ax_pb2.py move. - harness: `Drainer` optional capability; stream plumbing for terminal metadata (OnCompleteWithMetadata / OnFailWithMetadata). - controller: honor `CloseBeforeNextStart` so harnesses tracking per-conversation turn state close before the next turn starts, while everything else keeps upstream's deferred-close semantics; persist terminal metadata on both COMPLETED and FAILED paths. - registry: drain Drainer harnesses on Close so warm actors awaiting deferred idle suspension are suspended, not leaked, on process exit. - server: GracefulStop drains gRPC before closing the controller so no turn re-arms the idle timer after the drain. - substrate harness: warm-then-suspend idle policy, warm actor revalidation with cold-resume fallback, shutdown drain. - go.mod: substrate pin -> aa1d14a7 (imagecache Phase 1 plus the 07-24 fixes); internal/ate and harnesstest adapted to name-based actor addressing. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_017TK34TZokEcfPq5JdbUVWo --- go.mod | 8 +- go.sum | 24 +- internal/ate/client.go | 18 +- internal/controller/controller.go | 112 ++++- internal/controller/controller_test.go | 376 +++++++++++++- internal/controller/eventlog/sql_test.go | 9 +- internal/controller/registry.go | 18 +- internal/controller/registry_test.go | 32 ++ internal/harness/harness.go | 10 + internal/harness/harnesstest/harnesstest.go | 54 +- internal/harness/stream.go | 26 + internal/harness/stream_test.go | 157 ++++++ internal/harness/substrate/substrate.go | 385 +++++++++++++- internal/harness/substrate/substrate_test.go | 500 +++++++++++++++++++ internal/server/server.go | 14 +- proto/ax.pb.go | 62 ++- proto/ax.proto | 6 + python/proto/ax_pb2.py | 68 +-- 18 files changed, 1748 insertions(+), 131 deletions(-) create mode 100644 internal/harness/stream_test.go diff --git a/go.mod b/go.mod index 73c022fa..66cfeac8 100644 --- a/go.mod +++ b/go.mod @@ -8,7 +8,7 @@ require ( charm.land/huh/v2 v2.0.3 charm.land/lipgloss/v2 v2.0.3 cloud.google.com/go/compute/metadata v0.9.0 - github.com/agent-substrate/substrate v0.0.0-20260706222328-3cb7433bd8a8 + github.com/agent-substrate/substrate v0.0.0-20260725015935-aa1d14a7b33b github.com/google/uuid v1.6.0 github.com/jackc/pgx/v5 v5.10.0 github.com/spf13/cobra v1.10.2 @@ -62,9 +62,9 @@ require ( go.opentelemetry.io/otel/metric v1.44.0 // indirect go.opentelemetry.io/proto/otlp v1.10.0 // indirect golang.org/x/net v0.55.0 // indirect - golang.org/x/sync v0.20.0 // indirect - golang.org/x/sys v0.45.0 // indirect - golang.org/x/text v0.37.0 // indirect + golang.org/x/sync v0.21.0 // indirect + golang.org/x/sys v0.46.0 // indirect + golang.org/x/text v0.39.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect modernc.org/libc v1.70.0 // indirect diff --git a/go.sum b/go.sum index a5bc044b..7b506cc3 100644 --- a/go.sum +++ b/go.sum @@ -10,8 +10,8 @@ cloud.google.com/go/compute/metadata v0.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdB cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10= github.com/MakeNowJust/heredoc v1.0.0 h1:cXCdzVdstXyiTqTvfqk9SDHpKNjxuom+DOlyEeQ4pzQ= github.com/MakeNowJust/heredoc v1.0.0/go.mod h1:mG5amYoWBHf8vpLOuehzbGGw0EHxpZZ6lCpQ4fNJ8LE= -github.com/agent-substrate/substrate v0.0.0-20260706222328-3cb7433bd8a8 h1:tUeBjLs9TJgIWfML2dlZ3UOFG3dacbwrIYkas2ctBgY= -github.com/agent-substrate/substrate v0.0.0-20260706222328-3cb7433bd8a8/go.mod h1:2cvSnnHPwZRAhkrC2LH1bQd1tQBfL3p+zU6pZiHBUkA= +github.com/agent-substrate/substrate v0.0.0-20260725015935-aa1d14a7b33b h1:XJsiwQcqxzMQT4eBhdO+Bg4cU0HTRTM9EXFwkmWUc+Q= +github.com/agent-substrate/substrate v0.0.0-20260725015935-aa1d14a7b33b/go.mod h1:FtVc9AA++A1aEyIHUQdjCesf6dP4sxIxJc+Xc+p6Mok= 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= @@ -157,21 +157,21 @@ go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY= golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70= -golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM= -golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU= +golang.org/x/mod v0.37.0 h1:vF1DjpVEshcIqoEaauuHebaLk1O1forxjxBaVn884JQ= +golang.org/x/mod v0.37.0/go.mod h1:m8S8VeM9r4dzDwjrKO0a1sZP3YjeMamRRlD+fmR2Q/0= golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= +golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= -golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= -golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= -golang.org/x/tools v0.44.0 h1:UP4ajHPIcuMjT1GqzDWRlalUEoY+uzoZKnhOjbIPD2c= -golang.org/x/tools v0.44.0/go.mod h1:KA0AfVErSdxRZIsOVipbv3rQhVXTnlU6UhKxHd1seDI= +golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= +golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/text v0.39.0 h1:UbZz4pLOvn600D6Oh6GGEI6VAmndrEBLv8/6BEXzyus= +golang.org/x/text v0.39.0/go.mod h1:3UwRclnC2g0TU9x8PZiyfOajCd1zaUNHF9cvqcQZ+ZM= +golang.org/x/tools v0.47.0 h1:7Kn5x/d1svx/PzryTsqeoZN4TZwqeH5pGWjefhLi/1Q= +golang.org/x/tools v0.47.0/go.mod h1:dFHnyTvFWY212G+h7ZY4Vsp/K3U4/7W9TyVaAul8uCA= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa h1:Kjn0N0tCrDgiAFW+lGO4JZ3ck44CehvJQMAwj9QF0G8= diff --git a/internal/ate/client.go b/internal/ate/client.go index 81a5c061..d9697480 100644 --- a/internal/ate/client.go +++ b/internal/ate/client.go @@ -58,16 +58,20 @@ func NewClient(ns, template, target string, opts ...grpc.DialOption) (*Client, e } // CreateActor creates a new actor. -func (c *Client) CreateActor(ctx context.Context, id string) (*ateapipb.CreateActorResponse, error) { +func (c *Client) CreateActor(ctx context.Context, id string) (*ateapipb.Actor, error) { client := ateapipb.NewControlClient(c.conn) // TODO(wjjclaud): Configure atespace in manifests instead of reusing the namespace. - if _, err := client.CreateAtespace(ctx, &ateapipb.CreateAtespaceRequest{Name: c.namespace}); err != nil && status.Code(err) != codes.AlreadyExists { + if _, err := client.CreateAtespace(ctx, &ateapipb.CreateAtespaceRequest{ + Atespace: &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: c.namespace}}, + }); err != nil && status.Code(err) != codes.AlreadyExists { return nil, fmt.Errorf("error when calling Control.CreateAtespace: %w", err) } resp, err := client.CreateActor(ctx, &ateapipb.CreateActorRequest{ - ActorRef: &ateapipb.ActorRef{Atespace: c.namespace, Name: id}, - ActorTemplateNamespace: c.namespace, - ActorTemplateName: c.template, + Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: c.namespace, Name: id}, + ActorTemplateNamespace: c.namespace, + ActorTemplateName: c.template, + }, }) if err != nil { return nil, fmt.Errorf("error when calling Control.CreateActor: %w", err) @@ -80,7 +84,7 @@ func (c *Client) CreateActor(ctx context.Context, id string) (*ateapipb.CreateAc func (c *Client) ResumeActor(ctx context.Context, id string) (*ateapipb.ResumeActorResponse, error) { client := ateapipb.NewControlClient(c.conn) resp, err := client.ResumeActor(ctx, &ateapipb.ResumeActorRequest{ - ActorRef: &ateapipb.ActorRef{Atespace: c.namespace, Name: id}, + Actor: &ateapipb.ObjectRef{Atespace: c.namespace, Name: id}, }) if err != nil { return nil, fmt.Errorf("error when calling Control.ResumeActor: %w", err) @@ -92,7 +96,7 @@ func (c *Client) ResumeActor(ctx context.Context, id string) (*ateapipb.ResumeAc func (c *Client) SuspendActor(ctx context.Context, id string) (*ateapipb.SuspendActorResponse, error) { client := ateapipb.NewControlClient(c.conn) resp, err := client.SuspendActor(ctx, &ateapipb.SuspendActorRequest{ - ActorRef: &ateapipb.ActorRef{Atespace: c.namespace, Name: id}, + Actor: &ateapipb.ObjectRef{Atespace: c.namespace, Name: id}, }) if err != nil { return nil, fmt.Errorf("error when calling Control.SuspendActor: %w", err) diff --git a/internal/controller/controller.go b/internal/controller/controller.go index ac7a129f..da39ed39 100644 --- a/internal/controller/controller.go +++ b/internal/controller/controller.go @@ -22,6 +22,7 @@ import ( "log/slog" "github.com/google/ax/internal/controller/eventlog" + "github.com/google/ax/internal/harness" "github.com/google/ax/proto" "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/types/known/structpb" @@ -113,10 +114,18 @@ func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler E if err != nil { return fmt.Errorf("failed to start harness session: %w", err) } - defer exec.Close(ctx) - - if err := exec.Run(ctx, hhandler); err != nil { - return fmt.Errorf("harness execution failed: %w", err) + var runErr error + if closeBeforeNextStart(exec) { + runErr = func() error { + defer exec.Close(ctx) + return exec.Run(ctx, hhandler) + }() + } else { + defer exec.Close(ctx) + runErr = exec.Run(ctx, hhandler) + } + if runErr != nil { + return fmt.Errorf("harness execution failed: %w", runErr) } } @@ -144,6 +153,19 @@ func (d *Controller) Exec(ctx context.Context, req *proto.ExecRequest, handler E return nil } +// eagerCloseExecution is an optional Execution capability. Executions whose +// harness tracks per-conversation turn state (substrate warm mode) must be +// closed before the controller starts another execution for the same +// conversation; everything else keeps upstream's deferred-close semantics. +type eagerCloseExecution interface { + CloseBeforeNextStart() bool +} + +func closeBeforeNextStart(exec harness.Execution) bool { + ec, ok := exec.(eagerCloseExecution) + return ok && ec.CloseBeforeNextStart() +} + type harnessHandler struct { logger *logger execHandler ExecHandler @@ -152,7 +174,7 @@ type harnessHandler struct { func (a *harnessHandler) OnMessage(ctx context.Context, execID string, msg *proto.Message) error { // Log every response received from the harness // TODO(anj): The harness should send the full input sent to get this particular response. - step, err := a.logger.LogOutputs(ctx, []*proto.Message{msg}, proto.State_STATE_PENDING) + step, err := a.logger.LogOutputs(ctx, []*proto.Message{msg}, proto.State_STATE_PENDING, nil, "") if err != nil { slog.WarnContext(ctx, "Failed to log streamed message to event log", slog.String("conversation_id", a.logger.conversationID), @@ -170,14 +192,74 @@ func (a *harnessHandler) OnMessage(ctx context.Context, execID string, msg *prot } func (a *harnessHandler) OnComplete(ctx context.Context, execID string) error { + return a.complete(ctx, execID, nil) +} + +// OnCompleteWithMetadata retains opaque metadata on the existing terminal +// event without expanding harness.Handler. +func (a *harnessHandler) OnCompleteWithMetadata(ctx context.Context, execID string, metadata []byte) error { + return a.complete(ctx, execID, metadata) +} + +// OnFailWithMetadata persists a terminal FAILED event that still carries the +// harness's opaque metadata (e.g. token usage collected before the failure), +// then returns the original cause unchanged so the caller's error path is +// unaffected by whether the metadata could be persisted. +func (a *harnessHandler) OnFailWithMetadata(ctx context.Context, execID string, metadata []byte, cause error) error { + // Metadata-bearing terminal events are stamped with the stream's execID, + // mirroring complete()'s convention for the COMPLETED path. + terminalExecID := "" + if len(metadata) > 0 { + terminalExecID = execID + } + seq, err := a.logger.LogOutputs(ctx, nil, proto.State_STATE_FAILED, metadata, terminalExecID) + if err != nil { + slog.WarnContext(ctx, "Failed to log FAILED terminal metadata", + slog.String("conversation_id", a.logger.conversationID), + slog.Any("error", err), + ) + return cause + } + if a.execHandler != nil { + if err := a.execHandler(&proto.ExecResponse{ + Step: seq, + HarnessMetadata: metadata, + }); err != nil { + slog.WarnContext(ctx, "Failed to stream FAILED terminal metadata to exec handler", + slog.String("conversation_id", a.logger.conversationID), + slog.Any("error", err), + ) + } + } + return cause +} + +func (a *harnessHandler) complete(ctx context.Context, execID string, metadata []byte) error { + // Metadata-bearing terminal events are stamped with the stream's execID; + // the legacy no-metadata path keeps the logger's (empty) execID unchanged. + terminalExecID := "" + if len(metadata) > 0 { + terminalExecID = execID + } // Mark the execution turn as completed in the conversation log - if _, err := a.logger.LogOutputs(ctx, nil, proto.State_STATE_COMPLETED); err != nil { + seq, err := a.logger.LogOutputs(ctx, nil, proto.State_STATE_COMPLETED, metadata, terminalExecID) + if err != nil { slog.WarnContext(ctx, "Failed to log completion event to event log", slog.String("conversation_id", a.logger.conversationID), slog.Any("error", err), ) + if len(metadata) > 0 { + return fmt.Errorf("failed to persist terminal harness metadata: %w", err) + } + return nil } - return nil + if len(metadata) == 0 || a.execHandler == nil { + return nil + } + return a.execHandler(&proto.ExecResponse{ + Step: seq, + HarnessMetadata: metadata, + }) } // Delete deletes all events for a specific conversation ID. @@ -269,12 +351,18 @@ func (l *logger) LogInputs(ctx context.Context, inputs []*proto.Message, harness return l.el.Append(ctx, ev) } -func (l *logger) LogOutputs(ctx context.Context, outputs []*proto.Message, state proto.State) (int32, error) { +// LogOutputs appends an output event. A non-empty execID overrides the +// logger's own (which is never populated today) on the appended event. +func (l *logger) LogOutputs(ctx context.Context, outputs []*proto.Message, state proto.State, harnessMetadata []byte, execID string) (int32, error) { + if execID == "" { + execID = l.execID + } ev := &proto.ConversationEvent{ - ConversationId: l.conversationID, - ExecId: l.execID, - Messages: outputs, - State: state, + ConversationId: l.conversationID, + ExecId: execID, + Messages: outputs, + State: state, + HarnessMetadata: harnessMetadata, } return l.el.Append(ctx, ev) } diff --git a/internal/controller/controller_test.go b/internal/controller/controller_test.go index 24ab7559..99503932 100644 --- a/internal/controller/controller_test.go +++ b/internal/controller/controller_test.go @@ -15,8 +15,10 @@ package controller import ( + "bytes" "context" "fmt" + "slices" "strings" "testing" @@ -26,6 +28,14 @@ import ( "github.com/google/ax/proto" ) +type metadataCompleter interface { + OnCompleteWithMetadata(ctx context.Context, execID string, metadata []byte) error +} + +type failMetadataCompleter interface { + OnFailWithMetadata(ctx context.Context, execID string, metadata []byte, cause error) error +} + type fakeHarness struct{} func (f *fakeHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { @@ -171,6 +181,193 @@ func TestController2_ExecHelloWorld(t *testing.T) { } +func TestController2_ExecPersistsAndStreamsTerminalHarnessMetadata(t *testing.T) { + ctx := context.Background() + const conversationID = "metadata-conversation" + wantMetadata := []byte("agentfleet-metadata-fixture") + + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + h := &testHarness{ + startFunc: func(context.Context, string) (harness.Execution, error) { + return &testExecution{ + id: "metadata-exec", + runFunc: func(ctx context.Context, execID string, handler harness.Handler) error { + completer, ok := handler.(metadataCompleter) + if !ok { + return fmt.Errorf("handler does not support terminal metadata") + } + return completer.OnCompleteWithMetadata(ctx, execID, wantMetadata) + }, + }, nil + }, + } + if err := reg.RegisterHarness("metadata", h); 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() + + var responses []*proto.ExecResponse + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: conversationID, + HarnessId: "metadata", + Inputs: []*proto.Message{{Role: "user"}}, + }, func(resp *proto.ExecResponse) error { + responses = append(responses, resp) + return nil + }) + if err != nil { + t.Fatalf("Exec: %v", err) + } + if len(responses) != 1 { + t.Fatalf("responses = %d, want 1 terminal metadata response", len(responses)) + } + if !bytes.Equal(responses[0].GetHarnessMetadata(), wantMetadata) { + t.Fatalf("response metadata = %q, want %q", responses[0].GetHarnessMetadata(), wantMetadata) + } + if len(responses[0].GetOutputs()) != 0 { + t.Fatalf("terminal metadata response has %d outputs, want 0", len(responses[0].GetOutputs())) + } + + events, err := log.Events(ctx, conversationID) + if err != nil { + t.Fatalf("Events: %v", err) + } + if len(events) != 2 { + t.Fatalf("events = %d, want input plus terminal event", len(events)) + } + terminal := events[1] + if terminal.GetState() != proto.State_STATE_COMPLETED { + t.Fatalf("terminal state = %v, want COMPLETED", terminal.GetState()) + } + if terminal.GetConversationId() != conversationID || terminal.GetExecId() != "metadata-exec" { + t.Fatalf("terminal identity = (%q, %q), want (%q, %q)", terminal.GetConversationId(), terminal.GetExecId(), conversationID, "metadata-exec") + } + if !bytes.Equal(terminal.GetHarnessMetadata(), wantMetadata) { + t.Fatalf("event metadata = %q, want %q", terminal.GetHarnessMetadata(), wantMetadata) + } + if responses[0].GetStep() != terminal.GetStep() { + t.Fatalf("response seq = %d, terminal event seq = %d", responses[0].GetStep(), terminal.GetStep()) + } +} + +func TestController2_ExecPersistsAndStreamsFailedTerminalHarnessMetadata(t *testing.T) { + ctx := context.Background() + const conversationID = "failed-metadata-conversation" + wantMetadata := []byte("agentfleet-failed-metadata-fixture") + wantCause := fmt.Errorf("harness failed: [13] boom") + + log := &eventlogtest.MemoryEventLog{} + reg := NewRegistry() + h := &testHarness{ + startFunc: func(context.Context, string) (harness.Execution, error) { + return &testExecution{ + id: "failed-metadata-exec", + runFunc: func(ctx context.Context, execID string, handler harness.Handler) error { + completer, ok := handler.(failMetadataCompleter) + if !ok { + return fmt.Errorf("handler does not support FAILED terminal metadata") + } + return completer.OnFailWithMetadata(ctx, execID, wantMetadata, wantCause) + }, + }, nil + }, + } + if err := reg.RegisterHarness("failed-metadata", h); 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() + + var responses []*proto.ExecResponse + err = c.Exec(ctx, &proto.ExecRequest{ + ConversationId: conversationID, + HarnessId: "failed-metadata", + Inputs: []*proto.Message{{Role: "user"}}, + }, func(resp *proto.ExecResponse) error { + responses = append(responses, resp) + return nil + }) + if err == nil || !strings.Contains(err.Error(), wantCause.Error()) { + t.Fatalf("Exec error = %v, want it to wrap %v", err, wantCause) + } + + if len(responses) != 1 { + t.Fatalf("responses = %d, want 1 terminal metadata response", len(responses)) + } + if !bytes.Equal(responses[0].GetHarnessMetadata(), wantMetadata) { + t.Fatalf("response metadata = %q, want %q", responses[0].GetHarnessMetadata(), wantMetadata) + } + if len(responses[0].GetOutputs()) != 0 { + t.Fatalf("terminal metadata response has %d outputs, want 0", len(responses[0].GetOutputs())) + } + + events, err := log.Events(ctx, conversationID) + if err != nil { + t.Fatalf("Events: %v", err) + } + if len(events) != 2 { + t.Fatalf("events = %d, want input plus terminal event", len(events)) + } + terminal := events[1] + if terminal.GetState() != proto.State_STATE_FAILED { + t.Fatalf("terminal state = %v, want FAILED", terminal.GetState()) + } + if terminal.GetConversationId() != conversationID || terminal.GetExecId() != "failed-metadata-exec" { + t.Fatalf("terminal identity = (%q, %q), want (%q, %q)", terminal.GetConversationId(), terminal.GetExecId(), conversationID, "failed-metadata-exec") + } + if !bytes.Equal(terminal.GetHarnessMetadata(), wantMetadata) { + t.Fatalf("event metadata = %q, want %q", terminal.GetHarnessMetadata(), wantMetadata) + } + if responses[0].GetStep() != terminal.GetStep() { + t.Fatalf("response seq = %d, terminal event seq = %d", responses[0].GetStep(), terminal.GetStep()) + } +} + +func TestHarnessHandlerWithoutMetadataPreservesLegacyEmptyExecID(t *testing.T) { + ctx := context.Background() + log := &eventlogtest.MemoryEventLog{} + handler := &harnessHandler{logger: newLogger(log, "legacy-conversation", "legacy-harness")} + + if err := handler.OnMessage(ctx, "runtime-exec", &proto.Message{Role: "assistant"}); err != nil { + t.Fatalf("OnMessage: %v", err) + } + if err := handler.OnComplete(ctx, "runtime-exec"); err != nil { + t.Fatalf("OnComplete: %v", err) + } + + events, err := log.Events(ctx, "legacy-conversation") + if err != nil { + t.Fatalf("Events: %v", err) + } + if len(events) != 2 { + t.Fatalf("events = %d, want streamed message plus completion", len(events)) + } + for i, event := range events { + if event.GetExecId() != "" { + t.Errorf("event[%d].exec_id = %q, want legacy empty value", i, event.GetExecId()) + } + if len(event.GetHarnessMetadata()) != 0 { + t.Errorf("event[%d].harness_metadata = %q, want empty", i, event.GetHarnessMetadata()) + } + } +} + func TestController2_ExecWithAgentID(t *testing.T) { ctx := context.Background() cid := "test-conversation-id" @@ -313,6 +510,11 @@ func (c *testExecution) Close(ctx context.Context) error { return nil } +// eagerTestExecution opts a testExecution in to the eager-close capability. +type eagerTestExecution struct{ *testExecution } + +func (e *eagerTestExecution) CloseBeforeNextStart() bool { return true } + func TestController2_ExecResumptionFlow(t *testing.T) { // Subtest 1: New Execution with Inputs t.Run("NewExecutionWithInputs", func(t *testing.T) { @@ -456,9 +658,14 @@ func TestController2_ExecResumptionFlow(t *testing.T) { reg := NewRegistry() + // The fake models substrate warm mode: it declares the eager-close + // capability and rejects a new Start while a turn is still open. var execs []*testExecution h := &testHarness{ startFunc: func(ctx context.Context, conversationID string) (harness.Execution, error) { + if len(execs) > 0 && execs[len(execs)-1].closeCalls == 0 { + return nil, fmt.Errorf("previous execution was not closed before the next start") + } exec := &testExecution{ id: fmt.Sprintf("exec-%d", len(execs)+1), runFunc: func(ctx context.Context, execID string, handler harness.Handler) error { @@ -466,7 +673,7 @@ func TestController2_ExecResumptionFlow(t *testing.T) { }, } execs = append(execs, exec) - return exec, nil + return &eagerTestExecution{exec}, nil }, } if err := reg.RegisterHarness("test-agent", h); err != nil { @@ -663,3 +870,170 @@ func TestExec_NewConversationLogsCanonicalDefault(t *testing.T) { t.Errorf("logged harness id = %q, want canonical %q (not empty)", stored, "harness-a") } } + +// journalExecution records lifecycle events into a shared journal so tests can +// assert Close ordering across the two executions of a pending-resume Exec. +type journalExecution struct { + name string + journal *[]string + runFunc func(ctx context.Context, execID string, handler harness.Handler) error +} + +func (e *journalExecution) ID() string { return e.name } + +func (e *journalExecution) Queue(ctx context.Context, msg ...*proto.Message) error { return nil } + +func (e *journalExecution) Run(ctx context.Context, handler harness.Handler) error { + *e.journal = append(*e.journal, "run:"+e.name) + if e.runFunc != nil { + return e.runFunc(ctx, e.name, handler) + } + return handler.OnComplete(ctx, e.name) +} + +func (e *journalExecution) Close(ctx context.Context) error { + *e.journal = append(*e.journal, "close:"+e.name) + return nil +} + +// eagerJournalExecution additionally opts in to the eager-close capability. +type eagerJournalExecution struct{ journalExecution } + +func (e *eagerJournalExecution) CloseBeforeNextStart() bool { return true } + +func seedPendingConversation(t *testing.T, log *eventlogtest.MemoryEventLog, cid string) { + t.Helper() + _, err := log.Append(context.Background(), &proto.ConversationEvent{ + ConversationId: cid, + HarnessId: "test-agent", + State: proto.State_STATE_PENDING, + Messages: []*proto.Message{ + {Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "Initial"}}}}, + }, + }) + if err != nil { + t.Fatal(err) + } +} + +func newJournalController(t *testing.T, log *eventlogtest.MemoryEventLog, start func(ctx context.Context, conversationID string) (harness.Execution, error)) *Controller { + t.Helper() + reg := NewRegistry() + if err := reg.RegisterHarness("test-agent", &testHarness{startFunc: start}); err != nil { + t.Fatal(err) + } + c, err := New(context.Background(), Config{ + Registry: reg, + EventLogBuilder: func() (eventlog.EventLog, error) { return log, nil }, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = c.Close() }) + return c +} + +// Without the eager-close capability the controller must keep upstream's +// deferred-close semantics: in the default immediate-suspend substrate mode an +// inline Close would synchronously suspend and immediately re-resume the actor +// in the middle of a single Exec call. +func TestExec_PendingResumeKeepsDeferredCloseWithoutEagerCapability(t *testing.T) { + ctx := context.Background() + cid := "pending-deferred-close" + log := &eventlogtest.MemoryEventLog{} + seedPendingConversation(t, log, cid) + + var journal []string + starts := 0 + c := newJournalController(t, log, func(ctx context.Context, conversationID string) (harness.Execution, error) { + starts++ + name := fmt.Sprintf("exec-%d", starts) + journal = append(journal, "start:"+name) + return &journalExecution{name: name, journal: &journal}, nil + }) + + err := c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "test-agent", + Inputs: []*proto.Message{ + {Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "More"}}}}, + }, + }, func(*proto.ExecResponse) error { return nil }) + if err != nil { + t.Fatal(err) + } + + want := []string{"start:exec-1", "run:exec-1", "start:exec-2", "run:exec-2", "close:exec-2", "close:exec-1"} + if !slices.Equal(journal, want) { + t.Fatalf("journal = %v,\nwant %v (upstream deferred-close semantics)", journal, want) + } +} + +// Executions that opt in via CloseBeforeNextStart must be closed before the +// controller starts the next execution for the same conversation (warm-mode +// turn-slot bookkeeping). +func TestExec_PendingResumeClosesEagerlyWithCapability(t *testing.T) { + ctx := context.Background() + cid := "pending-eager-close" + log := &eventlogtest.MemoryEventLog{} + seedPendingConversation(t, log, cid) + + var journal []string + starts := 0 + c := newJournalController(t, log, func(ctx context.Context, conversationID string) (harness.Execution, error) { + starts++ + name := fmt.Sprintf("exec-%d", starts) + journal = append(journal, "start:"+name) + return &eagerJournalExecution{journalExecution{name: name, journal: &journal}}, nil + }) + + err := c.Exec(ctx, &proto.ExecRequest{ + ConversationId: cid, + HarnessId: "test-agent", + Inputs: []*proto.Message{ + {Role: "user", Content: &proto.Content{Type: &proto.Content_Text{Text: &proto.TextContent{Text: "More"}}}}, + }, + }, func(*proto.ExecResponse) error { return nil }) + if err != nil { + t.Fatal(err) + } + + want := []string{"start:exec-1", "run:exec-1", "close:exec-1", "start:exec-2", "run:exec-2", "close:exec-2"} + if !slices.Equal(journal, want) { + t.Fatalf("journal = %v,\nwant %v (eager close before the next Start)", journal, want) + } +} + +// A panic out of Run must not leak the execution: in warm mode a skipped Close +// leaves the conversation's turn slot taken forever. +func TestExec_PendingResumePanicStillClosesExecution(t *testing.T) { + ctx := context.Background() + cid := "pending-panic-close" + log := &eventlogtest.MemoryEventLog{} + seedPendingConversation(t, log, cid) + + var journal []string + c := newJournalController(t, log, func(ctx context.Context, conversationID string) (harness.Execution, error) { + return &eagerJournalExecution{journalExecution{ + name: "exec-1", + journal: &journal, + runFunc: func(ctx context.Context, execID string, handler harness.Handler) error { + panic("boom") + }, + }}, nil + }) + + func() { + defer func() { + if recover() == nil { + t.Fatal("expected the Run panic to propagate") + } + }() + _ = c.Exec(ctx, &proto.ExecRequest{ConversationId: cid, HarnessId: "test-agent"}, + func(*proto.ExecResponse) error { return nil }) + }() + + if !slices.Contains(journal, "close:exec-1") { + t.Fatalf("journal = %v, want Close to run despite the panic", journal) + } +} diff --git a/internal/controller/eventlog/sql_test.go b/internal/controller/eventlog/sql_test.go index 4f07b07f..4dacefda 100644 --- a/internal/controller/eventlog/sql_test.go +++ b/internal/controller/eventlog/sql_test.go @@ -15,6 +15,7 @@ package eventlog import ( + "bytes" "context" "os" "path/filepath" @@ -44,7 +45,8 @@ func testEventLog(t *testing.T, newLog func(t *testing.T) EventLog) { // 1. Conversation log. cev1 := &proto.ConversationEvent{ConversationId: conv, Step: 1, ExecId: task1} - cev2 := &proto.ConversationEvent{ConversationId: conv, Step: 2, ExecId: task2} + metadata := []byte("agentfleet-metadata-fixture") + cev2 := &proto.ConversationEvent{ConversationId: conv, Step: 2, ExecId: task2, HarnessMetadata: metadata} if _, err := log.Append(ctx, cev1); err != nil { t.Fatalf("failed to append cev1: %v", err) } @@ -65,7 +67,9 @@ func testEventLog(t *testing.T, newLog func(t *testing.T) EventLog) { if cEvents[0].ExecId != task1 || cEvents[1].ExecId != task2 { t.Errorf("conversation events mismatch: %q, %q", cEvents[0].ExecId, cEvents[1].ExecId) } - + if !bytes.Equal(cEvents[1].GetHarnessMetadata(), metadata) { + t.Errorf("metadata = %q, want %q", cEvents[1].GetHarnessMetadata(), metadata) + } }) @@ -104,7 +108,6 @@ func testEventLog(t *testing.T, newLog func(t *testing.T) EventLog) { t.Fatalf("append: %v", err) } - if err := log.DeleteAll(ctx, conv1); err != nil { t.Fatalf("failed to delete events: %v", err) } diff --git a/internal/controller/registry.go b/internal/controller/registry.go index 7c1ab57b..5088760d 100644 --- a/internal/controller/registry.go +++ b/internal/controller/registry.go @@ -15,6 +15,7 @@ package controller import ( + "context" "fmt" "sync" @@ -73,7 +74,22 @@ func (r *Registry) SetDefaultHarness(id string) error { return nil } -// Close releases resources held by the registry. +// Close releases resources held by the registry. It drains every registered +// harness that implements the optional harness.Drainer capability so warm +// actors awaiting deferred idle suspension are suspended rather than leaked on +// process exit. func (r *Registry) Close() error { + r.mu.RLock() + drainers := make([]harness.Drainer, 0, len(r.harnesses)) + for _, h := range r.harnesses { + if d, ok := h.(harness.Drainer); ok { + drainers = append(drainers, d) + } + } + r.mu.RUnlock() + + for _, d := range drainers { + d.Shutdown(context.Background()) + } return nil } diff --git a/internal/controller/registry_test.go b/internal/controller/registry_test.go index c3eb7516..8e952483 100644 --- a/internal/controller/registry_test.go +++ b/internal/controller/registry_test.go @@ -27,6 +27,38 @@ func (d *dummyHarness) Start(ctx context.Context, conversationID string, harness return nil, nil } +// drainableHarness implements the optional harness.Drainer capability so the +// registry test can assert that Close drains it. +type drainableHarness struct { + dummyHarness + shutdownCalls int +} + +func (d *drainableHarness) Shutdown(ctx context.Context) { + d.shutdownCalls++ +} + +// Close must invoke Shutdown on every registered harness that implements the +// Drainer capability so warm actors are suspended before the process exits. +// Harnesses without the capability are skipped without error. +func TestRegistry_CloseDrainsDrainerHarnesses(t *testing.T) { + r := NewRegistry() + drainable := &drainableHarness{} + if err := r.RegisterHarness("warm", drainable); err != nil { + t.Fatalf("RegisterHarness(warm): %v", err) + } + if err := r.RegisterHarness("plain", &dummyHarness{}); err != nil { + t.Fatalf("RegisterHarness(plain): %v", err) + } + + if err := r.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + if drainable.shutdownCalls != 1 { + t.Fatalf("Shutdown calls = %d, want 1", drainable.shutdownCalls) + } +} + func TestRegistry_RegisterHarness(t *testing.T) { r := NewRegistry() h := &dummyHarness{} diff --git a/internal/harness/harness.go b/internal/harness/harness.go index 716f0165..07031df6 100644 --- a/internal/harness/harness.go +++ b/internal/harness/harness.go @@ -46,6 +46,16 @@ type Harness interface { Start(ctx context.Context, conversationID string, harnessConfig []byte) (Execution, error) } +// Drainer is an optional Harness capability. Shutdown releases per-conversation +// resources that outlive individual turns -- for example warm actors kept +// running between turns and awaiting a deferred idle suspension -- so they are +// not leaked when the process exits. Harnesses that hold no such deferred state +// need not implement it. Callers should invoke Shutdown after in-flight turns +// have drained so no turn re-arms deferred state after the drain. +type Drainer interface { + Shutdown(ctx context.Context) +} + // Execution represents an active interactive session with an agent or planner. type Execution interface { // Run executes the session and streams events to the provided Handler. diff --git a/internal/harness/harnesstest/harnesstest.go b/internal/harness/harnesstest/harnesstest.go index 6328b063..3472c1e8 100644 --- a/internal/harness/harnesstest/harnesstest.go +++ b/internal/harness/harnesstest/harnesstest.go @@ -36,7 +36,7 @@ import ( // mockControlServer is an in-process ateapipb.ControlServer that records the // actor lifecycle calls SubstrateHarness makes and lets tests steer the -// CreateActor/ResumeActor responses. Only the three RPCs SubstrateHarness uses +// CreateActor/ResumeActor responses. Only the lifecycle RPCs SubstrateHarness uses // are implemented; the rest come from the embedded Unimplemented server. type MockControlServer struct { ateapipb.UnimplementedControlServer @@ -46,39 +46,49 @@ type MockControlServer struct { resumeCalls []string suspendCalls []string - CreateErr error // returned from CreateActor when non-nil - ResumeIP string // AteomPodIp returned from ResumeActor - ResumeNilActor bool // when true, ResumeActor returns a nil Actor + CreateErr error // returned from CreateActor when non-nil + ResumeIP string // AteomPodIp returned from ResumeActor + ResumeIPs []string // per-call AteomPodIp values; overrides ResumeIP when set + ResumeNilActor bool // when true, ResumeActor returns a nil Actor + SuspendErr error // returned from SuspendActor when non-nil } -func (f *MockControlServer) CreateAtespace(_ context.Context, req *ateapipb.CreateAtespaceRequest) (*ateapipb.CreateAtespaceResponse, error) { - return &ateapipb.CreateAtespaceResponse{Atespace: &ateapipb.Atespace{Name: req.GetName()}}, nil +func (f *MockControlServer) CreateAtespace(_ context.Context, req *ateapipb.CreateAtespaceRequest) (*ateapipb.Atespace, error) { + return &ateapipb.Atespace{Metadata: &ateapipb.ResourceMetadata{Name: req.GetAtespace().GetMetadata().GetName()}}, nil } -func (f *MockControlServer) CreateActor(_ context.Context, req *ateapipb.CreateActorRequest) (*ateapipb.CreateActorResponse, error) { +func (f *MockControlServer) CreateActor(_ context.Context, req *ateapipb.CreateActorRequest) (*ateapipb.Actor, error) { f.mu.Lock() - f.createCalls = append(f.createCalls, req.GetActorRef().GetName()) + f.createCalls = append(f.createCalls, req.GetActor().GetMetadata().GetName()) f.mu.Unlock() if f.CreateErr != nil { return nil, f.CreateErr } - return &ateapipb.CreateActorResponse{Actor: &ateapipb.Actor{ActorId: req.GetActorRef().GetName()}}, nil + return &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Name: req.GetActor().GetMetadata().GetName()}}, nil } func (f *MockControlServer) ResumeActor(_ context.Context, req *ateapipb.ResumeActorRequest) (*ateapipb.ResumeActorResponse, error) { f.mu.Lock() - f.resumeCalls = append(f.resumeCalls, req.GetActorRef().GetName()) + f.resumeCalls = append(f.resumeCalls, req.GetActor().GetName()) + resumeIP := f.ResumeIP + if len(f.ResumeIPs) > 0 { + index := min(len(f.resumeCalls)-1, len(f.ResumeIPs)-1) + resumeIP = f.ResumeIPs[index] + } f.mu.Unlock() if f.ResumeNilActor { return &ateapipb.ResumeActorResponse{}, nil } - return &ateapipb.ResumeActorResponse{Actor: &ateapipb.Actor{ActorId: req.GetActorRef().GetName(), AteomPodIp: f.ResumeIP}}, nil + return &ateapipb.ResumeActorResponse{Actor: &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Name: req.GetActor().GetName()}, AteomPodIp: resumeIP}}, nil } func (f *MockControlServer) SuspendActor(_ context.Context, req *ateapipb.SuspendActorRequest) (*ateapipb.SuspendActorResponse, error) { f.mu.Lock() - f.suspendCalls = append(f.suspendCalls, req.GetActorRef().GetName()) + f.suspendCalls = append(f.suspendCalls, req.GetActor().GetName()) f.mu.Unlock() + if f.SuspendErr != nil { + return nil, f.SuspendErr + } return &ateapipb.SuspendActorResponse{}, nil } @@ -109,8 +119,12 @@ type MockHarnessServer struct { ErrCode int32 // ErrMessage is the error text used by FailConnect/FailFrame. ErrMessage string + // HarnessMetadata is attached to the terminal HarnessEnd frame, whether + // COMPLETED or (via FailFrame) FAILED. + HarnessMetadata []byte mu sync.Mutex + connectCalls int gotConvID string gotHarnessID string gotHarnessConfig []byte @@ -118,6 +132,9 @@ type MockHarnessServer struct { } func (s *MockHarnessServer) Connect(stream proto.HarnessService_ConnectServer) error { + s.mu.Lock() + s.connectCalls++ + s.mu.Unlock() if s.FailConnect { return status.Error(codes.Internal, s.ErrMessage) } @@ -151,6 +168,7 @@ func (s *MockHarnessServer) Connect(stream proto.HarnessService_ConnectServer) e Code: s.ErrCode, Description: s.ErrMessage, }, + HarnessMetadata: s.HarnessMetadata, }, }, }) @@ -174,10 +192,20 @@ func (s *MockHarnessServer) Connect(stream proto.HarnessService_ConnectServer) e } return stream.Send(&proto.HarnessResponse{ ConversationId: convID, - Type: &proto.HarnessResponse_End{End: &proto.HarnessEnd{State: proto.State_STATE_COMPLETED}}, + Type: &proto.HarnessResponse_End{End: &proto.HarnessEnd{ + State: proto.State_STATE_COMPLETED, + HarnessMetadata: s.HarnessMetadata, + }}, }) } +// ConnectCalls returns how many per-turn HarnessService streams were opened. +func (s *MockHarnessServer) ConnectCalls() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.connectCalls +} + // Received returns a copy of the start frame the server received. func (s *MockHarnessServer) Received() (convID, harnessID string, harnessConfig []byte, inputs []string) { s.mu.Lock() diff --git a/internal/harness/stream.go b/internal/harness/stream.go index 7551da85..be26cf95 100644 --- a/internal/harness/stream.go +++ b/internal/harness/stream.go @@ -22,11 +22,26 @@ import ( "github.com/google/ax/proto" ) +// terminalMetadataHandler is an internal, optional completion capability. It +// keeps Handler unchanged while allowing the controller to retain opaque +// metadata from harnesses that provide it. +type terminalMetadataHandler interface { + OnCompleteWithMetadata(ctx context.Context, execID string, metadata []byte) error +} + +// failMetadataHandler is the failure counterpart to terminalMetadataHandler: +// it lets the controller retain opaque metadata (e.g. token usage) collected +// before a failed execution. +type failMetadataHandler interface { + OnFailWithMetadata(ctx context.Context, execID string, metadata []byte, cause error) error +} + // DrainStream reads from the harness gRPC stream until io.EOF, dispatching messages // to the handler, and returns the final execution status. func DrainStream(ctx context.Context, stream proto.HarnessService_ConnectClient, execID string, handler Handler) error { var endState proto.State var endErr error + var harnessMetadata []byte hasEnd := false for { @@ -48,6 +63,7 @@ func DrainStream(ctx context.Context, stream proto.HarnessService_ConnectClient, case *proto.HarnessResponse_End: hasEnd = true endState = payload.End.GetState() + harnessMetadata = payload.End.GetHarnessMetadata() if endState == proto.State_STATE_FAILED { if errDetail := payload.End.GetError(); errDetail != nil { endErr = fmt.Errorf("harness failed: [%d] %s", errDetail.GetCode(), errDetail.GetDescription()) @@ -62,7 +78,17 @@ func DrainStream(ctx context.Context, stream proto.HarnessService_ConnectClient, return fmt.Errorf("harness stream ended without HarnessEnd frame") } if endState == proto.State_STATE_FAILED { + if len(harnessMetadata) > 0 { + if fh, ok := handler.(failMetadataHandler); ok { + return fh.OnFailWithMetadata(ctx, execID, harnessMetadata, endErr) + } + } return endErr } + if len(harnessMetadata) > 0 { + if metadataHandler, ok := handler.(terminalMetadataHandler); ok { + return metadataHandler.OnCompleteWithMetadata(ctx, execID, harnessMetadata) + } + } return handler.OnComplete(ctx, execID) } diff --git a/internal/harness/stream_test.go b/internal/harness/stream_test.go new file mode 100644 index 00000000..435c8bf1 --- /dev/null +++ b/internal/harness/stream_test.go @@ -0,0 +1,157 @@ +// 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 +// +// https://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 harness_test + +import ( + "bytes" + "context" + "testing" + + "github.com/google/ax/internal/harness" + "github.com/google/ax/internal/harness/harnesstest" + "github.com/google/ax/proto" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" +) + +type terminalMetadataHandler struct { + completed bool + metadata []byte +} + +func (h *terminalMetadataHandler) OnMessage(context.Context, string, *proto.Message) error { + return nil +} + +func (h *terminalMetadataHandler) OnComplete(context.Context, string) error { + h.completed = true + return nil +} + +func (h *terminalMetadataHandler) OnCompleteWithMetadata(_ context.Context, _ string, metadata []byte) error { + h.metadata = append([]byte(nil), metadata...) + return nil +} + +// failMetadataHandler is the failure counterpart to terminalMetadataHandler +// above: it records the metadata delivered alongside a FAILED terminal event +// while still surfacing the original cause to the caller. +type failMetadataHandler struct { + completed bool + metadata []byte + cause error +} + +func (h *failMetadataHandler) OnMessage(context.Context, string, *proto.Message) error { + return nil +} + +func (h *failMetadataHandler) OnComplete(context.Context, string) error { + h.completed = true + return nil +} + +func (h *failMetadataHandler) OnFailWithMetadata(_ context.Context, _ string, metadata []byte, cause error) error { + h.metadata = append([]byte(nil), metadata...) + h.cause = cause + return cause +} + +func TestDrainStreamDispatchesFailedTerminalHarnessMetadata(t *testing.T) { + want := []byte("agentfleet-failed-metadata-fixture") + addr := harnesstest.StartHarnessServer(t, &harnesstest.MockHarnessServer{ + FailFrame: true, + ErrCode: 13, + ErrMessage: "boom", + HarnessMetadata: want, + }) + conn, err := grpc.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("grpc.NewClient: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + + ctx := context.Background() + stream, err := proto.NewHarnessServiceClient(conn).Connect(ctx) + if err != nil { + t.Fatalf("Connect: %v", err) + } + if err := stream.Send(&proto.HarnessRequest{ + ConversationId: "conv-1", + Type: &proto.HarnessRequest_Start{ + Start: &proto.HarnessStart{}, + }, + }); err != nil { + t.Fatalf("Send: %v", err) + } + if err := stream.CloseSend(); err != nil { + t.Fatalf("CloseSend: %v", err) + } + + handler := &failMetadataHandler{} + err = harness.DrainStream(ctx, stream, "exec-1", handler) + if err == nil { + t.Fatal("DrainStream: expected a non-nil error for a FAILED terminal event") + } + if handler.completed { + t.Fatal("OnComplete called on a FAILED terminal event") + } + if !bytes.Equal(handler.metadata, want) { + t.Fatalf("metadata = %q, want %q", handler.metadata, want) + } + if handler.cause == nil || err.Error() != handler.cause.Error() { + t.Fatalf("DrainStream error = %v, want the same cause the handler received (%v)", err, handler.cause) + } +} + +func TestDrainStreamDispatchesTerminalHarnessMetadata(t *testing.T) { + want := []byte("agentfleet-metadata-fixture") + addr := harnesstest.StartHarnessServer(t, &harnesstest.MockHarnessServer{ + HarnessMetadata: want, + }) + conn, err := grpc.NewClient(addr, grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("grpc.NewClient: %v", err) + } + t.Cleanup(func() { _ = conn.Close() }) + + ctx := context.Background() + stream, err := proto.NewHarnessServiceClient(conn).Connect(ctx) + if err != nil { + t.Fatalf("Connect: %v", err) + } + if err := stream.Send(&proto.HarnessRequest{ + ConversationId: "conv-1", + Type: &proto.HarnessRequest_Start{ + Start: &proto.HarnessStart{}, + }, + }); err != nil { + t.Fatalf("Send: %v", err) + } + if err := stream.CloseSend(); err != nil { + t.Fatalf("CloseSend: %v", err) + } + + handler := &terminalMetadataHandler{} + if err := harness.DrainStream(ctx, stream, "exec-1", handler); err != nil { + t.Fatalf("DrainStream: %v", err) + } + if handler.completed { + t.Fatal("OnComplete called instead of the metadata-aware completion capability") + } + if !bytes.Equal(handler.metadata, want) { + t.Fatalf("metadata = %q, want %q", handler.metadata, want) + } +} diff --git a/internal/harness/substrate/substrate.go b/internal/harness/substrate/substrate.go index 4743a602..82f83497 100644 --- a/internal/harness/substrate/substrate.go +++ b/internal/harness/substrate/substrate.go @@ -21,6 +21,7 @@ import ( "fmt" "io" "log/slog" + "os" "sync" "time" @@ -47,16 +48,48 @@ var _ harness.Execution = (*substrateExecution)(nil) // created/resumed actor's harness to become reachable and ready. const healthCheckTimeout = 60 * time.Second +const defaultWarmIdleTimeout = 30 * time.Second + +// defaultWarmProbeTimeout bounds the health probe against a reused warm +// address: the actor was serving moments ago, so unreachability within this +// window means it is gone and the cold resume path should take over. +const defaultWarmProbeTimeout = 5 * time.Second + +type idleMode uint8 + +const ( + idleModeImmediateSuspend idleMode = iota + idleModeWarmThenSuspend +) + +type warmActorState struct { + generation uint64 + workerAddr string + inTurn bool + timer *time.Timer + suspending chan struct{} +} + // SubstrateHarness manages execution in a SubstrATE sandboxed actor over gRPC HarnessService. type SubstrateHarness struct { harnessID string ateClient *ate.Client port int dialOpts []grpc.DialOption + + idleMode idleMode + idleTimeout time.Duration + warmProbeTimeout time.Duration + idleMu sync.Mutex + warmActors map[string]*warmActorState } // New creates a new SubstrateHarness. func New(harnessID string, endpoint string, namespace string, template string, port int, opts ...grpc.DialOption) (*SubstrateHarness, error) { + idleMode, idleTimeout, err := idlePolicyFromEnv() + if err != nil { + return nil, err + } if port == 0 { port = 50053 // Default HarnessService port } @@ -76,19 +109,93 @@ func New(harnessID string, endpoint string, namespace string, template string, p } opts = append(opts, grpc.WithStatsHandler(otelgrpc.NewClientHandler())) return &SubstrateHarness{ - harnessID: harnessID, - ateClient: client, - port: port, - dialOpts: opts, + harnessID: harnessID, + ateClient: client, + port: port, + dialOpts: opts, + idleMode: idleMode, + idleTimeout: idleTimeout, + warmProbeTimeout: defaultWarmProbeTimeout, }, nil } +func idlePolicyFromEnv() (idleMode, time.Duration, error) { + modeValue := os.Getenv("AX_SUBSTRATE_IDLE_MODE") + switch modeValue { + case "", "immediate-suspend": + return idleModeImmediateSuspend, 0, nil + case "warm-then-suspend": + timeoutValue := os.Getenv("AX_SUBSTRATE_IDLE_TIMEOUT") + if timeoutValue == "" { + return idleModeWarmThenSuspend, defaultWarmIdleTimeout, nil + } + timeout, err := time.ParseDuration(timeoutValue) + if err != nil { + return idleModeImmediateSuspend, 0, fmt.Errorf("invalid AX_SUBSTRATE_IDLE_TIMEOUT %q: %w", timeoutValue, err) + } + if timeout <= 0 { + return idleModeImmediateSuspend, 0, fmt.Errorf("AX_SUBSTRATE_IDLE_TIMEOUT must be positive") + } + return idleModeWarmThenSuspend, timeout, nil + default: + return idleModeImmediateSuspend, 0, fmt.Errorf("invalid AX_SUBSTRATE_IDLE_MODE %q", modeValue) + } +} + // Start implements Harness interface. It creates/resumes the target actor. -func (h *SubstrateHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (harness.Execution, error) { +func (h *SubstrateHarness) Start(ctx context.Context, conversationID string, harnessConfig []byte) (execution harness.Execution, err error) { if conversationID == "" { return nil, errors.New("SubstrateHarness needs valid conversationID") } + workerAddr, reusedWarmActor, err := h.beginWarmTurn(ctx, conversationID) + if err != nil { + return nil, err + } + if h.idleMode == idleModeWarmThenSuspend { + defer func() { + if err != nil { + h.abortWarmTurn(conversationID) + } + }() + } + + if reusedWarmActor { + // The cached address only proves that this process previously observed a + // running actor. Resolve through the authoritative control path before + // connecting because a worker IP can be reassigned while the actor is warm. + cachedWorkerAddr := workerAddr + workerAddr, err = h.resumeWorkerAddr(ctx, conversationID) + if err != nil { + return nil, err + } + h.rememberWarmActor(conversationID, workerAddr) + + exec, probeErr := h.connect(ctx, conversationID, harnessConfig, workerAddr, h.probeTimeout()) + if probeErr == nil { + return exec, nil + } + if ctx.Err() != nil { + return nil, probeErr + } + slog.WarnContext(ctx, "Warm SubstrATE actor unreachable; restarting before cold resume", + slog.String("conversation_id", conversationID), + slog.String("worker_addr", workerAddr), + slog.String("cached_worker_addr", cachedWorkerAddr), + slog.Any("error", probeErr), + ) + // ResumeActor is a no-op for an actor already marked RUNNING. Suspend it + // first so the cold path below performs a real restore. Keep the cached + // address until this succeeds so a failed reset still gets idle cleanup. + resetCtx, cancelReset := context.WithTimeout(ctx, 10*time.Second) + _, suspendErr := h.ateClient.SuspendActor(resetCtx, conversationID) + cancelReset() + if suspendErr != nil { + return nil, fmt.Errorf("failed to reset unreachable substrate actor %s after %v: %w", conversationID, probeErr, suspendErr) + } + h.forgetWarmAddr(conversationID) + } + // CreateActor is idempotent here: on follow-up turns the actor was created // (and suspended) on a previous turn, so AlreadyExists is expected and fine. if _, err := h.ateClient.CreateActor(ctx, conversationID); err != nil && status.Code(err) != codes.AlreadyExists { @@ -96,28 +203,45 @@ func (h *SubstrateHarness) Start(ctx context.Context, conversationID string, har } // Resume the actor so it is scheduled onto a worker and gets a routable IP. + workerAddr, err = h.resumeWorkerAddr(ctx, conversationID) + if err != nil { + return nil, err + } + h.rememberWarmActor(conversationID, workerAddr) + + return h.connect(ctx, conversationID, harnessConfig, workerAddr, healthCheckTimeout) +} + +// resumeWorkerAddr resolves the current actor through ATE and returns its +// authoritative worker address. ResumeActor is idempotent for RUNNING actors, +// so warm turns pay only the control-plane check and do not restore the actor. +func (h *SubstrateHarness) resumeWorkerAddr(ctx context.Context, conversationID string) (string, error) { resumeResp, err := h.ateClient.ResumeActor(ctx, conversationID) if err != nil { - return nil, fmt.Errorf("failed to resume substrate actor %s: %w", conversationID, err) + return "", fmt.Errorf("failed to resume substrate actor %s: %w", conversationID, err) } - actor := resumeResp.Actor + actor := resumeResp.GetActor() if actor == nil { - return nil, fmt.Errorf("received nil actor in response for %s", conversationID) + return "", fmt.Errorf("received nil actor in response for %s", conversationID) } - if actor.AteomPodIp == "" { - return nil, fmt.Errorf("actor %s has no active worker IP address", conversationID) + if actor.GetMetadata().GetName() != conversationID { + return "", fmt.Errorf("received actor %s while resuming %s", actor.GetMetadata().GetName(), conversationID) } + if actor.GetAteomPodIp() == "" { + return "", fmt.Errorf("actor %s has no active worker IP address", conversationID) + } + return fmt.Sprintf("%s:%d", actor.GetAteomPodIp(), h.port), nil +} - // Establish connection to the actor's worker IP - workerAddr := fmt.Sprintf("%s:%d", actor.AteomPodIp, h.port) +// connect dials the actor's worker address and waits for the harness to be +// reachable and ready before handing back the execution. +func (h *SubstrateHarness) connect(ctx context.Context, conversationID string, harnessConfig []byte, workerAddr string, healthTimeout time.Duration) (harness.Execution, error) { conn, err := grpc.NewClient(workerAddr, h.dialOpts...) if err != nil { return nil, fmt.Errorf("failed to dial remote harness service at %s: %w", workerAddr, err) } - // Wait for the harness to be reachable and ready before handing back the - // execution. - if err := waitForHealthy(ctx, conn, healthCheckTimeout); err != nil { + if err := waitForHealthy(ctx, conn, healthTimeout); err != nil { conn.Close() return nil, fmt.Errorf("harness for %s not ready at %s: %w", conversationID, workerAddr, err) } @@ -132,6 +256,109 @@ func (h *SubstrateHarness) Start(ctx context.Context, conversationID string, har }, nil } +func (h *SubstrateHarness) probeTimeout() time.Duration { + if h.warmProbeTimeout > 0 { + return h.warmProbeTimeout + } + return defaultWarmProbeTimeout +} + +// forgetWarmAddr drops the cached worker address for an in-turn conversation +// after a failed reuse probe, so the cold path re-resolves it and a failed +// turn deletes the entry instead of re-arming a timer around a dead address. +func (h *SubstrateHarness) forgetWarmAddr(conversationID string) { + h.idleMu.Lock() + defer h.idleMu.Unlock() + if state := h.warmActors[conversationID]; state != nil { + state.workerAddr = "" + } +} + +func (h *SubstrateHarness) beginWarmTurn(ctx context.Context, conversationID string) (string, bool, error) { + if h.idleMode != idleModeWarmThenSuspend { + return "", false, nil + } + + for { + h.idleMu.Lock() + if h.warmActors == nil { + h.warmActors = make(map[string]*warmActorState) + } + state := h.warmActors[conversationID] + if state == nil { + state = &warmActorState{inTurn: true} + h.warmActors[conversationID] = state + h.idleMu.Unlock() + return "", false, nil + } + if state.suspending != nil { + done := state.suspending + h.idleMu.Unlock() + select { + case <-ctx.Done(): + return "", false, ctx.Err() + case <-done: + continue + } + } + if state.inTurn { + h.idleMu.Unlock() + return "", false, fmt.Errorf("substrate actor %s already has an active turn", conversationID) + } + + state.generation++ + if state.timer != nil { + state.timer.Stop() + state.timer = nil + } + state.inTurn = true + workerAddr := state.workerAddr + h.idleMu.Unlock() + return workerAddr, workerAddr != "", nil + } +} + +func (h *SubstrateHarness) rememberWarmActor(conversationID, workerAddr string) { + if h.idleMode != idleModeWarmThenSuspend { + return + } + h.idleMu.Lock() + defer h.idleMu.Unlock() + state := h.warmActors[conversationID] + if state != nil { + state.workerAddr = workerAddr + } +} + +func (h *SubstrateHarness) abortWarmTurn(conversationID string) { + if h.idleMode != idleModeWarmThenSuspend { + return + } + h.idleMu.Lock() + defer h.idleMu.Unlock() + state := h.warmActors[conversationID] + if state == nil { + return + } + h.endWarmTurnLocked(state, conversationID, "") +} + +// endWarmTurnLocked marks the conversation's turn ended and either drops the +// entry (no usable worker address) or arms the idle suspend timer. The caller +// must hold idleMu. +func (h *SubstrateHarness) endWarmTurnLocked(state *warmActorState, conversationID, execID string) { + state.inTurn = false + if state.workerAddr == "" { + delete(h.warmActors, conversationID) + return + } + state.generation++ + generation := state.generation + state.timer = time.AfterFunc(h.idleTimeout, func() { + h.suspendWarmActor(conversationID, execID, generation) + }) +} + // waitForHealthy blocks until the harness behind conn reports SERVING via the // standard gRPC health protocol until timeout. A harness that is reachable // but does not implement the health service (Unimplemented) is treated as @@ -232,25 +459,137 @@ func (e *substrateExecution) Run(ctx context.Context, handler harness.Handler) e return harness.DrainStream(ctx, stream, e.execID, handler) } +// CloseBeforeNextStart implements the controller's optional eager-close +// capability. Warm mode tracks a per-conversation turn slot, so this execution +// must be closed before another Start for the same conversation; immediate +// mode keeps upstream's deferred-close semantics. +func (e *substrateExecution) CloseBeforeNextStart() bool { + return e.harness.idleMode == idleModeWarmThenSuspend +} + func (e *substrateExecution) Close(ctx context.Context) error { - // Close connection if e.conn != nil { - e.conn.Close() + _ = e.conn.Close() + } + if e.harness.idleMode == idleModeWarmThenSuspend { + e.harness.scheduleWarmSuspend(e.conversationID, e.execID) + return nil + } + + e.harness.suspendActor(ctx, e.conversationID, e.execID) + return nil +} + +func (h *SubstrateHarness) scheduleWarmSuspend(conversationID, execID string) { + h.idleMu.Lock() + defer h.idleMu.Unlock() + state := h.warmActors[conversationID] + if state == nil { + return } + h.endWarmTurnLocked(state, conversationID, execID) +} + +func (h *SubstrateHarness) suspendWarmActor(conversationID, execID string, generation uint64) { + h.idleMu.Lock() + state := h.warmActors[conversationID] + if state == nil || state.generation != generation || state.inTurn { + h.idleMu.Unlock() + return + } + state.timer = nil + done := make(chan struct{}) + state.suspending = done + h.idleMu.Unlock() + + h.suspendActor(context.Background(), conversationID, execID) + + h.idleMu.Lock() + state = h.warmActors[conversationID] + if state != nil && state.suspending == done { + delete(h.warmActors, conversationID) + } + close(done) + h.idleMu.Unlock() +} +// Shutdown drains warm actors awaiting idle suspension so a process exit does +// not leak them as RUNNING actors. A warm actor sits between turns with a +// pending idle timer whose only home is this process's memory; if the process +// dies before the timer fires, the actor is never suspended. Shutdown stops each +// pending timer and suspends the actor synchronously. +// +// Actors with an active turn are left untouched: the turn owns the actor and +// schedules its own suspension on Close. Callers should therefore invoke +// Shutdown only after in-flight turns have drained (e.g. after the gRPC server's +// GracefulStop returns) so no turn re-arms a timer after this drain. +func (h *SubstrateHarness) Shutdown(ctx context.Context) { + if h.idleMode != idleModeWarmThenSuspend { + return + } + + h.idleMu.Lock() + var ( + drain []string + inProgress []chan struct{} + ) + for conversationID, state := range h.warmActors { + switch { + case state.inTurn: + // An active turn owns the actor; it will suspend on Close. + continue + case state.suspending != nil: + // A fired timer is already suspending this actor; wait for it below. + inProgress = append(inProgress, state.suspending) + case state.timer != nil: + state.timer.Stop() + state.timer = nil + // Neutralize a timer callback that already fired but is still blocked + // on idleMu: the generation bump makes suspendWarmActor a no-op so it + // cannot suspend again after we do. + state.generation++ + if state.workerAddr == "" { + delete(h.warmActors, conversationID) + continue + } + drain = append(drain, conversationID) + default: + // No timer and not suspending: nothing deferred to clean up. + delete(h.warmActors, conversationID) + } + } + h.idleMu.Unlock() + + for _, conversationID := range drain { + h.suspendActor(ctx, conversationID, "") + h.idleMu.Lock() + delete(h.warmActors, conversationID) + h.idleMu.Unlock() + } + + // Wait for any timer-driven suspensions already in progress so they are not + // cut short by process exit. + for _, done := range inProgress { + select { + case <-ctx.Done(): + return + case <-done: + } + } +} + +func (h *SubstrateHarness) suspendActor(ctx context.Context, conversationID, execID string) { // Suspend actor to return resource to standard standby pool slog.InfoContext(ctx, "Suspending SubstrATE actor", - slog.String("conversation_id", e.conversationID), - slog.String("exec_id", e.execID), + slog.String("conversation_id", conversationID), + slog.String("exec_id", execID), ) suspendCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - if _, err := e.harness.ateClient.SuspendActor(suspendCtx, e.conversationID); err != nil { + if _, err := h.ateClient.SuspendActor(suspendCtx, conversationID); err != nil { slog.ErrorContext(ctx, "Failed to suspend SubstrATE actor", - slog.String("conversation_id", e.conversationID), + slog.String("conversation_id", conversationID), slog.Any("error", err), ) } - - return nil } diff --git a/internal/harness/substrate/substrate_test.go b/internal/harness/substrate/substrate_test.go index 4403dbd1..755dd6bf 100644 --- a/internal/harness/substrate/substrate_test.go +++ b/internal/harness/substrate/substrate_test.go @@ -26,6 +26,7 @@ import ( "github.com/google/ax/internal/harness/harnesstest" "github.com/google/ax/internal/ate" + "github.com/google/ax/proto" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" @@ -124,6 +125,49 @@ func TestWaitForHealthy_ServerDown(t *testing.T) { } } +func TestNewReadsWarmIdlePolicyFromEnvironment(t *testing.T) { + t.Setenv("AX_SUBSTRATE_IDLE_MODE", "warm-then-suspend") + t.Setenv("AX_SUBSTRATE_IDLE_TIMEOUT", "750ms") + + h, err := New("antigravity", "127.0.0.1:1", "ax", "antigravity-template", 50053) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = h.ateClient.Close() }) + if h.idleMode != idleModeWarmThenSuspend { + t.Fatalf("idle mode = %v, want warm-then-suspend", h.idleMode) + } + if h.idleTimeout != 750*time.Millisecond { + t.Fatalf("idle timeout = %v, want 750ms", h.idleTimeout) + } +} + +func TestNewDefaultsToImmediateSuspend(t *testing.T) { + t.Setenv("AX_SUBSTRATE_IDLE_MODE", "") + t.Setenv("AX_SUBSTRATE_IDLE_TIMEOUT", "") + + h, err := New("antigravity", "127.0.0.1:1", "ax", "antigravity-template", 50053) + if err != nil { + t.Fatalf("New: %v", err) + } + t.Cleanup(func() { _ = h.ateClient.Close() }) + if h.idleMode != idleModeImmediateSuspend || h.idleTimeout != 0 { + t.Fatalf("idle policy = (%v, %v), want immediate-suspend", h.idleMode, h.idleTimeout) + } +} + +func TestNewRejectsInvalidIdlePolicy(t *testing.T) { + for _, mode := range []string{"keep-forever", "pause-resume"} { + t.Run(mode, func(t *testing.T) { + t.Setenv("AX_SUBSTRATE_IDLE_MODE", mode) + + if _, err := New("antigravity", "127.0.0.1:1", "ax", "antigravity-template", 50053); err == nil { + t.Fatalf("New accepted invalid AX_SUBSTRATE_IDLE_MODE %q", mode) + } + }) + } +} + // newTestSubstrateHarness builds a SubstrateHarness wired to the mock control // server and the mock harness server. It constructs the struct directly (rather // than via NewSubstrateHarness) so the control client can use insecure @@ -213,6 +257,133 @@ func TestSubstrateHarness_EndToEnd(t *testing.T) { } } +func TestSubstrateHarness_WarmThenSuspendReusesActorUntilIdle(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv)) + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = 200 * time.Millisecond + + runTurn := func(input string) { + t.Helper() + ctx := context.Background() + exec, err := h.Start(ctx, "conv-warm", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start(%q): %v", input, err) + } + if err := exec.Queue(ctx, harnesstest.UserText(input)); err != nil { + t.Fatalf("Queue(%q): %v", input, err) + } + if err := exec.Run(ctx, &harnesstest.MockHandler{}); err != nil { + t.Fatalf("Run(%q): %v", input, err) + } + if err := exec.Close(ctx); err != nil { + t.Fatalf("Close(%q): %v", input, err) + } + } + + runTurn("one") + runTurn("two") + + create, resume, suspend := ctrl.Calls() + if !slices.Equal(create, []string{"conv-warm"}) { + t.Fatalf("create=%v, want one actor creation", create) + } + if !slices.Equal(resume, []string{"conv-warm", "conv-warm"}) { + t.Fatalf("resume=%v, want one authoritative resume check per turn", resume) + } + if len(suspend) != 0 { + t.Fatalf("suspend called before the final idle timeout: %v", suspend) + } + if got := srv.ConnectCalls(); got != 2 { + t.Fatalf("HarnessService Connect calls = %d, want one fresh stream per turn", got) + } + + deadline := time.Now().Add(2 * time.Second) + for { + _, _, suspend = ctrl.Calls() + if slices.Equal(suspend, []string{"conv-warm"}) { + break + } + if time.Now().After(deadline) { + t.Fatalf("suspend=%v, want one call after idle timeout", suspend) + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestSubstrateHarness_StaleWarmTimerCannotSuspendActiveTurn(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv)) + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = 100 * time.Millisecond + + ctx := context.Background() + first, err := h.Start(ctx, "conv-generation", substrateHarnessConfig) + if err != nil { + t.Fatalf("first Start: %v", err) + } + if err := first.Run(ctx, &harnesstest.MockHandler{}); err != nil { + t.Fatalf("first Run: %v", err) + } + if err := first.Close(ctx); err != nil { + t.Fatalf("first Close: %v", err) + } + + second, err := h.Start(ctx, "conv-generation", substrateHarnessConfig) + if err != nil { + t.Fatalf("second Start: %v", err) + } + time.Sleep(2 * h.idleTimeout) + if _, _, suspend := ctrl.Calls(); len(suspend) != 0 { + t.Fatalf("stale timer suspended an active turn: %v", suspend) + } + if err := second.Run(ctx, &harnesstest.MockHandler{}); err != nil { + t.Fatalf("second Run: %v", err) + } + if err := second.Close(ctx); err != nil { + t.Fatalf("second Close: %v", err) + } + + deadline := time.Now().Add(2 * time.Second) + for { + _, _, suspend := ctrl.Calls() + if slices.Equal(suspend, []string{"conv-generation"}) { + break + } + if time.Now().After(deadline) { + t.Fatalf("suspend=%v, want one call after the active turn closed", suspend) + } + time.Sleep(10 * time.Millisecond) + } +} + +func TestSubstrateHarness_WarmStartFailureStillSchedulesSuspend(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), "127.0.0.1:1") + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = 20 * time.Millisecond + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + if _, err := h.Start(ctx, "conv-start-failure", substrateHarnessConfig); err == nil { + t.Fatal("Start unexpectedly succeeded against an unavailable harness") + } + + deadline := time.Now().Add(2 * time.Second) + for { + _, _, suspend := ctrl.Calls() + if slices.Equal(suspend, []string{"conv-start-failure"}) { + break + } + if time.Now().After(deadline) { + t.Fatalf("suspend=%v, want failed warm Start to retain idle cleanup", suspend) + } + time.Sleep(10 * time.Millisecond) + } +} + func TestSubstrateHarness_CreateAlreadyExistsTolerated(t *testing.T) { ctrl := &harnesstest.MockControlServer{ ResumeIP: "127.0.0.1", @@ -288,3 +459,332 @@ func TestSubstrateHarness_HarnessFailedFrame(t *testing.T) { t.Errorf("error = %v, want it to mention 'harness failed'", err) } } + +func startStoppableHarnessServerOn(t *testing.T, srv *harnesstest.MockHarnessServer, listenAddr string) (string, func()) { + t.Helper() + lis, err := net.Listen("tcp4", listenAddr) + if err != nil { + t.Fatalf("failed to listen on %s: %v", listenAddr, err) + } + s := grpc.NewServer() + proto.RegisterHarnessServiceServer(s, srv) + hs := health.NewServer() + hs.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING) + grpc_health_v1.RegisterHealthServer(s, hs) + go func() { _ = s.Serve(lis) }() + t.Cleanup(s.Stop) + return lis.Addr().String(), s.Stop +} + +func TestSubstrateHarness_UnreachableWarmActorIsSuspendedBeforeColdResume(t *testing.T) { + srvA := &harnesstest.MockHarnessServer{} + addrA, stopA := startStoppableHarnessServerOn(t, srvA, "127.0.0.1:0") + srvB := &harnesstest.MockHarnessServer{} + addrB := harnesstest.StartHarnessServer(t, srvB) + + ctrl := &harnesstest.MockControlServer{ + ResumeIPs: []string{"127.0.0.10", "127.0.0.10", "127.0.0.11", "127.0.0.11"}, + } + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), addrA) + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = time.Minute + h.warmProbeTimeout = 100 * time.Millisecond + h.dialOpts = []grpc.DialOption{ + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithContextDialer(func(ctx context.Context, target string) (net.Conn, error) { + host, _, err := net.SplitHostPort(target) + if err != nil { + return nil, err + } + actualAddr := map[string]string{"127.0.0.10": addrA, "127.0.0.11": addrB}[host] + return (&net.Dialer{}).DialContext(ctx, "tcp", actualAddr) + }), + } + + runTurn := func(input string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + exec, err := h.Start(ctx, "conv-restart", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start(%q): %v", input, err) + } + if err := exec.Queue(ctx, harnesstest.UserText(input)); err != nil { + t.Fatalf("Queue(%q): %v", input, err) + } + if err := exec.Run(ctx, &harnesstest.MockHandler{}); err != nil { + t.Fatalf("Run(%q): %v", input, err) + } + if err := exec.Close(ctx); err != nil { + t.Fatalf("Close(%q): %v", input, err) + } + } + + runTurn("one") + stopA() + runTurn("two") + runTurn("three") + + create, resume, suspend := ctrl.Calls() + if !slices.Equal(create, []string{"conv-restart", "conv-restart"}) { + t.Fatalf("create=%v, want initial create plus cold recovery", create) + } + if !slices.Equal(resume, []string{"conv-restart", "conv-restart", "conv-restart", "conv-restart"}) { + t.Fatalf("resume=%v, want authoritative checks plus cold recovery", resume) + } + if !slices.Equal(suspend, []string{"conv-restart"}) { + t.Fatalf("suspend=%v, want one reset before cold recovery", suspend) + } + if got := srvA.ConnectCalls(); got != 1 { + t.Fatalf("first worker Connect calls = %d, want 1", got) + } + if got := srvB.ConnectCalls(); got != 2 { + t.Fatalf("replacement worker Connect calls = %d, want 2", got) + } +} + +// A cached worker address can still be healthy after its IP has been reassigned +// to another actor. Warm reuse must resolve the current worker through ATE before +// connecting instead of treating generic gRPC health as proof of actor identity. +func TestSubstrateHarness_WarmReuseUsesAuthoritativeWorkerAddress(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srvA := &harnesstest.MockHarnessServer{} + addrA := harnesstest.StartHarnessServer(t, srvA) + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), addrA) + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = time.Minute // must not fire during the test + + runTurn := func(input string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + exec, err := h.Start(ctx, "conv-fallback", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start(%q): %v", input, err) + } + if err := exec.Queue(ctx, harnesstest.UserText(input)); err != nil { + t.Fatalf("Queue(%q): %v", input, err) + } + if err := exec.Run(ctx, &harnesstest.MockHandler{}); err != nil { + t.Fatalf("Run(%q): %v", input, err) + } + if err := exec.Close(ctx); err != nil { + t.Fatalf("Close(%q): %v", input, err) + } + } + + runTurn("one") + + // ATE now reports a different worker while the cached endpoint remains + // healthy. Changing the test port models the authoritative worker address. + srvB := &harnesstest.MockHarnessServer{} + addrB := harnesstest.StartHarnessServer(t, srvB) + _, portStr, err := net.SplitHostPort(addrB) + if err != nil { + t.Fatalf("bad replacement addr %q: %v", addrB, err) + } + if h.port, err = strconv.Atoi(portStr); err != nil { + t.Fatalf("bad replacement port %q: %v", portStr, err) + } + + runTurn("two") + + if got := srvA.ConnectCalls(); got != 1 { + t.Fatalf("cached worker Connect calls = %d, want no reuse after ATE moved the actor", got) + } + if got := srvB.ConnectCalls(); got != 1 { + t.Fatalf("authoritative worker Connect calls = %d, want 1", got) + } + wantResumes := []string{"conv-fallback", "conv-fallback"} + if _, resume, _ := ctrl.Calls(); !slices.Equal(resume, wantResumes) { + t.Fatalf("resume=%v, want one authoritative resume check per turn", resume) + } + + // Every warm turn revalidates ownership; it does not trust the refreshed + // address across turns. + runTurn("three") + wantResumes = append(wantResumes, "conv-fallback") + if _, resume, _ := ctrl.Calls(); !slices.Equal(resume, wantResumes) { + t.Fatalf("resume=%v, want one authoritative resume check per turn", resume) + } + if got := srvB.ConnectCalls(); got != 2 { + t.Fatalf("authoritative worker Connect calls = %d, want 2", got) + } +} + +func TestSubstrateHarness_FailedWarmResetRetainsIdleCleanup(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ + ResumeIP: "127.0.0.1", + SuspendErr: status.Error(codes.Unavailable, "control unavailable"), + } + ctrlAddr := harnesstest.StartControlServer(t, ctrl) + + // Reserve and release a port so the authoritative worker address is known + // to be unreachable. + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to reserve port: %v", err) + } + deadAddr := lis.Addr().String() + if err := lis.Close(); err != nil { + t.Fatalf("failed to release port: %v", err) + } + + h := newTestSubstrateHarness(t, ctrlAddr, deadAddr) + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = 20 * time.Millisecond + h.warmProbeTimeout = 50 * time.Millisecond + h.warmActors = map[string]*warmActorState{ + "conv-reset-failure": {workerAddr: deadAddr}, + } + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if _, err := h.Start(ctx, "conv-reset-failure", substrateHarnessConfig); err == nil { + t.Fatal("Start unexpectedly succeeded against an unreachable warm actor") + } + + deadline := time.Now().Add(2 * time.Second) + for { + _, _, suspend := ctrl.Calls() + if len(suspend) >= 2 { + return + } + if time.Now().After(deadline) { + t.Fatalf("suspend=%v, want failed reset plus idle cleanup retry", suspend) + } + time.Sleep(10 * time.Millisecond) + } +} + +// A warm actor sits with a pending idle-suspend timer between turns. If the +// ax-server process exits before that timer fires, the timer dies with the +// process and the actor is never suspended -- it leaks as a RUNNING actor +// holding a worker slot. Shutdown must drain those pending timers by suspending +// the warm actors synchronously. +func TestSubstrateHarness_ShutdownDrainsWarmActorsAwaitingIdleSuspend(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv)) + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = time.Hour // must NOT fire on its own during the test + + ctx := context.Background() + exec, err := h.Start(ctx, "conv-drain", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start: %v", err) + } + if err := exec.Run(ctx, &harnesstest.MockHandler{}); err != nil { + t.Fatalf("Run: %v", err) + } + if err := exec.Close(ctx); err != nil { + t.Fatalf("Close: %v", err) + } + + // The actor is warm now with a one-hour pending timer; nothing suspended yet. + if _, _, suspend := ctrl.Calls(); len(suspend) != 0 { + t.Fatalf("suspend=%v, want none before shutdown", suspend) + } + + h.Shutdown(ctx) + + if _, _, suspend := ctrl.Calls(); !slices.Equal(suspend, []string{"conv-drain"}) { + t.Fatalf("suspend=%v, want the warm actor suspended on shutdown", suspend) + } + + // The bookkeeping entry is gone, so a lingering timer cannot suspend twice. + h.idleMu.Lock() + _, present := h.warmActors["conv-drain"] + h.idleMu.Unlock() + if present { + t.Fatalf("warm actor entry still present after shutdown drain") + } +} + +// Shutdown must not touch an actor that still has an active turn: the turn owns +// the actor and will schedule its own suspension on Close. +func TestSubstrateHarness_ShutdownLeavesActiveTurnAlone(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv)) + h.idleMode = idleModeWarmThenSuspend + h.idleTimeout = time.Hour + + ctx := context.Background() + exec, err := h.Start(ctx, "conv-active", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(func() { _ = exec.Close(ctx) }) + + // Turn is in-flight (Start ran, Close has not). Shutdown must skip it. + h.Shutdown(ctx) + + if _, _, suspend := ctrl.Calls(); len(suspend) != 0 { + t.Fatalf("suspend=%v, want none while a turn is active", suspend) + } +} + +// Shutdown is a no-op in immediate-suspend mode, which never tracks warm actors. +func TestSubstrateHarness_ShutdownImmediateModeIsNoOp(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{} + h := newTestSubstrateHarness(t, harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv)) + // idleMode defaults to immediate-suspend. + + ctx := context.Background() + exec, err := h.Start(ctx, "conv-immediate", substrateHarnessConfig) + if err != nil { + t.Fatalf("Start: %v", err) + } + if err := exec.Close(ctx); err != nil { + t.Fatalf("Close: %v", err) + } + + // Close already suspended once (immediate mode); Shutdown adds nothing. + _, _, before := ctrl.Calls() + h.Shutdown(ctx) + if _, _, after := ctrl.Calls(); !slices.Equal(after, before) { + t.Fatalf("suspend calls changed across Shutdown: before=%v after=%v", before, after) + } +} + +// The eager-close capability tells the controller whether an execution must be +// closed before the next Start for the same conversation. Only warm mode needs +// that (turn-slot bookkeeping); immediate mode must keep upstream's +// deferred-close semantics. +func TestSubstrateExecutionEagerCloseCapabilityTracksIdleMode(t *testing.T) { + ctrl := &harnesstest.MockControlServer{ResumeIP: "127.0.0.1"} + srv := &harnesstest.MockHarnessServer{} + ctrlAddr, harnessAddr := harnesstest.StartControlServer(t, ctrl), harnesstest.StartHarnessServer(t, srv) + + for _, tc := range []struct { + name string + mode idleMode + want bool + }{ + {name: "warm-then-suspend", mode: idleModeWarmThenSuspend, want: true}, + {name: "immediate-suspend", mode: idleModeImmediateSuspend, want: false}, + } { + t.Run(tc.name, func(t *testing.T) { + h := newTestSubstrateHarness(t, ctrlAddr, harnessAddr) + h.idleMode = tc.mode + h.idleTimeout = time.Minute + + ctx := context.Background() + exec, err := h.Start(ctx, "conv-capability-"+tc.name, substrateHarnessConfig) + if err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(func() { _ = exec.Close(ctx) }) + + ec, ok := exec.(interface{ CloseBeforeNextStart() bool }) + if !ok { + t.Fatal("substrateExecution does not implement the eager-close capability") + } + if got := ec.CloseBeforeNextStart(); got != tc.want { + t.Fatalf("CloseBeforeNextStart() = %v, want %v", got, tc.want) + } + }) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 267460ff..aa50e0e8 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -73,7 +73,6 @@ func (s *Server) Exec(req *proto.ExecRequest, stream grpc.ServerStreamingServer[ return s.controller.Exec(ctx, req, outputHandler) } - func (s *Server) DeleteConversation(ctx context.Context, req *proto.DeleteConversationRequest) (*proto.DeleteConversationResponse, error) { slog.InfoContext(ctx, "Deleting conversation...", slog.String("conversation_id", req.ConversationId)) @@ -121,15 +120,20 @@ func (s *Server) Serve(address string, opts ...grpc.ServerOption) error { return nil } -// GracefulStop stops the gRPC server gracefully. +// GracefulStop stops the gRPC server gracefully. The gRPC server is drained +// first so in-flight turns run to completion before the controller releases +// resources. This ordering also lets warm harnesses drain correctly: by the +// time the controller closes, every turn has finished and parked its actor for +// idle suspension, so the drain suspends those actors instead of racing a turn +// that would re-arm the idle timer afterward. func (s *Server) GracefulStop() { slog.Info("Stopping server gracefully...") - if s.controller != nil { - s.controller.Close() - } if s.grpcServer != nil { s.grpcServer.GracefulStop() } + if s.controller != nil { + s.controller.Close() + } } func (s *Server) markInFlight(id string) (exists bool, cleanup func()) { diff --git a/proto/ax.pb.go b/proto/ax.pb.go index 3fa254f9..2be18ae9 100644 --- a/proto/ax.pb.go +++ b/proto/ax.pb.go @@ -210,8 +210,10 @@ type ConversationEvent struct { HarnessConfig *structpb.Struct `protobuf:"bytes,5,opt,name=harness_config,json=harnessConfig,proto3" json:"harness_config,omitempty"` Messages []*Message `protobuf:"bytes,6,rep,name=messages,proto3" json:"messages,omitempty"` State State `protobuf:"varint,7,opt,name=state,proto3,enum=ax.State" json:"state,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + // Opaque metadata produced by the harness for this terminal event. + HarnessMetadata []byte `protobuf:"bytes,8,opt,name=harness_metadata,json=harnessMetadata,proto3" json:"harness_metadata,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *ConversationEvent) Reset() { @@ -293,6 +295,13 @@ func (x *ConversationEvent) GetState() State { return State_STATE_UNSPECIFIED } +func (x *ConversationEvent) GetHarnessMetadata() []byte { + if x != nil { + return x.HarnessMetadata + } + return nil +} + type HarnessStart struct { state protoimpl.MessageState `protogen:"open.v1"` // Per-execution harness configuration. @@ -592,9 +601,11 @@ type HarnessEnd struct { // Terminal state for the harness execution. State State `protobuf:"varint,1,opt,name=state,proto3,enum=ax.State" json:"state,omitempty"` // Optional error details. - Error *Error `protobuf:"bytes,2,opt,name=error,proto3" json:"error,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + Error *Error `protobuf:"bytes,2,opt,name=error,proto3" json:"error,omitempty"` + // Opaque metadata associated with this completed turn. + HarnessMetadata []byte `protobuf:"bytes,3,opt,name=harness_metadata,json=harnessMetadata,proto3" json:"harness_metadata,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *HarnessEnd) Reset() { @@ -641,6 +652,13 @@ func (x *HarnessEnd) GetError() *Error { return nil } +func (x *HarnessEnd) GetHarnessMetadata() []byte { + if x != nil { + return x.HarnessMetadata + } + return nil +} + type HarnessResponse struct { state protoimpl.MessageState `protogen:"open.v1"` ConversationId string `protobuf:"bytes,1,opt,name=conversation_id,json=conversationId,proto3" json:"conversation_id,omitempty"` @@ -810,11 +828,13 @@ func (x *ExecRequest) GetHarnessConfig() []byte { // ExecResponse contains the result of an execution. type ExecResponse struct { - state protoimpl.MessageState `protogen:"open.v1"` - Outputs []*Message `protobuf:"bytes,1,rep,name=outputs,proto3" json:"outputs,omitempty"` // Output content - Step int32 `protobuf:"varint,2,opt,name=step,proto3" json:"step,omitempty"` // Step of the outputs - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + state protoimpl.MessageState `protogen:"open.v1"` + Outputs []*Message `protobuf:"bytes,1,rep,name=outputs,proto3" json:"outputs,omitempty"` // Output content + Step int32 `protobuf:"varint,2,opt,name=step,proto3" json:"step,omitempty"` // Step of the outputs + // Opaque metadata associated with the terminal event at step. + HarnessMetadata []byte `protobuf:"bytes,3,opt,name=harness_metadata,json=harnessMetadata,proto3" json:"harness_metadata,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *ExecResponse) Reset() { @@ -861,6 +881,13 @@ func (x *ExecResponse) GetStep() int32 { return 0 } +func (x *ExecResponse) GetHarnessMetadata() []byte { + if x != nil { + return x.HarnessMetadata + } + return nil +} + type DeleteConversationRequest struct { state protoimpl.MessageState `protogen:"open.v1"` ConversationId string `protobuf:"bytes,1,opt,name=conversation_id,json=conversationId,proto3" json:"conversation_id,omitempty"` @@ -948,7 +975,7 @@ const file_proto_ax_proto_rawDesc = "" + "\x0eproto/ax.proto\x12\x02ax\x1a\x1cgoogle/protobuf/struct.proto\x1a\x13proto/content.proto\"D\n" + "\aMessage\x12\x12\n" + "\x04role\x18\x01 \x01(\tR\x04role\x12%\n" + - "\acontent\x18\x02 \x01(\v2\v.ax.ContentR\acontent\"\x92\x02\n" + + "\acontent\x18\x02 \x01(\v2\v.ax.ContentR\acontent\"\xbd\x02\n" + "\x11ConversationEvent\x12'\n" + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\x12\x12\n" + "\x04step\x18\x02 \x01(\x05R\x04step\x12\x17\n" + @@ -957,7 +984,8 @@ const file_proto_ax_proto_rawDesc = "" + "harness_id\x18\x04 \x01(\tR\tharnessId\x12>\n" + "\x0eharness_config\x18\x05 \x01(\v2\x17.google.protobuf.StructR\rharnessConfig\x12'\n" + "\bmessages\x18\x06 \x03(\v2\v.ax.MessageR\bmessages\x12\x1f\n" + - "\x05state\x18\a \x01(\x0e2\t.ax.StateR\x05state\"^\n" + + "\x05state\x18\a \x01(\x0e2\t.ax.StateR\x05state\x12)\n" + + "\x10harness_metadata\x18\b \x01(\fR\x0fharnessMetadata\"^\n" + "\fHarnessStart\x12%\n" + "\x0eharness_config\x18\x01 \x01(\fR\rharnessConfig\x12'\n" + "\bmessages\x18\x02 \x03(\v2\v.ax.MessageR\bmessages\"9\n" + @@ -974,11 +1002,12 @@ const file_proto_ax_proto_rawDesc = "" + "\bmessages\x18\x01 \x03(\v2\v.ax.MessageR\bmessages\"=\n" + "\x05Error\x12\x12\n" + "\x04code\x18\x01 \x01(\x05R\x04code\x12 \n" + - "\vdescription\x18\x02 \x01(\tR\vdescription\"N\n" + + "\vdescription\x18\x02 \x01(\tR\vdescription\"y\n" + "\n" + "HarnessEnd\x12\x1f\n" + "\x05state\x18\x01 \x01(\x0e2\t.ax.StateR\x05state\x12\x1f\n" + - "\x05error\x18\x02 \x01(\v2\t.ax.ErrorR\x05error\"\x96\x01\n" + + "\x05error\x18\x02 \x01(\v2\t.ax.ErrorR\x05error\x12)\n" + + "\x10harness_metadata\x18\x03 \x01(\fR\x0fharnessMetadata\"\x96\x01\n" + "\x0fHarnessResponse\x12'\n" + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\x12.\n" + "\aoutputs\x18\x02 \x01(\v2\x12.ax.HarnessOutputsH\x00R\aoutputs\x12\"\n" + @@ -990,10 +1019,11 @@ const file_proto_ax_proto_rawDesc = "" + "\tlast_step\x18\x03 \x01(\x05R\blastStep\x12\x1d\n" + "\n" + "harness_id\x18\x04 \x01(\tR\tharnessId\x12%\n" + - "\x0eharness_config\x18\x05 \x01(\fR\rharnessConfig\"I\n" + + "\x0eharness_config\x18\x05 \x01(\fR\rharnessConfig\"t\n" + "\fExecResponse\x12%\n" + "\aoutputs\x18\x01 \x03(\v2\v.ax.MessageR\aoutputs\x12\x12\n" + - "\x04step\x18\x02 \x01(\x05R\x04step\"D\n" + + "\x04step\x18\x02 \x01(\x05R\x04step\x12)\n" + + "\x10harness_metadata\x18\x03 \x01(\fR\x0fharnessMetadata\"D\n" + "\x19DeleteConversationRequest\x12'\n" + "\x0fconversation_id\x18\x01 \x01(\tR\x0econversationId\"\x1c\n" + "\x1aDeleteConversationResponse*l\n" + diff --git a/proto/ax.proto b/proto/ax.proto index 31b90a61..919eecd5 100644 --- a/proto/ax.proto +++ b/proto/ax.proto @@ -41,6 +41,8 @@ message ConversationEvent { google.protobuf.Struct harness_config = 5; repeated Message messages = 6; State state = 7; + // Opaque metadata produced by the harness for this terminal event. + bytes harness_metadata = 8; } message HarnessStart { @@ -79,6 +81,8 @@ message HarnessEnd { State state = 1; // Optional error details. Error error = 2; + // Opaque metadata associated with this completed turn. + bytes harness_metadata = 3; } message HarnessResponse { @@ -128,6 +132,8 @@ message ExecRequest { message ExecResponse { repeated Message outputs = 1; // Output content int32 step = 2; // Step of the outputs + // Opaque metadata associated with the terminal event at step. + bytes harness_metadata = 3; } service ExecutionService { diff --git a/python/proto/ax_pb2.py b/python/proto/ax_pb2.py index 64befc9e..a2f5b52d 100644 --- a/python/proto/ax_pb2.py +++ b/python/proto/ax_pb2.py @@ -30,7 +30,7 @@ from proto import content_pb2 as proto_dot_content__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0eproto/ax.proto\x12\x02\x61x\x1a\x1cgoogle/protobuf/struct.proto\x1a\x13proto/content.proto\"5\n\x07Message\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x1c\n\x07\x63ontent\x18\x02 \x01(\x0b\x32\x0b.ax.Content\"\xc9\x01\n\x11\x43onversationEvent\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0c\n\x04step\x18\x02 \x01(\x05\x12\x0f\n\x07\x65xec_id\x18\x03 \x01(\t\x12\x12\n\nharness_id\x18\x04 \x01(\t\x12/\n\x0eharness_config\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x1d\n\x08messages\x18\x06 \x03(\x0b\x32\x0b.ax.Message\x12\x18\n\x05state\x18\x07 \x01(\x0e\x32\t.ax.State\"E\n\x0cHarnessStart\x12\x16\n\x0eharness_config\x18\x01 \x01(\x0c\x12\x1d\n\x08messages\x18\x02 \x03(\x0b\x32\x0b.ax.Message\"1\n\rHarnessCancel\x12 \n\x06reason\x18\x01 \x01(\x0e\x32\x10.ax.CancelReason\"\x8d\x01\n\x0eHarnessRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x12\n\nharness_id\x18\x02 \x01(\t\x12!\n\x05start\x18\x03 \x01(\x0b\x32\x10.ax.HarnessStartH\x00\x12#\n\x06\x63\x61ncel\x18\x04 \x01(\x0b\x32\x11.ax.HarnessCancelH\x00\x42\x06\n\x04type\"/\n\x0eHarnessOutputs\x12\x1d\n\x08messages\x18\x01 \x03(\x0b\x32\x0b.ax.Message\"*\n\x05\x45rror\x12\x0c\n\x04\x63ode\x18\x01 \x01(\x05\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\"@\n\nHarnessEnd\x12\x18\n\x05state\x18\x01 \x01(\x0e\x32\t.ax.State\x12\x18\n\x05\x65rror\x18\x02 \x01(\x0b\x32\t.ax.Error\"x\n\x0fHarnessResponse\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12%\n\x07outputs\x18\x02 \x01(\x0b\x32\x12.ax.HarnessOutputsH\x00\x12\x1d\n\x03\x65nd\x18\x03 \x01(\x0b\x32\x0e.ax.HarnessEndH\x00\x42\x06\n\x04type\"\x82\x01\n\x0b\x45xecRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x1b\n\x06inputs\x18\x02 \x03(\x0b\x32\x0b.ax.Message\x12\x11\n\tlast_step\x18\x03 \x01(\x05\x12\x12\n\nharness_id\x18\x04 \x01(\t\x12\x16\n\x0eharness_config\x18\x05 \x01(\x0c\":\n\x0c\x45xecResponse\x12\x1c\n\x07outputs\x18\x01 \x03(\x0b\x32\x0b.ax.Message\x12\x0c\n\x04step\x18\x02 \x01(\x05\"4\n\x19\x44\x65leteConversationRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\"\x1c\n\x1a\x44\x65leteConversationResponse*l\n\x05State\x12\x15\n\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n\rSTATE_PENDING\x10\x01\x12\x10\n\x0cSTATE_FAILED\x10\x02\x12\x13\n\x0fSTATE_COMPLETED\x10\x03\x12\x12\n\x0eSTATE_CANCELED\x10\x04*\x8c\x01\n\x0c\x43\x61ncelReason\x12\x1d\n\x19\x43\x41NCEL_REASON_UNSPECIFIED\x10\x00\x12 \n\x1c\x43\x41NCEL_REASON_USER_REQUESTED\x10\x01\x12\x19\n\x15\x43\x41NCEL_REASON_TIMEOUT\x10\x02\x12 \n\x1c\x43\x41NCEL_REASON_INTERNAL_ERROR\x10\x03\x32H\n\x0eHarnessService\x12\x36\n\x07\x43onnect\x12\x12.ax.HarnessRequest\x1a\x13.ax.HarnessResponse(\x01\x30\x01\x32?\n\x10\x45xecutionService\x12+\n\x04\x45xec\x12\x0f.ax.ExecRequest\x1a\x10.ax.ExecResponse0\x01\x32j\n\x13\x43onversationService\x12S\n\x12\x44\x65leteConversation\x12\x1d.ax.DeleteConversationRequest\x1a\x1e.ax.DeleteConversationResponseB\x1cZ\x1agithub.com/google/ax/protob\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0eproto/ax.proto\x12\x02\x61x\x1a\x1cgoogle/protobuf/struct.proto\x1a\x13proto/content.proto\"5\n\x07Message\x12\x0c\n\x04role\x18\x01 \x01(\t\x12\x1c\n\x07\x63ontent\x18\x02 \x01(\x0b\x32\x0b.ax.Content\"\xe3\x01\n\x11\x43onversationEvent\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x0c\n\x04step\x18\x02 \x01(\x05\x12\x0f\n\x07\x65xec_id\x18\x03 \x01(\t\x12\x12\n\nharness_id\x18\x04 \x01(\t\x12/\n\x0eharness_config\x18\x05 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x1d\n\x08messages\x18\x06 \x03(\x0b\x32\x0b.ax.Message\x12\x18\n\x05state\x18\x07 \x01(\x0e\x32\t.ax.State\x12\x18\n\x10harness_metadata\x18\x08 \x01(\x0c\"E\n\x0cHarnessStart\x12\x16\n\x0eharness_config\x18\x01 \x01(\x0c\x12\x1d\n\x08messages\x18\x02 \x03(\x0b\x32\x0b.ax.Message\"1\n\rHarnessCancel\x12 \n\x06reason\x18\x01 \x01(\x0e\x32\x10.ax.CancelReason\"\x8d\x01\n\x0eHarnessRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x12\n\nharness_id\x18\x02 \x01(\t\x12!\n\x05start\x18\x03 \x01(\x0b\x32\x10.ax.HarnessStartH\x00\x12#\n\x06\x63\x61ncel\x18\x04 \x01(\x0b\x32\x11.ax.HarnessCancelH\x00\x42\x06\n\x04type\"/\n\x0eHarnessOutputs\x12\x1d\n\x08messages\x18\x01 \x03(\x0b\x32\x0b.ax.Message\"*\n\x05\x45rror\x12\x0c\n\x04\x63ode\x18\x01 \x01(\x05\x12\x13\n\x0b\x64\x65scription\x18\x02 \x01(\t\"Z\n\nHarnessEnd\x12\x18\n\x05state\x18\x01 \x01(\x0e\x32\t.ax.State\x12\x18\n\x05\x65rror\x18\x02 \x01(\x0b\x32\t.ax.Error\x12\x18\n\x10harness_metadata\x18\x03 \x01(\x0c\"x\n\x0fHarnessResponse\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12%\n\x07outputs\x18\x02 \x01(\x0b\x32\x12.ax.HarnessOutputsH\x00\x12\x1d\n\x03\x65nd\x18\x03 \x01(\x0b\x32\x0e.ax.HarnessEndH\x00\x42\x06\n\x04type\"\x82\x01\n\x0b\x45xecRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\x12\x1b\n\x06inputs\x18\x02 \x03(\x0b\x32\x0b.ax.Message\x12\x11\n\tlast_step\x18\x03 \x01(\x05\x12\x12\n\nharness_id\x18\x04 \x01(\t\x12\x16\n\x0eharness_config\x18\x05 \x01(\x0c\"T\n\x0c\x45xecResponse\x12\x1c\n\x07outputs\x18\x01 \x03(\x0b\x32\x0b.ax.Message\x12\x0c\n\x04step\x18\x02 \x01(\x05\x12\x18\n\x10harness_metadata\x18\x03 \x01(\x0c\"4\n\x19\x44\x65leteConversationRequest\x12\x17\n\x0f\x63onversation_id\x18\x01 \x01(\t\"\x1c\n\x1a\x44\x65leteConversationResponse*l\n\x05State\x12\x15\n\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n\rSTATE_PENDING\x10\x01\x12\x10\n\x0cSTATE_FAILED\x10\x02\x12\x13\n\x0fSTATE_COMPLETED\x10\x03\x12\x12\n\x0eSTATE_CANCELED\x10\x04*\x8c\x01\n\x0c\x43\x61ncelReason\x12\x1d\n\x19\x43\x41NCEL_REASON_UNSPECIFIED\x10\x00\x12 \n\x1c\x43\x41NCEL_REASON_USER_REQUESTED\x10\x01\x12\x19\n\x15\x43\x41NCEL_REASON_TIMEOUT\x10\x02\x12 \n\x1c\x43\x41NCEL_REASON_INTERNAL_ERROR\x10\x03\x32H\n\x0eHarnessService\x12\x36\n\x07\x43onnect\x12\x12.ax.HarnessRequest\x1a\x13.ax.HarnessResponse(\x01\x30\x01\x32?\n\x10\x45xecutionService\x12+\n\x04\x45xec\x12\x0f.ax.ExecRequest\x1a\x10.ax.ExecResponse0\x01\x32j\n\x13\x43onversationService\x12S\n\x12\x44\x65leteConversation\x12\x1d.ax.DeleteConversationRequest\x1a\x1e.ax.DeleteConversationResponseB\x1cZ\x1agithub.com/google/ax/protob\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -38,40 +38,40 @@ if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'Z\032github.com/google/ax/proto' - _globals['_STATE']._serialized_start=1156 - _globals['_STATE']._serialized_end=1264 - _globals['_CANCELREASON']._serialized_start=1267 - _globals['_CANCELREASON']._serialized_end=1407 + _globals['_STATE']._serialized_start=1234 + _globals['_STATE']._serialized_end=1342 + _globals['_CANCELREASON']._serialized_start=1345 + _globals['_CANCELREASON']._serialized_end=1485 _globals['_MESSAGE']._serialized_start=73 _globals['_MESSAGE']._serialized_end=126 _globals['_CONVERSATIONEVENT']._serialized_start=129 - _globals['_CONVERSATIONEVENT']._serialized_end=330 - _globals['_HARNESSSTART']._serialized_start=332 - _globals['_HARNESSSTART']._serialized_end=401 - _globals['_HARNESSCANCEL']._serialized_start=403 - _globals['_HARNESSCANCEL']._serialized_end=452 - _globals['_HARNESSREQUEST']._serialized_start=455 - _globals['_HARNESSREQUEST']._serialized_end=596 - _globals['_HARNESSOUTPUTS']._serialized_start=598 - _globals['_HARNESSOUTPUTS']._serialized_end=645 - _globals['_ERROR']._serialized_start=647 - _globals['_ERROR']._serialized_end=689 - _globals['_HARNESSEND']._serialized_start=691 - _globals['_HARNESSEND']._serialized_end=755 - _globals['_HARNESSRESPONSE']._serialized_start=757 - _globals['_HARNESSRESPONSE']._serialized_end=877 - _globals['_EXECREQUEST']._serialized_start=880 - _globals['_EXECREQUEST']._serialized_end=1010 - _globals['_EXECRESPONSE']._serialized_start=1012 - _globals['_EXECRESPONSE']._serialized_end=1070 - _globals['_DELETECONVERSATIONREQUEST']._serialized_start=1072 - _globals['_DELETECONVERSATIONREQUEST']._serialized_end=1124 - _globals['_DELETECONVERSATIONRESPONSE']._serialized_start=1126 - _globals['_DELETECONVERSATIONRESPONSE']._serialized_end=1154 - _globals['_HARNESSSERVICE']._serialized_start=1409 - _globals['_HARNESSSERVICE']._serialized_end=1481 - _globals['_EXECUTIONSERVICE']._serialized_start=1483 - _globals['_EXECUTIONSERVICE']._serialized_end=1546 - _globals['_CONVERSATIONSERVICE']._serialized_start=1548 - _globals['_CONVERSATIONSERVICE']._serialized_end=1654 + _globals['_CONVERSATIONEVENT']._serialized_end=356 + _globals['_HARNESSSTART']._serialized_start=358 + _globals['_HARNESSSTART']._serialized_end=427 + _globals['_HARNESSCANCEL']._serialized_start=429 + _globals['_HARNESSCANCEL']._serialized_end=478 + _globals['_HARNESSREQUEST']._serialized_start=481 + _globals['_HARNESSREQUEST']._serialized_end=622 + _globals['_HARNESSOUTPUTS']._serialized_start=624 + _globals['_HARNESSOUTPUTS']._serialized_end=671 + _globals['_ERROR']._serialized_start=673 + _globals['_ERROR']._serialized_end=715 + _globals['_HARNESSEND']._serialized_start=717 + _globals['_HARNESSEND']._serialized_end=807 + _globals['_HARNESSRESPONSE']._serialized_start=809 + _globals['_HARNESSRESPONSE']._serialized_end=929 + _globals['_EXECREQUEST']._serialized_start=932 + _globals['_EXECREQUEST']._serialized_end=1062 + _globals['_EXECRESPONSE']._serialized_start=1064 + _globals['_EXECRESPONSE']._serialized_end=1148 + _globals['_DELETECONVERSATIONREQUEST']._serialized_start=1150 + _globals['_DELETECONVERSATIONREQUEST']._serialized_end=1202 + _globals['_DELETECONVERSATIONRESPONSE']._serialized_start=1204 + _globals['_DELETECONVERSATIONRESPONSE']._serialized_end=1232 + _globals['_HARNESSSERVICE']._serialized_start=1487 + _globals['_HARNESSSERVICE']._serialized_end=1559 + _globals['_EXECUTIONSERVICE']._serialized_start=1561 + _globals['_EXECUTIONSERVICE']._serialized_end=1624 + _globals['_CONVERSATIONSERVICE']._serialized_start=1626 + _globals['_CONVERSATIONSERVICE']._serialized_end=1732 # @@protoc_insertion_point(module_scope)