From 98f80ffd6e2c0dd792fa4606100c9724baf4b79f Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 09:51:23 +0000 Subject: [PATCH 1/4] Bind Sessions through fenced Runtime assignments Every Session frame between Core and a Runtime carries an assignment reference. Core binds the Session before its first operation, records a release intent with each deletion, Environment release and archive, and reconciles pending releases after a reconnect. The Runtime fences stale, foreign and mismatched frames and settles a Session's work before it acknowledges a release. prompt_request had no producer and is deleted. --- .../agent/codex/preparation_router_test.go | 2 + .../internal/agenthost/admit_linux_test.go | 24 ++ .../agenthost/agenthost_linux_test.go | 66 ++- .../internal/agenthost/view_linux_test.go | 2 +- .../cli/claude_sdk_live_linux_test.go | 30 +- apps/daemon/internal/cli/connect.go | 1 + .../internal/cli/connect_cleanup_test.go | 28 ++ apps/daemon/internal/cli/connect_suspend.go | 66 +-- .../internal/cli/connect_suspend_test.go | 41 +- apps/daemon/internal/dispatch/assignment.go | 164 ++++++++ .../internal/dispatch/assignment_test.go | 170 ++++++++ apps/daemon/internal/dispatch/cancellation.go | 57 +-- .../internal/dispatch/cancellation_test.go | 139 ------- .../dispatch/capability_admission_test.go | 22 +- .../internal/dispatch/durable_only_test.go | 1 + .../internal/dispatch/environment_test.go | 68 +-- apps/daemon/internal/dispatch/executor.go | 28 +- .../dispatch/executor_cancel_receipt_test.go | 1 + .../dispatch/executor_handoff_test.go | 2 + .../daemon/internal/dispatch/executor_test.go | 12 +- apps/daemon/internal/dispatch/functions.go | 22 +- .../dispatch/functions_native_test.go | 193 --------- .../internal/dispatch/functions_test.go | 58 +-- .../dispatch/interaction_decisions.go | 66 +-- .../internal/dispatch/local_directory_test.go | 15 +- .../daemon/internal/dispatch/mcp_http_test.go | 43 +- .../dispatch/optional_interactions_test.go | 45 -- apps/daemon/internal/dispatch/output.go | 103 ----- apps/daemon/internal/dispatch/preparation.go | 42 +- .../internal/dispatch/preparation_start.go | 18 +- .../internal/dispatch/preparation_test.go | 7 +- .../internal/dispatch/prepared_handoff.go | 2 +- apps/daemon/internal/dispatch/prompt.go | 148 ------- .../internal/dispatch/receipt_order_test.go | 175 -------- .../dispatch/receipt_shutdown_test.go | 80 ---- apps/daemon/internal/dispatch/router.go | 162 ++------ apps/daemon/internal/dispatch/router_test.go | 388 +++--------------- .../internal/dispatch/runtime_preparation.go | 32 +- .../dispatch/runtime_preparation_test.go | 15 +- apps/daemon/internal/dispatch/shutdown.go | 47 +-- apps/daemon/internal/dispatch/steering.go | 33 +- .../dispatch/steering_lifetime_test.go | 58 +-- .../daemon/internal/dispatch/steering_test.go | 62 ++- apps/daemon/internal/dispatch/suspend.go | 43 +- apps/daemon/internal/dispatch/suspend_test.go | 61 ++- .../dispatch/workspace_directory_test.go | 2 +- .../internal/dispatch/workspace_export.go | 31 +- .../dispatch/workspace_export_test.go | 6 +- .../internal/dispatch/workspace_read.go | 26 +- .../internal/dispatch/workspace_write.go | 38 +- .../internal/dispatch/workspace_write_test.go | 3 +- apps/daemon/internal/transport/ws_test.go | 4 +- apps/daemon/testdata/onboarding/main.go | 5 +- docs/architecture.md | 2 +- docs/runtime-protocol.md | 31 +- docs/zh/architecture.md | 4 +- docs/zh/runtime-protocol.md | 33 +- internal/agentdaemon/proto/assignment.go | 76 ++++ internal/agentdaemon/proto/envelope.go | 21 +- internal/agentdaemon/proto/envelope_test.go | 8 +- internal/agentdaemon/proto/inbound.go | 2 + internal/agentdaemon/proto/outbound.go | 12 +- internal/agentdaemon/proto/prototest/wire.go | 20 +- .../agentdaemon/proto/prototest/wire_test.go | 5 + internal/agentdaemon/proto/runtime_prepare.go | 2 +- internal/harnessconfig/preparation_test.go | 2 +- .../core/internal/db/queries/core_metrics.sql | 4 +- services/core/internal/db/queries/devices.sql | 40 +- .../db/queries/environment_initialization.sql | 4 +- .../core/internal/db/queries/scheduling.sql | 4 +- .../db/queries/subagent_identities.sql | 4 +- .../core/internal/db/sqlc/core_metrics.sql.go | 4 +- services/core/internal/db/sqlc/devices.sql.go | 121 +++++- .../db/sqlc/environment_initialization.sql.go | 12 +- services/core/internal/db/sqlc/models.go | 17 +- .../core/internal/db/sqlc/scheduling.sql.go | 4 +- .../db/sqlc/subagent_identities.sql.go | 4 +- services/core/internal/deployment/storage.go | 8 +- services/core/internal/execution/artifacts.go | 2 +- .../internal/execution/assignment_releases.go | 80 ++++ services/core/internal/execution/delivery.go | 30 +- .../internal/execution/device_authority.go | 17 + .../execution/directory_preparation.go | 10 +- .../core/internal/execution/dispatcher.go | 4 + .../execution/environment_directory.go | 8 +- .../execution/environment_file_write.go | 4 +- .../execution/executor_preparation.go | 4 +- services/core/internal/execution/functions.go | 2 + .../core/internal/execution/preparation.go | 9 +- .../internal/execution/prepared_dispatch.go | 8 +- .../execution/runtime_capabilities_test.go | 8 +- .../internal/execution/runtime_compute.go | 20 +- .../execution/runtime_compute_wake.go | 8 +- .../execution/runtime_initialization.go | 6 +- .../core/internal/execution/runtime_setup.go | 10 +- .../internal/execution/runtime_setup_test.go | 8 +- services/core/internal/execution/worker.go | 5 + .../postgres/coremetricspg/store_test.go | 4 +- .../postgres/deploymentpg/allocations.go | 3 + .../postgres/deploymentpg/session_archive.go | 3 + .../persistence/postgres/sessionpg/devices.go | 26 +- .../postgres/sessionpg/devices_test.go | 2 +- .../postgres/sessionpg/environment.go | 3 +- .../sessionpg/execution_environment.go | 42 ++ .../sessionpg/execution_journal_test.go | 2 +- .../postgres/sessionpg/session_test.go | 2 +- .../postgres/sessionpg/session_writes.go | 3 + .../internal/runtimegateway/assignment.go | 117 ++++++ .../runtimegateway/cancellation_test.go | 2 +- .../runtimegateway/capabilities_test.go | 4 +- .../internal/runtimegateway/functions_test.go | 4 +- .../mcp_bearer_live_linux_test.go | 75 +++- .../core/internal/runtimegateway/mcp_test.go | 8 +- .../runtimegateway/preparation_test.go | 2 +- .../runtimegateway/runtime_prepare.go | 3 +- .../runtimegateway/runtime_prepare_test.go | 11 +- .../core/internal/runtimegateway/session.go | 53 ++- .../internal/runtimegateway/session_test.go | 13 +- .../internal/runtimegateway/suspension.go | 4 +- .../runtimegateway/suspension_test.go | 6 +- .../core/internal/runtimegateway/wire_test.go | 33 ++ .../workspace_directory_test.go | 4 +- .../runtimegateway/workspace_export.go | 7 +- .../runtimegateway/workspace_export_test.go | 6 +- .../internal/runtimegateway/workspace_read.go | 11 +- .../runtimegateway/workspace_read_test.go | 22 +- .../runtimegateway/workspace_write.go | 3 +- .../runtimegateway/workspace_write_test.go | 4 +- .../core/internal/sessions/environment.go | 3 + .../sessions/execution_environment.go | 35 ++ .../sessions/execution_environment_test.go | 12 + services/core/internal/sessions/session.go | 7 +- services/core/internal/sessions/turn.go | 4 + .../000092_session_runtime_assignments.sql | 38 ++ .../integration/archive_cancellation_test.go | 10 +- .../integration/assignment_release_test.go | 66 +++ .../integration/claude_execution_test.go | 2 +- .../core/tests/integration/claude_mcp_test.go | 6 +- .../deployment_model_providers_http_test.go | 2 +- .../core/tests/integration/devices_test.go | 2 +- .../core/tests/integration/dispatch_test.go | 7 +- .../integration/environment_directory_test.go | 2 +- .../environment_expiry_dispatch_test.go | 2 +- .../environment_runtime_fixture_test.go | 2 +- .../integration/environment_work_test.go | 2 +- .../environment_worker_helpers_test.go | 5 +- .../integration/execution_messages_test.go | 2 +- .../integration/executor_fixture_test.go | 24 ++ .../integration/function_execution_test.go | 2 +- .../function_input_execution_test.go | 3 + .../tests/integration/function_worker_test.go | 2 +- .../integration/harness_onboarding_test.go | 2 +- .../local_environment_devices_test.go | 2 +- .../local_environment_file_write_test.go | 9 +- .../local_environment_worker_test.go | 3 +- .../integration/public_execution_test.go | 2 +- .../runtime_compute_lifecycle_test.go | 8 +- .../runtime_initialization_peer_test.go | 8 +- .../structured_output_dispatch_test.go | 2 +- .../integration/subagent_identities_test.go | 2 +- .../tests/integration/worker_capacity_test.go | 3 + .../integration/worker_lease_loss_test.go | 2 +- 162 files changed, 2256 insertions(+), 2153 deletions(-) create mode 100644 apps/daemon/internal/dispatch/assignment.go create mode 100644 apps/daemon/internal/dispatch/assignment_test.go delete mode 100644 apps/daemon/internal/dispatch/cancellation_test.go delete mode 100644 apps/daemon/internal/dispatch/functions_native_test.go delete mode 100644 apps/daemon/internal/dispatch/optional_interactions_test.go delete mode 100644 apps/daemon/internal/dispatch/output.go delete mode 100644 apps/daemon/internal/dispatch/prompt.go delete mode 100644 apps/daemon/internal/dispatch/receipt_order_test.go delete mode 100644 apps/daemon/internal/dispatch/receipt_shutdown_test.go create mode 100644 internal/agentdaemon/proto/assignment.go create mode 100644 services/core/internal/execution/assignment_releases.go create mode 100644 services/core/internal/runtimegateway/assignment.go create mode 100644 services/core/migrations/000092_session_runtime_assignments.sql create mode 100644 services/core/tests/integration/assignment_release_test.go diff --git a/apps/daemon/internal/agent/codex/preparation_router_test.go b/apps/daemon/internal/agent/codex/preparation_router_test.go index 62d31d24c..2ba23bef6 100644 --- a/apps/daemon/internal/agent/codex/preparation_router_test.go +++ b/apps/daemon/internal/agent/codex/preparation_router_test.go @@ -88,10 +88,12 @@ func TestPreparationRouterRetainsActualNativeChild(t *testing.T) { if err != nil { t.Fatal(err) } + env.Assignment = proto.AssignmentRef{SessionID: session, AssignmentID: "assignment", Epoch: 1} if err = r.Handle(t.Context(), env); err != nil { t.Fatal(err) } } + send(proto.TypeAssignmentBind, "bind", proto.AssignmentBindPayload{EnvironmentID: environment}) await := func(state string) proto.PreparationStatusPayload { t.Helper() timer := time.NewTimer(4 * time.Second) diff --git a/apps/daemon/internal/agenthost/admit_linux_test.go b/apps/daemon/internal/agenthost/admit_linux_test.go index 054085479..a6d6f7d3b 100644 --- a/apps/daemon/internal/agenthost/admit_linux_test.go +++ b/apps/daemon/internal/agenthost/admit_linux_test.go @@ -251,5 +251,29 @@ func TestViewExecutorReceivesTheGatewayRequest(t *testing.T) { if _, err := os.Stat(f.session.Home.Host); dials.Load() != 0 || err != nil || len(leftEntries(t, f.cfg)) != 0 { t.Errorf("%d dials, home %v and transient entries %v after the preparation", dials.Load(), err, leftEntries(t, f.cfg)) } +} +// TestReleaseRemovesTheHome checks that the agent host removes a Session's +// home behind its assignment's release. +func TestReleaseRemovesTheHome(t *testing.T) { + f := newViewFixture(t) + var dials atomic.Int32 + d := newDaemon(t, f.cfg, deps{dial: countingDial(&dials), tasks: noTasks}) + b := newBinding(newResource()) + if _, p := d.prepare(t, b, request("viewed", "/workspace", "https://model.test", "sk-test")); p.State != "failed" { + t.Fatalf("the preparation is %s, want failed with the factory", p.State) + } + if _, err := os.Stat(f.session.Home.Host); err != nil { + t.Fatalf("the home: %v", err) + } + released := ref(b) + released.Epoch++ + id := "release" + d.handle(t, released, proto.TypeAssignmentRelease, id, proto.AssignmentReleasePayload{RemoveHome: true}) + if status := d.status(t, id); status.State != proto.AssignmentHomeRemoved { + t.Fatalf("the release is %s (%s), want home_removed", status.State, status.ErrorCode) + } + if _, err := os.Stat(filepath.Dir(f.session.Home.Host)); !errors.Is(err, os.ErrNotExist) { + t.Errorf("the Session directory remains: %v", err) + } } diff --git a/apps/daemon/internal/agenthost/agenthost_linux_test.go b/apps/daemon/internal/agenthost/agenthost_linux_test.go index d0d47654e..7316189c3 100644 --- a/apps/daemon/internal/agenthost/agenthost_linux_test.go +++ b/apps/daemon/internal/agenthost/agenthost_linux_test.go @@ -144,8 +144,9 @@ func leftEntries(t *testing.T, cfg Config) []string { } // daemon drives Sessions through a dispatch Router, as the daemon does. It -// binds each request to the Session its state key names and records the -// latest Executor the agent host opened for each Session. +// binds each request to the Session of the assignment the Router admitted it +// under, records the latest Executor the agent host opened for each Session, +// and removes a released Session's home. type daemon struct { router *dispatch.Router // mcp is the installed MCP that the Environment's preparation resolves @@ -175,7 +176,16 @@ func newDaemon(t *testing.T, cfg Config, d deps) *daemon { return e, err }) var err error - if dm.router, err = dispatch.New(dispatch.Config{Registry: reg, Sender: dm, SessionEnvironments: true, Log: slog.New(slog.DiscardHandler)}); err != nil { + removeHome := func(session string) error { + dm.mu.Lock() + b, ok := dm.bindings[session] + dm.mu.Unlock() + if !ok { + return fmt.Errorf("%w: no binding", ErrInvalidSession) + } + return (&Host{cfg: cfg}).RemoveHome(b.SessionID) + } + if dm.router, err = dispatch.New(dispatch.Config{Registry: reg, Sender: dm, SessionEnvironments: true, RemoveHome: removeHome, Log: slog.New(slog.DiscardHandler)}); err != nil { t.Fatal(err) } t.Cleanup(func() { dm.shutdown() }) @@ -188,13 +198,18 @@ const stateKeyPrefix = "agents-api-" func (dm *daemon) bind(req proto.PromptRequestPayload) (Binding, Environment, error) { dm.mu.Lock() defer dm.mu.Unlock() - b, ok := dm.bindings[strings.TrimPrefix(req.AgentStateKey, stateKeyPrefix)] - if !ok { + b, ok := dm.bindings[req.Assignment.SessionID] + if !ok || ref(b) != req.Assignment { return Binding{}, Environment{}, fmt.Errorf("%w: no binding", ErrInvalidSession) } return b, Environment{}, nil } +// ref is the reference of b's assignment. +func ref(b Binding) proto.AssignmentRef { + return proto.AssignmentRef{SessionID: b.SessionID.String(), AssignmentID: b.AssignmentID.String(), Epoch: b.AssignmentEpoch} +} + func (dm *daemon) Send(_ context.Context, e proto.Envelope) error { dm.frame(e.ID) <- e return nil @@ -211,13 +226,14 @@ func (dm *daemon) frame(id string) chan proto.Envelope { return ch } -// handle hands the Router an envelope from Core. -func (dm *daemon) handle(t *testing.T, typ, id string, payload any) { +// handle hands the Router an envelope from Core under the assignment ref. +func (dm *daemon) handle(t *testing.T, ref proto.AssignmentRef, typ, id string, payload any) { t.Helper() e, err := proto.NewEnvelope(typ, id, payload) if err != nil { t.Fatal(err) } + e.Assignment = ref if err := dm.router.Handle(context.Background(), e); err != nil { t.Fatalf("%s: %v", typ, err) } @@ -235,17 +251,39 @@ func (dm *daemon) next(t *testing.T, id string) proto.Envelope { } } -// prepare prepares an Executor of b's Session for req. It returns the -// request ID and the preparation's first status other than preparing. -func (dm *daemon) prepare(t *testing.T, b Binding, req proto.PromptRequestPayload) (string, proto.PreparationStatusPayload) { +// assign binds b's Session to the Router in environment. +func (dm *daemon) assign(t *testing.T, b Binding, environment string) { t.Helper() - session := b.SessionID.String() dm.mu.Lock() - dm.bindings[session] = b + dm.bindings[b.SessionID.String()] = b dm.mu.Unlock() + id := sandboxwire.NewID().String() + dm.handle(t, ref(b), proto.TypeAssignmentBind, id, proto.AssignmentBindPayload{EnvironmentID: environment}) + if status := dm.status(t, id); status.State != proto.AssignmentBound { + t.Fatalf("the bind is %s (%s), want bound", status.State, status.ErrorCode) + } +} + +// status returns the assignment status sent with id. +func (dm *daemon) status(t *testing.T, id string) proto.AssignmentStatusPayload { + t.Helper() + var status proto.AssignmentStatusPayload + if err := dm.next(t, id).DecodePayload(&status); err != nil { + t.Fatal(err) + } + return status +} + +// prepare binds b's Session and prepares an Executor of it for req. It +// returns the request ID and the preparation's first status other than +// preparing. +func (dm *daemon) prepare(t *testing.T, b Binding, req proto.PromptRequestPayload) (string, proto.PreparationStatusPayload) { + t.Helper() + dm.assign(t, b, req.EnvironmentID()) + session := b.SessionID.String() req.AgentStateKey = stateKeyPrefix + session id := sandboxwire.NewID().String() - dm.handle(t, proto.TypeExecutionPrepare, id, proto.ExecutionPreparePayload{SessionID: session, Configuration: req}) + dm.handle(t, ref(b), proto.TypeExecutionPrepare, id, proto.ExecutionPreparePayload{SessionID: session, Configuration: req}) for { var p proto.PreparationStatusPayload if err := dm.next(t, id).DecodePayload(&p); err != nil { @@ -266,7 +304,7 @@ func (dm *daemon) start(t *testing.T, b Binding, req proto.PromptRequestPayload, t.Fatalf("the preparation is %s (%s), want ready", p.State, p.ErrorCode) } run := "run-" + id - dm.handle(t, proto.TypeExecutionStart, id, proto.ExecutionStartPayload{Handle: p.Handle, ExecutorID: p.ExecutorID, RunID: run, Input: proto.TextInput(text)}) + dm.handle(t, ref(b), proto.TypeExecutionStart, id, proto.ExecutionStartPayload{Handle: p.Handle, ExecutorID: p.ExecutorID, RunID: run, Input: proto.TextInput(text)}) for { if err := dm.next(t, id).DecodePayload(&p); err != nil { t.Fatal(err) diff --git a/apps/daemon/internal/agenthost/view_linux_test.go b/apps/daemon/internal/agenthost/view_linux_test.go index d08db8c6b..f4bac56cb 100644 --- a/apps/daemon/internal/agenthost/view_linux_test.go +++ b/apps/daemon/internal/agenthost/view_linux_test.go @@ -278,7 +278,7 @@ func TestSessionRunsInAViewOverItsAttachment(t *testing.T) { t.Errorf("the other Session's command printed %q and exited %d; stderr %s", r.Stdout, r.Code, r.Stderr) } // Cancelling the waiting Turn ends it with its view. - d.handle(t, proto.TypePromptCancel, run, proto.PromptCancelPayload{}) + d.handle(t, ref(waiting), proto.TypePromptCancel, run, proto.PromptCancelPayload{}) d.done(t, run) if err := d.shutdown(); err != nil { t.Fatalf("Shutdown = %v", err) diff --git a/apps/daemon/internal/cli/claude_sdk_live_linux_test.go b/apps/daemon/internal/cli/claude_sdk_live_linux_test.go index 445688eee..78217d94b 100644 --- a/apps/daemon/internal/cli/claude_sdk_live_linux_test.go +++ b/apps/daemon/internal/cli/claude_sdk_live_linux_test.go @@ -15,6 +15,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/dispatch" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" "github.com/google/uuid" ) @@ -68,6 +69,7 @@ func TestLiveRegisteredClaudeSDK(t *testing.T) { Cancelled bool `json:"cancelled"` } nonce := "registered-function-" + uuid.NewString() + ref := proto.AssignmentRef{SessionID: prototest.SessionID, AssignmentID: "registered-acceptance", Epoch: 1} run := func(index int, prompt, resume string, callFunction, cancelOnText bool) execution { t.Helper() reg := agent.NewRegistry() @@ -87,21 +89,24 @@ func TestLiveRegisteredClaudeSDK(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 120*time.Second) defer cancel() id := uuid.NewString() - request := proto.PromptRequestPayload{RunID: id, AgentKind: "claude_sdk", Input: proto.TextInput(prompt), AgentStateKey: "registered-acceptance", AgentSessionID: resume, StrictResume: true, ReleaseOnCompletion: true, ObserveMessages: true, ObserveToolObservations: true, DisableExecutionEnvironment: true, DisableSubagents: true, ExecutionControls: &proto.ExecutionControls{WebSearch: "disabled", TextVerbosity: "medium"}, Model: "MiniMax-M3"} + request := proto.PromptRequestPayload{AgentKind: "claude_sdk", AgentStateKey: prototest.StateKey, AgentSessionID: resume, RequireExistingNativeSession: resume != "", StrictResume: true, ReleaseOnCompletion: true, ObserveMessages: true, ObserveToolObservations: true, DisableExecutionEnvironment: true, DisableSubagents: true, ExecutionControls: &proto.ExecutionControls{WebSearch: "disabled", TextVerbosity: "medium"}, Model: "MiniMax-M3"} if callFunction { request.FunctionTools = []proto.FunctionTool{{Name: "lookup", Description: "Return a verification value.", Parameters: json.RawMessage(`{"type":"object","properties":{"id":{"type":"string"}},"required":["id"],"additionalProperties":false}`)}} } - handle := func(kind string, payload any) { + send := func(kind, envID string, payload any) { t.Helper() - env, err := proto.NewEnvelope(kind, id, payload) + env, err := proto.NewEnvelope(kind, envID, payload) if err != nil { t.Fatal(err) } + env.Assignment = ref if err := router.Handle(ctx, env); err != nil { t.Fatal("registered router request failed", err) } } - handle(proto.TypePromptRequest, request) + handle := func(kind string, payload any) { t.Helper(); send(kind, id, payload) } + send(proto.TypeAssignmentBind, "bind", proto.AssignmentBindPayload{}) + send(proto.TypeExecutionPrepare, "prepare", proto.ExecutionPreparePayload{SessionID: ref.SessionID, Configuration: request}) proof := execution{} defer func() { data, _ := json.MarshalIndent(proof, "", " ") @@ -117,6 +122,23 @@ func TestLiveRegisteredClaudeSDK(t *testing.T) { case <-ctx.Done(): t.Fatal("registered execution timed out", ctx.Err()) } + if event.Type == proto.TypeAssignmentStatus { + continue + } + if event.Type == proto.TypePreparationStatus { + var status proto.PreparationStatusPayload + if err := event.DecodePayload(&status); err != nil { + t.Fatal(err) + } + switch status.State { + case "ready": + send(proto.TypeExecutionStart, "prepare", proto.ExecutionStartPayload{Handle: status.Handle, ExecutorID: status.ExecutorID, RunID: id, Input: proto.TextInput(prompt)}) + case "started": + default: + t.Fatal("registered preparation failed", status) + } + continue + } if event.ID != id { t.Fatal("event identity changed") } diff --git a/apps/daemon/internal/cli/connect.go b/apps/daemon/internal/cli/connect.go index f8718358d..71c8fe2e5 100644 --- a/apps/daemon/internal/cli/connect.go +++ b/apps/daemon/internal/cli/connect.go @@ -340,6 +340,7 @@ func pumpConn(parentCtx context.Context, conn *transport.Conn, registry *agent.R ActiveRequests: router.ActiveRuns(), DaemonVersion: Version, SupportedAgentKinds: kinds, + HomeRemoval: proto.CapabilityUnsupported, } }, obslog.Bg().With("component", "heartbeat")) diff --git a/apps/daemon/internal/cli/connect_cleanup_test.go b/apps/daemon/internal/cli/connect_cleanup_test.go index 24cf58f4c..c4658752d 100644 --- a/apps/daemon/internal/cli/connect_cleanup_test.go +++ b/apps/daemon/internal/cli/connect_cleanup_test.go @@ -145,11 +145,16 @@ func testDisconnectedPumpCleanup(t *testing.T, suspend bool) { t.Fatal("initial connection missing") } defer peer.Close() + ref := proto.AssignmentRef{SessionID: "cleanup", AssignmentID: "assignment", Epoch: 1} + if got := sendAssignment(t, peer, proto.TypeAssignmentBind, ref, proto.AssignmentBindPayload{}); got.State != proto.AssignmentBound { + t.Fatalf("bind = %+v", got) + } env, err := proto.NewEnvelope(proto.TypeExecutionPrepare, "prepare", proto.ExecutionPreparePayload{SessionID: "cleanup", Configuration: prototest.WithModel(proto.PromptRequestPayload{AgentKind: "cleanup", AgentStateKey: "agents-api-cleanup", StrictResume: true, DisableExecutionEnvironment: true})}) if err != nil { t.Fatal(err) } + env.Assignment = ref if err := peer.WriteJSON(env); err != nil { t.Fatal(err) } @@ -195,3 +200,26 @@ func testDisconnectedPumpCleanup(t *testing.T, suspend bool) { t.Fatalf("factory=%d close=%d", factories.Load(), owner.closes.Load()) } } + +func sendAssignment(t *testing.T, peer *websocket.Conn, kind string, ref proto.AssignmentRef, payload any) proto.AssignmentStatusPayload { + t.Helper() + env, err := proto.NewEnvelope(kind, kind, payload) + if err != nil { + t.Fatal(err) + } + env.Assignment = ref + if err := peer.WriteJSON(env); err != nil { + t.Fatal(err) + } + _ = peer.SetReadDeadline(time.Now().Add(3 * time.Second)) + for { + var reply proto.Envelope + if err := peer.ReadJSON(&reply); err != nil { + t.Fatal(err) + } + var status proto.AssignmentStatusPayload + if reply.Type == proto.TypeAssignmentStatus && reply.ID == kind && reply.Assignment == ref && reply.DecodePayload(&status) == nil { + return status + } + } +} diff --git a/apps/daemon/internal/cli/connect_suspend.go b/apps/daemon/internal/cli/connect_suspend.go index 21068ca9c..2fc97dba7 100644 --- a/apps/daemon/internal/cli/connect_suspend.go +++ b/apps/daemon/internal/cli/connect_suspend.go @@ -157,7 +157,7 @@ func (s *suspendedRouter) heartbeats(ctx context.Context, conn *transport.Conn, caps := &kinds[i].Capabilities caps.WorkspaceOutputExport = proto.CapabilityFromBool(s.local.CanExport() && caps.LocalEnvironment.IsSupported() && caps.WorkspaceReadPreparation.IsSupported()) } - return proto.HeartbeatPayload{Timestamp: time.Now().Unix(), ActiveRequests: s.router.ActiveRuns(), DaemonVersion: Version, SupportedAgentKinds: kinds} + return proto.HeartbeatPayload{Timestamp: time.Now().Unix(), ActiveRequests: s.router.ActiveRuns(), DaemonVersion: Version, SupportedAgentKinds: kinds, HomeRemoval: proto.CapabilityUnsupported} }, obslog.Bg()) } @@ -176,7 +176,7 @@ func (s *suspendedRouter) pump(ctx context.Context, conn *transport.Conn, boot * if env.Type == proto.TypeEnvironmentResume { request := rejectedResumeRequest(env) code := "not_suspended" - valid := env.ID != "" && len(env.ID) <= 128 && request.EnvironmentID == control.identity.EnvironmentID && request.SuspendID != "" && len(request.SuspendID) <= 128 + valid := env.ID != "" && len(env.ID) <= 128 && env.Assignment.Valid() && request.EnvironmentID == control.identity.EnvironmentID && request.SuspendID != "" && len(request.SuspendID) <= 128 if valid && ((control.lastResumed != nil && control.lastResumed.SameSuspension(request)) || request.Rollback) { code = "" control.lastResumed = &request @@ -195,10 +195,15 @@ func (s *suspendedRouter) pump(ctx context.Context, conn *transport.Conn, boot * continue } quiet, cancel := context.WithTimeout(ctx, 5*time.Second) - err := s.router.Quiesce(quiet, request) + err := s.router.Quiesce(quiet, env.Assignment, request) cancel() if err != nil { - if sendErr := sendSuspendResult(ctx, conn, env, proto.TypeEnvironmentQuiesced, request, "resource_busy"); sendErr != nil { + code := "resource_busy" + var rejected dispatch.AssignmentError + if errors.As(err, &rejected) { + code = string(rejected) + } + if sendErr := sendSuspendResult(ctx, conn, env, proto.TypeEnvironmentQuiesced, request, code); sendErr != nil { return nil, sendErr } if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { @@ -221,37 +226,48 @@ func (s *suspendedRouter) pump(ctx context.Context, conn *transport.Conn, boot * } } +// resume waits for Core to confirm the suspension on conn. The quiesced Router +// replies on conn from the start, because Core may release an assignment +// before it resumes. func (s *suspendedRouter) resume(ctx context.Context, conn *transport.Conn, request proto.EnvironmentSuspendPayload, control *suspendControl) error { - select { - case <-ctx.Done(): - return ctx.Err() - case <-conn.Done(): - if err := conn.Err(); err != nil { - return err - } - return transport.ErrConnClosed - case env, ok := <-conn.Recv(): - if !ok { + s.sender.replace(conn) + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-conn.Done(): if err := conn.Err(); err != nil { return err } return transport.ErrConnClosed + case env, ok := <-conn.Recv(): + if !ok { + if err := conn.Err(); err != nil { + return err + } + return transport.ErrConnClosed + } + if env.Type == proto.TypeAssignmentRelease { + if err := s.router.Handle(ctx, env); err != nil { + obslog.Bg().Error("router.Handle failed", "type", env.Type, "err", err) + } + continue + } + var echoed proto.EnvironmentSuspendPayload + if env.Type != proto.TypeEnvironmentResume || env.ID == "" || env.DecodeRequest(&echoed) != nil || !echoed.SameSuspension(request) { + return errors.Join(transport.ErrPermanent, errors.New("connect: expected authenticated suspension resume")) + } + if err := s.router.Resume(env.Assignment, echoed, s.sender); err != nil { + return errors.Join(transport.ErrPermanent, err) + } + control.lastResumed = &request + return sendSuspendResult(ctx, conn, env, proto.TypeEnvironmentResumed, request, "") } - var echoed proto.EnvironmentSuspendPayload - if env.Type != proto.TypeEnvironmentResume || env.ID == "" || env.DecodeRequest(&echoed) != nil || !echoed.SameSuspension(request) { - return errors.Join(transport.ErrPermanent, errors.New("connect: expected authenticated suspension resume")) - } - s.sender.replace(conn) - if err := s.router.Resume(echoed, s.sender); err != nil { - return errors.Join(transport.ErrPermanent, err) - } - control.lastResumed = &request - return sendSuspendResult(ctx, conn, env, proto.TypeEnvironmentResumed, request, "") } } func sendSuspendResult(ctx context.Context, conn *transport.Conn, env proto.Envelope, kind string, request proto.EnvironmentSuspendPayload, code string) error { - result, err := proto.NewEnvelope(kind, env.ID, proto.EnvironmentSuspendResultPayload{EnvironmentID: request.EnvironmentID, SuspendID: request.SuspendID, Accepted: code == "", ErrorCode: code}) + result, err := env.Reply(kind, proto.EnvironmentSuspendResultPayload{EnvironmentID: request.EnvironmentID, SuspendID: request.SuspendID, Accepted: code == "", ErrorCode: code}) if err != nil { return err } diff --git a/apps/daemon/internal/cli/connect_suspend_test.go b/apps/daemon/internal/cli/connect_suspend_test.go index 4c7177746..7c46007b5 100644 --- a/apps/daemon/internal/cli/connect_suspend_test.go +++ b/apps/daemon/internal/cli/connect_suspend_test.go @@ -61,6 +61,9 @@ func TestPlannedReconnectRequiresAuthenticatedMatchingResume(t *testing.T) { t.Fatal("initial connection missing") } defer first.Close() + if got := sendAssignment(t, first, proto.TypeAssignmentBind, lifecycleRef, proto.AssignmentBindPayload{EnvironmentID: "env"}); got.State != proto.AssignmentBound { + t.Fatalf("bind = %+v", got) + } request := proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "planned"} sendLifecycleFrame(t, first, proto.TypeEnvironmentQuiesce, request) result := readLifecycleResult(t, first, proto.TypeEnvironmentQuiesced) @@ -105,6 +108,15 @@ func TestPlannedReconnectRequiresAuthenticatedMatchingResume(t *testing.T) { if scenario == "wrong_operation" { request.SuspendID = "stale" } + if scenario == "resume" { + // Core may release the Session before it resumes; the reply + // arrives on the new connection. + released := lifecycleRef + released.Epoch++ + if got := sendAssignment(t, second, proto.TypeAssignmentRelease, released, proto.AssignmentReleasePayload{}); got.State != proto.AssignmentReleased { + t.Fatalf("release during suspension = %+v", got) + } + } sendLifecycleFrame(t, second, proto.TypeEnvironmentResume, request) if scenario == "resume" || scenario == "rollback" || scenario == "reconnect" { result = readLifecycleResult(t, second, proto.TypeEnvironmentResumed) @@ -153,16 +165,21 @@ func TestPlannedReconnectRequiresAuthenticatedMatchingResume(t *testing.T) { } } +// lifecycleRef is the assignment of the Session whose Environment suspends. +var lifecycleRef = proto.AssignmentRef{SessionID: "session", AssignmentID: "assignment", Epoch: 1} + func sendLifecycleFrame(t *testing.T, peer *websocket.Conn, kind string, payload proto.EnvironmentSuspendPayload) { t.Helper() env, err := proto.NewEnvelope(kind, "operation", payload) if err != nil { t.Fatal(err) } + env.Assignment = lifecycleRef if err := peer.WriteJSON(env); err != nil { t.Fatal(err) } } + func readLifecycleResult(t *testing.T, peer *websocket.Conn, kind string) proto.EnvironmentSuspendResultPayload { t.Helper() _ = peer.SetReadDeadline(time.Now().Add(3 * time.Second)) @@ -238,10 +255,11 @@ func TestSuspensionReconnectBeforeConfirmation(t *testing.T) { t.Run(scenario, func(t *testing.T) { t.Setenv("OAC_RUNTIME_WORKSPACE", "") var attempts atomic.Int32 - peers := make(chan *websocket.Conn, 2) + peers := make(chan *websocket.Conn, 3) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { attempt := attempts.Add(1) - if scenario == "revoked" && attempt > 1 { + // The first connection binds and quiesces the Router. + if scenario == "revoked" && attempt > 2 { http.Error(w, "revoked", http.StatusUnauthorized) return } @@ -254,15 +272,30 @@ func TestSuspensionReconnectBeforeConfirmation(t *testing.T) { dial := func(ctx context.Context) (*transport.Conn, error) { return transport.Dial(ctx, transport.DialOptions{WSURL: "ws" + strings.TrimPrefix(server.URL, "http"), DeviceID: "device", Credential: "credential", DaemonVersion: proto.Version}) } - state, err := newSuspendedRouter(nil, agent.NewRegistry()) + conn, err := dial(t.Context()) + if err != nil { + t.Fatal(err) + } + setup := <-peers + defer setup.Close() + state, err := newSuspendedRouter(conn, agent.NewRegistry()) if err != nil { t.Fatal(err) } defer state.shutdown() + bind, err := proto.NewEnvelope(proto.TypeAssignmentBind, "bind", proto.AssignmentBindPayload{EnvironmentID: "env"}) + if err != nil { + t.Fatal(err) + } + bind.Assignment = lifecycleRef + if err := state.router.Handle(t.Context(), bind); err != nil { + t.Fatal(err) + } request := proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "pending-confirmation"} - if err := state.router.Quiesce(t.Context(), request); err != nil { + if err := state.router.Quiesce(t.Context(), lifecycleRef, request); err != nil { t.Fatal(err) } + _ = conn.Close() control := &suspendControl{path: filepath.Join(t.TempDir(), "control.json"), identity: suspendIdentity{EnvironmentID: "env"}, signal: make(chan os.Signal, 1)} if err := control.Arm(request); err != nil { t.Fatal(err) diff --git a/apps/daemon/internal/dispatch/assignment.go b/apps/daemon/internal/dispatch/assignment.go new file mode 100644 index 000000000..3e71dc1c4 --- /dev/null +++ b/apps/daemon/internal/dispatch/assignment.go @@ -0,0 +1,164 @@ +package dispatch + +import ( + "context" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" +) + +// assignmentState is the Router's record of one Session's assignment. Router.mu +// protects it. A released assignment stays recorded, so its frames stay fenced. +type assignmentState struct { + ref proto.AssignmentRef + environmentID string + released bool +} + +// admitLocked returns why ref admits no new work of sessionID in +// environmentID, or "". Router.mu must be held. +func (r *Router) admitLocked(ref proto.AssignmentRef, sessionID, environmentID string) string { + a := r.assignments[ref.SessionID] + switch { + case !ref.Valid() || a == nil || a.ref.AssignmentID != ref.AssignmentID: + return proto.AssignmentConflict + case ref.Epoch < a.ref.Epoch || ref.Epoch == a.ref.Epoch && a.released: + return proto.AssignmentStale + case ref.Epoch != a.ref.Epoch || sessionID != ref.SessionID || environmentID != a.environmentID: + return proto.AssignmentConflict + } + return "" +} + +// admitRunLocked admits a frame for the run, which ref must have started. +// Router.mu must be held. +func (r *Router) admitRunLocked(ref proto.AssignmentRef, state *sessionState) string { + if state.assignment != ref { + return proto.AssignmentConflict + } + return r.admitLocked(ref, ref.SessionID, state.environmentID) +} + +func (r *Router) handleAssignmentBind(ctx context.Context, env proto.Envelope) error { + var input proto.AssignmentBindPayload + ref, code := env.Assignment, "" + if env.ID == "" || env.DecodeRequest(&input) != nil || !ref.Valid() { + code = "invalid_request" + } else { + r.mu.Lock() + a := r.assignments[ref.SessionID] + switch { + case a == nil: + r.assignments[ref.SessionID] = &assignmentState{ref: ref, environmentID: input.EnvironmentID} + case a.ref.AssignmentID == ref.AssignmentID && (ref.Epoch < a.ref.Epoch || ref.Epoch == a.ref.Epoch && a.released): + code = proto.AssignmentStale + case a.ref != ref || a.environmentID != input.EnvironmentID: + code = proto.AssignmentConflict + } + r.mu.Unlock() + } + return r.reply(ctx, env, proto.TypeAssignmentStatus, assignmentStatus(proto.AssignmentBound, code)) +} + +// handleAssignmentRelease fences the assignment, then settles the Session's +// Executor and removes its home before it replies. A retry at the same epoch +// repeats the cleanup. +func (r *Router) handleAssignmentRelease(ctx context.Context, env proto.Envelope) error { + var input proto.AssignmentReleasePayload + ref, code := env.Assignment, "" + if env.ID == "" || env.DecodeRequest(&input) != nil || !ref.Valid() { + code = "invalid_request" + } else if input.RemoveHome && r.removeHome == nil { + code = proto.UnsupportedOperation + } + if code != "" { + return r.reply(ctx, env, proto.TypeAssignmentStatus, assignmentStatus("", code)) + } + r.mu.Lock() + a := r.assignments[ref.SessionID] + switch { + case a == nil: + r.assignments[ref.SessionID] = &assignmentState{ref: ref, released: true} + case a.ref.AssignmentID != ref.AssignmentID || ref.Epoch == a.ref.Epoch && !a.released: + code = proto.AssignmentConflict + case ref.Epoch < a.ref.Epoch: + code = proto.AssignmentStale + default: + a.ref, a.released = ref, true + } + if code != "" { + r.mu.Unlock() + return r.reply(ctx, env, proto.TypeAssignmentStatus, assignmentStatus("", code)) + } + r.shutdownWG.Add(1) + r.mu.Unlock() + go func() { + defer r.shutdownWG.Done() + state, code := proto.AssignmentReleased, "" + err := r.closeSessionExecutor(ref.SessionID) + if err == nil && input.RemoveHome { + state, err = proto.AssignmentHomeRemoved, r.removeHome(ref.SessionID) + } + if err != nil { + r.log.Warn("assignment release cleanup unconfirmed", "session_id", ref.SessionID, "err", err) + code = proto.CleanupUnconfirmed + } + sendCtx, stop := r.shutdownContext(context.WithoutCancel(ctx)) + defer stop() + _ = r.reply(sendCtx, env, proto.TypeAssignmentStatus, assignmentStatus(state, code)) + }() + return nil +} + +// closeSessionExecutor ends the Session's Executor: it abandons a pending +// admission, releases a Turn, waits for a native preparation and closes the +// Executor. The released assignment admits no new preparation meanwhile. +func (r *Router) closeSessionExecutor(sessionID string) error { + r.mu.Lock() + owner := r.executors[sessionID] + if owner == nil { + r.mu.Unlock() + return nil + } + owner.invalid = true + owner.closeReason = "assignment_released" + if owner.preparing { + owner.cancel() + } + var release *preparedRelease + var attempt *preparedReleaseAttempt + admission := owner.admission + if owner.run != nil { + release, attempt = r.claimPreparedReleaseLocked(owner.run, true, "", true) + admission = nil + } + r.mu.Unlock() + if admission != nil { + r.abandonExecutorAdmission(admission, "failed", proto.AssignmentStale, true) + } + if release != nil { + if err := r.awaitPreparedNativeRelease(context.Background(), release, attempt); err != nil { + return err + } + } + <-owner.prepared + return r.closeExecutor(owner) +} + +func assignmentStatus(state, code string) proto.AssignmentStatusPayload { + if code != "" { + state = proto.AssignmentFailed + } + return proto.AssignmentStatusPayload{State: state, ErrorCode: code} +} + +// reply sends the reply to the request env within the send budget. +func (r *Router) reply(ctx context.Context, env proto.Envelope, typ string, payload any) error { + reply, err := env.Reply(typ, payload) + if err != nil { + return err + } + ctx, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + return r.sender.Send(ctx, reply) +} diff --git a/apps/daemon/internal/dispatch/assignment_test.go b/apps/daemon/internal/dispatch/assignment_test.go new file mode 100644 index 000000000..665e99c7e --- /dev/null +++ b/apps/daemon/internal/dispatch/assignment_test.go @@ -0,0 +1,170 @@ +package dispatch_test + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" + "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/dispatch" + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" +) + +// frameFor returns the last frame of kind correlated with id. +func frameFor(sender *recSender, kind, id string) (proto.Envelope, bool) { + frames := sender.snapshot() + for i := len(frames) - 1; i >= 0; i-- { + if frames[i].Type == kind && frames[i].ID == id { + return frames[i], true + } + } + return proto.Envelope{}, false +} + +func waitAssignmentStatus(t *testing.T, sender *recSender, id string) proto.AssignmentStatusPayload { + t.Helper() + waitFor(t, func() bool { return hasFrame(sender, proto.TypeAssignmentStatus, id) }, "assignment_status "+id) + frame, _ := frameFor(sender, proto.TypeAssignmentStatus, id) + var status proto.AssignmentStatusPayload + if err := frame.DecodePayload(&status); err != nil { + t.Fatal(err) + } + return status +} + +// release sends assignment_release under the session's assignment at epoch. +func release(t *testing.T, r *dispatch.Router, session, id string, epoch uint64, removeHome bool) { + t.Helper() + env := scoped(t, session, proto.TypeAssignmentRelease, id, proto.AssignmentReleasePayload{RemoveHome: removeHome}) + env.Assignment.Epoch = epoch + if err := r.Handle(t.Context(), env); err != nil { + t.Fatal(err) + } +} + +// observedExecutor runs closed before it closes. +type observedExecutor struct { + *reusableExecutor + closed func() +} + +func (e *observedExecutor) Close(ctx context.Context) error { + e.closed() + return e.reusableExecutor.Close(ctx) +} + +func TestAssignmentRejectsStaleAndForeignFrames(t *testing.T) { + h := newHarness(t) + defer h.router.Shutdown(context.Background()) + assign(t, h.router, "s", "") + release(t, h.router, "s", "release", 2, false) + if got := waitAssignmentStatus(t, h.sender, "release"); got.State != proto.AssignmentReleased || got.ErrorCode != "" { + t.Fatalf("release = %+v", got) + } + foreign := ref("s") + foreign.AssignmentID = "foreign" + prepare := proto.ExecutionPreparePayload{SessionID: "s", Configuration: proto.PromptRequestPayload{AgentKind: "fake_alpha", AgentStateKey: stateKey("s"), StrictResume: true, DisableExecutionEnvironment: true}} + for id, test := range map[string]struct { + ref proto.AssignmentRef + code string + }{ + "stale": {ref("s"), proto.AssignmentStale}, + "foreign": {foreign, proto.AssignmentConflict}, + } { + env := scoped(t, "s", proto.TypeExecutionPrepare, id, prepare) + env.Assignment = test.ref + if err := h.router.Handle(t.Context(), env); err == nil { + t.Fatalf("%s preparation admitted", id) + } + if got := waitPreparationStatus(t, h.sender, id, "rejected", ""); got.ErrorCode != test.code { + t.Fatalf("%s preparation = %+v", id, got) + } + bind := scoped(t, "s", proto.TypeAssignmentBind, id+"-bind", proto.AssignmentBindPayload{}) + bind.Assignment = test.ref + if err := h.router.Handle(t.Context(), bind); err != nil { + t.Fatal(err) + } + if got := waitAssignmentStatus(t, h.sender, id+"-bind"); got.State != proto.AssignmentFailed || got.ErrorCode != test.code { + t.Fatalf("%s bind = %+v", id, got) + } + } +} + +func TestAssignmentReleaseWaitsForRacingPreparation(t *testing.T) { + var calls atomic.Int32 + var sender *recSender + var replyBeforeClose atomic.Bool + owner := &observedExecutor{reusableExecutor: &reusableExecutor{}, closed: func() { + replyBeforeClose.Store(hasFrame(sender, proto.TypeAssignmentStatus, "release")) + }} + entered, cancelled, unblock := make(chan struct{}), make(chan struct{}), make(chan struct{}) + r, sender := poolRouter(t, func(ctx context.Context, _ proto.PromptRequestPayload) (agent.Executor, error) { + calls.Add(1) + close(entered) + <-ctx.Done() + close(cancelled) + <-unblock + return owner, nil + }) + defer close(unblock) + assign(t, r, preparationSessionID, "") + if err := r.Handle(t.Context(), scoped(t, preparationSessionID, proto.TypeExecutionPrepare, "prepare", executorRequest())); err != nil { + t.Fatal(err) + } + <-entered + release(t, r, preparationSessionID, "release", 2, false) + select { + case <-cancelled: + case <-time.After(3 * time.Second): + t.Fatal("release did not cancel the preparation") + } + if hasFrame(sender, proto.TypeAssignmentStatus, "release") { + t.Fatal("release replied before the Executor closed") + } + unblock <- struct{}{} + if got := waitAssignmentStatus(t, sender, "release"); got.State != proto.AssignmentReleased || got.ErrorCode != "" { + t.Fatalf("release = %+v", got) + } + if owner.closes.Load() != 1 || replyBeforeClose.Load() { + t.Fatalf("Executor closes = %d, replied before close = %t", owner.closes.Load(), replyBeforeClose.Load()) + } + for _, frame := range sender.snapshot() { + var status proto.PreparationStatusPayload + if frame.Type == proto.TypePreparationStatus && frame.ID == "prepare" && frame.DecodePayload(&status) == nil && status.State == "ready" { + t.Fatal("released preparation became ready") + } + } + late := scoped(t, preparationSessionID, proto.TypeExecutionPrepare, "late", executorRequest()) + late.Assignment.Epoch = 2 + if err := r.Handle(t.Context(), late); err == nil { + t.Fatal("released assignment admitted a preparation") + } + if got := waitPreparationStatus(t, sender, "late", "rejected", ""); got.ErrorCode != proto.AssignmentStale || calls.Load() != 1 { + t.Fatalf("late preparation = %+v, factory calls = %d", got, calls.Load()) + } +} + +func TestReleaseWithoutHomeRemovalKeepsAssignment(t *testing.T) { + h := newHarness(t) + defer h.router.Shutdown(context.Background()) + assign(t, h.router, "s", "") + release(t, h.router, "s", "remove", 2, true) + if got := waitAssignmentStatus(t, h.sender, "remove"); got.State != proto.AssignmentFailed || got.ErrorCode != proto.UnsupportedOperation { + t.Fatalf("release = %+v", got) + } + startRun(t, h.router, h.sender, "fake_alpha", "s") +} + +func TestUnknownEnvelopeGetsCorrelatedProtocolError(t *testing.T) { + h := newHarness(t) + defer h.router.Shutdown(context.Background()) + if err := h.router.Handle(t.Context(), scoped(t, "s", "future_operation", "request", nil)); err != nil { + t.Fatal(err) + } + frame, ok := frameFor(h.sender, proto.TypeProtocolError, "request") + var got proto.ProtocolErrorPayload + if !ok || frame.DecodePayload(&got) != nil || got != (proto.ProtocolErrorPayload{Type: "future_operation", ErrorCode: proto.UnsupportedOperation}) || frame.Assignment != ref("s") { + t.Fatalf("protocol_error = %+v %+v", frame, got) + } +} diff --git a/apps/daemon/internal/dispatch/cancellation.go b/apps/daemon/internal/dispatch/cancellation.go index 4e65412d9..c54e9eb45 100644 --- a/apps/daemon/internal/dispatch/cancellation.go +++ b/apps/daemon/internal/dispatch/cancellation.go @@ -2,22 +2,12 @@ package dispatch import ( "context" - "errors" - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" ) -func (r *Router) releaseCompletedSession(state *sessionState) error { - r.mu.Lock() - state.retain = false - r.mu.Unlock() - receiptErr := r.finishSteering(state) - err := state.session.Cancel(context.Background()) - state.ctxCancel() - return errors.Join(receiptErr, err) -} - +// handlePromptCancel cancels the run that env.Assignment started. It admits a +// released assignment, so a cancellation of admitted work keeps its receipt. func (r *Router) handlePromptCancel(ctx context.Context, env proto.Envelope) error { var request proto.PromptCancelPayload if err := env.DecodeRequest(&request); err != nil { @@ -29,46 +19,29 @@ func (r *Router) handlePromptCancel(ctx context.Context, env proto.Envelope) err return ErrRouterClosed } state := r.sessions[env.ID] - if state != nil { - state.retain = false - } - var owner agent.Session - if state != nil && state.preparedHandoff != nil { - handoff := state.preparedHandoff - release, attempt := r.claimPreparedReleaseLocked(state, true, "", true) - if request.DeliveryID != "" { - r.shutdownWG.Add(1) - go r.sendPreparedCancellation(state, handoff, release, attempt, env, request.DeliveryID) - } + if state == nil || state.assignment != env.Assignment { r.mu.Unlock() - return nil + ack := proto.InteractionDecisionAckPayload{DeliveryID: request.DeliveryID, ErrorCode: "run_inactive"} + if state != nil { + ack.ErrorCode = proto.AssignmentConflict + } + return r.sendCancellationAck(ctx, env, ack) } - if state != nil && state.session != nil { - owner = state.session + handoff := state.preparedHandoff + release, attempt := r.claimPreparedReleaseLocked(state, true, "", true) + if request.DeliveryID != "" { + r.shutdownWG.Add(1) + go r.sendPreparedCancellation(state, handoff, release, attempt, env, request.DeliveryID) } r.mu.Unlock() - ack := proto.InteractionDecisionAckPayload{DeliveryID: request.DeliveryID, ErrorCode: "run_inactive"} - if state != nil && owner == nil { - ack.ErrorCode = "not_ready" - } else if owner != nil { - if err := owner.Cancel(ctx); err != nil { - r.log.WarnContext(ctx, "session.Cancel failed", "run_id", env.ID, "err", err) - ack.ErrorCode = "cancel_failed" - } else { - ack.Applied, ack.ErrorCode = true, "" - outcome := owner.CancellationOutcome() - ack.Outcome = &outcome - } - state.ctxCancel() - } - return r.sendCancellationAck(ctx, env, ack) + return nil } func (r *Router) sendCancellationAck(ctx context.Context, env proto.Envelope, ack proto.InteractionDecisionAckPayload) error { if ack.DeliveryID == "" { return nil } - reply, err := proto.NewEnvelopeWithTrace(proto.TypeInteractionDecisionAck, env.ID, ack, env.Trace) + reply, err := env.Reply(proto.TypeInteractionDecisionAck, ack) if err != nil { return err } diff --git a/apps/daemon/internal/dispatch/cancellation_test.go b/apps/daemon/internal/dispatch/cancellation_test.go deleted file mode 100644 index 4c5d32e69..000000000 --- a/apps/daemon/internal/dispatch/cancellation_test.go +++ /dev/null @@ -1,139 +0,0 @@ -package dispatch_test - -import ( - "context" - "errors" - "reflect" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" -) - -type cancelReceiptSession struct { - *fakeSession - entered chan struct{} - release chan struct{} - err error - outcome proto.DonePayload -} - -func (s *cancelReceiptSession) CancellationOutcome() proto.DonePayload { - return s.outcome -} - -func TestCompletionWaitsForNativeWriterRelease(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - sess := &cancelReceiptSession{entered: make(chan struct{}), release: make(chan struct{})} - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - sess.fakeSession = &fakeSession{out: out, closeOutOnCancel: true} - return sess, nil - }) - if err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "release", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", AgentStateKey: "stable", ReleaseOnCompletion: true}))); err != nil { - t.Fatal(err) - } - sess.out <- mustEnv(t, proto.TypeDone, "release", proto.DonePayload{Content: "Finished"}) - <-sess.entered - if len(h.sender.snapshot()) != 0 { - t.Fatal("completion acknowledged before native writer was released") - } - close(sess.release) - waitFor(t, func() bool { return h.router.ActiveRuns() == 0 }, "release completion") - frames := h.sender.snapshot() - if len(frames) != 1 || frames[0].Type != proto.TypeDone || sess.cancels() != 1 { - t.Fatal("completion or native release missing") - } -} - -func (s *cancelReceiptSession) Cancel(ctx context.Context) error { - if s.entered != nil { - close(s.entered) - <-s.release - s.entered = nil - } - _ = s.fakeSession.Cancel(ctx) - return s.err -} - -func TestCancellationReceiptFollowsAdapterOutcome(t *testing.T) { - observed := proto.DonePayload{Content: "partial output", Metadata: map[string]any{proto.DoneMetaAgentSessionID: "native-cancelled"}} - for _, test := range []struct { - name string - outcome proto.DonePayload - err error - }{ - {name: "observed", outcome: observed}, - {name: "unknown"}, - {name: "failed", outcome: observed, err: errors.New("adapter could not cancel")}, - {name: "unsupported", outcome: observed, err: agent.ErrUnsupportedOperation}, - {name: "deadline", outcome: observed, err: context.DeadlineExceeded}, - } { - t.Run(test.name, func(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - sess := &cancelReceiptSession{entered: make(chan struct{}), release: make(chan struct{}), outcome: test.outcome, err: test.err} - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - sess.fakeSession = &fakeSession{out: out, closeOutOnCancel: true} - return sess, nil - }) - if err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "run", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex"}))); err != nil { - t.Fatal(err) - } - done := make(chan error, 1) - go func() { - done <- h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptCancel, "run", proto.PromptCancelPayload{DeliveryID: "cancel-1"})) - }() - <-sess.entered - for _, env := range h.sender.snapshot() { - if env.Type == proto.TypeInteractionDecisionAck { - t.Fatal("cancellation acknowledged before adapter returned") - } - } - close(sess.release) - if err := <-done; err != nil { - t.Fatal(err) - } - found := false - for _, env := range h.sender.snapshot() { - if env.Type == proto.TypeInteractionDecisionAck { - found = true - var ack proto.InteractionDecisionAckPayload - _ = env.DecodePayload(&ack) - if ack.Applied != (test.err == nil) || ack.DeliveryID != "cancel-1" { - t.Fatalf("wrong receipt: %+v", ack) - } - if test.err == nil { - if ack.ErrorCode != "" || ack.Outcome == nil || !reflect.DeepEqual(*ack.Outcome, test.outcome) { - t.Fatalf("cancellation receipt changed observed evidence: %+v", ack) - } - } else if ack.ErrorCode != "cancel_failed" || ack.Outcome != nil { - t.Fatalf("failed cancellation supplied a success outcome: %+v", ack) - } - } - } - if !found { - t.Fatal("missing cancellation receipt") - } - }) - } -} - -func TestLegacyCancellationDoesNotEmitNewFrames(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - if err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "legacy", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"}))); err != nil { - t.Fatal(err) - } - sess := <-h.gotSess - sess.closeOutOnCancel = true - if err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptCancel, "legacy", proto.PromptCancelPayload{})); err != nil { - t.Fatal(err) - } - for _, env := range h.sender.snapshot() { - if env.Type == proto.TypeInteractionDecisionAck { - t.Fatal("legacy cancellation emitted new receipt") - } - } -} diff --git a/apps/daemon/internal/dispatch/capability_admission_test.go b/apps/daemon/internal/dispatch/capability_admission_test.go index e6975aef5..8cdb2240d 100644 --- a/apps/daemon/internal/dispatch/capability_admission_test.go +++ b/apps/daemon/internal/dispatch/capability_admission_test.go @@ -24,15 +24,13 @@ func TestSteeringUsesAdmittedDeclarationAndDoesNotReplayUnsupportedImplementatio }}, nil } info := proto.SupportedAgentKind{Kind: "fixture", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilityFromBool(supported)})} - h.reg.RegisterKind(info, prototest.ModelConfiguration(), factory) - if err := h.router.Handle(t.Context(), mustEnv(t, proto.TypePromptRequest, "run", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fixture"}))); err != nil { - t.Fatal(err) - } + registerSession(h.reg, info, factory) + startRun(t, h.router, h.sender, "fixture", "run") // A new registration cannot rewrite the already admitted owner's contract. info.Capabilities.Steering = proto.CapabilityFromBool(!supported) - h.reg.RegisterKind(info, prototest.ModelConfiguration(), factory) + registerSession(h.reg, info, factory) for range 2 { - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run", proto.PromptSteerPayload{InputID: "input", Input: proto.TextInput("hello")})); err != nil { + if err := handleSteeringAndWait(t, h, scoped(t, "run", proto.TypePromptSteer, "run", proto.PromptSteerPayload{InputID: "input", Input: proto.TextInput("hello")})); err != nil { t.Fatal(err) } ack := lastSteeringAck(t, h.sender, "run", "input") @@ -63,21 +61,19 @@ func TestInteractionDeclarationPrecedesResponderMethods(t *testing.T) { defer h.router.Shutdown(context.Background()) var session *fakeSession info := proto.SupportedAgentKind{Kind: "fixture", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Permissions: proto.CapabilityFromBool(supported)})} - h.reg.RegisterKind(info, prototest.ModelConfiguration(), func(_ context.Context, _ proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(h.reg, info, func(_ context.Context, _ proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { session = &fakeSession{out: out, closeOutOnCancel: true, submitErr: agent.ErrUnsupportedOperation, askErr: agent.ErrUnsupportedOperation} return session, nil }) - if err := h.router.Handle(t.Context(), mustEnv(t, proto.TypePromptRequest, "run", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fixture"}))); err != nil { - t.Fatal(err) - } + startRun(t, h.router, h.sender, "fixture", "run") event := mustEnv(t, proto.TypePermissionRequest, "run", proto.PermissionRequestPayload{RequestID: "interaction", Tool: "fixture"}) - decision := mustEnv(t, proto.TypePermissionDecision, "interaction", proto.PermissionDecisionPayload{DeliveryID: "decision", Approved: true}) + decision := scoped(t, "run", proto.TypePermissionDecision, "interaction", proto.PermissionDecisionPayload{DeliveryID: "decision", Approved: true}) if ask { event = mustEnv(t, proto.TypePromptForUserChoice, "run", proto.PromptForUserChoicePayload{AskID: "interaction"}) - decision = mustEnv(t, proto.TypePromptForUserChoiceDecision, "interaction", proto.PromptForUserChoiceDecisionPayload{DeliveryID: "decision"}) + decision = scoped(t, "run", proto.TypePromptForUserChoiceDecision, "interaction", proto.PromptForUserChoiceDecisionPayload{DeliveryID: "decision"}) } session.out <- event - waitFor(t, func() bool { return len(h.sender.snapshot()) > 0 }, "interaction indexed") + waitFor(t, func() bool { return hasFrame(h.sender, event.Type, "run") }, "interaction indexed") if err := h.router.Handle(t.Context(), decision); err != nil { t.Fatal(err) } diff --git a/apps/daemon/internal/dispatch/durable_only_test.go b/apps/daemon/internal/dispatch/durable_only_test.go index ffb5e4c98..d8707a13d 100644 --- a/apps/daemon/internal/dispatch/durable_only_test.go +++ b/apps/daemon/internal/dispatch/durable_only_test.go @@ -33,6 +33,7 @@ func (t *durableOnlyTurn) SteerWithReceipt(_ context.Context, _ proto.PromptStee func TestPublicTextSteeringDoesNotRequireOptionalSteerer(t *testing.T) { owner := &durableOnlyExecutor{reusableExecutor: &reusableExecutor{starts: make(chan *reusableTurn, 1)}} router, sender := poolRouter(t, func(context.Context, proto.PromptRequestPayload) (agent.Executor, error) { return owner, nil }) + assign(t, router, preparationSessionID, "") admission := executorAdmission(t, router, sender, "prepare", executorRequest()) startExecutorTurn(t, router, sender, "prepare", "run", admission) turn := <-owner.starts diff --git a/apps/daemon/internal/dispatch/environment_test.go b/apps/daemon/internal/dispatch/environment_test.go index 1fe4a940e..03e63552a 100644 --- a/apps/daemon/internal/dispatch/environment_test.go +++ b/apps/daemon/internal/dispatch/environment_test.go @@ -10,40 +10,6 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" ) -func TestNoEnvironmentRejectsOtherEngineBeforeFactory(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - called := false - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "fake_alpha", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}, prototest.ModelConfiguration(), func(context.Context, proto.PromptRequestPayload, chan<- proto.Envelope) (agent.Session, error) { - called = true - return nil, nil - }) - err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "none", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha", DisableExecutionEnvironment: true}))) - if err == nil || called { - t.Fatal("unsupported engine was started", err) - } - frames := h.sender.snapshot() - if len(frames) != 2 || frames[0].Type != proto.TypeError || frames[1].Type != proto.TypeDone { - t.Fatal(frames) - } -} - -func TestNoEnvironmentUsesAvailableCapability(t *testing.T) { - for _, available := range []bool{false, true} { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - called := false - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "claude_sdk", Available: available, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{EnvironmentNone: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(context.Context, proto.PromptRequestPayload, chan<- proto.Envelope) (agent.Session, error) { - called = true - return nil, errors.New("controlled factory stop") - }) - _ = h.router.Handle(t.Context(), mustEnv(t, proto.TypePromptRequest, "sdk", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "claude_sdk", DisableExecutionEnvironment: true}))) - if called != available { - t.Fatalf("factory called=%t, available=%t", called, available) - } - } -} - func TestLocalEnvironmentRequiresAvailableCapability(t *testing.T) { for _, mode := range []string{"unsupported", "unavailable", "none conflict", "supported"} { t.Run(mode, func(t *testing.T) { @@ -51,25 +17,29 @@ func TestLocalEnvironmentRequiresAvailableCapability(t *testing.T) { defer h.router.Shutdown(context.Background()) called := false h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: mode != "unavailable", - Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{LocalEnvironment: proto.CapabilityFromBool(mode != "unsupported")})}, - prototest.ModelConfiguration(), func(_ context.Context, req proto.PromptRequestPayload, _ chan<- proto.Envelope) (agent.Session, error) { - called = true - if req.LocalEnvironment == nil || req.LocalEnvironment.ID != preparationEnvironmentID { - t.Error("local descriptor lost before factory") - } - return nil, errors.New("controlled factory stop") + Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{LocalEnvironment: proto.CapabilityFromBool(mode != "unsupported"), EnvironmentNone: proto.CapabilitySupported})}, + prototest.ModelConfiguration(), func(context.Context, proto.PromptRequestPayload, chan<- proto.Envelope) (agent.Session, error) { + return nil, errors.New("ordinary factory is forbidden") }) - req := preparationRequest().Configuration - req.AgentKind = "codex" - req.DisableExecutionEnvironment = mode == "none conflict" - _ = h.router.Handle(t.Context(), mustEnv(t, proto.TypePromptRequest, "local", req)) + h.reg.RegisterExecutor("codex", func(_ context.Context, req proto.PromptRequestPayload) (agent.Executor, error) { + called = true + if req.LocalEnvironment == nil || req.LocalEnvironment.ID != preparationEnvironmentID { + t.Error("local descriptor lost before factory") + } + return nil, errors.New("controlled factory stop") + }) + req := preparationRequest() + req.Configuration.AgentKind = "codex" + req.Configuration.DisableExecutionEnvironment = mode == "none conflict" + _ = h.router.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "local", req)) + state := "rejected" + if mode == "supported" { + state = "failed" + } + waitPreparationStatus(t, h.sender, "local", state, "") if called != (mode == "supported") { t.Fatalf("unexpected factory call for %s", mode) } - frames := h.sender.snapshot() - if len(frames) != 2 || frames[0].Type != proto.TypeError || frames[1].Type != proto.TypeDone { - t.Fatal("missing terminal error frames") - } }) } } diff --git a/apps/daemon/internal/dispatch/executor.go b/apps/daemon/internal/dispatch/executor.go index 1cade41ee..82c5d227e 100644 --- a/apps/daemon/internal/dispatch/executor.go +++ b/apps/daemon/internal/dispatch/executor.go @@ -26,12 +26,14 @@ type executorState struct { admission *preparationState run *sessionState preparing bool - invalid bool - closeDone chan struct{} - closeReason string - closeErr error - timer *time.Timer - idleLease uint64 + // prepared closes once the native preparation returns. + prepared chan struct{} + invalid bool + closeDone chan struct{} + closeReason string + closeErr error + timer *time.Timer + idleLease uint64 } func executorFingerprint(req proto.PromptRequestPayload) ([32]byte, error) { @@ -73,11 +75,16 @@ func (r *Router) handleExecutorPrepare(ctx context.Context, env proto.Envelope, return r.rejectPreparation(env, "invalid_configuration") } requestFingerprint := sha256.Sum256(encoded) + req.Assignment = env.Assignment r.mu.Lock() if r.closed || r.suspension != nil { r.mu.Unlock() return ErrRouterClosed } + if code := r.admitLocked(env.Assignment, input.SessionID, req.EnvironmentID()); code != "" { + r.mu.Unlock() + return r.rejectPreparation(env, code) + } if r.workspaceWrite != nil || r.workspaceExport != nil || r.runtimePreparation != nil { r.mu.Unlock() return r.rejectPreparation(env, "resource_unavailable") @@ -142,12 +149,12 @@ func (r *Router) handleExecutorPrepare(ctx context.Context, env proto.Envelope, return r.rejectPreparation(env, "executor_capacity") } ownerCtx, cancel := context.WithCancel(context.WithoutCancel(ctx)) - owner = &executorState{capabilities: caps, id: uuid.NewString(), sessionID: input.SessionID, environmentID: req.EnvironmentID(), stateKey: req.AgentStateKey, fingerprint: fingerprint, ctx: ownerCtx, cancel: cancel, preparing: true, nativeID: req.AgentSessionID} + owner = &executorState{capabilities: caps, id: uuid.NewString(), sessionID: input.SessionID, environmentID: req.EnvironmentID(), stateKey: req.AgentStateKey, fingerprint: fingerprint, ctx: ownerCtx, cancel: cancel, preparing: true, prepared: make(chan struct{}), nativeID: req.AgentSessionID} r.executors[input.SessionID] = owner r.log.Info("executor owner_created", "executor_id", owner.id, "session_id", owner.sessionID) } operation, cancel := context.WithCancel(context.WithoutCancel(ctx)) - p := &preparationState{capabilities: owner.capabilities, requestID: env.ID, trace: env.Trace, fingerprint: requestFingerprint, ctx: operation, cancel: cancel, stateKey: req.AgentStateKey, environmentID: req.EnvironmentID(), executor: owner, owns: true, busy: !reused, deadline: time.Now().Add(r.preparationTimeout)} + p := &preparationState{capabilities: owner.capabilities, request: proto.Envelope{ID: env.ID, Trace: env.Trace, Assignment: env.Assignment}, fingerprint: requestFingerprint, ctx: operation, cancel: cancel, environmentID: req.EnvironmentID(), executor: owner, owns: true, busy: !reused, deadline: time.Now().Add(r.preparationTimeout)} state := "preparing" if reused { state = "ready" @@ -176,7 +183,7 @@ func (r *Router) prepareExecutor(p *preparationState, req proto.PromptRequestPay r.mu.Lock() initial := p.status r.mu.Unlock() - if !r.sendPreparation(p.requestID, p.trace, initial) { + if !r.sendPreparation(p.request, initial) { r.abandonExecutorAdmission(p, "failed", "status_delivery_failed", false) } var native agent.Executor @@ -193,6 +200,7 @@ func (r *Router) prepareExecutor(p *preparationState, req proto.PromptRequestPay } r.mu.Lock() owner.native, owner.preparing = native, false + close(owner.prepared) r.log.Info("executor native_prepare", "executor_id", owner.id, "session_id", owner.sessionID, "duration_ms", time.Since(started).Milliseconds(), "success", err == nil && native != nil) ready := native != nil && err == nil && !owner.invalid && !r.closed && r.suspension == nil && p.status.State == "preparing" && p.ctx.Err() == nil p.busy = false @@ -218,7 +226,7 @@ func (r *Router) prepareExecutor(p *preparationState, req proto.PromptRequestPay status = p.status r.mu.Unlock() } - if !r.sendPreparation(p.requestID, p.trace, status) && ready { + if !r.sendPreparation(p.request, status) && ready { r.abandonExecutorAdmission(p, "failed", "status_delivery_failed", false) } } diff --git a/apps/daemon/internal/dispatch/executor_cancel_receipt_test.go b/apps/daemon/internal/dispatch/executor_cancel_receipt_test.go index 30aef2660..496e4dd4b 100644 --- a/apps/daemon/internal/dispatch/executor_cancel_receipt_test.go +++ b/apps/daemon/internal/dispatch/executor_cancel_receipt_test.go @@ -133,6 +133,7 @@ func TestExecutorCancellationReachesNativeBeforeDurableReceiptJoin(t *testing.T) if err != nil { t.Fatal(err) } + assign(t, r, preparationSessionID, "") var receiptOnce, closeOnce sync.Once releaseReceipt := func() { receiptOnce.Do(func() { close(sender.release) }) } releaseClose := func() { closeOnce.Do(func() { close(owner.closeRelease) }) } diff --git a/apps/daemon/internal/dispatch/executor_handoff_test.go b/apps/daemon/internal/dispatch/executor_handoff_test.go index 9f3b233c2..632edbde3 100644 --- a/apps/daemon/internal/dispatch/executor_handoff_test.go +++ b/apps/daemon/internal/dispatch/executor_handoff_test.go @@ -104,6 +104,7 @@ func TestPreparedDonePublishesAfterExecutorHandoff(t *testing.T) { if err != nil { t.Fatal(err) } + bindAssignment(r, suspendRef, "") var releaseOnce sync.Once unblock := func() { releaseOnce.Do(func() { close(sender.release) }) } t.Cleanup(func() { @@ -120,6 +121,7 @@ func TestPreparedDonePublishesAfterExecutorHandoff(t *testing.T) { if err != nil { t.Fatal(err) } + env.Assignment = suspendRef if err = r.Handle(t.Context(), env); err != nil { t.Fatal(err) } diff --git a/apps/daemon/internal/dispatch/executor_test.go b/apps/daemon/internal/dispatch/executor_test.go index 74274af6d..affb5d60d 100644 --- a/apps/daemon/internal/dispatch/executor_test.go +++ b/apps/daemon/internal/dispatch/executor_test.go @@ -68,7 +68,7 @@ func (t *reusableTurn) AwaitSettlement(ctx context.Context) (agent.TurnSettlemen } func executorRequest() proto.ExecutionPreparePayload { - return proto.ExecutionPreparePayload{SessionID: "session", Configuration: prototest.WithModel(proto.PromptRequestPayload{AgentKind: "reusable", AgentStateKey: "agents-api-session", StrictResume: true, DisableExecutionEnvironment: true})} + return proto.ExecutionPreparePayload{SessionID: preparationSessionID, Configuration: prototest.WithModel(proto.PromptRequestPayload{AgentKind: "reusable", AgentStateKey: "agents-api-" + preparationSessionID, StrictResume: true, DisableExecutionEnvironment: true})} } func executorRouter(t *testing.T, owner *reusableExecutor, idle time.Duration) (*dispatch.Router, *recSender, *atomic.Int32) { t.Helper() @@ -86,6 +86,7 @@ func executorRouter(t *testing.T, owner *reusableExecutor, idle time.Duration) ( if err != nil { t.Fatal(err) } + assign(t, router, preparationSessionID, "") t.Cleanup(func() { owner.mu.Lock() owner.closeErr = nil @@ -100,7 +101,7 @@ func executorRouter(t *testing.T, owner *reusableExecutor, idle time.Duration) ( } func executorAdmission(t *testing.T, r *dispatch.Router, s *recSender, key string, req proto.ExecutionPreparePayload) proto.PreparationStatusPayload { t.Helper() - if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, key, req)); err != nil { + if err := r.Handle(t.Context(), scoped(t, req.SessionID, proto.TypeExecutionPrepare, key, req)); err != nil { t.Fatal(err) } return waitPreparationStatus(t, s, key, "ready", "") @@ -255,8 +256,9 @@ func TestExecutorIdleAndActiveCapacitiesAreIndependent(t *testing.T) { }) for i := 0; i < 17; i++ { id := fmt.Sprint("idle-", i) + assign(t, r, id, "") ready := executorAdmission(t, r, s, id, poolRequest(id)) - if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionRelease, id, proto.ExecutionReleasePayload{Handle: ready.Handle})); err != nil { + if err := r.Handle(t.Context(), scoped(t, id, proto.TypeExecutionRelease, id, proto.ExecutionReleasePayload{Handle: ready.Handle})); err != nil { t.Fatal(err) } } @@ -271,9 +273,11 @@ func TestExecutorIdleAndActiveCapacitiesAreIndependent(t *testing.T) { }, "idle capacity eviction") for i := 0; i < 4; i++ { id := fmt.Sprint("active-", i) + assign(t, r, id, "") executorAdmission(t, r, s, id, poolRequest(id)) } - if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "overflow", poolRequest("overflow"))); err == nil { + assign(t, r, "overflow", "") + if err := r.Handle(t.Context(), scoped(t, "overflow", proto.TypeExecutionPrepare, "overflow", poolRequest("overflow"))); err == nil { t.Fatal("active capacity exceeded") } } diff --git a/apps/daemon/internal/dispatch/functions.go b/apps/daemon/internal/dispatch/functions.go index 01bfa71f7..b19ec7282 100644 --- a/apps/daemon/internal/dispatch/functions.go +++ b/apps/daemon/internal/dispatch/functions.go @@ -19,7 +19,7 @@ func (r *Router) handleFunctionResult(ctx context.Context, env proto.Envelope) e return errors.New("function result requires run, call and delivery identities") } if err := result.ValidateContent(); err != nil { - return r.sendInteractionDecisionAck(ctx, env.ID, result.DeliveryID, false, "invalid_result", err.Error()) + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, "invalid_result", err.Error()) } decision := result decision.DeliveryID = "" @@ -30,11 +30,17 @@ func (r *Router) handleFunctionResult(ctx context.Context, env proto.Envelope) e fingerprint := sha256.Sum256(encoded) // Scope receipt replay to both identities, even when native call IDs repeat across Runs. kind := proto.TypeFunctionResult + "\x00" + result.CallID - if handled, err := r.replayAppliedInteractionDecision(ctx, env.ID, result.DeliveryID, kind, fingerprint); handled { + if handled, err := r.replayAppliedInteractionDecision(ctx, env, result.DeliveryID, kind, fingerprint); handled { return err } r.mu.Lock() state := r.sessions[env.ID] + if state != nil { + if code := r.admitRunLocked(env.Assignment, state); code != "" { + r.mu.Unlock() + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, code, "The run's assignment does not admit this result.") + } + } session, finishOperation, ready := r.preparedOperationLocked(state) var submitter agent.FunctionResultSubmitter if ready { @@ -42,7 +48,7 @@ func (r *Router) handleFunctionResult(ctx context.Context, env proto.Envelope) e } r.mu.Unlock() if state != nil && !ready { - return r.sendInteractionDecisionAck(ctx, env.ID, result.DeliveryID, false, "not_ready", "function call is waiting for the native session") + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, "not_ready", "function call is waiting for the native session") } if finishOperation != nil { defer finishOperation() @@ -51,13 +57,13 @@ func (r *Router) handleFunctionResult(ctx context.Context, env proto.Envelope) e defer stop() } if state != nil && !state.capabilities.FunctionTools.IsSupported() { - return r.sendInteractionDecisionAck(ctx, env.ID, result.DeliveryID, false, "unsupported", "The runtime declaration does not support function results.") + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, "unsupported", "The runtime declaration does not support function results.") } if ready && submitter == nil { - return r.sendInteractionDecisionAck(ctx, env.ID, result.DeliveryID, false, "contract_violation", "Declared function capability has no implementation.") + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, "contract_violation", "Declared function capability has no implementation.") } if submitter == nil { - return r.sendInteractionDecisionAck(ctx, env.ID, result.DeliveryID, false, "not_pending", "function call is no longer pending") + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, "not_pending", "function call is no longer pending") } if err := submitter.SubmitFunctionResult(ctx, result); err != nil { code := "runtime_error" @@ -67,8 +73,8 @@ func (r *Router) handleFunctionResult(ctx context.Context, env proto.Envelope) e if errors.Is(err, agent.ErrUnknownFunctionCall) { code = "not_pending" } - return r.sendInteractionDecisionAck(ctx, env.ID, result.DeliveryID, false, code, "function result was not applied") + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, code, "function result was not applied") } r.rememberAppliedInteractionDecision(env.ID, kind, fingerprint) - return r.sendInteractionDecisionAck(ctx, env.ID, result.DeliveryID, true, "", "") + return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, true, "", "") } diff --git a/apps/daemon/internal/dispatch/functions_native_test.go b/apps/daemon/internal/dispatch/functions_native_test.go deleted file mode 100644 index a17d0f7a7..000000000 --- a/apps/daemon/internal/dispatch/functions_native_test.go +++ /dev/null @@ -1,193 +0,0 @@ -package dispatch_test - -import ( - "context" - "encoding/json" - "fmt" - "net/http" - "net/http/httptest" - "os" - "path/filepath" - "reflect" - "strings" - "sync/atomic" - "testing" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent/codex" - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/dispatch" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" - "github.com/MiniMax-AI/OpenAgentCore/internal/modelprovider" -) - -type nativeFunctionSender chan proto.Envelope - -func (s nativeFunctionSender) Send(ctx context.Context, e proto.Envelope) error { - select { - case s <- e: - return nil - case <-ctx.Done(): - return ctx.Err() - } -} - -func TestNativeFunctionBridge(t *testing.T) { - root := os.Getenv("OAC_TEST_NATIVE_PROOF_DIR") - if root == "" { - t.Skip("explicit native Codex binary and proof directory required") - } - home, err := os.MkdirTemp(root, "daemon-functions-") - if err != nil { - t.Fatal(err) - } - if err := os.Chmod(home, 0o700); err != nil { - t.Fatal(err) - } - t.Setenv("OAC_RUNTIME_HOME", home) - var count atomic.Int32 - model := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - var body map[string]any - if err := json.NewDecoder(r.Body).Decode(&body); err != nil { - t.Error(err) - return - } - n := count.Add(1) - raw, _ := json.MarshalIndent(body, "", " ") - _ = os.WriteFile(filepath.Join(home, fmt.Sprintf("request-%d.json", n)), raw, 0600) - var item map[string]any - if n%2 == 1 { - if !strings.Contains(string(raw), "lookup_ticket") { - t.Error("tool was not registered") - } - item = map[string]any{"id": fmt.Sprintf("fc_%d", n), "type": "function_call", "call_id": fmt.Sprintf("call_%d", n), "name": "lookup_ticket", "arguments": `{"ticket":"42"}`, "status": "completed"} - } else { - - var request struct { - Input []struct { - Type string `json:"type"` - CallID string `json:"call_id"` - Output json.RawMessage `json:"output"` - } `json:"input"` - } - if err := json.Unmarshal(raw, &request); err != nil { - t.Error(err) - } - found := false - for _, entry := range request.Input { - if entry.Type != "function_call_output" || entry.CallID != fmt.Sprintf("call_%d", n-1) { - continue - } - found = true - var parts []proto.InputContent - if err := json.Unmarshal(entry.Output, &parts); err != nil { - t.Error(err) - continue - } - expected := functionResultContent("TICKET-RESULT") - if !reflect.DeepEqual(parts, expected) { - t.Errorf("native result lost text/image content or order: %s", entry.Output) - } - } - if !found { - t.Error("native model did not receive function result") - } - item = map[string]any{"id": fmt.Sprintf("msg_%d", n), "type": "message", "role": "assistant", "phase": "final_answer", "status": "completed", "content": []any{map[string]any{"type": "output_text", "text": "FUNCTION-OK", "annotations": []any{}}}} - } - w.Header().Set("Content-Type", "text/event-stream") - send := func(kind string, data map[string]any) { - data["type"] = kind - b, _ := json.Marshal(data) - fmt.Fprintf(w, "event: %s\ndata: %s\n\n", kind, b) - w.(http.Flusher).Flush() - } - send("response.created", map[string]any{"response": map[string]any{"id": fmt.Sprintf("r_%d", n), "status": "in_progress", "output": []any{}}}) - send("response.output_item.added", map[string]any{"output_index": 0, "item": item}) - send("response.output_item.done", map[string]any{"output_index": 0, "item": item}) - send("response.completed", map[string]any{"response": map[string]any{"id": fmt.Sprintf("r_%d", n), "object": "response", "created_at": time.Now().Unix(), "status": "completed", "model": "gpt-5.5", "output": []any{item}}}) - })) - defer model.Close() - reg := agent.NewRegistry() - reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{FunctionTools: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported})}, prototest.ModelConfiguration(), codex.Factory) - sender := make(nativeFunctionSender, 256) - router, err := dispatch.New(dispatch.Config{Registry: reg, Sender: sender}) - if err != nil { - t.Fatal(err) - } - defer router.Shutdown(context.Background()) - ctx, cancel := context.WithTimeout(t.Context(), 60*time.Second) - defer cancel() - await := func(kind string) proto.Envelope { - t.Helper() - for { - select { - case env := <-sender: - if env.Type == proto.TypeError { - t.Fatalf("native error: %s", env.Payload) - } - if env.Type == kind { - return env - } - case <-ctx.Done(): - t.Fatalf("waiting for %s; evidence %s", kind, home) - } - } - } - nativeID := "" - for index := 0; index < 3; index++ { - run := fmt.Sprintf("run-%d", index) - request := proto.PromptRequestPayload{AgentKind: "codex", Input: proto.TextInput("Look up ticket 42."), RunID: run, AgentStateKey: "native-functions", AgentSessionID: nativeID, StrictResume: true, ReleaseOnCompletion: true, DisableExecutionEnvironment: true, ObserveToolObservations: true, - FunctionTools: []proto.FunctionTool{{Name: "lookup_ticket", Description: "Read a synthetic ticket", Parameters: json.RawMessage(`{"type":"object","properties":{"ticket":{"type":"string"}},"required":["ticket"],"additionalProperties":false}`)}}, - Model: "gpt-5.5", ModelProvider: &modelprovider.Provider{Protocol: modelprovider.Responses, BaseURL: model.URL + "/v1", APIKey: "synthetic-local-token"}} - env, _ := proto.NewEnvelope(proto.TypePromptRequest, run, request) - if err := router.Handle(ctx, env); err != nil { - t.Fatal(err) - } - call := await(proto.TypeFunctionCall) - var payload proto.FunctionCallPayload - if err := call.DecodePayload(&payload); err != nil || call.ID != run || payload.Name != "lookup_ticket" { - t.Fatal(call, err) - } - if index == 2 { - cancelFrame, _ := proto.NewEnvelope(proto.TypePromptCancel, run, proto.PromptCancelPayload{DeliveryID: "cancel"}) - if err := router.Handle(ctx, cancelFrame); err != nil { - t.Fatal(err) - } - ack := await(proto.TypeInteractionDecisionAck) - var receipt proto.InteractionDecisionAckPayload - _ = ack.DecodePayload(&receipt) - if !receipt.Applied { - t.Fatal(receipt) - } - late, _ := proto.NewEnvelope(proto.TypeFunctionResult, run, proto.FunctionResultPayload{CallID: payload.CallID, Success: true, Content: functionResultContent("late"), DeliveryID: "late"}) - if err := router.Handle(ctx, late); err != nil { - t.Fatal(err) - } - _ = await(proto.TypeInteractionDecisionAck).DecodePayload(&receipt) - if receipt.Applied || receipt.ErrorCode != "not_pending" { - t.Fatal(receipt) - } - break - } - result, _ := proto.NewEnvelope(proto.TypeFunctionResult, run, proto.FunctionResultPayload{CallID: payload.CallID, Success: index == 0, Content: functionResultContent("TICKET-RESULT"), DeliveryID: "result"}) - if err := router.Handle(ctx, result); err != nil { - t.Fatal(err) - } - ack := await(proto.TypeInteractionDecisionAck) - var receipt proto.InteractionDecisionAckPayload - _ = ack.DecodePayload(&receipt) - if !receipt.Applied { - t.Fatal(receipt) - } - done := await(proto.TypeDone) - var output proto.DonePayload - _ = done.DecodePayload(&output) - id, _ := output.Metadata[proto.DoneMetaAgentSessionID].(string) - if output.Content != "FUNCTION-OK" || id == "" || (nativeID != "" && id != nativeID) { - t.Fatal(output) - } - nativeID = id - } - t.Logf("Native function success/failure, fresh-process resume, cancellation and late-result rejection passed; evidence %s", home) -} diff --git a/apps/daemon/internal/dispatch/functions_test.go b/apps/daemon/internal/dispatch/functions_test.go index 663d7a3c0..e668436b8 100644 --- a/apps/daemon/internal/dispatch/functions_test.go +++ b/apps/daemon/internal/dispatch/functions_test.go @@ -5,6 +5,7 @@ import ( "context" "encoding/base64" "encoding/json" + "errors" "image" "image/color" "image/png" @@ -36,7 +37,7 @@ func TestFunctionReceiptsScopeRetriesAndConflicts(t *testing.T) { reg := agent.NewRegistry() sender := &recSender{} sessions := map[string]*functionSession{} - reg.RegisterKind(proto.SupportedAgentKind{Kind: "function-test", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{FunctionTools: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, p proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(reg, proto.SupportedAgentKind{Kind: "function-test", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{FunctionTools: proto.CapabilitySupported})}, func(ctx context.Context, p proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { s := &functionSession{fakeSession: &fakeSession{out: out, ctx: ctx, closeOutOnCancel: true}} sessions[p.RunID] = s return s, nil @@ -47,14 +48,11 @@ func TestFunctionReceiptsScopeRetriesAndConflicts(t *testing.T) { } defer router.Shutdown(context.Background()) for _, id := range []string{"one", "two"} { - env, _ := proto.NewEnvelope(proto.TypePromptRequest, id, prototest.WithModel(proto.PromptRequestPayload{AgentKind: "function-test", Input: proto.TextInput("lookup"), FunctionTools: []proto.FunctionTool{{Name: "lookup", Parameters: json.RawMessage(`{}`)}}})) - if err := router.Handle(t.Context(), env); err != nil { - t.Fatal(err) - } + startRun(t, router, sender, "function-test", id) } submit := func(run, call, text, delivery string) proto.InteractionDecisionAckPayload { t.Helper() - env, _ := proto.NewEnvelope(proto.TypeFunctionResult, run, proto.FunctionResultPayload{CallID: call, Success: true, Content: functionResultContent(text), DeliveryID: delivery}) + env := scoped(t, run, proto.TypeFunctionResult, run, proto.FunctionResultPayload{CallID: call, Success: true, Content: functionResultContent(text), DeliveryID: delivery}) if err := router.Handle(t.Context(), env); err != nil { t.Fatal(err) } @@ -73,7 +71,7 @@ func TestFunctionReceiptsScopeRetriesAndConflicts(t *testing.T) { t.Fatal(a) } - invalid, _ := proto.NewEnvelope(proto.TypeFunctionResult, "one", proto.FunctionResultPayload{CallID: "call", DeliveryID: "invalid", Content: []proto.InputContent{{Type: "input_audio"}}}) + invalid := scoped(t, "one", proto.TypeFunctionResult, "one", proto.FunctionResultPayload{CallID: "call", DeliveryID: "invalid", Content: []proto.InputContent{{Type: "input_audio"}}}) if err := router.Handle(t.Context(), invalid); err != nil { t.Fatal(err) } @@ -87,7 +85,7 @@ func TestFunctionReceiptsScopeRetriesAndConflicts(t *testing.T) { sender.mu.Lock() sender.failNow = true sender.mu.Unlock() - first, _ := proto.NewEnvelope(proto.TypeFunctionResult, "one", proto.FunctionResultPayload{CallID: "call", Success: true, Content: functionResultContent("answer"), DeliveryID: "lost"}) + first := scoped(t, "one", proto.TypeFunctionResult, "one", proto.FunctionResultPayload{CallID: "call", Success: true, Content: functionResultContent("answer"), DeliveryID: "lost"}) if err := router.Handle(t.Context(), first); err == nil { t.Fatal("receipt send failure was hidden") } @@ -109,7 +107,7 @@ func TestFunctionReceiptsScopeRetriesAndConflicts(t *testing.T) { case "success": result.Success = false } - env, _ := proto.NewEnvelope(proto.TypeFunctionResult, "one", result) + env := scoped(t, "one", proto.TypeFunctionResult, "one", result) if err := router.Handle(t.Context(), env); err != nil { t.Fatal(err) } @@ -133,19 +131,31 @@ func TestFunctionReceiptsScopeRetriesAndConflicts(t *testing.T) { } } -func TestFunctionToolsRequireAdvertisedSupport(t *testing.T) { +// rejectsBeforeFactory reports whether a Router rejects preparing req, whose +// kind declares caps, before the kind's factory. +func rejectsBeforeFactory(t *testing.T, caps proto.AgentKindCapabilities, req proto.PromptRequestPayload) bool { + t.Helper() reg := agent.NewRegistry() called := false - reg.RegisterKind(proto.SupportedAgentKind{Kind: "unsupported", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}, prototest.ModelConfiguration(), func(context.Context, proto.PromptRequestPayload, chan<- proto.Envelope) (agent.Session, error) { + reg.RegisterKind(proto.SupportedAgentKind{Kind: req.AgentKind, Available: true, Capabilities: caps}, prototest.ModelConfiguration(), func(context.Context, proto.PromptRequestPayload, chan<- proto.Envelope) (agent.Session, error) { + return nil, errors.New("ordinary factory is forbidden") + }) + reg.RegisterExecutor(req.AgentKind, func(context.Context, proto.PromptRequestPayload) (agent.Executor, error) { called = true - return nil, nil + return nil, errors.New("controlled factory stop") }) - sender := &recSender{} - router, _ := dispatch.New(dispatch.Config{Registry: reg, Sender: sender}) + router, _ := dispatch.New(dispatch.Config{Registry: reg, Sender: &recSender{}}) defer router.Shutdown(context.Background()) - env, _ := proto.NewEnvelope(proto.TypePromptRequest, "run", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "unsupported", FunctionTools: []proto.FunctionTool{{Name: "lookup", Parameters: json.RawMessage(`{}`)}}})) - if err := router.Handle(t.Context(), env); err == nil || called { - t.Fatal("unsupported engine silently ignored tools", err) + assign(t, router, preparationSessionID, "") + req.AgentStateKey, req.StrictResume, req.DisableExecutionEnvironment = stateKey(preparationSessionID), true, true + err := router.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "prepare", proto.ExecutionPreparePayload{SessionID: preparationSessionID, Configuration: req})) + return err != nil && !called +} + +func TestFunctionToolsRequireAdvertisedSupport(t *testing.T) { + caps := prototest.Capabilities(proto.AgentKindCapabilities{EnvironmentNone: proto.CapabilitySupported}) + if !rejectsBeforeFactory(t, caps, prototest.WithModel(proto.PromptRequestPayload{AgentKind: "unsupported", FunctionTools: []proto.FunctionTool{{Name: "lookup", Parameters: json.RawMessage(`{}`)}}})) { + t.Fatal("unsupported engine silently ignored tools") } } @@ -162,18 +172,10 @@ func functionResultContent(text string) []proto.InputContent { } func TestDiscoveryCannotReachAnEagerOnlyAdapter(t *testing.T) { - reg := agent.NewRegistry() - called := false - reg.RegisterKind(proto.SupportedAgentKind{Kind: "eager-only", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{FunctionTools: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(context.Context, proto.PromptRequestPayload, chan<- proto.Envelope) (agent.Session, error) { - called = true - return nil, nil - }) - router, _ := dispatch.New(dispatch.Config{Registry: reg, Sender: &recSender{}}) - defer router.Shutdown(context.Background()) + caps := prototest.Capabilities(proto.AgentKindCapabilities{EnvironmentNone: proto.CapabilitySupported, FunctionTools: proto.CapabilitySupported}) for _, search := range []bool{false, true} { - env, _ := proto.NewEnvelope(proto.TypePromptRequest, "discovery", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "eager-only", ToolSearch: search, FunctionTools: []proto.FunctionTool{{Name: "lookup", Parameters: json.RawMessage(`{"type":"object"}`), DeferLoading: true}}})) - if err := router.Handle(t.Context(), env); err == nil || called { - t.Fatal("deferred definitions reached an eager-only adapter", err) + if !rejectsBeforeFactory(t, caps, prototest.WithModel(proto.PromptRequestPayload{AgentKind: "eager-only", ToolSearch: search, FunctionTools: []proto.FunctionTool{{Name: "lookup", Parameters: json.RawMessage(`{"type":"object"}`), DeferLoading: true}}})) { + t.Fatal("deferred definitions reached an eager-only adapter") } } } diff --git a/apps/daemon/internal/dispatch/interaction_decisions.go b/apps/daemon/internal/dispatch/interaction_decisions.go index 3de5b1186..0376d915c 100644 --- a/apps/daemon/internal/dispatch/interaction_decisions.go +++ b/apps/daemon/internal/dispatch/interaction_decisions.go @@ -28,7 +28,7 @@ func (r *Router) handlePermissionDecision(ctx context.Context, env proto.Envelop if err != nil { return fmt.Errorf("dispatch: fingerprint permission_decision: %w", err) } - if handled, err := r.replayAppliedInteractionDecision(ctx, env.ID, payload.DeliveryID, proto.TypePermissionDecision, fingerprint); handled { + if handled, err := r.replayAppliedInteractionDecision(ctx, env, payload.DeliveryID, proto.TypePermissionDecision, fingerprint); handled { return err } @@ -39,9 +39,13 @@ func (r *Router) handlePermissionDecision(ctx context.Context, env proto.Envelop var finishOperation func() if known { state = r.sessions[runID] - if state != nil { - session, finishOperation, _ = r.preparedOperationLocked(state) + } + if state != nil { + if code := r.admitRunLocked(env.Assignment, state); code != "" { + r.mu.Unlock() + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, code, "The run's assignment does not admit this decision.") } + session, finishOperation, _ = r.preparedOperationLocked(state) } r.mu.Unlock() @@ -49,38 +53,38 @@ func (r *Router) handlePermissionDecision(ctx context.Context, env proto.Envelop // Server's perm timeout / cancel race; common enough that info // is right. r.log.InfoContext(ctx, "permission_decision for unknown perm (run gone)", "perm_id", env.ID) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "not_pending", "permission request is no longer pending") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "not_pending", "permission request is no longer pending") } if session == nil { - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "not_ready", "permission request is waiting for the native session") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "not_ready", "permission request is waiting for the native session") } defer finishOperation() ctx, stop := r.shutdownContext(ctx) defer stop() if !state.capabilities.Permissions.IsSupported() { - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "unsupported", "The runtime declaration does not support interaction decisions.") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "unsupported", "The runtime declaration does not support interaction decisions.") } responder, supported := session.(agent.PermissionResponder) if !supported { r.dropPermission(state, env.ID) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "contract_violation", "Declared runtime capability does not implement permission responses") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "contract_violation", "Declared runtime capability does not implement permission responses") } if err := responder.SubmitPermission(ctx, env.ID, payload); err != nil { if errors.Is(err, agent.ErrUnsupportedOperation) { - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "contract_violation", "Declared permission capability has no implementation.") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "contract_violation", "Declared permission capability has no implementation.") } if errors.Is(err, agent.ErrUnknownPermission) { r.log.InfoContext(ctx, "agent reports unknown perm (race with cancel)", "perm_id", env.ID, "run_id", runID) r.dropPermission(state, env.ID) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "not_pending", err.Error()) + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "not_pending", err.Error()) } r.log.WarnContext(ctx, "agent rejected permission decision", "perm_id", env.ID, "run_id", runID, "err", err) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "runtime_error", err.Error()) + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "runtime_error", err.Error()) } r.dropPermission(state, env.ID) r.rememberAppliedInteractionDecision(env.ID, proto.TypePermissionDecision, fingerprint) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, true, "", "") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, true, "", "") } func (r *Router) dropPermission(s *sessionState, permissionID string) { @@ -117,7 +121,7 @@ func (r *Router) handlePromptForUserChoiceDecision(ctx context.Context, env prot if err != nil { return fmt.Errorf("dispatch: fingerprint prompt_for_user_choice_decision: %w", err) } - if handled, err := r.replayAppliedInteractionDecision(ctx, env.ID, payload.DeliveryID, proto.TypePromptForUserChoiceDecision, fingerprint); handled { + if handled, err := r.replayAppliedInteractionDecision(ctx, env, payload.DeliveryID, proto.TypePromptForUserChoiceDecision, fingerprint); handled { return err } @@ -128,65 +132,69 @@ func (r *Router) handlePromptForUserChoiceDecision(ctx context.Context, env prot var finishOperation func() if known { state = r.sessions[runID] - if state != nil { - session, finishOperation, _ = r.preparedOperationLocked(state) + } + if state != nil { + if code := r.admitRunLocked(env.Assignment, state); code != "" { + r.mu.Unlock() + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, code, "The run's assignment does not admit this decision.") } + session, finishOperation, _ = r.preparedOperationLocked(state) } r.mu.Unlock() if !known || state == nil { r.log.InfoContext(ctx, "prompt_for_user_choice_decision for unknown ask (run gone)", "ask_id", env.ID) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "not_pending", "user-input request is no longer pending") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "not_pending", "user-input request is no longer pending") } if session == nil { - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "not_ready", "user-input request is waiting for the native session") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "not_ready", "user-input request is waiting for the native session") } defer finishOperation() ctx, stop := r.shutdownContext(ctx) defer stop() if !state.capabilities.Permissions.IsSupported() { - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "unsupported", "The runtime declaration does not support interaction decisions.") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "unsupported", "The runtime declaration does not support interaction decisions.") } responder, supported := session.(agent.UserChoiceResponder) if !supported { r.dropAsk(state, env.ID) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "contract_violation", "Declared runtime capability does not implement user-choice responses") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "contract_violation", "Declared runtime capability does not implement user-choice responses") } err = responder.SubmitPromptForUserChoice(ctx, env.ID, payload) if err != nil { if errors.Is(err, agent.ErrUnsupportedOperation) { - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "contract_violation", "Declared user-choice capability has no implementation.") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "contract_violation", "Declared user-choice capability has no implementation.") } if errors.Is(err, agent.ErrUnknownAsk) { r.log.InfoContext(ctx, "agent reports unknown ask (race with cancel)", "ask_id", env.ID, "run_id", runID) r.dropAsk(state, env.ID) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "not_pending", err.Error()) + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "not_pending", err.Error()) } // Keep the routing entry for transient runtime failures. Codex, for // example, restores its pending request when a JSON-RPC reply write // fails, so dropping the ask here would turn a retryable error into a // permanent not_pending response on the next attempt. r.log.WarnContext(ctx, "agent rejected user-input decision", "ask_id", env.ID, "run_id", runID, "err", err) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, false, "runtime_error", err.Error()) + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "runtime_error", err.Error()) } r.dropAsk(state, env.ID) r.rememberAppliedInteractionDecision(env.ID, proto.TypePromptForUserChoiceDecision, fingerprint) - return r.sendInteractionDecisionAck(ctx, env.ID, payload.DeliveryID, true, "", "") + return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, true, "", "") } -func (r *Router) replayAppliedInteractionDecision(ctx context.Context, requestID, deliveryID, kind string, fingerprint [32]byte) (bool, error) { - key := appliedInteractionDecisionKey(requestID, kind) +func (r *Router) replayAppliedInteractionDecision(ctx context.Context, env proto.Envelope, deliveryID, kind string, fingerprint [32]byte) (bool, error) { + key := appliedInteractionDecisionKey(env.ID, kind) r.mu.Lock() applied, ok := r.applied[key] r.mu.Unlock() if !ok { return false, nil } - if applied.requestID != requestID || applied.kind != kind || applied.fingerprint != fingerprint { - return true, r.sendInteractionDecisionAck(ctx, requestID, deliveryID, false, "decision_conflict", "request was already applied with a different decision") + if applied.requestID != env.ID || applied.kind != kind || applied.fingerprint != fingerprint { + return true, r.sendInteractionDecisionAck(ctx, env, deliveryID, false, "decision_conflict", "request was already applied with a different decision") } - return true, r.sendInteractionDecisionAck(ctx, requestID, deliveryID, true, "", "") + return true, r.sendInteractionDecisionAck(ctx, env, deliveryID, true, "", "") } func (r *Router) rememberAppliedInteractionDecision(requestID, kind string, fingerprint [32]byte) { @@ -241,8 +249,8 @@ func interactionDecisionFingerprint(payload any) ([32]byte, error) { } } -func (r *Router) sendInteractionDecisionAck(ctx context.Context, requestID, deliveryID string, applied bool, errorCode, message string) error { - env, err := proto.NewEnvelope(proto.TypeInteractionDecisionAck, requestID, proto.InteractionDecisionAckPayload{ +func (r *Router) sendInteractionDecisionAck(ctx context.Context, request proto.Envelope, deliveryID string, applied bool, errorCode, message string) error { + env, err := request.Reply(proto.TypeInteractionDecisionAck, proto.InteractionDecisionAckPayload{ DeliveryID: deliveryID, Applied: applied, ErrorCode: errorCode, diff --git a/apps/daemon/internal/dispatch/local_directory_test.go b/apps/daemon/internal/dispatch/local_directory_test.go index 467a3c13f..1a0a786aa 100644 --- a/apps/daemon/internal/dispatch/local_directory_test.go +++ b/apps/daemon/internal/dispatch/local_directory_test.go @@ -1,7 +1,5 @@ package dispatch_test -import "github.com/MiniMax-AI/OpenAgentCore/internal/harnessconfig" - import ( "context" "errors" @@ -16,13 +14,14 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/localworkspace" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" + "github.com/MiniMax-AI/OpenAgentCore/internal/harnessconfig" "github.com/google/uuid" ) func TestLocalDirectoryPreparationNeedsNoHarnessAndRejectsOtherOwners(t *testing.T) { workspace := t.TempDir() - environment, session := uuid.NewString(), uuid.NewString() + environment, session := uuid.NewString(), preparationSessionID binding, err := localworkspace.New(environment, session, workspace) if err != nil { t.Fatal(err) @@ -43,8 +42,9 @@ func TestLocalDirectoryPreparationNeedsNoHarnessAndRejectsOtherOwners(t *testing t.Fatal(err) } t.Cleanup(func() { _ = r.Shutdown(context.Background()) }) + assign(t, r, session, environment) request := proto.PromptRequestPayload{AgentKind: "native", LocalEnvironment: &proto.LocalEnvironment{ID: environment}, AgentStateKey: "agents-api-" + session, StrictResume: true, ReleaseOnCompletion: true, WorkspaceReadOnly: true} - if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "idle", proto.ExecutionPreparePayload{Configuration: request})); err != nil { + if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "idle", proto.ExecutionPreparePayload{SessionID: session, Configuration: request})); err != nil { t.Fatal(err) } ready := waitPreparationStatus(t, sender, "idle", "ready", "") @@ -71,7 +71,7 @@ func TestLocalDirectoryPreparationNeedsNoHarnessAndRejectsOtherOwners(t *testing } bad := request bad.AgentStateKey = "agents-api-" + uuid.NewString() - if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "wrong-session", proto.ExecutionPreparePayload{Configuration: bad})); err == nil { + if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "wrong-session", proto.ExecutionPreparePayload{SessionID: session, Configuration: bad})); err == nil { t.Fatal("wrong Session accepted") } _ = r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionRelease, "idle", proto.ExecutionReleasePayload{Handle: ready.Handle})) @@ -105,7 +105,7 @@ func TestLocalDirectoryKeepsNotDirectorySeparateFromFailures(t *testing.T) { if err := os.WriteFile(filepath.Join(workspace, "file"), nil, 0600); err != nil { t.Fatal(err) } - environment, session := uuid.NewString(), uuid.NewString() + environment, session := uuid.NewString(), preparationSessionID binding, err := localworkspace.New(environment, session, workspace) if err != nil { t.Fatal(err) @@ -123,8 +123,9 @@ func TestLocalDirectoryKeepsNotDirectorySeparateFromFailures(t *testing.T) { t.Fatal(err) } t.Cleanup(func() { _ = r.Shutdown(context.Background()) }) + assign(t, r, session, environment) request := proto.PromptRequestPayload{AgentKind: "native", LocalEnvironment: &proto.LocalEnvironment{ID: environment}, AgentStateKey: "agents-api-" + session, StrictResume: true, ReleaseOnCompletion: true, WorkspaceReadOnly: true} - if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "idle", proto.ExecutionPreparePayload{Configuration: request})); err != nil { + if err := r.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "idle", proto.ExecutionPreparePayload{SessionID: session, Configuration: request})); err != nil { t.Fatal(err) } ready := waitPreparationStatus(t, sender, "idle", "ready", "") diff --git a/apps/daemon/internal/dispatch/mcp_http_test.go b/apps/daemon/internal/dispatch/mcp_http_test.go index e809961a5..f5ebe8431 100644 --- a/apps/daemon/internal/dispatch/mcp_http_test.go +++ b/apps/daemon/internal/dispatch/mcp_http_test.go @@ -20,7 +20,7 @@ func TestMCPHTTPBearerRejectsUnsupportedRequestsBeforeFactory(t *testing.T) { defer h.router.Shutdown(context.Background()) token := "synthetic-private-token" servers := []proto.MCPHTTPServer{{ConnectionOrigin: "service", ServerLabel: "tools", ServerURL: "https://tools.example/mcp", BearerToken: &token}} - req := prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", DisableExecutionEnvironment: true, MCPHTTPServers: &servers}) + req := prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", AgentStateKey: stateKey(preparationSessionID), StrictResume: true, DisableExecutionEnvironment: true, MCPHTTPServers: &servers}) caps := prototest.Capabilities(proto.AgentKindCapabilities{EnvironmentNone: proto.CapabilitySupported, MCPHTTPTools: proto.CapabilitySupported, MCPHTTPBearerAuth: proto.CapabilitySupported}) switch mode { case "claude", "claude old peer", "claude local": @@ -54,24 +54,33 @@ func TestMCPHTTPBearerRejectsUnsupportedRequestsBeforeFactory(t *testing.T) { } called := false h.reg.RegisterKind(proto.SupportedAgentKind{Kind: req.AgentKind, Available: mode != "unavailable", Capabilities: caps}, - prototest.ModelConfiguration(), func(_ context.Context, got proto.PromptRequestPayload, _ chan<- proto.Envelope) (agent.Session, error) { - called = true - if mode == "required" && !(*got.MCPHTTPServers)[0].Required { - t.Error("required initialization lost before adapter") - } - if (mode == "supported" || mode == "claude") && (got.MCPHTTPServers == nil || (*got.MCPHTTPServers)[0].BearerToken == nil || *(*got.MCPHTTPServers)[0].BearerToken != token) { - t.Error("token lost before adapter") - } - return nil, errors.New("controlled factory stop") + prototest.ModelConfiguration(), func(context.Context, proto.PromptRequestPayload, chan<- proto.Envelope) (agent.Session, error) { + return nil, errors.New("ordinary factory is forbidden") }) - err := h.router.Handle(t.Context(), mustEnv(t, proto.TypePromptRequest, "mcp-bearer", req)) - if called != (mode == "supported" || mode == "claude" || mode == "other engine" || mode == "credential-free" || mode == "product" || mode == "required" || mode == "optional old peer") { - t.Fatal("wrong factory admission") + h.reg.RegisterExecutor(req.AgentKind, func(_ context.Context, got proto.PromptRequestPayload) (agent.Executor, error) { + called = true + if mode == "required" && !(*got.MCPHTTPServers)[0].Required { + t.Error("required initialization lost before adapter") + } + if (mode == "supported" || mode == "claude") && (got.MCPHTTPServers == nil || (*got.MCPHTTPServers)[0].BearerToken == nil || *(*got.MCPHTTPServers)[0].BearerToken != token) { + t.Error("token lost before adapter") + } + return nil, errors.New("controlled factory stop") + }) + assign(t, h.router, preparationSessionID, "") + err := h.router.Handle(t.Context(), mustEnv(t, proto.TypeExecutionPrepare, "mcp-bearer", proto.ExecutionPreparePayload{SessionID: preparationSessionID, Configuration: req})) + admitted := mode == "supported" || mode == "claude" || mode == "other engine" || mode == "credential-free" || mode == "product" || mode == "required" || mode == "optional old peer" + state := "rejected" + if admitted { + state = "failed" + } + waitPreparationStatus(t, h.sender, "mcp-bearer", state, "") + if called != admitted || (err == nil) != admitted { + t.Fatal("wrong factory admission", err) } - frames := h.sender.snapshot() - raw, _ := json.Marshal(frames) - if err == nil || strings.Contains(err.Error(), token) || strings.Contains(string(raw), token) || len(frames) != 2 || frames[0].Type != proto.TypeError || frames[1].Type != proto.TypeDone { - t.Fatal("terminal rejection missing or exposed credential") + raw, _ := json.Marshal(h.sender.snapshot()) + if err != nil && strings.Contains(err.Error(), token) || strings.Contains(string(raw), token) { + t.Fatal("rejection exposed the credential") } }) } diff --git a/apps/daemon/internal/dispatch/optional_interactions_test.go b/apps/daemon/internal/dispatch/optional_interactions_test.go deleted file mode 100644 index 84f73984b..000000000 --- a/apps/daemon/internal/dispatch/optional_interactions_test.go +++ /dev/null @@ -1,45 +0,0 @@ -package dispatch_test - -import ( - "context" - "testing" - - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" -) - -// Wrapping Session proves that no responder stubs are required for a Session. -type lifecycleOnly struct{ agent.Session } - -func TestOptionalInteractionResponders(t *testing.T) { - for _, ask := range []bool{false, true} { - t.Run(map[bool]string{false: "permission", true: "user choice"}[ask], func(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - var output chan<- proto.Envelope - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "minimal", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}, prototest.ModelConfiguration(), func(_ context.Context, _ proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - output = out - s := &fakeSession{out: out, closeOutOnCancel: true} - return lifecycleOnly{Session: s}, nil - }) - if err := h.router.Handle(t.Context(), mustEnv(t, proto.TypePromptRequest, "run", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "minimal"}))); err != nil { - t.Fatal(err) - } - event := mustEnv(t, proto.TypePermissionRequest, "run", proto.PermissionRequestPayload{RequestID: "interaction", Tool: "fixture"}) - decision := mustEnv(t, proto.TypePermissionDecision, "interaction", proto.PermissionDecisionPayload{DeliveryID: "decision", Approved: true}) - if ask { - event = mustEnv(t, proto.TypePromptForUserChoice, "run", proto.PromptForUserChoicePayload{AskID: "interaction"}) - decision = mustEnv(t, proto.TypePromptForUserChoiceDecision, "interaction", proto.PromptForUserChoiceDecisionPayload{DeliveryID: "decision"}) - } - // An inconsistent adapter emitted an interaction it cannot answer: reject it, - // never acknowledge application or call a fabricated responder. - output <- event - waitFor(t, func() bool { return len(h.sender.snapshot()) > 0 }, "interaction indexed") - if err := h.router.Handle(t.Context(), decision); err != nil { - t.Fatal(err) - } - assertDecisionAck(t, h.sender, "decision", false, "unsupported") - }) - } -} diff --git a/apps/daemon/internal/dispatch/output.go b/apps/daemon/internal/dispatch/output.go deleted file mode 100644 index e0e96dca9..000000000 --- a/apps/daemon/internal/dispatch/output.go +++ /dev/null @@ -1,103 +0,0 @@ -package dispatch - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - obslog "github.com/MiniMax-AI/OpenAgentCore/internal/obs/log" -) - -// pump forwards every Envelope the session writes onto out to the -// upstream sender, then cleans up when the session closes out. -func (r *Router) pump(s *sessionState) { - defer r.shutdownWG.Done() - defer r.cleanupSession(s) - _ = r.forwardSessionOutput(s, true) -} - -// wait is false only after a failed Start and its cancellation have returned: -// no Session owns the output sink, so only already queued frames can be read. -func (r *Router) forwardSessionOutput(s *sessionState, wait bool) error { - // Logging-only ctx carrying the run's trace; sends use their own - // ctx tied to shutdownCh. - pumpCtx := context.Background() - if s.traceparent != "" { - if carrier, err := obslog.ParseTraceparent(s.traceparent); err == nil { - pumpCtx = obslog.WithTrace(pumpCtx, carrier) - } - } - r.log.InfoContext(pumpCtx, "pump: started", "run_id", s.runID) - - // Long-lived send ctx — must keep forwarding even after the - // session's ctx is cancelled (session might emit a final "done" - // in response to cancel). Stops on out close or router shutdown. - for { - if !wait && len(s.out) == 0 { - return nil - } - select { - case env, ok := <-s.out: - if !ok { - r.log.InfoContext(pumpCtx, "pump: out channel closed", "run_id", s.runID) - return nil - } - if s.session != nil { - r.indexPermissionFrame(s, env) - } - if env.Type == proto.TypeDone && s.releaseOnCompletion { - if err := r.releaseCompletedSession(s); err != nil { - sendCtx, stop := r.shutdownContext(pumpCtx) - r.emitTerminalError(sendCtx, s.runID, "failed to release completed executor") - stop() - continue - } - } - r.log.InfoContext(pumpCtx, "pump: forwarding envelope", "run_id", s.runID, "type", env.Type, "env_id", env.ID) - // Stamp the run's trace onto outbound frames so the - // gateway attributes daemon-emitted lines to the same - // trace_id as the original prompt_request. - if env.Trace == "" && s.traceparent != "" { - env.Trace = s.traceparent - } - // Short-ish send ctx that respects router shutdown — if - // the transport is wedged we don't want to block forever. - sendCtx, cancel := context.WithCancel(context.Background()) - stopOnShutdown := make(chan struct{}) - go func() { - select { - case <-r.shutdownCh: - cancel() - case <-stopOnShutdown: - } - }() - err := r.sender.Send(sendCtx, env) - close(stopOnShutdown) - cancel() - if err != nil { - // Sender failed — log, ask the session to wind down, - // but KEEP draining out so the agent's goroutines - // don't block on a full channel. - r.log.ErrorContext(pumpCtx, "send envelope failed", "type", env.Type, "run_id", env.ID, "err", err) - r.mu.Lock() - s.retain = false - r.mu.Unlock() - s.ctxCancel() - if wait { - r.drain(s.out) - } - return err - } - case <-r.shutdownCh: - // Router shutdown — cancel + drain so the session's - // goroutines unblock and close out cleanly. - r.mu.Lock() - s.retain = false - r.mu.Unlock() - s.ctxCancel() - if wait { - r.drain(s.out) - } - return ErrRouterClosed - } - } -} diff --git a/apps/daemon/internal/dispatch/preparation.go b/apps/daemon/internal/dispatch/preparation.go index 9f72d48d1..0bcbaa630 100644 --- a/apps/daemon/internal/dispatch/preparation.go +++ b/apps/daemon/internal/dispatch/preparation.go @@ -19,10 +19,11 @@ const preparationRecords = 64 // All mutable fields are protected by Router.mu. owns includes resources whose // cancellation is underway; a slow close cannot bypass the capacity bound. type preparationState struct { - capabilities proto.AgentKindCapabilities - executor *executorState - requestID string - trace string + capabilities proto.AgentKindCapabilities + executor *executorState + // request is the execution_prepare's ID, trace and assignment, which + // every status echoes. + request proto.Envelope fingerprint [32]byte startFingerprint [32]byte status proto.PreparationStatusPayload @@ -31,7 +32,6 @@ type preparationState struct { ctx context.Context cancel context.CancelFunc prepared io.Closer - stateKey string environmentID string busy bool owns bool @@ -79,7 +79,11 @@ func (r *Router) handleExecutionPrepare(ctx context.Context, env proto.Envelope) r.mu.Unlock() return ErrRouterClosed } - if r.runtimePreparation != nil || ((r.workspaceWrite != nil || r.workspaceExport != nil) && !req.WorkspaceReadOnly) { + if code := r.admitLocked(env.Assignment, input.SessionID, req.EnvironmentID()); code != "" { + r.mu.Unlock() + return r.rejectPreparation(env, code) + } + if r.runtimePreparation != nil { r.mu.Unlock() return r.rejectPreparation(env, "resource_unavailable") } @@ -105,9 +109,9 @@ func (r *Router) handleExecutionPrepare(ctx context.Context, env proto.Envelope) return r.rejectPreparation(env, "preparation_capacity") } owner, cancel := context.WithCancel(context.WithoutCancel(ctx)) - p := &preparationState{capabilities: caps, requestID: env.ID, trace: env.Trace, fingerprint: fingerprint, ctx: owner, cancel: cancel, stateKey: req.AgentStateKey, environmentID: req.EnvironmentID(), workspaceReadOnly: req.WorkspaceReadOnly, busy: true, owns: true, deadline: time.Now().Add(r.preparationTimeout)} + p := &preparationState{capabilities: caps, request: proto.Envelope{ID: env.ID, Trace: env.Trace, Assignment: env.Assignment}, fingerprint: fingerprint, ctx: owner, cancel: cancel, environmentID: req.EnvironmentID(), workspaceReadOnly: req.WorkspaceReadOnly, busy: true, owns: true, deadline: time.Now().Add(r.preparationTimeout)} p.status = proto.PreparationStatusPayload{Handle: uuid.NewString(), Revision: 1, State: "preparing", ExpiresAt: p.deadline.UnixMilli()} - r.preparations[p.status.Handle], r.preparationRequests[p.requestID] = p, p + r.preparations[p.status.Handle], r.preparationRequests[p.request.ID] = p, p p.timer = time.AfterFunc(r.preparationTimeout, func() { r.releasePreparation(p, "expired", "", true, true) }) r.shutdownWG.Add(1) r.mu.Unlock() @@ -118,7 +122,7 @@ func (r *Router) handleExecutionPrepare(ctx context.Context, env proto.Envelope) // prepareExecution readies a read-only preparation without starting a Harness. func (r *Router) prepareExecution(p *preparationState) { defer r.shutdownWG.Done() - if !r.sendPreparation(p.requestID, p.trace, proto.PreparationStatusPayload{Handle: p.status.Handle, Revision: 1, State: "preparing", ExpiresAt: p.deadline.UnixMilli()}) { + if !r.sendPreparation(p.request, proto.PreparationStatusPayload{Handle: p.status.Handle, Revision: 1, State: "preparing", ExpiresAt: p.deadline.UnixMilli()}) { r.releasePreparation(p, "failed", "status_delivery_failed", false, false) } r.mu.Lock() @@ -141,7 +145,7 @@ func (r *Router) prepareExecution(p *preparationState) { r.closePreparationResource(p) return } - if !r.sendPreparation(p.requestID, p.trace, status) { + if !r.sendPreparation(p.request, status) { r.releasePreparation(p, "failed", "status_delivery_failed", false, false) } } @@ -153,7 +157,7 @@ func (r *Router) handleExecutionRelease(_ context.Context, env proto.Envelope) e } r.mu.Lock() p := r.preparations[input.Handle] - valid := p != nil && p.requestID == env.ID + valid := p != nil && p.request.ID == env.ID && p.request.Assignment == env.Assignment r.mu.Unlock() if !valid { return r.rejectPreparation(env, "unknown_preparation") @@ -227,14 +231,14 @@ func (r *Router) prunePreparationsLocked() { } if time.Now().After(p.deadline) { delete(r.preparations, handle) - delete(r.preparationRequests, p.requestID) + delete(r.preparationRequests, p.request.ID) } else if oldest == nil || p.deadline.Before(oldest.deadline) { oldest = p } } if len(r.preparations) >= preparationRecords && oldest != nil { delete(r.preparations, oldest.status.Handle) - delete(r.preparationRequests, oldest.requestID) + delete(r.preparationRequests, oldest.request.ID) } } @@ -255,22 +259,22 @@ func (r *Router) publishPreparation(p *preparationState, status proto.Preparatio go func() { defer r.shutdownWG.Done() // A failed terminal notification must not restart incomplete cleanup. - if !r.sendPreparation(p.requestID, p.trace, status) && (!p.workspaceReadOnly || status.State == "preparing" || status.State == "ready") { + if !r.sendPreparation(p.request, status) && (!p.workspaceReadOnly || status.State == "preparing" || status.State == "ready") { r.releasePreparation(p, "failed", "status_delivery_failed", false, false) } }() } -func (r *Router) sendPreparation(requestID, trace string, status proto.PreparationStatusPayload) bool { - return r.sendPreparationUntil(requestID, trace, status, time.Now().Add(5*time.Second)) +func (r *Router) sendPreparation(request proto.Envelope, status proto.PreparationStatusPayload) bool { + return r.sendPreparationUntil(request, status, time.Now().Add(5*time.Second)) } -func (r *Router) sendPreparationUntil(requestID, trace string, status proto.PreparationStatusPayload, deadline time.Time) bool { +func (r *Router) sendPreparationUntil(request proto.Envelope, status proto.PreparationStatusPayload, deadline time.Time) bool { ctx, stop := r.shutdownContext(context.Background()) defer stop() ctx, cancel := context.WithDeadline(ctx, deadline) defer cancel() - env, err := proto.NewEnvelopeWithTrace(proto.TypePreparationStatus, requestID, status, trace) + env, err := request.Reply(proto.TypePreparationStatus, status) return err == nil && r.sender.Send(ctx, env) == nil } @@ -280,7 +284,7 @@ func (r *Router) rejectPreparation(env proto.Envelope, code string) error { r.shutdownWG.Add(1) go func() { defer r.shutdownWG.Done() - r.sendPreparation(env.ID, env.Trace, proto.PreparationStatusPayload{State: "rejected", ErrorCode: code, Operation: env.Type}) + r.sendPreparation(env, proto.PreparationStatusPayload{State: "rejected", ErrorCode: code, Operation: env.Type}) }() } r.mu.Unlock() diff --git a/apps/daemon/internal/dispatch/preparation_start.go b/apps/daemon/internal/dispatch/preparation_start.go index d1bd51d76..089338a70 100644 --- a/apps/daemon/internal/dispatch/preparation_start.go +++ b/apps/daemon/internal/dispatch/preparation_start.go @@ -29,10 +29,18 @@ func (r *Router) handleExecutionStart(_ context.Context, env proto.Envelope) err return r.rejectPreparation(env, "resource_unavailable") } p := r.preparations[input.Handle] - if p == nil || p.requestID != env.ID { + if p == nil || p.request.ID != env.ID { r.mu.Unlock() return r.rejectPreparation(env, "unknown_preparation") } + code := proto.AssignmentConflict + if p.request.Assignment == env.Assignment { + code = r.admitLocked(env.Assignment, env.Assignment.SessionID, p.environmentID) + } + if code != "" { + r.mu.Unlock() + return r.rejectPreparation(env, code) + } if p.workspaceReadOnly { r.mu.Unlock() return r.rejectPreparation(env, "read_only_preparation") @@ -70,7 +78,7 @@ func (r *Router) handleExecutionStart(_ context.Context, env proto.Envelope) err } p.status.State, p.status.RunID, p.status.Revision = "starting", input.RunID, p.status.Revision+1 p.startFingerprint, p.busy = fingerprint, true - state := &sessionState{capabilities: p.capabilities, runID: input.RunID, stateKey: p.stateKey, environmentID: p.environmentID, out: make(chan proto.Envelope, 64), ctx: p.ctx, ctxCancel: p.cancel, pendingIDs: make(map[string]struct{}), pendingAsks: make(map[string]struct{}), traceparent: env.Trace} + state := &sessionState{assignment: env.Assignment, capabilities: p.capabilities, runID: input.RunID, environmentID: p.environmentID, out: make(chan proto.Envelope, 64), ctx: p.ctx, ctxCancel: p.cancel, pendingIDs: make(map[string]struct{}), pendingAsks: make(map[string]struct{}), traceparent: env.Trace} state.preparedHandoff = newPreparedHandoff(p, owner.native) p.handoff, owner.run = state.preparedHandoff, state r.sessions[input.RunID] = state @@ -86,7 +94,7 @@ func (r *Router) startPreparedExecution(p *preparationState, state *sessionState handoff := state.preparedHandoff go r.forwardPreparedOutput(state) <-handoff.outputReady - delivered := r.sendPreparation(p.requestID, p.trace, starting) + delivered := r.sendPreparation(p.request, starting) r.mu.Lock() blocked := !delivered || handoff.release != nil && handoff.release.aborted() || p.ctx.Err() != nil r.mu.Unlock() @@ -129,7 +137,7 @@ func (r *Router) startPreparedExecution(p *preparationState, state *sessionState } status = proto.PreparationStatusPayload{Handle: p.status.Handle, ExecutorID: p.executor.id, State: "rejected", ErrorCode: code, Operation: proto.TypeExecutionStart} if delivered { - r.sendPreparation(p.requestID, p.trace, status) + r.sendPreparation(p.request, status) } return } @@ -139,7 +147,7 @@ func (r *Router) startPreparedExecution(p *preparationState, state *sessionState if aborted { status.State, status.ErrorCode = "failed", "start_cancelled" } - delivered = delivered && !aborted && r.sendPreparationUntil(p.requestID, p.trace, status, p.deadline) + delivered = delivered && !aborted && r.sendPreparationUntil(p.request, status, p.deadline) r.mu.Lock() if !delivered && !aborted { handoff.outputErr = errors.Join(handoff.outputErr, errPreparedStatusDelivery) diff --git a/apps/daemon/internal/dispatch/preparation_test.go b/apps/daemon/internal/dispatch/preparation_test.go index 3fd82eed1..57e89950c 100644 --- a/apps/daemon/internal/dispatch/preparation_test.go +++ b/apps/daemon/internal/dispatch/preparation_test.go @@ -107,11 +107,15 @@ func localPreparationHarness(t *testing.T) *harness { if err != nil { t.Fatal(err) } + assign(t, h.router, preparationSessionID, preparationEnvironmentID) return h } +// stateKey is the AgentStateKey Core derives for the Session. +func stateKey(session string) string { return "agents-api-" + session } + func preparationRequest() proto.ExecutionPreparePayload { - return proto.ExecutionPreparePayload{SessionID: preparationSessionID, Configuration: prototest.WithModel(proto.PromptRequestPayload{AgentKind: "prepared", AgentStateKey: "agents-api-" + preparationSessionID, StrictResume: true, ReleaseOnCompletion: true, LocalEnvironment: &proto.LocalEnvironment{ID: preparationEnvironmentID, NetworkAccess: "enabled", WorkspaceDirectory: "/workspace", CapabilitySources: &agentcapabilities.Input{}}})} + return proto.ExecutionPreparePayload{SessionID: preparationSessionID, Configuration: prototest.WithModel(proto.PromptRequestPayload{AgentKind: "prepared", AgentStateKey: stateKey(preparationSessionID), StrictResume: true, ReleaseOnCompletion: true, LocalEnvironment: &proto.LocalEnvironment{ID: preparationEnvironmentID, NetworkAccess: "enabled", WorkspaceDirectory: "/workspace", CapabilitySources: &agentcapabilities.Input{}}})} } func preparationRouter(t *testing.T, sender dispatch.Sender, timeout time.Duration, factory preparationFactory) *dispatch.Router { @@ -125,6 +129,7 @@ func preparationRouter(t *testing.T, sender dispatch.Sender, timeout time.Durati if err != nil { t.Fatal(err) } + assign(t, r, preparationSessionID, preparationEnvironmentID) t.Cleanup(func() { ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) defer cancel() diff --git a/apps/daemon/internal/dispatch/prepared_handoff.go b/apps/daemon/internal/dispatch/prepared_handoff.go index 966f93d36..d4c232b18 100644 --- a/apps/daemon/internal/dispatch/prepared_handoff.go +++ b/apps/daemon/internal/dispatch/prepared_handoff.go @@ -111,7 +111,6 @@ func (r *Router) claimPreparedReleaseLocked(state *sessionState, abort bool, fai if release == nil { release = &preparedRelease{abort: make(chan struct{}), failure: failure, settled: make(chan struct{})} handoff.release = release - state.retain = false state.steeringClosed = true state.session = nil r.clearInteractionRoutesLocked(state) @@ -364,6 +363,7 @@ func (r *Router) sendSessionOutput(pumpCtx context.Context, state *sessionState, if env.Trace == "" && state.traceparent != "" { env.Trace = state.traceparent } + env.Assignment = state.assignment r.log.InfoContext(pumpCtx, "pump: forwarding envelope", "run_id", state.runID, "type", env.Type, "env_id", env.ID) sendCtx, cancel := context.WithCancel(context.Background()) stopOnShutdown := make(chan struct{}) diff --git a/apps/daemon/internal/dispatch/prompt.go b/apps/daemon/internal/dispatch/prompt.go deleted file mode 100644 index 7cfc864c6..000000000 --- a/apps/daemon/internal/dispatch/prompt.go +++ /dev/null @@ -1,148 +0,0 @@ -package dispatch - -import ( - "context" - "errors" - "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - obslog "github.com/MiniMax-AI/OpenAgentCore/internal/obs/log" -) - -func (r *Router) handlePromptRequest(callerCtx context.Context, env proto.Envelope) error { - r.log.InfoContext(callerCtx, "handlePromptRequest: decoding payload", "env_id", env.ID, "env_type", env.Type) - var req proto.PromptRequestPayload - if err := env.DecodeRequest(&req); err != nil { - r.log.ErrorContext(callerCtx, "handlePromptRequest: decode failed", "env_id", env.ID, "err", err) - return fmt.Errorf("dispatch: decode prompt_request: %w", err) - } - // Envelope.ID is the run id; payload mirrors it but envelope wins. - runID := env.ID - if runID == "" { - runID = req.RunID - } - if runID == "" { - r.log.ErrorContext(callerCtx, "handlePromptRequest: missing run id") - return errors.New("dispatch: prompt_request missing run id (Envelope.ID and Payload.RunID both empty)") - } - req.RunID = runID - var err error - if req, err = r.localWorkspace.Configure(req); err != nil { - r.emitTerminalError(callerCtx, runID, err.Error()) - return err - } - if req.AgentKind == "" { - r.log.ErrorContext(callerCtx, "handlePromptRequest: missing agent_kind", "run_id", runID) - return errors.New("dispatch: prompt_request missing agent_kind") - } - if req.WorkspaceReadOnly { - err := errors.New("read-only preparation cannot accept a prompt") - r.emitTerminalError(callerCtx, runID, err.Error()) - return err - } - caps, available := r.availableCapabilities(req.AgentKind) - if !available { - _, err := r.registry.Resolve(req.AgentKind) - if err == nil { - err = errors.New("engine is unavailable on this runtime") - } - r.emitTerminalError(callerCtx, runID, err.Error()) - return err - } - if len(req.FunctionTools) > 0 && !caps.FunctionTools.IsSupported() { - err := errors.New("engine does not support function tools") - r.emitTerminalError(callerCtx, runID, err.Error()) - return err - } - if err := validateExecutionEnvironment(req, caps); err != nil { - r.emitTerminalError(callerCtx, runID, err.Error()) - return err - } - r.log.InfoContext(callerCtx, "handlePromptRequest: decoded", - "run_id", runID, "agent_kind", req.AgentKind, - "message_count", len(req.Input), - "agent_session_id", req.AgentSessionID, - "agent_state_key", req.AgentStateKey) - - factory, err := r.registry.Resolve(req.AgentKind) - if err != nil { - r.log.ErrorContext(callerCtx, "handlePromptRequest: registry.Resolve failed", "run_id", runID, "agent_kind", req.AgentKind, "err", err) - // Synthesize error+done so the server-side stream closes - // cleanly instead of waiting for a done that never comes. - r.emitTerminalError(callerCtx, runID, fmt.Sprintf("unsupported agent_kind %q on this daemon", req.AgentKind)) - return err - } - r.log.InfoContext(callerCtx, "handlePromptRequest: factory resolved", "run_id", runID, "agent_kind", req.AgentKind) - - // Lock-protect duplicate-run check + insert so two prompt_requests - // with the same RunID can't both start sessions. - r.mu.Lock() - if r.closed { - r.mu.Unlock() - r.log.ErrorContext(callerCtx, "handlePromptRequest: router closed", "run_id", runID) - return ErrRouterClosed - } - if r.runtimePreparation != nil || r.workspaceWrite != nil || r.workspaceExport != nil { - r.mu.Unlock() - err := errors.New("local workspace has an unsettled write") - r.emitTerminalError(callerCtx, runID, err.Error()) - return err - } - if _, dup := r.sessions[runID]; dup { - r.mu.Unlock() - r.log.WarnContext(callerCtx, "ignoring duplicate prompt_request", "run_id", runID) - return nil - } - stateKey := sessionStateKey(req) - r.touchIdleLocked(stateKey) - sessionCtx, sessionCancel := context.WithCancel(context.Background()) - // Re-attach the inbound trace so every log under this run shows - // the same trace_id as the prompt_request that started it. - if carrier, ok := obslog.TraceFromContext(callerCtx); ok { - sessionCtx = obslog.WithTrace(sessionCtx, carrier) - } - out := make(chan proto.Envelope, 64) - state := &sessionState{ - runID: runID, - capabilities: caps, - environmentID: req.EnvironmentID(), - stateKey: stateKey, - out: out, - ctx: sessionCtx, - ctxCancel: sessionCancel, - pendingIDs: make(map[string]struct{}), - pendingAsks: make(map[string]struct{}), - traceparent: env.Trace, - retain: stateKey != "", - releaseOnCompletion: req.ReleaseOnCompletion, - } - r.sessions[runID] = state - r.mu.Unlock() - - r.log.InfoContext(callerCtx, "handlePromptRequest: calling factory", "run_id", runID) - req, err = r.localWorkspace.Prepare(sessionCtx, req) - var sess agent.Session - if err == nil { - sess, err = factory(sessionCtx, req, out) - } - if err != nil { - r.log.ErrorContext(callerCtx, "handlePromptRequest: factory call failed", "run_id", runID, "agent_kind", req.AgentKind, "err", err) - // Roll back the registration, cancel ctx, surface error+done - // so the server doesn't hang on a phantom run. - r.mu.Lock() - delete(r.sessions, runID) - r.mu.Unlock() - sessionCancel() - r.emitTerminalError(callerCtx, runID, fmt.Sprintf("agent factory failed: %v", err)) - return fmt.Errorf("dispatch: factory %q: %w", req.AgentKind, err) - } - r.log.InfoContext(callerCtx, "handlePromptRequest: session created, starting pump", "run_id", runID) - r.mu.Lock() - state.session = sess - r.mu.Unlock() - - r.shutdownWG.Add(1) - go r.pump(state) - return nil -} diff --git a/apps/daemon/internal/dispatch/receipt_order_test.go b/apps/daemon/internal/dispatch/receipt_order_test.go deleted file mode 100644 index 4f43bcc62..000000000 --- a/apps/daemon/internal/dispatch/receipt_order_test.go +++ /dev/null @@ -1,175 +0,0 @@ -package dispatch_test - -import ( - "context" - "errors" - "sync" - "sync/atomic" - "testing" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/dispatch" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" -) - -type blockedReceiptSender struct { - recSender - once sync.Once - entered chan struct{} - release chan struct{} - exited chan struct{} - fail bool -} - -func (s *blockedReceiptSender) Send(ctx context.Context, env proto.Envelope) error { - blocked := false - if env.Type == proto.TypePromptSteerAck { - s.once.Do(func() { blocked = true }) - } - if blocked { - close(s.entered) - defer close(s.exited) - select { - case <-s.release: - case <-ctx.Done(): - return ctx.Err() - } - if s.fail { - return errors.New("receipt transport failed") - } - } - return s.recSender.Send(ctx, env) -} - -func TestDurableCompletionWaitsForSteeringReceiptSend(t *testing.T) { - for _, mode := range []string{"consumed", "unknown", "send_failure"} { - t.Run(mode, func(t *testing.T) { - sender := &blockedReceiptSender{entered: make(chan struct{}), release: make(chan struct{}), exited: make(chan struct{}), fail: mode == "send_failure"} - registry := agent.NewRegistry() - var session *fakeSession - var calls atomic.Int32 - registry.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - session = &fakeSession{out: out, closeOutOnCancel: true} - return &steeringSession{fakeSession: session, steer: func(context.Context, proto.PromptSteerPayload) error { - calls.Add(1) - if mode == "unknown" { - return errors.New("native outcome unknown") - } - return nil - }}, nil - }) - router, err := dispatch.New(dispatch.Config{Registry: registry, Sender: sender}) - if err != nil { - t.Fatal(err) - } - defer router.Shutdown(context.Background()) - var release sync.Once - defer release.Do(func() { close(sender.release) }) - handle := func(kind string, payload any) { - t.Helper() - if err := router.Handle(context.Background(), mustEnv(t, kind, "ordered", payload)); err != nil { - t.Fatal(err) - } - } - handle(proto.TypePromptRequest, prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", AgentStateKey: "stable", ReleaseOnCompletion: true})) - input := proto.PromptSteerPayload{InputID: "input-1", Input: proto.TextInput("original")} - handle(proto.TypePromptSteer, input) - <-sender.entered - session.out <- mustEnv(t, proto.TypeDone, "ordered", proto.DonePayload{Content: "finished"}) - waitFor(t, func() bool { return router.SteeringClosedForTest("ordered") }, "closed steering admission") - if session.cancels() != 0 || len(sender.snapshot()) != 0 { - t.Fatal("native release or Done overtook receipt") - } - if mode != "send_failure" { - handle(proto.TypePromptSteer, input) - frames := sender.snapshot() - var ack proto.PromptSteerAckPayload - if len(frames) != 1 || frames[0].DecodePayload(&ack) != nil || ack.Accepted != (mode == "consumed") { - t.Fatalf("cached receipt lost: %+v", ack) - } - if mode == "unknown" && ack.ErrorCode != "outcome_unknown" { - t.Fatal("uncertainty lost") - } - input.Input = proto.TextInput("changed") - handle(proto.TypePromptSteer, input) - input.InputID = "new" - handle(proto.TypePromptSteer, input) - frames = sender.snapshot() - for i, want := range []string{"input_conflict", "run_inactive"} { - if frames[i+1].DecodePayload(&ack) != nil || ack.ErrorCode != want { - t.Fatalf("receipt %d: %+v", i, ack) - } - } - } - release.Do(func() { close(sender.release) }) - waitFor(t, func() bool { return router.ActiveRuns() == 0 }, "durable completion") - frames := sender.snapshot() - if len(frames) == 0 || frames[len(frames)-1].Type != proto.TypeDone || calls.Load() != 1 || session.cancels() != 1 { - t.Fatalf("invalid terminal order/calls: %+v, calls=%d cancels=%d", frames, calls.Load(), session.cancels()) - } - if mode == "send_failure" && len(frames) != 1 { - t.Fatal("failed send fabricated an applied receipt") - } - }) - } -} - -func TestShutdownReleasesSteeringWorkerAndCompletionBarrier(t *testing.T) { - for _, phase := range []string{"native", "receipt_send"} { - t.Run(phase, func(t *testing.T) { - sender := &blockedReceiptSender{entered: make(chan struct{}), release: make(chan struct{}), exited: make(chan struct{})} - registry := agent.NewRegistry() - entered, exited := make(chan struct{}), make(chan struct{}) - var session *fakeSession - registry.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - session = &fakeSession{out: out, closeOutOnCancel: true} - return &steeringSession{fakeSession: session, steer: func(ctx context.Context, _ proto.PromptSteerPayload) error { - close(entered) - defer close(exited) - if phase == "native" { - <-ctx.Done() - return ctx.Err() - } - return nil - }}, nil - }) - router, err := dispatch.New(dispatch.Config{Registry: registry, Sender: sender}) - if err != nil { - t.Fatal(err) - } - defer router.Shutdown(context.Background()) - if err = router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "shutdown", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", ReleaseOnCompletion: true}))); err != nil { - t.Fatal(err) - } - if err = router.Handle(context.Background(), mustEnv(t, proto.TypePromptSteer, "shutdown", proto.PromptSteerPayload{InputID: "one", Input: proto.TextInput("text")})); err != nil { - t.Fatal(err) - } - <-entered - if phase == "receipt_send" { - <-sender.entered - session.out <- mustEnv(t, proto.TypeDone, "shutdown", proto.DonePayload{}) - waitFor(t, func() bool { return router.SteeringClosedForTest("shutdown") }, "completion barrier") - } - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - if err = router.Shutdown(ctx); err != nil { - t.Fatal(err) - } - select { - case <-exited: - default: - t.Fatal("native receipt worker leaked") - } - select { - case <-sender.exited: - default: - t.Fatal("receipt send worker leaked") - } - if router.ActiveRuns() != 0 { - t.Fatal("run remained after shutdown") - } - }) - } -} diff --git a/apps/daemon/internal/dispatch/receipt_shutdown_test.go b/apps/daemon/internal/dispatch/receipt_shutdown_test.go deleted file mode 100644 index 882aeb994..000000000 --- a/apps/daemon/internal/dispatch/receipt_shutdown_test.go +++ /dev/null @@ -1,80 +0,0 @@ -package dispatch_test - -import ( - "context" - "errors" - "testing" - "time" - - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/dispatch" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto/prototest" -) - -type shutdownAllSendsBlockSender struct { - recSender - entered chan struct{} - terminal chan context.Context - rescue chan struct{} -} - -func (s *shutdownAllSendsBlockSender) Send(ctx context.Context, env proto.Envelope) error { - if env.Type == proto.TypePromptSteerAck { - close(s.entered) - select { - case <-ctx.Done(): - <-time.After(50 * time.Millisecond) - return ctx.Err() - case <-s.rescue: - return errors.New("test cleanup") - } - } - if env.Type == proto.TypeError { - s.terminal <- ctx - select { - case <-ctx.Done(): - return ctx.Err() - case <-s.rescue: - return errors.New("test cleanup") - } - } - return s.recSender.Send(ctx, env) -} - -func TestShutdownCancelsCompletionErrorSend(t *testing.T) { - sender := &shutdownAllSendsBlockSender{entered: make(chan struct{}), terminal: make(chan context.Context, 1), rescue: make(chan struct{})} - registry := agent.NewRegistry() - var session *fakeSession - registry.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - session = &fakeSession{out: out, closeOutOnCancel: true} - return &steeringSession{fakeSession: session, steer: func(context.Context, proto.PromptSteerPayload) error { return nil }}, nil - }) - router, err := dispatch.New(dispatch.Config{Registry: registry, Sender: sender}) - if err != nil { - t.Fatal(err) - } - if err = router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "shutdown-terminal", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", ReleaseOnCompletion: true}))); err != nil { - t.Fatal(err) - } - if err = router.Handle(context.Background(), mustEnv(t, proto.TypePromptSteer, "shutdown-terminal", proto.PromptSteerPayload{InputID: "one", Input: proto.TextInput("text")})); err != nil { - t.Fatal(err) - } - <-sender.entered - session.out <- mustEnv(t, proto.TypeDone, "shutdown-terminal", proto.DonePayload{}) - waitFor(t, func() bool { return router.SteeringClosedForTest("shutdown-terminal") }, "completion barrier") - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - shutdownErr := router.Shutdown(ctx) - active := router.ActiveRuns() - select { - case terminalCtx := <-sender.terminal: - t.Logf("error send context: Done is nil=%t, Err=%v", terminalCtx.Done() == nil, terminalCtx.Err()) - default: - } - close(sender.rescue) - waitFor(t, func() bool { return router.ActiveRuns() == 0 }, "test cleanup") - if shutdownErr != nil || active != 0 { - t.Fatalf("cooperating receipt sender canceled, but shutdown failed: err=%v active_runs=%d", shutdownErr, active) - } -} diff --git a/apps/daemon/internal/dispatch/router.go b/apps/daemon/internal/dispatch/router.go index 5cbb71ea9..bf67b6902 100644 --- a/apps/daemon/internal/dispatch/router.go +++ b/apps/daemon/internal/dispatch/router.go @@ -1,11 +1,11 @@ // Package dispatch wires inbound WebSocket frames to the agent layer. -// It owns one Session per active RunID, a per-session pump goroutine -// that forwards the agent's events to the transport, and a -// permission_id → run_id index so permission_decision frames route -// back to the right session. +// It owns the Session assignments of its connection, one prepared run per +// active RunID with a goroutine that forwards the run's events to the +// transport, and a permission_id → run_id index so permission_decision +// frames route back to the right run. // // Concurrency: Handle is safe for one goroutine (typically the read -// loop). Each session runs its own goroutine. Internal state is +// loop). Each run has its own goroutines. Internal state is // mutex-protected. package dispatch @@ -37,10 +37,10 @@ type Router struct { admission sync.RWMutex suspension *proto.EnvironmentSuspendPayload mu sync.Mutex - sessions map[string]*sessionState // RunID → state - idle map[string]map[*sessionState]struct{} - permIndex map[string]string // permID → RunID - askIndex map[string]string // askID → RunID + assignments map[string]*assignmentState // SessionID → assignment + sessions map[string]*sessionState // RunID → state + permIndex map[string]string // permID → RunID + askIndex map[string]string // askID → RunID applied map[string]appliedInteractionDecision shutdownAttempt *shutdownAttempt shutdownCh chan struct{} // closed by Shutdown @@ -57,6 +57,7 @@ type Router struct { workspaceReads map[string]struct{} localWorkspace *localworkspace.Binding sessionEnvironments bool + removeHome func(sessionID string) error } type appliedInteractionDecision struct { @@ -68,30 +69,25 @@ type appliedInteractionDecision struct { // sessionState is the dispatcher's per-run bookkeeping. The agent // owns the close of out; the dispatcher cancels ctxCancel to wind -// down. traceparent captures the prompt_request's W3C trace so every +// down. traceparent captures the execution_start's W3C trace so every // outbound frame stamps env.Trace with the same value, completing // frontend → server → daemon → agent → server attribution. type sessionState struct { - capabilities proto.AgentKindCapabilities - runID string - environmentID string - stateKey string - session agent.Session - out chan proto.Envelope - ctx context.Context - ctxCancel context.CancelFunc - pendingIDs map[string]struct{} - pendingAsks map[string]struct{} - traceparent string - idleTimer *time.Timer - idleLease uint64 - retain bool - releaseOnCompletion bool - steering map[string]steeringReceipt - steerBusy bool - steeringClosed bool - steeringDone chan struct{} - preparedHandoff *preparedHandoff + assignment proto.AssignmentRef + capabilities proto.AgentKindCapabilities + runID string + environmentID string + session agent.Session + out chan proto.Envelope + ctx context.Context + ctxCancel context.CancelFunc + pendingIDs map[string]struct{} + pendingAsks map[string]struct{} + traceparent string + steering map[string]steeringReceipt + steerBusy bool + steeringClosed bool + preparedHandoff *preparedHandoff } // Config is the constructor input. Registry and Sender are required; @@ -107,6 +103,9 @@ type Config struct { // is its Session's Environment, which the Executor factory binds, and not // a local workspace of this daemon. It excludes LocalWorkspace. SessionEnvironments bool + // RemoveHome removes the Session's native home once its Executors have + // closed. Nil declares that assignment_release does not accept RemoveHome. + RemoveHome func(sessionID string) error } const defaultIdleTimeout = time.Hour @@ -139,8 +138,8 @@ func New(cfg Config) (*Router, error) { registry: cfg.Registry, sender: cfg.Sender, log: log, + assignments: make(map[string]*assignmentState), sessions: make(map[string]*sessionState), - idle: make(map[string]map[*sessionState]struct{}), permIndex: make(map[string]string), askIndex: make(map[string]string), applied: make(map[string]appliedInteractionDecision), @@ -152,12 +151,14 @@ func New(cfg Config) (*Router, error) { preparationTimeout: preparationTimeout, localWorkspace: cfg.LocalWorkspace, sessionEnvironments: cfg.SessionEnvironments, + removeHome: cfg.RemoveHome, }, nil } // Handle dispatches one inbound Envelope. Errors are returned for // programmer-visible problems (bad shape, registry miss); transient -// session-level failures are logged and swallowed. +// session-level failures are logged and swallowed. A quiesced Router +// still handles assignment_release. // // Adopts env.Trace into ctx so every downstream log under it inherits // the same trace_id, making a single grep cover both sides. @@ -169,7 +170,7 @@ func (r *Router) Handle(ctx context.Context, env proto.Envelope) error { r.mu.Unlock() return ErrRouterClosed } - if r.suspension != nil { + if r.suspension != nil && env.Type != proto.TypeAssignmentRelease { r.mu.Unlock() return ErrRouterQuiesced } @@ -178,6 +179,10 @@ func (r *Router) Handle(ctx context.Context, env proto.Envelope) error { ctx = adoptEnvelopeTrace(ctx, env) switch env.Type { + case proto.TypeAssignmentBind: + return r.handleAssignmentBind(ctx, env) + case proto.TypeAssignmentRelease: + return r.handleAssignmentRelease(ctx, env) case proto.TypeRuntimePrepare: return r.handleRuntimePrepare(ctx, env) case proto.TypeWorkspaceExport: @@ -192,8 +197,6 @@ func (r *Router) Handle(ctx context.Context, env proto.Envelope) error { return r.handleExecutionStart(ctx, env) case proto.TypeExecutionRelease: return r.handleExecutionRelease(ctx, env) - case proto.TypePromptRequest: - return r.handlePromptRequest(ctx, env) case proto.TypePromptCancel: return r.handlePromptCancel(ctx, env) case proto.TypeFunctionResult: @@ -205,10 +208,8 @@ func (r *Router) Handle(ctx context.Context, env proto.Envelope) error { case proto.TypePromptForUserChoiceDecision: return r.handlePromptForUserChoiceDecision(ctx, env) default: - // Unknown types are logged and dropped — keeps the daemon - // forward-compatible with server-side additions. - r.log.WarnContext(ctx, "dropping unknown envelope type", "type", env.Type, "id", env.ID) - return nil + r.log.WarnContext(ctx, "unsupported envelope type", "type", env.Type, "id", env.ID) + return r.reply(ctx, env, proto.TypeProtocolError, proto.ProtocolErrorPayload{Type: env.Type, ErrorCode: proto.UnsupportedOperation}) } } @@ -244,8 +245,8 @@ func (r *Router) drain(ch <-chan proto.Envelope) { } } -// cleanupSession removes the session from registry maps. Called from -// pump's defer so it runs exactly once. +// cleanupSession removes the run from the registry maps. The run's prepared +// release calls it once. func (r *Router) cleanupSession(s *sessionState) { r.mu.Lock() delete(r.sessions, s.runID) @@ -255,84 +256,5 @@ func (r *Router) cleanupSession(s *sessionState) { for askID := range s.pendingAsks { delete(r.askIndex, askID) } - if !s.retain || s.session == nil || s.stateKey == "" || r.closed { - r.mu.Unlock() - return - } - states := r.idle[s.stateKey] - if states == nil { - states = make(map[*sessionState]struct{}) - r.idle[s.stateKey] = states - } - states[s] = struct{}{} - r.scheduleIdleLocked(s) - r.mu.Unlock() -} - -func sessionStateKey(req proto.PromptRequestPayload) string { - return req.AgentStateKey -} - -func (r *Router) touchIdleLocked(stateKey string) { - if stateKey == "" { - return - } - for state := range r.idle[stateKey] { - r.scheduleIdleLocked(state) - } -} - -func (r *Router) scheduleIdleLocked(state *sessionState) { - if state.idleTimer != nil { - state.idleTimer.Stop() - } - state.idleLease++ - lease := state.idleLease - state.idleTimer = time.AfterFunc(r.idleTimeout, func() { - r.expireIdle(state, lease) - }) -} - -func (r *Router) expireIdle(state *sessionState, lease uint64) { - r.mu.Lock() - states := r.idle[state.stateKey] - if _, ok := states[state]; !ok || state.idleLease != lease || r.suspension != nil { - r.mu.Unlock() - return - } - r.shutdownWG.Add(1) - defer r.shutdownWG.Done() - delete(states, state) - if len(states) == 0 { - delete(r.idle, state.stateKey) - } - state.retain = false r.mu.Unlock() - - state.ctxCancel() - if err := state.session.Cancel(context.Background()); err != nil { - r.log.Warn("idle session cancel failed", "run_id", state.runID, "state_key", state.stateKey, "err", err) - } -} - -// emitTerminalError synthesises error + done for a run that couldn't -// even be started. Stamps env.Trace from ctx so the gateway can -// attribute these frames to the same trace_id. -func (r *Router) emitTerminalError(ctx context.Context, runID, msg string) { - traceparent := "" - if carrier, ok := obslog.TraceFromContext(ctx); ok { - traceparent = carrier.String() - } - errEnv, err := proto.NewEnvelopeWithTrace(proto.TypeError, runID, proto.ErrorPayload{Error: msg}, traceparent) - if err == nil { - if sendErr := r.sender.Send(ctx, errEnv); sendErr != nil { - r.log.ErrorContext(ctx, "emit terminal error frame failed", "run_id", runID, "err", sendErr) - } - } - doneEnv, err := proto.NewEnvelopeWithTrace(proto.TypeDone, runID, proto.DonePayload{}, traceparent) - if err == nil { - if sendErr := r.sender.Send(ctx, doneEnv); sendErr != nil { - r.log.ErrorContext(ctx, "emit terminal done frame failed", "run_id", runID, "err", sendErr) - } - } } diff --git a/apps/daemon/internal/dispatch/router_test.go b/apps/daemon/internal/dispatch/router_test.go index 9b05035a8..99a9a0d6a 100644 --- a/apps/daemon/internal/dispatch/router_test.go +++ b/apps/daemon/internal/dispatch/router_test.go @@ -3,7 +3,6 @@ package dispatch_test import ( "context" "errors" - "slices" "sync" "testing" "time" @@ -18,8 +17,7 @@ import ( // test doubles // --------------------------------------------------------------------- -// recSender records every Envelope. failNow makes the next Send fail -// (used to exercise the pump's error path). +// recSender records every Envelope. failNow makes the next Send fail. type recSender struct { mu sync.Mutex frames []proto.Envelope @@ -131,9 +129,7 @@ func (s *fakeSession) submissions() []permCall { // helpers // --------------------------------------------------------------------- -// newHarness builds a Router whose registry exposes a single -// fake_alpha factory that records inputs and exposes the in-flight -// session. +// harness is a Router whose fake_alpha kind runs each Turn as a fakeSession. type harness struct { router *dispatch.Router sender *recSender @@ -143,10 +139,6 @@ type harness struct { } func newHarness(t *testing.T) *harness { - return newHarnessWithIdleTimeout(t, time.Hour) -} - -func newHarnessWithIdleTimeout(t *testing.T, idleTimeout time.Duration) *harness { t.Helper() h := &harness{ sender: &recSender{}, @@ -154,13 +146,13 @@ func newHarnessWithIdleTimeout(t *testing.T, idleTimeout time.Duration) *harness gotReq: make(chan proto.PromptRequestPayload, 16), gotSess: make(chan *fakeSession, 16), } - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "fake_alpha", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Permissions: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - sess := &fakeSession{out: out, ctx: ctx} + registerSession(h.reg, proto.SupportedAgentKind{Kind: "fake_alpha", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Permissions: proto.CapabilitySupported})}, func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + sess := &fakeSession{out: out, ctx: ctx, closeOutOnCancel: true} h.gotReq <- req h.gotSess <- sess return sess, nil }) - r, err := dispatch.New(dispatch.Config{Registry: h.reg, Sender: h.sender, IdleTimeout: idleTimeout}) + r, err := dispatch.New(dispatch.Config{Registry: h.reg, Sender: h.sender}) if err != nil { t.Fatalf("dispatch.New: %v", err) } @@ -168,171 +160,70 @@ func newHarnessWithIdleTimeout(t *testing.T, idleTimeout time.Duration) *harness return h } -func TestCompletedSessionCancelsAfterIdleTimeout(t *testing.T) { - h := newHarnessWithIdleTimeout(t, 40*time.Millisecond) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "run_idle", prototest.WithModel(proto.PromptRequestPayload{ - AgentKind: "fake_alpha", ConversationID: "conv-idle", AgentStateKey: "conv-idle/agent/fake_alpha", +// registerSession declares info, without an Environment, and starts each +// Turn of the kind with factory. +func registerSession(reg *agent.Registry, info proto.SupportedAgentKind, factory agent.Factory) { + info.Capabilities.EnvironmentNone = proto.CapabilitySupported + reg.RegisterKind(info, prototest.ModelConfiguration(), factory) + reg.RegisterExecutor(info.Kind, preparationExecutorFixture(func(_ context.Context, req proto.PromptRequestPayload) (preparedFixture, error) { + return &controlledPreparation{start: func(ctx context.Context, id string, input proto.MessageInput, out chan<- proto.Envelope) (agent.Session, error) { + req.RunID, req.Input = id, input + return factory(ctx, req, out) + }}, nil })) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("Handle prompt_request: %v", err) - } - sess := <-h.gotSess - close(sess.out) - waitFor(t, func() bool { return h.router.ActiveRuns() == 0 }, "active run cleanup") - - select { - case <-sess.ctx.Done(): - t.Fatal("completed session cancelled before idle timeout") - case <-time.After(15 * time.Millisecond): - } - waitFor(t, func() bool { return sess.cancels() == 1 }, "idle session cancellation") } -func TestNewPromptResetsCompletedSessionIdleTimeout(t *testing.T) { - h := newHarnessWithIdleTimeout(t, 80*time.Millisecond) - defer h.router.Shutdown(context.Background()) - - stateKey := "conv-renew/agent/fake_alpha" - first := mustEnv(t, proto.TypePromptRequest, "run_first", prototest.WithModel(proto.PromptRequestPayload{ - AgentKind: "fake_alpha", ConversationID: "conv-renew", AgentStateKey: stateKey, - })) - if err := h.router.Handle(context.Background(), first); err != nil { - t.Fatalf("Handle first prompt: %v", err) +// startRun binds the Session run to r, prepares its Executor of kind and +// starts run. sender records r's frames. +func startRun(t *testing.T, r *dispatch.Router, sender *recSender, kind, run string) { + t.Helper() + assign(t, r, run, "") + prepare := scoped(t, run, proto.TypeExecutionPrepare, "prepare-"+run, proto.ExecutionPreparePayload{SessionID: run, Configuration: prototest.WithModel(proto.PromptRequestPayload{AgentKind: kind, AgentStateKey: stateKey(run), StrictResume: true, DisableExecutionEnvironment: true})}) + if err := r.Handle(t.Context(), prepare); err != nil { + t.Fatalf("execution_prepare: %v", err) } - firstSession := <-h.gotSess - close(firstSession.out) - waitFor(t, func() bool { return h.router.ActiveRuns() == 0 }, "first run cleanup") - time.Sleep(50 * time.Millisecond) - - second := mustEnv(t, proto.TypePromptRequest, "run_second", prototest.WithModel(proto.PromptRequestPayload{ - AgentKind: "fake_alpha", ConversationID: "conv-renew", AgentStateKey: stateKey, - })) - if err := h.router.Handle(context.Background(), second); err != nil { - t.Fatalf("Handle second prompt: %v", err) + ready := waitPreparationStatus(t, sender, prepare.ID, "ready", "") + if err := r.Handle(t.Context(), scoped(t, run, proto.TypeExecutionStart, prepare.ID, proto.ExecutionStartPayload{Handle: ready.Handle, ExecutorID: ready.ExecutorID, RunID: run, Input: proto.TextInput("input")})); err != nil { + t.Fatalf("execution_start: %v", err) } - secondSession := <-h.gotSess + waitPreparationStatus(t, sender, prepare.ID, "started", "") +} + +// ref is the assignment the tests bind session to. +func ref(session string) proto.AssignmentRef { + return proto.AssignmentRef{SessionID: session, AssignmentID: "assignment-" + session, Epoch: 1} +} - time.Sleep(45 * time.Millisecond) - if firstSession.cancels() != 0 { - t.Fatal("new prompt did not renew the completed session idle timeout") +// assign binds session to r in environment. +func assign(t *testing.T, r *dispatch.Router, session, environment string) { + t.Helper() + if err := r.Handle(t.Context(), scoped(t, session, proto.TypeAssignmentBind, "bind-"+session, proto.AssignmentBindPayload{EnvironmentID: environment})); err != nil { + t.Fatalf("assignment_bind: %v", err) } - close(secondSession.out) - waitFor(t, func() bool { return firstSession.cancels() == 1 }, "renewed idle session cancellation") } -func mustEnv(t *testing.T, typ, id string, payload any) proto.Envelope { +// scoped is a frame of session's work under its assignment. +func scoped(t *testing.T, session, typ, id string, payload any) proto.Envelope { t.Helper() env, err := proto.NewEnvelope(typ, id, payload) if err != nil { t.Fatalf("NewEnvelope %s: %v", typ, err) } + env.Assignment = ref(session) return env } +// mustEnv is a frame of preparationSessionID's work, the Session most tests +// run. +func mustEnv(t *testing.T, typ, id string, payload any) proto.Envelope { + t.Helper() + return scoped(t, preparationSessionID, typ, id, payload) +} + // --------------------------------------------------------------------- // tests // --------------------------------------------------------------------- -func TestHandlePromptRequestInvokesFactoryAndForwardsOutput(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "run_1", prototest.WithModel(proto.PromptRequestPayload{ - AgentKind: "fake_alpha", Input: proto.TextInput("hi"), ConversationID: "c1", - })) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("Handle prompt_request: %v", err) - } - - req := <-h.gotReq - if req.RunID != "run_1" || req.AgentKind != "fake_alpha" || *req.Input[0].Content[0].Text != "hi" { - t.Errorf("factory got %+v, want run_1/fake_alpha/hi", req) - } - sess := <-h.gotSess - - // Session emits a delta + done; both should reach the sender. - sess.out <- mustEnv(t, proto.TypeDelta, "run_1", proto.DeltaPayload{Delta: "hello", Sequence: 1}) - sess.out <- mustEnv(t, proto.TypeDone, "run_1", proto.DonePayload{Content: "hello"}) - close(sess.out) - - waitForTypes(t, h.sender, "run_1", []string{proto.TypeDelta, proto.TypeDone}) - - // Pump should have removed the session. - waitFor(t, func() bool { return h.router.ActiveRuns() == 0 }, "active runs to drop to 0") -} - -func TestHandlePromptRequestRejectsDuplicateRunID(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "run_dup", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("first Handle: %v", err) - } - <-h.gotReq // drain first - sess := <-h.gotSess - - // Second prompt_request with same RunID should NOT spin up a second factory call. - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("duplicate Handle: %v", err) - } - - select { - case extra := <-h.gotReq: - t.Fatalf("factory invoked twice for duplicate run, second req=%+v", extra) - case <-time.After(50 * time.Millisecond): - } - close(sess.out) -} - -func TestHandlePromptRequestUnsupportedKindEmitsErrorDone(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "run_x", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_beta"})) - err := h.router.Handle(context.Background(), env) - if !errors.Is(err, agent.ErrUnsupportedKind) { - t.Errorf("Handle unsupported = %v, want ErrUnsupportedKind", err) - } - got := h.sender.typesFor("run_x") - want := []string{proto.TypeError, proto.TypeDone} - if !slices.Equal(got, want) { - t.Errorf("sender frames for run_x = %v, want %v", got, want) - } -} - -func TestHandlePromptRequestMissingRunIDIsError(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - if err := h.router.Handle(context.Background(), env); err == nil { - t.Fatal("expected error on missing run id") - } -} - -func TestHandlePromptCancelInvokesSessionCancel(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "run_2", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("prompt_request: %v", err) - } - <-h.gotReq - sess := <-h.gotSess - sess.closeOutOnCancel = true - - if err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptCancel, "run_2", nil)); err != nil { - t.Fatalf("prompt_cancel: %v", err) - } - - waitFor(t, func() bool { return sess.cancels() == 1 }, "session.Cancel to fire once") - waitFor(t, func() bool { return h.router.ActiveRuns() == 0 }, "session to be cleaned up") -} - func TestHandlePromptCancelUnknownRunIsNoop(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) @@ -346,22 +237,17 @@ func TestPermissionRequestIsIndexedAndDecisionRoutes(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) - env := mustEnv(t, proto.TypePromptRequest, "run_p", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("prompt_request: %v", err) - } + startRun(t, h.router, h.sender, "fake_alpha", "run_p") <-h.gotReq sess := <-h.gotSess - // Session emits a permission_request; pump should index it. - permEnv := mustEnv(t, proto.TypePermissionRequest, "run_p", proto.PermissionRequestPayload{ + // Session emits a permission_request; the Router indexes it before it forwards it. + sess.out <- mustEnv(t, proto.TypePermissionRequest, "run_p", proto.PermissionRequestPayload{ RequestID: "perm_abcd1234", Tool: "Bash", Title: "rm -rf /", }) - sess.out <- permEnv - // Wait until sender records — indexing happens before send. - waitFor(t, func() bool { return len(h.sender.snapshot()) >= 1 }, "permission_request to be forwarded") + waitFor(t, func() bool { return hasFrame(h.sender, proto.TypePermissionRequest, "run_p") }, "permission_request to be forwarded") - dec := mustEnv(t, proto.TypePermissionDecision, "perm_abcd1234", proto.PermissionDecisionPayload{DeliveryID: "delivery-perm-1", Approved: true}) + dec := scoped(t, "run_p", proto.TypePermissionDecision, "perm_abcd1234", proto.PermissionDecisionPayload{DeliveryID: "delivery-perm-1", Approved: true}) if err := h.router.Handle(context.Background(), dec); err != nil { t.Fatalf("permission_decision: %v", err) } @@ -370,7 +256,7 @@ func TestPermissionRequestIsIndexedAndDecisionRoutes(t *testing.T) { t.Errorf("submissions = %+v, want one approved perm_abcd1234", calls) } assertDecisionAck(t, h.sender, "delivery-perm-1", true, "") - retry := mustEnv(t, proto.TypePermissionDecision, "perm_abcd1234", proto.PermissionDecisionPayload{DeliveryID: "delivery-perm-2", Approved: true}) + retry := scoped(t, "run_p", proto.TypePermissionDecision, "perm_abcd1234", proto.PermissionDecisionPayload{DeliveryID: "delivery-perm-2", Approved: true}) if err := h.router.Handle(context.Background(), retry); err != nil { t.Fatalf("idempotent permission replay: %v", err) } @@ -378,39 +264,11 @@ func TestPermissionRequestIsIndexedAndDecisionRoutes(t *testing.T) { t.Fatalf("idempotent replay reached agent twice: %+v", calls) } assertDecisionAck(t, h.sender, "delivery-perm-2", true, "") - conflict := mustEnv(t, proto.TypePermissionDecision, "perm_abcd1234", proto.PermissionDecisionPayload{DeliveryID: "delivery-perm-3", Approved: false}) + conflict := scoped(t, "run_p", proto.TypePermissionDecision, "perm_abcd1234", proto.PermissionDecisionPayload{DeliveryID: "delivery-perm-3", Approved: false}) if err := h.router.Handle(context.Background(), conflict); err != nil { t.Fatalf("conflicting permission replay: %v", err) } assertDecisionAck(t, h.sender, "delivery-perm-3", false, "decision_conflict") - - close(sess.out) -} - -func TestPermissionCancelDeindexes(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "run_p2", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - _ = h.router.Handle(context.Background(), env) - <-h.gotReq - sess := <-h.gotSess - - sess.out <- mustEnv(t, proto.TypePermissionRequest, "run_p2", proto.PermissionRequestPayload{RequestID: "perm_xx", Tool: "Bash"}) - waitFor(t, func() bool { return len(h.sender.snapshot()) >= 1 }, "perm forwarded") - sess.out <- mustEnv(t, proto.TypePermissionCancel, "perm_xx", nil) - waitFor(t, func() bool { return len(h.sender.snapshot()) >= 2 }, "perm_cancel forwarded") - - // A decision for the cancelled perm should be a no-op (session never sees it). - if err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePermissionDecision, "perm_xx", proto.PermissionDecisionPayload{DeliveryID: "delivery-cancelled"})); err != nil { - t.Fatalf("decision: %v", err) - } - if calls := sess.submissions(); len(calls) != 0 { - t.Errorf("expected zero submissions after cancel, got %+v", calls) - } - assertDecisionAck(t, h.sender, "delivery-cancelled", false, "not_pending") - - close(sess.out) } func TestPermissionDecisionUnknownPermIsNoop(t *testing.T) { @@ -423,84 +281,27 @@ func TestPermissionDecisionUnknownPermIsNoop(t *testing.T) { assertDecisionAck(t, h.sender, "delivery-unknown-perm", false, "not_pending") } -func TestPromptForUserChoiceDecisionRoutesToSession(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - env := mustEnv(t, proto.TypePromptRequest, "run_ask", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("prompt_request: %v", err) - } - <-h.gotReq - sess := <-h.gotSess - - // Envelope.ID is the run id (server-side dispatch fans on it); the - // ask id rides on the payload. Daemon's indexPermissionFrame reads - // payload.AskID to seed askIndex. - askEnv := mustEnv(t, proto.TypePromptForUserChoice, "run_ask", proto.PromptForUserChoicePayload{ - AskID: "ask_abcd1234", - Questions: []proto.PromptForUserChoiceQuestion{{Question: "?", Options: []proto.PromptForUserChoiceOption{{Label: "yes"}, {Label: "no"}}}}, - ToolUseID: "toolu_42", - }) - sess.out <- askEnv - waitFor(t, func() bool { return len(h.sender.snapshot()) >= 1 }, "prompt_for_user_choice forwarded") - - dec := mustEnv(t, proto.TypePromptForUserChoiceDecision, "ask_abcd1234", proto.PromptForUserChoiceDecisionPayload{ - DeliveryID: "delivery-ask-1", QuestionAnswers: []proto.PromptForUserChoiceQuestionAnswer{{QuestionID: "q0", Answers: []string{"yes"}}}, - }) - if err := h.router.Handle(context.Background(), dec); err != nil { - t.Fatalf("prompt_for_user_choice_decision: %v", err) - } - - sess.askMu.Lock() - calls := append([]askCall(nil), sess.askCalls...) - sess.askMu.Unlock() - if len(calls) != 1 || calls[0].id != "ask_abcd1234" { - t.Fatalf("askCalls = %+v, want one ask_abcd1234", calls) - } - if len(calls[0].decision.QuestionAnswers[0].Answers) != 1 || calls[0].decision.QuestionAnswers[0].Answers[0] != "yes" { - t.Errorf("answer payload mismatch: %+v", calls[0].decision) - } - assertDecisionAck(t, h.sender, "delivery-ask-1", true, "") - - // Cleanup contract: a successful decision drops the ask from both - // the router-level index and the session's pendingAsks set, so a - // stale retry short-circuits as "run gone". - if got := h.router.AskIndexLenForTest(); got != 0 { - t.Errorf("askIndex len = %d, want 0 after decision", got) - } - if got := h.router.PendingAsksLenForTest("run_ask"); got != 0 { - t.Errorf("pendingAsks len = %d, want 0 after decision", got) - } - - close(sess.out) -} - // TestPromptForUserChoiceDecisionClearsIndexOnAgentUnknown locks in the -// other cleanup branch: when the session returns ErrUnknownAsk (timer -// already consumed the entry), the router still drops the index so a -// retry doesn't loop into Submit again. +// cleanup branch: when the session returns ErrUnknownAsk (timer already +// consumed the entry), the router still drops the index so a retry doesn't +// loop into Submit again. func TestPromptForUserChoiceDecisionClearsIndexOnAgentUnknown(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) - env := mustEnv(t, proto.TypePromptRequest, "run_ask_u", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("prompt_request: %v", err) - } + startRun(t, h.router, h.sender, "fake_alpha", "run_ask_u") <-h.gotReq sess := <-h.gotSess sess.askErr = agent.ErrUnknownAsk - askEnv := mustEnv(t, proto.TypePromptForUserChoice, "run_ask_u", proto.PromptForUserChoicePayload{ + sess.out <- mustEnv(t, proto.TypePromptForUserChoice, "run_ask_u", proto.PromptForUserChoicePayload{ AskID: "ask_xxxxxxxx", Questions: []proto.PromptForUserChoiceQuestion{{Question: "?", Options: []proto.PromptForUserChoiceOption{{Label: "yes"}}}}, ToolUseID: "toolu_y", }) - sess.out <- askEnv - waitFor(t, func() bool { return len(h.sender.snapshot()) >= 1 }, "prompt_for_user_choice forwarded") + waitFor(t, func() bool { return hasFrame(h.sender, proto.TypePromptForUserChoice, "run_ask_u") }, "prompt_for_user_choice forwarded") - dec := mustEnv(t, proto.TypePromptForUserChoiceDecision, "ask_xxxxxxxx", proto.PromptForUserChoiceDecisionPayload{ + dec := scoped(t, "run_ask_u", proto.TypePromptForUserChoiceDecision, "ask_xxxxxxxx", proto.PromptForUserChoiceDecisionPayload{ DeliveryID: "delivery-ask-gone", QuestionAnswers: []proto.PromptForUserChoiceQuestionAnswer{{QuestionID: "q0", Answers: []string{"yes"}}}, }) if err := h.router.Handle(context.Background(), dec); err != nil { @@ -514,18 +315,13 @@ func TestPromptForUserChoiceDecisionClearsIndexOnAgentUnknown(t *testing.T) { t.Errorf("pendingAsks len = %d, want 0 after ErrUnknownAsk", got) } assertDecisionAck(t, h.sender, "delivery-ask-gone", false, "not_pending") - - close(sess.out) } func TestPromptForUserChoiceDecisionKeepsIndexOnTransientAgentError(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) - env := mustEnv(t, proto.TypePromptRequest, "run_ask_retry", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"})) - if err := h.router.Handle(context.Background(), env); err != nil { - t.Fatalf("prompt_request: %v", err) - } + startRun(t, h.router, h.sender, "fake_alpha", "run_ask_retry") <-h.gotReq sess := <-h.gotSess sess.askErr = errors.New("temporary stdin failure") @@ -533,9 +329,9 @@ func TestPromptForUserChoiceDecisionKeepsIndexOnTransientAgentError(t *testing.T sess.out <- mustEnv(t, proto.TypePromptForUserChoice, "run_ask_retry", proto.PromptForUserChoicePayload{ AskID: "ask_retry", Questions: []proto.PromptForUserChoiceQuestion{{ID: "q0", Question: "Retry?"}}, }) - waitFor(t, func() bool { return len(h.sender.snapshot()) >= 1 }, "prompt_for_user_choice forwarded") + waitFor(t, func() bool { return hasFrame(h.sender, proto.TypePromptForUserChoice, "run_ask_retry") }, "prompt_for_user_choice forwarded") - decision := mustEnv(t, proto.TypePromptForUserChoiceDecision, "ask_retry", proto.PromptForUserChoiceDecisionPayload{ + decision := scoped(t, "run_ask_retry", proto.TypePromptForUserChoiceDecision, "ask_retry", proto.PromptForUserChoiceDecisionPayload{ DeliveryID: "delivery-ask-retry", QuestionAnswers: []proto.PromptForUserChoiceQuestionAnswer{{QuestionID: "q0", Answers: []string{"yes"}}}, }) if err := h.router.Handle(context.Background(), decision); err != nil { @@ -549,7 +345,9 @@ func TestPromptForUserChoiceDecisionKeepsIndexOnTransientAgentError(t *testing.T t.Fatalf("pendingAsks len = %d, want 1 after transient error", got) } + sess.askMu.Lock() sess.askErr = nil + sess.askMu.Unlock() if err := h.router.Handle(context.Background(), decision); err != nil { t.Fatalf("retry decision: %v", err) } @@ -557,7 +355,6 @@ func TestPromptForUserChoiceDecisionKeepsIndexOnTransientAgentError(t *testing.T if got := h.router.AskIndexLenForTest(); got != 0 { t.Fatalf("askIndex len = %d, want 0 after successful retry", got) } - close(sess.out) } func TestPromptForUserChoiceDecisionUnknownAskIsNoop(t *testing.T) { @@ -591,70 +388,21 @@ func assertDecisionAck(t *testing.T, sender *recSender, deliveryID string, appli t.Fatalf("no decision ack for delivery %q in %+v", deliveryID, frames) } -func TestHandleUnknownTypeIsNoop(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - - if err := h.router.Handle(context.Background(), proto.Envelope{Type: "fancy_new_event"}); err != nil { - t.Errorf("unknown type Handle = %v, want nil", err) - } -} - func TestHandleAfterShutdownReturnsErrRouterClosed(t *testing.T) { h := newHarness(t) if err := h.router.Shutdown(context.Background()); err != nil { t.Fatalf("Shutdown: %v", err) } - err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "r", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"}))) + err := h.router.Handle(context.Background(), mustEnv(t, proto.TypeExecutionPrepare, "r", proto.ExecutionPreparePayload{})) if !errors.Is(err, dispatch.ErrRouterClosed) { t.Errorf("post-shutdown Handle = %v, want ErrRouterClosed", err) } } -func TestShutdownWaitsForPumpDrain(t *testing.T) { - h := newHarness(t) - - if err := h.router.Handle(context.Background(), mustEnv(t, proto.TypePromptRequest, "rs", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"}))); err != nil { - t.Fatalf("prompt_request: %v", err) - } - <-h.gotReq - sess := <-h.gotSess - - // Background: emit one frame then close out shortly after - // shutdown is asked for. - go func() { - sess.out <- mustEnv(t, proto.TypeDelta, "rs", proto.DeltaPayload{Delta: "x"}) - time.Sleep(20 * time.Millisecond) - close(sess.out) - }() - - ctx, cancel := context.WithTimeout(context.Background(), time.Second) - defer cancel() - if err := h.router.Shutdown(ctx); err != nil { - t.Fatalf("Shutdown returned %v before pump drained", err) - } - if h.router.ActiveRuns() != 0 { - t.Errorf("ActiveRuns after Shutdown = %d, want 0", h.router.ActiveRuns()) - } -} - // --------------------------------------------------------------------- // helpers // --------------------------------------------------------------------- -func waitForTypes(t *testing.T, s *recSender, runID string, want []string) { - t.Helper() - deadline := time.Now().Add(2 * time.Second) - for time.Now().Before(deadline) { - got := s.typesFor(runID) - if slices.Equal(got, want) { - return - } - time.Sleep(5 * time.Millisecond) - } - t.Fatalf("never observed types %v for run %s; got %v", want, runID, s.typesFor(runID)) -} - func waitFor(t *testing.T, cond func() bool, what string) { t.Helper() deadline := time.Now().Add(2 * time.Second) diff --git a/apps/daemon/internal/dispatch/runtime_preparation.go b/apps/daemon/internal/dispatch/runtime_preparation.go index ba8866e43..42be3ba95 100644 --- a/apps/daemon/internal/dispatch/runtime_preparation.go +++ b/apps/daemon/internal/dispatch/runtime_preparation.go @@ -45,7 +45,7 @@ func (r *Router) handleRuntimePrepare(ctx context.Context, env proto.Envelope) e // A late malformed frame cannot report rejection of an earlier commit. return errors.New("dispatch: malformed pending Runtime preparation frame") } - return r.sendRuntimePrepareResult(ctx, env.ID, rejectedRuntimePreparation("invalid_request")) + return r.sendRuntimePrepareResult(ctx, env, rejectedRuntimePreparation("invalid_request")) } r.mu.Lock() if r.closed || r.suspension != nil { @@ -59,11 +59,15 @@ func (r *Router) handleRuntimePrepare(ctx context.Context, env proto.Envelope) e if duplicate { return errors.New("dispatch: Runtime preparation already admitted") } - return r.sendRuntimePrepareResult(ctx, env.ID, rejectedRuntimePreparation("runtime_preparation_capacity")) + return r.sendRuntimePrepareResult(ctx, env, rejectedRuntimePreparation("runtime_preparation_capacity")) + } + if code := r.admitLocked(env.Assignment, request.SessionID, request.EnvironmentID); code != "" { + r.mu.Unlock() + return r.sendRuntimePrepareResult(ctx, env, rejectedRuntimePreparation(code)) } if r.localWorkspace == nil || !r.localWorkspace.Matches(request.EnvironmentID, request.SessionID) || r.runtimePreparationResourcesBusyLocked() { r.mu.Unlock() - return r.sendRuntimePrepareResult(ctx, env.ID, rejectedRuntimePreparation("resource_unavailable")) + return r.sendRuntimePrepareResult(ctx, env, rejectedRuntimePreparation("resource_unavailable")) } owner, cancel := context.WithTimeout(context.WithoutCancel(ctx), runtimePreparationTimeout) u := &runtimePreparationTransfer{ @@ -74,16 +78,16 @@ func (r *Router) handleRuntimePrepare(ctx context.Context, env proto.Envelope) e r.shutdownWG.Add(1) r.mu.Unlock() go r.runRuntimePreparationTransfer(owner, u, r.localWorkspace.ApplyRuntimePreparation) - if err := r.sendRuntimePrepareResult(ctx, env.ID, proto.RuntimePrepareResultPayload{Outcome: "ready"}); err != nil { + if err := r.sendRuntimePrepareResult(ctx, env, proto.RuntimePrepareResultPayload{Outcome: "ready"}); err != nil { cancel() return err } return nil } u := r.runtimePreparation - if u == nil || u.envelope.ID != env.ID { + if u == nil || u.envelope.ID != env.ID || u.envelope.Assignment != env.Assignment { r.mu.Unlock() - return r.sendRuntimePrepareResult(ctx, env.ID, rejectedRuntimePreparation("resource_unavailable")) + return r.sendRuntimePrepareResult(ctx, env, rejectedRuntimePreparation("resource_unavailable")) } if u.finished { r.mu.Unlock() @@ -93,7 +97,7 @@ func (r *Router) handleRuntimePrepare(ctx context.Context, env proto.Envelope) e u.data = append(u.data, request.Data...) offset := len(u.data) r.mu.Unlock() - if err := r.sendRuntimePrepareResult(ctx, env.ID, proto.RuntimePrepareResultPayload{Outcome: "received", Offset: offset}); err != nil { + if err := r.sendRuntimePrepareResult(ctx, env, proto.RuntimePrepareResultPayload{Outcome: "received", Offset: offset}); err != nil { u.cancel() return err } @@ -114,7 +118,7 @@ func (r *Router) handleRuntimePrepare(ctx context.Context, env proto.Envelope) e } func (r *Router) runtimePreparationResourcesBusyLocked() bool { - if r.workspaceWrite != nil || r.workspaceExport != nil || len(r.workspaceReads) != 0 || len(r.sessions) != 0 || len(r.idle) != 0 || len(r.executors) != 0 { + if r.workspaceWrite != nil || r.workspaceExport != nil || len(r.workspaceReads) != 0 || len(r.sessions) != 0 || len(r.executors) != 0 { return true } for _, p := range r.preparations { @@ -159,7 +163,7 @@ func (r *Router) runRuntimePreparationTransfer(ctx context.Context, u *runtimePr } r.mu.Unlock() // The result has a separate send budget, independent of an installation timeout. - _ = r.sendRuntimePrepareResult(context.WithoutCancel(ctx), u.envelope.ID, result) + _ = r.sendRuntimePrepareResult(context.WithoutCancel(ctx), u.envelope, result) } func rejectedRuntimePreparation(code string) proto.RuntimePrepareResultPayload { @@ -186,12 +190,6 @@ func runtimePreparationResult(err error, size int) proto.RuntimePrepareResultPay return proto.RuntimePrepareResultPayload{Outcome: "unknown", ErrorCode: "runtime_preparation_unconfirmed"} } -func (r *Router) sendRuntimePrepareResult(ctx context.Context, id string, result proto.RuntimePrepareResultPayload) error { - ctx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - env, err := proto.NewEnvelope(proto.TypeRuntimePrepareResult, id, result) - if err != nil { - return err - } - return r.sender.Send(ctx, env) +func (r *Router) sendRuntimePrepareResult(ctx context.Context, request proto.Envelope, result proto.RuntimePrepareResultPayload) error { + return r.reply(ctx, request, proto.TypeRuntimePrepareResult, result) } diff --git a/apps/daemon/internal/dispatch/runtime_preparation_test.go b/apps/daemon/internal/dispatch/runtime_preparation_test.go index feb7fbf6e..d1eadabda 100644 --- a/apps/daemon/internal/dispatch/runtime_preparation_test.go +++ b/apps/daemon/internal/dispatch/runtime_preparation_test.go @@ -30,9 +30,12 @@ func (s *capabilitiesTestSender) Send(ctx context.Context, env proto.Envelope) e } } +// capabilityRef is the assignment every capabilitiesTestRouter binds. +var capabilityRef = proto.AssignmentRef{SessionID: "0b6f1f3e-6f0a-4d38-9c1e-2f5d7a8b9c10", AssignmentID: "assignment", Epoch: 1} + func capabilitiesTestRouter(t *testing.T) (*Router, *capabilitiesTestSender, string, string) { t.Helper() - environment, session := uuid.NewString(), uuid.NewString() + environment, session := uuid.NewString(), capabilityRef.SessionID binding, err := localworkspace.New(environment, session, t.TempDir()) if err != nil { t.Fatal(err) @@ -42,6 +45,7 @@ func capabilitiesTestRouter(t *testing.T) (*Router, *capabilitiesTestSender, str if err != nil { t.Fatal(err) } + bindAssignment(router, capabilityRef, environment) return router, sender, environment, session } @@ -51,6 +55,7 @@ func capabilityEnvelope(t *testing.T, id string, request proto.RuntimePreparePay if err != nil { t.Fatal(err) } + env.Assignment = capabilityRef return env } @@ -169,7 +174,7 @@ func TestRuntimePreparationBeginRequiresExactBindingAndBounds(t *testing.T) { } func TestRuntimePreparationPreparationExcludesOwnedResources(t *testing.T) { - for _, mode := range []string{"write", "export", "read", "run", "idle", "executor", "preparation"} { + for _, mode := range []string{"write", "export", "read", "run", "executor", "preparation"} { t.Run(mode, func(t *testing.T) { r, sender, environment, session := capabilitiesTestRouter(t) switch mode { @@ -181,8 +186,6 @@ func TestRuntimePreparationPreparationExcludesOwnedResources(t *testing.T) { r.workspaceReads = map[string]struct{}{"read": {}} case "run": r.sessions["run"] = &sessionState{} - case "idle": - r.idle["state"] = map[*sessionState]struct{}{} case "executor": r.executors[session] = &executorState{} case "preparation": @@ -202,7 +205,6 @@ func TestRuntimePreparationPreparationExcludesOwnedResources(t *testing.T) { r.workspaceExport = nil r.workspaceReads = nil clear(r.sessions) - clear(r.idle) clear(r.executors) clear(r.preparations) shutdownCapabilitiesRouter(t, r) @@ -217,7 +219,7 @@ func TestRuntimePreparationUploadBlocksWorkspaceWriteAndSuspension(t *testing.T) t.Fatal(err) } capabilitiesReceipt(t, sender, id, "ready") - if err := r.Quiesce(t.Context(), proto.EnvironmentSuspendPayload{EnvironmentID: environment, SuspendID: uuid.NewString()}); !errors.Is(err, ErrRouterBusy) { + if err := r.Quiesce(t.Context(), capabilityRef, proto.EnvironmentSuspendPayload{EnvironmentID: environment, SuspendID: uuid.NewString()}); !errors.Is(err, ErrRouterBusy) { t.Fatal(err) } digest := sha256.Sum256([]byte("abc")) @@ -225,6 +227,7 @@ func TestRuntimePreparationUploadBlocksWorkspaceWriteAndSuspension(t *testing.T) if err != nil { t.Fatal(err) } + write.Assignment = capabilityRef if err := r.Handle(t.Context(), write); err != nil { t.Fatal(err) } diff --git a/apps/daemon/internal/dispatch/shutdown.go b/apps/daemon/internal/dispatch/shutdown.go index 585b806be..de4dd9133 100644 --- a/apps/daemon/internal/dispatch/shutdown.go +++ b/apps/daemon/internal/dispatch/shutdown.go @@ -4,8 +4,6 @@ import ( "context" "errors" "fmt" - - "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" ) type shutdownAttempt struct { @@ -33,29 +31,14 @@ func (r *Router) Shutdown(ctx context.Context) error { first := !r.closed var victims []sessionCancellation for _, state := range r.sessions { - state.retain = false - if state.preparedHandoff != nil { - release, attempt := r.claimPreparedReleaseLocked(state, true, "", true) - victims = append(victims, sessionCancellation{runID: state.runID, handoff: state.preparedHandoff, release: release, attempt: attempt}) - } else if first { - victims = append(victims, sessionCancellation{runID: state.runID, ctxCancel: state.ctxCancel, session: state.session}) - } + release, attempt := r.claimPreparedReleaseLocked(state, true, "", true) + victims = append(victims, sessionCancellation{runID: state.runID, release: release, attempt: attempt}) } if first { r.closed = true if r.runtimePreparation != nil { r.runtimePreparation.cancel() } - for _, states := range r.idle { - for state := range states { - state.retain = false - if state.idleTimer != nil { - state.idleTimer.Stop() - } - victims = append(victims, sessionCancellation{runID: state.runID, ctxCancel: state.ctxCancel, session: state.session}) - } - } - r.idle = make(map[string]map[*sessionState]struct{}) // Prepared release claims exist before this signal can interrupt output. close(r.shutdownCh) } @@ -78,19 +61,10 @@ func (r *Router) Shutdown(ctx context.Context) error { func (r *Router) runShutdownAttempt(attempt *shutdownAttempt, victims []sessionCancellation) { var releaseErr error for _, victim := range victims { - if victim.handoff != nil { - // The cleanup operation has its own fixed native deadline. The - // Shutdown caller's deadline only bounds its wait for this attempt. - if err := r.awaitPreparedNativeRelease(context.Background(), victim.release, victim.attempt); err != nil { - releaseErr = errors.Join(releaseErr, fmt.Errorf("dispatch: prepared run %s: %w", victim.runID, err)) - } - continue - } - victim.ctxCancel() - if victim.session != nil { - if err := victim.session.Cancel(context.Background()); err != nil { - r.log.Warn("session.Cancel failed", "run_id", victim.runID, "err", err) - } + // The cleanup operation has its own fixed native deadline. The + // Shutdown caller's deadline only bounds its wait for this attempt. + if err := r.awaitPreparedNativeRelease(context.Background(), victim.release, victim.attempt); err != nil { + releaseErr = errors.Join(releaseErr, fmt.Errorf("dispatch: prepared run %s: %w", victim.runID, err)) } } r.shutdownWG.Done() @@ -142,10 +116,7 @@ func waitShutdown(ctx context.Context, attempt *shutdownAttempt) error { } type sessionCancellation struct { - runID string - ctxCancel context.CancelFunc - session agent.Session - handoff *preparedHandoff - release *preparedRelease - attempt *preparedReleaseAttempt + runID string + release *preparedRelease + attempt *preparedReleaseAttempt } diff --git a/apps/daemon/internal/dispatch/steering.go b/apps/daemon/internal/dispatch/steering.go index df7b098c5..b4cd28629 100644 --- a/apps/daemon/internal/dispatch/steering.go +++ b/apps/daemon/internal/dispatch/steering.go @@ -47,7 +47,7 @@ func (r *Router) handlePromptSteer(ctx context.Context, env proto.Envelope) erro } func (r *Router) sendSteeringAck(ctx context.Context, env proto.Envelope, ack proto.PromptSteerAckPayload) error { - reply, err := proto.NewEnvelopeWithTrace(proto.TypePromptSteerAck, env.ID, ack, env.Trace) + reply, err := env.Reply(proto.TypePromptSteerAck, ack) if err != nil { return err } @@ -63,6 +63,10 @@ func (r *Router) queueSteering(ctx context.Context, env proto.Envelope, input pr ack.ErrorCode, ack.Error = "run_inactive", "The run is no longer active." return &ack } + if code := r.admitRunLocked(env.Assignment, state); code != "" { + ack.ErrorCode, ack.Error = code, "The run's assignment does not admit this input." + return &ack + } encoded, _ := json.Marshal(input.Input) fingerprint := sha256.Sum256(encoded) if previous, ok := state.steering[input.InputID]; ok { @@ -99,7 +103,7 @@ func (r *Router) queueSteering(ctx context.Context, env proto.Envelope, input pr session := state.session steerer, supportsSteering := session.(agent.Steerer) if input.DurableReceipt { - if _, ok := session.(agent.DurableSteerer); !ok || (!state.releaseOnCompletion && state.preparedHandoff == nil) { + if _, ok := session.(agent.DurableSteerer); !ok { ack.ErrorCode, ack.Error = "unsupported", "Durable input receipts require a supported Turn settlement contract." return &ack } @@ -120,8 +124,6 @@ func (r *Router) queueSteering(ctx context.Context, env proto.Envelope, input pr ack.ErrorCode, ack.Error = "in_flight", "This input is awaiting an engine receipt." state.steering[input.InputID] = steeringReceipt{fingerprint: fingerprint, ack: ack, durable: input.DurableReceipt} state.steerBusy = true - finished := make(chan struct{}) - state.steeringDone = finished r.shutdownWG.Add(1) go func() { defer r.shutdownWG.Done() @@ -129,7 +131,6 @@ func (r *Router) queueSteering(ctx context.Context, env proto.Envelope, input pr defer func() { r.mu.Lock() state.steerBusy = false - close(finished) r.mu.Unlock() }() ctx, stop := r.shutdownContext(ctx) @@ -180,25 +181,3 @@ func (r *Router) shutdownContext(parent context.Context) (context.Context, conte }() return ctx, cancel } - -// Close admission before taking the last worker: its lifetime includes the -// receipt send, so a durable Done cannot close the gateway subscription first. -func (r *Router) finishSteering(state *sessionState) error { - r.mu.Lock() - state.steeringClosed = true - finished := state.steeringDone - r.mu.Unlock() - if finished == nil { - return nil - } - timer := time.NewTimer(steeringCallTimeout + steeringSendTimeout) - defer timer.Stop() - select { - case <-finished: - return nil - case <-r.shutdownCh: - return context.Canceled - case <-timer.C: - return context.DeadlineExceeded - } -} diff --git a/apps/daemon/internal/dispatch/steering_lifetime_test.go b/apps/daemon/internal/dispatch/steering_lifetime_test.go index b6879868c..0c4c47115 100644 --- a/apps/daemon/internal/dispatch/steering_lifetime_test.go +++ b/apps/daemon/internal/dispatch/steering_lifetime_test.go @@ -26,7 +26,7 @@ func TestDurableSteeringWaitsBeyondTransportDeadline(t *testing.T) { var session *fakeSession var calls atomic.Int32 release := make(chan struct{}) - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(h.reg, proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { session = &fakeSession{out: out, closeOutOnCancel: true} return &durableSteeringSession{steeringSession: &steeringSession{fakeSession: session}, phased: func(ctx context.Context, input proto.PromptSteerPayload, written func()) error { calls.Add(1) @@ -39,12 +39,9 @@ func TestDurableSteeringWaitsBeyondTransportDeadline(t *testing.T) { } }}, nil }) - ctx := context.Background() - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, "durable", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", ReleaseOnCompletion: true}))); err != nil { - t.Fatal(err) - } + startRun(t, h.router, h.sender, "codex", "durable") input := proto.PromptSteerPayload{InputID: "extra", Input: proto.TextInput("additional"), DurableReceipt: true} - env := mustEnv(t, proto.TypePromptSteer, "durable", input) + env := scoped(t, "durable", proto.TypePromptSteer, "durable", input) if err := handleSteeringAndWait(t, h, env); err != nil { t.Fatal(err) } @@ -52,7 +49,7 @@ func TestDurableSteeringWaitsBeyondTransportDeadline(t *testing.T) { t.Fatalf("write phase: %+v", ack) } time.Sleep(11 * time.Second) - if len(h.sender.snapshot()) != 1 { + if len(h.sender.typesFor("durable")) != 1 { t.Fatal("native wait ended at transport deadline") } if err := handleSteeringAndWait(t, h, env); err != nil { @@ -62,14 +59,14 @@ func TestDurableSteeringWaitsBeyondTransportDeadline(t *testing.T) { t.Fatalf("cached phase: %+v", ack) } close(release) - waitFor(t, func() bool { return len(h.sender.snapshot()) == 3 }, "native acceptance") + waitFor(t, func() bool { return len(h.sender.typesFor("durable")) == 3 }, "native acceptance") if ack := lastSteeringAck(t, h.sender, "durable", "extra"); !ack.Accepted || ack.Written { t.Fatalf("final phase: %+v", ack) } session.out <- mustEnv(t, proto.TypeDone, "durable", proto.DonePayload{}) waitFor(t, func() bool { return h.router.ActiveRuns() == 0 }, "completion") - frames := h.sender.snapshot() - if calls.Load() != 1 || frames[len(frames)-1].Type != proto.TypeDone { + frames := h.sender.typesFor("durable") + if calls.Load() != 1 || frames[len(frames)-1] != proto.TypeDone { t.Fatal("replayed input or incorrect completion order") } } @@ -80,7 +77,7 @@ func TestDurableSteeringTransportTimeoutAndShutdown(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) exited := make(chan struct{}) - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(_ context.Context, _ proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(h.reg, proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, func(_ context.Context, _ proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { return &durableSteeringSession{steeringSession: &steeringSession{fakeSession: &fakeSession{out: out, closeOutOnCancel: true}}, phased: func(ctx context.Context, _ proto.PromptSteerPayload, written func()) error { defer close(exited) if phase == "written" { @@ -91,10 +88,8 @@ func TestDurableSteeringTransportTimeoutAndShutdown(t *testing.T) { }}, nil }) ctx := context.Background() - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, "run", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", ReleaseOnCompletion: true}))); err != nil { - t.Fatal(err) - } - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptSteer, "run", proto.PromptSteerPayload{InputID: "one", Input: proto.TextInput("text"), DurableReceipt: true})); err != nil { + startRun(t, h.router, h.sender, "codex", "run") + if err := h.router.Handle(ctx, scoped(t, "run", proto.TypePromptSteer, "run", proto.PromptSteerPayload{InputID: "one", Input: proto.TextInput("text"), DurableReceipt: true})); err != nil { t.Fatal(err) } if phase == "blocked-write" { @@ -103,12 +98,12 @@ func TestDurableSteeringTransportTimeoutAndShutdown(t *testing.T) { case <-time.After(12 * time.Second): t.Fatal("blocked write was not bounded") } - waitFor(t, func() bool { return len(h.sender.snapshot()) == 1 }, "unknown receipt") + waitFor(t, func() bool { return len(h.sender.typesFor("run")) == 1 }, "unknown receipt") if ack := lastSteeringAck(t, h.sender, "run", "one"); ack.Written || ack.Accepted || ack.ErrorCode != "outcome_unknown" { t.Fatalf("transport uncertainty: %+v", ack) } } else { - waitFor(t, func() bool { return len(h.sender.snapshot()) == 1 }, "written phase") + waitFor(t, func() bool { return len(h.sender.typesFor("run")) == 1 }, "written phase") stopCtx, cancel := context.WithTimeout(ctx, time.Second) defer cancel() if err := h.router.Shutdown(stopCtx); err != nil { @@ -123,32 +118,3 @@ func TestDurableSteeringTransportTimeoutAndShutdown(t *testing.T) { }) } } - -func TestDurableSteeringRequiresOptInAndAdapter(t *testing.T) { - for _, supported := range []bool{false, true} { - t.Run(map[bool]string{false: "old-adapter", true: "retained-run"}[supported], func(t *testing.T) { - h := newHarness(t) - defer h.router.Shutdown(context.Background()) - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(_ context.Context, _ proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { - s := &steeringSession{fakeSession: &fakeSession{out: out, closeOutOnCancel: true}, steer: func(context.Context, proto.PromptSteerPayload) error { t.Error("unexpected legacy call"); return nil }} - if !supported { - return s, nil - } - return &durableSteeringSession{steeringSession: s, phased: func(context.Context, proto.PromptSteerPayload, func()) error { - t.Error("unexpected phased call") - return nil - }}, nil - }) - ctx := context.Background() - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, "run", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex", ReleaseOnCompletion: !supported}))); err != nil { - t.Fatal(err) - } - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run", proto.PromptSteerPayload{InputID: "one", Input: proto.TextInput("text"), DurableReceipt: true})); err != nil { - t.Fatal(err) - } - if ack := lastSteeringAck(t, h.sender, "run", "one"); ack.ErrorCode != "unsupported" { - t.Fatalf("capability gate: %+v", ack) - } - }) - } -} diff --git a/apps/daemon/internal/dispatch/steering_test.go b/apps/daemon/internal/dispatch/steering_test.go index 1ac57dfec..d9c6eef4a 100644 --- a/apps/daemon/internal/dispatch/steering_test.go +++ b/apps/daemon/internal/dispatch/steering_test.go @@ -31,7 +31,7 @@ func TestSteeringReceiptsAndRetries(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) calls, starts := 0, 0 - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(h.reg, proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { starts++ return &steeringSession{ fakeSession: &fakeSession{out: out, closeOutOnCancel: true}, @@ -43,7 +43,7 @@ func TestSteeringReceiptsAndRetries(t *testing.T) { if input.InputID != "input-1" || *input.Input[0].Content[0].Text != "additional text" { t.Errorf("input lost: %+v", input) } - if len(h.sender.snapshot()) != 0 { + if hasFrame(h.sender, proto.TypePromptSteerAck, "run-1") { t.Error("ack sent before engine accepted input") } return engineError @@ -51,11 +51,9 @@ func TestSteeringReceiptsAndRetries(t *testing.T) { }, nil }) ctx := context.Background() - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, "run-1", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex"}))); err != nil { - t.Fatal(err) - } + startRun(t, h.router, h.sender, "codex", "run-1") input := proto.PromptSteerPayload{InputID: "input-1", Input: proto.TextInput("additional text")} - env := mustEnv(t, proto.TypePromptSteer, "run-1", input) + env := scoped(t, "run-1", proto.TypePromptSteer, "run-1", input) // An ack transport failure must not cause another native invocation. h.sender.failNow = true if err := h.router.Handle(ctx, env); err != nil { @@ -75,7 +73,7 @@ func TestSteeringReceiptsAndRetries(t *testing.T) { } } input.Input = proto.TextInput("changed text") - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run-1", input)); err != nil { + if err := handleSteeringAndWait(t, h, scoped(t, "run-1", proto.TypePromptSteer, "run-1", input)); err != nil { t.Fatal(err) } if ack := lastSteeringAck(t, h.sender, "run-1", "input-1"); ack.ErrorCode != "input_conflict" { @@ -92,7 +90,7 @@ func TestSteeringReadinessAndUnsupportedRuns(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) calls := 0 - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(h.reg, proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { return &steeringSession{ fakeSession: &fakeSession{out: out, closeOutOnCancel: true}, steer: func(context.Context, proto.PromptSteerPayload) error { @@ -105,11 +103,9 @@ func TestSteeringReadinessAndUnsupportedRuns(t *testing.T) { }, nil }) ctx := context.Background() - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, "run-1", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex"}))); err != nil { - t.Fatal(err) - } + startRun(t, h.router, h.sender, "codex", "run-1") input := proto.PromptSteerPayload{InputID: "input-1", Input: proto.TextInput("extra")} - env := mustEnv(t, proto.TypePromptSteer, "run-1", input) + env := scoped(t, "run-1", proto.TypePromptSteer, "run-1", input) for _, expected := range []string{"not_ready", ""} { if err := handleSteeringAndWait(t, h, env); err != nil { t.Fatal(err) @@ -119,7 +115,7 @@ func TestSteeringReadinessAndUnsupportedRuns(t *testing.T) { } if expected == "not_ready" { changed := proto.PromptSteerPayload{InputID: "input-1", Input: proto.TextInput("different during startup")} - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run-1", changed)); err != nil { + if err := handleSteeringAndWait(t, h, scoped(t, "run-1", proto.TypePromptSteer, "run-1", changed)); err != nil { t.Fatal(err) } if ack := lastSteeringAck(t, h.sender, "run-1", "input-1"); ack.ErrorCode != "input_conflict" { @@ -127,7 +123,7 @@ func TestSteeringReadinessAndUnsupportedRuns(t *testing.T) { } } } - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptCancel, "run-1", nil)); err != nil { + if err := h.router.Handle(ctx, scoped(t, "run-1", proto.TypePromptCancel, "run-1", nil)); err != nil { t.Fatal(err) } waitFor(t, func() bool { return h.router.ActiveRuns() == 0 }, "cancel cleanup") @@ -137,19 +133,15 @@ func TestSteeringReadinessAndUnsupportedRuns(t *testing.T) { if ack := lastSteeringAck(t, h.sender, "run-1", "input-1"); ack.ErrorCode != "run_inactive" { t.Fatalf("inactive: %+v", ack) } - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, "run-2", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "fake_alpha"}))); err != nil { - t.Fatal(err) - } - session := <-h.gotSess - defer close(session.out) - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run-2", input)); err != nil { + startRun(t, h.router, h.sender, "fake_alpha", "run-2") + if err := handleSteeringAndWait(t, h, scoped(t, "run-2", proto.TypePromptSteer, "run-2", input)); err != nil { t.Fatal(err) } if ack := lastSteeringAck(t, h.sender, "run-2", "input-1"); ack.ErrorCode != "unsupported" { t.Fatalf("unsupported: %+v", ack) } input.Input = proto.TextInput("") - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run-2", input)); err != nil { + if err := handleSteeringAndWait(t, h, scoped(t, "run-2", proto.TypePromptSteer, "run-2", input)); err != nil { t.Fatal(err) } if ack := lastSteeringAck(t, h.sender, "run-2", "input-1"); ack.ErrorCode != "invalid_input" { @@ -157,7 +149,7 @@ func TestSteeringReadinessAndUnsupportedRuns(t *testing.T) { } // Whitespace-only text is content: dispatch forwards it like any other text. input.InputID, input.Input = "input-2", proto.TextInput(" \n ") - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run-2", input)); err != nil { + if err := handleSteeringAndWait(t, h, scoped(t, "run-2", proto.TypePromptSteer, "run-2", input)); err != nil { t.Fatal(err) } if ack := lastSteeringAck(t, h.sender, "run-2", "input-2"); ack.ErrorCode != "unsupported" { @@ -170,7 +162,7 @@ func TestSteeringDoesNotBlockOtherRunCancellation(t *testing.T) { defer h.router.Shutdown(context.Background()) entered, release := make(chan struct{}), make(chan struct{}) defer close(release) - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(h.reg, proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { return &steeringSession{ fakeSession: &fakeSession{out: out, closeOutOnCancel: true}, steer: func(context.Context, proto.PromptSteerPayload) error { @@ -182,15 +174,12 @@ func TestSteeringDoesNotBlockOtherRunCancellation(t *testing.T) { }) ctx := context.Background() for _, run := range []struct{ id, engine string }{{"run-1", "codex"}, {"run-2", "fake_alpha"}} { - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, run.id, prototest.WithModel(proto.PromptRequestPayload{AgentKind: run.engine}))); err != nil { - t.Fatal(err) - } + startRun(t, h.router, h.sender, run.engine, run.id) } other := <-h.gotSess - other.closeOutOnCancel = true returned := make(chan error, 1) go func() { - returned <- h.router.Handle(ctx, mustEnv(t, proto.TypePromptSteer, "run-1", proto.PromptSteerPayload{InputID: "slow", Input: proto.TextInput("extra")})) + returned <- h.router.Handle(ctx, scoped(t, "run-1", proto.TypePromptSteer, "run-1", proto.PromptSteerPayload{InputID: "slow", Input: proto.TextInput("extra")})) }() select { case err := <-returned: @@ -201,12 +190,10 @@ func TestSteeringDoesNotBlockOtherRunCancellation(t *testing.T) { t.Fatal("steering blocked dispatch") } <-entered - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptCancel, "run-2", nil)); err != nil { + if err := h.router.Handle(ctx, scoped(t, "run-2", proto.TypePromptCancel, "run-2", nil)); err != nil { t.Fatal(err) } - if other.cancels() != 1 { - t.Fatal("other run cancellation blocked") - } + waitFor(t, func() bool { return other.cancels() == 1 }, "other run cancellation") } func lastSteeringAck(t *testing.T, sender *recSender, runID, inputID string) proto.PromptSteerAckPayload { @@ -233,7 +220,7 @@ func TestSteeringCapacityPreservesExistingReceipts(t *testing.T) { h := newHarness(t) defer h.router.Shutdown(context.Background()) calls := 0 - h.reg.RegisterKind(proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, prototest.ModelConfiguration(), func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { + registerSession(h.reg, proto.SupportedAgentKind{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Steering: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported})}, func(ctx context.Context, req proto.PromptRequestPayload, out chan<- proto.Envelope) (agent.Session, error) { return &steeringSession{ fakeSession: &fakeSession{out: out, closeOutOnCancel: true}, steer: func(context.Context, proto.PromptSteerPayload) error { @@ -242,13 +229,10 @@ func TestSteeringCapacityPreservesExistingReceipts(t *testing.T) { }, }, nil }) - ctx := context.Background() - if err := h.router.Handle(ctx, mustEnv(t, proto.TypePromptRequest, "run-1", prototest.WithModel(proto.PromptRequestPayload{AgentKind: "codex"}))); err != nil { - t.Fatal(err) - } + startRun(t, h.router, h.sender, "codex", "run-1") for i := range 257 { input := proto.PromptSteerPayload{InputID: fmt.Sprintf("input-%d", i), Input: proto.TextInput("extra")} - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run-1", input)); err != nil { + if err := handleSteeringAndWait(t, h, scoped(t, "run-1", proto.TypePromptSteer, "run-1", input)); err != nil { t.Fatal(err) } ack := lastSteeringAck(t, h.sender, "run-1", input.InputID) @@ -256,7 +240,7 @@ func TestSteeringCapacityPreservesExistingReceipts(t *testing.T) { t.Fatalf("input %d: %+v", i, ack) } } - if err := handleSteeringAndWait(t, h, mustEnv(t, proto.TypePromptSteer, "run-1", proto.PromptSteerPayload{InputID: "input-0", Input: proto.TextInput("extra")})); err != nil { + if err := handleSteeringAndWait(t, h, scoped(t, "run-1", proto.TypePromptSteer, "run-1", proto.PromptSteerPayload{InputID: "input-0", Input: proto.TextInput("extra")})); err != nil { t.Fatal(err) } if ack := lastSteeringAck(t, h.sender, "run-1", "input-0"); !ack.Accepted || calls != 256 { diff --git a/apps/daemon/internal/dispatch/suspend.go b/apps/daemon/internal/dispatch/suspend.go index 003c9ce2f..077ca9b98 100644 --- a/apps/daemon/internal/dispatch/suspend.go +++ b/apps/daemon/internal/dispatch/suspend.go @@ -11,10 +11,17 @@ import ( var ErrRouterQuiesced = errors.New("dispatch: router quiesced") var ErrRouterBusy = errors.New("dispatch: router has unsettled work") +// AssignmentError rejects a frame whose assignment does not admit it. Its +// value is the error code. +type AssignmentError string + +func (e AssignmentError) Error() string { return "dispatch: " + string(e) } + // Quiesce serializes against admission, then drains every admitted output and // receipt before acknowledging suspension. Busy rejection leaves admission open; // a drain timeout keeps it closed until the caller shuts the connection down. -func (r *Router) Quiesce(ctx context.Context, request proto.EnvironmentSuspendPayload) error { +// ref must admit work in the suspended Environment. +func (r *Router) Quiesce(ctx context.Context, ref proto.AssignmentRef, request proto.EnvironmentSuspendPayload) error { if strings.TrimSpace(request.EnvironmentID) == "" || strings.TrimSpace(request.SuspendID) == "" || len(request.SuspendID) > 128 { return errors.New("dispatch: invalid suspension identity") } @@ -31,6 +38,10 @@ func (r *Router) Quiesce(ctx context.Context, request proto.EnvironmentSuspendPa r.mu.Unlock() return ErrRouterQuiesced } + if code := r.admitLocked(ref, ref.SessionID, request.EnvironmentID); code != "" { + r.mu.Unlock() + return AssignmentError(code) + } if r.runtimePreparation != nil || len(r.sessions) != 0 || len(r.workspaceReads) != 0 || r.workspaceWrite != nil || r.workspaceExport != nil || len(r.permIndex) != 0 || len(r.askIndex) != 0 { r.mu.Unlock() return ErrRouterBusy @@ -41,14 +52,6 @@ func (r *Router) Quiesce(ctx context.Context, request proto.EnvironmentSuspendPa return ErrRouterBusy } } - for _, states := range r.idle { - for state := range states { - if state.environmentID != request.EnvironmentID || state.steerBusy || len(state.pendingIDs) != 0 || len(state.pendingAsks) != 0 { - r.mu.Unlock() - return ErrRouterBusy - } - } - } for _, owner := range r.executors { if owner.invalid || owner.preparing || owner.run != nil || owner.admission != nil || owner.environmentID != request.EnvironmentID { r.mu.Unlock() @@ -61,14 +64,6 @@ func (r *Router) Quiesce(ctx context.Context, request proto.EnvironmentSuspendPa p.timer.Stop() } } - for _, states := range r.idle { - for state := range states { - state.idleLease++ - if state.idleTimer != nil { - state.idleTimer.Stop() - } - } - } for _, owner := range r.executors { owner.idleLease++ if owner.timer != nil { @@ -87,8 +82,10 @@ func (r *Router) Quiesce(ctx context.Context, request proto.EnvironmentSuspendPa } // Resume opens admission only after the caller authenticated a new connection -// and Core confirmed the exact suspension identity on that connection. -func (r *Router) Resume(request proto.EnvironmentSuspendPayload, sender Sender) error { +// and Core confirmed the exact suspension identity on that connection. ref +// names the assignment, which a release during the suspension may have +// advanced. +func (r *Router) Resume(ref proto.AssignmentRef, request proto.EnvironmentSuspendPayload, sender Sender) error { r.admission.Lock() defer r.admission.Unlock() r.mu.Lock() @@ -99,15 +96,13 @@ func (r *Router) Resume(request proto.EnvironmentSuspendPayload, sender Sender) if sender == nil || r.suspension == nil || !r.suspension.SameSuspension(request) { return errors.New("dispatch: suspension identity mismatch") } + if a := r.assignments[ref.SessionID]; !ref.Valid() || a == nil || a.ref.AssignmentID != ref.AssignmentID || ref.Epoch > a.ref.Epoch { + return AssignmentError(proto.AssignmentConflict) + } r.sender = sender r.suspension = nil for _, owner := range r.executors { r.scheduleExecutorIdleLocked(owner) } - for _, states := range r.idle { - for state := range states { - r.scheduleIdleLocked(state) - } - } return nil } diff --git a/apps/daemon/internal/dispatch/suspend_test.go b/apps/daemon/internal/dispatch/suspend_test.go index 64049b954..3a05102f6 100644 --- a/apps/daemon/internal/dispatch/suspend_test.go +++ b/apps/daemon/internal/dispatch/suspend_test.go @@ -3,7 +3,6 @@ package dispatch import ( "context" "errors" - "sync/atomic" "testing" "time" @@ -15,11 +14,14 @@ type suspendSender func(context.Context, proto.Envelope) error func (s suspendSender) Send(ctx context.Context, env proto.Envelope) error { return s(ctx, env) } -type suspendedSession struct{ cancelled atomic.Int32 } +var suspendRef = proto.AssignmentRef{SessionID: "session", AssignmentID: "assignment", Epoch: 1} -func (s *suspendedSession) CancellationOutcome() proto.DonePayload { return proto.DonePayload{} } - -func (s *suspendedSession) Cancel(context.Context) error { s.cancelled.Add(1); return nil } +// bindAssignment records ref as bound in environmentID, as assignment_bind does. +func bindAssignment(r *Router, ref proto.AssignmentRef, environmentID string) { + r.mu.Lock() + r.assignments[ref.SessionID] = &assignmentState{ref: ref, environmentID: environmentID} + r.mu.Unlock() +} func suspensionRouter(t *testing.T, sender Sender) *Router { t.Helper() @@ -27,6 +29,7 @@ func suspensionRouter(t *testing.T, sender Sender) *Router { if err != nil { t.Fatal(err) } + bindAssignment(r, suspendRef, "env") t.Cleanup(func() { if err := r.Shutdown(context.Background()); err != nil { t.Error(err) @@ -50,7 +53,7 @@ func TestQuiesceRejectsEveryUnsettledResource(t *testing.T) { t.Run(name, func(t *testing.T) { r := suspensionRouter(t, suspendSender(func(context.Context, proto.Envelope) error { return nil })) setup(r) - err := r.Quiesce(context.Background(), proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) + err := r.Quiesce(context.Background(), suspendRef, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) if !errors.Is(err, ErrRouterBusy) { t.Fatalf("quiesce=%v", err) } @@ -72,7 +75,7 @@ func TestQuiesceDrainsPendingReceiptAndFencesConcurrentAdmission(t *testing.T) { <-entered request := proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"} quiet := make(chan error, 1) - go func() { quiet <- r.Quiesce(context.Background(), request) }() + go func() { quiet <- r.Quiesce(context.Background(), suspendRef, request) }() deadline := time.After(time.Second) for { r.mu.Lock() @@ -90,7 +93,7 @@ func TestQuiesceDrainsPendingReceiptAndFencesConcurrentAdmission(t *testing.T) { } admitted := make(chan error, 1) go func() { - admitted <- r.Handle(context.Background(), proto.Envelope{Type: proto.TypePromptRequest, ID: "late"}) + admitted <- r.Handle(context.Background(), proto.Envelope{Type: proto.TypeExecutionPrepare, ID: "late", Assignment: suspendRef}) }() select { case err := <-quiet: @@ -106,38 +109,28 @@ func TestQuiesceDrainsPendingReceiptAndFencesConcurrentAdmission(t *testing.T) { } } -func TestQuiescePreservesIdleOwnerAgainstExpiredTimerAndRequiresExactResume(t *testing.T) { +func TestResumeRequiresExactSuspensionAndAssignment(t *testing.T) { sender := suspendSender(func(context.Context, proto.Envelope) error { return nil }) r := suspensionRouter(t, sender) - session := &suspendedSession{} - state := &sessionState{runID: "run", environmentID: "env", stateKey: "state", session: session, ctxCancel: func() {}, retain: true} - r.mu.Lock() - r.idle["state"] = map[*sessionState]struct{}{state: {}} - r.scheduleIdleLocked(state) - oldLease := state.idleLease - r.mu.Unlock() request := proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"} - if err := r.Quiesce(context.Background(), request); err != nil { - t.Fatal(err) + foreign := suspendRef + foreign.AssignmentID = "other" + if err := r.Quiesce(context.Background(), foreign, request); !errors.Is(err, AssignmentError(proto.AssignmentConflict)) { + t.Fatalf("foreign quiesce = %v", err) } - r.expireIdle(state, oldLease) - if session.cancelled.Load() != 0 { - t.Fatal("pre-snapshot timer killed retained owner") + if err := r.Quiesce(context.Background(), suspendRef, request); err != nil { + t.Fatal(err) } wrong := request wrong.SuspendID = "obsolete" - if err := r.Resume(wrong, sender); err == nil { + if err := r.Resume(suspendRef, wrong, sender); err == nil { t.Fatal("stale operation reopened admission") } - if err := r.Resume(request, sender); err != nil { - t.Fatal(err) + if err := r.Resume(foreign, request, sender); err == nil { + t.Fatal("foreign assignment reopened admission") } - r.mu.Lock() - newLease := state.idleLease - r.mu.Unlock() - r.expireIdle(state, newLease) - if session.cancelled.Load() != 1 { - t.Fatal("normal idle expiration was not restored") + if err := r.Resume(suspendRef, request, sender); err != nil { + t.Fatal(err) } } @@ -145,13 +138,13 @@ func TestShutdownDestroysQuiescedOwnerAndCannotResume(t *testing.T) { sender := suspendSender(func(context.Context, proto.Envelope) error { return nil }) r := suspensionRouter(t, sender) request := proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"} - if err := r.Quiesce(context.Background(), request); err != nil { + if err := r.Quiesce(context.Background(), suspendRef, request); err != nil { t.Fatal(err) } if err := r.Shutdown(context.Background()); err != nil { t.Fatal(err) } - if !errors.Is(r.Resume(request, sender), ErrRouterClosed) { + if !errors.Is(r.Resume(suspendRef, request, sender), ErrRouterClosed) { t.Fatal("closed Router resurrected") } } @@ -161,10 +154,10 @@ func TestQuiesceDrainDeadlineCannotReopenAdmission(t *testing.T) { r.shutdownWG.Add(1) ctx, cancel := context.WithCancel(context.Background()) cancel() - if err := r.Quiesce(ctx, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}); !errors.Is(err, context.Canceled) { + if err := r.Quiesce(ctx, suspendRef, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}); !errors.Is(err, context.Canceled) { t.Fatalf("quiesce=%v", err) } - if err := r.Handle(context.Background(), proto.Envelope{Type: proto.TypePromptRequest, ID: "late"}); !errors.Is(err, ErrRouterQuiesced) { + if err := r.Handle(context.Background(), proto.Envelope{Type: proto.TypeExecutionPrepare, ID: "late", Assignment: suspendRef}); !errors.Is(err, ErrRouterQuiesced) { t.Fatalf("deadline reopened admission: %v", err) } r.shutdownWG.Done() diff --git a/apps/daemon/internal/dispatch/workspace_directory_test.go b/apps/daemon/internal/dispatch/workspace_directory_test.go index 1c584484f..49686cca6 100644 --- a/apps/daemon/internal/dispatch/workspace_directory_test.go +++ b/apps/daemon/internal/dispatch/workspace_directory_test.go @@ -36,7 +36,7 @@ func TestWorkspaceDirectoryRetainsEnvironmentAndTransferredOwner(t *testing.T) { bad := request bad.EnvironmentID = "another-environment" _ = r.Handle(t.Context(), mustEnv(t, proto.TypeWorkspaceRead, phase+"-foreign", bad)) - if got := waitWorkspaceRead(t, sender, phase+"-foreign"); got.ErrorCode != "resource_unavailable" { + if got := waitWorkspaceRead(t, sender, phase+"-foreign"); got.ErrorCode != proto.AssignmentConflict { t.Fatal(got) } if phase == "idle" { diff --git a/apps/daemon/internal/dispatch/workspace_export.go b/apps/daemon/internal/dispatch/workspace_export.go index 55af959f4..9366ed690 100644 --- a/apps/daemon/internal/dispatch/workspace_export.go +++ b/apps/daemon/internal/dispatch/workspace_export.go @@ -10,7 +10,7 @@ import ( ) type workspaceExport struct { - id string + request proto.Envelope requests chan proto.WorkspaceExportPayload cancel context.CancelFunc } @@ -27,9 +27,9 @@ func (r *Router) handleWorkspaceExport(ctx context.Context, env proto.Envelope) } u := r.workspaceExport if request.Step != "begin" { - if u == nil || u.id != env.ID { + if u == nil || u.request.ID != env.ID || u.request.Assignment != env.Assignment { r.mu.Unlock() - return r.sendWorkspaceExport(ctx, env.ID, proto.WorkspaceExportResultPayload{Outcome: "rejected", ErrorCode: "resource_unavailable"}) + return r.sendWorkspaceExport(ctx, env, proto.WorkspaceExportResultPayload{Outcome: "rejected", ErrorCode: "resource_unavailable"}) } if request.Step == "cancel" { u.cancel() @@ -46,15 +46,18 @@ func (r *Router) handleWorkspaceExport(ctx context.Context, env proto.Envelope) return errors.New("dispatch: workspace export request already pending") } } - _, code := r.workspaceResourceLocked(proto.WorkspaceReadPayload{Handle: request.Handle, EnvironmentID: request.EnvironmentID}) + _, code := r.workspaceResourceLocked(env.Assignment, proto.WorkspaceReadPayload{Handle: request.Handle, EnvironmentID: request.EnvironmentID}) p := r.preparations[request.Handle] - if u != nil || r.workspaceWrite != nil || !r.localWorkspace.CanExport() || code != "" || p == nil || !p.workspaceReadOnly { + if code == "" && (u != nil || r.workspaceWrite != nil || !r.localWorkspace.CanExport() || p == nil || !p.workspaceReadOnly) { + code = "resource_unavailable" + } + if code != "" { r.mu.Unlock() - return r.sendWorkspaceExport(ctx, env.ID, proto.WorkspaceExportResultPayload{Outcome: "rejected", ErrorCode: "resource_unavailable"}) + return r.sendWorkspaceExport(ctx, env, proto.WorkspaceExportResultPayload{Outcome: "rejected", ErrorCode: code}) } owner, stop := r.shutdownContext(p.ctx) owner, cancel := context.WithTimeout(owner, 180*time.Second) - u = &workspaceExport{id: env.ID, requests: make(chan proto.WorkspaceExportPayload, 1), cancel: func() { cancel(); stop() }} + u = &workspaceExport{request: env, requests: make(chan proto.WorkspaceExportPayload, 1), cancel: func() { cancel(); stop() }} u.requests <- request r.workspaceExport = u r.shutdownWG.Add(1) @@ -91,7 +94,7 @@ func (r *Router) runWorkspaceExport(ctx context.Context, u *workspaceExport) { select { case request := <-u.requests: if request.Offset != offset { - _ = r.sendWorkspaceExport(ctx, u.id, proto.WorkspaceExportResultPayload{Outcome: "failed", Offset: offset, ErrorCode: "invalid_request"}) + _ = r.sendWorkspaceExport(ctx, u.request, proto.WorkspaceExportResultPayload{Outcome: "failed", Offset: offset, ErrorCode: "invalid_request"}) return } case <-ctx.Done(): @@ -117,18 +120,12 @@ func (r *Router) runWorkspaceExport(ctx context.Context, u *workspaceExport) { } r.mu.Unlock() } - if r.sendWorkspaceExport(ctx, u.id, result) != nil || result.Outcome != "chunk" { + if r.sendWorkspaceExport(ctx, u.request, result) != nil || result.Outcome != "chunk" { return } } } -func (r *Router) sendWorkspaceExport(ctx context.Context, id string, result proto.WorkspaceExportResultPayload) error { - ctx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - env, err := proto.NewEnvelope(proto.TypeWorkspaceExportResult, id, result) - if err != nil { - return err - } - return r.sender.Send(ctx, env) +func (r *Router) sendWorkspaceExport(ctx context.Context, request proto.Envelope, result proto.WorkspaceExportResultPayload) error { + return r.reply(ctx, request, proto.TypeWorkspaceExportResult, result) } diff --git a/apps/daemon/internal/dispatch/workspace_export_test.go b/apps/daemon/internal/dispatch/workspace_export_test.go index 90dba1e5e..64520ff74 100644 --- a/apps/daemon/internal/dispatch/workspace_export_test.go +++ b/apps/daemon/internal/dispatch/workspace_export_test.go @@ -46,7 +46,7 @@ func exporterRouter(t *testing.T, program string) (*Router, exportSender, proto. } f.Close() } - environment, session := uuid.NewString(), uuid.NewString() + environment, session := uuid.NewString(), capabilityRef.SessionID binding, err := localworkspace.New(environment, session, workspace) if err != nil { t.Fatal(err) @@ -56,8 +56,9 @@ func exporterRouter(t *testing.T, program string) (*Router, exportSender, proto. if err != nil { t.Fatal(err) } + bindAssignment(r, capabilityRef, environment) handle := uuid.NewString() - r.preparations[handle] = &preparationState{workspaceReadOnly: true, environmentID: environment, owns: true, ctx: context.Background(), deadline: time.Now().Add(time.Hour), status: proto.PreparationStatusPayload{State: "ready"}} + r.preparations[handle] = &preparationState{request: proto.Envelope{Assignment: capabilityRef}, workspaceReadOnly: true, environmentID: environment, owns: true, ctx: context.Background(), deadline: time.Now().Add(time.Hour), status: proto.PreparationStatusPayload{State: "ready"}} t.Cleanup(func() { r.mu.Lock() delete(r.preparations, handle) @@ -77,6 +78,7 @@ func sendExport(t *testing.T, r *Router, id string, p proto.WorkspaceExportPaylo if err != nil { t.Fatal(err) } + env.Assignment = capabilityRef if err := r.Handle(t.Context(), env); err != nil { t.Fatal(err) } diff --git a/apps/daemon/internal/dispatch/workspace_read.go b/apps/daemon/internal/dispatch/workspace_read.go index 157ff0f32..528999397 100644 --- a/apps/daemon/internal/dispatch/workspace_read.go +++ b/apps/daemon/internal/dispatch/workspace_read.go @@ -36,7 +36,7 @@ func (r *Router) handleWorkspaceRead(ctx context.Context, env proto.Envelope) er r.mu.Unlock() return r.sendWorkspaceRead(ctx, env, rejectedWorkspaceRead("read_capacity")) } - resource, code := r.workspaceResourceLocked(request) + resource, code := r.workspaceResourceLocked(env.Assignment, request) if code != "" { r.mu.Unlock() return r.sendWorkspaceRead(ctx, env, rejectedWorkspaceRead(code)) @@ -59,11 +59,16 @@ func (r *Router) handleWorkspaceRead(ctx context.Context, env proto.Envelope) er return nil } -func (r *Router) workspaceResourceLocked(request proto.WorkspaceReadPayload) (any, string) { +// workspaceResourceLocked returns what ref reads: the preparation or the run +// it admitted. +func (r *Router) workspaceResourceLocked(ref proto.AssignmentRef, request proto.WorkspaceReadPayload) (any, string) { + if code := r.admitLocked(ref, ref.SessionID, request.EnvironmentID); code != "" { + return nil, code + } var resource any if request.Handle != "" { p := r.preparations[request.Handle] - if p == nil || p.environmentID != request.EnvironmentID || p.status.State != "ready" || + if p == nil || p.request.Assignment != ref || p.environmentID != request.EnvironmentID || p.status.State != "ready" || !p.owns || p.busy || p.ctx.Err() != nil || !time.Now().Before(p.deadline) { return nil, "resource_unavailable" } @@ -73,7 +78,7 @@ func (r *Router) workspaceResourceLocked(request proto.WorkspaceReadPayload) (an } } else { s := r.sessions[request.RunID] - if s == nil || s.environmentID != request.EnvironmentID || s.session == nil || + if s == nil || s.assignment != ref || s.environmentID != request.EnvironmentID || s.session == nil || !r.interactionRouteOpenLocked(s) { return nil, "resource_unavailable" } @@ -143,15 +148,8 @@ func workspaceReadResult(read agent.WorkspaceReadResult, err error, limit int) p } func (r *Router) sendWorkspaceRead(ctx context.Context, request proto.Envelope, result proto.WorkspaceReadResultPayload) error { - ctx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - trace := request.Trace - if len(trace) > 256 { - trace = "" - } - env, err := proto.NewEnvelopeWithTrace(proto.TypeWorkspaceReadResult, request.ID, result, trace) - if err != nil { - return err + if len(request.Trace) > 256 { + request.Trace = "" } - return r.sender.Send(ctx, env) + return r.reply(ctx, request, proto.TypeWorkspaceReadResult, result) } diff --git a/apps/daemon/internal/dispatch/workspace_write.go b/apps/daemon/internal/dispatch/workspace_write.go index 38a613278..78d367f11 100644 --- a/apps/daemon/internal/dispatch/workspace_write.go +++ b/apps/daemon/internal/dispatch/workspace_write.go @@ -38,7 +38,7 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e if pending { return errors.New("dispatch: malformed pending write frame") } - return r.sendWorkspaceWrite(ctx, env.ID, rejectedWorkspaceWrite("invalid_request")) + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("invalid_request")) } r.mu.Lock() if r.closed { @@ -48,7 +48,7 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e if request.Step == "begin" { if r.workspaceExport != nil || r.runtimePreparation != nil { r.mu.Unlock() - return r.sendWorkspaceWrite(ctx, env.ID, rejectedWorkspaceWrite("resource_unavailable")) + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("resource_unavailable")) } if r.workspaceWrite != nil { duplicate := r.workspaceWrite.envelope.ID == env.ID @@ -56,22 +56,26 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e if duplicate { return errors.New("dispatch: workspace write already admitted") } - return r.sendWorkspaceWrite(ctx, env.ID, rejectedWorkspaceWrite("write_capacity")) + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("write_capacity")) } - if !r.localWorkspace.AcceptsFileWrite(request.EnvironmentID, request.SessionID) || len(r.sessions) != 0 || len(r.idle) != 0 || len(r.workspaceReads) != 0 { + if code := r.admitLocked(env.Assignment, request.SessionID, request.EnvironmentID); code != "" { r.mu.Unlock() - return r.sendWorkspaceWrite(ctx, env.ID, rejectedWorkspaceWrite("resource_unavailable")) + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite(code)) + } + if !r.localWorkspace.AcceptsFileWrite(request.EnvironmentID, request.SessionID) || len(r.sessions) != 0 || len(r.workspaceReads) != 0 { + r.mu.Unlock() + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("resource_unavailable")) } for _, owner := range r.executors { if owner.preparing || owner.admission != nil || owner.run != nil || owner.invalid { r.mu.Unlock() - return r.sendWorkspaceWrite(ctx, env.ID, rejectedWorkspaceWrite("resource_unavailable")) + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("resource_unavailable")) } } for _, p := range r.preparations { if p.owns { r.mu.Unlock() - return r.sendWorkspaceWrite(ctx, env.ID, rejectedWorkspaceWrite("resource_unavailable")) + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("resource_unavailable")) } } u := &workspaceUpload{envelope: env, request: request, data: make([]byte, 0, request.SizeBytes), ready: make(chan struct{})} @@ -79,12 +83,12 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e r.shutdownWG.Add(1) r.mu.Unlock() go r.runWorkspaceUpload(context.WithoutCancel(ctx), u) - return r.sendWorkspaceWrite(ctx, env.ID, proto.WorkspaceWriteResultPayload{Outcome: "ready"}) + return r.sendWorkspaceWrite(ctx, env, proto.WorkspaceWriteResultPayload{Outcome: "ready"}) } u := r.workspaceWrite - if u == nil || u.envelope.ID != env.ID { + if u == nil || u.envelope.ID != env.ID || u.envelope.Assignment != env.Assignment { r.mu.Unlock() - return r.sendWorkspaceWrite(ctx, env.ID, rejectedWorkspaceWrite("resource_unavailable")) + return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("resource_unavailable")) } if u.finished { r.mu.Unlock() @@ -94,7 +98,7 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e u.data = append(u.data, request.Data...) offset := len(u.data) r.mu.Unlock() - return r.sendWorkspaceWrite(ctx, env.ID, proto.WorkspaceWriteResultPayload{Outcome: "received", Offset: offset}) + return r.sendWorkspaceWrite(ctx, env, proto.WorkspaceWriteResultPayload{Outcome: "received", Offset: offset}) } if request.Step == "commit" && len(u.data) == u.request.SizeBytes { digest := sha256.Sum256(u.data) @@ -132,7 +136,7 @@ func (r *Router) runWorkspaceUpload(ctx context.Context, u *workspaceUpload) { r.workspaceWrite = nil } r.mu.Unlock() - _ = r.sendWorkspaceWrite(ctx, u.envelope.ID, result) + _ = r.sendWorkspaceWrite(ctx, u.envelope, result) } func rejectedWorkspaceWrite(code string) proto.WorkspaceWriteResultPayload { @@ -171,12 +175,6 @@ func workspaceWriteResult(write agent.WorkspaceWriteResult, err error, size int) return proto.WorkspaceWriteResultPayload{Outcome: "unknown", ErrorCode: "write_unconfirmed"} } -func (r *Router) sendWorkspaceWrite(ctx context.Context, id string, result proto.WorkspaceWriteResultPayload) error { - ctx, cancel := context.WithTimeout(ctx, 5*time.Second) - defer cancel() - env, err := proto.NewEnvelope(proto.TypeWorkspaceWriteResult, id, result) - if err != nil { - return err - } - return r.sender.Send(ctx, env) +func (r *Router) sendWorkspaceWrite(ctx context.Context, request proto.Envelope, result proto.WorkspaceWriteResultPayload) error { + return r.reply(ctx, request, proto.TypeWorkspaceWriteResult, result) } diff --git a/apps/daemon/internal/dispatch/workspace_write_test.go b/apps/daemon/internal/dispatch/workspace_write_test.go index e6e710fe9..132773cd4 100644 --- a/apps/daemon/internal/dispatch/workspace_write_test.go +++ b/apps/daemon/internal/dispatch/workspace_write_test.go @@ -19,7 +19,7 @@ import ( func localWriterRouter(t *testing.T) (*dispatch.Router, *recSender, proto.WorkspaceWritePayload, string) { t.Helper() workspace := t.TempDir() - environment, session := uuid.NewString(), uuid.NewString() + environment, session := uuid.NewString(), preparationSessionID binding, err := localworkspace.New(environment, session, workspace) if err != nil { t.Fatal(err) @@ -34,6 +34,7 @@ func localWriterRouter(t *testing.T) (*dispatch.Router, *recSender, proto.Worksp defer cancel() _ = r.Shutdown(ctx) }) + assign(t, r, session, environment) digest := sha256.Sum256([]byte("abc")) return r, sender, proto.WorkspaceWritePayload{Step: "begin", EnvironmentID: environment, SessionID: session, Path: "file", SizeBytes: 3, SHA256: hex.EncodeToString(digest[:])}, workspace } diff --git a/apps/daemon/internal/transport/ws_test.go b/apps/daemon/internal/transport/ws_test.go index 13d5e2390..01cee4ff9 100644 --- a/apps/daemon/internal/transport/ws_test.go +++ b/apps/daemon/internal/transport/ws_test.go @@ -269,7 +269,7 @@ func TestStartHeartbeatsTicks(t *testing.T) { var calls atomic.Int32 conn.StartHeartbeats(context.Background(), 30*time.Millisecond, func() proto.HeartbeatPayload { calls.Add(1) - return proto.HeartbeatPayload{Timestamp: time.Now().Unix(), DaemonVersion: "0.0.0-dev"} + return proto.HeartbeatPayload{Timestamp: time.Now().Unix(), DaemonVersion: "0.0.0-dev", HomeRemoval: proto.CapabilityUnsupported} }, nil) deadline := time.Now().Add(time.Second) @@ -311,7 +311,7 @@ func TestStartHeartbeatsSendsImmediately(t *testing.T) { defer cancel() conn.StartHeartbeats(hbCtx, time.Hour, func() proto.HeartbeatPayload { calls.Add(1) - return proto.HeartbeatPayload{Timestamp: time.Now().Unix(), DaemonVersion: "0.0.0-dev"} + return proto.HeartbeatPayload{Timestamp: time.Now().Unix(), DaemonVersion: "0.0.0-dev", HomeRemoval: proto.CapabilityUnsupported} }, nil) deadline := time.Now().Add(2 * time.Second) diff --git a/apps/daemon/testdata/onboarding/main.go b/apps/daemon/testdata/onboarding/main.go index e63a63ed1..4c35cc0cc 100644 --- a/apps/daemon/testdata/onboarding/main.go +++ b/apps/daemon/testdata/onboarding/main.go @@ -174,7 +174,10 @@ func run() error { return err } defer router.Shutdown(context.Background()) - heartbeat, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{SupportedAgentKinds: registry.SupportedAgentKinds()}) + heartbeat, err := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{SupportedAgentKinds: registry.SupportedAgentKinds(), HomeRemoval: proto.CapabilityUnsupported}) + if err != nil { + return err + } if err = sink.Send(context.Background(), heartbeat); err != nil { return err } diff --git a/docs/architecture.md b/docs/architecture.md index 889874b55..18f8bf4bf 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -44,6 +44,6 @@ The [repository map](./development.md#repository-map) locates these components. An application creates a Session through the Agents API. Core resolves its configuration and execution location. A managed Session obtains compute through the selected Sandbox Provider; a self-hosted Session waits for the user to run its installation command. A Session with `environment: none` uses a connected execution device without a workspace. The [application guide](./api/public-agent-api.md#create-a-session) describes these choices. -After the daemon connects, Core checks the available Harness and requested capabilities. The Runtime prepares a workspace Environment and its capability snapshot, then prepares or reuses the Session Executor. Each Turn runs through the native Harness. Core persists the output, tool interactions and receipts for application reads and events. Completion or cancellation settles the Turn; a healthy Executor can serve the next Turn in the same Environment. +After the daemon connects, Core checks the available Harness and requested capabilities and binds the Session to that Runtime through a fenced [assignment](./runtime-protocol.md#session-assignments), which deleting the Session or releasing its Environment releases. The Runtime prepares a workspace Environment and its capability snapshot, then prepares or reuses the Session Executor. Each Turn runs through the native Harness. Core persists the output, tool interactions and receipts for application reads and events. Completion or cancellation settles the Turn; a healthy Executor can serve the next Turn in the same Environment. Execution and compute have separate lifetimes: closing an Executor preserves its allocation and workspace until the Provider reclaims them. Preparation, connection and execution readiness have distinct states. The [Environment contract](../contracts/agents-api/environments.md) owns preparation, and the [Core–Runtime protocol](./runtime-protocol.md) owns ordering, receipts and failure handling. diff --git a/docs/runtime-protocol.md b/docs/runtime-protocol.md index 050cbb4d7..78f1d9bc2 100644 --- a/docs/runtime-protocol.md +++ b/docs/runtime-protocol.md @@ -15,12 +15,12 @@ A Runtime connects in this order: 1. Obtain a daemon credential and device ID. The [machine connection API](../contracts/agents-api/machine-api.md#credentials) lists the credential kinds; Project API keys and the Core key are never Runtime credentials. 2. Call `POST /api/v1/agent-daemon/bootstrap` with the credential as a Bearer header and the device ID. Use the connection URL it returns. 3. Dial the WebSocket at `/api/v1/agent-daemon/ws` with `device_id` and `version` query parameters and the Bearer header. Never put a credential in a URL, a payload log or a trace. -4. Send a heartbeat at once, then at the interval bootstrap returned. Each heartbeat declares `supported_agent_kinds`, their availability and their [capabilities](#capability-declarations). Before the first heartbeat, capabilities are unknown; a kind missing from a heartbeat is not advertised. Neither permits inference. +4. Send a heartbeat at once, then at the interval bootstrap returned. Each heartbeat declares `supported_agent_kinds`, their availability and their [capabilities](#capability-declarations), and whether the Runtime removes a Session's native home on release (`home_removal`, [Session assignments](#session-assignments)). Before the first heartbeat, capabilities are unknown; a kind missing from a heartbeat is not advertised. Neither permits inference. 5. Exchange ordered JSON [envelopes](#envelope-and-identity). Heartbeats establish liveness only, never execution progress or a receipt for an earlier message. The wire version is [`proto.Version`](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/version.go), independent of the Runtime build version that heartbeats report. Core accepts only an exact match, including the patch component. A mismatch returns HTTP 426 `incompatible_version` before any dispatch; the daemon treats it as permanent and stops reconnecting. Deploy matching peers together. -Each physical connection has fresh routing, admission handles and transfer state. A newer connection for the same device replaces the previous one: Core closes the previous connection and evicts its Run and interaction routes, and the new connection inherits none of them. A valid credential and connection are never authority to choose another Session or Environment binding. +Each physical connection has fresh routing, admission handles and transfer state. A newer connection for the same device replaces the previous one: Core closes the previous connection and evicts its Run and interaction routes, and the new connection inherits none of them. A valid credential and connection are never authority to choose another Session or Environment binding; [assignments](#session-assignments) fence which Session a frame may act for. ## Capability declarations @@ -53,7 +53,7 @@ A declaration describes what the Runtime can do. Core admits a public feature on `permissions` gates permission decisions inside the Runtime. Core has no admission rule for `usage` and `resume`. -The prompt request (`prompt_request`, or the configuration of `execution_prepare`) carries the Session's model configuration and the opt-ins Core sets for each Run: +The configuration of `execution_prepare` carries the Session's model configuration and the opt-ins Core sets for each Run: | Field | Set by Core | | --- | --- | @@ -71,11 +71,12 @@ Requests without an opt-in keep the frames and fields they had without it. ## Envelope and identity -Every data frame is one JSON [`Envelope`](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/envelope.go): `type`, a type-dependent `id`, a typed `payload` and an optional W3C `trace`. The trace is diagnostic correlation only; missing or invalid trace data creates a local trace and never changes ownership. Never use a trace ID as a request ID. +Every data frame is one JSON [`Envelope`](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/envelope.go): `type`, a type-dependent `id`, a typed `payload`, the Session's `assignment` on Session frames and an optional W3C `trace`. The trace is diagnostic correlation only; missing or invalid trace data creates a local trace and never changes ownership. Never use a trace ID as a request ID. | Identity | Scope and meaning | | --- | --- | | Device ID and connection | Authenticated Runtime routing and connection ownership | +| Assignment | `Envelope.assignment`: the Session, assignment ID and epoch that fence the frame; see [Session assignments](#session-assignments) | | Session ID / Environment ID | Core-owned configuration and workspace binding; canonical UUIDs where the payload validator requires them | | Executor ID | Runtime-owned native resource, possibly retained across settled Turns with identical configuration | | Preparation request ID | `Envelope.id` for prepare, start, release and status; distinct from a Run | @@ -95,9 +96,11 @@ The linked source files define the required fields, validators, limits and finit | Core → Runtime | Runtime → Core | Definition | | --- | --- | --- | +| `assignment_bind`, `assignment_release` | `assignment_status` | [Session assignments](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/assignment.go) | +| Any frame the Runtime cannot route | `protocol_error` | [Session assignments](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/assignment.go) | | `runtime_prepare` | `runtime_prepare_result` | [Initialization and capability transfer](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/runtime_prepare.go) | | `execution_prepare`, `execution_start`, `execution_release` | `preparation_status` | [Execution admission](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/preparation.go) | -| `prompt_request`, `prompt_cancel` | `delta`, `thinking`, `output_message`, `tool_call`, `usage`, `error`, `done`, `heartbeat` | [Requests](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/outbound.go), [events and capabilities](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/inbound.go) | +| `prompt_cancel` | `delta`, `thinking`, `output_message`, `tool_call`, `usage`, `error`, `done`, `heartbeat` | [Requests](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/outbound.go), [events and capabilities](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/inbound.go) | | `permission_decision`, `prompt_for_user_choice_decision` | `permission_request`, `permission_cancel`, `prompt_for_user_choice`, `interaction_decision_ack` | [Requests](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/outbound.go), [interactions](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/inbound.go) | | `prompt_steer` | `prompt_steer_ack` | [Active input receipts](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/steering.go) | | `function_result` | `function_call`, `interaction_decision_ack` | [Function calls](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/functions.go) | @@ -108,6 +111,18 @@ Initial, prepared and active input use the same [ordered MessageInput](https://g Usage frames and the final usage snapshot each carry the cumulative measurement of the current execution and replace the previous snapshot; never add them. An absent measurement is unknown, not zero. +## Session assignments + +An assignment binds one Session to the Runtime that runs it. `Envelope.assignment` names it as `session_id`, `assignment_id` and `epoch`, and is the only place a frame carries it. Core advances the epoch whenever it changes the assignment's desired state, so a lower epoch is stale. + +Every Session frame carries the assignment: `execution_prepare`, `execution_start` and `execution_release`; `prompt_cancel`, `prompt_steer`, `function_result`, `permission_decision` and `prompt_for_user_choice_decision`; every frame of `runtime_prepare`, `workspace_read`, `workspace_write` and `workspace_export`; and `environment_quiesce` and `environment_resume`. A reply echoes its request's assignment. Heartbeats carry none. + +Before a Session's first operation on a connection, including Environment initialization and file work without a Turn, Core sends `assignment_bind` with the Session's Environment ID and waits for `assignment_status` `bound`. A repeated bind of the same assignment is `bound` again. The Runtime admits a Session frame only under the assignment it bound: an older epoch, or a released one, fails with `assignment_stale`; another assignment, Session or Environment fails with `assignment_conflict`. A started Run's frames, including its cancellation receipt, stay admissible under the assignment that started it until the release. + +Core records a release and advances the epoch before it sends anything. Deleting a Session releases its assignment with `remove_home: true`; releasing its Environment sends `false`. A deletion never revokes a shared Runtime credential. `assignment_release` fences the assignment at once. The Runtime then stops the Session's work, closes its Executors and, when asked, removes the native home; only then does it reply `released` or `home_removed`. Unfinished cleanup replies `failed` with `cleanup_unconfirmed`, and a retry at the same epoch repeats it. A Runtime that declares `home_removal` unsupported answers `remove_home: true` with `unsupported_operation`, and Core asks it only to release. Core records the release as applied only from a matching `released` or `home_removed`, and resends every unacknowledged release to a Runtime when it connects. A quiesced Runtime admits only a release and the matching `environment_resume`, which carries the assignment that quiesced it. + +The Runtime answers a Core frame it cannot route with `protocol_error`, which echoes the request's ID and carries its type and an error code. + ## Preparation and execution order Environment initialization uses `runtime_prepare` on every connection, managed or user-owned; the [Environment contract](../contracts/agents-api/environments.md#runtime-capability-preparation) owns what is prepared and when. For a file or archive, send `begin`, wait for `ready`, send ordered chunks and await each matching `received` offset, then send `commit` and await `completed`. Initialization and finalization have typed headers without file data. Validate the expected outcome, offset, size and finite error code with the shared validator. One transfer is allowed per connection. A chunk receipt confirms staged bytes, not installation; a completed commit confirms that operation, not that a later Turn ran. @@ -126,8 +141,6 @@ Preparation and start run outside the receive loop and router lock. An admission Idle expiry of an Executor is a Runtime resource policy, separate from Core's active-Turn concurrency. On shutdown the Runtime closes active and idle Executors, keeps any target whose close failed and allows a later serialized retry. An ordinary disconnection closes the failed transport and keeps the exact router until shutdown succeeds; a wait timeout or failed cleanup never authorizes reconnection, and process shutdown keeps waiting rather than discarding owned native resources. Workspace operations keep their binding and settlement rules across Turn boundaries and Executor closure. -`prompt_request` starts a Run directly, without an admission handle. It is not a fallback after a failed prepared start. - ## Active input receipts Core delivers active input as `prompt_steer` with `durable_receipt: true`, one input at a time per Run, and waits for its receipt before sending the next: @@ -168,7 +181,7 @@ The public Turn status is a separate projection. [`execution/delivery.go`](https | Unsupported capability or invalid binding | Reject before starting the operation; never select another Harness | | Confirmed preparation or execution failure | Keep the finite error category and any observed result; the Runtime settles its resources | | Deadline or connection loss after dispatch | The effect is unknown unless an application receipt proves otherwise; do not convert it to an execution failure | -| Reconnection | Reestablish the transport and the capability declaration; never replay input, initialization, transfers or unresolved mutations | +| Reconnection | Reestablish the transport and the capability declaration, bind again before the next Session operation and resend unacknowledged releases; never replay input, initialization, transfers or unresolved mutations | | Duplicate preparation or start | Connection-local identity and fingerprint rules apply; a conflicting request rejects, and an old handle cannot start replacement work | | Duplicate input, function result or decision | That family's receipt identity and conflict rules apply; there is no transport-wide deduplication or exactly-once promise | | Cleanup failure | Keep resource ownership and report unconfirmed cleanup; a waiter's timeout does not make a resource reusable | @@ -198,7 +211,7 @@ Core runs an idle directory read on the Worker's Session scheduling reservation ## MCP connection authority -Every public `MCPHTTPServer` in a prompt request carries an explicit `connection_origin`; a missing or unknown value rejects rather than selecting a default, and Core freezes the public default before dispatch. The Runtime validates the origin with the common validator before selecting a factory and resolves public and installed MCP into transient effective bindings. The [Environment contract](../contracts/agents-api/environments.md#public-mcp-connection-origin) owns the supported combinations, native limits and failure ownership. +Every public `MCPHTTPServer` in an execution configuration carries an explicit `connection_origin`; a missing or unknown value rejects rather than selecting a default, and Core freezes the public default before dispatch. The Runtime validates the origin with the common validator before selecting a factory and resolves public and installed MCP into transient effective bindings. The [Environment contract](../contracts/agents-api/environments.md#public-mcp-connection-origin) owns the supported combinations, native limits and failure ownership. ## Contract verification diff --git a/docs/zh/architecture.md b/docs/zh/architecture.md index 8777d50e1..f74dca545 100644 --- a/docs/zh/architecture.md +++ b/docs/zh/architecture.md @@ -1,7 +1,7 @@ --- title: "架构" source: docs/architecture.md -source_hash: b8a00701caa83314115f28c3fa754291e8eaf734e88de8e0012e45d8b1893338 +source_hash: 2ee1d5723194cb689b2540167c093b03f738ffc2f5e766922416dd7b954b31d0 --- OpenAgentCore 将编排、计算资源和原生执行分开。Core 负责 API 和持久状态。Sandbox Provider 管理计算资源。Runtime daemon 准备 Environment 并运行选定的 Harness；Harness 的原生 SDK 或协议负责模型与工具循环。 @@ -46,6 +46,6 @@ flowchart TB 应用通过 Agents API 创建 Session。Core 解析其配置与执行位置。托管 Session 通过选定的 Sandbox Provider 获取计算资源；自托管 Session 等待用户运行安装命令。`environment: none` 的 Session 使用已连接的执行设备，不提供工作区。[应用指南](api/public-agent-api.md#create-a-session) 说明这些选项。 -daemon 连接后，Core 检查可用 Harness 和请求的能力。Runtime 准备工作区 Environment 及其能力快照，然后准备或复用 Session Executor。每个 Turn 通过原生 Harness 运行。Core 持久化输出、工具交互和回执，供应用读取和接收事件。完成或取消使 Turn 结算；健康的 Executor 可以在同一 Environment 中执行下一个 Turn。 +daemon 连接后，Core 检查可用 Harness 和请求的能力，并通过受约束的[分配](./runtime-protocol.md#session-assignments)把 Session 绑定到该 Runtime；删除 Session 或释放其 Environment 会释放该分配。Runtime 准备工作区 Environment 及其能力快照，然后准备或复用 Session Executor。每个 Turn 通过原生 Harness 运行。Core 持久化输出、工具交互和回执，供应用读取和接收事件。完成或取消使 Turn 结算；健康的 Executor 可以在同一 Environment 中执行下一个 Turn。 执行与计算资源拥有独立生命周期：关闭 Executor 后，其 allocation 和工作区保留到 Provider 回收为止。准备、连接和执行就绪具有不同状态。[Environment 契约](../../contracts/agents-api/zh/environments.md) 负责准备规则，[Core–Runtime 协议](runtime-protocol.md) 负责顺序、回执和故障处理。 diff --git a/docs/zh/runtime-protocol.md b/docs/zh/runtime-protocol.md index d5eabb785..d40b11f62 100644 --- a/docs/zh/runtime-protocol.md +++ b/docs/zh/runtime-protocol.md @@ -1,7 +1,7 @@ --- title: "Core–Runtime 协议" source: docs/runtime-protocol.md -source_hash: b9e759ff6a1a9f24ec2281130a60d609a17022e082bf2ff5bbc081720c919f49 +source_hash: ca1f05d7f95e6a91f57e44919da83d05e94b146809d727a4e87ad38c94a9d6fd --- 此协议在 Runtime daemon 获取机器凭据后连接 Core 与 daemon，定义 daemon 连接上消息的含义和顺序。wire 类型、限制和验证器仅在 [`internal/agentdaemon/proto`](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/internal/agentdaemon/proto) 中定义一次；Core 的 [gateway](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/services/core/internal/runtimegateway) 与参考 Runtime 的 [dispatcher](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/apps/daemon/internal/dispatch) 都使用它们，因此无需同步第二套 payload schema。签发凭据和打开连接的 HTTP 路由见[机器连接 API](../../contracts/agents-api/zh/machine-api.md)。 @@ -17,12 +17,12 @@ Runtime 按以下顺序连接： 1. 获取 daemon 凭据和 device ID。[机器连接 API](../../contracts/agents-api/zh/machine-api.md#credentials) 列出凭据类型；Project API key 和 Core key 都不是 Runtime 凭据。 2. 调用 `POST /api/v1/agent-daemon/bootstrap`，通过 Bearer header 提供凭据，并提供 device ID。使用返回的连接 URL。 3. 连接 `/api/v1/agent-daemon/ws` 的 WebSocket，传入 `device_id`、`version` 查询参数和 Bearer header。凭据不得出现在 URL、payload 日志或 trace 中。 -4. 立即发送 heartbeat，之后按 bootstrap 返回的间隔发送。每个 heartbeat 声明 `supported_agent_kinds`、其可用性和[能力](#capability-declarations)。第一个 heartbeat 之前能力未知；heartbeat 中缺失的 kind 视为未声明。两者都不允许推断。 +4. 立即发送 heartbeat，之后按 bootstrap 返回的间隔发送。每个 heartbeat 声明 `supported_agent_kinds`、其可用性和[能力](#capability-declarations)，以及 Runtime 是否在释放时删除 Session 的原生 home（`home_removal`，见 [Session 分配](#session-assignments)）。第一个 heartbeat 之前能力未知；heartbeat 中缺失的 kind 视为未声明。两者都不允许推断。 5. 交换有序 JSON [envelope](#envelope-and-identity)。Heartbeat 仅证明存活，不证明执行进度，也不充当此前消息的回执。 wire 版本为 [`proto.Version`](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/version.go)，独立于 heartbeat 报告的 Runtime 构建版本。Core 仅接受精确匹配，包括 patch 部分。不匹配时，在任何 dispatch 前返回 HTTP 426 `incompatible_version`；daemon 将其视为永久错误并停止重连。应一起部署版本匹配的两端。 -每条物理连接拥有新的路由、admission handle 和传输状态。同一设备的新连接替代旧连接：Core 关闭旧连接并移除其 Run 和 interaction 路由，新连接不继承这些状态。有效凭据和连接不授权选择其他 Session 或 Environment 绑定。 +每条物理连接拥有新的路由、admission handle 和传输状态。同一设备的新连接替代旧连接：Core 关闭旧连接并移除其 Run 和 interaction 路由，新连接不继承这些状态。有效凭据和连接不授权选择其他 Session 或 Environment 绑定；[分配](#session-assignments)约束 frame 可代表哪个 Session 行事。 ## 能力声明 {#capability-declarations} @@ -55,7 +55,7 @@ wire 上每个字段都是 JSON boolean，所有字段都必须出现，包括 ` `permissions` 控制 Runtime 内部权限决定。Core 对 `usage` 和 `resume` 没有准入规则。 -prompt 请求（`prompt_request` 或 `execution_prepare` 的配置）携带 Session 的模型配置和 Core 为各 Run 设置的显式启用项： +`execution_prepare` 的配置携带 Session 的模型配置和 Core 为各 Run 设置的显式启用项： | 字段 | Core 设置方式 | | --- | --- | @@ -73,11 +73,12 @@ prompt 请求（`prompt_request` 或 `execution_prepare` 的配置）携带 Sess ## Envelope 与身份 {#envelope-and-identity} -每个数据 frame 都是一个 JSON [`Envelope`](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/envelope.go)：`type`、随类型变化的 `id`、类型化 `payload` 和可选 W3C `trace`。trace 仅用于诊断关联；trace 数据缺失或无效时创建本地 trace，不改变所有权。不要将 trace ID 用作 request ID。 +每个数据 frame 都是一个 JSON [`Envelope`](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/envelope.go)：`type`、随类型变化的 `id`、类型化 `payload`、Session frame 上的 Session `assignment`，以及可选 W3C `trace`。trace 仅用于诊断关联；trace 数据缺失或无效时创建本地 trace，不改变所有权。不要将 trace ID 用作 request ID。 | 身份 | 范围与含义 | | --- | --- | | Device ID 和连接 | 经认证的 Runtime 路由与连接所有权 | +| Assignment | `Envelope.assignment`：约束 frame 的 Session、assignment ID 和 epoch；见 [Session 分配](#session-assignments) | | Session ID / Environment ID | Core 拥有的配置与工作区绑定；payload validator 要求时使用规范 UUID | | Executor ID | Runtime 拥有的原生资源，可在配置相同的已结算 Turn 之间保留 | | Preparation request ID | prepare、start、release 和 status 的 `Envelope.id`；与 Run 不同 | @@ -97,9 +98,11 @@ User-choice decision 携带 `question_answers`：每个已提供回答都有明 | Core → Runtime | Runtime → Core | 定义 | | --- | --- | --- | +| `assignment_bind`, `assignment_release` | `assignment_status` | [Session 分配](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/assignment.go) | +| Runtime 无法路由的任意 frame | `protocol_error` | [Session 分配](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/assignment.go) | | `runtime_prepare` | `runtime_prepare_result` | [初始化与能力传输](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/runtime_prepare.go) | | `execution_prepare`, `execution_start`, `execution_release` | `preparation_status` | [执行准入](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/preparation.go) | -| `prompt_request`, `prompt_cancel` | `delta`, `thinking`, `output_message`, `tool_call`, `usage`, `error`, `done`, `heartbeat` | [请求](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/outbound.go)、[事件与能力](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/inbound.go) | +| `prompt_cancel` | `delta`, `thinking`, `output_message`, `tool_call`, `usage`, `error`, `done`, `heartbeat` | [请求](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/outbound.go)、[事件与能力](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/inbound.go) | | `permission_decision`, `prompt_for_user_choice_decision` | `permission_request`, `permission_cancel`, `prompt_for_user_choice`, `interaction_decision_ack` | [请求](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/outbound.go)、[交互](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/inbound.go) | | `prompt_steer` | `prompt_steer_ack` | [活动输入回执](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/steering.go) | | `function_result` | `function_call`, `interaction_decision_ack` | [Function 调用](https://github.com/MiniMax-AI/OpenAgentCore/blob/main/internal/agentdaemon/proto/functions.go) | @@ -110,6 +113,18 @@ User-choice decision 携带 `question_answers`：每个已提供回答都有明 Usage frame 和最终 usage snapshot 都携带当前执行的累计测量，替换之前的快照；不要相加。缺失的测量表示未知，不是零。 +## Session 分配 {#session-assignments} + +分配（assignment）把一个 Session 绑定到运行它的 Runtime。`Envelope.assignment` 以 `session_id`、`assignment_id` 和 `epoch` 命名它，这是 frame 携带分配的唯一位置。Core 每次改变分配的期望状态时推进 epoch，因此较低的 epoch 是陈旧的。 + +每个 Session frame 都携带分配：`execution_prepare`、`execution_start` 和 `execution_release`；`prompt_cancel`、`prompt_steer`、`function_result`、`permission_decision` 和 `prompt_for_user_choice_decision`；`runtime_prepare`、`workspace_read`、`workspace_write` 和 `workspace_export` 的每个 frame；以及 `environment_quiesce` 和 `environment_resume`。回复回显请求的分配。heartbeat 不携带分配。 + +在一条连接上执行 Session 的第一个操作之前，包括没有 Turn 的 Environment 初始化和文件操作，Core 发送带 Session 的 Environment ID 的 `assignment_bind`，并等待 `assignment_status` `bound`。重复绑定同一分配仍得到 `bound`。Runtime 只在其已绑定的分配下准入 Session frame：较旧的 epoch 或已释放的分配以 `assignment_stale` 失败；其他分配、Session 或 Environment 以 `assignment_conflict` 失败。已启动 Run 的 frame，包括其取消回执，在释放前仍可在启动它的分配下准入。 + +Core 先记录释放并推进 epoch，再发送任何消息。删除 Session 以 `remove_home: true` 释放其分配；释放其 Environment 发送 `false`。删除从不吊销共享的 Runtime 凭据。`assignment_release` 立即约束该分配。随后 Runtime 停止 Session 的工作，关闭其 Executor，并在要求时删除原生 home；此后才回复 `released` 或 `home_removed`。未完成的清理回复 `failed` 和 `cleanup_unconfirmed`，同一 epoch 的重试会重复清理。声明 `home_removal` 不支持的 Runtime 以 `unsupported_operation` 回答 `remove_home: true`，Core 只要求它释放。Core 只根据匹配的 `released` 或 `home_removed` 记录释放已应用，并在 Runtime 连接时重发所有未确认的释放。已 quiesce 的 Runtime 只准入释放和匹配的 `environment_resume`，后者携带使其 quiesce 的分配。 + +Runtime 对无法路由的 Core frame 回复 `protocol_error`，回显请求 ID，并携带其类型和错误码。 + ## 准备与执行顺序 {#preparation-and-execution-order} 无论托管还是用户自有环境，每条连接都通过 `runtime_prepare` 初始化 Environment；[Environment 契约](../../contracts/agents-api/zh/environments.md#runtime-capability-preparation)负责准备内容和时机。传输文件或 archive 时，发送 `begin`，等待 `ready`，发送有序 chunk 并等待每个匹配的 `received` offset，再发送 `commit` 并等待 `completed`。初始化和终结阶段使用不含文件数据的类型化 header。使用共享 validator 验证预期结果、offset、size 和有限错误 code。每条连接允许一个 transfer。chunk 回执确认暂存字节，不确认安装；完成的 commit 确认该操作，不证明后续 Turn 已运行。 @@ -128,8 +143,6 @@ preparation 和 start 在 receive loop 与 router lock 之外运行。admission Executor 空闲到期属于 Runtime 资源策略，与 Core 的活动 Turn 并发限制独立。关闭时 Runtime 关闭活动和空闲 Executor，保留关闭失败的目标，并允许稍后串行重试。普通断连会关闭失败的 transport 并保留原 router，直到 shutdown 成功；等待超时或清理失败不授权重连，进程 shutdown 继续等待，不丢弃自己拥有的原生资源。工作区操作在跨 Turn 和 Executor 关闭后仍保留绑定与结算规则。 -`prompt_request` 不经过 admission handle，直接启动 Run。它不是 prepared start 失败后的回退。 - ## 活动输入回执 {#active-input-receipts} Core 通过 `prompt_steer` 交付活动输入，设置 `durable_receipt: true`，每个 Run 一次交付一个输入，并等待回执后再发送下一个： @@ -170,7 +183,7 @@ transport 与 execution 结果分开。Core 唯一 Run 订阅入口是 `Subscrib | 不支持的能力或无效绑定 | 在开始操作前拒绝；不选择其他 Harness | | 已确认准备或执行失败 | 保留有限错误类别和已观测结果；Runtime 结算自己的资源 | | dispatch 后截止时间到达或连接丢失 | 除非应用回执另有证明，效果未知；不转换为执行失败 | -| 重连 | 重新建立 transport 和能力声明；不重放输入、初始化、传输或未决修改 | +| 重连 | 重新建立 transport 和能力声明，在下一次 Session 操作前重新绑定，并重发未确认的释放；不重放输入、初始化、传输或未决修改 | | 重复 preparation 或 start | 应用连接本地身份与 fingerprint 规则；冲突请求被拒绝，旧 handle 不能启动替代工作 | | 重复 input、function result 或 decision | 应用对应消息族的回执身份和冲突规则；不存在 transport 全局去重或 exactly-once 保证 | | 清理失败 | 保留资源所有权，报告清理未确认；等待方超时不使资源可复用 | @@ -200,7 +213,7 @@ Core 在 Worker 的 Session 调度预约上运行空闲目录读取，活动执 ## MCP 连接权限 {#mcp-connection-authority} -prompt request 中每个公开 `MCPHTTPServer` 都携带明确的 `connection_origin`；值缺失或未知时拒绝，不选择默认值，Core 在 dispatch 前冻结公开默认值。Runtime 在选择 factory 前用公共 validator 验证 origin，并将公开和已安装 MCP 解析为临时 effective binding。[Environment 契约](../../contracts/agents-api/zh/environments.md#public-mcp-connection-origin)负责支持的组合、原生限制和故障所有权。 +执行配置中每个公开 `MCPHTTPServer` 都携带明确的 `connection_origin`；值缺失或未知时拒绝，不选择默认值，Core 在 dispatch 前冻结公开默认值。Runtime 在选择 factory 前用公共 validator 验证 origin，并将公开和已安装 MCP 解析为临时 effective binding。[Environment 契约](../../contracts/agents-api/zh/environments.md#public-mcp-connection-origin)负责支持的组合、原生限制和故障所有权。 ## 契约验证 {#contract-verification} diff --git a/internal/agentdaemon/proto/assignment.go b/internal/agentdaemon/proto/assignment.go new file mode 100644 index 000000000..84470cda0 --- /dev/null +++ b/internal/agentdaemon/proto/assignment.go @@ -0,0 +1,76 @@ +package proto + +// Assignment traffic binds a Session to the Runtime that runs it. Envelope.ID +// correlates a request with its assignment_status, and Envelope.Assignment +// names the assignment on both. +const ( + // TypeAssignmentBind binds the Session to this Runtime. The Runtime + // replies bound before Core sends the Session's first operation. + TypeAssignmentBind = "assignment_bind" + // TypeAssignmentRelease ends the assignment. The Runtime stops the + // Session's work and closes its Executors before it replies. + TypeAssignmentRelease = "assignment_release" + TypeAssignmentStatus = "assignment_status" + // TypeProtocolError answers a request the Runtime cannot route. It + // carries the request's ID and type. + TypeProtocolError = "protocol_error" +) + +// Assignment states an assignment_status reports. +const ( + AssignmentBound = "bound" + AssignmentReleased = "released" + AssignmentHomeRemoved = "home_removed" + AssignmentFailed = "failed" +) + +// Error codes of assignment fencing. A frame whose reference is older than +// the Runtime's assignment, or names a released one, is stale; a frame that +// names no assignment the Runtime holds, or another Session or Environment, +// conflicts. +const ( + AssignmentStale = "assignment_stale" + AssignmentConflict = "assignment_conflict" + // UnsupportedOperation rejects an operation the Runtime does not declare, + // before any effect. + UnsupportedOperation = "unsupported_operation" + // CleanupUnconfirmed reports a release whose cleanup did not finish; the + // assignment stays released and a retry repeats the cleanup. + CleanupUnconfirmed = "cleanup_unconfirmed" +) + +// AssignmentRef names one assignment of a Session to a Runtime. Core advances +// Epoch whenever it changes the assignment's desired state. +type AssignmentRef struct { + SessionID string `json:"session_id"` + AssignmentID string `json:"assignment_id"` + Epoch uint64 `json:"epoch"` +} + +// Valid reports whether the reference names an assignment. +func (r AssignmentRef) Valid() bool { + return r.SessionID != "" && len(r.SessionID) <= 128 && r.AssignmentID != "" && len(r.AssignmentID) <= 128 && r.Epoch > 0 +} + +// AssignmentBindPayload is what the Runtime fences the Session's frames with. +// EnvironmentID is empty for a Session without an Environment. +type AssignmentBindPayload struct { + EnvironmentID string `json:"environment_id,omitempty"` +} + +// AssignmentReleasePayload asks the Runtime to remove the Session's native +// home after its work stops. Only a Runtime that declares home removal accepts +// RemoveHome. +type AssignmentReleasePayload struct { + RemoveHome bool `json:"remove_home"` +} + +type AssignmentStatusPayload struct { + State string `json:"state"` + ErrorCode string `json:"error_code,omitempty"` +} + +type ProtocolErrorPayload struct { + Type string `json:"type"` + ErrorCode string `json:"error_code"` +} diff --git a/internal/agentdaemon/proto/envelope.go b/internal/agentdaemon/proto/envelope.go index 94fde404d..25a8591ab 100644 --- a/internal/agentdaemon/proto/envelope.go +++ b/internal/agentdaemon/proto/envelope.go @@ -12,14 +12,18 @@ // a lookup table. // // Envelope.ID correlation: -// - prompt_request / prompt_cancel: ID = RunID. +// - prompt_cancel / prompt_steer / function_result: ID = RunID. // - delta / tool_call / usage / error / done: ID = originating RunID. // - permission_request: ID = RunID; payload.request_id is the interaction ID. // - permission_decision / permission_cancel: ID = interaction ID. // - execution_prepare / execution_start / execution_release and // preparation_status: ID = preparation request ID, never RunID. // - runtime_prepare / runtime_prepare_result: ID = connection-local transfer ID. +// - assignment_bind / assignment_release / assignment_status: ID = request ID. // - heartbeats carry no ID. +// +// Envelope.Assignment names the Session's assignment on every frame of a +// Session or Environment operation and on the Runtime's replies to it. package proto import ( @@ -51,6 +55,10 @@ type Envelope struct { // Receivers MUST tolerate missing/unparseable values — both mean // "mint a fresh trace locally", never reject. Trace string `json:"trace,omitempty"` + + // Assignment is the only location of a frame's assignment reference. + // Heartbeats carry none. + Assignment AssignmentRef `json:"assignment,omitzero"` } // NewEnvelope marshals payload into an Envelope. A nil payload yields @@ -81,6 +89,17 @@ func NewEnvelopeWithTrace(typ string, id string, payload any, traceparent string return env, nil } +// Reply builds the Runtime's reply to the request e. It carries the request's +// ID, trace and assignment. +func (e Envelope) Reply(typ string, payload any) (Envelope, error) { + reply, err := NewEnvelopeWithTrace(typ, e.ID, payload, e.Trace) + if err != nil { + return Envelope{}, err + } + reply.Assignment = e.Assignment + return reply, nil +} + // DecodePayload unpacks Envelope.Payload into out. An empty Payload is // a non-error so bodyless types (prompt_cancel, permission_cancel) // decode cleanly. diff --git a/internal/agentdaemon/proto/envelope_test.go b/internal/agentdaemon/proto/envelope_test.go index ee8f6ab5f..6c36fc215 100644 --- a/internal/agentdaemon/proto/envelope_test.go +++ b/internal/agentdaemon/proto/envelope_test.go @@ -18,6 +18,8 @@ func TestEnvelopeRoundTrip(t *testing.T) { if env.ID != "run-123" { t.Fatalf("ID = %q, want %q", env.ID, "run-123") } + ref := AssignmentRef{SessionID: "session", AssignmentID: "assignment", Epoch: 2} + env.Assignment = ref raw, err := json.Marshal(env) if err != nil { t.Fatalf("Marshal envelope: %v", err) @@ -33,6 +35,9 @@ func TestEnvelopeRoundTrip(t *testing.T) { if out.Delta != "hello" || out.Sequence != 42 { t.Fatalf("payload round-trip lost data: %+v", out) } + if reply, err := got.Reply(TypeDone, nil); err != nil || got.Assignment != ref || reply.Assignment != ref || reply.ID != "run-123" { + t.Fatalf("assignment round trip = %+v, reply = %+v, %v", got.Assignment, reply, err) + } } func TestEnvelopeOmitsEmptyPayload(t *testing.T) { @@ -66,8 +71,7 @@ func TestDecodeRequestRejectsUndeclaredMembers(t *testing.T) { payload string out any }{ - "unknown key": {TypePromptRequest, `{"agent_kind":"codex","model":"m","unknown":true}`, &PromptRequestPayload{}}, - "agent_options": {TypePromptRequest, `{"agent_kind":"codex","agent_options":{"model":"m"}}`, &PromptRequestPayload{}}, + "unknown key": {TypeExecutionPrepare, `{"session_id":"s","configuration":{"agent_kind":"codex","model":"m","unknown":true}}`, &ExecutionPreparePayload{}}, "agent_options in preparation": {TypeExecutionPrepare, `{"session_id":"s","configuration":{"agent_kind":"codex","agent_options":{}}}`, &ExecutionPreparePayload{}}, } { env := Envelope{Type: c.typ, Payload: json.RawMessage(c.payload)} diff --git a/internal/agentdaemon/proto/inbound.go b/internal/agentdaemon/proto/inbound.go index 73a49f6a8..84d369d18 100644 --- a/internal/agentdaemon/proto/inbound.go +++ b/internal/agentdaemon/proto/inbound.go @@ -266,9 +266,11 @@ type SupportedAgentKind struct { // HeartbeatPayload advertises only explicit engine descriptors. Missing // supported_agent_kinds establishes no engine availability or capabilities. +// HomeRemoval declares whether assignment_release accepts RemoveHome. type HeartbeatPayload struct { Timestamp int64 `json:"ts"` ActiveRequests int `json:"active_requests"` DaemonVersion string `json:"daemon_version,omitempty"` SupportedAgentKinds []SupportedAgentKind `json:"supported_agent_kinds,omitempty"` + HomeRemoval CapabilitySupport `json:"home_removal"` } diff --git a/internal/agentdaemon/proto/outbound.go b/internal/agentdaemon/proto/outbound.go index d5e49e3ba..f4bc5774e 100644 --- a/internal/agentdaemon/proto/outbound.go +++ b/internal/agentdaemon/proto/outbound.go @@ -8,12 +8,6 @@ import ( // Type constants for server → daemon frames. const ( - // TypePromptRequest triggers one prompt cycle. Envelope.ID = RunID; - // the daemon stamps every resulting upstream frame with the same - // ID so the gateway can fan them back to the matching StreamPrompt - // subscriber. - TypePromptRequest = "prompt_request" - // TypePromptCancel aborts an in-flight prompt. Envelope.ID = // RunID. Idempotent — cancelling an unknown / already-finished // run is a no-op on the daemon side. @@ -86,6 +80,12 @@ type PromptRequestPayload struct { ToolSearch bool `json:"tool_search,omitempty"` DisableExecutionEnvironment bool `json:"disable_execution_environment,omitempty"` DisableSubagents bool `json:"disable_subagents,omitempty"` + + // Assignment is the assignment the request runs under: Core sends its + // frames under it, and the Runtime records the Envelope.Assignment it + // admitted the request under for the Executor factory. It is never + // encoded. + Assignment AssignmentRef `json:"-"` } // PromptCancelPayload optionally requests an application receipt; identity is on Envelope.ID. diff --git a/internal/agentdaemon/proto/prototest/wire.go b/internal/agentdaemon/proto/prototest/wire.go index 2f05c524c..001fe6910 100644 --- a/internal/agentdaemon/proto/prototest/wire.go +++ b/internal/agentdaemon/proto/prototest/wire.go @@ -28,11 +28,17 @@ const ( const ( SessionID = "session" StateKey = "agents-api-session" + AssignmentID = "assignment" + BindID = "bind" PreparationID = "prepare" RunID = "run" CancelDeliveryID = "cancel" ) +// Assignment is the Session's assignment. Every scenario frame is +// Session-scoped and carries it. +var Assignment = proto.AssignmentRef{SessionID: SessionID, AssignmentID: AssignmentID, Epoch: 1} + // Placeholders for values the Runtime generates. Core's side sends them as // written; the Runtime's side binds each to the value its Runtime sends. const ( @@ -112,7 +118,13 @@ func WireScenarios() []WireScenario { SessionID: SessionID, Configuration: WithModel(proto.PromptRequestPayload{AgentKind: HarnessKind, AgentStateKey: StateKey, StrictResume: true, DisableExecutionEnvironment: true}), }) + // Core binds the Session before its first Session-scoped frame. + bind := []Step{ + send(Core, proto.TypeAssignmentBind, BindID, proto.AssignmentBindPayload{}), + send(Runtime, proto.TypeAssignmentStatus, BindID, proto.AssignmentStatusPayload{State: proto.AssignmentBound}), + } started := []Step{ + bind[0], bind[1], prepare, status(1, "preparing", "", ""), status(2, "ready", "", ""), @@ -135,6 +147,7 @@ func WireScenarios() []WireScenario { Name: "preparation_failure_cleans_up_without_run_completion", NativeSetupFails: true, Steps: []Step{ + bind[0], bind[1], prepare, status(1, "preparing", "", ""), status(2, "failed", "", "preparation_failed"), @@ -157,6 +170,7 @@ func send(from Peer, kind, id string, payload any) Step { if err != nil { panic(err) } + frame.Assignment = Assignment return Step{Action: Send, From: from, Frame: frame} } @@ -166,7 +180,7 @@ func status(revision uint64, state, runID, errorCode string) Step { }) } -// SameFrame checks that got carries want's type, correlation ID and payload. +// SameFrame checks that got carries want's type, correlation ID, assignment and payload. // Trace is diagnostic correlation and is not part of a scenario. func SameFrame(want, got proto.Envelope) error { return Bindings{}.match(want, got, false) @@ -213,8 +227,8 @@ func (b Bindings) match(want, got proto.Envelope, bind bool) error { wantFields[field] = value } } - if want.Type != got.Type || want.ID != got.ID || !reflect.DeepEqual(wantFields, gotFields) { - return fmt.Errorf("got %s %q %s, want %s %q %s", got.Type, got.ID, got.Payload, want.Type, want.ID, want.Payload) + if want.Type != got.Type || want.ID != got.ID || want.Assignment != got.Assignment || !reflect.DeepEqual(wantFields, gotFields) { + return fmt.Errorf("got %s %q %+v %s, want %s %q %+v %s", got.Type, got.ID, got.Assignment, got.Payload, want.Type, want.ID, want.Assignment, want.Payload) } for key, value := range bound { b[key] = value diff --git a/internal/agentdaemon/proto/prototest/wire_test.go b/internal/agentdaemon/proto/prototest/wire_test.go index c98cba861..19812a04c 100644 --- a/internal/agentdaemon/proto/prototest/wire_test.go +++ b/internal/agentdaemon/proto/prototest/wire_test.go @@ -15,6 +15,11 @@ func TestBindingsBindRuntimeValuesOnce(t *testing.T) { if err := b.Match(ready, sent("h-1", 3)); err == nil { t.Fatal("a payload difference matched") } + stale := sent("h-1", 2) + stale.Assignment.Epoch++ + if err := b.Match(ready, stale); err == nil { + t.Fatal("another assignment matched") + } if len(b) != 0 { t.Fatal("a failed match bound values") } diff --git a/internal/agentdaemon/proto/runtime_prepare.go b/internal/agentdaemon/proto/runtime_prepare.go index bbb2549c6..01d13224d 100644 --- a/internal/agentdaemon/proto/runtime_prepare.go +++ b/internal/agentdaemon/proto/runtime_prepare.go @@ -146,7 +146,7 @@ func ValidRuntimePrepareResult(r RuntimePrepareResultPayload, expected string, o return false } switch r.ErrorCode { - case "invalid_request", "resource_unavailable", "runtime_preparation_capacity", "runtime_preparation_unsupported", "runtime_preparation_rejected": + case "invalid_request", "resource_unavailable", "runtime_preparation_capacity", "runtime_preparation_unsupported", "runtime_preparation_rejected", AssignmentStale, AssignmentConflict: return true default: return false diff --git a/internal/harnessconfig/preparation_test.go b/internal/harnessconfig/preparation_test.go index 485efaed0..bea7305f2 100644 --- a/internal/harnessconfig/preparation_test.go +++ b/internal/harnessconfig/preparation_test.go @@ -88,7 +88,7 @@ func TestPrepareRejectsDecodedRequestsWithoutModelOrProvider(t *testing.T) { `{"model":null,"model_provider":` + provider + `}`: ErrModel, } { var req proto.PromptRequestPayload - if err := (proto.Envelope{Type: proto.TypePromptRequest, Payload: json.RawMessage(payload)}).DecodeRequest(&req); err != nil { + if err := (proto.Envelope{Payload: json.RawMessage(payload)}).DecodeRequest(&req); err != nil { t.Fatalf("%s: %v", payload, err) } if _, err := c.Prepare(req); !errors.Is(err, want) { diff --git a/services/core/internal/db/queries/core_metrics.sql b/services/core/internal/db/queries/core_metrics.sql index 19cde907a..d30175625 100644 --- a/services/core/internal/db/queries/core_metrics.sql +++ b/services/core/internal/db/queries/core_metrics.sql @@ -1,11 +1,11 @@ -- name: CoreExecutionSnapshot :one SELECT count(*) FILTER (WHERE t.status = 'queued')::bigint AS queued_turns, count(*) FILTER (WHERE t.status = 'queued' AND - (b.device_id IS NULL OR NOT (b.device_id = ANY(sqlc.arg(connected_device_ids)::uuid[]))))::bigint AS waiting_for_daemon, + (b.runtime_id IS NULL OR NOT (b.runtime_id = ANY(sqlc.arg(connected_device_ids)::uuid[]))))::bigint AS waiting_for_daemon, count(*) FILTER (WHERE t.status = 'in_progress')::bigint AS in_progress_turns, COALESCE(GREATEST(0, extract(epoch FROM (sqlc.arg(observed_at)::timestamptz - min(t.created_at) FILTER (WHERE t.status = 'queued')))), 0)::double precision AS oldest_queued_seconds -FROM turns t LEFT JOIN session_devices b ON b.session_id = t.session_id +FROM turns t LEFT JOIN session_runtime_assignments b ON b.session_id = t.session_id WHERE t.status IN ('queued', 'in_progress'); -- name: CoreInterruptedTurns :one diff --git a/services/core/internal/db/queries/devices.sql b/services/core/internal/db/queries/devices.sql index 80262106e..1af723a27 100644 --- a/services/core/internal/db/queries/devices.sql +++ b/services/core/internal/db/queries/devices.sql @@ -20,37 +20,53 @@ UPDATE devices SET last_seen_at = clock_timestamp() WHERE devices.id = $1 AND EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = devices.id); -- name: BindSessionDevice :one -INSERT INTO session_devices (session_id, device_id) +INSERT INTO session_runtime_assignments (session_id, runtime_id) SELECT s.id, d.id FROM sessions s JOIN devices d ON d.tenant_id = s.tenant_id WHERE s.tenant_id = $1 AND s.id = $2 AND d.id = $3 AND d.revoked_at IS NULL AND (d.environment_id IS NULL OR EXISTS ( SELECT 1 FROM environments e WHERE e.id = d.environment_id AND e.session_id = s.id )) -ON CONFLICT (session_id) DO UPDATE SET device_id = session_devices.device_id -WHERE session_devices.device_id = EXCLUDED.device_id -RETURNING device_id; +ON CONFLICT (session_id) DO UPDATE SET runtime_id = session_runtime_assignments.runtime_id +WHERE session_runtime_assignments.runtime_id = EXCLUDED.runtime_id AND session_runtime_assignments.desired_state = 'bound' +RETURNING runtime_id; -- name: GetSessionDevice :one -SELECT d.id, d.name, d.environment_id FROM session_devices b +SELECT d.id, d.name, d.environment_id, b.assignment_id, b.epoch FROM session_runtime_assignments b JOIN sessions s ON s.id = b.session_id -JOIN devices d ON d.id = b.device_id AND d.tenant_id = s.tenant_id -WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL +JOIN devices d ON d.id = b.runtime_id AND d.tenant_id = s.tenant_id +WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL AND b.desired_state = 'bound' AND EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = d.id) AND (d.environment_id IS NULL OR EXISTS ( SELECT 1 FROM environments e WHERE e.id = d.environment_id AND e.session_id = s.id )); -- name: GetSessionExecutionBinding :one -SELECT d.id, d.name, b.native_session_id, d.environment_id, +SELECT d.id, d.name, b.native_session_id, d.environment_id, b.assignment_id, b.epoch, EXISTS (SELECT 1 FROM turns t WHERE t.session_id = s.id AND t.started_at IS NOT NULL) AS has_started_turn -FROM session_devices b +FROM session_runtime_assignments b JOIN sessions s ON s.id = b.session_id -JOIN devices d ON d.id = b.device_id AND d.tenant_id = s.tenant_id -WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL +JOIN devices d ON d.id = b.runtime_id AND d.tenant_id = s.tenant_id +WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL AND b.desired_state = 'bound' AND EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = d.id) AND (d.environment_id IS NULL OR EXISTS ( SELECT 1 FROM environments e WHERE e.id = d.environment_id AND e.session_id = s.id )); -- name: RememberNativeSession :execrows -UPDATE session_devices SET native_session_id = $2 WHERE session_id = $1; +UPDATE session_runtime_assignments SET native_session_id = $2 WHERE session_id = $1; + +-- name: ReleaseSessionAssignment :exec +-- An identical release keeps its epoch; a release that adds home removal +-- advances it. +UPDATE session_runtime_assignments +SET desired_state = 'released', epoch = epoch + 1, remove_home = remove_home OR sqlc.arg(remove_home)::boolean +WHERE session_id = sqlc.arg(session_id) AND (desired_state = 'bound' OR (sqlc.arg(remove_home)::boolean AND NOT remove_home)); + +-- name: ListPendingAssignmentReleases :many +SELECT session_id, runtime_id, assignment_id, epoch, remove_home FROM session_runtime_assignments +WHERE desired_state = 'released' AND applied_epoch < epoch AND runtime_id = ANY(sqlc.arg(runtime_ids)::uuid[]) +ORDER BY runtime_id, session_id; + +-- name: AcknowledgeAssignmentRelease :execrows +UPDATE session_runtime_assignments SET applied_epoch = epoch +WHERE session_id = $1 AND assignment_id = $2 AND epoch = $3 AND desired_state = 'released'; diff --git a/services/core/internal/db/queries/environment_initialization.sql b/services/core/internal/db/queries/environment_initialization.sql index d1452e781..30d7ca519 100644 --- a/services/core/internal/db/queries/environment_initialization.sql +++ b/services/core/internal/db/queries/environment_initialization.sql @@ -1,7 +1,7 @@ -- name: ListEnvironmentInitializations :many -SELECT e.id, e.session_id, s.tenant_id, s.engine, e.initialization, b.device_id +SELECT e.id, e.session_id, s.tenant_id, s.engine, e.initialization, b.runtime_id, b.assignment_id, b.epoch FROM environments e JOIN sessions s ON s.id = e.session_id -LEFT JOIN session_devices b ON b.session_id = s.id +LEFT JOIN session_runtime_assignments b ON b.session_id = s.id AND b.desired_state = 'bound' WHERE e.id > $1 AND s.deleted_at IS NULL AND e.status NOT IN ('failed', 'expired') AND e.initialization IN ('pending', 'running') ORDER BY e.id LIMIT 32; diff --git a/services/core/internal/db/queries/scheduling.sql b/services/core/internal/db/queries/scheduling.sql index 54da1333c..7dac84e27 100644 --- a/services/core/internal/db/queries/scheduling.sql +++ b/services/core/internal/db/queries/scheduling.sql @@ -16,13 +16,13 @@ ORDER BY t.id LIMIT 100; SELECT r.id, r.session_id, s.tenant_id FROM environment_input_reservations r JOIN sessions s ON s.id = r.session_id -LEFT JOIN session_devices b ON b.session_id = s.id +LEFT JOIN session_runtime_assignments b ON b.session_id = s.id WHERE r.state = 'pending' AND r.deadline > clock_timestamp() AND r.id > sqlc.arg(after_id)::uuid AND s.deleted_at IS NULL AND EXISTS ( SELECT 1 FROM devices d WHERE d.tenant_id = s.tenant_id AND d.revoked_at IS NULL AND d.id = ANY(sqlc.arg(connected_devices)::uuid[]) - AND (b.device_id IS NULL OR d.id = b.device_id) + AND (b.runtime_id IS NULL OR d.id = b.runtime_id) ) ORDER BY r.id LIMIT 100; diff --git a/services/core/internal/db/queries/subagent_identities.sql b/services/core/internal/db/queries/subagent_identities.sql index 40d481bd7..c55aec02e 100644 --- a/services/core/internal/db/queries/subagent_identities.sql +++ b/services/core/internal/db/queries/subagent_identities.sql @@ -3,9 +3,9 @@ INSERT INTO subagent_identities ( id, session_id, device_id, engine, native_id, parent_native_id, native_created_at, first_turn_id, first_event_ordinal ) -SELECT sqlc.arg(id), s.id, b.device_id, s.engine, sqlc.arg(native_id), +SELECT sqlc.arg(id), s.id, b.runtime_id, s.engine, sqlc.arg(native_id), sqlc.arg(parent_native_id), sqlc.arg(native_created_at), sqlc.arg(first_turn_id), sqlc.arg(first_event_ordinal) -FROM sessions s JOIN session_devices b ON b.session_id = s.id +FROM sessions s JOIN session_runtime_assignments b ON b.session_id = s.id WHERE s.id = sqlc.arg(session_id) AND (b.native_session_id = sqlc.arg(parent_native_id) OR EXISTS (SELECT 1 FROM subagent_identities p diff --git a/services/core/internal/db/sqlc/core_metrics.sql.go b/services/core/internal/db/sqlc/core_metrics.sql.go index 46fb6f12d..b1f1f8fc8 100644 --- a/services/core/internal/db/sqlc/core_metrics.sql.go +++ b/services/core/internal/db/sqlc/core_metrics.sql.go @@ -25,11 +25,11 @@ func (q *Queries) CoreDatabaseSize(ctx context.Context) (int64, error) { const coreExecutionSnapshot = `-- name: CoreExecutionSnapshot :one SELECT count(*) FILTER (WHERE t.status = 'queued')::bigint AS queued_turns, count(*) FILTER (WHERE t.status = 'queued' AND - (b.device_id IS NULL OR NOT (b.device_id = ANY($1::uuid[]))))::bigint AS waiting_for_daemon, + (b.runtime_id IS NULL OR NOT (b.runtime_id = ANY($1::uuid[]))))::bigint AS waiting_for_daemon, count(*) FILTER (WHERE t.status = 'in_progress')::bigint AS in_progress_turns, COALESCE(GREATEST(0, extract(epoch FROM ($2::timestamptz - min(t.created_at) FILTER (WHERE t.status = 'queued')))), 0)::double precision AS oldest_queued_seconds -FROM turns t LEFT JOIN session_devices b ON b.session_id = t.session_id +FROM turns t LEFT JOIN session_runtime_assignments b ON b.session_id = t.session_id WHERE t.status IN ('queued', 'in_progress') ` diff --git a/services/core/internal/db/sqlc/devices.sql.go b/services/core/internal/db/sqlc/devices.sql.go index fbe8cbd2c..da39d4fc0 100644 --- a/services/core/internal/db/sqlc/devices.sql.go +++ b/services/core/internal/db/sqlc/devices.sql.go @@ -11,16 +11,35 @@ import ( "github.com/jackc/pgx/v5/pgtype" ) +const acknowledgeAssignmentRelease = `-- name: AcknowledgeAssignmentRelease :execrows +UPDATE session_runtime_assignments SET applied_epoch = epoch +WHERE session_id = $1 AND assignment_id = $2 AND epoch = $3 AND desired_state = 'released' +` + +type AcknowledgeAssignmentReleaseParams struct { + SessionID pgtype.UUID `json:"session_id"` + AssignmentID pgtype.UUID `json:"assignment_id"` + Epoch int64 `json:"epoch"` +} + +func (q *Queries) AcknowledgeAssignmentRelease(ctx context.Context, arg AcknowledgeAssignmentReleaseParams) (int64, error) { + result, err := q.db.Exec(ctx, acknowledgeAssignmentRelease, arg.SessionID, arg.AssignmentID, arg.Epoch) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil +} + const bindSessionDevice = `-- name: BindSessionDevice :one -INSERT INTO session_devices (session_id, device_id) +INSERT INTO session_runtime_assignments (session_id, runtime_id) SELECT s.id, d.id FROM sessions s JOIN devices d ON d.tenant_id = s.tenant_id WHERE s.tenant_id = $1 AND s.id = $2 AND d.id = $3 AND d.revoked_at IS NULL AND (d.environment_id IS NULL OR EXISTS ( SELECT 1 FROM environments e WHERE e.id = d.environment_id AND e.session_id = s.id )) -ON CONFLICT (session_id) DO UPDATE SET device_id = session_devices.device_id -WHERE session_devices.device_id = EXCLUDED.device_id -RETURNING device_id +ON CONFLICT (session_id) DO UPDATE SET runtime_id = session_runtime_assignments.runtime_id +WHERE session_runtime_assignments.runtime_id = EXCLUDED.runtime_id AND session_runtime_assignments.desired_state = 'bound' +RETURNING runtime_id ` type BindSessionDeviceParams struct { @@ -31,9 +50,9 @@ type BindSessionDeviceParams struct { func (q *Queries) BindSessionDevice(ctx context.Context, arg BindSessionDeviceParams) (pgtype.UUID, error) { row := q.db.QueryRow(ctx, bindSessionDevice, arg.TenantID, arg.ID, arg.ID_2) - var device_id pgtype.UUID - err := row.Scan(&device_id) - return device_id, err + var runtime_id pgtype.UUID + err := row.Scan(&runtime_id) + return runtime_id, err } const createDevice = `-- name: CreateDevice :one @@ -110,10 +129,10 @@ func (q *Queries) GetDeviceCredential(ctx context.Context, id pgtype.UUID) (GetD } const getSessionDevice = `-- name: GetSessionDevice :one -SELECT d.id, d.name, d.environment_id FROM session_devices b +SELECT d.id, d.name, d.environment_id, b.assignment_id, b.epoch FROM session_runtime_assignments b JOIN sessions s ON s.id = b.session_id -JOIN devices d ON d.id = b.device_id AND d.tenant_id = s.tenant_id -WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL +JOIN devices d ON d.id = b.runtime_id AND d.tenant_id = s.tenant_id +WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL AND b.desired_state = 'bound' AND EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = d.id) AND (d.environment_id IS NULL OR EXISTS ( SELECT 1 FROM environments e WHERE e.id = d.environment_id AND e.session_id = s.id @@ -129,22 +148,30 @@ type GetSessionDeviceRow struct { ID pgtype.UUID `json:"id"` Name string `json:"name"` EnvironmentID pgtype.UUID `json:"environment_id"` + AssignmentID pgtype.UUID `json:"assignment_id"` + Epoch int64 `json:"epoch"` } func (q *Queries) GetSessionDevice(ctx context.Context, arg GetSessionDeviceParams) (GetSessionDeviceRow, error) { row := q.db.QueryRow(ctx, getSessionDevice, arg.TenantID, arg.ID) var i GetSessionDeviceRow - err := row.Scan(&i.ID, &i.Name, &i.EnvironmentID) + err := row.Scan( + &i.ID, + &i.Name, + &i.EnvironmentID, + &i.AssignmentID, + &i.Epoch, + ) return i, err } const getSessionExecutionBinding = `-- name: GetSessionExecutionBinding :one -SELECT d.id, d.name, b.native_session_id, d.environment_id, +SELECT d.id, d.name, b.native_session_id, d.environment_id, b.assignment_id, b.epoch, EXISTS (SELECT 1 FROM turns t WHERE t.session_id = s.id AND t.started_at IS NOT NULL) AS has_started_turn -FROM session_devices b +FROM session_runtime_assignments b JOIN sessions s ON s.id = b.session_id -JOIN devices d ON d.id = b.device_id AND d.tenant_id = s.tenant_id -WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL +JOIN devices d ON d.id = b.runtime_id AND d.tenant_id = s.tenant_id +WHERE s.tenant_id = $1 AND s.id = $2 AND d.revoked_at IS NULL AND b.desired_state = 'bound' AND EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = d.id) AND (d.environment_id IS NULL OR EXISTS ( SELECT 1 FROM environments e WHERE e.id = d.environment_id AND e.session_id = s.id @@ -161,6 +188,8 @@ type GetSessionExecutionBindingRow struct { Name string `json:"name"` NativeSessionID string `json:"native_session_id"` EnvironmentID pgtype.UUID `json:"environment_id"` + AssignmentID pgtype.UUID `json:"assignment_id"` + Epoch int64 `json:"epoch"` HasStartedTurn bool `json:"has_started_turn"` } @@ -172,13 +201,73 @@ func (q *Queries) GetSessionExecutionBinding(ctx context.Context, arg GetSession &i.Name, &i.NativeSessionID, &i.EnvironmentID, + &i.AssignmentID, + &i.Epoch, &i.HasStartedTurn, ) return i, err } +const listPendingAssignmentReleases = `-- name: ListPendingAssignmentReleases :many +SELECT session_id, runtime_id, assignment_id, epoch, remove_home FROM session_runtime_assignments +WHERE desired_state = 'released' AND applied_epoch < epoch AND runtime_id = ANY($1::uuid[]) +ORDER BY runtime_id, session_id +` + +type ListPendingAssignmentReleasesRow struct { + SessionID pgtype.UUID `json:"session_id"` + RuntimeID pgtype.UUID `json:"runtime_id"` + AssignmentID pgtype.UUID `json:"assignment_id"` + Epoch int64 `json:"epoch"` + RemoveHome bool `json:"remove_home"` +} + +func (q *Queries) ListPendingAssignmentReleases(ctx context.Context, runtimeIds []pgtype.UUID) ([]ListPendingAssignmentReleasesRow, error) { + rows, err := q.db.Query(ctx, listPendingAssignmentReleases, runtimeIds) + if err != nil { + return nil, err + } + defer rows.Close() + items := []ListPendingAssignmentReleasesRow{} + for rows.Next() { + var i ListPendingAssignmentReleasesRow + if err := rows.Scan( + &i.SessionID, + &i.RuntimeID, + &i.AssignmentID, + &i.Epoch, + &i.RemoveHome, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const releaseSessionAssignment = `-- name: ReleaseSessionAssignment :exec +UPDATE session_runtime_assignments +SET desired_state = 'released', epoch = epoch + 1, remove_home = remove_home OR $1::boolean +WHERE session_id = $2 AND (desired_state = 'bound' OR ($1::boolean AND NOT remove_home)) +` + +type ReleaseSessionAssignmentParams struct { + RemoveHome bool `json:"remove_home"` + SessionID pgtype.UUID `json:"session_id"` +} + +// An identical release keeps its epoch; a release that adds home removal +// advances it. +func (q *Queries) ReleaseSessionAssignment(ctx context.Context, arg ReleaseSessionAssignmentParams) error { + _, err := q.db.Exec(ctx, releaseSessionAssignment, arg.RemoveHome, arg.SessionID) + return err +} + const rememberNativeSession = `-- name: RememberNativeSession :execrows -UPDATE session_devices SET native_session_id = $2 WHERE session_id = $1 +UPDATE session_runtime_assignments SET native_session_id = $2 WHERE session_id = $1 ` type RememberNativeSessionParams struct { diff --git a/services/core/internal/db/sqlc/environment_initialization.sql.go b/services/core/internal/db/sqlc/environment_initialization.sql.go index 102d4837d..21c56093d 100644 --- a/services/core/internal/db/sqlc/environment_initialization.sql.go +++ b/services/core/internal/db/sqlc/environment_initialization.sql.go @@ -47,9 +47,9 @@ func (q *Queries) FailEnvironmentInitialization(ctx context.Context, id pgtype.U } const listEnvironmentInitializations = `-- name: ListEnvironmentInitializations :many -SELECT e.id, e.session_id, s.tenant_id, s.engine, e.initialization, b.device_id +SELECT e.id, e.session_id, s.tenant_id, s.engine, e.initialization, b.runtime_id, b.assignment_id, b.epoch FROM environments e JOIN sessions s ON s.id = e.session_id -LEFT JOIN session_devices b ON b.session_id = s.id +LEFT JOIN session_runtime_assignments b ON b.session_id = s.id AND b.desired_state = 'bound' WHERE e.id > $1 AND s.deleted_at IS NULL AND e.status NOT IN ('failed', 'expired') AND e.initialization IN ('pending', 'running') ORDER BY e.id LIMIT 32 @@ -61,7 +61,9 @@ type ListEnvironmentInitializationsRow struct { TenantID pgtype.UUID `json:"tenant_id"` Engine string `json:"engine"` Initialization string `json:"initialization"` - DeviceID pgtype.UUID `json:"device_id"` + RuntimeID pgtype.UUID `json:"runtime_id"` + AssignmentID pgtype.UUID `json:"assignment_id"` + Epoch pgtype.Int8 `json:"epoch"` } func (q *Queries) ListEnvironmentInitializations(ctx context.Context, id pgtype.UUID) ([]ListEnvironmentInitializationsRow, error) { @@ -79,7 +81,9 @@ func (q *Queries) ListEnvironmentInitializations(ctx context.Context, id pgtype. &i.TenantID, &i.Engine, &i.Initialization, - &i.DeviceID, + &i.RuntimeID, + &i.AssignmentID, + &i.Epoch, ); err != nil { return nil, err } diff --git a/services/core/internal/db/sqlc/models.go b/services/core/internal/db/sqlc/models.go index eca6abc6f..14083ca76 100644 --- a/services/core/internal/db/sqlc/models.go +++ b/services/core/internal/db/sqlc/models.go @@ -394,12 +394,6 @@ type SessionArtifact struct { CreatedAt pgtype.Timestamptz `json:"created_at"` } -type SessionDevice struct { - SessionID pgtype.UUID `json:"session_id"` - DeviceID pgtype.UUID `json:"device_id"` - NativeSessionID string `json:"native_session_id"` -} - type SessionEvent struct { SessionID pgtype.UUID `json:"session_id"` Sequence int64 `json:"sequence"` @@ -429,6 +423,17 @@ type SessionModelExecution struct { EncryptedConfig []byte `json:"encrypted_config"` } +type SessionRuntimeAssignment struct { + SessionID pgtype.UUID `json:"session_id"` + RuntimeID pgtype.UUID `json:"runtime_id"` + NativeSessionID string `json:"native_session_id"` + AssignmentID pgtype.UUID `json:"assignment_id"` + Epoch int64 `json:"epoch"` + DesiredState string `json:"desired_state"` + RemoveHome bool `json:"remove_home"` + AppliedEpoch int64 `json:"applied_epoch"` +} + type Skill struct { ID pgtype.UUID `json:"id"` TenantID pgtype.UUID `json:"tenant_id"` diff --git a/services/core/internal/db/sqlc/scheduling.sql.go b/services/core/internal/db/sqlc/scheduling.sql.go index 4a748867c..9cdc8213b 100644 --- a/services/core/internal/db/sqlc/scheduling.sql.go +++ b/services/core/internal/db/sqlc/scheduling.sql.go @@ -39,13 +39,13 @@ const listEnvironmentInputWork = `-- name: ListEnvironmentInputWork :many SELECT r.id, r.session_id, s.tenant_id FROM environment_input_reservations r JOIN sessions s ON s.id = r.session_id -LEFT JOIN session_devices b ON b.session_id = s.id +LEFT JOIN session_runtime_assignments b ON b.session_id = s.id WHERE r.state = 'pending' AND r.deadline > clock_timestamp() AND r.id > $1::uuid AND s.deleted_at IS NULL AND EXISTS ( SELECT 1 FROM devices d WHERE d.tenant_id = s.tenant_id AND d.revoked_at IS NULL AND d.id = ANY($2::uuid[]) - AND (b.device_id IS NULL OR d.id = b.device_id) + AND (b.runtime_id IS NULL OR d.id = b.runtime_id) ) ORDER BY r.id LIMIT 100 ` diff --git a/services/core/internal/db/sqlc/subagent_identities.sql.go b/services/core/internal/db/sqlc/subagent_identities.sql.go index 112e42424..7afb4ef92 100644 --- a/services/core/internal/db/sqlc/subagent_identities.sql.go +++ b/services/core/internal/db/sqlc/subagent_identities.sql.go @@ -73,9 +73,9 @@ INSERT INTO subagent_identities ( id, session_id, device_id, engine, native_id, parent_native_id, native_created_at, first_turn_id, first_event_ordinal ) -SELECT $1, s.id, b.device_id, s.engine, $2, +SELECT $1, s.id, b.runtime_id, s.engine, $2, $3, $4, $5, $6 -FROM sessions s JOIN session_devices b ON b.session_id = s.id +FROM sessions s JOIN session_runtime_assignments b ON b.session_id = s.id WHERE s.id = $7 AND (b.native_session_id = $3 OR EXISTS (SELECT 1 FROM subagent_identities p diff --git a/services/core/internal/deployment/storage.go b/services/core/internal/deployment/storage.go index 20cff0768..f00548625 100644 --- a/services/core/internal/deployment/storage.go +++ b/services/core/internal/deployment/storage.go @@ -132,7 +132,8 @@ type AllocationTx interface { type AllocationCleanupTx interface { AllocationTx sessions.EnvironmentTerminationTx - // RevokeDevice revokes the allocation's device. + // RevokeDevice releases the Session's Runtime assignment without home + // removal, then revokes the allocation's device. RevokeDevice(current Allocation) error // RequestCleanup records that the allocation's resources await cleanup. RequestCleanup(current Allocation) (Allocation, error) @@ -163,8 +164,9 @@ type SessionArchiveTx interface { // FindAllocation returns the Environment's allocation and whether it has // one. FindAllocation(environment string) (Allocation, bool, error) - // RequestArchiveCleanup revokes the allocation's device and records that - // its resources await cleanup. The device's first revocation records the + // RequestArchiveCleanup releases the Session's Runtime assignment without + // home removal, revokes the allocation's device and records that its + // resources await cleanup. The device's first revocation records the // Session's active Turn whose cancellation was requested, which the // archived cancellation receipt reports. RequestArchiveCleanup(current Allocation) error diff --git a/services/core/internal/execution/artifacts.go b/services/core/internal/execution/artifacts.go index f1aa98504..4c5148a6d 100644 --- a/services/core/internal/execution/artifacts.go +++ b/services/core/internal/execution/artifacts.go @@ -20,7 +20,7 @@ func (d *Dispatcher) captureCompletedArtifacts(ctx context.Context, peer *runtim err := d.sessionExecution.BeginTurnArtifactCapture(owner, session.TenantID, session.ID, turnID, result.AppliedThrough) if err == nil { err = d.withPreparedWorkspace(owner, peer, session, environment, bound, func(ctx context.Context, handle string) error { - return peer.ExportWorkspaceOutputs(ctx, proto.WorkspaceExportPayload{Handle: handle, EnvironmentID: environment.ID}, func(body io.Reader) error { + return peer.ExportWorkspaceOutputs(ctx, bound.Assignment, proto.WorkspaceExportPayload{Handle: handle, EnvironmentID: environment.ID}, func(body io.Reader) error { return d.Sessions.StageTurnArtifacts(ctx, sessions.StageTurnArtifactsCommand{TenantID: session.TenantID, SessionID: session.ID, TurnID: turnID, EnvironmentID: environment.ID, Export: body}) }) }) diff --git a/services/core/internal/execution/assignment_releases.go b/services/core/internal/execution/assignment_releases.go new file mode 100644 index 000000000..cfce17413 --- /dev/null +++ b/services/core/internal/execution/assignment_releases.go @@ -0,0 +1,80 @@ +package execution + +import ( + "context" + "sync" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/internal/obs/log" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// runAssignmentReleases delivers each released assignment to its connected +// Runtime until the Runtime acknowledges it, so a Runtime that reconnects +// receives the releases it missed. Releases run one per Session, bounded like +// executions. +func (w *Worker) runAssignmentReleases(ctx context.Context) error { + active := make(map[string]bool) + done := make(chan string, w.executionConcurrency()) + var running sync.WaitGroup + ctx, stop := context.WithCancel(ctx) + defer func() { stop(); running.Wait() }() + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return ctx.Err() + case session := <-done: + delete(active, session) + continue + case <-ticker.C: + } + releases, err := w.dispatcher.sessionExecution.ListAssignmentReleases(ctx, w.dispatcher.Registry.Devices()) + if err != nil { + return err + } + for _, release := range releases { + session := release.Assignment.SessionID + if active[session] || len(active) >= w.executionConcurrency() { + continue + } + active[session] = true + running.Add(1) + go func() { + defer running.Done() + w.releaseAssignment(ctx, release) + done <- session + }() + } + } +} + +// releaseAssignment sends one release and records its acknowledgement. Home +// removal is requested only from a Runtime that declares it. +func (w *Worker) releaseAssignment(ctx context.Context, release sessions.AssignmentRelease) { + ctx, cancel := context.WithTimeout(ctx, 2*time.Minute) + defer cancel() + peer, err := w.dispatcher.authorizedPeer(ctx, release.RuntimeID) + if err != nil { + return + } + supported, known := peer.RemovesHomes() + if !known { + return + } + removeHome := release.RemoveHome && supported + want := proto.AssignmentReleased + if removeHome { + want = proto.AssignmentHomeRemoved + } + status, err := peer.Release(ctx, release.Assignment, removeHome) + if err != nil || status.State != want { + log.Warn(ctx, "Runtime assignment release unconfirmed", "session_id", release.Assignment.SessionID, "runtime_id", release.RuntimeID, "error_code", status.ErrorCode) + return + } + if err := w.dispatcher.sessionExecution.AcknowledgeAssignmentRelease(ctx, release.Assignment); err != nil { + log.Warn(ctx, "Runtime assignment release not recorded", "session_id", release.Assignment.SessionID, "runtime_id", release.RuntimeID) + } +} diff --git a/services/core/internal/execution/delivery.go b/services/core/internal/execution/delivery.go index 3cd1e1833..9eda34790 100644 --- a/services/core/internal/execution/delivery.go +++ b/services/core/internal/execution/delivery.go @@ -25,29 +25,32 @@ type cancellationResult struct { err error } -func requestCancellation(ctx context.Context, peer *runtimegateway.Session, runID string) <-chan cancellationResult { +func requestCancellation(ctx context.Context, peer *runtimegateway.Session, ref proto.AssignmentRef, runID string) <-chan cancellationResult { out := make(chan cancellationResult, 1) go func() { id := "cancel:" + runID env, _ := proto.NewEnvelope(proto.TypePromptCancel, runID, proto.PromptCancelPayload{DeliveryID: id}) + env.Assignment = ref ack, err := peer.SendAndWaitInteractionAck(ctx, env, id) out <- cancellationResult{ack: ack, err: err} }() return out } -func send(ctx context.Context, peer *runtimegateway.Session, kind, runID string, payload any) error { - env, err := proto.NewEnvelope(kind, runID, payload) +// send sends one frame of the assignment's Session. +func send(ctx context.Context, peer *runtimegateway.Session, ref proto.AssignmentRef, kind, id string, payload any) error { + env, err := proto.NewEnvelope(kind, id, payload) if err != nil { return err } + env.Assignment = ref ctx, cancel := context.WithTimeout(ctx, 3*time.Second) defer cancel() return peer.Send(ctx, env) } -func abort(peer *runtimegateway.Session, runID string) { - _ = send(context.Background(), peer, proto.TypePromptCancel, runID, proto.PromptCancelPayload{}) +func abort(peer *runtimegateway.Session, ref proto.AssignmentRef, runID string) { + _ = send(context.Background(), peer, ref, proto.TypePromptCancel, runID, proto.PromptCancelPayload{}) } func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, peer *runtimegateway.Session, request proto.PromptRequestPayload, first int64, prepared *preparedStart) (result Result, status string) { @@ -64,7 +67,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe defer peer.Unsubscribe(request.RunID) defer func() { if status == sessions.TurnFailed { - abort(peer, request.RunID) + abort(peer, request.Assignment, request.RunID) } }() journal := &journal{writer: d.sessionExecution, tenant: tenantID, session: sessionID, turn: request.RunID, next: 1, @@ -86,13 +89,8 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe var executorRetried, nativeObserved bool inputStarted := time.Now() firstTextObserved := false - if prepared != nil { - preparationEvents = prepared.sub.Events - err = prepared.start(ctx, request) - } else { - err = send(ctx, peer, proto.TypePromptRequest, request.RunID, request) - } - if err != nil { + preparationEvents = prepared.sub.Events + if err = prepared.start(ctx, request); err != nil { result.ErrorCode = "delivery_unknown" return } @@ -105,7 +103,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe var pending *pendingInput var cancelSent time.Time var cancelReply <-chan cancellationResult - functions := &functionExchange{kind: request.AgentKind, turns: d.SessionsReader, sessions: d.sessionExecution, tenant: tenantID, session: sessionID, turn: request.RunID, tools: request.FunctionTools} + functions := &functionExchange{assignment: request.Assignment, kind: request.AgentKind, turns: d.SessionsReader, sessions: d.sessionExecution, tenant: tenantID, session: sessionID, turn: request.RunID, tools: request.FunctionTools} done := false cancelCtx, stopCancellation := context.WithCancel(ctx) defer stopCancellation() @@ -291,7 +289,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe return } if !turn.CancelRequestedAt.IsZero() { - cancelReply = requestCancellation(cancelCtx, peer, request.RunID) + cancelReply = requestCancellation(cancelCtx, peer, request.Assignment, request.RunID) cancelSent = time.Now() continue } @@ -338,7 +336,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe result.ErrorCode = "message_input_unsupported" return } - if send(ctx, peer, proto.TypePromptSteer, request.RunID, proto.PromptSteerPayload{InputID: strconv.FormatInt(pending.sequence, 10), Input: pending.input, DurableReceipt: true}) != nil { + if send(ctx, peer, request.Assignment, proto.TypePromptSteer, request.RunID, proto.PromptSteerPayload{InputID: strconv.FormatInt(pending.sequence, 10), Input: pending.input, DurableReceipt: true}) != nil { result.ErrorCode = "input_outcome_unknown" return } diff --git a/services/core/internal/execution/device_authority.go b/services/core/internal/execution/device_authority.go index fb13e8283..5600e383c 100644 --- a/services/core/internal/execution/device_authority.go +++ b/services/core/internal/execution/device_authority.go @@ -35,3 +35,20 @@ func authorizedRuntimePeer(ctx context.Context, devices sessions.DeviceReader, r func (d *Dispatcher) authorizedPeer(ctx context.Context, id string) (*runtimegateway.Session, error) { return authorizedRuntimePeer(ctx, d.SessionsReader, d.Registry, id) } + +// assignedRuntimePeer returns the authorized peer of the Session's bound +// Runtime after the Runtime acknowledged the Session's assignment. +func assignedRuntimePeer(ctx context.Context, devices sessions.DeviceReader, registry *runtimegateway.Registry, bound sessions.ExecutionDevice) (*runtimegateway.Session, error) { + peer, err := authorizedRuntimePeer(ctx, devices, registry, bound.ID) + if err != nil { + return nil, err + } + if err := peer.Bind(ctx, bound.Assignment, bound.EnvironmentID); err != nil { + return nil, err + } + return peer, nil +} + +func (d *Dispatcher) assignedPeer(ctx context.Context, bound sessions.ExecutionDevice) (*runtimegateway.Session, error) { + return assignedRuntimePeer(ctx, d.SessionsReader, d.Registry, bound) +} diff --git a/services/core/internal/execution/directory_preparation.go b/services/core/internal/execution/directory_preparation.go index 427e400cb..586b208e8 100644 --- a/services/core/internal/execution/directory_preparation.go +++ b/services/core/internal/execution/directory_preparation.go @@ -15,7 +15,7 @@ func (d *Dispatcher) readPreparedDirectory(ctx context.Context, peer *runtimegat var result directoryReadResult err := d.withPreparedWorkspace(owner, peer, session, environment, bound, func(ctx context.Context, handle string) error { read.Handle = handle - result = readEnvironmentDirectory(ctx, peer, read) + result = readEnvironmentDirectory(ctx, peer, bound.Assignment, read) return result.err }) if err != nil { @@ -25,11 +25,11 @@ func (d *Dispatcher) readPreparedDirectory(ctx context.Context, peer *runtimegat } func (d *Dispatcher) withPreparedWorkspace(owner context.Context, peer *runtimegateway.Session, session sessions.Session, environment sessions.Environment, bound sessions.ExecutionDevice, consume func(context.Context, string) error) error { - req := proto.PromptRequestPayload{AgentKind: session.Engine, AgentStateKey: "agents-api-" + session.ID, StrictResume: true, ReleaseOnCompletion: true, WorkspaceReadOnly: true} + req := proto.PromptRequestPayload{Assignment: bound.Assignment, AgentKind: session.Engine, AgentStateKey: "agents-api-" + session.ID, StrictResume: true, ReleaseOnCompletion: true, WorkspaceReadOnly: true} if err := d.configurePreparedEnvironment(session, environment, bound, &req); err != nil { return ErrExecutionUnavailable } - prepared, err := newPreparedStart(peer) + prepared, err := newPreparedStart(peer, bound.Assignment) if err != nil { return ErrExecutionUnavailable } @@ -42,7 +42,7 @@ func (d *Dispatcher) withPreparedWorkspace(owner context.Context, peer *runtimeg } }() prepare, stop := context.WithTimeout(owner, 10*time.Second) - err = send(prepare, peer, proto.TypeExecutionPrepare, prepared.requestID, proto.ExecutionPreparePayload{SessionID: session.ID, Configuration: req}) + err = send(prepare, peer, bound.Assignment, proto.TypeExecutionPrepare, prepared.requestID, proto.ExecutionPreparePayload{SessionID: session.ID, Configuration: req}) if err == nil { err = prepared.awaitDirectoryReady(prepare) } @@ -84,7 +84,7 @@ func (p *preparedStart) awaitDirectoryReady(ctx context.Context) error { func (p *preparedStart) releaseDirectory() error { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() - if send(ctx, p.peer, proto.TypeExecutionRelease, p.requestID, proto.ExecutionReleasePayload{Handle: p.handle}) != nil { + if send(ctx, p.peer, p.assignment, proto.TypeExecutionRelease, p.requestID, proto.ExecutionReleasePayload{Handle: p.handle}) != nil { return ErrExecutionUnavailable } for { diff --git a/services/core/internal/execution/dispatcher.go b/services/core/internal/execution/dispatcher.go index 7a1545402..8eaa489b2 100644 --- a/services/core/internal/execution/dispatcher.go +++ b/services/core/internal/execution/dispatcher.go @@ -107,7 +107,11 @@ func (d *Dispatcher) Run(ctx context.Context, tenantID, sessionID, turnID string if err != nil { return sessions.Turn{}, err } + req.Assignment = bound.Device.Assignment req.DisableExecutionEnvironment = true + if err := peer.Bind(ctx, req.Assignment, bound.Device.EnvironmentID); err != nil { + return sessions.Turn{}, err + } prepared, err := d.prepareTurnExecutor(ctx, peer, tenantID, sessionID, turnID, req, sessions.TurnQueued) if err != nil { return sessions.Turn{}, err diff --git a/services/core/internal/execution/environment_directory.go b/services/core/internal/execution/environment_directory.go index 5d1286efb..a189005dc 100644 --- a/services/core/internal/execution/environment_directory.go +++ b/services/core/internal/execution/environment_directory.go @@ -109,14 +109,14 @@ func (w *Worker) runDirectoryRead(owner context.Context, request directoryReadRe if err != nil || !environmentDeviceMatches(session, environment, bound) || !w.directoryDeviceReady(check, bound.ID, session.Engine, placement, prepare) { return } - peer, err := w.dispatcher.authorizedPeer(check, bound.ID) + peer, err := w.dispatcher.assignedPeer(check, bound) if err != nil { return } read := proto.WorkspaceReadPayload{EnvironmentID: environment.ID, Path: request.path, MaxEntries: proto.WorkspaceDirectoryMaxEntries} if !prepare { read.RunID = run - result = readEnvironmentDirectory(owner, peer, read) + result = readEnvironmentDirectory(owner, peer, bound.Assignment, read) return } result = w.dispatcher.readPreparedDirectory(owner, peer, session, environment, bound, read) @@ -136,8 +136,8 @@ func (w *Worker) directoryDeviceReady(ctx context.Context, id, engine string, pl return known && found && info.Available && placementReady && (!prepare || (info.Capabilities.Preparation && info.Capabilities.WorkspaceReadPreparation)) } -func readEnvironmentDirectory(ctx context.Context, peer *runtimegateway.Session, request proto.WorkspaceReadPayload) directoryReadResult { - result, err := peer.ListWorkspaceDirectory(ctx, request) +func readEnvironmentDirectory(ctx context.Context, peer *runtimegateway.Session, ref proto.AssignmentRef, request proto.WorkspaceReadPayload) directoryReadResult { + result, err := peer.ListWorkspaceDirectory(ctx, ref, request) if err == nil && result.Outcome == "completed" && result.Directory != nil && !result.Directory.Truncated && proto.ValidWorkspaceDirectory(result.Directory, request.MaxEntries) { return directoryReadResult{directory: *result.Directory} } diff --git a/services/core/internal/execution/environment_file_write.go b/services/core/internal/execution/environment_file_write.go index d608265fd..6967eff8a 100644 --- a/services/core/internal/execution/environment_file_write.go +++ b/services/core/internal/execution/environment_file_write.go @@ -93,7 +93,7 @@ func (w *Worker) runFileWrite(owner context.Context, request fileWriteRequest) f if err != nil || !environmentDeviceMatches(session, environment, bound) || w.dispatcher.Registry == nil { return unavailable } - peer, err := w.dispatcher.authorizedPeer(ctx, bound.ID) + peer, err := w.dispatcher.assignedPeer(ctx, bound) if err != nil { return unavailable } @@ -116,7 +116,7 @@ func (w *Worker) runFileWrite(owner context.Context, request fileWriteRequest) f if intent.Replayed { return unavailable } - result, err := peer.WriteWorkspaceFile(ctx, key.ID, proto.WorkspaceWritePayload{EnvironmentID: environment.ID, SessionID: session.ID, Path: request.path}, request.data) + result, err := peer.WriteWorkspaceFile(ctx, key.ID, bound.Assignment, proto.WorkspaceWritePayload{EnvironmentID: environment.ID, SessionID: session.ID, Path: request.path}, request.data) if err != nil || (result.Outcome != "completed" && result.Outcome != "rejected") { return unavailable } diff --git a/services/core/internal/execution/executor_preparation.go b/services/core/internal/execution/executor_preparation.go index 11fafdf9d..efce9c708 100644 --- a/services/core/internal/execution/executor_preparation.go +++ b/services/core/internal/execution/executor_preparation.go @@ -14,12 +14,12 @@ import ( // Core does not cache native ownership. A replacement requires the Runtime to // confirm cleanup and recover the exact Session history before returning ready. func (d *Dispatcher) prepareTurnExecutor(ctx context.Context, peer *runtimegateway.Session, tenant, session, turn string, request proto.PromptRequestPayload, expectedStatus string) (*preparedStart, error) { - prepared, err := newPreparedStart(peer) + prepared, err := newPreparedStart(peer, request.Assignment) if err != nil { return nil, err } request.RunID, request.ConversationID, request.Input = "", "", nil - if err = send(ctx, peer, proto.TypeExecutionPrepare, prepared.requestID, proto.ExecutionPreparePayload{SessionID: session, Configuration: request}); err == nil { + if err = send(ctx, peer, request.Assignment, proto.TypeExecutionPrepare, prepared.requestID, proto.ExecutionPreparePayload{SessionID: session, Configuration: request}); err == nil { err = d.awaitTurnExecutor(ctx, tenant, session, turn, expectedStatus, prepared) } if err != nil { diff --git a/services/core/internal/execution/functions.go b/services/core/internal/execution/functions.go index f3804de11..b784c47b0 100644 --- a/services/core/internal/execution/functions.go +++ b/services/core/internal/execution/functions.go @@ -52,6 +52,7 @@ type functionExchange struct { turns sessions.TurnReader sessions *sessions.ExecutionOperations tenant, session, turn string + assignment proto.AssignmentRef kind string tools []proto.FunctionTool callID string @@ -99,6 +100,7 @@ func (f *functionExchange) start(ctx context.Context, peer *runtimegateway.Sessi if err != nil { return err } + env.Assignment = f.assignment replies := make(chan functionReply, 1) f.callID, f.reply = call.CallID, replies go func() { diff --git a/services/core/internal/execution/preparation.go b/services/core/internal/execution/preparation.go index ca815d02e..29b22cfa2 100644 --- a/services/core/internal/execution/preparation.go +++ b/services/core/internal/execution/preparation.go @@ -26,24 +26,25 @@ type preparedStart struct { startObserved bool readyObserved bool peer *runtimegateway.Session + assignment proto.AssignmentRef requestID string handle string executorID string sub *runtimegateway.Subscription } -func newPreparedStart(peer *runtimegateway.Session) (*preparedStart, error) { +func newPreparedStart(peer *runtimegateway.Session, ref proto.AssignmentRef) (*preparedStart, error) { id := uuid.NewString() sub, err := peer.SubscribePreparation(id) if err != nil { return nil, err } - return &preparedStart{peer: peer, requestID: id, sub: sub, createdAt: time.Now()}, nil + return &preparedStart{peer: peer, assignment: ref, requestID: id, sub: sub, createdAt: time.Now()}, nil } func (p *preparedStart) close() { if p.handle != "" { - _ = send(context.Background(), p.peer, proto.TypeExecutionRelease, p.requestID, proto.ExecutionReleasePayload{Handle: p.handle}) + _ = send(context.Background(), p.peer, p.assignment, proto.TypeExecutionRelease, p.requestID, proto.ExecutionReleasePayload{Handle: p.handle}) } p.peer.UnsubscribePreparation(p.requestID) } @@ -130,7 +131,7 @@ func (d *Dispatcher) awaitPreparation(ctx context.Context, tenant, session strin func (p *preparedStart) start(ctx context.Context, request proto.PromptRequestPayload) error { p.startSentAt = time.Now() - return send(ctx, p.peer, proto.TypeExecutionStart, p.requestID, proto.ExecutionStartPayload{Handle: p.handle, ExecutorID: p.executorID, RunID: request.RunID, Input: request.Input}) + return send(ctx, p.peer, p.assignment, proto.TypeExecutionStart, p.requestID, proto.ExecutionStartPayload{Handle: p.handle, ExecutorID: p.executorID, RunID: request.RunID, Input: request.Input}) } func (p *preparedStart) started(env proto.Envelope, runID string) (bool, error) { diff --git a/services/core/internal/execution/prepared_dispatch.go b/services/core/internal/execution/prepared_dispatch.go index 5bbce163b..0cd73dd90 100644 --- a/services/core/internal/execution/prepared_dispatch.go +++ b/services/core/internal/execution/prepared_dispatch.go @@ -59,6 +59,7 @@ func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, t if err != nil { return run, err } + req.Assignment = bound.Device.Assignment var messages proto.MessageInput for _, input := range run.Reservation.Inputs { if input.Kind != "message" { @@ -76,12 +77,15 @@ func (d *Dispatcher) RunEnvironmentInput(ctx context.Context, lease Ownership, t if err := d.configurePreparedEnvironment(session, environment, bound.Device, &req); err != nil { return run, err } - prepared, err := newPreparedStart(peer) + if err := peer.Bind(owner, req.Assignment, bound.Device.EnvironmentID); err != nil { + return run, err + } + prepared, err := newPreparedStart(peer, req.Assignment) if err != nil { return run, err } defer prepared.close() - if err = send(owner, peer, proto.TypeExecutionPrepare, prepared.requestID, proto.ExecutionPreparePayload{SessionID: sessionID, Configuration: req}); err != nil { + if err = send(owner, peer, req.Assignment, proto.TypeExecutionPrepare, prepared.requestID, proto.ExecutionPreparePayload{SessionID: sessionID, Configuration: req}); err != nil { return run, err } run.Reservation, err = d.awaitPreparation(owner, tenantID, sessionID, run.Reservation, prepared) diff --git a/services/core/internal/execution/runtime_capabilities_test.go b/services/core/internal/execution/runtime_capabilities_test.go index c2c29bca1..7d29fcbf8 100644 --- a/services/core/internal/execution/runtime_capabilities_test.go +++ b/services/core/internal/execution/runtime_capabilities_test.go @@ -21,7 +21,7 @@ type capabilityFixture struct { err error } -func (f *capabilityFixture) PrepareRuntime(_ context.Context, id string, request proto.RuntimePreparePayload, body []byte) (proto.RuntimePrepareResultPayload, error) { +func (f *capabilityFixture) PrepareRuntime(_ context.Context, id string, _ proto.AssignmentRef, request proto.RuntimePreparePayload, body []byte) (proto.RuntimePrepareResultPayload, error) { f.requestID, f.request, f.body = id, request, append([]byte(nil), body...) return proto.RuntimePrepareResultPayload{Outcome: f.outcome, ErrorCode: "safe_failure"}, f.err } @@ -40,7 +40,7 @@ func TestRuntimeCapabilitiesPreserveRawBundlesAndSetupOrdering(t *testing.T) { owner := agentcapabilities.Identity{EnvironmentID: uuid.NewString(), SessionID: uuid.NewString()} for _, op := range operations { peer := &capabilityFixture{outcome: "completed"} - if err := runRuntimeSetup(t.Context(), peer, owner, op); err != nil { + if err := runRuntimeSetup(t.Context(), peer, proto.AssignmentRef{}, owner, op); err != nil { t.Fatal(err) } if _, err := uuid.Parse(peer.requestID); err != nil { @@ -64,14 +64,14 @@ func TestRuntimeCapabilitiesPreserveRawBundlesAndSetupOrdering(t *testing.T) { func TestRuntimeCapabilitiesConfirmedAndUnknownFailures(t *testing.T) { for _, outcome := range []string{"failed", "rejected", "unknown", "unexpected"} { peer := &capabilityFixture{outcome: outcome} - err := runRuntimeSetup(t.Context(), peer, agentcapabilities.Identity{}, runtimeSetupOperation{Request: proto.RuntimePreparePayload{Action: "finalize"}}) + err := runRuntimeSetup(t.Context(), peer, proto.AssignmentRef{}, agentcapabilities.Identity{}, runtimeSetupOperation{Request: proto.RuntimePreparePayload{Action: "finalize"}}) var confirmed *runtimeStepFailure if err == nil || errors.As(err, &confirmed) != (outcome == "failed" || outcome == "rejected") { t.Fatal(outcome, err) } } peer := &capabilityFixture{outcome: "completed", err: errors.New(setupCanary)} - err := runRuntimeSetup(t.Context(), peer, agentcapabilities.Identity{}, runtimeSetupOperation{Request: proto.RuntimePreparePayload{}}) + err := runRuntimeSetup(t.Context(), peer, proto.AssignmentRef{}, agentcapabilities.Identity{}, runtimeSetupOperation{Request: proto.RuntimePreparePayload{}}) if err == nil || bytes.Contains([]byte(err.Error()), []byte(setupCanary)) { t.Fatal("transport error leaked or succeeded", err) } diff --git a/services/core/internal/execution/runtime_compute.go b/services/core/internal/execution/runtime_compute.go index 7d08a1b33..9c328c01c 100644 --- a/services/core/internal/execution/runtime_compute.go +++ b/services/core/internal/execution/runtime_compute.go @@ -9,6 +9,7 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/deployment" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sandbox" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" "github.com/google/uuid" ) @@ -122,6 +123,17 @@ func (r *runtimeLifecycle) observeCompute(ctx context.Context, owner deployment. } } +// allocationAssignment reads the assignment of the Session whose Environment +// the allocation runs; it is ErrNotFound unless the allocation's Runtime holds +// it. +func (r *runtimeLifecycle) allocationAssignment(ctx context.Context, owner deployment.Allocation) (sessions.ExecutionDevice, error) { + bound, err := r.sessions.GetSessionRuntimeDevice(ctx, owner.TenantID, owner.SessionID) + if err == nil && bound.ID != owner.DeviceID { + err = sessions.ErrNotFound + } + return bound, err +} + func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.CheckpointProvider, owner deployment.Allocation, state runtimeCompute) error { compute, err := p.GetCompute(ctx, runtimeReference(owner), state.Current) if err != nil { @@ -130,7 +142,11 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint if compute.Status != "running" || !compute.BootstrapComplete { return sandbox.ErrComputeUnconfirmed } - peer, err := authorizedRuntimePeer(ctx, r.sessions, r.registry, owner.DeviceID) + bound, err := r.allocationAssignment(ctx, owner) + if err != nil { + return err + } + peer, err := assignedRuntimePeer(ctx, r.sessions, r.registry, bound) if err != nil { return err } @@ -155,7 +171,7 @@ func (r *runtimeLifecycle) idleCompute(ctx context.Context, p sandbox.Checkpoint if err != nil { return err } - result, err := peer.SuspendControl(ctx, proto.TypeEnvironmentQuiesce, proto.EnvironmentSuspendPayload{EnvironmentID: owner.EnvironmentID, SuspendID: state.SuspendID}) + result, err := peer.SuspendControl(ctx, proto.TypeEnvironmentQuiesce, bound.Assignment, proto.EnvironmentSuspendPayload{EnvironmentID: owner.EnvironmentID, SuspendID: state.SuspendID}) if err != nil { return err } diff --git a/services/core/internal/execution/runtime_compute_wake.go b/services/core/internal/execution/runtime_compute_wake.go index 552ede1ed..bb53c71bd 100644 --- a/services/core/internal/execution/runtime_compute_wake.go +++ b/services/core/internal/execution/runtime_compute_wake.go @@ -46,7 +46,13 @@ func (r *runtimeLifecycle) wakeCompute(ctx context.Context, p sandbox.Checkpoint } } } - result, err := peer.SuspendControl(ctx, proto.TypeEnvironmentResume, proto.EnvironmentSuspendPayload{EnvironmentID: owner.EnvironmentID, SuspendID: state.SuspendID, Rollback: state.Rollback}) + // Resume carries the reference that quiesced the Runtime; a suspended + // Runtime admits nothing else, so it is not bound again. + bound, err := r.allocationAssignment(ctx, owner) + if err != nil { + return err + } + result, err := peer.SuspendControl(ctx, proto.TypeEnvironmentResume, bound.Assignment, proto.EnvironmentSuspendPayload{EnvironmentID: owner.EnvironmentID, SuspendID: state.SuspendID, Rollback: state.Rollback}) if err != nil { return err } diff --git a/services/core/internal/execution/runtime_initialization.go b/services/core/internal/execution/runtime_initialization.go index 4f94d4d52..85dc81fbb 100644 --- a/services/core/internal/execution/runtime_initialization.go +++ b/services/core/internal/execution/runtime_initialization.go @@ -122,7 +122,7 @@ func (w *Worker) prepareEnvironment(ctx context.Context, owner sessions.Environm if err != nil { return err } - peer, err := w.dispatcher.authorizedPeer(ctx, owner.DeviceID) + peer, err := w.dispatcher.assignedPeer(ctx, sessions.ExecutionDevice{ID: owner.DeviceID, EnvironmentID: owner.EnvironmentID, Assignment: owner.Assignment}) if err != nil { return err } @@ -144,12 +144,12 @@ func (w *Worker) prepareEnvironment(ctx context.Context, owner sessions.Environm var body []byte metadata, body, err = w.dispatcher.SessionsReader.ReadInitialEnvironmentFile(step, owner.TenantID, owner.SessionID, index) if err == nil { - err = installInitialFile(step, peer, identity, metadata, body) + err = installInitialFile(step, peer, owner.Assignment, identity, metadata, body) } } else if err == nil { command := operations[index-len(cfg.Files)] candidate = command.provisioningFailure(0) - err = runRuntimeSetup(step, peer, identity, command) + err = runRuntimeSetup(step, peer, owner.Assignment, identity, command) } stop() if err != nil { diff --git a/services/core/internal/execution/runtime_setup.go b/services/core/internal/execution/runtime_setup.go index ccfcb795f..f02994f70 100644 --- a/services/core/internal/execution/runtime_setup.go +++ b/services/core/internal/execution/runtime_setup.go @@ -78,16 +78,16 @@ func setupOperations(setup environmentconfig.Setup) []runtimeSetupOperation { } type runtimePreparer interface { - PrepareRuntime(context.Context, string, proto.RuntimePreparePayload, []byte) (proto.RuntimePrepareResultPayload, error) + PrepareRuntime(context.Context, string, proto.AssignmentRef, proto.RuntimePreparePayload, []byte) (proto.RuntimePrepareResultPayload, error) } -func runRuntimeSetup(ctx context.Context, peer runtimePreparer, identity agentcapabilities.Identity, operation runtimeSetupOperation) error { +func runRuntimeSetup(ctx context.Context, peer runtimePreparer, ref proto.AssignmentRef, identity agentcapabilities.Identity, operation runtimeSetupOperation) error { if peer == nil { return errors.New("environment initialization request unavailable") } request := operation.Request request.EnvironmentID, request.SessionID = identity.EnvironmentID, identity.SessionID - result, err := peer.PrepareRuntime(ctx, uuid.NewString(), request, operation.Data) + result, err := peer.PrepareRuntime(ctx, uuid.NewString(), ref, request, operation.Data) if err == nil && result.Outcome == "completed" { return nil } @@ -97,9 +97,9 @@ func runRuntimeSetup(ctx context.Context, peer runtimePreparer, identity agentca return errors.New("environment initialization operation unconfirmed") } -func installInitialFile(ctx context.Context, peer runtimePreparer, identity agentcapabilities.Identity, file environmentconfig.InitialFileMetadata, body []byte) error { +func installInitialFile(ctx context.Context, peer runtimePreparer, ref proto.AssignmentRef, identity agentcapabilities.Identity, file environmentconfig.InitialFileMetadata, body []byte) error { if file.SizeBytes == nil || *file.SizeBytes != int64(len(body)) || len(body) > environmentconfig.MaxInitialFileBytes { return errors.New("environment initialization request unavailable") } - return runRuntimeSetup(ctx, peer, identity, runtimeSetupOperation{Request: proto.RuntimePreparePayload{Action: "file", File: &proto.RuntimeInitialFile{Path: file.Path}}, Data: body}) + return runRuntimeSetup(ctx, peer, ref, identity, runtimeSetupOperation{Request: proto.RuntimePreparePayload{Action: "file", File: &proto.RuntimeInitialFile{Path: file.Path}}, Data: body}) } diff --git a/services/core/internal/execution/runtime_setup_test.go b/services/core/internal/execution/runtime_setup_test.go index be2a99461..21431179b 100644 --- a/services/core/internal/execution/runtime_setup_test.go +++ b/services/core/internal/execution/runtime_setup_test.go @@ -21,7 +21,7 @@ type receiptRuntime struct { data []byte } -func (p *receiptRuntime) PrepareRuntime(_ context.Context, _ string, request proto.RuntimePreparePayload, data []byte) (proto.RuntimePrepareResultPayload, error) { +func (p *receiptRuntime) PrepareRuntime(_ context.Context, _ string, _ proto.AssignmentRef, request proto.RuntimePreparePayload, data []byte) (proto.RuntimePrepareResultPayload, error) { p.request, p.data = request, data return p.result, p.err } @@ -36,7 +36,7 @@ func TestRuntimeSetupReceiptOutcomes(t *testing.T) { } { t.Run(test.name, func(t *testing.T) { peer := &receiptRuntime{result: proto.RuntimePrepareResultPayload{Outcome: test.outcome, ExitCode: test.code}, err: test.err} - err := runRuntimeSetup(t.Context(), peer, agentcapabilities.Identity{}, runtimeSetupOperation{Request: proto.RuntimePreparePayload{Action: "initialize", Initialization: &proto.RuntimeInitialization{Action: "setup", Command: setupCanary}}}) + err := runRuntimeSetup(t.Context(), peer, proto.AssignmentRef{}, agentcapabilities.Identity{}, runtimeSetupOperation{Request: proto.RuntimePreparePayload{Action: "initialize", Initialization: &proto.RuntimeInitialization{Action: "setup", Command: setupCanary}}}) if test.outcome == "completed" && test.err == nil { if err != nil { t.Fatal(err) @@ -73,14 +73,14 @@ func TestInitialFileUsesTypedRuntimeBytes(t *testing.T) { size := int64(len(body)) owner := agentcapabilities.Identity{EnvironmentID: "environment", SessionID: "session"} peer := &receiptRuntime{result: proto.RuntimePrepareResultPayload{Outcome: "completed"}} - if err := installInitialFile(t.Context(), peer, owner, environmentconfig.InitialFileMetadata{Path: "/workspace/a", SizeBytes: &size}, body); err != nil { + if err := installInitialFile(t.Context(), peer, proto.AssignmentRef{}, owner, environmentconfig.InitialFileMetadata{Path: "/workspace/a", SizeBytes: &size}, body); err != nil { t.Fatal(err) } if peer.request.Action != "file" || peer.request.File.Path != "/workspace/a" || peer.request.EnvironmentID != owner.EnvironmentID || peer.request.SessionID != owner.SessionID || string(peer.data) != setupCanary { t.Fatal("file transport changed") } size++ - if err := installInitialFile(t.Context(), peer, owner, environmentconfig.InitialFileMetadata{SizeBytes: &size}, body); err == nil { + if err := installInitialFile(t.Context(), peer, proto.AssignmentRef{}, owner, environmentconfig.InitialFileMetadata{SizeBytes: &size}, body); err == nil { t.Fatal("mismatched source size accepted") } } diff --git a/services/core/internal/execution/worker.go b/services/core/internal/execution/worker.go index 026a0a015..0766df958 100644 --- a/services/core/internal/execution/worker.go +++ b/services/core/internal/execution/worker.go @@ -187,6 +187,9 @@ func (w *Worker) Run(ctx context.Context) (runErr error) { preparationDone := make(chan error, 1) running.Add(1) go func() { defer running.Done(); preparationDone <- w.runEnvironmentInitializations(ctx) }() + releasesDone := make(chan error, 1) + running.Add(1) + go func() { defer running.Done(); releasesDone <- w.runAssignmentReleases(ctx) }() lifecycleDone := make(chan error, 1) if w.runtimes != nil { running.Add(1) @@ -221,6 +224,8 @@ func (w *Worker) Run(ctx context.Context) (runErr error) { return ctx.Err() case err := <-preparationDone: return err + case err := <-releasesDone: + return err case err := <-lifecycleDone: return err case request := <-w.fileWrites: diff --git a/services/core/internal/persistence/postgres/coremetricspg/store_test.go b/services/core/internal/persistence/postgres/coremetricspg/store_test.go index db5d8edba..4485d4386 100644 --- a/services/core/internal/persistence/postgres/coremetricspg/store_test.go +++ b/services/core/internal/persistence/postgres/coremetricspg/store_test.go @@ -24,7 +24,7 @@ func coreMetricsSession(t *testing.T, pool *pgxpool.Pool, deleted bool) string { t.Fatal(err) } t.Cleanup(func() { - for _, query := range []string{`DELETE FROM turns WHERE session_id=$1`, `DELETE FROM session_devices WHERE session_id=$1`, `DELETE FROM sessions WHERE id=$1`} { + for _, query := range []string{`DELETE FROM turns WHERE session_id=$1`, `DELETE FROM session_runtime_assignments WHERE session_id=$1`, `DELETE FROM sessions WHERE id=$1`} { if _, err := pool.Exec(context.Background(), query, id); err != nil { t.Error(err) } @@ -63,7 +63,7 @@ func TestCoreMetricsSnapshot(t *testing.T) { session := coreMetricsSession(t, pool, false) coreMetricsTurn(t, pool, session, status, "", oldest.Add(time.Duration(i)*time.Second), nil, nil) if i == 0 { - if _, err := pool.Exec(t.Context(), `INSERT INTO session_devices(session_id,device_id) VALUES($1,$2)`, session, connected); err != nil { + if _, err := pool.Exec(t.Context(), `INSERT INTO session_runtime_assignments(session_id,runtime_id) VALUES($1,$2)`, session, connected); err != nil { t.Fatal(err) } } diff --git a/services/core/internal/persistence/postgres/deploymentpg/allocations.go b/services/core/internal/persistence/postgres/deploymentpg/allocations.go index 08e2569cb..d1fb3044c 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/allocations.go +++ b/services/core/internal/persistence/postgres/deploymentpg/allocations.go @@ -340,6 +340,9 @@ func (t *cleanupTx) RevokeDevice(current deployment.Allocation) error { if err != nil { return err } + if err := t.ReleaseAssignment(t.ctx, false); err != nil { + return err + } _, err = t.q.RevokeRuntimeCleanupDevice(t.ctx, sqlc.RevokeRuntimeCleanupDeviceParams{TenantID: t.tenant, DeviceID: device}) return err } diff --git a/services/core/internal/persistence/postgres/deploymentpg/session_archive.go b/services/core/internal/persistence/postgres/deploymentpg/session_archive.go index a361597ee..e576f5013 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/session_archive.go +++ b/services/core/internal/persistence/postgres/deploymentpg/session_archive.go @@ -69,6 +69,9 @@ func (t *archiveTx) RequestArchiveCleanup(current deployment.Allocation) error { if err != nil { return err } + if err := t.ReleaseAssignment(t.ctx, false); err != nil { + return err + } if _, err := t.q.RevokeArchivedRuntimeDevice(t.ctx, sqlc.RevokeArchivedRuntimeDeviceParams{TenantID: t.tenant, DeviceID: device, SessionID: t.session}); err != nil { return err } diff --git a/services/core/internal/persistence/postgres/sessionpg/devices.go b/services/core/internal/persistence/postgres/sessionpg/devices.go index 9bf5e67e5..887aa50bb 100644 --- a/services/core/internal/persistence/postgres/sessionpg/devices.go +++ b/services/core/internal/persistence/postgres/sessionpg/devices.go @@ -9,6 +9,7 @@ import ( "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" @@ -43,7 +44,8 @@ func (s *Store) GetSessionExecutionBinding(ctx context.Context, tenant, session return sessions.ExecutionBinding{}, err } return sessions.ExecutionBinding{ - Device: sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, EnvironmentID: optionalID(row.EnvironmentID)}, + Device: sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, EnvironmentID: optionalID(row.EnvironmentID), + Assignment: assignmentRef(lookup.ID, row.AssignmentID, row.Epoch)}, NativeSessionID: row.NativeSessionID, HasStartedTurn: row.HasStartedTurn, }, nil @@ -98,7 +100,17 @@ func loadSessionDevice(ctx context.Context, q *sqlc.Queries, tenant, session pgt if err != nil { return sessions.ExecutionDevice{}, false, err } - return sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, EnvironmentID: optionalID(row.EnvironmentID)}, true, nil + return sessions.ExecutionDevice{ID: uuid.UUID(row.ID.Bytes).String(), Name: row.Name, EnvironmentID: optionalID(row.EnvironmentID), + Assignment: assignmentRef(session, row.AssignmentID, row.Epoch)}, true, nil +} + +// assignmentRef is the reference of a Session's assignment; an absent +// assignment is the zero reference. +func assignmentRef(session, assignment pgtype.UUID, epoch int64) proto.AssignmentRef { + if !assignment.Valid { + return proto.AssignmentRef{} + } + return proto.AssignmentRef{SessionID: optionalID(session), AssignmentID: optionalID(assignment), Epoch: uint64(epoch)} } // GetDeviceCredential reads a device's credential for the Runtime gateway. A @@ -279,8 +291,16 @@ func (t *SessionTx) LoadDevice(ctx context.Context, device string) (bool, error) return err == nil, err } +// ReleaseAssignment records that the Session's Runtime assignment is +// released, with home removal if removeHome, and advances its epoch unless +// that release is already recorded. A Session without an assignment has +// nothing to release. +func (t *SessionTx) ReleaseAssignment(ctx context.Context, removeHome bool) error { + return t.q.ReleaseSessionAssignment(ctx, sqlc.ReleaseSessionAssignmentParams{SessionID: t.session, RemoveHome: removeHome}) +} + // BindDevice binds the device to the Session; a Session bound to another -// device is sessions.ErrDeviceBindingConflict. +// device, or a released assignment, is sessions.ErrDeviceBindingConflict. func (t *SessionTx) BindDevice(ctx context.Context, device string) error { id, err := parseID(device) if err != nil { diff --git a/services/core/internal/persistence/postgres/sessionpg/devices_test.go b/services/core/internal/persistence/postgres/sessionpg/devices_test.go index 71c17500a..55ff5bcc9 100644 --- a/services/core/internal/persistence/postgres/sessionpg/devices_test.go +++ b/services/core/internal/persistence/postgres/sessionpg/devices_test.go @@ -65,7 +65,7 @@ func TestBindSessionDeviceTranslatesTheBindingOutcome(t *testing.T) { t.Fatal("rebinding the Session to another device", err) } var bound string - if err := pool.QueryRow(t.Context(), `SELECT device_id::text FROM session_devices WHERE session_id = $1`, sessionID).Scan(&bound); err != nil || bound != device { + if err := pool.QueryRow(t.Context(), `SELECT runtime_id::text FROM session_runtime_assignments WHERE session_id = $1`, sessionID).Scan(&bound); err != nil || bound != device { t.Fatal("Session bound to", bound, err) } } diff --git a/services/core/internal/persistence/postgres/sessionpg/environment.go b/services/core/internal/persistence/postgres/sessionpg/environment.go index 830d095ab..25e56c911 100644 --- a/services/core/internal/persistence/postgres/sessionpg/environment.go +++ b/services/core/internal/persistence/postgres/sessionpg/environment.go @@ -63,7 +63,8 @@ func (s *Store) ListEnvironmentInitializations(ctx context.Context, after string for _, row := range rows { result = append(result, sessions.EnvironmentInitialization{ EnvironmentID: optionalID(row.ID), SessionID: optionalID(row.SessionID), TenantID: optionalID(row.TenantID), - DeviceID: optionalID(row.DeviceID), State: row.Initialization, Engine: row.Engine, + DeviceID: optionalID(row.RuntimeID), State: row.Initialization, Engine: row.Engine, + Assignment: assignmentRef(row.SessionID, row.AssignmentID, row.Epoch.Int64), }) } return result, nil diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_environment.go b/services/core/internal/persistence/postgres/sessionpg/execution_environment.go index 505f859ec..bc2b8ed03 100644 --- a/services/core/internal/persistence/postgres/sessionpg/execution_environment.go +++ b/services/core/internal/persistence/postgres/sessionpg/execution_environment.go @@ -6,7 +6,9 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" @@ -73,6 +75,46 @@ func (e *Execution) ListEnvironmentConnections(ctx context.Context, after string return keys, nil } +func (e *Execution) ListAssignmentReleases(ctx context.Context, runtimes []string) ([]sessions.AssignmentRelease, error) { + ids := make([]pgtype.UUID, 0, len(runtimes)) + for _, runtime := range runtimes { + id, err := parseID(runtime) + if err != nil { + return nil, err + } + ids = append(ids, id) + } + var rows []sqlc.ListPendingAssignmentReleasesRow + err := e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + var err error + rows, err = sqlc.New(tx).ListPendingAssignmentReleases(ctx, ids) + return err + }) + if err != nil { + return nil, err + } + releases := make([]sessions.AssignmentRelease, 0, len(rows)) + for _, row := range rows { + releases = append(releases, sessions.AssignmentRelease{RuntimeID: optionalID(row.RuntimeID), Assignment: assignmentRef(row.SessionID, row.AssignmentID, row.Epoch), RemoveHome: row.RemoveHome}) + } + return releases, nil +} + +func (e *Execution) AcknowledgeAssignmentRelease(ctx context.Context, assignment proto.AssignmentRef) error { + session, err := parseID(assignment.SessionID) + if err != nil { + return err + } + id, err := parseID(assignment.AssignmentID) + if err != nil { + return err + } + return e.lease.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + _, err := sqlc.New(tx).AcknowledgeAssignmentRelease(ctx, sqlc.AcknowledgeAssignmentReleaseParams{SessionID: session, AssignmentID: id, Epoch: int64(assignment.Epoch)}) + return err + }) +} + // environmentTx is the Session's Environment inside the Session transaction // of an initialization or connection operation. type environmentTx struct{ *SessionTx } diff --git a/services/core/internal/persistence/postgres/sessionpg/execution_journal_test.go b/services/core/internal/persistence/postgres/sessionpg/execution_journal_test.go index a7a1ffe7a..42adfa930 100644 --- a/services/core/internal/persistence/postgres/sessionpg/execution_journal_test.go +++ b/services/core/internal/persistence/postgres/sessionpg/execution_journal_test.go @@ -121,7 +121,7 @@ func TestSubagentProjectionStoresTheSessionRows(t *testing.T) { VALUES ($1, $2, 'codex', 'key', 'hash', '{"agent":{"id":"agent_root"}}')`, session, tenant) exec(t, pool, `INSERT INTO turns(id, session_id, status) VALUES ($1, $2, 'in_progress')`, turn, session) exec(t, pool, `INSERT INTO devices(id, tenant_id, name, credential_hash) VALUES ($1, $2, 'device', repeat('a', 64))`, device, tenant) - exec(t, pool, `INSERT INTO session_devices(session_id, device_id, native_session_id) VALUES ($1, $2, 'root')`, session, device) + exec(t, pool, `INSERT INTO session_runtime_assignments(session_id, runtime_id, native_session_id) VALUES ($1, $2, 'root')`, session, device) encode := func(kind string, value any) sessions.ExecutionEvent { raw, err := json.Marshal(value) if err != nil { diff --git a/services/core/internal/persistence/postgres/sessionpg/session_test.go b/services/core/internal/persistence/postgres/sessionpg/session_test.go index 51d3e56e7..d192899d7 100644 --- a/services/core/internal/persistence/postgres/sessionpg/session_test.go +++ b/services/core/internal/persistence/postgres/sessionpg/session_test.go @@ -264,7 +264,7 @@ func TestCreateEnvironmentDeviceBindsOneDevice(t *testing.T) { }) } bound := func(session pgtype.UUID) []uuid.UUID { - rows, err := pool.Query(t.Context(), `SELECT device_id FROM session_devices WHERE session_id = $1`, session) + rows, err := pool.Query(t.Context(), `SELECT runtime_id FROM session_runtime_assignments WHERE session_id = $1`, session) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/persistence/postgres/sessionpg/session_writes.go b/services/core/internal/persistence/postgres/sessionpg/session_writes.go index ac1ad07be..344e44b9a 100644 --- a/services/core/internal/persistence/postgres/sessionpg/session_writes.go +++ b/services/core/internal/persistence/postgres/sessionpg/session_writes.go @@ -36,6 +36,9 @@ type deletionTx struct { } func (t *deletionTx) ApplyDeletion(ctx context.Context) error { + if err := t.ReleaseAssignment(ctx, true); err != nil { + return err + } if err := t.q.DeleteSessionArtifacts(ctx, t.session); err != nil { return err } diff --git a/services/core/internal/runtimegateway/assignment.go b/services/core/internal/runtimegateway/assignment.go new file mode 100644 index 000000000..56c1fe331 --- /dev/null +++ b/services/core/internal/runtimegateway/assignment.go @@ -0,0 +1,117 @@ +package runtimegateway + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" +) + +// Bind binds the Session's assignment to this connection's Runtime: it sends +// assignment_bind and waits for bound, once per connection and reference. +// Every Session operation on the connection follows its Bind. +func (s *Session) Bind(ctx context.Context, ref proto.AssignmentRef, environmentID string) error { + s.assignmentMu.Lock() + bound := s.assignments[ref.SessionID] == ref + s.assignmentMu.Unlock() + if bound { + return nil + } + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() + status, err := s.exchangeAssignment(ctx, proto.TypeAssignmentBind, ref, proto.AssignmentBindPayload{EnvironmentID: environmentID}) + if err != nil { + return err + } + if status.State != proto.AssignmentBound { + return fmt.Errorf("agentdaemon gateway: assignment bind failed: %s", status.ErrorCode) + } + s.assignmentMu.Lock() + s.assignments[ref.SessionID] = ref + s.assignmentMu.Unlock() + return nil +} + +// Release ends the assignment on this connection's Runtime and returns the +// Runtime's status. Only a Runtime that declares home removal accepts +// removeHome. +func (s *Session) Release(ctx context.Context, ref proto.AssignmentRef, removeHome bool) (proto.AssignmentStatusPayload, error) { + s.assignmentMu.Lock() + delete(s.assignments, ref.SessionID) + s.assignmentMu.Unlock() + return s.exchangeAssignment(ctx, proto.TypeAssignmentRelease, ref, proto.AssignmentReleasePayload{RemoveHome: removeHome}) +} + +// RemovesHomes reports the Runtime's declaration of home removal; known is +// false before its first heartbeat. +func (s *Session) RemovesHomes() (supported, known bool) { + s.kindsMu.RLock() + defer s.kindsMu.RUnlock() + return s.homeRemoval.IsSupported(), s.kindsSeen +} + +func (s *Session) exchangeAssignment(ctx context.Context, typ string, ref proto.AssignmentRef, payload any) (proto.AssignmentStatusPayload, error) { + var status proto.AssignmentStatusPayload + if !ref.Valid() { + return status, errors.New("agentdaemon gateway: invalid assignment") + } + id := uuid.NewString() + env, err := proto.NewEnvelope(typ, id, payload) + if err != nil { + return status, err + } + env.Assignment = ref + replies := make(chan proto.Envelope, 1) + s.assignmentMu.Lock() + if s.IsClosed() { + s.assignmentMu.Unlock() + return status, ErrSessionClosed + } + s.assignmentReplies[id] = replies + s.assignmentMu.Unlock() + defer func() { s.assignmentMu.Lock(); delete(s.assignmentReplies, id); s.assignmentMu.Unlock() }() + reply, err := s.exchangeChunkFrame(ctx, env, replies) + if err != nil { + return status, err + } + if reply.Type != proto.TypeAssignmentStatus || reply.Assignment != ref || reply.DecodePayload(&status) != nil || !validAssignmentStatus(typ, status) { + return proto.AssignmentStatusPayload{}, errors.New("agentdaemon gateway: invalid assignment status") + } + return status, nil +} + +func validAssignmentStatus(typ string, status proto.AssignmentStatusPayload) bool { + switch status.State { + case proto.AssignmentFailed: + return status.ErrorCode != "" + case proto.AssignmentBound: + return typ == proto.TypeAssignmentBind && status.ErrorCode == "" + case proto.AssignmentReleased, proto.AssignmentHomeRemoved: + return typ == proto.TypeAssignmentRelease && status.ErrorCode == "" + } + return false +} + +func (s *Session) dispatchAssignmentStatus(env proto.Envelope) { + s.assignmentMu.Lock() + defer s.assignmentMu.Unlock() + if replies := s.assignmentReplies[env.ID]; replies != nil { + select { + case replies <- env: + default: + } + } +} + +func (s *Session) closeAssignmentReplies() { + s.assignmentMu.Lock() + defer s.assignmentMu.Unlock() + for id, replies := range s.assignmentReplies { + close(replies) + delete(s.assignmentReplies, id) + } +} diff --git a/services/core/internal/runtimegateway/cancellation_test.go b/services/core/internal/runtimegateway/cancellation_test.go index aeae810b7..ab019ad63 100644 --- a/services/core/internal/runtimegateway/cancellation_test.go +++ b/services/core/internal/runtimegateway/cancellation_test.go @@ -58,7 +58,7 @@ func TestArchivedReceiptTracksDeliveryBeyondDoneAndRejectsNewWork(t *testing.T) if _, err := peer.TrackExecutionDelivery("new-run"); err == nil { t.Fatal("new delivery admitted while draining") } - for _, kind := range []string{proto.TypePromptRequest, proto.TypePromptSteer, proto.TypeExecutionPrepare, proto.TypeRuntimePrepare, proto.TypeWorkspaceWrite, proto.TypeWorkspaceRead} { + for _, kind := range []string{proto.TypePromptSteer, proto.TypeExecutionPrepare, proto.TypeRuntimePrepare, proto.TypeWorkspaceWrite, proto.TypeWorkspaceRead} { env, _ := proto.NewEnvelope(kind, "run", nil) if err := peer.Send(t.Context(), env); err == nil { t.Fatal("drain permitted new operation", kind) diff --git a/services/core/internal/runtimegateway/capabilities_test.go b/services/core/internal/runtimegateway/capabilities_test.go index da0a6af11..e7e7b66f4 100644 --- a/services/core/internal/runtimegateway/capabilities_test.go +++ b/services/core/internal/runtimegateway/capabilities_test.go @@ -19,7 +19,7 @@ func TestCapabilityProjectionCoversEveryField(t *testing.T) { t.Run(field.Name, func(t *testing.T) { caps := declaration reflect.ValueOf(&caps).Elem().Field(i).Set(reflect.ValueOf(proto.CapabilitySupported)) - kinds := deviceKindsFromHeartbeat(proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "fixture", Available: true, Capabilities: caps}}}) + kinds := deviceKindsFromHeartbeat(proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "fixture", Available: true, Capabilities: caps}}}) actual := reflect.ValueOf(kinds[0].Capabilities) for j := 0; j < deviceType.NumField(); j++ { if actual.Field(j).Bool() != (deviceType.Field(j).Name == field.Name) { @@ -32,7 +32,7 @@ func TestCapabilityProjectionCoversEveryField(t *testing.T) { func TestInvalidHeartbeatDiscardsPreviousDeclarationAndClosesTransport(t *testing.T) { session := NewSession(newFakeConn(), "runtime", "workspace", proto.Version, NewRegistry(), nil) - valid, err := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "fixture", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}}}) + valid, err := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "fixture", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}}}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/runtimegateway/functions_test.go b/services/core/internal/runtimegateway/functions_test.go index e39a7fa83..03e8635d1 100644 --- a/services/core/internal/runtimegateway/functions_test.go +++ b/services/core/internal/runtimegateway/functions_test.go @@ -8,9 +8,9 @@ import ( func TestFunctionCapabilitySurvivesHeartbeatMapping(t *testing.T) { for _, supported := range []bool{false, true} { - kinds := deviceKindsFromHeartbeat(proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{FunctionTools: proto.CapabilityFromBool(supported), FunctionResultImages: proto.CapabilityFromBool(supported)})}}}) + kinds := deviceKindsFromHeartbeat(proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{FunctionTools: proto.CapabilityFromBool(supported), FunctionResultImages: proto.CapabilityFromBool(supported)})}}}) s := &Session{} - s.setSupportedAgentKinds(kinds) + s.setDeclarations(kinds, proto.CapabilityUnsupported) info, found, known := s.AgentKindStatus("codex") if !found || !known || info.Capabilities.FunctionTools != supported || info.Capabilities.FunctionResultImages != supported { t.Fatal(info, found, known) diff --git a/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go b/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go index e03385c68..259463cc9 100644 --- a/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go +++ b/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go @@ -82,22 +82,24 @@ func TestLiveMCPBearerGatewayColdContinuation(t *testing.T) { } allowed, anonymousTools := []string{"remember", "fail"}, []string{"ping"} servers := []proto.MCPHTTPServer{{ConnectionOrigin: "service", ServerLabel: "private_mcp", ServerURL: fixture.private.URL, AllowedTools: &allowed, BearerToken: &token}, {ConnectionOrigin: "service", ServerLabel: "anonymous_mcp", ServerURL: fixture.anonymous.URL, AllowedTools: &anonymousTools}} + assignment := proto.AssignmentRef{SessionID: uuid.NewString(), AssignmentID: uuid.NewString(), Epoch: 1} + if err := peer.Bind(ctx, assignment, ""); err != nil { + t.Fatal("built daemon did not bind the Session's assignment") + } run := func(prompt, resume string, expected map[string]string) *mcpBearerTurn { t.Helper() turn := &mcpBearerTurn{} turns = append(turns, turn) runID := uuid.NewString() - request := proto.PromptRequestPayload{AgentKind: "codex", ConversationID: "mcp-bearer-acceptance", RunID: runID, Input: proto.TextInput(prompt), AgentStateKey: "mcp-bearer-acceptance", AgentSessionID: resume, StrictResume: true, ReleaseOnCompletion: true, ObserveMessages: true, ObserveToolObservations: true, DisableExecutionEnvironment: true, DisableSubagents: true, MCPHTTPServers: &servers, Model: "MiniMax-M3", ExecutionControls: &proto.ExecutionControls{WebSearch: "disabled", TextVerbosity: "medium"}} + request := proto.PromptRequestPayload{AgentKind: "codex", AgentStateKey: "mcp-bearer-acceptance", AgentSessionID: resume, StrictResume: true, ReleaseOnCompletion: true, ObserveMessages: true, ObserveToolObservations: true, DisableExecutionEnvironment: true, DisableSubagents: true, MCPHTTPServers: &servers, Model: "MiniMax-M3", ExecutionControls: &proto.ExecutionControls{WebSearch: "disabled", TextVerbosity: "medium"}} sub, err := peer.SubscribeDurable(runID) if err != nil { t.Fatal("cannot subscribe before real daemon dispatch") } defer peer.Unsubscribe(runID) - envelope, err := proto.NewEnvelope(proto.TypePromptRequest, runID, request) - if err != nil || peer.Send(ctx, envelope) != nil { - t.Fatal("cannot dispatch the private MCP request") - } + control := mcpBearerStart(t, ctx, peer, assignment, request, runID, proto.TextInput(prompt)) mcpBearerCollectTurn(t, ctx, sub, runID, turn, expected, token, provider) + control.release() turn.NativeLaunches = mcpBearerReleased(t, root) turn.BearerEnvironmentReference = mcpBearerConfigReference(t, root, token) return turn @@ -122,6 +124,69 @@ func TestLiveMCPBearerGatewayColdContinuation(t *testing.T) { } } +// mcpBearerControl is one prepared Executor admission of the acceptance +// Session. +type mcpBearerControl struct { + t *testing.T + ctx context.Context + peer *Session + assignment proto.AssignmentRef + id string + events *Subscription + status proto.PreparationStatusPayload +} + +// mcpBearerStart prepares the Session's Executor and starts the Run on it, as +// Core does. +func mcpBearerStart(t *testing.T, ctx context.Context, peer *Session, assignment proto.AssignmentRef, request proto.PromptRequestPayload, runID string, input proto.MessageInput) *mcpBearerControl { + t.Helper() + c := &mcpBearerControl{t: t, ctx: ctx, peer: peer, assignment: assignment, id: uuid.NewString()} + events, err := peer.SubscribePreparation(c.id) + if err != nil { + t.Fatal("cannot subscribe to the preparation") + } + c.events = events + c.send(proto.TypeExecutionPrepare, proto.ExecutionPreparePayload{SessionID: assignment.SessionID, Configuration: request}) + c.await("ready") + c.send(proto.TypeExecutionStart, proto.ExecutionStartPayload{Handle: c.status.Handle, ExecutorID: c.status.ExecutorID, RunID: runID, Input: input}) + c.await("started") + return c +} + +func (c *mcpBearerControl) send(kind string, payload any) { + c.t.Helper() + env, err := proto.NewEnvelope(kind, c.id, payload) + if err != nil { + c.t.Fatal("cannot encode the preparation control") + } + env.Assignment = c.assignment + if c.peer.Send(c.ctx, env) != nil { + c.t.Fatal("cannot dispatch the preparation control") + } +} + +func (c *mcpBearerControl) await(state string) { + c.t.Helper() + for c.status.State != state { + select { + case env, ok := <-c.events.Events: + if !ok || env.DecodePayload(&c.status) != nil || c.status.State == "rejected" || c.status.State == "failed" { + c.t.Fatal("built daemon did not admit the prepared Run") + } + case <-c.ctx.Done(): + c.t.Fatal("prepared Run admission timed out") + } + } +} + +// release closes the Executor, so the next Run starts a fresh native process. +func (c *mcpBearerControl) release() { + c.t.Helper() + defer c.peer.UnsubscribePreparation(c.id) + c.send(proto.TypeExecutionRelease, proto.ExecutionReleasePayload{Handle: c.status.Handle}) + c.await("released") +} + func mcpBearerCollectTurn(t *testing.T, ctx context.Context, sub *Subscription, runID string, turn *mcpBearerTurn, expected map[string]string, secrets ...string) { t.Helper() before, after := make(map[string]string), make(map[string]int) diff --git a/services/core/internal/runtimegateway/mcp_test.go b/services/core/internal/runtimegateway/mcp_test.go index c78c9d6ee..13d7cefa6 100644 --- a/services/core/internal/runtimegateway/mcp_test.go +++ b/services/core/internal/runtimegateway/mcp_test.go @@ -11,7 +11,7 @@ import ( func TestMCPHTTPBearerCapabilitySurvivesHeartbeatMapping(t *testing.T) { for _, supported := range []bool{false, true} { - heartbeat := proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, + heartbeat := proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{MCPHTTPTools: proto.CapabilitySupported, MCPHTTPBearerAuth: proto.CapabilityFromBool(supported)})}}} raw, err := json.Marshal(heartbeat) if err != nil || strings.Contains(string(raw), `"mcp_http_bearer_auth":true`) != supported { @@ -22,7 +22,7 @@ func TestMCPHTTPBearerCapabilitySurvivesHeartbeatMapping(t *testing.T) { t.Fatal(err) } s := &Session{} - s.setSupportedAgentKinds(deviceKindsFromHeartbeat(decoded)) + s.setDeclarations(deviceKindsFromHeartbeat(decoded), proto.CapabilityUnsupported) info, found, known := s.AgentKindStatus("codex") if !found || !known || !info.Capabilities.MCPHTTPTools || info.Capabilities.MCPHTTPBearerAuth != supported { t.Fatal("bearer capability was lost or inferred from credential-free MCP") @@ -32,7 +32,7 @@ func TestMCPHTTPBearerCapabilitySurvivesHeartbeatMapping(t *testing.T) { func TestMCPRequiredCapabilityIsExplicit(t *testing.T) { for _, supported := range []bool{false, true} { - heartbeat := proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{MCPHTTPTools: proto.CapabilitySupported, MCPHTTPBearerAuth: proto.CapabilitySupported, MCPHTTPRequired: proto.CapabilityFromBool(supported)})}}} + heartbeat := proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{MCPHTTPTools: proto.CapabilitySupported, MCPHTTPBearerAuth: proto.CapabilitySupported, MCPHTTPRequired: proto.CapabilityFromBool(supported)})}}} raw, err := json.Marshal(heartbeat) if err != nil || strings.Contains(string(raw), `"mcp_http_required":true`) != supported { t.Fatal("wire capability differs", err) @@ -42,7 +42,7 @@ func TestMCPRequiredCapabilityIsExplicit(t *testing.T) { t.Fatal(err) } s := &Session{} - s.setSupportedAgentKinds(deviceKindsFromHeartbeat(decoded)) + s.setDeclarations(deviceKindsFromHeartbeat(decoded), proto.CapabilityUnsupported) info, found, known := s.AgentKindStatus("codex") if !found || !known || info.Capabilities.MCPHTTPRequired != supported { t.Fatal("combination capability lost or inferred") diff --git a/services/core/internal/runtimegateway/preparation_test.go b/services/core/internal/runtimegateway/preparation_test.go index b4383a498..28e58c854 100644 --- a/services/core/internal/runtimegateway/preparation_test.go +++ b/services/core/internal/runtimegateway/preparation_test.go @@ -78,7 +78,7 @@ func TestPreparationCloseAndOverflowDoNotInventRunEvents(t *testing.T) { } func TestPreparationCapabilitySurvivesHeartbeatMapping(t *testing.T) { - kinds := deviceKindsFromHeartbeat(proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Preparation: proto.CapabilitySupported, LocalEnvironment: proto.CapabilitySupported, WorkspaceReadPreparation: proto.CapabilitySupported, NativeSessionRecovery: proto.CapabilitySupported})}}}) + kinds := deviceKindsFromHeartbeat(proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Preparation: proto.CapabilitySupported, LocalEnvironment: proto.CapabilitySupported, WorkspaceReadPreparation: proto.CapabilitySupported, NativeSessionRecovery: proto.CapabilitySupported})}}}) if len(kinds) != 1 || (!kinds[0].Capabilities.Preparation || !kinds[0].Capabilities.LocalEnvironment || !kinds[0].Capabilities.WorkspaceReadPreparation || !kinds[0].Capabilities.NativeSessionRecovery) { t.Fatal("preparation capability lost") } diff --git a/services/core/internal/runtimegateway/runtime_prepare.go b/services/core/internal/runtimegateway/runtime_prepare.go index 77e6493f9..4f6562390 100644 --- a/services/core/internal/runtimegateway/runtime_prepare.go +++ b/services/core/internal/runtimegateway/runtime_prepare.go @@ -14,7 +14,7 @@ import ( // PrepareRuntime sends one typed preparation operation once. An // interrupted or unconfirmed transfer reports unknown; callers must not replay it. -func (s *Session) PrepareRuntime(ctx context.Context, id string, request proto.RuntimePreparePayload, data []byte) (proto.RuntimePrepareResultPayload, error) { +func (s *Session) PrepareRuntime(ctx context.Context, id string, ref proto.AssignmentRef, request proto.RuntimePreparePayload, data []byte) (proto.RuntimePrepareResultPayload, error) { unknown := proto.RuntimePrepareResultPayload{Outcome: "unknown", ErrorCode: "runtime_preparation_unconfirmed"} parsed, err := uuid.Parse(id) if err != nil || parsed == uuid.Nil || parsed.String() != id || len(data) > proto.RuntimePrepareMaxBytes || @@ -57,6 +57,7 @@ func (s *Session) PrepareRuntime(ctx context.Context, id string, request proto.R if err != nil { return unknown, errors.New("agentdaemon gateway: invalid Runtime frame") } + env.Assignment = ref encoded, err := json.Marshal(env) if err != nil || len(encoded) > proto.RuntimePrepareMaxFrameBytes { return unknown, errors.New("agentdaemon gateway: invalid Runtime frame") diff --git a/services/core/internal/runtimegateway/runtime_prepare_test.go b/services/core/internal/runtimegateway/runtime_prepare_test.go index 02834d87a..6ee8347b8 100644 --- a/services/core/internal/runtimegateway/runtime_prepare_test.go +++ b/services/core/internal/runtimegateway/runtime_prepare_test.go @@ -29,7 +29,10 @@ type capabilityOutcome struct { func beginCapabilities(s *Session, ctx context.Context, id string, request proto.RuntimePreparePayload, data []byte) <-chan capabilityOutcome { done := make(chan capabilityOutcome, 1) - go func() { r, err := s.PrepareRuntime(ctx, id, request, data); done <- capabilityOutcome{r, err} }() + go func() { + r, err := s.PrepareRuntime(ctx, id, testAssignment, request, data) + done <- capabilityOutcome{r, err} + }() return done } func nextCapabilityFrame(t *testing.T, s *Session) proto.Envelope { @@ -145,7 +148,7 @@ func TestCapabilitiesRefusesConflictingArchiveBeforeSending(t *testing.T) { } { request := skillPreparation() mutate(&request) - if result, err := s.PrepareRuntime(t.Context(), uuid.NewString(), request, []byte("data")); err == nil || result.Outcome != "unknown" { + if result, err := s.PrepareRuntime(t.Context(), uuid.NewString(), testAssignment, request, []byte("data")); err == nil || result.Outcome != "unknown" { t.Fatal("conflict admitted", result, err) } noCapabilityFrame(t, s) @@ -206,7 +209,7 @@ func TestCapabilitiesConnectionOwnershipAndUnknownInterruption(t *testing.T) { defer cancel() done := beginCapabilities(s, ctx, uuid.NewString(), skillPreparation(), []byte("data")) nextCapabilityFrame(t, s) - if _, err := s.PrepareRuntime(t.Context(), uuid.NewString(), skillPreparation(), []byte("second")); err == nil { + if _, err := s.PrepareRuntime(t.Context(), uuid.NewString(), testAssignment, skillPreparation(), []byte("second")); err == nil { t.Fatal("concurrent transfer admitted") } if closeConnection { @@ -290,7 +293,7 @@ func TestRuntimeInitializationNoDataAndExitReceipt(t *testing.T) { s := NewSession(newFakeConn(), "device", "tenant", "test", nil, nil) defer s.Close("test") request := proto.RuntimePreparePayload{EnvironmentID: uuid.NewString(), SessionID: uuid.NewString(), Action: "initialize", Initialization: &proto.RuntimeInitialization{Action: "setup", Command: "echo test"}} - if _, err := s.PrepareRuntime(t.Context(), uuid.NewString(), request, []byte("forbidden")); err == nil { + if _, err := s.PrepareRuntime(t.Context(), uuid.NewString(), testAssignment, request, []byte("forbidden")); err == nil { t.Fatal("initialization body accepted") } noCapabilityFrame(t, s) diff --git a/services/core/internal/runtimegateway/session.go b/services/core/internal/runtimegateway/session.go index 8d7e83202..bcffabd39 100644 --- a/services/core/internal/runtimegateway/session.go +++ b/services/core/internal/runtimegateway/session.go @@ -97,6 +97,13 @@ type Session struct { kindsMu sync.RWMutex kindsSeen bool supportedKinds []runtimedevice.SupportedAgentKind + homeRemoval proto.CapabilitySupport + + // assignments holds each Session's reference the Runtime acknowledged + // bound on this connection. + assignmentMu sync.Mutex + assignments map[string]proto.AssignmentRef + assignmentReplies map[string]chan proto.Envelope // Subscribers keyed by runID. The read loop only sends on these // channels; Unsubscribe is the only place that closes them. @@ -146,19 +153,21 @@ func NewSession(conn WSConn, deviceID, workspaceID, daemonVersion string, reg *R } now := time.Now() return &Session{ - DeviceID: deviceID, - WorkspaceID: workspaceID, - DaemonVersion: daemonVersion, - ConnectedAt: now, - conn: conn, - log: log, - reg: reg, - lastSeenAt: now, - subs: map[string]*Subscription{}, - preparations: map[string]*preparationSubscription{}, - ackWaiters: map[string]chan proto.InteractionDecisionAckPayload{}, - sendCh: make(chan proto.Envelope, 64), - closed: make(chan struct{}), + DeviceID: deviceID, + WorkspaceID: workspaceID, + DaemonVersion: daemonVersion, + ConnectedAt: now, + conn: conn, + log: log, + reg: reg, + lastSeenAt: now, + subs: map[string]*Subscription{}, + assignments: map[string]proto.AssignmentRef{}, + assignmentReplies: map[string]chan proto.Envelope{}, + preparations: map[string]*preparationSubscription{}, + ackWaiters: map[string]chan proto.InteractionDecisionAckPayload{}, + sendCh: make(chan proto.Envelope, 64), + closed: make(chan struct{}), } } @@ -215,12 +224,13 @@ func (s *Session) AgentKindStatus(kind string) (info runtimedevice.SupportedAgen return runtimedevice.SupportedAgentKind{}, false, true } -func (s *Session) setSupportedAgentKinds(kinds []runtimedevice.SupportedAgentKind) { +func (s *Session) setDeclarations(kinds []runtimedevice.SupportedAgentKind, homeRemoval proto.CapabilitySupport) { copyKinds := make([]runtimedevice.SupportedAgentKind, len(kinds)) copy(copyKinds, kinds) s.kindsMu.Lock() s.kindsSeen = true s.supportedKinds = copyKinds + s.homeRemoval = homeRemoval s.kindsMu.Unlock() } @@ -247,6 +257,7 @@ func (s *Session) Close(reason string) { s.closeWorkspaceWrites() s.closeCapabilities() s.closeWorkspaceExports() + s.closeAssignmentReplies() }) } @@ -395,14 +406,14 @@ func (s *Session) markSeen() { func (s *Session) handleHeartbeat(env proto.Envelope) { var p proto.HeartbeatPayload - if err := env.DecodePayload(&p); err != nil { + if err := env.DecodePayload(&p); err != nil || p.HomeRemoval == proto.CapabilityUnspecified { s.log("agentdaemon gateway: invalid heartbeat declaration device=%s", s.DeviceID) - s.setSupportedAgentKinds(nil) + s.setDeclarations(nil, proto.CapabilityUnspecified) s.Close("invalid heartbeat declaration") return } kinds := deviceKindsFromHeartbeat(p) - s.setSupportedAgentKinds(kinds) + s.setDeclarations(kinds, p.HomeRemoval) if s.heartbeat == nil { return } @@ -503,6 +514,14 @@ func (s *Session) dispatch(env proto.Envelope) { case proto.TypeHeartbeat: s.handleHeartbeat(env) return + case proto.TypeAssignmentStatus: + s.dispatchAssignmentStatus(env) + return + case proto.TypeProtocolError: + var p proto.ProtocolErrorPayload + _ = env.DecodePayload(&p) + s.log("agentdaemon gateway: Runtime rejected %s request=%s device=%s: %s", p.Type, env.ID, s.DeviceID, p.ErrorCode) + return case proto.TypePermissionRequest: var p proto.PermissionRequestPayload requestID := "" diff --git a/services/core/internal/runtimegateway/session_test.go b/services/core/internal/runtimegateway/session_test.go index a30e8a66d..20dad7072 100644 --- a/services/core/internal/runtimegateway/session_test.go +++ b/services/core/internal/runtimegateway/session_test.go @@ -14,6 +14,9 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/runtimedevice" ) +// testAssignment is the assignment the Session operations under test name. +var testAssignment = proto.AssignmentRef{SessionID: "session", AssignmentID: "assignment", Epoch: 1} + // fakeConn is the WSConn implementation used by session + registry // tests. Concurrency-safe. type fakeConn struct { @@ -342,11 +345,7 @@ func TestSession_SendWritesToWire(t *testing.T) { sess.Start() defer sess.Close("test done") - env, _ := proto.NewEnvelope(proto.TypePromptRequest, "run-1", proto.PromptRequestPayload{ - AgentKind: "fake_alpha", - RunID: "run-1", - Input: proto.TextInput("hello"), - }) + env, _ := proto.NewEnvelope(proto.TypePromptSteer, "run-1", proto.PromptSteerPayload{InputID: "1", Input: proto.TextInput("hello")}) if err := sess.Send(context.Background(), env); err != nil { t.Fatalf("Send: %v", err) } @@ -371,7 +370,7 @@ func TestSession_HeartbeatPersistsSupportedAgentKinds(t *testing.T) { sess.Start() defer sess.Close("test done") - env, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{ + env, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, Timestamp: 1710000000, ActiveRequests: 2, DaemonVersion: "0.2.0-test", @@ -445,7 +444,7 @@ func TestSession_HeartbeatDoesNotInferCapabilities(t *testing.T) { sess.Start() defer sess.Close("test done") - conn.Feed([]byte(`{"type":"heartbeat","payload":{"ts":1710000100,"claude_available":true}}`)) + conn.Feed([]byte(`{"type":"heartbeat","payload":{"ts":1710000100,"claude_available":true,"home_removal":false}}`)) got := heartbeat.waitDaemonHeartbeat(t) if len(got.SupportedAgentKinds) != 0 { t.Fatalf("undeclared capabilities inferred: %#v", got.SupportedAgentKinds) diff --git a/services/core/internal/runtimegateway/suspension.go b/services/core/internal/runtimegateway/suspension.go index a62aaaf1c..9c57dbf48 100644 --- a/services/core/internal/runtimegateway/suspension.go +++ b/services/core/internal/runtimegateway/suspension.go @@ -9,7 +9,8 @@ import ( // SuspendControl correlates one lifecycle operation on this authenticated socket. // A closed socket is not an acknowledgement; callers retain their durable intent. -func (s *Session) SuspendControl(ctx context.Context, kind string, request proto.EnvironmentSuspendPayload) (proto.EnvironmentSuspendResultPayload, error) { +// The reference names the Session's assignment. +func (s *Session) SuspendControl(ctx context.Context, kind string, ref proto.AssignmentRef, request proto.EnvironmentSuspendPayload) (proto.EnvironmentSuspendResultPayload, error) { var result proto.EnvironmentSuspendResultPayload expected := proto.TypeEnvironmentQuiesced if kind == proto.TypeEnvironmentResume { @@ -25,6 +26,7 @@ func (s *Session) SuspendControl(ctx context.Context, kind string, request proto if err != nil { return result, err } + envelope.Assignment = ref s.suspendMu.Lock() if s.IsClosed() { s.suspendMu.Unlock() diff --git a/services/core/internal/runtimegateway/suspension_test.go b/services/core/internal/runtimegateway/suspension_test.go index 081d2332d..e76bb9e62 100644 --- a/services/core/internal/runtimegateway/suspension_test.go +++ b/services/core/internal/runtimegateway/suspension_test.go @@ -12,7 +12,7 @@ func TestSuspendAcknowledgementSurvivesImmediateConnectionClose(t *testing.T) { s := NewSession(newFakeConn(), "device", "tenant", "test", nil, nil) done := make(chan error, 1) go func() { - result, err := s.SuspendControl(t.Context(), proto.TypeEnvironmentQuiesce, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) + result, err := s.SuspendControl(t.Context(), proto.TypeEnvironmentQuiesce, testAssignment, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) if err == nil && !result.Accepted { err = errors.New("ack lost") } @@ -36,7 +36,7 @@ func TestSuspendControlRejectsForeignIdentityAndDoesNotReplay(t *testing.T) { s := NewSession(newFakeConn(), "device", "tenant", "test", nil, nil) done := make(chan error, 1) go func() { - _, err := s.SuspendControl(t.Context(), proto.TypeEnvironmentResume, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) + _, err := s.SuspendControl(t.Context(), proto.TypeEnvironmentResume, testAssignment, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) done <- err }() request := <-s.sendCh @@ -59,7 +59,7 @@ func TestSuspendObserverCancellationRetainsUnknownOutcome(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { - _, err := s.SuspendControl(ctx, proto.TypeEnvironmentQuiesce, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) + _, err := s.SuspendControl(ctx, proto.TypeEnvironmentQuiesce, testAssignment, proto.EnvironmentSuspendPayload{EnvironmentID: "env", SuspendID: "attempt"}) done <- err }() <-s.sendCh diff --git a/services/core/internal/runtimegateway/wire_test.go b/services/core/internal/runtimegateway/wire_test.go index 7b09c5e9b..eb20e7ec9 100644 --- a/services/core/internal/runtimegateway/wire_test.go +++ b/services/core/internal/runtimegateway/wire_test.go @@ -111,6 +111,9 @@ type coreSide struct { preparations map[string]*runtimegateway.Subscription runs map[string]*runtimegateway.Subscription receipts map[string]chan receipt + // bind is the pending Bind and bindID the request ID it generated. + bind chan error + bindID string } func (c *coreSide) connect() { @@ -193,6 +196,20 @@ func (c *coreSide) coreSends(frame proto.Envelope) { t := c.t t.Helper() switch frame.Type { + case proto.TypeAssignmentBind: + var request proto.AssignmentBindPayload + if err := frame.DecodePayload(&request); err != nil { + t.Fatal(err) + } + c.bind = make(chan error, 1) + go func() { c.bind <- c.session.Bind(t.Context(), frame.Assignment, request.EnvironmentID) }() + got := c.peer.receive(t) + // Bind generates its request ID; the Runtime's status echoes it. + c.bindID, got.ID = got.ID, frame.ID + if err := prototest.SameFrame(frame, got); err != nil { + t.Fatal(err) + } + return case proto.TypeExecutionPrepare: sub, err := c.session.SubscribePreparation(frame.ID) if err != nil { @@ -257,6 +274,22 @@ func (c *coreSide) runtimeSends(frame proto.Envelope) { t.Fatal("missing receipt") } delete(c.receipts, want.DeliveryID) + case proto.TypeAssignmentStatus: + select { + case err := <-c.bind: + t.Fatalf("Bind returned before the Runtime answered: %v", err) + default: + } + frame.ID = c.bindID + c.write(frame) + select { + case err := <-c.bind: + if err != nil { + t.Fatal(err) + } + case <-time.After(wait): + t.Fatal("Bind did not complete") + } case proto.TypePreparationStatus: c.write(frame) c.delivered(c.preparations[frame.ID], frame) diff --git a/services/core/internal/runtimegateway/workspace_directory_test.go b/services/core/internal/runtimegateway/workspace_directory_test.go index 3f6ea3c0d..e95e7aae9 100644 --- a/services/core/internal/runtimegateway/workspace_directory_test.go +++ b/services/core/internal/runtimegateway/workspace_directory_test.go @@ -15,7 +15,7 @@ func TestWorkspaceDirectorySharesReadCorrelationAndFrameBound(t *testing.T) { request := proto.WorkspaceReadPayload{Handle: "prepared", EnvironmentID: "environment", MaxEntries: proto.WorkspaceDirectoryMaxEntries} done := make(chan error, 1) go func() { - _, err := s.ListWorkspaceDirectory(t.Context(), request) + _, err := s.ListWorkspaceDirectory(t.Context(), testAssignment, request) done <- err }() message := <-s.sendCh @@ -79,7 +79,7 @@ func TestWorkspaceDirectoryRequiresExplicitWireTruncation(t *testing.T) { defer s.Close("test") done := make(chan error, 1) go func() { - _, err := s.ListWorkspaceDirectory(t.Context(), proto.WorkspaceReadPayload{Handle: "prepared", EnvironmentID: "environment", MaxEntries: 1}) + _, err := s.ListWorkspaceDirectory(t.Context(), testAssignment, proto.WorkspaceReadPayload{Handle: "prepared", EnvironmentID: "environment", MaxEntries: 1}) done <- err }() request := <-s.sendCh diff --git a/services/core/internal/runtimegateway/workspace_export.go b/services/core/internal/runtimegateway/workspace_export.go index 2b9549ff5..273e504dd 100644 --- a/services/core/internal/runtimegateway/workspace_export.go +++ b/services/core/internal/runtimegateway/workspace_export.go @@ -11,7 +11,7 @@ import ( ) // ExportWorkspaceOutputs consumes bounded chunks and requires the exporter completion receipt. -func (s *Session) ExportWorkspaceOutputs(ctx context.Context, request proto.WorkspaceExportPayload, consume func(io.Reader) error) error { +func (s *Session) ExportWorkspaceOutputs(ctx context.Context, ref proto.AssignmentRef, request proto.WorkspaceExportPayload, consume func(io.Reader) error) error { request.Step = "begin" if !proto.ValidWorkspaceExportRequest(request) || consume == nil { return errors.New("agentdaemon gateway: invalid workspace export") @@ -32,12 +32,13 @@ func (s *Session) ExportWorkspaceOutputs(ctx context.Context, request proto.Work defer func() { s.workspaceExportMu.Lock(); delete(s.workspaceExports, id); s.workspaceExportMu.Unlock() }() ctx, cancel := context.WithTimeout(ctx, 180*time.Second) defer cancel() - r := &workspaceExportReader{ctx: ctx, peer: s, id: id, replies: replies, request: request} + r := &workspaceExportReader{ctx: ctx, peer: s, id: id, ref: ref, replies: replies, request: request} defer func() { if !r.completed { stop, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() env, _ := proto.NewEnvelope(proto.TypeWorkspaceExport, id, proto.WorkspaceExportPayload{Step: "cancel"}) + env.Assignment = ref _ = s.Send(stop, env) } }() @@ -53,6 +54,7 @@ type workspaceExportReader struct { ctx context.Context peer *Session id string + ref proto.AssignmentRef replies <-chan proto.Envelope request proto.WorkspaceExportPayload data []byte @@ -76,6 +78,7 @@ func (r *workspaceExportReader) Read(p []byte) (int, error) { if err != nil { return 0, err } + env.Assignment = r.ref if err = r.peer.Send(r.ctx, env); err != nil { return 0, err } diff --git a/services/core/internal/runtimegateway/workspace_export_test.go b/services/core/internal/runtimegateway/workspace_export_test.go index dc1b0bd4d..6b6efeae0 100644 --- a/services/core/internal/runtimegateway/workspace_export_test.go +++ b/services/core/internal/runtimegateway/workspace_export_test.go @@ -27,7 +27,7 @@ func TestWorkspaceExportPullsOnlyAfterConsumedAndRequiresCompletion(t *testing.T done := make(chan error, 1) body := bytes.Repeat([]byte{0, 255, 7}, 100) go func() { - done <- s.ExportWorkspaceOutputs(t.Context(), proto.WorkspaceExportPayload{Handle: "prepared", EnvironmentID: "env"}, func(r io.Reader) error { + done <- s.ExportWorkspaceOutputs(t.Context(), testAssignment, proto.WorkspaceExportPayload{Handle: "prepared", EnvironmentID: "env"}, func(r io.Reader) error { got := make([]byte, len(body)) if _, err := io.ReadFull(r, got); err != nil { return err @@ -77,7 +77,7 @@ func TestWorkspaceExportRejectsInvalidChunksAndNativeFailure(t *testing.T) { s := NewSession(newFakeConn(), "device", "tenant", "test", nil, nil) done := make(chan error, 1) go func() { - done <- s.ExportWorkspaceOutputs(t.Context(), proto.WorkspaceExportPayload{Handle: "prepared", EnvironmentID: "env"}, func(r io.Reader) error { _, err := io.Copy(io.Discard, r); return err }) + done <- s.ExportWorkspaceOutputs(t.Context(), testAssignment, proto.WorkspaceExportPayload{Handle: "prepared", EnvironmentID: "env"}, func(r io.Reader) error { _, err := io.Copy(io.Discard, r); return err }) }() first := <-s.sendCh exportReply(t, s, first.ID, result) @@ -99,7 +99,7 @@ func TestWorkspaceExportDisconnectAndCancellationSettleRead(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { - done <- s.ExportWorkspaceOutputs(ctx, proto.WorkspaceExportPayload{Handle: "prepared", EnvironmentID: "env"}, func(r io.Reader) error { _, err := io.Copy(io.Discard, r); return err }) + done <- s.ExportWorkspaceOutputs(ctx, testAssignment, proto.WorkspaceExportPayload{Handle: "prepared", EnvironmentID: "env"}, func(r io.Reader) error { _, err := io.Copy(io.Discard, r); return err }) }() <-s.sendCh if disconnect { diff --git a/services/core/internal/runtimegateway/workspace_read.go b/services/core/internal/runtimegateway/workspace_read.go index ca7159086..f0846ac05 100644 --- a/services/core/internal/runtimegateway/workspace_read.go +++ b/services/core/internal/runtimegateway/workspace_read.go @@ -10,19 +10,19 @@ import ( ) // ReadWorkspaceFile observes one private operation; cancellation never retries or cancels native work. -func (s *Session) ReadWorkspaceFile(ctx context.Context, request proto.WorkspaceReadPayload) (proto.WorkspaceReadResultPayload, error) { +func (s *Session) ReadWorkspaceFile(ctx context.Context, ref proto.AssignmentRef, request proto.WorkspaceReadPayload) (proto.WorkspaceReadResultPayload, error) { if request.Operation != "" { return proto.WorkspaceReadResultPayload{}, errors.New("agentdaemon gateway: invalid byte read operation") } - return s.readWorkspace(ctx, request) + return s.readWorkspace(ctx, ref, request) } -func (s *Session) ListWorkspaceDirectory(ctx context.Context, request proto.WorkspaceReadPayload) (proto.WorkspaceReadResultPayload, error) { +func (s *Session) ListWorkspaceDirectory(ctx context.Context, ref proto.AssignmentRef, request proto.WorkspaceReadPayload) (proto.WorkspaceReadResultPayload, error) { request.Operation = "directory" - return s.readWorkspace(ctx, request) + return s.readWorkspace(ctx, ref, request) } -func (s *Session) readWorkspace(ctx context.Context, request proto.WorkspaceReadPayload) (proto.WorkspaceReadResultPayload, error) { +func (s *Session) readWorkspace(ctx context.Context, ref proto.AssignmentRef, request proto.WorkspaceReadPayload) (proto.WorkspaceReadResultPayload, error) { var result proto.WorkspaceReadResultPayload if !proto.ValidWorkspaceReadRequest(request) { return result, errors.New("agentdaemon gateway: invalid workspace read") @@ -32,6 +32,7 @@ func (s *Session) readWorkspace(ctx context.Context, request proto.WorkspaceRead if err != nil || len(env.Payload) > proto.WorkspaceReadMaxRequestBytes { return result, errors.New("agentdaemon gateway: invalid workspace read") } + env.Assignment = ref s.workspaceReadMu.Lock() if s.IsClosed() { s.workspaceReadMu.Unlock() diff --git a/services/core/internal/runtimegateway/workspace_read_test.go b/services/core/internal/runtimegateway/workspace_read_test.go index f9df6a871..c7de9567f 100644 --- a/services/core/internal/runtimegateway/workspace_read_test.go +++ b/services/core/internal/runtimegateway/workspace_read_test.go @@ -22,7 +22,7 @@ func TestWorkspaceReadCorrelatesOneBoundedResult(t *testing.T) { done := make(chan error, 1) data := bytes.Repeat([]byte{0, 127, 255, 3}, proto.WorkspaceReadMaxBytes/4) go func() { - result, err := s.ReadWorkspaceFile(t.Context(), workspaceReadRequest()) + result, err := s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()) if err == nil && (!bytes.Equal(result.Data, data) || !result.Truncated || !result.CloseAcknowledged) { err = errors.New("read data or acknowledgment differs") } @@ -56,7 +56,10 @@ func TestWorkspaceReadRejectsIncompleteOrContradictoryReplies(t *testing.T) { } { s := NewSession(newFakeConn(), "device", "tenant", "test", nil, nil) done := make(chan error, 1) - go func() { _, err := s.ReadWorkspaceFile(t.Context(), workspaceReadRequest()); done <- err }() + go func() { + _, err := s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()) + done <- err + }() request := <-s.sendCh reply, _ := proto.NewEnvelope(proto.TypeWorkspaceReadResult, request.ID, result) s.dispatch(reply) @@ -87,9 +90,9 @@ func TestWorkspaceDirectoryAcceptsNotDirectoryOnlyForDirectoryReads(t *testing.T go func() { var err error if test.directory { - _, err = s.ListWorkspaceDirectory(t.Context(), directory) + _, err = s.ListWorkspaceDirectory(t.Context(), testAssignment, directory) } else { - _, err = s.ReadWorkspaceFile(t.Context(), workspaceReadRequest()) + _, err = s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()) } done <- err }() @@ -108,7 +111,7 @@ func TestWorkspaceReadObserverCancellationDoesNotSendCancelOrRetry(t *testing.T) defer s.Close("test") ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) - go func() { _, err := s.ReadWorkspaceFile(ctx, workspaceReadRequest()); done <- err }() + go func() { _, err := s.ReadWorkspaceFile(ctx, testAssignment, workspaceReadRequest()); done <- err }() request := <-s.sendCh cancel() if err := <-done; !errors.Is(err, context.Canceled) { @@ -133,14 +136,17 @@ func TestWorkspaceReadCapacityAndConnectionLoss(t *testing.T) { s := NewSession(newFakeConn(), "device", "tenant", "test", nil, nil) done := make(chan error, 4) for i := 0; i < 4; i++ { - go func() { _, err := s.ReadWorkspaceFile(t.Context(), workspaceReadRequest()); done <- err }() + go func() { + _, err := s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()) + done <- err + }() select { case <-s.sendCh: case <-time.After(time.Second): t.Fatal("read not sent") } } - if _, err := s.ReadWorkspaceFile(t.Context(), workspaceReadRequest()); err == nil { + if _, err := s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()); err == nil { t.Fatal("capacity bypassed") } s.Close("connection lost") @@ -170,7 +176,7 @@ func TestWorkspaceReadRejectsOversizedRequestsBeforeQueueing(t *testing.T) { } ctx, cancel := context.WithTimeout(t.Context(), time.Second) defer cancel() - if _, err := s.ReadWorkspaceFile(ctx, request); err == nil || errors.Is(err, context.DeadlineExceeded) { + if _, err := s.ReadWorkspaceFile(ctx, testAssignment, request); err == nil || errors.Is(err, context.DeadlineExceeded) { t.Fatal("oversized request not rejected before send", err) } select { diff --git a/services/core/internal/runtimegateway/workspace_write.go b/services/core/internal/runtimegateway/workspace_write.go index a69223bbd..109c5b960 100644 --- a/services/core/internal/runtimegateway/workspace_write.go +++ b/services/core/internal/runtimegateway/workspace_write.go @@ -13,7 +13,7 @@ import ( // WriteWorkspaceFile sends an already durably owned mutation once. A transport // error is not a rejection and must retain the caller's unknown-outcome gate. -func (s *Session) WriteWorkspaceFile(ctx context.Context, id string, request proto.WorkspaceWritePayload, data []byte) (proto.WorkspaceWriteResultPayload, error) { +func (s *Session) WriteWorkspaceFile(ctx context.Context, id string, ref proto.AssignmentRef, request proto.WorkspaceWritePayload, data []byte) (proto.WorkspaceWriteResultPayload, error) { var empty proto.WorkspaceWriteResultPayload parsed, err := uuid.Parse(id) if err != nil || parsed == uuid.Nil || parsed.String() != id || len(data) > proto.WorkspaceWriteMaxBytes { @@ -44,6 +44,7 @@ func (s *Session) WriteWorkspaceFile(ctx context.Context, id string, request pro if err != nil || len(env.Payload) > proto.WorkspaceWriteMaxFrameBytes { return empty, errors.New("agentdaemon gateway: invalid write frame") } + env.Assignment = ref reply, err := s.exchangeChunkFrame(ctx, env, replies) if err != nil { return empty, err diff --git a/services/core/internal/runtimegateway/workspace_write_test.go b/services/core/internal/runtimegateway/workspace_write_test.go index 924566f94..17fca00be 100644 --- a/services/core/internal/runtimegateway/workspace_write_test.go +++ b/services/core/internal/runtimegateway/workspace_write_test.go @@ -20,7 +20,7 @@ func TestWorkspaceWriteChunksAndCorrelatesReceipt(t *testing.T) { request := proto.WorkspaceWritePayload{EnvironmentID: uuid.NewString(), SessionID: uuid.NewString(), Path: "file"} done := make(chan error, 1) go func() { - result, err := s.WriteWorkspaceFile(t.Context(), id, request, data) + result, err := s.WriteWorkspaceFile(t.Context(), id, testAssignment, request, data) if err == nil && (result.Outcome != "completed" || result.SizeBytes != len(data)) { err = errors.New("missing commit") } @@ -72,7 +72,7 @@ func TestWorkspaceWriteDoesNotCommitAfterLostObservation(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) done := make(chan error, 1) go func() { - _, err := s.WriteWorkspaceFile(ctx, uuid.NewString(), proto.WorkspaceWritePayload{EnvironmentID: uuid.NewString(), SessionID: uuid.NewString(), Path: "file"}, []byte("value")) + _, err := s.WriteWorkspaceFile(ctx, uuid.NewString(), testAssignment, proto.WorkspaceWritePayload{EnvironmentID: uuid.NewString(), SessionID: uuid.NewString(), Path: "file"}, []byte("value")) done <- err }() <-s.sendCh diff --git a/services/core/internal/sessions/environment.go b/services/core/internal/sessions/environment.go index 3ec6c03fe..14f3edd3d 100644 --- a/services/core/internal/sessions/environment.go +++ b/services/core/internal/sessions/environment.go @@ -9,6 +9,7 @@ import ( "github.com/google/uuid" + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/environmentconfig" ) @@ -76,6 +77,8 @@ type Environment struct { // A running record without its process-local owner is unknown, never replayable. type EnvironmentInitialization struct { EnvironmentID, SessionID, TenantID, DeviceID, State, Engine string + // Assignment is the Session's bound assignment to DeviceID. + Assignment proto.AssignmentRef } // EnvironmentInputActivity is the reservation-owned override before a newer Turn exists. diff --git a/services/core/internal/sessions/execution_environment.go b/services/core/internal/sessions/execution_environment.go index 8974f4955..107b8ed5b 100644 --- a/services/core/internal/sessions/execution_environment.go +++ b/services/core/internal/sessions/execution_environment.go @@ -5,6 +5,8 @@ import ( "errors" "github.com/google/uuid" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" ) // EnvironmentKey names one Environment of a tenant. @@ -17,6 +19,14 @@ type EnvironmentConnection struct { Revision int64 } +// AssignmentRelease is a released assignment whose Runtime has not +// acknowledged the release. +type AssignmentRelease struct { + RuntimeID string + Assignment proto.AssignmentRef + RemoveHome bool +} + // EnvironmentExecution is the lease-bound storage of the Environment // execution operations. type EnvironmentExecution interface { @@ -47,6 +57,13 @@ type EnvironmentExecution interface { // Environments of Sessions that were not publicly deleted and that are // connected or have a connection generation. ListEnvironmentConnections(ctx context.Context, after string) ([]EnvironmentKey, error) + // ListAssignmentReleases lists the releases the given Runtimes have not + // acknowledged. + ListAssignmentReleases(ctx context.Context, runtimes []string) ([]AssignmentRelease, error) + // AcknowledgeAssignmentRelease records that the Runtime applied the + // release of the assignment at its epoch. An assignment Core has since + // changed keeps its pending release. + AcknowledgeAssignmentRelease(ctx context.Context, assignment proto.AssignmentRef) error } // InitializationTx is the Session transaction the Environment initialization @@ -342,6 +359,24 @@ func (o *ExecutionOperations) ReconcileEnvironmentConnections(ctx context.Contex } } +// ListAssignmentReleases lists the releases the given Runtimes still owe an +// acknowledgement. +func (o *ExecutionOperations) ListAssignmentReleases(ctx context.Context, runtimes []string) ([]AssignmentRelease, error) { + if len(runtimes) == 0 { + return nil, nil + } + return o.storage.ListAssignmentReleases(ctx, runtimes) +} + +// AcknowledgeAssignmentRelease records the Runtime's acknowledgement of the +// assignment's release. +func (o *ExecutionOperations) AcknowledgeAssignmentRelease(ctx context.Context, assignment proto.AssignmentRef) error { + if !assignment.Valid() || !validID(assignment.SessionID) || !validID(assignment.AssignmentID) { + return ErrInvalidInput + } + return o.storage.AcknowledgeAssignmentRelease(ctx, assignment) +} + // BindSessionDevice binds the tenant's device to the Session, including a // publicly deleted one. A retry with the same device succeeds and another // device is ErrDeviceBindingConflict, so a Session's filesystem never moves diff --git a/services/core/internal/sessions/execution_environment_test.go b/services/core/internal/sessions/execution_environment_test.go index 73461e26b..adb4ab649 100644 --- a/services/core/internal/sessions/execution_environment_test.go +++ b/services/core/internal/sessions/execution_environment_test.go @@ -7,6 +7,8 @@ import ( "strings" "testing" "time" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" ) func (f *fakeTx) LoadSessionDevice(context.Context) (ExecutionDevice, bool, error) { @@ -111,6 +113,16 @@ func (s *fakeExecutionStorage) ListEnvironmentConnections(_ context.Context, aft return page, nil } +func (s *fakeExecutionStorage) ListAssignmentReleases(context.Context, []string) ([]AssignmentRelease, error) { + s.t.Fatal("unexpected ListAssignmentReleases") + return nil, nil +} + +func (s *fakeExecutionStorage) AcknowledgeAssignmentRelease(context.Context, proto.AssignmentRef) error { + s.t.Fatal("unexpected AcknowledgeAssignmentRelease") + return nil +} + // loads is a fake method that reads value and reports whether it was found. func loads[T any](value T, found bool) func() (T, bool, error) { return func() (T, bool, error) { return value, found, nil } diff --git a/services/core/internal/sessions/session.go b/services/core/internal/sessions/session.go index da31907f2..821aaab46 100644 --- a/services/core/internal/sessions/session.go +++ b/services/core/internal/sessions/session.go @@ -101,9 +101,10 @@ type SessionDeletionTx interface { // reservation that no Turn has admitted or superseded, nil when there is // none. LoadEnvironmentInput(ctx context.Context) (*EnvironmentInputState, error) - // ApplyDeletion deletes the Session's Artifacts, releases the node - // placement of its Environments that have no allocation and removes the - // Session from public access. Its row stays so that execution can settle. + // ApplyDeletion releases the Session's Runtime assignment with home + // removal, deletes the Session's Artifacts, releases the node placement of + // its Environments that have no allocation and removes the Session from + // public access. Its row stays so that execution can settle. ApplyDeletion(ctx context.Context) error // RecordDeletionAudit records the write audit of the deletion. RecordDeletionAudit(ctx context.Context) error diff --git a/services/core/internal/sessions/turn.go b/services/core/internal/sessions/turn.go index f6a4574c5..c4323dacb 100644 --- a/services/core/internal/sessions/turn.go +++ b/services/core/internal/sessions/turn.go @@ -4,6 +4,8 @@ import ( "context" "encoding/json" "time" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" ) const ( @@ -82,10 +84,12 @@ type ExecutionEvent struct { type ExecutionWork struct{ TenantID, SessionID, TurnID, Status string } // ExecutionDevice contains safe identity only, never a device credential. +// Assignment is the Session's bound assignment to the device. type ExecutionDevice struct { ID string Name string EnvironmentID string + Assignment proto.AssignmentRef } // ExecutionBinding identifies the Runtime and native history selected for one API Session. diff --git a/services/core/migrations/000092_session_runtime_assignments.sql b/services/core/migrations/000092_session_runtime_assignments.sql new file mode 100644 index 000000000..9878d1fbc --- /dev/null +++ b/services/core/migrations/000092_session_runtime_assignments.sql @@ -0,0 +1,38 @@ +-- +goose Up +-- A Session's binding to its Runtime becomes a fenced assignment. Core +-- advances epoch with each change of desired_state; applied_epoch is the +-- latest released epoch the Runtime acknowledged. +ALTER TABLE session_devices RENAME TO session_runtime_assignments; +ALTER TABLE session_runtime_assignments RENAME COLUMN device_id TO runtime_id; +ALTER INDEX session_devices_pkey RENAME TO session_runtime_assignments_pkey; +ALTER TABLE session_runtime_assignments RENAME CONSTRAINT session_devices_session_id_fkey TO session_runtime_assignments_session_id_fkey; +ALTER TABLE session_runtime_assignments RENAME CONSTRAINT session_devices_device_id_fkey TO session_runtime_assignments_runtime_id_fkey; +ALTER TABLE session_runtime_assignments + ADD COLUMN assignment_id uuid NOT NULL DEFAULT gen_random_uuid() UNIQUE, + ADD COLUMN epoch bigint NOT NULL DEFAULT 1 CHECK (epoch > 0), + ADD COLUMN desired_state text NOT NULL DEFAULT 'bound' CHECK (desired_state IN ('bound', 'released')), + ADD COLUMN remove_home boolean NOT NULL DEFAULT false, + ADD COLUMN applied_epoch bigint NOT NULL DEFAULT 0, + ADD CONSTRAINT session_runtime_assignments_applied CHECK (applied_epoch BETWEEN 0 AND epoch), + ADD CONSTRAINT session_runtime_assignments_release CHECK (desired_state = 'released' OR NOT remove_home); +-- A deleted Session owes its Runtime the release that removes its home. +UPDATE session_runtime_assignments a SET desired_state = 'released', remove_home = true, epoch = 2 +FROM sessions s WHERE s.id = a.session_id AND s.deleted_at IS NOT NULL; +CREATE INDEX session_runtime_assignments_pending_idx ON session_runtime_assignments (runtime_id, session_id) + WHERE desired_state = 'released' AND applied_epoch < epoch; + +-- +goose Down +DROP INDEX session_runtime_assignments_pending_idx; +ALTER TABLE session_runtime_assignments + DROP CONSTRAINT session_runtime_assignments_release, + DROP CONSTRAINT session_runtime_assignments_applied, + DROP COLUMN applied_epoch, + DROP COLUMN remove_home, + DROP COLUMN desired_state, + DROP COLUMN epoch, + DROP COLUMN assignment_id; +ALTER TABLE session_runtime_assignments RENAME CONSTRAINT session_runtime_assignments_runtime_id_fkey TO session_devices_device_id_fkey; +ALTER TABLE session_runtime_assignments RENAME CONSTRAINT session_runtime_assignments_session_id_fkey TO session_devices_session_id_fkey; +ALTER INDEX session_runtime_assignments_pkey RENAME TO session_devices_pkey; +ALTER TABLE session_runtime_assignments RENAME COLUMN runtime_id TO device_id; +ALTER TABLE session_runtime_assignments RENAME TO session_devices; diff --git a/services/core/tests/integration/archive_cancellation_test.go b/services/core/tests/integration/archive_cancellation_test.go index 0b64fef50..4430f0eee 100644 --- a/services/core/tests/integration/archive_cancellation_test.go +++ b/services/core/tests/integration/archive_cancellation_test.go @@ -108,7 +108,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { h.d = h.bound() capabilities := workerEnvironmentCapabilities() capabilities.FunctionTools = proto.CapabilitySupported - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: capabilities}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: capabilities}}}) var peer *runtimegateway.Session for deadline := time.Now().Add(3 * time.Second); ; { peer, err = registry.LookupDevice(owner.DeviceID) @@ -234,7 +234,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { if err != nil { t.Fatal(err) } - if _, err := tx.Exec(t.Context(), "SELECT session_id FROM session_devices WHERE session_id=$1 FOR UPDATE", session.ID); err != nil { + if _, err := tx.Exec(t.Context(), "SELECT session_id FROM session_runtime_assignments WHERE session_id=$1 FOR UPDATE", session.ID); err != nil { t.Fatal(err) } unlockCommit = func() { @@ -252,7 +252,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { if scenario == "transport_lost" { h.conn.Close() } else if heartbeat && scenario != "ack_commit_blocked" { - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported}) } if !closedCase { ack := proto.InteractionDecisionAckPayload{DeliveryID: request.DeliveryID, Applied: true, Outcome: &proto.DonePayload{Usage: proto.Usage{InputTokens: 17}, Metadata: map[string]any{proto.DoneMetaAgentSessionID: "cancelled-native"}}} @@ -270,7 +270,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { // Observe actual SQL lock contention, not an assumed timing delay. for deadline := time.Now().Add(3 * time.Second); ; { var blocked bool - if err := s.pool.QueryRow(t.Context(), "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock' AND query ILIKE '%session_devices%')").Scan(&blocked); err != nil { + if err := s.pool.QueryRow(t.Context(), "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE datname=current_database() AND wait_event_type='Lock' AND query ILIKE '%session_runtime_assignments%')").Scan(&blocked); err != nil { t.Fatal(err) } if blocked { @@ -281,7 +281,7 @@ func TestArchiveWaitingCancellationReceipts(t *testing.T) { } time.Sleep(time.Millisecond) } - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported}) draining, err := peer.DrainArchivedCancellation(t.Context()) if err != nil || !draining || peer.IsClosed() { t.Fatal("ACK lost drain before terminal commit", draining, err) diff --git a/services/core/tests/integration/assignment_release_test.go b/services/core/tests/integration/assignment_release_test.go new file mode 100644 index 000000000..b496ef0e3 --- /dev/null +++ b/services/core/tests/integration/assignment_release_test.go @@ -0,0 +1,66 @@ +package integration + +import ( + "context" + "net/http" + "net/url" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +// TestDeletionReleaseReachesReconnectedRuntime checks that a release recorded +// while the Runtime was offline reaches it after it reconnects, and that only +// its acknowledgement records the release as applied. +func TestDeletionReleaseReachesReconnectedRuntime(t *testing.T) { + h := newDispatchHarness(t) + h.conn.Close() + awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "Runtime disconnected", func() bool { + _, err := h.registry.LookupDevice(h.device.ID) + return err != nil + }) + if err := sessionService(t, h.s).DeleteSession(t.Context(), sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); err != nil { + t.Fatal(err) + } + worker := startWorker(t, t.Context(), h.s, h.d) + ctx, cancel := context.WithCancel(t.Context()) + done := make(chan error, 1) + go func() { done <- worker.Run(ctx) }() + t.Cleanup(func() { cancel(); <-done }) + + u, _ := url.Parse(strings.Replace(h.url, "http", "ws", 1) + "/api/v1/agent-daemon/ws") + u.RawQuery = url.Values{"device_id": {h.device.ID}, "version": {proto.Version}}.Encode() + conn, _, err := websocket.DefaultDialer.Dial(u.String(), http.Header{"Authorization": {"Bearer " + h.credential}}) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { conn.Close() }) + h.conn = conn + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilitySupported, SupportedAgentKinds: []proto.SupportedAgentKind{}}) + + release := h.read(proto.TypeAssignmentRelease) + var request proto.AssignmentReleasePayload + if release.DecodePayload(&request) != nil || release.Assignment.SessionID != h.session.ID || release.Assignment.Epoch != 2 || !request.RemoveHome { + t.Fatal("deletion did not release the assignment with home removal", release.Assignment, request) + } + applied := func() int64 { + var epoch int64 + if err := h.s.pool.QueryRow(t.Context(), "SELECT applied_epoch FROM session_runtime_assignments WHERE session_id=$1", h.session.ID).Scan(&epoch); err != nil { + t.Fatal(err) + } + return epoch + } + if applied() != 0 { + t.Fatal("release recorded before the Runtime acknowledged it") + } + reply, err := release.Reply(proto.TypeAssignmentStatus, proto.AssignmentStatusPayload{State: proto.AssignmentHomeRemoved}) + if err != nil || conn.WriteJSON(reply) != nil { + t.Fatal("cannot acknowledge the release", err) + } + awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "release acknowledged", func() bool { return applied() == 2 }) +} diff --git a/services/core/tests/integration/claude_execution_test.go b/services/core/tests/integration/claude_execution_test.go index 3c3d2ab68..25fa8035f 100644 --- a/services/core/tests/integration/claude_execution_test.go +++ b/services/core/tests/integration/claude_execution_test.go @@ -31,7 +31,7 @@ func claudeSession(t *testing.T, h *dispatchHarness, configuration string, prebo func claudeHeartbeat(t *testing.T, h *dispatchHarness, ready bool) { t.Helper() caps := prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilityFromBool(ready), ExecutionControls: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, FunctionTools: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, Preparation: proto.CapabilitySupported}) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) deadline := time.Now().Add(3 * time.Second) for { peer, _ := h.registry.LookupDevice(h.device.ID) diff --git a/services/core/tests/integration/claude_mcp_test.go b/services/core/tests/integration/claude_mcp_test.go index 32a18777a..feeea8e5c 100644 --- a/services/core/tests/integration/claude_mcp_test.go +++ b/services/core/tests/integration/claude_mcp_test.go @@ -29,7 +29,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) { claudeSession(t, h, configuration, prebound) // Base MCP support does not imply authentication or required initialization. caps := prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, MCPHTTPTools: proto.CapabilityFromBool(authenticated || required), Preparation: proto.CapabilitySupported}) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "Claude MCP heartbeat", func() bool { peer, _ := h.registry.LookupDevice(h.device.ID) info, found, known := peer.AgentKindStatus("claude_sdk") @@ -66,7 +66,7 @@ func TestClaudeMCPWaitsForCapableRuntime(t *testing.T) { } caps.MCPHTTPTools, caps.MCPHTTPBearerAuth = proto.CapabilitySupported, proto.CapabilityFromBool(authenticated) caps.MCPHTTPRequired = proto.CapabilityFromBool(required) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) var prompt proto.PromptRequestPayload if h.read(testExecutionRequest).DecodePayload(&prompt) != nil || prompt.AgentKind != "claude_sdk" || prompt.MCPHTTPServers == nil || len(*prompt.MCPHTTPServers) != 1 { t.Fatal("missing typed MCP dispatch") @@ -105,7 +105,7 @@ func TestClaudeMCPUnsupportedSnapshotRejectedBeforeClaim(t *testing.T) { if profile == "missing required capability" { caps.MCPHTTPRequired = proto.CapabilityUnsupported } - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "claude_sdk", Available: true, Capabilities: caps}}}) deadline := time.Now().Add(3 * time.Second) for { peer, _ := h.registry.LookupDevice(h.device.ID) diff --git a/services/core/tests/integration/deployment_model_providers_http_test.go b/services/core/tests/integration/deployment_model_providers_http_test.go index 059160423..4ee6290df 100644 --- a/services/core/tests/integration/deployment_model_providers_http_test.go +++ b/services/core/tests/integration/deployment_model_providers_http_test.go @@ -292,7 +292,7 @@ func TestLegacySessionWithoutProviderCannotStartWork(t *testing.T) { if executor.conn.ReadJSON(&frame) != nil { break } - if frame.Type == proto.TypeExecutionPrepare || frame.Type == proto.TypePromptRequest { + if frame.Type == proto.TypeExecutionPrepare { t.Fatal("provider-free work reached the executor", frame.Type) } } diff --git a/services/core/tests/integration/devices_test.go b/services/core/tests/integration/devices_test.go index 385c4571f..270e48ae7 100644 --- a/services/core/tests/integration/devices_test.go +++ b/services/core/tests/integration/devices_test.go @@ -187,7 +187,7 @@ func TestStandaloneGatewayUsesExecutionCredentials(t *testing.T) { if err := sessionService(t, s).RevokeDevice(ctx, tenant, a.ID); err != nil { t.Fatal(err) } - if err := second.WriteJSON(map[string]any{"type": proto.TypeHeartbeat, "payload": map[string]any{"version": "test"}}); err != nil { + if err := second.WriteJSON(map[string]any{"type": proto.TypeHeartbeat, "payload": map[string]any{"version": "test", "home_removal": false}}); err != nil { t.Fatal(err) } _ = second.SetReadDeadline(time.Now().Add(2 * time.Second)) diff --git a/services/core/tests/integration/dispatch_test.go b/services/core/tests/integration/dispatch_test.go index 0f340ad9c..f675cbcb2 100644 --- a/services/core/tests/integration/dispatch_test.go +++ b/services/core/tests/integration/dispatch_test.go @@ -100,7 +100,7 @@ func newDispatchHarnessForSession(t *testing.T, configuration []byte, local bool t.Fatal("device connection failed") } t.Cleanup(func() { h.conn.Close() }) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, SubagentObservations: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, NativeSessionRecovery: proto.CapabilitySupported, Preparation: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported})}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, SubagentObservations: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, NativeSessionRecovery: proto.CapabilitySupported, Preparation: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported})}}}) deadline := time.Now().Add(3 * time.Second) for { peer, e := h.registry.LookupDevice(h.device.ID) @@ -158,6 +158,9 @@ func (h *dispatchHarness) read(kind string) proto.Envelope { if err := h.conn.ReadJSON(&env); err != nil { h.t.Fatal(err) } + if kind != proto.TypeAssignmentBind && h.assignmentFrame(env) { + continue + } if kind == testExecutionRequest { var keep bool env, keep = h.executionFrame(env) @@ -430,7 +433,7 @@ func TestExecutionRejectsRuntimeMissingCapabilityBeforeClaim(t *testing.T) { missing := tc.missing t.Run(missing, func(t *testing.T) { h := newDispatchHarness(t) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilityFromBool(missing != "durable_turns"), DurableInputReceipts: proto.CapabilityFromBool(missing != "durable_input_receipts"), Preparation: proto.CapabilityFromBool(missing != "preparation"), WebSearchControl: proto.CapabilityFromBool(missing != "web_search_control"), TextVerbosity: proto.CapabilityFromBool(missing != "text_verbosity"), ExecutionControls: proto.CapabilityFromBool(missing != "execution_controls"), SubagentControl: proto.CapabilityFromBool(missing != "subagent_control"), ToolObservations: proto.CapabilityFromBool(missing != "tool_observations"), EnvironmentNone: proto.CapabilityFromBool(missing != "environment_none")})}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilityFromBool(missing != "durable_turns"), DurableInputReceipts: proto.CapabilityFromBool(missing != "durable_input_receipts"), Preparation: proto.CapabilityFromBool(missing != "preparation"), WebSearchControl: proto.CapabilityFromBool(missing != "web_search_control"), TextVerbosity: proto.CapabilityFromBool(missing != "text_verbosity"), ExecutionControls: proto.CapabilityFromBool(missing != "execution_controls"), SubagentControl: proto.CapabilityFromBool(missing != "subagent_control"), ToolObservations: proto.CapabilityFromBool(missing != "tool_observations"), EnvironmentNone: proto.CapabilityFromBool(missing != "environment_none")})}}}) deadline := time.Now().Add(3 * time.Second) for { peer, _ := h.registry.LookupDevice(h.device.ID) diff --git a/services/core/tests/integration/environment_directory_test.go b/services/core/tests/integration/environment_directory_test.go index 815203094..deb43aeda 100644 --- a/services/core/tests/integration/environment_directory_test.go +++ b/services/core/tests/integration/environment_directory_test.go @@ -25,7 +25,7 @@ func directoryWorker(t *testing.T) (*dispatchHarness, *execution.Worker, session if err != nil { t.Fatal(err) } - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{LocalEnvironment: proto.CapabilitySupported, Preparation: proto.CapabilitySupported, WorkspaceReadPreparation: proto.CapabilitySupported})}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{LocalEnvironment: proto.CapabilitySupported, Preparation: proto.CapabilitySupported, WorkspaceReadPreparation: proto.CapabilitySupported})}}}) peer, err := h.registry.LookupDevice(h.device.ID) if err != nil { t.Fatal(err) diff --git a/services/core/tests/integration/environment_expiry_dispatch_test.go b/services/core/tests/integration/environment_expiry_dispatch_test.go index 623fe0ea5..c96c13f73 100644 --- a/services/core/tests/integration/environment_expiry_dispatch_test.go +++ b/services/core/tests/integration/environment_expiry_dispatch_test.go @@ -10,7 +10,7 @@ import ( ) func enableEnvironmentExpiryDispatch(h *dispatchHarness) { - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{ + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{ Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) } diff --git a/services/core/tests/integration/environment_runtime_fixture_test.go b/services/core/tests/integration/environment_runtime_fixture_test.go index 4a047b13a..f2703e39d 100644 --- a/services/core/tests/integration/environment_runtime_fixture_test.go +++ b/services/core/tests/integration/environment_runtime_fixture_test.go @@ -111,7 +111,7 @@ func assertNoRuntimeAllocation(t *testing.T, h *dispatchHarness) { func awaitFixtureCapabilities(t *testing.T, h *dispatchHarness, caps proto.AgentKindCapabilities) { t.Helper() - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: caps}}}) awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "updated Runtime capabilities", func() bool { peer, err := h.registry.LookupDevice(h.device.ID) if err != nil { diff --git a/services/core/tests/integration/environment_work_test.go b/services/core/tests/integration/environment_work_test.go index fae8b6abc..3f6ee16cd 100644 --- a/services/core/tests/integration/environment_work_test.go +++ b/services/core/tests/integration/environment_work_test.go @@ -34,7 +34,7 @@ func TestEnvironmentInputWorkFiltersAndPagesDevices(t *testing.T) { } case "unbound": wanted[pending.ID] = pending.SessionID - if _, err := pool.Exec(t.Context(), "DELETE FROM session_devices WHERE session_id=$1", pending.SessionID); err != nil { + if _, err := pool.Exec(t.Context(), "DELETE FROM session_runtime_assignments WHERE session_id=$1", pending.SessionID); err != nil { t.Fatal(err) } default: diff --git a/services/core/tests/integration/environment_worker_helpers_test.go b/services/core/tests/integration/environment_worker_helpers_test.go index d587a3ca4..86c34df7a 100644 --- a/services/core/tests/integration/environment_worker_helpers_test.go +++ b/services/core/tests/integration/environment_worker_helpers_test.go @@ -15,7 +15,7 @@ import ( func enableWorkerEnvironment(t *testing.T, h *dispatchHarness) { t.Helper() - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: workerEnvironmentCapabilities()}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: workerEnvironmentCapabilities()}}}) awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "worker preparation capability", func() bool { peer, err := h.registry.LookupDevice(h.device.ID) if err != nil { @@ -64,6 +64,9 @@ func workerFrames(t *testing.T, runtimes ...*dispatchHarness) <-chan proto.Envel if h.conn.ReadJSON(&env) != nil { return } + if h.assignmentFrame(env) { + continue + } var keep bool env, keep = h.executionFrame(env) if !keep { diff --git a/services/core/tests/integration/execution_messages_test.go b/services/core/tests/integration/execution_messages_test.go index 860192ca2..22ffa33b5 100644 --- a/services/core/tests/integration/execution_messages_test.go +++ b/services/core/tests/integration/execution_messages_test.go @@ -24,7 +24,7 @@ func TestExecutionNegotiatesAndPersistsMessageObservations(t *testing.T) { } h.write(first.TurnID, proto.TypeDone, proto.DonePayload{}) h.finished(result, sessions.TurnCompleted) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, MessageItems: proto.CapabilitySupported, NativeSessionRecovery: proto.CapabilitySupported, Preparation: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported})}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, MessageItems: proto.CapabilitySupported, NativeSessionRecovery: proto.CapabilitySupported, Preparation: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported})}}}) deadline := time.Now().Add(3 * time.Second) for { peer, err := h.registry.LookupDevice(h.device.ID) diff --git a/services/core/tests/integration/executor_fixture_test.go b/services/core/tests/integration/executor_fixture_test.go index 36008c6e6..3e75c0e25 100644 --- a/services/core/tests/integration/executor_fixture_test.go +++ b/services/core/tests/integration/executor_fixture_test.go @@ -16,6 +16,30 @@ type fixtureAdmission struct { executor string } +// assignmentReply is the reply of a Runtime that binds every assignment to +// env, when env is an assignment_bind. +func assignmentReply(env proto.Envelope) (proto.Envelope, bool) { + if env.Type != proto.TypeAssignmentBind { + return proto.Envelope{}, false + } + reply, err := env.Reply(proto.TypeAssignmentStatus, proto.AssignmentStatusPayload{State: proto.AssignmentBound}) + return reply, err == nil +} + +// assignmentFrame answers an assignment_bind and reports whether env was one. +func (h *dispatchHarness) assignmentFrame(env proto.Envelope) bool { + reply, ok := assignmentReply(env) + if !ok { + return false + } + h.writeMu.Lock() + defer h.writeMu.Unlock() + if err := h.conn.WriteJSON(reply); err != nil { + h.t.Fatal(err) + } + return true +} + func (h *dispatchHarness) executionFrame(env proto.Envelope) (proto.Envelope, bool) { if h.admissions == nil { h.admissions = make(map[string]fixtureAdmission) diff --git a/services/core/tests/integration/function_execution_test.go b/services/core/tests/integration/function_execution_test.go index 73a8ed394..76367912f 100644 --- a/services/core/tests/integration/function_execution_test.go +++ b/services/core/tests/integration/function_execution_test.go @@ -27,7 +27,7 @@ func newFunctionHarness(t *testing.T) *dispatchHarness { if err := bindSessionDevice(t, h.s, h.tenant, h.session.ID, h.device.ID); err != nil { t.Fatal(err) } - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, FunctionTools: proto.CapabilitySupported, FunctionResultImages: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, Resume: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, FunctionTools: proto.CapabilitySupported, FunctionResultImages: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) deadline := time.Now().Add(time.Second) for { peer, _ := h.registry.LookupDevice(h.device.ID) diff --git a/services/core/tests/integration/function_input_execution_test.go b/services/core/tests/integration/function_input_execution_test.go index 63dbe8d17..aea59abac 100644 --- a/services/core/tests/integration/function_input_execution_test.go +++ b/services/core/tests/integration/function_input_execution_test.go @@ -33,6 +33,9 @@ func TestExecutionFunctionInputBatchStillSteersMessages(t *testing.T) { if err := h.conn.ReadJSON(&env); err != nil { t.Fatal(err) } + if h.assignmentFrame(env) { + continue + } switch env.Type { case proto.TypeFunctionResult: var result proto.FunctionResultPayload diff --git a/services/core/tests/integration/function_worker_test.go b/services/core/tests/integration/function_worker_test.go index 71582ac3d..c5d28736d 100644 --- a/services/core/tests/integration/function_worker_test.go +++ b/services/core/tests/integration/function_worker_test.go @@ -48,7 +48,7 @@ func TestWorkerWaitsForToolCapabilities(t *testing.T) { } caps := prototest.Capabilities(proto.AgentKindCapabilities{Preparation: proto.CapabilityFromBool(missing != "preparation"), Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilityFromBool(missing != "durable_input_receipts"), EnvironmentNone: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilityFromBool(missing != "execution_controls"), SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilityFromBool(missing != "tool_observations"), MCPHTTPTools: proto.CapabilityFromBool(missing != "mcp_http_tools"), MCPHTTPRequired: proto.CapabilityFromBool(missing != "mcp_http_required"), MCPHTTPBearerAuth: proto.CapabilityFromBool(missing != "mcp_http_bearer_auth"), FunctionTools: proto.CapabilityFromBool(missing != "function_tools" && !isMCP)}) heartbeat := func() { - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: caps}}}) } heartbeat() deadline := time.Now().Add(3 * time.Second) diff --git a/services/core/tests/integration/harness_onboarding_test.go b/services/core/tests/integration/harness_onboarding_test.go index 6272a3f57..9f120e8cf 100644 --- a/services/core/tests/integration/harness_onboarding_test.go +++ b/services/core/tests/integration/harness_onboarding_test.go @@ -128,7 +128,7 @@ func TestThirdHarnessPublicOnboarding(t *testing.T) { changed := declaration changed.Capabilities.DurableInputReceipts = proto.CapabilityUnsupported // A separate unbound Session is used, without changing public handler behavior. - update, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{changed}}) + update, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{changed}}) if err := write(update); err != nil { t.Fatal(err) } diff --git a/services/core/tests/integration/local_environment_devices_test.go b/services/core/tests/integration/local_environment_devices_test.go index 88a04c240..aa479f8b7 100644 --- a/services/core/tests/integration/local_environment_devices_test.go +++ b/services/core/tests/integration/local_environment_devices_test.go @@ -53,7 +53,7 @@ func TestEnvironmentDeviceAuthorityAndLifecycle(t *testing.T) { } reopened, _ := testStore(t) got, err := sessionAdapter(reopened).GetSessionDevice(t.Context(), tenant, session.ID) - if err != nil || got != bound { + if err != nil || got.ID != bound.ID || got.EnvironmentID != bound.EnvironmentID || got.Assignment.SessionID != session.ID || got.Assignment.Epoch != 1 || !got.Assignment.Valid() { t.Fatalf("durable exact binding: %+v %v", got, err) } if _, ok, err := sessionAdapter(s).GetDeviceCredential(t.Context(), bound.ID); err != nil || !ok { diff --git a/services/core/tests/integration/local_environment_file_write_test.go b/services/core/tests/integration/local_environment_file_write_test.go index 0b802a30f..57d1d5e26 100644 --- a/services/core/tests/integration/local_environment_file_write_test.go +++ b/services/core/tests/integration/local_environment_file_write_test.go @@ -47,9 +47,16 @@ func TestLocalEnvironmentFileWriteOwnsMutationBeforeDispatch(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) defer cancel() done := startLocalWrite(ctx, w, environment) + // Environment work without a Turn binds the Session's assignment first. + bind := h.read(proto.TypeAssignmentBind) + var binding proto.AssignmentBindPayload + if bind.DecodePayload(&binding) != nil || bind.Assignment.SessionID != h.session.ID || bind.Assignment.Epoch != 1 || binding.EnvironmentID != environment.ID { + t.Fatal("upload did not bind the Session's assignment") + } + h.assignmentFrame(bind) begin := h.read(proto.TypeWorkspaceWrite) var request proto.WorkspaceWritePayload - if begin.DecodePayload(&request) != nil || request.Step != "begin" || request.EnvironmentID != environment.ID || request.SessionID != h.session.ID || request.Path != "input" || request.SizeBytes != 3 { + if begin.Assignment != bind.Assignment || begin.DecodePayload(&request) != nil || request.Step != "begin" || request.EnvironmentID != environment.ID || request.SessionID != h.session.ID || request.Path != "input" || request.SizeBytes != 3 { t.Fatal("upload identity changed") } intent, err := FixtureFileWrite(t.Context(), h.s.pool, h.tenant, environment.ID, begin.ID) diff --git a/services/core/tests/integration/local_environment_worker_test.go b/services/core/tests/integration/local_environment_worker_test.go index 494efc624..f26451725 100644 --- a/services/core/tests/integration/local_environment_worker_test.go +++ b/services/core/tests/integration/local_environment_worker_test.go @@ -32,7 +32,7 @@ func localWorker(t *testing.T, scoped, execute bool) (*dispatchHarness, *executi caps.WebSearchControl, caps.TextVerbosity, caps.ExecutionControls = proto.CapabilitySupported, proto.CapabilitySupported, proto.CapabilitySupported caps.SubagentControl, caps.ToolObservations = proto.CapabilitySupported, proto.CapabilitySupported } - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: caps}}}) awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "local capability", func() bool { peer, err := h.registry.LookupDevice(h.device.ID) if err != nil { @@ -120,6 +120,7 @@ func TestLocalEnvironmentWorkerSchedulesPreparationWithoutRemoteResolver(t *test if h.conn.ReadJSON(&frame) != nil { t.Fatal("local reservation was not scheduled") } + h.assignmentFrame(frame) } var prepare proto.ExecutionPreparePayload if frame.DecodePayload(&prepare) != nil || prepare.Configuration.LocalEnvironment == nil || prepare.Configuration.LocalEnvironment.ID != environment.ID { diff --git a/services/core/tests/integration/public_execution_test.go b/services/core/tests/integration/public_execution_test.go index a51476e35..92ea75caf 100644 --- a/services/core/tests/integration/public_execution_test.go +++ b/services/core/tests/integration/public_execution_test.go @@ -23,7 +23,7 @@ func publicSession(t *testing.T, h *dispatchHarness, key string) sessions.Sessio func TestExecutionWorkerAdmissionBindingAndRecovery(t *testing.T) { h := newDispatchHarness(t) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) h.session = publicSession(t, h, "public") ctx, cancel := context.WithCancel(context.Background()) defer cancel() diff --git a/services/core/tests/integration/runtime_compute_lifecycle_test.go b/services/core/tests/integration/runtime_compute_lifecycle_test.go index c6949a6de..f24a56fdb 100644 --- a/services/core/tests/integration/runtime_compute_lifecycle_test.go +++ b/services/core/tests/integration/runtime_compute_lifecycle_test.go @@ -185,7 +185,7 @@ func (p *fakeCheckpointProvider) connect(ctx context.Context, b sandbox.Bootstra if err != nil { return err } - heartbeat, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}}}) + heartbeat, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}}}) if err := conn.WriteJSON(heartbeat); err != nil { conn.Close() return err @@ -205,6 +205,12 @@ func (p *fakeCheckpointProvider) connect(ctx context.Context, b sandbox.Bootstra if conn.ReadJSON(&env) != nil { return } + if reply, ok := assignmentReply(env); ok { + if conn.WriteJSON(reply) != nil { + return + } + continue + } if env.Type == proto.TypeRuntimePrepare && p.preparation != nil { reply, err := transfer.receive(env) if err != nil { diff --git a/services/core/tests/integration/runtime_initialization_peer_test.go b/services/core/tests/integration/runtime_initialization_peer_test.go index e67d099db..9fed3cc5b 100644 --- a/services/core/tests/integration/runtime_initialization_peer_test.go +++ b/services/core/tests/integration/runtime_initialization_peer_test.go @@ -44,7 +44,7 @@ func (p *initializationPeer) connect(b sandbox.Bootstrap) error { return err } p.t.Cleanup(func() { _ = c.Close() }) - heartbeat, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: !p.unavailable, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}}}) + heartbeat, _ := proto.NewEnvelope(proto.TypeHeartbeat, "", proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: !p.unavailable, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{})}}}) if err := c.WriteJSON(heartbeat); err != nil { return err } @@ -56,6 +56,12 @@ func (p *initializationPeer) connect(b sandbox.Bootstrap) error { if c.ReadJSON(&env) != nil { return } + if reply, ok := assignmentReply(env); ok { + if c.WriteJSON(reply) != nil { + return + } + continue + } if env.Type != proto.TypeRuntimePrepare { continue } diff --git a/services/core/tests/integration/structured_output_dispatch_test.go b/services/core/tests/integration/structured_output_dispatch_test.go index 2ae4b27be..0b32c1194 100644 --- a/services/core/tests/integration/structured_output_dispatch_test.go +++ b/services/core/tests/integration/structured_output_dispatch_test.go @@ -31,7 +31,7 @@ func TestStructuredOutputDispatchRechecksOperationQualification(t *testing.T) { t.Fatal(err) } caps := prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, StructuredOutput: proto.CapabilitySupported, MessageItems: proto.CapabilitySupported, Preparation: proto.CapabilitySupported}) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "fixture_harness", Available: true, Capabilities: caps}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "fixture_harness", Available: true, Capabilities: caps}}}) peer, _ := h.registry.LookupDevice(h.device.ID) for deadline := time.Now().Add(3 * time.Second); ; { info, found, known := peer.AgentKindStatus("fixture_harness") diff --git a/services/core/tests/integration/subagent_identities_test.go b/services/core/tests/integration/subagent_identities_test.go index 4a7fccf6d..f84c309fa 100644 --- a/services/core/tests/integration/subagent_identities_test.go +++ b/services/core/tests/integration/subagent_identities_test.go @@ -97,7 +97,7 @@ func TestSubagentIdentityIsAtomicScopedAndImmutable(t *testing.T) { if err = journal.AppendTurnEvents(ctx, tenant, foreign.ID, foreignInput.TurnID, 1, []sessions.ExecutionEvent{a}); !errors.Is(err, sessions.ErrIdempotencyConflict) { t.Fatal("same device/native child reassigned to another Session", err) } - if _, err = pool.Exec(ctx, "UPDATE session_devices SET native_session_id='known-root' WHERE session_id=$1", foreign.ID); err != nil { + if _, err = pool.Exec(ctx, "UPDATE session_runtime_assignments SET native_session_id='known-root' WHERE session_id=$1", foreign.ID); err != nil { t.Fatal(err) } if err = journal.AppendTurnEvents(ctx, tenant, foreign.ID, foreignInput.TurnID, 1, []sessions.ExecutionEvent{subagentIdentityEvent("other-child", "root", 101)}); !errors.Is(err, sessions.ErrIdempotencyConflict) { diff --git a/services/core/tests/integration/worker_capacity_test.go b/services/core/tests/integration/worker_capacity_test.go index 0f11cef48..957cd8ccd 100644 --- a/services/core/tests/integration/worker_capacity_test.go +++ b/services/core/tests/integration/worker_capacity_test.go @@ -21,6 +21,9 @@ func capacityWorkerFrames(t *testing.T, h *dispatchHarness) <-chan proto.Envelop if h.conn.ReadJSON(&frame) != nil { return } + if h.assignmentFrame(frame) { + continue + } select { case frames <- frame: case <-t.Context().Done(): diff --git a/services/core/tests/integration/worker_lease_loss_test.go b/services/core/tests/integration/worker_lease_loss_test.go index d2b75293f..45b5d0950 100644 --- a/services/core/tests/integration/worker_lease_loss_test.go +++ b/services/core/tests/integration/worker_lease_loss_test.go @@ -14,7 +14,7 @@ import ( func TestWorkerLeaseLossLeavesUncertainWorkForSuccessor(t *testing.T) { h := newDispatchHarness(t) _, pool := testStore(t) - h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) + h.write("", proto.TypeHeartbeat, proto.HeartbeatPayload{HomeRemoval: proto.CapabilityUnsupported, SupportedAgentKinds: []proto.SupportedAgentKind{{Kind: "codex", Available: true, Capabilities: prototest.Capabilities(proto.AgentKindCapabilities{Streaming: proto.CapabilitySupported, Steering: proto.CapabilitySupported, DurableTurns: proto.CapabilitySupported, DurableInputReceipts: proto.CapabilitySupported, WebSearchControl: proto.CapabilitySupported, TextVerbosity: proto.CapabilitySupported, ExecutionControls: proto.CapabilitySupported, SubagentControl: proto.CapabilitySupported, ToolObservations: proto.CapabilitySupported, EnvironmentNone: proto.CapabilitySupported, Preparation: proto.CapabilitySupported})}}}) h.session = publicSession(t, h, "active") queued := publicSession(t, h, "queued") worker := startWorker(t, t.Context(), h.s, h.d) From 0b75b12415710e54c93f3883219d3d1cf7e19db9 Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 11:50:35 +0000 Subject: [PATCH 2/4] Fence every Session operation on release and check echoed assignments Release settles a Session's non-Executor work before it replies; cached interaction receipts and Resume answer only their own assignment. Core checks each reply's assignment in the common exchange and subscriptions, accepts assignment rejections in workspace receipts, settles releases no Runtime can act on, and backs off failing releases. --- apps/daemon/internal/dispatch/assignment.go | 49 ++++++++++++- apps/daemon/internal/dispatch/functions.go | 2 +- .../internal/dispatch/functions_test.go | 11 +++ .../dispatch/interaction_decisions.go | 14 ++-- apps/daemon/internal/dispatch/router.go | 4 +- .../internal/dispatch/runtime_preparation.go | 12 +++- .../dispatch/runtime_preparation_test.go | 6 +- apps/daemon/internal/dispatch/suspend.go | 11 ++- apps/daemon/internal/dispatch/suspend_test.go | 8 ++- .../internal/dispatch/workspace_export.go | 5 +- .../internal/dispatch/workspace_read.go | 12 +++- .../internal/dispatch/workspace_write.go | 12 +++- services/core/internal/db/queries/devices.sql | 21 ++++-- services/core/internal/db/sqlc/devices.sql.go | 21 ++++-- .../internal/execution/assignment_releases.go | 71 +++++++++++++++---- .../execution/assignment_releases_test.go | 30 ++++++++ services/core/internal/execution/delivery.go | 2 +- .../core/internal/execution/preparation.go | 2 +- .../postgres/deploymentpg/allocations.go | 6 +- .../postgres/deploymentpg/session_archive.go | 5 +- .../postgres/sessionpg/session_writes.go | 10 +-- .../internal/runtimegateway/assignment.go | 4 +- .../runtimegateway/cancellation_test.go | 2 +- .../internal/runtimegateway/chunk_exchange.go | 26 ------- .../core/internal/runtimegateway/exchange.go | 44 ++++++++++++ .../internal/runtimegateway/exchange_test.go | 52 ++++++++++++++ .../mcp_bearer_live_linux_test.go | 4 +- .../internal/runtimegateway/preparation.go | 12 ++-- .../runtimegateway/preparation_test.go | 4 +- .../runtimegateway/runtime_prepare.go | 2 +- .../runtimegateway/runtime_prepare_test.go | 1 + .../core/internal/runtimegateway/session.go | 37 ++++------ .../internal/runtimegateway/session_test.go | 10 +-- .../internal/runtimegateway/subscription.go | 34 ++++++--- .../runtimegateway/subscription_test.go | 6 +- .../internal/runtimegateway/suspension.go | 31 ++------ .../runtimegateway/suspension_test.go | 4 +- .../core/internal/runtimegateway/wire_test.go | 4 +- .../workspace_directory_test.go | 4 +- .../runtimegateway/workspace_export.go | 41 +++++------ .../runtimegateway/workspace_export_test.go | 1 + .../internal/runtimegateway/workspace_read.go | 21 ++---- .../runtimegateway/workspace_read_test.go | 10 +-- .../runtimegateway/workspace_write.go | 4 +- .../runtimegateway/workspace_write_test.go | 7 ++ .../000092_session_runtime_assignments.sql | 6 +- .../integration/assignment_release_test.go | 41 +++++++++++ .../core/tests/integration/dispatch_test.go | 5 ++ .../integration/executor_fixture_test.go | 14 +++- .../runtime_compute_lifecycle_test.go | 2 +- .../runtime_initialization_peer_test.go | 2 +- 51 files changed, 525 insertions(+), 224 deletions(-) create mode 100644 services/core/internal/execution/assignment_releases_test.go delete mode 100644 services/core/internal/runtimegateway/chunk_exchange.go create mode 100644 services/core/internal/runtimegateway/exchange.go create mode 100644 services/core/internal/runtimegateway/exchange_test.go diff --git a/apps/daemon/internal/dispatch/assignment.go b/apps/daemon/internal/dispatch/assignment.go index 3e71dc1c4..1669dddcd 100644 --- a/apps/daemon/internal/dispatch/assignment.go +++ b/apps/daemon/internal/dispatch/assignment.go @@ -61,8 +61,8 @@ func (r *Router) handleAssignmentBind(ctx context.Context, env proto.Envelope) e } // handleAssignmentRelease fences the assignment, then settles the Session's -// Executor and removes its home before it replies. A retry at the same epoch -// repeats the cleanup. +// work and Executor and removes its home before it replies. A retry at the +// same epoch repeats the cleanup. func (r *Router) handleAssignmentRelease(ctx context.Context, env proto.Envelope) error { var input proto.AssignmentReleasePayload ref, code := env.Assignment, "" @@ -90,10 +90,17 @@ func (r *Router) handleAssignmentRelease(ctx context.Context, env proto.Envelope r.mu.Unlock() return r.reply(ctx, env, proto.TypeAssignmentStatus, assignmentStatus("", code)) } + work, preparations := r.fenceSessionWorkLocked(ref.SessionID) r.shutdownWG.Add(1) r.mu.Unlock() go func() { defer r.shutdownWG.Done() + for _, p := range preparations { + r.releasePreparation(p, "failed", proto.AssignmentStale, true, false) + } + for _, done := range work { + <-done + } state, code := proto.AssignmentReleased, "" err := r.closeSessionExecutor(ref.SessionID) if err == nil && input.RemoveHome { @@ -110,6 +117,44 @@ func (r *Router) handleAssignmentRelease(ctx context.Context, env proto.Envelope return nil } +// fenceSessionWorkLocked ends the Session's admitted work outside its +// Executor. A transfer still receiving its body ends without applying it, and +// a transfer that already committed checks the released assignment before it +// applies. It returns the work the release waits for and the read-only +// preparations it releases, whose release also cancels their exports. +// Router.mu must be held. +func (r *Router) fenceSessionWorkLocked(sessionID string) ([]chan struct{}, []*preparationState) { + var work []chan struct{} + if u := r.workspaceWrite; u != nil && u.envelope.Assignment.SessionID == sessionID { + if !u.finished { + u.finished = true + close(u.ready) + } + work = append(work, u.done) + } + if u := r.runtimePreparation; u != nil && u.envelope.Assignment.SessionID == sessionID { + if !u.finished { + r.finishRuntimePreparationTransferLocked(u, false) + } + work = append(work, u.done) + } + if u := r.workspaceExport; u != nil && u.request.Assignment.SessionID == sessionID { + work = append(work, u.done) + } + for _, read := range r.workspaceReads { + if read.sessionID == sessionID { + work = append(work, read.done) + } + } + var preparations []*preparationState + for _, p := range r.preparations { + if p.executor == nil && p.workspaceReadOnly && p.owns && p.request.Assignment.SessionID == sessionID { + preparations = append(preparations, p) + } + } + return work, preparations +} + // closeSessionExecutor ends the Session's Executor: it abandons a pending // admission, releases a Turn, waits for a native preparation and closes the // Executor. The released assignment admits no new preparation meanwhile. diff --git a/apps/daemon/internal/dispatch/functions.go b/apps/daemon/internal/dispatch/functions.go index b19ec7282..6aed14c06 100644 --- a/apps/daemon/internal/dispatch/functions.go +++ b/apps/daemon/internal/dispatch/functions.go @@ -75,6 +75,6 @@ func (r *Router) handleFunctionResult(ctx context.Context, env proto.Envelope) e } return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, false, code, "function result was not applied") } - r.rememberAppliedInteractionDecision(env.ID, kind, fingerprint) + r.rememberAppliedInteractionDecision(env, kind, fingerprint) return r.sendInteractionDecisionAck(ctx, env, result.DeliveryID, true, "", "") } diff --git a/apps/daemon/internal/dispatch/functions_test.go b/apps/daemon/internal/dispatch/functions_test.go index e668436b8..f44e8cac0 100644 --- a/apps/daemon/internal/dispatch/functions_test.go +++ b/apps/daemon/internal/dispatch/functions_test.go @@ -95,6 +95,17 @@ func TestFunctionReceiptsScopeRetriesAndConflicts(t *testing.T) { if a := submit("one", "call", "changed", "conflict"); a.Applied || a.ErrorCode != "decision_conflict" { t.Fatal(a) } + // The cached receipt answers only the assignment that applied the result. + foreign := scoped(t, "one", proto.TypeFunctionResult, "one", proto.FunctionResultPayload{CallID: "call", Success: true, Content: functionResultContent("answer"), DeliveryID: "foreign"}) + foreign.Assignment.AssignmentID = "foreign" + if err := router.Handle(t.Context(), foreign); err != nil { + t.Fatal(err) + } + frames = sender.snapshot() + var foreignAck proto.InteractionDecisionAckPayload + if last := frames[len(frames)-1]; last.DecodePayload(&foreignAck) != nil || foreignAck.Applied || foreignAck.ErrorCode != proto.AssignmentConflict || last.Assignment != foreign.Assignment { + t.Fatal(last, foreignAck) + } for _, mutation := range []string{"image", "order", "success"} { result := proto.FunctionResultPayload{CallID: "call", Success: true, Content: functionResultContent("answer"), DeliveryID: mutation} diff --git a/apps/daemon/internal/dispatch/interaction_decisions.go b/apps/daemon/internal/dispatch/interaction_decisions.go index 0376d915c..1e296e324 100644 --- a/apps/daemon/internal/dispatch/interaction_decisions.go +++ b/apps/daemon/internal/dispatch/interaction_decisions.go @@ -83,7 +83,7 @@ func (r *Router) handlePermissionDecision(ctx context.Context, env proto.Envelop return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "runtime_error", err.Error()) } r.dropPermission(state, env.ID) - r.rememberAppliedInteractionDecision(env.ID, proto.TypePermissionDecision, fingerprint) + r.rememberAppliedInteractionDecision(env, proto.TypePermissionDecision, fingerprint) return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, true, "", "") } @@ -179,7 +179,7 @@ func (r *Router) handlePromptForUserChoiceDecision(ctx context.Context, env prot return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, false, "runtime_error", err.Error()) } r.dropAsk(state, env.ID) - r.rememberAppliedInteractionDecision(env.ID, proto.TypePromptForUserChoiceDecision, fingerprint) + r.rememberAppliedInteractionDecision(env, proto.TypePromptForUserChoiceDecision, fingerprint) return r.sendInteractionDecisionAck(ctx, env, payload.DeliveryID, true, "", "") } @@ -191,15 +191,19 @@ func (r *Router) replayAppliedInteractionDecision(ctx context.Context, env proto if !ok { return false, nil } + // The receipt answers only the assignment that applied the decision. + if applied.assignment != env.Assignment { + return true, r.sendInteractionDecisionAck(ctx, env, deliveryID, false, proto.AssignmentConflict, "The decision was applied under another assignment.") + } if applied.requestID != env.ID || applied.kind != kind || applied.fingerprint != fingerprint { return true, r.sendInteractionDecisionAck(ctx, env, deliveryID, false, "decision_conflict", "request was already applied with a different decision") } return true, r.sendInteractionDecisionAck(ctx, env, deliveryID, true, "", "") } -func (r *Router) rememberAppliedInteractionDecision(requestID, kind string, fingerprint [32]byte) { +func (r *Router) rememberAppliedInteractionDecision(env proto.Envelope, kind string, fingerprint [32]byte) { now := time.Now().UTC() - key := appliedInteractionDecisionKey(requestID, kind) + key := appliedInteractionDecisionKey(env.ID, kind) r.mu.Lock() if len(r.applied) >= 1024 { cutoff := now.Add(-time.Hour) @@ -216,7 +220,7 @@ func (r *Router) rememberAppliedInteractionDecision(requestID, kind string, fing } } r.applied[key] = appliedInteractionDecision{ - requestID: requestID, kind: kind, fingerprint: fingerprint, recordedAt: now, + requestID: env.ID, kind: kind, fingerprint: fingerprint, assignment: env.Assignment, recordedAt: now, } r.mu.Unlock() } diff --git a/apps/daemon/internal/dispatch/router.go b/apps/daemon/internal/dispatch/router.go index bf67b6902..0255d2762 100644 --- a/apps/daemon/internal/dispatch/router.go +++ b/apps/daemon/internal/dispatch/router.go @@ -36,6 +36,7 @@ type Router struct { admission sync.RWMutex suspension *proto.EnvironmentSuspendPayload + suspendedBy proto.AssignmentRef // the assignment that quiesced mu sync.Mutex assignments map[string]*assignmentState // SessionID → assignment sessions map[string]*sessionState // RunID → state @@ -54,7 +55,7 @@ type Router struct { runtimePreparation *runtimePreparationTransfer workspaceWrite *workspaceUpload workspaceExport *workspaceExport - workspaceReads map[string]struct{} + workspaceReads map[string]workspaceRead localWorkspace *localworkspace.Binding sessionEnvironments bool removeHome func(sessionID string) error @@ -64,6 +65,7 @@ type appliedInteractionDecision struct { requestID string kind string fingerprint [32]byte + assignment proto.AssignmentRef recordedAt time.Time } diff --git a/apps/daemon/internal/dispatch/runtime_preparation.go b/apps/daemon/internal/dispatch/runtime_preparation.go index 42be3ba95..08e1201dc 100644 --- a/apps/daemon/internal/dispatch/runtime_preparation.go +++ b/apps/daemon/internal/dispatch/runtime_preparation.go @@ -17,11 +17,13 @@ const runtimePreparationTimeout = 120 * time.Second // Router.mu protects one connection-local transfer. Partial installation data // belongs to the bound Environment and is never removed by transfer cleanup. +// done closes once its result is sent. type runtimePreparationTransfer struct { envelope proto.Envelope request proto.RuntimePreparePayload data []byte ready chan struct{} + done chan struct{} cancel context.CancelFunc finished bool apply bool @@ -72,7 +74,7 @@ func (r *Router) handleRuntimePrepare(ctx context.Context, env proto.Envelope) e owner, cancel := context.WithTimeout(context.WithoutCancel(ctx), runtimePreparationTimeout) u := &runtimePreparationTransfer{ envelope: env, request: request, data: make([]byte, 0, request.SizeBytes), - ready: make(chan struct{}), cancel: cancel, + ready: make(chan struct{}), done: make(chan struct{}), cancel: cancel, } r.runtimePreparation = u r.shutdownWG.Add(1) @@ -138,6 +140,7 @@ func (r *Router) finishRuntimePreparationTransferLocked(u *runtimePreparationTra // shutdown, but cannot release ownership while that call is still running. func (r *Router) runRuntimePreparationTransfer(ctx context.Context, u *runtimePreparationTransfer, apply func(context.Context, proto.RuntimePreparePayload, []byte) error) { defer r.shutdownWG.Done() + defer close(u.done) defer u.cancel() select { case <-u.ready: @@ -145,12 +148,17 @@ func (r *Router) runRuntimePreparationTransfer(ctx context.Context, u *runtimePr case <-ctx.Done(): } r.mu.Lock() - admitted := u.apply && !r.closed && ctx.Err() == nil + // A release of the assignment before the preparation applies fences it. + fenced := r.admitLocked(u.envelope.Assignment, u.request.SessionID, u.request.EnvironmentID) + admitted := u.apply && !r.closed && ctx.Err() == nil && fenced == "" u.finished = true data := u.data u.data = nil r.mu.Unlock() result := rejectedRuntimePreparation("invalid_request") + if fenced != "" { + result = rejectedRuntimePreparation(fenced) + } if admitted { result = runtimePreparationResult(apply(ctx, u.request, data), u.request.SizeBytes) } diff --git a/apps/daemon/internal/dispatch/runtime_preparation_test.go b/apps/daemon/internal/dispatch/runtime_preparation_test.go index d1eadabda..1a98b5b4f 100644 --- a/apps/daemon/internal/dispatch/runtime_preparation_test.go +++ b/apps/daemon/internal/dispatch/runtime_preparation_test.go @@ -183,7 +183,7 @@ func TestRuntimePreparationPreparationExcludesOwnedResources(t *testing.T) { case "export": r.workspaceExport = &workspaceExport{} case "read": - r.workspaceReads = map[string]struct{}{"read": {}} + r.workspaceReads = map[string]workspaceRead{"read": {}} case "run": r.sessions["run"] = &sessionState{} case "executor": @@ -254,7 +254,7 @@ func TestRuntimePreparationCancellationKeepsOwnershipUntilApplyStops(t *testing. ctx, cancel := context.WithCancel(context.Background()) id := uuid.NewString() request := proto.RuntimePreparePayload{Step: "begin", Action: "finalize", EnvironmentID: environment, SessionID: session, Sources: &agentcapabilities.Input{}} - owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, ready: make(chan struct{}), cancel: cancel, finished: true, apply: true} + owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, ready: make(chan struct{}), done: make(chan struct{}), cancel: cancel, finished: true, apply: true} close(owner.ready) r.runtimePreparation = owner r.shutdownWG.Add(1) @@ -334,7 +334,7 @@ func TestRuntimePreparationResultCategoriesAndUnknownOwnership(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) id := uuid.NewString() request := capabilityBegin(environment, session, []byte("abc")) - owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, data: []byte("abc"), ready: make(chan struct{}), cancel: cancel, finished: true, apply: true} + owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, data: []byte("abc"), ready: make(chan struct{}), done: make(chan struct{}), cancel: cancel, finished: true, apply: true} close(owner.ready) r.runtimePreparation = owner r.shutdownWG.Add(1) diff --git a/apps/daemon/internal/dispatch/suspend.go b/apps/daemon/internal/dispatch/suspend.go index 077ca9b98..9220b7c6a 100644 --- a/apps/daemon/internal/dispatch/suspend.go +++ b/apps/daemon/internal/dispatch/suspend.go @@ -58,7 +58,7 @@ func (r *Router) Quiesce(ctx context.Context, ref proto.AssignmentRef, request p return ErrRouterBusy } } - r.suspension = &request + r.suspension, r.suspendedBy = &request, ref for _, p := range r.preparations { if p.timer != nil { p.timer.Stop() @@ -82,9 +82,8 @@ func (r *Router) Quiesce(ctx context.Context, ref proto.AssignmentRef, request p } // Resume opens admission only after the caller authenticated a new connection -// and Core confirmed the exact suspension identity on that connection. ref -// names the assignment, which a release during the suspension may have -// advanced. +// and Core confirmed the exact suspension identity on that connection under the +// assignment that quiesced. func (r *Router) Resume(ref proto.AssignmentRef, request proto.EnvironmentSuspendPayload, sender Sender) error { r.admission.Lock() defer r.admission.Unlock() @@ -96,11 +95,11 @@ func (r *Router) Resume(ref proto.AssignmentRef, request proto.EnvironmentSuspen if sender == nil || r.suspension == nil || !r.suspension.SameSuspension(request) { return errors.New("dispatch: suspension identity mismatch") } - if a := r.assignments[ref.SessionID]; !ref.Valid() || a == nil || a.ref.AssignmentID != ref.AssignmentID || ref.Epoch > a.ref.Epoch { + if ref != r.suspendedBy { return AssignmentError(proto.AssignmentConflict) } r.sender = sender - r.suspension = nil + r.suspension, r.suspendedBy = nil, proto.AssignmentRef{} for _, owner := range r.executors { r.scheduleExecutorIdleLocked(owner) } diff --git a/apps/daemon/internal/dispatch/suspend_test.go b/apps/daemon/internal/dispatch/suspend_test.go index 3a05102f6..2ebd342cf 100644 --- a/apps/daemon/internal/dispatch/suspend_test.go +++ b/apps/daemon/internal/dispatch/suspend_test.go @@ -43,7 +43,7 @@ func TestQuiesceRejectsEveryUnsettledResource(t *testing.T) { "active": func(r *Router) { r.sessions["run"] = &sessionState{ctxCancel: func() {}} }, "preparing": func(r *Router) { r.preparations["p"] = &preparationState{owns: true} }, "receipt": func(r *Router) { r.preparations["p"] = &preparationState{busy: true} }, - "read": func(r *Router) { r.workspaceReads = map[string]struct{}{"read": {}} }, + "read": func(r *Router) { r.workspaceReads = map[string]workspaceRead{"read": {}} }, "write": func(r *Router) { r.workspaceWrite = &workspaceUpload{} }, "export": func(r *Router) { r.workspaceExport = &workspaceExport{} }, "permission": func(r *Router) { r.permIndex["permission"] = "run" }, @@ -129,6 +129,12 @@ func TestResumeRequiresExactSuspensionAndAssignment(t *testing.T) { if err := r.Resume(foreign, request, sender); err == nil { t.Fatal("foreign assignment reopened admission") } + // Another bound assignment did not quiesce the Runtime. + other := proto.AssignmentRef{SessionID: "other", AssignmentID: "other", Epoch: 1} + bindAssignment(r, other, "env") + if err := r.Resume(other, request, sender); !errors.Is(err, AssignmentError(proto.AssignmentConflict)) { + t.Fatalf("other assignment resume = %v", err) + } if err := r.Resume(suspendRef, request, sender); err != nil { t.Fatal(err) } diff --git a/apps/daemon/internal/dispatch/workspace_export.go b/apps/daemon/internal/dispatch/workspace_export.go index 9366ed690..5689bcaa0 100644 --- a/apps/daemon/internal/dispatch/workspace_export.go +++ b/apps/daemon/internal/dispatch/workspace_export.go @@ -9,10 +9,12 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" ) +// workspaceExport is one admitted export; done closes once it ends. type workspaceExport struct { request proto.Envelope requests chan proto.WorkspaceExportPayload cancel context.CancelFunc + done chan struct{} } func (r *Router) handleWorkspaceExport(ctx context.Context, env proto.Envelope) error { @@ -57,7 +59,7 @@ func (r *Router) handleWorkspaceExport(ctx context.Context, env proto.Envelope) } owner, stop := r.shutdownContext(p.ctx) owner, cancel := context.WithTimeout(owner, 180*time.Second) - u = &workspaceExport{request: env, requests: make(chan proto.WorkspaceExportPayload, 1), cancel: func() { cancel(); stop() }} + u = &workspaceExport{request: env, requests: make(chan proto.WorkspaceExportPayload, 1), cancel: func() { cancel(); stop() }, done: make(chan struct{})} u.requests <- request r.workspaceExport = u r.shutdownWG.Add(1) @@ -68,6 +70,7 @@ func (r *Router) handleWorkspaceExport(ctx context.Context, env proto.Envelope) func (r *Router) runWorkspaceExport(ctx context.Context, u *workspaceExport) { defer r.shutdownWG.Done() + defer close(u.done) reader, writer := io.Pipe() done := make(chan struct{}) go func() { diff --git a/apps/daemon/internal/dispatch/workspace_read.go b/apps/daemon/internal/dispatch/workspace_read.go index 528999397..d78c5962f 100644 --- a/apps/daemon/internal/dispatch/workspace_read.go +++ b/apps/daemon/internal/dispatch/workspace_read.go @@ -13,6 +13,12 @@ import ( const workspaceReadCapacity = 4 +// workspaceRead is one admitted read; done closes once its result is sent. +type workspaceRead struct { + sessionID string + done chan struct{} +} + func (r *Router) handleWorkspaceRead(ctx context.Context, env proto.Envelope) error { // Never echo an unbounded correlation ID onto the shared connection. if len(env.ID) > proto.WorkspaceReadMaxIDBytes { @@ -42,13 +48,15 @@ func (r *Router) handleWorkspaceRead(ctx context.Context, env proto.Envelope) er return r.sendWorkspaceRead(ctx, env, rejectedWorkspaceRead(code)) } if r.workspaceReads == nil { - r.workspaceReads = make(map[string]struct{}) + r.workspaceReads = make(map[string]workspaceRead) } - r.workspaceReads[env.ID] = struct{}{} + read := workspaceRead{sessionID: env.Assignment.SessionID, done: make(chan struct{})} + r.workspaceReads[env.ID] = read r.shutdownWG.Add(1) r.mu.Unlock() go func() { defer r.shutdownWG.Done() + defer close(read.done) defer func() { r.mu.Lock(); delete(r.workspaceReads, env.ID); r.mu.Unlock() }() // Observer loss does not discard an admitted native wait or replay it. operation, cancel := context.WithTimeout(context.WithoutCancel(ctx), 12*time.Second) diff --git a/apps/daemon/internal/dispatch/workspace_write.go b/apps/daemon/internal/dispatch/workspace_write.go index 78d367f11..2383218e5 100644 --- a/apps/daemon/internal/dispatch/workspace_write.go +++ b/apps/daemon/internal/dispatch/workspace_write.go @@ -13,11 +13,13 @@ import ( ) // Router.mu protects this single bounded transfer for the dedicated Environment. +// done closes once its result is sent. type workspaceUpload struct { envelope proto.Envelope request proto.WorkspaceWritePayload data []byte ready chan struct{} + done chan struct{} finished bool apply bool uncertain bool @@ -78,7 +80,7 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("resource_unavailable")) } } - u := &workspaceUpload{envelope: env, request: request, data: make([]byte, 0, request.SizeBytes), ready: make(chan struct{})} + u := &workspaceUpload{envelope: env, request: request, data: make([]byte, 0, request.SizeBytes), ready: make(chan struct{}), done: make(chan struct{})} r.workspaceWrite = u r.shutdownWG.Add(1) r.mu.Unlock() @@ -112,6 +114,7 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e func (r *Router) runWorkspaceUpload(ctx context.Context, u *workspaceUpload) { defer r.shutdownWG.Done() + defer close(u.done) timer := time.NewTimer(120 * time.Second) defer timer.Stop() select { @@ -120,12 +123,17 @@ func (r *Router) runWorkspaceUpload(ctx context.Context, u *workspaceUpload) { case <-timer.C: } r.mu.Lock() - apply := u.apply && !r.closed + // A release of the assignment before the write applies fences it. + fenced := r.admitLocked(u.envelope.Assignment, u.request.SessionID, u.request.EnvironmentID) + apply := u.apply && !r.closed && fenced == "" u.finished = true data := u.data u.data = nil r.mu.Unlock() result := rejectedWorkspaceWrite("invalid_request") + if fenced != "" { + result = rejectedWorkspaceWrite(fenced) + } if apply { write, err := r.localWorkspace.WriteWorkspaceFile(ctx, u.request.Path, data) result = workspaceWriteResult(write, err, u.request.SizeBytes) diff --git a/services/core/internal/db/queries/devices.sql b/services/core/internal/db/queries/devices.sql index 1af723a27..798815b49 100644 --- a/services/core/internal/db/queries/devices.sql +++ b/services/core/internal/db/queries/devices.sql @@ -12,8 +12,15 @@ LEFT JOIN runtime_allocations a ON a.device_id = d.id WHERE d.id = $1; -- name: RevokeDevice :execrows -UPDATE devices SET revoked_at = COALESCE(revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL -WHERE tenant_id = $1 AND id = $2; +-- No Runtime is left to act on a revoked device's releases, so revocation +-- settles them. +WITH settled AS ( + UPDATE session_runtime_assignments b SET applied_epoch = b.epoch + WHERE b.runtime_id = sqlc.arg(id) AND b.desired_state = 'released' + AND EXISTS (SELECT 1 FROM devices d WHERE d.tenant_id = sqlc.arg(tenant_id) AND d.id = sqlc.arg(id)) +) +UPDATE devices v SET revoked_at = COALESCE(v.revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL +WHERE v.tenant_id = sqlc.arg(tenant_id) AND v.id = sqlc.arg(id); -- name: TouchDevice :execrows UPDATE devices SET last_seen_at = clock_timestamp() @@ -57,10 +64,12 @@ UPDATE session_runtime_assignments SET native_session_id = $2 WHERE session_id = -- name: ReleaseSessionAssignment :exec -- An identical release keeps its epoch; a release that adds home removal --- advances it. -UPDATE session_runtime_assignments -SET desired_state = 'released', epoch = epoch + 1, remove_home = remove_home OR sqlc.arg(remove_home)::boolean -WHERE session_id = sqlc.arg(session_id) AND (desired_state = 'bound' OR (sqlc.arg(remove_home)::boolean AND NOT remove_home)); +-- advances it. A release whose Runtime has no authority is settled, since no +-- Runtime can act on it. +UPDATE session_runtime_assignments b +SET desired_state = 'released', epoch = b.epoch + 1, remove_home = b.remove_home OR sqlc.arg(remove_home)::boolean, + applied_epoch = CASE WHEN EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = b.runtime_id) THEN b.applied_epoch ELSE b.epoch + 1 END +WHERE b.session_id = sqlc.arg(session_id) AND (b.desired_state = 'bound' OR (sqlc.arg(remove_home)::boolean AND NOT b.remove_home)); -- name: ListPendingAssignmentReleases :many SELECT session_id, runtime_id, assignment_id, epoch, remove_home FROM session_runtime_assignments diff --git a/services/core/internal/db/sqlc/devices.sql.go b/services/core/internal/db/sqlc/devices.sql.go index da39d4fc0..cec299d14 100644 --- a/services/core/internal/db/sqlc/devices.sql.go +++ b/services/core/internal/db/sqlc/devices.sql.go @@ -249,9 +249,10 @@ func (q *Queries) ListPendingAssignmentReleases(ctx context.Context, runtimeIds } const releaseSessionAssignment = `-- name: ReleaseSessionAssignment :exec -UPDATE session_runtime_assignments -SET desired_state = 'released', epoch = epoch + 1, remove_home = remove_home OR $1::boolean -WHERE session_id = $2 AND (desired_state = 'bound' OR ($1::boolean AND NOT remove_home)) +UPDATE session_runtime_assignments b +SET desired_state = 'released', epoch = b.epoch + 1, remove_home = b.remove_home OR $1::boolean, + applied_epoch = CASE WHEN EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = b.runtime_id) THEN b.applied_epoch ELSE b.epoch + 1 END +WHERE b.session_id = $2 AND (b.desired_state = 'bound' OR ($1::boolean AND NOT b.remove_home)) ` type ReleaseSessionAssignmentParams struct { @@ -260,7 +261,8 @@ type ReleaseSessionAssignmentParams struct { } // An identical release keeps its epoch; a release that adds home removal -// advances it. +// advances it. A release whose Runtime has no authority is settled, since no +// Runtime can act on it. func (q *Queries) ReleaseSessionAssignment(ctx context.Context, arg ReleaseSessionAssignmentParams) error { _, err := q.db.Exec(ctx, releaseSessionAssignment, arg.RemoveHome, arg.SessionID) return err @@ -284,8 +286,13 @@ func (q *Queries) RememberNativeSession(ctx context.Context, arg RememberNativeS } const revokeDevice = `-- name: RevokeDevice :execrows -UPDATE devices SET revoked_at = COALESCE(revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL -WHERE tenant_id = $1 AND id = $2 +WITH settled AS ( + UPDATE session_runtime_assignments b SET applied_epoch = b.epoch + WHERE b.runtime_id = $2 AND b.desired_state = 'released' + AND EXISTS (SELECT 1 FROM devices d WHERE d.tenant_id = $1 AND d.id = $2) +) +UPDATE devices v SET revoked_at = COALESCE(v.revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL +WHERE v.tenant_id = $1 AND v.id = $2 ` type RevokeDeviceParams struct { @@ -293,6 +300,8 @@ type RevokeDeviceParams struct { ID pgtype.UUID `json:"id"` } +// No Runtime is left to act on a revoked device's releases, so revocation +// settles them. func (q *Queries) RevokeDevice(ctx context.Context, arg RevokeDeviceParams) (int64, error) { result, err := q.db.Exec(ctx, revokeDevice, arg.TenantID, arg.ID) if err != nil { diff --git a/services/core/internal/execution/assignment_releases.go b/services/core/internal/execution/assignment_releases.go index cfce17413..113b11d3b 100644 --- a/services/core/internal/execution/assignment_releases.go +++ b/services/core/internal/execution/assignment_releases.go @@ -13,10 +13,16 @@ import ( // runAssignmentReleases delivers each released assignment to its connected // Runtime until the Runtime acknowledges it, so a Runtime that reconnects // receives the releases it missed. Releases run one per Session, bounded like -// executions. +// executions. A failed release backs off, so it cannot hold a slot that later +// releases need. func (w *Worker) runAssignmentReleases(ctx context.Context) error { + type outcome struct { + ref proto.AssignmentRef + acknowledged bool + } active := make(map[string]bool) - done := make(chan string, w.executionConcurrency()) + retries := releaseRetries{} + done := make(chan outcome, w.executionConcurrency()) var running sync.WaitGroup ctx, stop := context.WithCancel(ctx) defer func() { stop(); running.Wait() }() @@ -26,8 +32,9 @@ func (w *Worker) runAssignmentReleases(ctx context.Context) error { select { case <-ctx.Done(): return ctx.Err() - case session := <-done: - delete(active, session) + case result := <-done: + delete(active, result.ref.SessionID) + retries.record(result.ref, result.acknowledged, time.Now()) continue case <-ticker.C: } @@ -35,34 +42,70 @@ func (w *Worker) runAssignmentReleases(ctx context.Context) error { if err != nil { return err } + retries.keep(releases) + now := time.Now() for _, release := range releases { - session := release.Assignment.SessionID - if active[session] || len(active) >= w.executionConcurrency() { + ref := release.Assignment + if active[ref.SessionID] || len(active) >= w.executionConcurrency() || !retries.due(ref, now) { continue } - active[session] = true + active[ref.SessionID] = true running.Add(1) go func() { defer running.Done() - w.releaseAssignment(ctx, release) - done <- session + done <- outcome{ref: ref, acknowledged: w.releaseAssignment(ctx, release)} }() } } } +// releaseRetries holds each failed release's next attempt. The delay doubles +// from a second up to a minute. +type releaseRetries map[proto.AssignmentRef]releaseRetry + +type releaseRetry struct { + at time.Time + delay time.Duration +} + +func (r releaseRetries) due(ref proto.AssignmentRef, now time.Time) bool { + return !now.Before(r[ref].at) +} + +func (r releaseRetries) record(ref proto.AssignmentRef, acknowledged bool, now time.Time) { + if acknowledged { + delete(r, ref) + return + } + delay := min(max(2*r[ref].delay, time.Second), time.Minute) + r[ref] = releaseRetry{at: now.Add(delay), delay: delay} +} + +// keep forgets releases that are no longer pending for a connected Runtime. +func (r releaseRetries) keep(releases []sessions.AssignmentRelease) { + pending := make(map[proto.AssignmentRef]bool, len(releases)) + for _, release := range releases { + pending[release.Assignment] = true + } + for ref := range r { + if !pending[ref] { + delete(r, ref) + } + } +} + // releaseAssignment sends one release and records its acknowledgement. Home // removal is requested only from a Runtime that declares it. -func (w *Worker) releaseAssignment(ctx context.Context, release sessions.AssignmentRelease) { +func (w *Worker) releaseAssignment(ctx context.Context, release sessions.AssignmentRelease) bool { ctx, cancel := context.WithTimeout(ctx, 2*time.Minute) defer cancel() peer, err := w.dispatcher.authorizedPeer(ctx, release.RuntimeID) if err != nil { - return + return false } supported, known := peer.RemovesHomes() if !known { - return + return false } removeHome := release.RemoveHome && supported want := proto.AssignmentReleased @@ -72,9 +115,11 @@ func (w *Worker) releaseAssignment(ctx context.Context, release sessions.Assignm status, err := peer.Release(ctx, release.Assignment, removeHome) if err != nil || status.State != want { log.Warn(ctx, "Runtime assignment release unconfirmed", "session_id", release.Assignment.SessionID, "runtime_id", release.RuntimeID, "error_code", status.ErrorCode) - return + return false } if err := w.dispatcher.sessionExecution.AcknowledgeAssignmentRelease(ctx, release.Assignment); err != nil { log.Warn(ctx, "Runtime assignment release not recorded", "session_id", release.Assignment.SessionID, "runtime_id", release.RuntimeID) + return false } + return true } diff --git a/services/core/internal/execution/assignment_releases_test.go b/services/core/internal/execution/assignment_releases_test.go new file mode 100644 index 000000000..633095ee9 --- /dev/null +++ b/services/core/internal/execution/assignment_releases_test.go @@ -0,0 +1,30 @@ +package execution + +import ( + "testing" + "time" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" +) + +func TestFailedReleaseBacksOffAndLeavesLaterReleasesDue(t *testing.T) { + failing := proto.AssignmentRef{SessionID: "failing", AssignmentID: "failing", Epoch: 2} + later := proto.AssignmentRef{SessionID: "later", AssignmentID: "later", Epoch: 2} + retries, now := releaseRetries{}, time.Unix(0, 0) + retries.record(failing, false, now) + if retries.due(failing, now) || !retries.due(failing, now.Add(time.Second)) || !retries.due(later, now) { + t.Fatal("a failed release did not wait a second while later releases stayed due") + } + for range 8 { + retries.record(failing, false, now) + } + if retries.due(failing, now.Add(59*time.Second)) || !retries.due(failing, now.Add(time.Minute)) { + t.Fatalf("delay = %s, want a minute", retries[failing].delay) + } + retries.keep([]sessions.AssignmentRelease{{Assignment: later}}) + retries.record(later, true, now) + if len(retries) != 0 { + t.Fatal("settled or acknowledged releases kept a retry", retries) + } +} diff --git a/services/core/internal/execution/delivery.go b/services/core/internal/execution/delivery.go index 9eda34790..fa04dbb4a 100644 --- a/services/core/internal/execution/delivery.go +++ b/services/core/internal/execution/delivery.go @@ -58,7 +58,7 @@ func (d *Dispatcher) deliver(ctx context.Context, tenantID, sessionID string, pe defer unsubscribeChanges() status = sessions.TurnFailed result.AppliedThrough = first - subscription, err := peer.SubscribeDurable(request.RunID) + subscription, err := peer.SubscribeDurable(request.RunID, request.Assignment) if err != nil { result.ErrorCode = "device_disconnected" return diff --git a/services/core/internal/execution/preparation.go b/services/core/internal/execution/preparation.go index 29b22cfa2..4f59b1916 100644 --- a/services/core/internal/execution/preparation.go +++ b/services/core/internal/execution/preparation.go @@ -35,7 +35,7 @@ type preparedStart struct { func newPreparedStart(peer *runtimegateway.Session, ref proto.AssignmentRef) (*preparedStart, error) { id := uuid.NewString() - sub, err := peer.SubscribePreparation(id) + sub, err := peer.SubscribePreparation(id, ref) if err != nil { return nil, err } diff --git a/services/core/internal/persistence/postgres/deploymentpg/allocations.go b/services/core/internal/persistence/postgres/deploymentpg/allocations.go index d1fb3044c..4d24c722f 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/allocations.go +++ b/services/core/internal/persistence/postgres/deploymentpg/allocations.go @@ -340,11 +340,11 @@ func (t *cleanupTx) RevokeDevice(current deployment.Allocation) error { if err != nil { return err } - if err := t.ReleaseAssignment(t.ctx, false); err != nil { + // The release follows the revocation, which leaves no Runtime to deliver it to. + if _, err := t.q.RevokeRuntimeCleanupDevice(t.ctx, sqlc.RevokeRuntimeCleanupDeviceParams{TenantID: t.tenant, DeviceID: device}); err != nil { return err } - _, err = t.q.RevokeRuntimeCleanupDevice(t.ctx, sqlc.RevokeRuntimeCleanupDeviceParams{TenantID: t.tenant, DeviceID: device}) - return err + return t.ReleaseAssignment(t.ctx, false) } func (t *cleanupTx) RequestCleanup(current deployment.Allocation) (deployment.Allocation, error) { diff --git a/services/core/internal/persistence/postgres/deploymentpg/session_archive.go b/services/core/internal/persistence/postgres/deploymentpg/session_archive.go index e576f5013..b4f3693e2 100644 --- a/services/core/internal/persistence/postgres/deploymentpg/session_archive.go +++ b/services/core/internal/persistence/postgres/deploymentpg/session_archive.go @@ -69,10 +69,11 @@ func (t *archiveTx) RequestArchiveCleanup(current deployment.Allocation) error { if err != nil { return err } - if err := t.ReleaseAssignment(t.ctx, false); err != nil { + // The release follows the revocation, which leaves no Runtime to deliver it to. + if _, err := t.q.RevokeArchivedRuntimeDevice(t.ctx, sqlc.RevokeArchivedRuntimeDeviceParams{TenantID: t.tenant, DeviceID: device, SessionID: t.session}); err != nil { return err } - if _, err := t.q.RevokeArchivedRuntimeDevice(t.ctx, sqlc.RevokeArchivedRuntimeDeviceParams{TenantID: t.tenant, DeviceID: device, SessionID: t.session}); err != nil { + if err := t.ReleaseAssignment(t.ctx, false); err != nil { return err } _, err = t.q.RequestRuntimeCleanup(t.ctx, id) diff --git a/services/core/internal/persistence/postgres/sessionpg/session_writes.go b/services/core/internal/persistence/postgres/sessionpg/session_writes.go index 344e44b9a..f473c62d4 100644 --- a/services/core/internal/persistence/postgres/sessionpg/session_writes.go +++ b/services/core/internal/persistence/postgres/sessionpg/session_writes.go @@ -36,16 +36,18 @@ type deletionTx struct { } func (t *deletionTx) ApplyDeletion(ctx context.Context) error { - if err := t.ReleaseAssignment(ctx, true); err != nil { - return err - } if err := t.q.DeleteSessionArtifacts(ctx, t.session); err != nil { return err } if err := t.q.ReleaseUnallocatedRuntimePlacement(ctx, t.session); err != nil { return err } - return t.q.MarkSessionDeleted(ctx, t.session) + if err := t.q.MarkSessionDeleted(ctx, t.session); err != nil { + return err + } + // The release follows the deletion, which ends an Environment Runtime's + // authority and so settles its release. + return t.ReleaseAssignment(ctx, true) } func (t *deletionTx) RecordDeletionAudit(ctx context.Context) error { diff --git a/services/core/internal/runtimegateway/assignment.go b/services/core/internal/runtimegateway/assignment.go index 56c1fe331..701e922a8 100644 --- a/services/core/internal/runtimegateway/assignment.go +++ b/services/core/internal/runtimegateway/assignment.go @@ -74,11 +74,11 @@ func (s *Session) exchangeAssignment(ctx context.Context, typ string, ref proto. s.assignmentReplies[id] = replies s.assignmentMu.Unlock() defer func() { s.assignmentMu.Lock(); delete(s.assignmentReplies, id); s.assignmentMu.Unlock() }() - reply, err := s.exchangeChunkFrame(ctx, env, replies) + reply, err := s.exchangeFrame(ctx, env, replies) if err != nil { return status, err } - if reply.Type != proto.TypeAssignmentStatus || reply.Assignment != ref || reply.DecodePayload(&status) != nil || !validAssignmentStatus(typ, status) { + if reply.Type != proto.TypeAssignmentStatus || reply.DecodePayload(&status) != nil || !validAssignmentStatus(typ, status) { return proto.AssignmentStatusPayload{}, errors.New("agentdaemon gateway: invalid assignment status") } return status, nil diff --git a/services/core/internal/runtimegateway/cancellation_test.go b/services/core/internal/runtimegateway/cancellation_test.go index ab019ad63..217c482ef 100644 --- a/services/core/internal/runtimegateway/cancellation_test.go +++ b/services/core/internal/runtimegateway/cancellation_test.go @@ -44,7 +44,7 @@ func TestArchivedReceiptTracksDeliveryBeyondDoneAndRejectsNewWork(t *testing.T) if err != nil { t.Fatal(err) } - subscription, err := peer.SubscribeDurable("run") + subscription, err := peer.SubscribeDurable("run", proto.AssignmentRef{}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/runtimegateway/chunk_exchange.go b/services/core/internal/runtimegateway/chunk_exchange.go deleted file mode 100644 index 05b50165a..000000000 --- a/services/core/internal/runtimegateway/chunk_exchange.go +++ /dev/null @@ -1,26 +0,0 @@ -package runtimegateway - -import ( - "context" - - "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" -) - -// exchangeChunkFrame sends once and waits for the owning transfer's next receipt. -// Callers validate their own frame and receipt; no transport outcome is retried. -func (s *Session) exchangeChunkFrame(ctx context.Context, env proto.Envelope, replies <-chan proto.Envelope) (proto.Envelope, error) { - if err := s.Send(ctx, env); err != nil { - return proto.Envelope{}, err - } - select { - case reply, ok := <-replies: - if !ok { - return proto.Envelope{}, ErrSessionClosed - } - return reply, nil - case <-ctx.Done(): - return proto.Envelope{}, ctx.Err() - case <-s.closed: - return proto.Envelope{}, ErrSessionClosed - } -} diff --git a/services/core/internal/runtimegateway/exchange.go b/services/core/internal/runtimegateway/exchange.go new file mode 100644 index 000000000..e8a7c0eeb --- /dev/null +++ b/services/core/internal/runtimegateway/exchange.go @@ -0,0 +1,44 @@ +package runtimegateway + +import ( + "context" + "errors" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" +) + +var errAssignmentEcho = errors.New("agentdaemon gateway: frame names another assignment") + +// exchangeFrame sends env once and waits for the reply its owner routes to +// replies. A reply that arrived before the deadline or the connection closed +// still counts. The reply must echo env's assignment. Callers validate their +// own frame and receipt; no transport outcome is retried. +func (s *Session) exchangeFrame(ctx context.Context, env proto.Envelope, replies <-chan proto.Envelope) (proto.Envelope, error) { + if err := s.Send(ctx, env); err != nil { + return proto.Envelope{}, err + } + var reply proto.Envelope + ok := true + select { + case reply, ok = <-replies: + case <-ctx.Done(): + select { + case reply, ok = <-replies: + default: + return proto.Envelope{}, ctx.Err() + } + case <-s.closed: + select { + case reply, ok = <-replies: + default: + return proto.Envelope{}, ErrSessionClosed + } + } + if !ok { + return proto.Envelope{}, ErrSessionClosed + } + if reply.Assignment != env.Assignment { + return proto.Envelope{}, errAssignmentEcho + } + return reply, nil +} diff --git a/services/core/internal/runtimegateway/exchange_test.go b/services/core/internal/runtimegateway/exchange_test.go new file mode 100644 index 000000000..a43650da4 --- /dev/null +++ b/services/core/internal/runtimegateway/exchange_test.go @@ -0,0 +1,52 @@ +package runtimegateway + +import ( + "errors" + "testing" + + "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" +) + +func TestFramesMustNameTheRequestAssignment(t *testing.T) { + foreign := testAssignment + foreign.AssignmentID = "foreign" + for name, ref := range map[string]proto.AssignmentRef{"missing": {}, "foreign": foreign} { + t.Run(name, func(t *testing.T) { + s := NewSession(newFakeConn(), "device", "tenant", "test", NewRegistry(), nil) + defer s.Close("test") + read := make(chan error, 1) + go func() { _, err := s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()); read <- err }() + request := <-s.sendCh + reply, _ := request.Reply(proto.TypeWorkspaceReadResult, proto.WorkspaceReadResultPayload{Outcome: "completed", CloseAcknowledged: true}) + reply.Assignment = ref + s.dispatch(reply) + if err := <-read; !errors.Is(err, errAssignmentEcho) { + t.Fatalf("read = %v", err) + } + + ack := make(chan error, 1) + decision, _ := proto.NewEnvelope(proto.TypePromptCancel, "run", proto.PromptCancelPayload{DeliveryID: "delivery"}) + decision.Assignment = testAssignment + go func() { _, err := s.SendAndWaitInteractionAck(t.Context(), decision, "delivery"); ack <- err }() + request = <-s.sendCh + reply, _ = request.Reply(proto.TypeInteractionDecisionAck, proto.InteractionDecisionAckPayload{DeliveryID: "delivery", Applied: true}) + reply.Assignment = ref + s.dispatch(reply) + if err := <-ack; !errors.Is(err, errAssignmentEcho) { + t.Fatalf("ack = %v", err) + } + + sub, err := s.SubscribeDurable("run", testAssignment) + if err != nil { + t.Fatal(err) + } + s.dispatch(proto.Envelope{Type: proto.TypeDelta, ID: "run", Assignment: ref}) + for range sub.Events { + t.Fatal("frame of another assignment delivered") + } + if !errors.Is(sub.Err(), errAssignmentEcho) { + t.Fatalf("subscription = %v", sub.Err()) + } + }) + } +} diff --git a/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go b/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go index 259463cc9..e059977d0 100644 --- a/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go +++ b/services/core/internal/runtimegateway/mcp_bearer_live_linux_test.go @@ -92,7 +92,7 @@ func TestLiveMCPBearerGatewayColdContinuation(t *testing.T) { turns = append(turns, turn) runID := uuid.NewString() request := proto.PromptRequestPayload{AgentKind: "codex", AgentStateKey: "mcp-bearer-acceptance", AgentSessionID: resume, StrictResume: true, ReleaseOnCompletion: true, ObserveMessages: true, ObserveToolObservations: true, DisableExecutionEnvironment: true, DisableSubagents: true, MCPHTTPServers: &servers, Model: "MiniMax-M3", ExecutionControls: &proto.ExecutionControls{WebSearch: "disabled", TextVerbosity: "medium"}} - sub, err := peer.SubscribeDurable(runID) + sub, err := peer.SubscribeDurable(runID, assignment) if err != nil { t.Fatal("cannot subscribe before real daemon dispatch") } @@ -141,7 +141,7 @@ type mcpBearerControl struct { func mcpBearerStart(t *testing.T, ctx context.Context, peer *Session, assignment proto.AssignmentRef, request proto.PromptRequestPayload, runID string, input proto.MessageInput) *mcpBearerControl { t.Helper() c := &mcpBearerControl{t: t, ctx: ctx, peer: peer, assignment: assignment, id: uuid.NewString()} - events, err := peer.SubscribePreparation(c.id) + events, err := peer.SubscribePreparation(c.id, assignment) if err != nil { t.Fatal("cannot subscribe to the preparation") } diff --git a/services/core/internal/runtimegateway/preparation.go b/services/core/internal/runtimegateway/preparation.go index 31820f5f4..2e00bb38b 100644 --- a/services/core/internal/runtimegateway/preparation.go +++ b/services/core/internal/runtimegateway/preparation.go @@ -14,8 +14,9 @@ type preparationSubscription struct { // SubscribePreparation correlates private control responses without registering // a Run. Unsubscribe on abandonment; a terminal resource status closes the stream. -// This subscription belongs to this physical daemon connection only. -func (s *Session) SubscribePreparation(requestID string) (*Subscription, error) { +// This subscription belongs to this physical daemon connection only. ref is the +// assignment that prepares. +func (s *Session) SubscribePreparation(requestID string, ref proto.AssignmentRef) (*Subscription, error) { s.preparationMu.Lock() defer s.preparationMu.Unlock() if s.IsClosed() { @@ -25,7 +26,7 @@ func (s *Session) SubscribePreparation(requestID string) (*Subscription, error) return nil, errors.New("agentdaemon gateway: invalid, duplicate or excess preparation subscription") } ch := make(chan proto.Envelope, 16) - sub := &Subscription{Events: ch, ch: ch} + sub := &Subscription{Events: ch, ch: ch, ref: ref} s.preparations[requestID] = &preparationSubscription{sub: sub} return sub, nil } @@ -60,14 +61,11 @@ func (s *Session) dispatchPreparation(env proto.Envelope) { } p.sub.mu.Lock() defer p.sub.mu.Unlock() - select { - case p.sub.ch <- env: + if p.sub.deliverLocked(env) { switch status.State { case "started", "released", "expired", "failed": p.sub.closeLocked(nil) } - default: - p.sub.closeLocked(ErrSubscriberOverflow) } if p.sub.closed { delete(s.preparations, env.ID) diff --git a/services/core/internal/runtimegateway/preparation_test.go b/services/core/internal/runtimegateway/preparation_test.go index 28e58c854..99eb254a3 100644 --- a/services/core/internal/runtimegateway/preparation_test.go +++ b/services/core/internal/runtimegateway/preparation_test.go @@ -12,7 +12,7 @@ func TestPreparationSubscriptionHasNoRunIdentityAndOrdersRevisions(t *testing.T) registry := NewRegistry() s := NewSession(newFakeConn(), "device", "tenant", "test", registry, nil) defer s.Close("test") - sub, err := s.SubscribePreparation("request") + sub, err := s.SubscribePreparation("request", proto.AssignmentRef{}) if err != nil { t.Fatal(err) } @@ -49,7 +49,7 @@ func TestPreparationSubscriptionHasNoRunIdentityAndOrdersRevisions(t *testing.T) func TestPreparationCloseAndOverflowDoNotInventRunEvents(t *testing.T) { for _, overflow := range []bool{false, true} { s := NewSession(newFakeConn(), "device", "tenant", "test", NewRegistry(), nil) - sub, err := s.SubscribePreparation("request") + sub, err := s.SubscribePreparation("request", proto.AssignmentRef{}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/runtimegateway/runtime_prepare.go b/services/core/internal/runtimegateway/runtime_prepare.go index 4f6562390..6fe0e38da 100644 --- a/services/core/internal/runtimegateway/runtime_prepare.go +++ b/services/core/internal/runtimegateway/runtime_prepare.go @@ -62,7 +62,7 @@ func (s *Session) PrepareRuntime(ctx context.Context, id string, ref proto.Assig if err != nil || len(encoded) > proto.RuntimePrepareMaxFrameBytes { return unknown, errors.New("agentdaemon gateway: invalid Runtime frame") } - reply, err := s.exchangeChunkFrame(ctx, env, replies) + reply, err := s.exchangeFrame(ctx, env, replies) if err != nil { return unknown, err } diff --git a/services/core/internal/runtimegateway/runtime_prepare_test.go b/services/core/internal/runtimegateway/runtime_prepare_test.go index 6ee8347b8..2ef92bc92 100644 --- a/services/core/internal/runtimegateway/runtime_prepare_test.go +++ b/services/core/internal/runtimegateway/runtime_prepare_test.go @@ -57,6 +57,7 @@ func finishCapabilities(t *testing.T, done <-chan capabilityOutcome) capabilityO } func replyCapabilities(s *Session, id string, result proto.RuntimePrepareResultPayload) { reply, _ := proto.NewEnvelope(proto.TypeRuntimePrepareResult, id, result) + reply.Assignment = testAssignment s.dispatch(reply) } func noCapabilityFrame(t *testing.T, s *Session) { diff --git a/services/core/internal/runtimegateway/session.go b/services/core/internal/runtimegateway/session.go index bcffabd39..6638e8276 100644 --- a/services/core/internal/runtimegateway/session.go +++ b/services/core/internal/runtimegateway/session.go @@ -123,7 +123,7 @@ type Session struct { workspaceExports map[string]chan proto.Envelope ackMu sync.Mutex - ackWaiters map[string]chan proto.InteractionDecisionAckPayload + ackWaiters map[string]chan proto.Envelope // sendCh feeds the WS write loop. Capacity is bounded so a slow // peer can't queue unbounded outbound frames; once full, Send @@ -165,7 +165,7 @@ func NewSession(conn WSConn, deviceID, workspaceID, daemonVersion string, reg *R assignments: map[string]proto.AssignmentRef{}, assignmentReplies: map[string]chan proto.Envelope{}, preparations: map[string]*preparationSubscription{}, - ackWaiters: map[string]chan proto.InteractionDecisionAckPayload{}, + ackWaiters: map[string]chan proto.Envelope{}, sendCh: make(chan proto.Envelope, 64), closed: make(chan struct{}), } @@ -305,7 +305,7 @@ func (s *Session) SendAndWaitInteractionAck(ctx context.Context, env proto.Envel if deliveryID == "" { return proto.InteractionDecisionAckPayload{}, errors.New("agentdaemon gateway: interaction delivery id is required") } - waiter := make(chan proto.InteractionDecisionAckPayload, 1) + waiter := make(chan proto.Envelope, 1) s.ackMu.Lock() if _, exists := s.ackWaiters[deliveryID]; exists { s.ackMu.Unlock() @@ -319,27 +319,20 @@ func (s *Session) SendAndWaitInteractionAck(ctx context.Context, env proto.Envel s.ackMu.Unlock() }() - if err := s.Send(ctx, env); err != nil { - return proto.InteractionDecisionAckPayload{}, err - } + // The exchange prefers an ack that raced the deadline; treating an + // already-applied decision as retryable can trigger a contradictory second + // human response. waitCtx, cancel := context.WithTimeout(ctx, InteractionAckTimeout) defer cancel() - select { - case ack := <-waiter: - return ack, nil - case <-s.closed: - return proto.InteractionDecisionAckPayload{}, ErrSessionClosed - case <-waitCtx.Done(): - // If the ack raced the deadline, prefer the application receipt; - // treating an already-applied decision as retryable can trigger a - // contradictory second human response. - select { - case ack := <-waiter: - return ack, nil - default: - return proto.InteractionDecisionAckPayload{}, waitCtx.Err() - } + reply, err := s.exchangeFrame(waitCtx, env, waiter) + if err != nil { + return proto.InteractionDecisionAckPayload{}, err + } + var ack proto.InteractionDecisionAckPayload + if reply.DecodePayload(&ack) != nil { + return proto.InteractionDecisionAckPayload{}, errors.New("agentdaemon gateway: invalid interaction decision ack") } + return ack, nil } // writeLoop is the single writer goroutine that gorilla/websocket @@ -559,7 +552,7 @@ func (s *Session) dispatch(env proto.Envelope) { return } select { - case waiter <- ack: + case waiter <- env: default: } return diff --git a/services/core/internal/runtimegateway/session_test.go b/services/core/internal/runtimegateway/session_test.go index 20dad7072..2b68bf637 100644 --- a/services/core/internal/runtimegateway/session_test.go +++ b/services/core/internal/runtimegateway/session_test.go @@ -206,7 +206,7 @@ func TestSession_DispatchDeliversToSubscriber(t *testing.T) { sess.Start() defer sess.Close("test done") - sub, err := sess.SubscribeDurable("run-1") + sub, err := sess.SubscribeDurable("run-1", proto.AssignmentRef{}) ch := sub.Events if err != nil { t.Fatalf("Subscribe: %v", err) @@ -232,7 +232,7 @@ func TestSession_DoneFrameAutoUnsubscribes(t *testing.T) { sess.Start() defer sess.Close("test done") - sub, _ := sess.SubscribeDurable("run-1") + sub, _ := sess.SubscribeDurable("run-1", proto.AssignmentRef{}) ch := sub.Events env, _ := proto.NewEnvelope(proto.TypeDone, "run-1", proto.DonePayload{Content: "ok"}) raw, _ := jsonMarshal(env) @@ -267,7 +267,7 @@ func TestSession_PermissionRequestIndexedInRegistry(t *testing.T) { sess.Start() defer sess.Close("test done") - sub, err := sess.SubscribeDurable("run-1") + sub, err := sess.SubscribeDurable("run-1", proto.AssignmentRef{}) ch := sub.Events if err != nil { t.Fatalf("subscribe: %v", err) @@ -304,7 +304,7 @@ func TestSession_PermissionRequestIndexedInRegistry(t *testing.T) { func TestSession_CloseReportsUnknownWithoutExecutionEvents(t *testing.T) { sess := NewSession(newFakeConn(), "device", "tenant", proto.Version, NewRegistry(), nil) - sub, err := sess.SubscribeDurable("run") + sub, err := sess.SubscribeDurable("run", proto.AssignmentRef{}) if err != nil { t.Fatal(err) } @@ -455,7 +455,7 @@ func TestSession_PermissionRequiresPayloadIdentity(t *testing.T) { reg := NewRegistry() sess := NewSession(newFakeConn(), "device", "tenant", proto.Version, reg, nil) defer sess.Close("test done") - sub, err := sess.SubscribeDurable("run") + sub, err := sess.SubscribeDurable("run", proto.AssignmentRef{}) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/runtimegateway/subscription.go b/services/core/internal/runtimegateway/subscription.go index 26864a24b..567d81a58 100644 --- a/services/core/internal/runtimegateway/subscription.go +++ b/services/core/internal/runtimegateway/subscription.go @@ -10,9 +10,11 @@ import ( var ErrSubscriberOverflow = errors.New("execution subscriber buffer overflow") +// Subscription streams one operation's frames, which all name ref. type Subscription struct { Events <-chan proto.Envelope ch chan proto.Envelope + ref proto.AssignmentRef mu sync.Mutex err error closed bool @@ -31,14 +33,31 @@ func (s *Subscription) closeLocked(err error) { } } +// deliverLocked queues env. A frame of another assignment or a full buffer +// ends the subscription with an error. +func (s *Subscription) deliverLocked(env proto.Envelope) bool { + if env.Assignment != s.ref { + s.closeLocked(errAssignmentEcho) + return false + } + select { + case s.ch <- env: + return true + default: + s.closeLocked(ErrSubscriberOverflow) + return false + } +} + // SubscribeDurable reports transport loss and overflow separately from native -// execution events. Callers must inspect Err after Events closes. -func (s *Session) SubscribeDurable(runID string) (*Subscription, error) { +// execution events. Callers must inspect Err after Events closes. ref is the +// assignment that starts the run. +func (s *Session) SubscribeDurable(runID string, ref proto.AssignmentRef) (*Subscription, error) { if runID == "" { return nil, fmt.Errorf("agentdaemon gateway: Subscribe requires non-empty runID") } ch := make(chan proto.Envelope, 256) - sub := &Subscription{Events: ch, ch: ch} + sub := &Subscription{Events: ch, ch: ch, ref: ref} s.subsMu.Lock() defer s.subsMu.Unlock() if s.IsClosed() { @@ -84,13 +103,8 @@ func (s *Session) dispatchToSubscriber(env proto.Envelope) { } sub.mu.Lock() defer sub.mu.Unlock() - select { - case sub.ch <- env: - if env.Type == proto.TypeDone { - sub.closeLocked(nil) - } - default: - sub.closeLocked(ErrSubscriberOverflow) + if sub.deliverLocked(env) && env.Type == proto.TypeDone { + sub.closeLocked(nil) } if sub.closed { delete(s.subs, env.ID) diff --git a/services/core/internal/runtimegateway/subscription_test.go b/services/core/internal/runtimegateway/subscription_test.go index 1fbece704..ae21d3716 100644 --- a/services/core/internal/runtimegateway/subscription_test.go +++ b/services/core/internal/runtimegateway/subscription_test.go @@ -11,11 +11,11 @@ import ( func TestDurableSubscriptionOverflowIsExplicitAndIsolated(t *testing.T) { s := NewSession(newFakeConn(), "device", "tenant", proto.Version, NewRegistry(), nil) defer s.Close("test finished") - sub, err := s.SubscribeDurable("slow") + sub, err := s.SubscribeDurable("slow", proto.AssignmentRef{}) if err != nil { t.Fatal(err) } - other, _ := s.SubscribeDurable("other") + other, _ := s.SubscribeDurable("other", proto.AssignmentRef{}) for range 257 { s.dispatchToSubscriber(proto.Envelope{ID: "slow", Type: proto.TypeDelta}) } @@ -36,7 +36,7 @@ func TestDurableSubscriptionOverflowIsExplicitAndIsolated(t *testing.T) { func TestSubscriptionCloseAndDispatchAreSerialized(t *testing.T) { for range 100 { s := NewSession(newFakeConn(), "device", "tenant", proto.Version, NewRegistry(), nil) - sub, _ := s.SubscribeDurable("run") + sub, _ := s.SubscribeDurable("run", proto.AssignmentRef{}) var wg sync.WaitGroup wg.Add(3) go func() { diff --git a/services/core/internal/runtimegateway/suspension.go b/services/core/internal/runtimegateway/suspension.go index 9c57dbf48..71bfa54c5 100644 --- a/services/core/internal/runtimegateway/suspension.go +++ b/services/core/internal/runtimegateway/suspension.go @@ -43,33 +43,16 @@ func (s *Session) SuspendControl(ctx context.Context, kind string, ref proto.Ass s.suspendReplies[id] = replies s.suspendMu.Unlock() defer func() { s.suspendMu.Lock(); delete(s.suspendReplies, id); s.suspendMu.Unlock() }() - if err := s.Send(ctx, envelope); err != nil { + // The daemon closes after quiesced; the exchange reads a result that + // arrived before the closure. + reply, err := s.exchangeFrame(ctx, envelope, replies) + if err != nil { return result, err } - // The daemon closes after quiesced. Read the buffered result before treating - // connection closure as unknown, including when both channels become ready. - decode := func(reply proto.Envelope, ok bool) (proto.EnvironmentSuspendResultPayload, error) { - if !ok { - return result, ErrSessionClosed - } - if reply.Type != expected || reply.DecodePayload(&result) != nil || result.EnvironmentID != request.EnvironmentID || result.SuspendID != request.SuspendID || (result.Accepted && result.ErrorCode != "") { - return proto.EnvironmentSuspendResultPayload{}, errors.New("invalid suspension acknowledgement") - } - return result, nil - } - select { - case reply, ok := <-replies: - return decode(reply, ok) - case <-ctx.Done(): - return result, ctx.Err() - case <-s.closed: - select { - case reply, ok := <-replies: - return decode(reply, ok) - default: - return result, ErrSessionClosed - } + if reply.Type != expected || reply.DecodePayload(&result) != nil || result.EnvironmentID != request.EnvironmentID || result.SuspendID != request.SuspendID || (result.Accepted && result.ErrorCode != "") { + return proto.EnvironmentSuspendResultPayload{}, errors.New("invalid suspension acknowledgement") } + return result, nil } func (s *Session) dispatchSuspendReply(env proto.Envelope) { diff --git a/services/core/internal/runtimegateway/suspension_test.go b/services/core/internal/runtimegateway/suspension_test.go index e76bb9e62..cb0bf0c56 100644 --- a/services/core/internal/runtimegateway/suspension_test.go +++ b/services/core/internal/runtimegateway/suspension_test.go @@ -19,7 +19,7 @@ func TestSuspendAcknowledgementSurvivesImmediateConnectionClose(t *testing.T) { done <- err }() request := <-s.sendCh - reply, _ := proto.NewEnvelope(proto.TypeEnvironmentQuiesced, request.ID, proto.EnvironmentSuspendResultPayload{EnvironmentID: "env", SuspendID: "attempt", Accepted: true}) + reply, _ := request.Reply(proto.TypeEnvironmentQuiesced, proto.EnvironmentSuspendResultPayload{EnvironmentID: "env", SuspendID: "attempt", Accepted: true}) s.dispatch(reply) s.Close("parked") if err := <-done; err != nil { @@ -40,7 +40,7 @@ func TestSuspendControlRejectsForeignIdentityAndDoesNotReplay(t *testing.T) { done <- err }() request := <-s.sendCh - reply, _ := proto.NewEnvelope(proto.TypeEnvironmentResumed, request.ID, result) + reply, _ := request.Reply(proto.TypeEnvironmentResumed, result) s.dispatch(reply) if err := <-done; err == nil { t.Fatal("foreign receipt accepted") diff --git a/services/core/internal/runtimegateway/wire_test.go b/services/core/internal/runtimegateway/wire_test.go index eb20e7ec9..41151dc76 100644 --- a/services/core/internal/runtimegateway/wire_test.go +++ b/services/core/internal/runtimegateway/wire_test.go @@ -159,7 +159,7 @@ func TestWireScenarios(t *testing.T) { c.connect() // Subscribing before any frame also observes that a failed // preparation produces no Run result. - run, err := c.session.SubscribeDurable(prototest.RunID) + run, err := c.session.SubscribeDurable(prototest.RunID, prototest.Assignment) if err != nil { t.Fatal(err) } @@ -211,7 +211,7 @@ func (c *coreSide) coreSends(frame proto.Envelope) { } return case proto.TypeExecutionPrepare: - sub, err := c.session.SubscribePreparation(frame.ID) + sub, err := c.session.SubscribePreparation(frame.ID, frame.Assignment) if err != nil { t.Fatal(err) } diff --git a/services/core/internal/runtimegateway/workspace_directory_test.go b/services/core/internal/runtimegateway/workspace_directory_test.go index e95e7aae9..aa104ea22 100644 --- a/services/core/internal/runtimegateway/workspace_directory_test.go +++ b/services/core/internal/runtimegateway/workspace_directory_test.go @@ -28,7 +28,7 @@ func TestWorkspaceDirectorySharesReadCorrelationAndFrameBound(t *testing.T) { directory.Entries[i] = proto.WorkspaceDirectoryEntry{Name: strings.Repeat("\x01", 250) + fmt.Sprintf("%04d", i), Kind: "directory"} } result := proto.WorkspaceReadResultPayload{Outcome: "completed", CloseAcknowledged: true, Directory: directory} - reply, _ := proto.NewEnvelope(proto.TypeWorkspaceReadResult, message.ID, result) + reply, _ := message.Reply(proto.TypeWorkspaceReadResult, result) encoded, err := json.Marshal(reply) if err != nil || int64(len(encoded)) >= ReadLimit { t.Fatal("directory result exceeds frame", len(encoded), err) @@ -83,7 +83,7 @@ func TestWorkspaceDirectoryRequiresExplicitWireTruncation(t *testing.T) { done <- err }() request := <-s.sendCh - reply := proto.Envelope{Type: proto.TypeWorkspaceReadResult, ID: request.ID, + reply := proto.Envelope{Type: proto.TypeWorkspaceReadResult, ID: request.ID, Assignment: request.Assignment, Payload: json.RawMessage(`{"outcome":"completed","close_acknowledged":true,"directory":` + tc.directory + `}`)} s.dispatch(reply) if err := <-done; (err == nil) != tc.valid { diff --git a/services/core/internal/runtimegateway/workspace_export.go b/services/core/internal/runtimegateway/workspace_export.go index 273e504dd..69069b8ae 100644 --- a/services/core/internal/runtimegateway/workspace_export.go +++ b/services/core/internal/runtimegateway/workspace_export.go @@ -79,34 +79,25 @@ func (r *workspaceExportReader) Read(p []byte) (int, error) { return 0, err } env.Assignment = r.ref - if err = r.peer.Send(r.ctx, env); err != nil { + reply, err := r.peer.exchangeFrame(r.ctx, env, r.replies) + if err != nil { return 0, err } - select { - case env, ok := <-r.replies: - if !ok { - return 0, ErrSessionClosed - } - var result proto.WorkspaceExportResultPayload - if len(env.Payload) > proto.WorkspaceExportMaxFrameBytes || env.DecodePayload(&result) != nil || result.Offset != r.offset { - return 0, errors.New("agentdaemon gateway: invalid export receipt") - } - if result.Outcome == "completed" && len(result.Data) == 0 && result.ErrorCode == "" { - r.completed = true - return 0, io.EOF - } - if result.Outcome != "chunk" || result.ErrorCode != "" || len(result.Data) == 0 || len(result.Data) > proto.WorkspaceExportChunkBytes || int64(len(result.Data)) > proto.WorkspaceExportMaxBytes-r.offset { - return 0, errors.New("agentdaemon gateway: workspace export incomplete") - } - r.data = result.Data - r.offset += int64(len(result.Data)) - r.request = proto.WorkspaceExportPayload{Step: "next", Offset: r.offset} - return r.Read(p) - case <-r.ctx.Done(): - return 0, r.ctx.Err() - case <-r.peer.closed: - return 0, ErrSessionClosed + var result proto.WorkspaceExportResultPayload + if len(reply.Payload) > proto.WorkspaceExportMaxFrameBytes || reply.DecodePayload(&result) != nil || result.Offset != r.offset { + return 0, errors.New("agentdaemon gateway: invalid export receipt") + } + if result.Outcome == "completed" && len(result.Data) == 0 && result.ErrorCode == "" { + r.completed = true + return 0, io.EOF + } + if result.Outcome != "chunk" || result.ErrorCode != "" || len(result.Data) == 0 || len(result.Data) > proto.WorkspaceExportChunkBytes || int64(len(result.Data)) > proto.WorkspaceExportMaxBytes-r.offset { + return 0, errors.New("agentdaemon gateway: workspace export incomplete") } + r.data = result.Data + r.offset += int64(len(result.Data)) + r.request = proto.WorkspaceExportPayload{Step: "next", Offset: r.offset} + return r.Read(p) } func (s *Session) dispatchWorkspaceExport(env proto.Envelope) { diff --git a/services/core/internal/runtimegateway/workspace_export_test.go b/services/core/internal/runtimegateway/workspace_export_test.go index 6b6efeae0..6a83e6bbc 100644 --- a/services/core/internal/runtimegateway/workspace_export_test.go +++ b/services/core/internal/runtimegateway/workspace_export_test.go @@ -17,6 +17,7 @@ func exportReply(t *testing.T, s *Session, id string, result proto.WorkspaceExpo if err != nil { t.Fatal(err) } + env.Assignment = testAssignment s.dispatch(env) } diff --git a/services/core/internal/runtimegateway/workspace_read.go b/services/core/internal/runtimegateway/workspace_read.go index f0846ac05..cbb912f66 100644 --- a/services/core/internal/runtimegateway/workspace_read.go +++ b/services/core/internal/runtimegateway/workspace_read.go @@ -51,23 +51,14 @@ func (s *Session) readWorkspace(ctx context.Context, ref proto.AssignmentRef, re defer func() { s.workspaceReadMu.Lock(); delete(s.workspaceReads, id); s.workspaceReadMu.Unlock() }() ctx, cancel := context.WithTimeout(ctx, 17*time.Second) defer cancel() - if err = s.Send(ctx, env); err != nil { + reply, err := s.exchangeFrame(ctx, env, replies) + if err != nil { return result, err } - select { - case reply, ok := <-replies: - if !ok { - return result, ErrSessionClosed - } - if reply.DecodePayload(&result) != nil || !validWorkspaceOperationResult(result, request) { - return proto.WorkspaceReadResultPayload{}, errors.New("agentdaemon gateway: invalid workspace read response") - } - return result, nil - case <-ctx.Done(): - return result, ctx.Err() - case <-s.closed: - return result, ErrSessionClosed + if reply.DecodePayload(&result) != nil || !validWorkspaceOperationResult(result, request) { + return proto.WorkspaceReadResultPayload{}, errors.New("agentdaemon gateway: invalid workspace read response") } + return result, nil } func validWorkspaceOperationResult(result proto.WorkspaceReadResultPayload, request proto.WorkspaceReadPayload) bool { @@ -99,7 +90,7 @@ func validWorkspaceReadResult(result proto.WorkspaceReadResultPayload, limit int return false } switch result.ErrorCode { - case "invalid_request", "resource_unavailable", "read_capacity", "read_unsupported", "not_found", "permission_denied": + case "invalid_request", "resource_unavailable", "read_capacity", "read_unsupported", "not_found", "permission_denied", proto.AssignmentStale, proto.AssignmentConflict: return true default: return false diff --git a/services/core/internal/runtimegateway/workspace_read_test.go b/services/core/internal/runtimegateway/workspace_read_test.go index c7de9567f..6d3c89550 100644 --- a/services/core/internal/runtimegateway/workspace_read_test.go +++ b/services/core/internal/runtimegateway/workspace_read_test.go @@ -32,7 +32,7 @@ func TestWorkspaceReadCorrelatesOneBoundedResult(t *testing.T) { result := proto.WorkspaceReadResultPayload{Outcome: "completed", Data: data, Truncated: true, CloseAcknowledged: true} foreign, _ := proto.NewEnvelope(proto.TypeWorkspaceReadResult, "other-operation", result) s.dispatch(foreign) - reply, _ := proto.NewEnvelope(proto.TypeWorkspaceReadResult, request.ID, result) + reply, _ := request.Reply(proto.TypeWorkspaceReadResult, result) encoded, err := json.Marshal(reply) if err != nil || int64(len(encoded)) >= ReadLimit { t.Fatal("result exceeds existing transport frame", err, len(encoded)) @@ -61,7 +61,7 @@ func TestWorkspaceReadRejectsIncompleteOrContradictoryReplies(t *testing.T) { done <- err }() request := <-s.sendCh - reply, _ := proto.NewEnvelope(proto.TypeWorkspaceReadResult, request.ID, result) + reply, _ := request.Reply(proto.TypeWorkspaceReadResult, result) s.dispatch(reply) if err := <-done; err == nil { t.Fatal("invalid result accepted", result.Outcome) @@ -84,6 +84,8 @@ func TestWorkspaceDirectoryAcceptsNotDirectoryOnlyForDirectoryReads(t *testing.T {true, proto.WorkspaceReadResultPayload{Outcome: "unknown", ErrorCode: proto.WorkspaceReadNotDirectory}, false}, {true, proto.WorkspaceReadResultPayload{Outcome: "completed", CloseAcknowledged: true, ErrorCode: proto.WorkspaceReadNotDirectory, Directory: &proto.WorkspaceDirectoryResult{Entries: []proto.WorkspaceDirectoryEntry{}}}, false}, {false, proto.WorkspaceReadResultPayload{Outcome: "rejected", ErrorCode: proto.WorkspaceReadNotDirectory}, false}, + {false, proto.WorkspaceReadResultPayload{Outcome: "rejected", ErrorCode: proto.AssignmentStale}, true}, + {true, proto.WorkspaceReadResultPayload{Outcome: "rejected", ErrorCode: proto.AssignmentConflict}, true}, } { s := NewSession(newFakeConn(), "device", "tenant", "test", nil, nil) done := make(chan error, 1) @@ -97,7 +99,7 @@ func TestWorkspaceDirectoryAcceptsNotDirectoryOnlyForDirectoryReads(t *testing.T done <- err }() request := <-s.sendCh - reply, _ := proto.NewEnvelope(proto.TypeWorkspaceReadResult, request.ID, test.result) + reply, _ := request.Reply(proto.TypeWorkspaceReadResult, test.result) s.dispatch(reply) if err := <-done; (err == nil) != test.accepted { t.Fatal("directory result validation changed", test.directory, test.result, err) @@ -117,7 +119,7 @@ func TestWorkspaceReadObserverCancellationDoesNotSendCancelOrRetry(t *testing.T) if err := <-done; !errors.Is(err, context.Canceled) { t.Fatal(err) } - reply, _ := proto.NewEnvelope(proto.TypeWorkspaceReadResult, request.ID, proto.WorkspaceReadResultPayload{Outcome: "unknown", ErrorCode: "read_unconfirmed"}) + reply, _ := request.Reply(proto.TypeWorkspaceReadResult, proto.WorkspaceReadResultPayload{Outcome: "unknown", ErrorCode: "read_unconfirmed"}) s.dispatch(reply) select { case extra := <-s.sendCh: diff --git a/services/core/internal/runtimegateway/workspace_write.go b/services/core/internal/runtimegateway/workspace_write.go index 109c5b960..3c8d888b8 100644 --- a/services/core/internal/runtimegateway/workspace_write.go +++ b/services/core/internal/runtimegateway/workspace_write.go @@ -45,7 +45,7 @@ func (s *Session) WriteWorkspaceFile(ctx context.Context, id string, ref proto.A return empty, errors.New("agentdaemon gateway: invalid write frame") } env.Assignment = ref - reply, err := s.exchangeChunkFrame(ctx, env, replies) + reply, err := s.exchangeFrame(ctx, env, replies) if err != nil { return empty, err } @@ -79,7 +79,7 @@ func validWorkspaceWriteResult(r proto.WorkspaceWriteResultPayload, expected str return false } switch r.ErrorCode { - case "invalid_request", "resource_unavailable", "write_capacity", "write_unsupported", "write_rejected": + case "invalid_request", "resource_unavailable", "write_capacity", "write_unsupported", "write_rejected", proto.AssignmentStale, proto.AssignmentConflict: return true default: return false diff --git a/services/core/internal/runtimegateway/workspace_write_test.go b/services/core/internal/runtimegateway/workspace_write_test.go index 17fca00be..dfa983b30 100644 --- a/services/core/internal/runtimegateway/workspace_write_test.go +++ b/services/core/internal/runtimegateway/workspace_write_test.go @@ -56,6 +56,7 @@ func TestWorkspaceWriteChunksAndCorrelatesReceipt(t *testing.T) { foreign, _ := proto.NewEnvelope(proto.TypeWorkspaceWriteResult, uuid.NewString(), result) s.dispatch(foreign) reply, _ := proto.NewEnvelope(proto.TypeWorkspaceWriteResult, id, result) + reply.Assignment = testAssignment s.dispatch(reply) if p.Step == "commit" { break @@ -112,6 +113,12 @@ func TestWorkspaceWriteRejectsPrematureOrContradictoryReceipts(t *testing.T) { t.Fatal("known rejection refused", reason) } } + // A release that fences the write settles it as rejected. + for _, code := range []string{proto.AssignmentStale, proto.AssignmentConflict} { + if !validWorkspaceWriteResult(proto.WorkspaceWriteResultPayload{Outcome: "rejected", ErrorCode: code}, "ready", 0, 4) { + t.Fatal("assignment rejection refused", code) + } + } if validWorkspaceWriteResult(proto.WorkspaceWriteResultPayload{Outcome: "received", Offset: 1}, "received", 2, 4) { t.Fatal("wrong offset accepted") } diff --git a/services/core/migrations/000092_session_runtime_assignments.sql b/services/core/migrations/000092_session_runtime_assignments.sql index 9878d1fbc..dd9243aac 100644 --- a/services/core/migrations/000092_session_runtime_assignments.sql +++ b/services/core/migrations/000092_session_runtime_assignments.sql @@ -1,7 +1,8 @@ -- +goose Up -- A Session's binding to its Runtime becomes a fenced assignment. Core -- advances epoch with each change of desired_state; applied_epoch is the --- latest released epoch the Runtime acknowledged. +-- latest released epoch the Runtime acknowledged, or that no Runtime can act +-- on because its Runtime lost authority. ALTER TABLE session_devices RENAME TO session_runtime_assignments; ALTER TABLE session_runtime_assignments RENAME COLUMN device_id TO runtime_id; ALTER INDEX session_devices_pkey RENAME TO session_runtime_assignments_pkey; @@ -16,7 +17,8 @@ ALTER TABLE session_runtime_assignments ADD CONSTRAINT session_runtime_assignments_applied CHECK (applied_epoch BETWEEN 0 AND epoch), ADD CONSTRAINT session_runtime_assignments_release CHECK (desired_state = 'released' OR NOT remove_home); -- A deleted Session owes its Runtime the release that removes its home. -UPDATE session_runtime_assignments a SET desired_state = 'released', remove_home = true, epoch = 2 +UPDATE session_runtime_assignments a SET desired_state = 'released', remove_home = true, epoch = 2, + applied_epoch = CASE WHEN EXISTS (SELECT 1 FROM runtime_device_authority r WHERE r.id = a.runtime_id) THEN 0 ELSE 2 END FROM sessions s WHERE s.id = a.session_id AND s.deleted_at IS NOT NULL; CREATE INDEX session_runtime_assignments_pending_idx ON session_runtime_assignments (runtime_id, session_id) WHERE desired_state = 'released' AND applied_epoch < epoch; diff --git a/services/core/tests/integration/assignment_release_test.go b/services/core/tests/integration/assignment_release_test.go index b496ef0e3..bc1e2960a 100644 --- a/services/core/tests/integration/assignment_release_test.go +++ b/services/core/tests/integration/assignment_release_test.go @@ -64,3 +64,44 @@ func TestDeletionReleaseReachesReconnectedRuntime(t *testing.T) { } awaitDaemonRemoteCondition(t, t.Context(), 3*time.Second, "release acknowledged", func() bool { return applied() == 2 }) } + +// TestRevocationSettlesUndeliverableReleases checks that a release whose +// Runtime was revoked is settled with the revocation instead of staying +// pending: no Runtime is left to act on it. +func TestRevocationSettlesUndeliverableReleases(t *testing.T) { + h := newDispatchHarness(t) + ctx := t.Context() + second, err := h.s.CreateSession(ctx, h.tenant, WithFixtureModelProvider(sessions.CreateSession{Creator: FixtureCreator(), Engine: "codex", IdempotencyKey: "second", Configuration: []byte(`{"agent":{"model":"test-model"},"environment":{"type":"none"}}`)})) + if err != nil { + t.Fatal(err) + } + if err := bindSessionDevice(t, h.s, h.tenant, second.ID, h.device.ID); err != nil { + t.Fatal(err) + } + service := sessionService(t, h.s) + settled := func(session string) bool { + var epoch, applied int64 + if err := h.s.pool.QueryRow(ctx, "SELECT epoch, applied_epoch FROM session_runtime_assignments WHERE session_id=$1", session).Scan(&epoch, &applied); err != nil { + t.Fatal(err) + } + return epoch == 2 && applied == 2 + } + if err := service.DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}); err != nil { + t.Fatal(err) + } + if settled(h.session.ID) { + t.Fatal("a release to an authorized Runtime was settled before delivery") + } + if err := service.RevokeDevice(ctx, h.tenant, h.device.ID); err != nil { + t.Fatal(err) + } + if !settled(h.session.ID) { + t.Fatal("revocation left the Runtime's release pending") + } + if err := service.DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: second.ID}); err != nil { + t.Fatal(err) + } + if !settled(second.ID) { + t.Fatal("a release to a revoked Runtime was left pending") + } +} diff --git a/services/core/tests/integration/dispatch_test.go b/services/core/tests/integration/dispatch_test.go index f675cbcb2..e573ffbbb 100644 --- a/services/core/tests/integration/dispatch_test.go +++ b/services/core/tests/integration/dispatch_test.go @@ -27,6 +27,7 @@ import ( type dispatchHarness struct { writeMu sync.Mutex + assignment proto.AssignmentRef // the latest assignment Core named; writeMu guards it admissions map[string]fixtureAdmission t *testing.T s *Store @@ -145,6 +146,9 @@ func (h *dispatchHarness) write(run, kind string, payload any) { if err != nil { h.t.Fatal(err) } + if kind != proto.TypeHeartbeat { + env.Assignment = h.assignment + } if err = h.conn.WriteJSON(env); err != nil { h.t.Fatal(err) } @@ -158,6 +162,7 @@ func (h *dispatchHarness) read(kind string) proto.Envelope { if err := h.conn.ReadJSON(&env); err != nil { h.t.Fatal(err) } + h.observe(env) if kind != proto.TypeAssignmentBind && h.assignmentFrame(env) { continue } diff --git a/services/core/tests/integration/executor_fixture_test.go b/services/core/tests/integration/executor_fixture_test.go index 3e75c0e25..1650b6522 100644 --- a/services/core/tests/integration/executor_fixture_test.go +++ b/services/core/tests/integration/executor_fixture_test.go @@ -26,8 +26,20 @@ func assignmentReply(env proto.Envelope) (proto.Envelope, bool) { return reply, err == nil } -// assignmentFrame answers an assignment_bind and reports whether env was one. +// observe remembers the assignment env names. The fixture Runtime writes its +// frames under it, as a Runtime echoes its request's assignment. +func (h *dispatchHarness) observe(env proto.Envelope) { + if env.Assignment.Valid() { + h.writeMu.Lock() + h.assignment = env.Assignment + h.writeMu.Unlock() + } +} + +// assignmentFrame observes env, answers an assignment_bind and reports whether +// env was one. func (h *dispatchHarness) assignmentFrame(env proto.Envelope) bool { + h.observe(env) reply, ok := assignmentReply(env) if !ok { return false diff --git a/services/core/tests/integration/runtime_compute_lifecycle_test.go b/services/core/tests/integration/runtime_compute_lifecycle_test.go index f24a56fdb..10eec20b5 100644 --- a/services/core/tests/integration/runtime_compute_lifecycle_test.go +++ b/services/core/tests/integration/runtime_compute_lifecycle_test.go @@ -248,7 +248,7 @@ func (p *fakeCheckpointProvider) connect(ctx context.Context, b sandbox.Bootstra if env.Type == proto.TypeEnvironmentQuiesce { kind = proto.TypeEnvironmentQuiesced } - reply, _ := proto.NewEnvelope(kind, env.ID, proto.EnvironmentSuspendResultPayload{EnvironmentID: request.EnvironmentID, SuspendID: request.SuspendID, Accepted: !reject}) + reply, _ := env.Reply(kind, proto.EnvironmentSuspendResultPayload{EnvironmentID: request.EnvironmentID, SuspendID: request.SuspendID, Accepted: !reject}) if conn.WriteJSON(reply) != nil { return } diff --git a/services/core/tests/integration/runtime_initialization_peer_test.go b/services/core/tests/integration/runtime_initialization_peer_test.go index 9fed3cc5b..35e6163dd 100644 --- a/services/core/tests/integration/runtime_initialization_peer_test.go +++ b/services/core/tests/integration/runtime_initialization_peer_test.go @@ -128,5 +128,5 @@ func (x *initializationTransfer) receive(env proto.Envelope) (proto.Envelope, er result.SizeBytes = len(x.data) } } - return proto.NewEnvelope(proto.TypeRuntimePrepareResult, env.ID, result) + return env.Reply(proto.TypeRuntimePrepareResult, result) } From 4fa91ee2f47a13c5d37b1069b8df40bcc35adffd Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 12:05:28 +0000 Subject: [PATCH 3/4] Document release fencing and echo assignments in integration fakes Cover the release fence on an unfinished workspace write, make the integration Runtime fakes echo the assignment of each request and Run, and document echoed assignments, release fencing, settlement without authority and release backoff. --- .../internal/dispatch/workspace_write_test.go | 38 +++++++++++++++++++ docs/runtime-protocol.md | 6 +-- docs/zh/runtime-protocol.md | 8 ++-- .../internal/runtimegateway/exchange_test.go | 5 ++- .../core/tests/integration/dispatch_test.go | 7 +++- .../integration/executor_fixture_test.go | 21 +++++++--- 6 files changed, 69 insertions(+), 16 deletions(-) diff --git a/apps/daemon/internal/dispatch/workspace_write_test.go b/apps/daemon/internal/dispatch/workspace_write_test.go index 132773cd4..4be853b32 100644 --- a/apps/daemon/internal/dispatch/workspace_write_test.go +++ b/apps/daemon/internal/dispatch/workspace_write_test.go @@ -168,3 +168,41 @@ func TestLocalUploadReportsDestinationConflictsAndReleasesOwner(t *testing.T) { }) } } + +func TestReleaseFencesUnfinishedWorkspaceWrite(t *testing.T) { + r, sender, request, workspace := localWriterRouter(t) + id := uuid.NewString() + if err := r.Handle(t.Context(), mustEnv(t, proto.TypeWorkspaceWrite, id, request)); err != nil { + t.Fatal(err) + } + waitWorkspaceWrite(t, sender, id, "ready") + if err := r.Handle(t.Context(), mustEnv(t, proto.TypeWorkspaceWrite, id, proto.WorkspaceWritePayload{Step: "chunk", Data: []byte("abc")})); err != nil { + t.Fatal(err) + } + waitWorkspaceWrite(t, sender, id, "received") + release(t, r, preparationSessionID, "release", 2, false) + if got := waitWorkspaceWrite(t, sender, id, "rejected"); got.ErrorCode != proto.AssignmentStale { + t.Fatal("release did not fence the write", got) + } + if got := waitAssignmentStatus(t, sender, "release"); got.State != proto.AssignmentReleased { + t.Fatal(got) + } + // The release replies only after the write's result. + for _, frame := range sender.snapshot() { + if frame.Type == proto.TypeAssignmentStatus && frame.ID == "release" { + t.Fatal("release replied before the write settled") + } + if frame.Type == proto.TypeWorkspaceWriteResult && frame.ID == id && frame.Assignment == ref(preparationSessionID) { + var result proto.WorkspaceWriteResultPayload + if frame.DecodePayload(&result) == nil && result.Outcome == "rejected" { + break + } + } + } + if err := r.Handle(t.Context(), mustEnv(t, proto.TypeWorkspaceWrite, id, proto.WorkspaceWritePayload{Step: "commit"})); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(workspace, "file")); !os.IsNotExist(err) { + t.Fatal("a released assignment's write applied", err) + } +} diff --git a/docs/runtime-protocol.md b/docs/runtime-protocol.md index 78f1d9bc2..62f9a2e5b 100644 --- a/docs/runtime-protocol.md +++ b/docs/runtime-protocol.md @@ -115,11 +115,11 @@ Usage frames and the final usage snapshot each carry the cumulative measurement An assignment binds one Session to the Runtime that runs it. `Envelope.assignment` names it as `session_id`, `assignment_id` and `epoch`, and is the only place a frame carries it. Core advances the epoch whenever it changes the assignment's desired state, so a lower epoch is stale. -Every Session frame carries the assignment: `execution_prepare`, `execution_start` and `execution_release`; `prompt_cancel`, `prompt_steer`, `function_result`, `permission_decision` and `prompt_for_user_choice_decision`; every frame of `runtime_prepare`, `workspace_read`, `workspace_write` and `workspace_export`; and `environment_quiesce` and `environment_resume`. A reply echoes its request's assignment. Heartbeats carry none. +Every Session frame carries the assignment: `execution_prepare`, `execution_start` and `execution_release`; `prompt_cancel`, `prompt_steer`, `function_result`, `permission_decision` and `prompt_for_user_choice_decision`; every frame of `runtime_prepare`, `workspace_read`, `workspace_write` and `workspace_export`; and `environment_quiesce` and `environment_resume`. A reply echoes its request's assignment, and a Run's frames carry the assignment that started it; Core rejects a reply or Run frame that names another. Heartbeats carry none. -Before a Session's first operation on a connection, including Environment initialization and file work without a Turn, Core sends `assignment_bind` with the Session's Environment ID and waits for `assignment_status` `bound`. A repeated bind of the same assignment is `bound` again. The Runtime admits a Session frame only under the assignment it bound: an older epoch, or a released one, fails with `assignment_stale`; another assignment, Session or Environment fails with `assignment_conflict`. A started Run's frames, including its cancellation receipt, stay admissible under the assignment that started it until the release. +Before a Session's first operation on a connection, including Environment initialization and file work without a Turn, Core sends `assignment_bind` with the Session's Environment ID and waits for `assignment_status` `bound`. A repeated bind of the same assignment is `bound` again. The Runtime admits a Session frame only under the assignment it bound: an older epoch, or a released one, fails with `assignment_stale`; another assignment, Session or Environment fails with `assignment_conflict`. A started Run's frames, including its cancellation receipt, stay admissible under the assignment that started it until the release. A repeated function result or decision whose receipt the Runtime already recorded is answered only under the assignment that applied it; another fails with `assignment_conflict`. -Core records a release and advances the epoch before it sends anything. Deleting a Session releases its assignment with `remove_home: true`; releasing its Environment sends `false`. A deletion never revokes a shared Runtime credential. `assignment_release` fences the assignment at once. The Runtime then stops the Session's work, closes its Executors and, when asked, removes the native home; only then does it reply `released` or `home_removed`. Unfinished cleanup replies `failed` with `cleanup_unconfirmed`, and a retry at the same epoch repeats it. A Runtime that declares `home_removal` unsupported answers `remove_home: true` with `unsupported_operation`, and Core asks it only to release. Core records the release as applied only from a matching `released` or `home_removed`, and resends every unacknowledged release to a Runtime when it connects. A quiesced Runtime admits only a release and the matching `environment_resume`, which carries the assignment that quiesced it. +Core records a release and advances the epoch before it sends anything. Deleting a Session releases its assignment with `remove_home: true`; releasing its Environment sends `false`. A deletion never revokes a shared Runtime credential. `assignment_release` fences the assignment at once. The Runtime then stops the Session's work: a transfer still receiving its body, or committed but not yet applied, ends with `assignment_stale`; it releases read-only preparations and waits until every workspace read, write, export and Runtime preparation has sent its result. It closes the Session's Executors and, when asked, removes the native home; only then does it reply `released` or `home_removed`. Unfinished cleanup replies `failed` with `cleanup_unconfirmed`, and a retry at the same epoch repeats it. A Runtime that declares `home_removal` unsupported answers `remove_home: true` with `unsupported_operation`, and Core asks it only to release. Core records the release as applied from a matching `released` or `home_removed`, or at once when no Runtime is left to act on it: a release to a Runtime without authority is settled when recorded, and revoking a Runtime settles its releases. Core resends every unacknowledged release to a Runtime when it connects and backs off a release that fails. A quiesced Runtime admits only a release and the matching `environment_resume`, which carries the assignment that quiesced it. The Runtime answers a Core frame it cannot route with `protocol_error`, which echoes the request's ID and carries its type and an error code. diff --git a/docs/zh/runtime-protocol.md b/docs/zh/runtime-protocol.md index d40b11f62..6c8cbd392 100644 --- a/docs/zh/runtime-protocol.md +++ b/docs/zh/runtime-protocol.md @@ -1,7 +1,7 @@ --- title: "Core–Runtime 协议" source: docs/runtime-protocol.md -source_hash: ca1f05d7f95e6a91f57e44919da83d05e94b146809d727a4e87ad38c94a9d6fd +source_hash: 84a2ec5b9308a5e1a118b05c1ea1c92f8741871785524f3878321286b7e2d25b --- 此协议在 Runtime daemon 获取机器凭据后连接 Core 与 daemon，定义 daemon 连接上消息的含义和顺序。wire 类型、限制和验证器仅在 [`internal/agentdaemon/proto`](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/internal/agentdaemon/proto) 中定义一次；Core 的 [gateway](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/services/core/internal/runtimegateway) 与参考 Runtime 的 [dispatcher](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/apps/daemon/internal/dispatch) 都使用它们，因此无需同步第二套 payload schema。签发凭据和打开连接的 HTTP 路由见[机器连接 API](../../contracts/agents-api/zh/machine-api.md)。 @@ -117,11 +117,11 @@ Usage frame 和最终 usage snapshot 都携带当前执行的累计测量，替 分配（assignment）把一个 Session 绑定到运行它的 Runtime。`Envelope.assignment` 以 `session_id`、`assignment_id` 和 `epoch` 命名它，这是 frame 携带分配的唯一位置。Core 每次改变分配的期望状态时推进 epoch，因此较低的 epoch 是陈旧的。 -每个 Session frame 都携带分配：`execution_prepare`、`execution_start` 和 `execution_release`；`prompt_cancel`、`prompt_steer`、`function_result`、`permission_decision` 和 `prompt_for_user_choice_decision`；`runtime_prepare`、`workspace_read`、`workspace_write` 和 `workspace_export` 的每个 frame；以及 `environment_quiesce` 和 `environment_resume`。回复回显请求的分配。heartbeat 不携带分配。 +每个 Session frame 都携带分配：`execution_prepare`、`execution_start` 和 `execution_release`；`prompt_cancel`、`prompt_steer`、`function_result`、`permission_decision` 和 `prompt_for_user_choice_decision`；`runtime_prepare`、`workspace_read`、`workspace_write` 和 `workspace_export` 的每个 frame；以及 `environment_quiesce` 和 `environment_resume`。回复回显请求的分配，Run 的 frame 携带启动它的分配；Core 拒绝指明其他分配的回复或 Run frame。heartbeat 不携带分配。 -在一条连接上执行 Session 的第一个操作之前，包括没有 Turn 的 Environment 初始化和文件操作，Core 发送带 Session 的 Environment ID 的 `assignment_bind`，并等待 `assignment_status` `bound`。重复绑定同一分配仍得到 `bound`。Runtime 只在其已绑定的分配下准入 Session frame：较旧的 epoch 或已释放的分配以 `assignment_stale` 失败；其他分配、Session 或 Environment 以 `assignment_conflict` 失败。已启动 Run 的 frame，包括其取消回执，在释放前仍可在启动它的分配下准入。 +在一条连接上执行 Session 的第一个操作之前，包括没有 Turn 的 Environment 初始化和文件操作，Core 发送带 Session 的 Environment ID 的 `assignment_bind`，并等待 `assignment_status` `bound`。重复绑定同一分配仍得到 `bound`。Runtime 只在其已绑定的分配下准入 Session frame：较旧的 epoch 或已释放的分配以 `assignment_stale` 失败；其他分配、Session 或 Environment 以 `assignment_conflict` 失败。已启动 Run 的 frame，包括其取消回执，在释放前仍可在启动它的分配下准入。Runtime 已记录回执的重复函数结果或决策只在应用它的分配下得到回答；其他分配以 `assignment_conflict` 失败。 -Core 先记录释放并推进 epoch，再发送任何消息。删除 Session 以 `remove_home: true` 释放其分配；释放其 Environment 发送 `false`。删除从不吊销共享的 Runtime 凭据。`assignment_release` 立即约束该分配。随后 Runtime 停止 Session 的工作，关闭其 Executor，并在要求时删除原生 home；此后才回复 `released` 或 `home_removed`。未完成的清理回复 `failed` 和 `cleanup_unconfirmed`，同一 epoch 的重试会重复清理。声明 `home_removal` 不支持的 Runtime 以 `unsupported_operation` 回答 `remove_home: true`，Core 只要求它释放。Core 只根据匹配的 `released` 或 `home_removed` 记录释放已应用，并在 Runtime 连接时重发所有未确认的释放。已 quiesce 的 Runtime 只准入释放和匹配的 `environment_resume`，后者携带使其 quiesce 的分配。 +Core 先记录释放并推进 epoch，再发送任何消息。删除 Session 以 `remove_home: true` 释放其分配；释放其 Environment 发送 `false`。删除从不吊销共享的 Runtime 凭据。`assignment_release` 立即约束该分配。随后 Runtime 停止 Session 的工作：仍在接收内容、或已提交但尚未应用的传输以 `assignment_stale` 结束；它释放只读准备，并等待每个 workspace 读取、写入、导出和 Runtime 准备发送结果。它关闭 Session 的 Executor，并在要求时删除原生 home；此后才回复 `released` 或 `home_removed`。未完成的清理回复 `failed` 和 `cleanup_unconfirmed`，同一 epoch 的重试会重复清理。声明 `home_removal` 不支持的 Runtime 以 `unsupported_operation` 回答 `remove_home: true`，Core 只要求它释放。Core 根据匹配的 `released` 或 `home_removed` 记录释放已应用；没有 Runtime 能处理该释放时立即记录：发给无授权 Runtime 的释放在记录时即结清，吊销 Runtime 会结清它的释放。Core 在 Runtime 连接时重发所有未确认的释放，并对失败的释放退避重试。已 quiesce 的 Runtime 只准入释放和匹配的 `environment_resume`，后者携带使其 quiesce 的分配。 Runtime 对无法路由的 Core frame 回复 `protocol_error`，回显请求 ID，并携带其类型和错误码。 diff --git a/services/core/internal/runtimegateway/exchange_test.go b/services/core/internal/runtimegateway/exchange_test.go index a43650da4..87d5db347 100644 --- a/services/core/internal/runtimegateway/exchange_test.go +++ b/services/core/internal/runtimegateway/exchange_test.go @@ -15,7 +15,10 @@ func TestFramesMustNameTheRequestAssignment(t *testing.T) { s := NewSession(newFakeConn(), "device", "tenant", "test", NewRegistry(), nil) defer s.Close("test") read := make(chan error, 1) - go func() { _, err := s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()); read <- err }() + go func() { + _, err := s.ReadWorkspaceFile(t.Context(), testAssignment, workspaceReadRequest()) + read <- err + }() request := <-s.sendCh reply, _ := request.Reply(proto.TypeWorkspaceReadResult, proto.WorkspaceReadResultPayload{Outcome: "completed", CloseAcknowledged: true}) reply.Assignment = ref diff --git a/services/core/tests/integration/dispatch_test.go b/services/core/tests/integration/dispatch_test.go index e573ffbbb..44f5a9833 100644 --- a/services/core/tests/integration/dispatch_test.go +++ b/services/core/tests/integration/dispatch_test.go @@ -27,7 +27,8 @@ import ( type dispatchHarness struct { writeMu sync.Mutex - assignment proto.AssignmentRef // the latest assignment Core named; writeMu guards it + assignments map[string]proto.AssignmentRef // by frame and Run ID; writeMu guards it + assignment proto.AssignmentRef // the latest assignment Core named admissions map[string]fixtureAdmission t *testing.T s *Store @@ -146,7 +147,9 @@ func (h *dispatchHarness) write(run, kind string, payload any) { if err != nil { h.t.Fatal(err) } - if kind != proto.TypeHeartbeat { + if ref, ok := h.assignments[run]; ok { + env.Assignment = ref + } else if kind != proto.TypeHeartbeat { env.Assignment = h.assignment } if err = h.conn.WriteJSON(env); err != nil { diff --git a/services/core/tests/integration/executor_fixture_test.go b/services/core/tests/integration/executor_fixture_test.go index 1650b6522..0cde028d1 100644 --- a/services/core/tests/integration/executor_fixture_test.go +++ b/services/core/tests/integration/executor_fixture_test.go @@ -26,13 +26,22 @@ func assignmentReply(env proto.Envelope) (proto.Envelope, bool) { return reply, err == nil } -// observe remembers the assignment env names. The fixture Runtime writes its -// frames under it, as a Runtime echoes its request's assignment. +// observe remembers the assignment env names for its ID and started Run. The +// fixture Runtime writes its frames under it, as a Runtime echoes its +// request's assignment. func (h *dispatchHarness) observe(env proto.Envelope) { - if env.Assignment.Valid() { - h.writeMu.Lock() - h.assignment = env.Assignment - h.writeMu.Unlock() + if !env.Assignment.Valid() { + return + } + h.writeMu.Lock() + defer h.writeMu.Unlock() + if h.assignments == nil { + h.assignments = make(map[string]proto.AssignmentRef) + } + h.assignment, h.assignments[env.ID] = env.Assignment, env.Assignment + var start proto.ExecutionStartPayload + if env.Type == proto.TypeExecutionStart && env.DecodePayload(&start) == nil && start.RunID != "" { + h.assignments[start.RunID] = env.Assignment } } From 063c6202ecfc2d271a2a726f7a58cb84d391583e Mon Sep 17 00:00:00 2001 From: SaladDay <1203511142@qq.com> Date: Wed, 7 Oct 2026 12:43:30 +0000 Subject: [PATCH 4/4] Order releases after results, serialize them with revocation and run them fairly A release now waits for every read, write, export and Runtime preparation of its Session until that work has sent its terminal result, counted on the assignment, and a canceled export still answers the request Core awaits on its own send budget. Revocation and release creation take the device row lock before either reads the other's state, so a release racing a revocation of its shared Runtime is settled. Due releases run longest-due first, so releases that keep failing cannot starve the rest. --- apps/daemon/internal/dispatch/assignment.go | 62 +++++++------- apps/daemon/internal/dispatch/router.go | 2 +- apps/daemon/internal/dispatch/router_test.go | 4 + .../internal/dispatch/runtime_preparation.go | 12 +-- .../dispatch/runtime_preparation_test.go | 10 +-- apps/daemon/internal/dispatch/suspend_test.go | 2 +- .../internal/dispatch/workspace_export.go | 37 ++++++--- .../dispatch/workspace_export_test.go | 26 ++++++ .../internal/dispatch/workspace_read.go | 14 +--- .../internal/dispatch/workspace_write.go | 11 ++- .../internal/dispatch/workspace_write_test.go | 35 ++++++++ docs/runtime-protocol.md | 2 +- docs/zh/runtime-protocol.md | 4 +- services/core/internal/db/queries/devices.sql | 24 ++++-- services/core/internal/db/sqlc/devices.sql.go | 38 +++++++-- .../internal/execution/assignment_releases.go | 22 +++-- .../execution/assignment_releases_test.go | 43 +++++++--- .../persistence/postgres/sessionpg/devices.go | 19 +++-- .../integration/assignment_release_test.go | 83 +++++++++++++++++++ 19 files changed, 332 insertions(+), 118 deletions(-) diff --git a/apps/daemon/internal/dispatch/assignment.go b/apps/daemon/internal/dispatch/assignment.go index 1669dddcd..209fc4be6 100644 --- a/apps/daemon/internal/dispatch/assignment.go +++ b/apps/daemon/internal/dispatch/assignment.go @@ -2,6 +2,7 @@ package dispatch import ( "context" + "sync" "time" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" @@ -13,6 +14,10 @@ type assignmentState struct { ref proto.AssignmentRef environmentID string released bool + // work counts the Session's admitted reads, writes, exports and Runtime + // preparations until each has sent its terminal result. A release waits + // for it, and the released assignment admits no more. + work sync.WaitGroup } // admitLocked returns why ref admits no new work of sessionID in @@ -30,6 +35,14 @@ func (r *Router) admitLocked(ref proto.AssignmentRef, sessionID, environmentID s return "" } +// trackWorkLocked counts work that ref admitted until the returned func runs, +// after the work has sent its terminal result. Router.mu must be held. +func (r *Router) trackWorkLocked(ref proto.AssignmentRef) func() { + work := &r.assignments[ref.SessionID].work + work.Add(1) + return work.Done +} + // admitRunLocked admits a frame for the run, which ref must have started. // Router.mu must be held. func (r *Router) admitRunLocked(ref proto.AssignmentRef, state *sessionState) string { @@ -90,7 +103,8 @@ func (r *Router) handleAssignmentRelease(ctx context.Context, env proto.Envelope r.mu.Unlock() return r.reply(ctx, env, proto.TypeAssignmentStatus, assignmentStatus("", code)) } - work, preparations := r.fenceSessionWorkLocked(ref.SessionID) + preparations := r.fenceSessionWorkLocked(ref.SessionID) + work := &r.assignments[ref.SessionID].work r.shutdownWG.Add(1) r.mu.Unlock() go func() { @@ -98,9 +112,7 @@ func (r *Router) handleAssignmentRelease(ctx context.Context, env proto.Envelope for _, p := range preparations { r.releasePreparation(p, "failed", proto.AssignmentStale, true, false) } - for _, done := range work { - <-done - } + work.Wait() state, code := proto.AssignmentReleased, "" err := r.closeSessionExecutor(ref.SessionID) if err == nil && input.RemoveHome { @@ -117,34 +129,18 @@ func (r *Router) handleAssignmentRelease(ctx context.Context, env proto.Envelope return nil } -// fenceSessionWorkLocked ends the Session's admitted work outside its -// Executor. A transfer still receiving its body ends without applying it, and -// a transfer that already committed checks the released assignment before it -// applies. It returns the work the release waits for and the read-only -// preparations it releases, whose release also cancels their exports. -// Router.mu must be held. -func (r *Router) fenceSessionWorkLocked(sessionID string) ([]chan struct{}, []*preparationState) { - var work []chan struct{} - if u := r.workspaceWrite; u != nil && u.envelope.Assignment.SessionID == sessionID { - if !u.finished { - u.finished = true - close(u.ready) - } - work = append(work, u.done) - } - if u := r.runtimePreparation; u != nil && u.envelope.Assignment.SessionID == sessionID { - if !u.finished { - r.finishRuntimePreparationTransferLocked(u, false) - } - work = append(work, u.done) - } - if u := r.workspaceExport; u != nil && u.request.Assignment.SessionID == sessionID { - work = append(work, u.done) - } - for _, read := range r.workspaceReads { - if read.sessionID == sessionID { - work = append(work, read.done) - } +// fenceSessionWorkLocked ends the Session's transfers: one still receiving +// its body ends without applying it, and one that already committed checks the +// released assignment before it applies. It returns the read-only preparations +// the release releases, which also cancels their exports. Router.mu must be +// held. +func (r *Router) fenceSessionWorkLocked(sessionID string) []*preparationState { + if u := r.workspaceWrite; u != nil && u.envelope.Assignment.SessionID == sessionID && !u.finished { + u.finished = true + close(u.ready) + } + if u := r.runtimePreparation; u != nil && u.envelope.Assignment.SessionID == sessionID && !u.finished { + r.finishRuntimePreparationTransferLocked(u, false) } var preparations []*preparationState for _, p := range r.preparations { @@ -152,7 +148,7 @@ func (r *Router) fenceSessionWorkLocked(sessionID string) ([]chan struct{}, []*p preparations = append(preparations, p) } } - return work, preparations + return preparations } // closeSessionExecutor ends the Session's Executor: it abandons a pending diff --git a/apps/daemon/internal/dispatch/router.go b/apps/daemon/internal/dispatch/router.go index 0255d2762..6b8a1caac 100644 --- a/apps/daemon/internal/dispatch/router.go +++ b/apps/daemon/internal/dispatch/router.go @@ -55,7 +55,7 @@ type Router struct { runtimePreparation *runtimePreparationTransfer workspaceWrite *workspaceUpload workspaceExport *workspaceExport - workspaceReads map[string]workspaceRead + workspaceReads map[string]struct{} localWorkspace *localworkspace.Binding sessionEnvironments bool removeHome func(sessionID string) error diff --git a/apps/daemon/internal/dispatch/router_test.go b/apps/daemon/internal/dispatch/router_test.go index 99a9a0d6a..dc805e279 100644 --- a/apps/daemon/internal/dispatch/router_test.go +++ b/apps/daemon/internal/dispatch/router_test.go @@ -22,9 +22,13 @@ type recSender struct { mu sync.Mutex frames []proto.Envelope failNow bool + hold func(proto.Envelope) // runs before a frame is sent } func (s *recSender) Send(_ context.Context, env proto.Envelope) error { + if s.hold != nil { + s.hold(env) + } s.mu.Lock() defer s.mu.Unlock() if s.failNow { diff --git a/apps/daemon/internal/dispatch/runtime_preparation.go b/apps/daemon/internal/dispatch/runtime_preparation.go index 08e1201dc..866eb710f 100644 --- a/apps/daemon/internal/dispatch/runtime_preparation.go +++ b/apps/daemon/internal/dispatch/runtime_preparation.go @@ -17,13 +17,11 @@ const runtimePreparationTimeout = 120 * time.Second // Router.mu protects one connection-local transfer. Partial installation data // belongs to the bound Environment and is never removed by transfer cleanup. -// done closes once its result is sent. type runtimePreparationTransfer struct { envelope proto.Envelope request proto.RuntimePreparePayload data []byte ready chan struct{} - done chan struct{} cancel context.CancelFunc finished bool apply bool @@ -74,12 +72,13 @@ func (r *Router) handleRuntimePrepare(ctx context.Context, env proto.Envelope) e owner, cancel := context.WithTimeout(context.WithoutCancel(ctx), runtimePreparationTimeout) u := &runtimePreparationTransfer{ envelope: env, request: request, data: make([]byte, 0, request.SizeBytes), - ready: make(chan struct{}), done: make(chan struct{}), cancel: cancel, + ready: make(chan struct{}), cancel: cancel, } r.runtimePreparation = u + done := r.trackWorkLocked(env.Assignment) r.shutdownWG.Add(1) r.mu.Unlock() - go r.runRuntimePreparationTransfer(owner, u, r.localWorkspace.ApplyRuntimePreparation) + go r.runRuntimePreparationTransfer(owner, u, r.localWorkspace.ApplyRuntimePreparation, done) if err := r.sendRuntimePrepareResult(ctx, env, proto.RuntimePrepareResultPayload{Outcome: "ready"}); err != nil { cancel() return err @@ -138,9 +137,10 @@ func (r *Router) finishRuntimePreparationTransferLocked(u *runtimePreparationTra // apply must return only after its local mutations stop. Cancellation requests // shutdown, but cannot release ownership while that call is still running. -func (r *Router) runRuntimePreparationTransfer(ctx context.Context, u *runtimePreparationTransfer, apply func(context.Context, proto.RuntimePreparePayload, []byte) error) { +// done runs once the result is sent. +func (r *Router) runRuntimePreparationTransfer(ctx context.Context, u *runtimePreparationTransfer, apply func(context.Context, proto.RuntimePreparePayload, []byte) error, done func()) { defer r.shutdownWG.Done() - defer close(u.done) + defer done() defer u.cancel() select { case <-u.ready: diff --git a/apps/daemon/internal/dispatch/runtime_preparation_test.go b/apps/daemon/internal/dispatch/runtime_preparation_test.go index 1a98b5b4f..6dde480cd 100644 --- a/apps/daemon/internal/dispatch/runtime_preparation_test.go +++ b/apps/daemon/internal/dispatch/runtime_preparation_test.go @@ -183,7 +183,7 @@ func TestRuntimePreparationPreparationExcludesOwnedResources(t *testing.T) { case "export": r.workspaceExport = &workspaceExport{} case "read": - r.workspaceReads = map[string]workspaceRead{"read": {}} + r.workspaceReads = map[string]struct{}{"read": {}} case "run": r.sessions["run"] = &sessionState{} case "executor": @@ -254,7 +254,7 @@ func TestRuntimePreparationCancellationKeepsOwnershipUntilApplyStops(t *testing. ctx, cancel := context.WithCancel(context.Background()) id := uuid.NewString() request := proto.RuntimePreparePayload{Step: "begin", Action: "finalize", EnvironmentID: environment, SessionID: session, Sources: &agentcapabilities.Input{}} - owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, ready: make(chan struct{}), done: make(chan struct{}), cancel: cancel, finished: true, apply: true} + owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, ready: make(chan struct{}), cancel: cancel, finished: true, apply: true} close(owner.ready) r.runtimePreparation = owner r.shutdownWG.Add(1) @@ -269,7 +269,7 @@ func TestRuntimePreparationCancellationKeepsOwnershipUntilApplyStops(t *testing. close(interrupted) <-release return os.WriteFile(retained, []byte("retained"), 0400) - }) + }, func() {}) <-started wait, stop := context.WithTimeout(context.Background(), 20*time.Millisecond) err := r.Shutdown(wait) @@ -334,11 +334,11 @@ func TestRuntimePreparationResultCategoriesAndUnknownOwnership(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) id := uuid.NewString() request := capabilityBegin(environment, session, []byte("abc")) - owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, data: []byte("abc"), ready: make(chan struct{}), done: make(chan struct{}), cancel: cancel, finished: true, apply: true} + owner := &runtimePreparationTransfer{envelope: capabilityEnvelope(t, id, request), request: request, data: []byte("abc"), ready: make(chan struct{}), cancel: cancel, finished: true, apply: true} close(owner.ready) r.runtimePreparation = owner r.shutdownWG.Add(1) - go r.runRuntimePreparationTransfer(ctx, owner, func(context.Context, proto.RuntimePreparePayload, []byte) error { return context.DeadlineExceeded }) + go r.runRuntimePreparationTransfer(ctx, owner, func(context.Context, proto.RuntimePreparePayload, []byte) error { return context.DeadlineExceeded }, func() {}) capabilitiesReceipt(t, sender, id, "unknown") r.mu.Lock() owned := r.runtimePreparation == owner && owner.uncertain && owner.data == nil diff --git a/apps/daemon/internal/dispatch/suspend_test.go b/apps/daemon/internal/dispatch/suspend_test.go index 2ebd342cf..4ae654dd4 100644 --- a/apps/daemon/internal/dispatch/suspend_test.go +++ b/apps/daemon/internal/dispatch/suspend_test.go @@ -43,7 +43,7 @@ func TestQuiesceRejectsEveryUnsettledResource(t *testing.T) { "active": func(r *Router) { r.sessions["run"] = &sessionState{ctxCancel: func() {}} }, "preparing": func(r *Router) { r.preparations["p"] = &preparationState{owns: true} }, "receipt": func(r *Router) { r.preparations["p"] = &preparationState{busy: true} }, - "read": func(r *Router) { r.workspaceReads = map[string]workspaceRead{"read": {}} }, + "read": func(r *Router) { r.workspaceReads = map[string]struct{}{"read": {}} }, "write": func(r *Router) { r.workspaceWrite = &workspaceUpload{} }, "export": func(r *Router) { r.workspaceExport = &workspaceExport{} }, "permission": func(r *Router) { r.permIndex["permission"] = "run" }, diff --git a/apps/daemon/internal/dispatch/workspace_export.go b/apps/daemon/internal/dispatch/workspace_export.go index 5689bcaa0..ca9ebc2c1 100644 --- a/apps/daemon/internal/dispatch/workspace_export.go +++ b/apps/daemon/internal/dispatch/workspace_export.go @@ -9,12 +9,11 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" ) -// workspaceExport is one admitted export; done closes once it ends. +// workspaceExport is one admitted export. type workspaceExport struct { request proto.Envelope requests chan proto.WorkspaceExportPayload cancel context.CancelFunc - done chan struct{} } func (r *Router) handleWorkspaceExport(ctx context.Context, env proto.Envelope) error { @@ -59,45 +58,57 @@ func (r *Router) handleWorkspaceExport(ctx context.Context, env proto.Envelope) } owner, stop := r.shutdownContext(p.ctx) owner, cancel := context.WithTimeout(owner, 180*time.Second) - u = &workspaceExport{request: env, requests: make(chan proto.WorkspaceExportPayload, 1), cancel: func() { cancel(); stop() }, done: make(chan struct{})} + u = &workspaceExport{request: env, requests: make(chan proto.WorkspaceExportPayload, 1), cancel: func() { cancel(); stop() }} u.requests <- request r.workspaceExport = u + done := r.trackWorkLocked(env.Assignment) r.shutdownWG.Add(1) r.mu.Unlock() - go r.runWorkspaceExport(owner, u) + go r.runWorkspaceExport(owner, u, done) return nil } -func (r *Router) runWorkspaceExport(ctx context.Context, u *workspaceExport) { +// runWorkspaceExport answers each request it admitted, even once the export is +// canceled, and runs done after its last result is sent. +func (r *Router) runWorkspaceExport(ctx context.Context, u *workspaceExport, done func()) { defer r.shutdownWG.Done() - defer close(u.done) + defer done() + // A result has its own send budget, independent of the export's cancellation. + send := func(result proto.WorkspaceExportResultPayload) error { + return r.sendWorkspaceExport(context.WithoutCancel(ctx), u.request, result) + } reader, writer := io.Pipe() - done := make(chan struct{}) + exported := make(chan struct{}) go func() { - defer close(done) + defer close(exported) err := r.localWorkspace.ExportOutputs(ctx, writer) _ = writer.CloseWithError(err) }() + var offset int64 defer func() { u.cancel() _ = reader.Close() - <-done + <-exported r.mu.Lock() if r.workspaceExport == u { r.workspaceExport = nil } r.mu.Unlock() + select { + case <-u.requests: + _ = send(proto.WorkspaceExportResultPayload{Outcome: "failed", Offset: offset, ErrorCode: "export_failed"}) + default: + } }() // Closing the reader unblocks a pending pipe read on cancellation or shutdown. stopRead := context.AfterFunc(ctx, func() { _ = reader.CloseWithError(ctx.Err()) }) defer stopRead() - var offset int64 buffer := make([]byte, proto.WorkspaceExportChunkBytes) for { select { case request := <-u.requests: if request.Offset != offset { - _ = r.sendWorkspaceExport(ctx, u.request, proto.WorkspaceExportResultPayload{Outcome: "failed", Offset: offset, ErrorCode: "invalid_request"}) + _ = send(proto.WorkspaceExportResultPayload{Outcome: "failed", Offset: offset, ErrorCode: "invalid_request"}) return } case <-ctx.Done(): @@ -116,14 +127,14 @@ func (r *Router) runWorkspaceExport(ctx context.Context, u *workspaceExport) { } if result.Outcome == "completed" { // The next owner may start immediately after receiving completion. - <-done + <-exported r.mu.Lock() if r.workspaceExport == u { r.workspaceExport = nil } r.mu.Unlock() } - if r.sendWorkspaceExport(ctx, u.request, result) != nil || result.Outcome != "chunk" { + if send(result) != nil || result.Outcome != "chunk" { return } } diff --git a/apps/daemon/internal/dispatch/workspace_export_test.go b/apps/daemon/internal/dispatch/workspace_export_test.go index 64520ff74..7a1e2e5b7 100644 --- a/apps/daemon/internal/dispatch/workspace_export_test.go +++ b/apps/daemon/internal/dispatch/workspace_export_test.go @@ -18,7 +18,11 @@ import ( type exportSender struct{ replies chan proto.Envelope } +// Send fails with a canceled context, as the connection does. func (s exportSender) Send(ctx context.Context, env proto.Envelope) error { + if err := ctx.Err(); err != nil { + return err + } select { case s.replies <- env: return nil @@ -188,3 +192,25 @@ func TestWorkspaceExportCancelUnblocksWriterAndReleasesCapacity(t *testing.T) { } t.Fatal("export cancellation did not release capacity") } + +func TestCanceledExportAnswersTheRequestCoreAwaits(t *testing.T) { + r, s, request := exporterRouter(t, "normal") + preparation, release := context.WithCancel(context.Background()) + r.mu.Lock() + r.preparations[request.Handle].ctx = preparation + r.mu.Unlock() + id := uuid.NewString() + sendExport(t, r, id, request) + first := readExport(t, s) + if first.Outcome != "chunk" { + t.Fatal(first) + } + // Core asks for the next chunk as the preparation that owns the export is released. + r.mu.Lock() + release() + r.workspaceExport.requests <- proto.WorkspaceExportPayload{Step: "next", Offset: int64(len(first.Data))} + r.mu.Unlock() + if got := readExport(t, s); got.Outcome != "failed" && got.Outcome != "chunk" { + t.Fatal("the canceled export answered with", got) + } +} diff --git a/apps/daemon/internal/dispatch/workspace_read.go b/apps/daemon/internal/dispatch/workspace_read.go index d78c5962f..7e4874bb8 100644 --- a/apps/daemon/internal/dispatch/workspace_read.go +++ b/apps/daemon/internal/dispatch/workspace_read.go @@ -13,12 +13,6 @@ import ( const workspaceReadCapacity = 4 -// workspaceRead is one admitted read; done closes once its result is sent. -type workspaceRead struct { - sessionID string - done chan struct{} -} - func (r *Router) handleWorkspaceRead(ctx context.Context, env proto.Envelope) error { // Never echo an unbounded correlation ID onto the shared connection. if len(env.ID) > proto.WorkspaceReadMaxIDBytes { @@ -48,15 +42,15 @@ func (r *Router) handleWorkspaceRead(ctx context.Context, env proto.Envelope) er return r.sendWorkspaceRead(ctx, env, rejectedWorkspaceRead(code)) } if r.workspaceReads == nil { - r.workspaceReads = make(map[string]workspaceRead) + r.workspaceReads = make(map[string]struct{}) } - read := workspaceRead{sessionID: env.Assignment.SessionID, done: make(chan struct{})} - r.workspaceReads[env.ID] = read + r.workspaceReads[env.ID] = struct{}{} + done := r.trackWorkLocked(env.Assignment) r.shutdownWG.Add(1) r.mu.Unlock() go func() { defer r.shutdownWG.Done() - defer close(read.done) + defer done() defer func() { r.mu.Lock(); delete(r.workspaceReads, env.ID); r.mu.Unlock() }() // Observer loss does not discard an admitted native wait or replay it. operation, cancel := context.WithTimeout(context.WithoutCancel(ctx), 12*time.Second) diff --git a/apps/daemon/internal/dispatch/workspace_write.go b/apps/daemon/internal/dispatch/workspace_write.go index 2383218e5..f3f80e968 100644 --- a/apps/daemon/internal/dispatch/workspace_write.go +++ b/apps/daemon/internal/dispatch/workspace_write.go @@ -13,13 +13,11 @@ import ( ) // Router.mu protects this single bounded transfer for the dedicated Environment. -// done closes once its result is sent. type workspaceUpload struct { envelope proto.Envelope request proto.WorkspaceWritePayload data []byte ready chan struct{} - done chan struct{} finished bool apply bool uncertain bool @@ -80,11 +78,12 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e return r.sendWorkspaceWrite(ctx, env, rejectedWorkspaceWrite("resource_unavailable")) } } - u := &workspaceUpload{envelope: env, request: request, data: make([]byte, 0, request.SizeBytes), ready: make(chan struct{}), done: make(chan struct{})} + u := &workspaceUpload{envelope: env, request: request, data: make([]byte, 0, request.SizeBytes), ready: make(chan struct{})} r.workspaceWrite = u + done := r.trackWorkLocked(env.Assignment) r.shutdownWG.Add(1) r.mu.Unlock() - go r.runWorkspaceUpload(context.WithoutCancel(ctx), u) + go r.runWorkspaceUpload(context.WithoutCancel(ctx), u, done) return r.sendWorkspaceWrite(ctx, env, proto.WorkspaceWriteResultPayload{Outcome: "ready"}) } u := r.workspaceWrite @@ -112,9 +111,9 @@ func (r *Router) handleWorkspaceWrite(ctx context.Context, env proto.Envelope) e return nil } -func (r *Router) runWorkspaceUpload(ctx context.Context, u *workspaceUpload) { +func (r *Router) runWorkspaceUpload(ctx context.Context, u *workspaceUpload, done func()) { defer r.shutdownWG.Done() - defer close(u.done) + defer done() timer := time.NewTimer(120 * time.Second) defer timer.Stop() select { diff --git a/apps/daemon/internal/dispatch/workspace_write_test.go b/apps/daemon/internal/dispatch/workspace_write_test.go index 4be853b32..85e74bbd0 100644 --- a/apps/daemon/internal/dispatch/workspace_write_test.go +++ b/apps/daemon/internal/dispatch/workspace_write_test.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "testing" + "testing/synctest" "time" "github.com/MiniMax-AI/OpenAgentCore/apps/daemon/internal/agent" @@ -206,3 +207,37 @@ func TestReleaseFencesUnfinishedWorkspaceWrite(t *testing.T) { t.Fatal("a released assignment's write applied", err) } } + +func TestReleaseWaitsUntilTheWriteResultIsSent(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + r, sender, request, workspace := localWriterRouter(t) + sent := make(chan struct{}) + sender.hold = func(env proto.Envelope) { + var result proto.WorkspaceWriteResultPayload + if env.Type == proto.TypeWorkspaceWriteResult && env.DecodePayload(&result) == nil && result.Outcome == "completed" { + <-sent + } + } + id := uuid.NewString() + for _, step := range []proto.WorkspaceWritePayload{request, {Step: "chunk", Data: []byte("abc")}, {Step: "commit"}} { + if err := r.Handle(t.Context(), mustEnv(t, proto.TypeWorkspaceWrite, id, step)); err != nil { + t.Fatal(err) + } + } + // The write has applied and released its owner; its result is not sent yet. + synctest.Wait() + release(t, r, preparationSessionID, "release", 2, false) + synctest.Wait() + if hasFrame(sender, proto.TypeAssignmentStatus, "release") { + t.Error("the release replied before the write's result was sent") + } + close(sent) + if got := waitAssignmentStatus(t, sender, "release"); got.State != proto.AssignmentReleased { + t.Fatal(got) + } + waitWorkspaceWrite(t, sender, id, "completed") + if data, err := os.ReadFile(filepath.Join(workspace, "file")); err != nil || string(data) != "abc" { + t.Fatal("the committed write did not apply", err) + } + }) +} diff --git a/docs/runtime-protocol.md b/docs/runtime-protocol.md index 62f9a2e5b..472454db7 100644 --- a/docs/runtime-protocol.md +++ b/docs/runtime-protocol.md @@ -119,7 +119,7 @@ Every Session frame carries the assignment: `execution_prepare`, `execution_star Before a Session's first operation on a connection, including Environment initialization and file work without a Turn, Core sends `assignment_bind` with the Session's Environment ID and waits for `assignment_status` `bound`. A repeated bind of the same assignment is `bound` again. The Runtime admits a Session frame only under the assignment it bound: an older epoch, or a released one, fails with `assignment_stale`; another assignment, Session or Environment fails with `assignment_conflict`. A started Run's frames, including its cancellation receipt, stay admissible under the assignment that started it until the release. A repeated function result or decision whose receipt the Runtime already recorded is answered only under the assignment that applied it; another fails with `assignment_conflict`. -Core records a release and advances the epoch before it sends anything. Deleting a Session releases its assignment with `remove_home: true`; releasing its Environment sends `false`. A deletion never revokes a shared Runtime credential. `assignment_release` fences the assignment at once. The Runtime then stops the Session's work: a transfer still receiving its body, or committed but not yet applied, ends with `assignment_stale`; it releases read-only preparations and waits until every workspace read, write, export and Runtime preparation has sent its result. It closes the Session's Executors and, when asked, removes the native home; only then does it reply `released` or `home_removed`. Unfinished cleanup replies `failed` with `cleanup_unconfirmed`, and a retry at the same epoch repeats it. A Runtime that declares `home_removal` unsupported answers `remove_home: true` with `unsupported_operation`, and Core asks it only to release. Core records the release as applied from a matching `released` or `home_removed`, or at once when no Runtime is left to act on it: a release to a Runtime without authority is settled when recorded, and revoking a Runtime settles its releases. Core resends every unacknowledged release to a Runtime when it connects and backs off a release that fails. A quiesced Runtime admits only a release and the matching `environment_resume`, which carries the assignment that quiesced it. +Core records a release and advances the epoch before it sends anything. Deleting a Session releases its assignment with `remove_home: true`; releasing its Environment sends `false`. A deletion never revokes a shared Runtime credential. `assignment_release` fences the assignment at once. The Runtime then stops the Session's work: a transfer still receiving its body, or committed but not yet applied, ends with `assignment_stale`; it releases read-only preparations and waits until every workspace read, write, export and Runtime preparation has sent its result. It closes the Session's Executors and, when asked, removes the native home; only then does it reply `released` or `home_removed`. Unfinished cleanup replies `failed` with `cleanup_unconfirmed`, and a retry at the same epoch repeats it. A Runtime that declares `home_removal` unsupported answers `remove_home: true` with `unsupported_operation`, and Core asks it only to release. Core records the release as applied from a matching `released` or `home_removed`, or at once when no Runtime is left to act on it: a release to a Runtime without authority is settled when recorded, and revoking a Runtime settles its releases. Core resends every unacknowledged release to a Runtime when it connects; a release that fails backs off, and the release due longest goes first, so failing releases cannot delay the rest. A quiesced Runtime admits only a release and the matching `environment_resume`, which carries the assignment that quiesced it. The Runtime answers a Core frame it cannot route with `protocol_error`, which echoes the request's ID and carries its type and an error code. diff --git a/docs/zh/runtime-protocol.md b/docs/zh/runtime-protocol.md index 6c8cbd392..d3d75d540 100644 --- a/docs/zh/runtime-protocol.md +++ b/docs/zh/runtime-protocol.md @@ -1,7 +1,7 @@ --- title: "Core–Runtime 协议" source: docs/runtime-protocol.md -source_hash: 84a2ec5b9308a5e1a118b05c1ea1c92f8741871785524f3878321286b7e2d25b +source_hash: 6410b49f1b818f05409c87905301722ce0c8d60fc0073c0e554309288f4cdafd --- 此协议在 Runtime daemon 获取机器凭据后连接 Core 与 daemon，定义 daemon 连接上消息的含义和顺序。wire 类型、限制和验证器仅在 [`internal/agentdaemon/proto`](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/internal/agentdaemon/proto) 中定义一次；Core 的 [gateway](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/services/core/internal/runtimegateway) 与参考 Runtime 的 [dispatcher](https://github.com/MiniMax-AI/OpenAgentCore/tree/main/apps/daemon/internal/dispatch) 都使用它们，因此无需同步第二套 payload schema。签发凭据和打开连接的 HTTP 路由见[机器连接 API](../../contracts/agents-api/zh/machine-api.md)。 @@ -121,7 +121,7 @@ Usage frame 和最终 usage snapshot 都携带当前执行的累计测量，替 在一条连接上执行 Session 的第一个操作之前，包括没有 Turn 的 Environment 初始化和文件操作，Core 发送带 Session 的 Environment ID 的 `assignment_bind`，并等待 `assignment_status` `bound`。重复绑定同一分配仍得到 `bound`。Runtime 只在其已绑定的分配下准入 Session frame：较旧的 epoch 或已释放的分配以 `assignment_stale` 失败；其他分配、Session 或 Environment 以 `assignment_conflict` 失败。已启动 Run 的 frame，包括其取消回执，在释放前仍可在启动它的分配下准入。Runtime 已记录回执的重复函数结果或决策只在应用它的分配下得到回答；其他分配以 `assignment_conflict` 失败。 -Core 先记录释放并推进 epoch，再发送任何消息。删除 Session 以 `remove_home: true` 释放其分配；释放其 Environment 发送 `false`。删除从不吊销共享的 Runtime 凭据。`assignment_release` 立即约束该分配。随后 Runtime 停止 Session 的工作：仍在接收内容、或已提交但尚未应用的传输以 `assignment_stale` 结束；它释放只读准备，并等待每个 workspace 读取、写入、导出和 Runtime 准备发送结果。它关闭 Session 的 Executor，并在要求时删除原生 home；此后才回复 `released` 或 `home_removed`。未完成的清理回复 `failed` 和 `cleanup_unconfirmed`，同一 epoch 的重试会重复清理。声明 `home_removal` 不支持的 Runtime 以 `unsupported_operation` 回答 `remove_home: true`，Core 只要求它释放。Core 根据匹配的 `released` 或 `home_removed` 记录释放已应用；没有 Runtime 能处理该释放时立即记录：发给无授权 Runtime 的释放在记录时即结清，吊销 Runtime 会结清它的释放。Core 在 Runtime 连接时重发所有未确认的释放，并对失败的释放退避重试。已 quiesce 的 Runtime 只准入释放和匹配的 `environment_resume`，后者携带使其 quiesce 的分配。 +Core 先记录释放并推进 epoch，再发送任何消息。删除 Session 以 `remove_home: true` 释放其分配；释放其 Environment 发送 `false`。删除从不吊销共享的 Runtime 凭据。`assignment_release` 立即约束该分配。随后 Runtime 停止 Session 的工作：仍在接收内容、或已提交但尚未应用的传输以 `assignment_stale` 结束；它释放只读准备，并等待每个 workspace 读取、写入、导出和 Runtime 准备发送结果。它关闭 Session 的 Executor，并在要求时删除原生 home；此后才回复 `released` 或 `home_removed`。未完成的清理回复 `failed` 和 `cleanup_unconfirmed`，同一 epoch 的重试会重复清理。声明 `home_removal` 不支持的 Runtime 以 `unsupported_operation` 回答 `remove_home: true`，Core 只要求它释放。Core 根据匹配的 `released` 或 `home_removed` 记录释放已应用；没有 Runtime 能处理该释放时立即记录：发给无授权 Runtime 的释放在记录时即结清，吊销 Runtime 会结清它的释放。Core 在 Runtime 连接时重发所有未确认的释放；失败的释放退避重试，等待最久的释放先发送，因此失败的释放不会拖延其他释放。已 quiesce 的 Runtime 只准入释放和匹配的 `environment_resume`，后者携带使其 quiesce 的分配。 Runtime 对无法路由的 Core frame 回复 `protocol_error`，回显请求 ID，并携带其类型和错误码。 diff --git a/services/core/internal/db/queries/devices.sql b/services/core/internal/db/queries/devices.sql index 798815b49..fc0f5f362 100644 --- a/services/core/internal/db/queries/devices.sql +++ b/services/core/internal/db/queries/devices.sql @@ -12,15 +12,17 @@ LEFT JOIN runtime_allocations a ON a.device_id = d.id WHERE d.id = $1; -- name: RevokeDevice :execrows +UPDATE devices SET revoked_at = COALESCE(revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL +WHERE tenant_id = $1 AND id = $2; + +-- name: SettleRevokedRuntimeReleases :exec -- No Runtime is left to act on a revoked device's releases, so revocation --- settles them. -WITH settled AS ( - UPDATE session_runtime_assignments b SET applied_epoch = b.epoch - WHERE b.runtime_id = sqlc.arg(id) AND b.desired_state = 'released' - AND EXISTS (SELECT 1 FROM devices d WHERE d.tenant_id = sqlc.arg(tenant_id) AND d.id = sqlc.arg(id)) -) -UPDATE devices v SET revoked_at = COALESCE(v.revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL -WHERE v.tenant_id = sqlc.arg(tenant_id) AND v.id = sqlc.arg(id); +-- settles them. It runs after the revocation in the same transaction: the +-- revocation holds the device row, which a release locks before it reads the +-- device's authority, so every release is either seen here or sees the +-- revocation. +UPDATE session_runtime_assignments SET applied_epoch = epoch +WHERE runtime_id = $1 AND desired_state = 'released'; -- name: TouchDevice :execrows UPDATE devices SET last_seen_at = clock_timestamp() @@ -62,6 +64,12 @@ AND (d.environment_id IS NULL OR EXISTS ( -- name: RememberNativeSession :execrows UPDATE session_runtime_assignments SET native_session_id = $2 WHERE session_id = $1; +-- name: LockAssignmentRuntime :exec +-- Locks the device row of the Session's Runtime before a release reads its +-- authority; see SettleRevokedRuntimeReleases. +SELECT 1 FROM session_runtime_assignments b JOIN devices d ON d.id = b.runtime_id +WHERE b.session_id = $1 FOR SHARE OF d; + -- name: ReleaseSessionAssignment :exec -- An identical release keeps its epoch; a release that adds home removal -- advances it. A release whose Runtime has no authority is settled, since no diff --git a/services/core/internal/db/sqlc/devices.sql.go b/services/core/internal/db/sqlc/devices.sql.go index cec299d14..274b848d3 100644 --- a/services/core/internal/db/sqlc/devices.sql.go +++ b/services/core/internal/db/sqlc/devices.sql.go @@ -248,6 +248,18 @@ func (q *Queries) ListPendingAssignmentReleases(ctx context.Context, runtimeIds return items, nil } +const lockAssignmentRuntime = `-- name: LockAssignmentRuntime :exec +SELECT 1 FROM session_runtime_assignments b JOIN devices d ON d.id = b.runtime_id +WHERE b.session_id = $1 FOR SHARE OF d +` + +// Locks the device row of the Session's Runtime before a release reads its +// authority; see SettleRevokedRuntimeReleases. +func (q *Queries) LockAssignmentRuntime(ctx context.Context, sessionID pgtype.UUID) error { + _, err := q.db.Exec(ctx, lockAssignmentRuntime, sessionID) + return err +} + const releaseSessionAssignment = `-- name: ReleaseSessionAssignment :exec UPDATE session_runtime_assignments b SET desired_state = 'released', epoch = b.epoch + 1, remove_home = b.remove_home OR $1::boolean, @@ -286,13 +298,8 @@ func (q *Queries) RememberNativeSession(ctx context.Context, arg RememberNativeS } const revokeDevice = `-- name: RevokeDevice :execrows -WITH settled AS ( - UPDATE session_runtime_assignments b SET applied_epoch = b.epoch - WHERE b.runtime_id = $2 AND b.desired_state = 'released' - AND EXISTS (SELECT 1 FROM devices d WHERE d.tenant_id = $1 AND d.id = $2) -) -UPDATE devices v SET revoked_at = COALESCE(v.revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL -WHERE v.tenant_id = $1 AND v.id = $2 +UPDATE devices SET revoked_at = COALESCE(revoked_at, clock_timestamp()), archive_cancel_turn_id = NULL +WHERE tenant_id = $1 AND id = $2 ` type RevokeDeviceParams struct { @@ -300,8 +307,6 @@ type RevokeDeviceParams struct { ID pgtype.UUID `json:"id"` } -// No Runtime is left to act on a revoked device's releases, so revocation -// settles them. func (q *Queries) RevokeDevice(ctx context.Context, arg RevokeDeviceParams) (int64, error) { result, err := q.db.Exec(ctx, revokeDevice, arg.TenantID, arg.ID) if err != nil { @@ -310,6 +315,21 @@ func (q *Queries) RevokeDevice(ctx context.Context, arg RevokeDeviceParams) (int return result.RowsAffected(), nil } +const settleRevokedRuntimeReleases = `-- name: SettleRevokedRuntimeReleases :exec +UPDATE session_runtime_assignments SET applied_epoch = epoch +WHERE runtime_id = $1 AND desired_state = 'released' +` + +// No Runtime is left to act on a revoked device's releases, so revocation +// settles them. It runs after the revocation in the same transaction: the +// revocation holds the device row, which a release locks before it reads the +// device's authority, so every release is either seen here or sees the +// revocation. +func (q *Queries) SettleRevokedRuntimeReleases(ctx context.Context, runtimeID pgtype.UUID) error { + _, err := q.db.Exec(ctx, settleRevokedRuntimeReleases, runtimeID) + return err +} + const touchDevice = `-- name: TouchDevice :execrows UPDATE devices SET last_seen_at = clock_timestamp() WHERE devices.id = $1 AND EXISTS (SELECT 1 FROM runtime_device_authority a WHERE a.id = devices.id) diff --git a/services/core/internal/execution/assignment_releases.go b/services/core/internal/execution/assignment_releases.go index 113b11d3b..f75f6d1ab 100644 --- a/services/core/internal/execution/assignment_releases.go +++ b/services/core/internal/execution/assignment_releases.go @@ -2,6 +2,7 @@ package execution import ( "context" + "slices" "sync" "time" @@ -13,8 +14,8 @@ import ( // runAssignmentReleases delivers each released assignment to its connected // Runtime until the Runtime acknowledges it, so a Runtime that reconnects // receives the releases it missed. Releases run one per Session, bounded like -// executions. A failed release backs off, so it cannot hold a slot that later -// releases need. +// executions. A failed release backs off, and the release due longest goes +// first, so failing releases cannot hold the slots that others need. func (w *Worker) runAssignmentReleases(ctx context.Context) error { type outcome struct { ref proto.AssignmentRef @@ -43,10 +44,9 @@ func (w *Worker) runAssignmentReleases(ctx context.Context) error { return err } retries.keep(releases) - now := time.Now() - for _, release := range releases { + for _, release := range retries.due(releases, time.Now()) { ref := release.Assignment - if active[ref.SessionID] || len(active) >= w.executionConcurrency() || !retries.due(ref, now) { + if active[ref.SessionID] || len(active) >= w.executionConcurrency() { continue } active[ref.SessionID] = true @@ -68,8 +68,16 @@ type releaseRetry struct { delay time.Duration } -func (r releaseRetries) due(ref proto.AssignmentRef, now time.Time) bool { - return !now.Before(r[ref].at) +// due returns the releases whose backoff has passed, the one due longest +// first; a release never attempted is due from the start. +func (r releaseRetries) due(releases []sessions.AssignmentRelease, now time.Time) []sessions.AssignmentRelease { + due := slices.DeleteFunc(slices.Clone(releases), func(release sessions.AssignmentRelease) bool { + return now.Before(r[release.Assignment].at) + }) + slices.SortStableFunc(due, func(a, b sessions.AssignmentRelease) int { + return r[a.Assignment].at.Compare(r[b.Assignment].at) + }) + return due } func (r releaseRetries) record(ref proto.AssignmentRef, acknowledged bool, now time.Time) { diff --git a/services/core/internal/execution/assignment_releases_test.go b/services/core/internal/execution/assignment_releases_test.go index 633095ee9..1f9c61858 100644 --- a/services/core/internal/execution/assignment_releases_test.go +++ b/services/core/internal/execution/assignment_releases_test.go @@ -8,23 +8,44 @@ import ( "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) -func TestFailedReleaseBacksOffAndLeavesLaterReleasesDue(t *testing.T) { - failing := proto.AssignmentRef{SessionID: "failing", AssignmentID: "failing", Epoch: 2} - later := proto.AssignmentRef{SessionID: "later", AssignmentID: "later", Epoch: 2} +func pendingRelease(session string) sessions.AssignmentRelease { + return sessions.AssignmentRelease{Assignment: proto.AssignmentRef{SessionID: session, AssignmentID: session, Epoch: 2}} +} + +func TestFailedReleaseBacksOff(t *testing.T) { + failing := pendingRelease("failing") retries, now := releaseRetries{}, time.Unix(0, 0) - retries.record(failing, false, now) - if retries.due(failing, now) || !retries.due(failing, now.Add(time.Second)) || !retries.due(later, now) { - t.Fatal("a failed release did not wait a second while later releases stayed due") + due := func(at time.Time) bool { return len(retries.due([]sessions.AssignmentRelease{failing}, at)) == 1 } + retries.record(failing.Assignment, false, now) + if due(now) || !due(now.Add(time.Second)) { + t.Fatal("a failed release did not wait a second") } for range 8 { - retries.record(failing, false, now) + retries.record(failing.Assignment, false, now) } - if retries.due(failing, now.Add(59*time.Second)) || !retries.due(failing, now.Add(time.Minute)) { - t.Fatalf("delay = %s, want a minute", retries[failing].delay) + if due(now.Add(59*time.Second)) || !due(now.Add(time.Minute)) { + t.Fatalf("delay = %s, want a minute", retries[failing.Assignment].delay) } - retries.keep([]sessions.AssignmentRelease{{Assignment: later}}) - retries.record(later, true, now) + retries.record(failing.Assignment, true, now) + retries.record(pendingRelease("settled").Assignment, false, now) + retries.keep(nil) if len(retries) != 0 { t.Fatal("settled or acknowledged releases kept a retry", retries) } } + +// With one slot, two releases whose Runtimes never acknowledge alternate +// while each takes the full acknowledgement timeout; a third still runs. +func TestFailingReleasesCannotStarveAnother(t *testing.T) { + releases := []sessions.AssignmentRelease{pendingRelease("a"), pendingRelease("b"), pendingRelease("c")} + retries, now := releaseRetries{}, time.Unix(0, 0) + for range 3 { + next := retries.due(releases, now)[0].Assignment + if next.SessionID == "c" { + return + } + now = now.Add(2 * time.Minute) + retries.record(next, false, now) + } + t.Fatal("releases that keep failing starved a release never attempted") +} diff --git a/services/core/internal/persistence/postgres/sessionpg/devices.go b/services/core/internal/persistence/postgres/sessionpg/devices.go index 887aa50bb..ade3dfabe 100644 --- a/services/core/internal/persistence/postgres/sessionpg/devices.go +++ b/services/core/internal/persistence/postgres/sessionpg/devices.go @@ -192,11 +192,17 @@ func (s *Store) RevokeDevice(ctx context.Context, tenant, device string) error { if err != nil { return err } - n, err := s.units.Queries().RevokeDevice(ctx, sqlc.RevokeDeviceParams(lookup)) - if err == nil && n == 0 { - return sessions.ErrNotFound - } - return err + return s.units.Transaction(ctx, func(ctx context.Context, tx pgx.Tx) error { + q := sqlc.New(tx) + n, err := q.RevokeDevice(ctx, sqlc.RevokeDeviceParams(lookup)) + if err != nil { + return err + } + if n == 0 { + return sessions.ErrNotFound + } + return q.SettleRevokedRuntimeReleases(ctx, lookup.ID) + }) } func (s *Store) TouchDevice(ctx context.Context, device string) (bool, error) { @@ -296,6 +302,9 @@ func (t *SessionTx) LoadDevice(ctx context.Context, device string) (bool, error) // that release is already recorded. A Session without an assignment has // nothing to release. func (t *SessionTx) ReleaseAssignment(ctx context.Context, removeHome bool) error { + if err := t.q.LockAssignmentRuntime(ctx, t.session); err != nil { + return err + } return t.q.ReleaseSessionAssignment(ctx, sqlc.ReleaseSessionAssignmentParams{SessionID: t.session, RemoveHome: removeHome}) } diff --git a/services/core/tests/integration/assignment_release_test.go b/services/core/tests/integration/assignment_release_test.go index bc1e2960a..f273ccda4 100644 --- a/services/core/tests/integration/assignment_release_test.go +++ b/services/core/tests/integration/assignment_release_test.go @@ -11,6 +11,9 @@ import ( "github.com/gorilla/websocket" "github.com/MiniMax-AI/OpenAgentCore/internal/agentdaemon/proto" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/db/sqlc" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/pgunit" + "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/persistence/postgres/sessionpg" "github.com/MiniMax-AI/OpenAgentCore/services/core/internal/sessions" ) @@ -105,3 +108,83 @@ func TestRevocationSettlesUndeliverableReleases(t *testing.T) { t.Fatal("a release to a revoked Runtime was left pending") } } + +// TestReleaseRacingRevocationIsSettled checks that a release racing the +// revocation of its shared Runtime is settled whichever reaches the device row +// first: the revocation settles the releases it sees, and a release sees the +// revocation. +func TestReleaseRacingRevocationIsSettled(t *testing.T) { + for _, first := range []string{"revocation", "release"} { + t.Run(first, func(t *testing.T) { + h := newDispatchHarness(t) + ctx := t.Context() + h.conn.Close() + awaitDaemonRemoteCondition(t, ctx, 3*time.Second, "Runtime disconnected", func() bool { + _, err := h.registry.LookupDevice(h.device.ID) + return err != nil + }) + // waiting reports whether the named query waits for a lock. + waiting := func(query string) bool { + var n int + if err := h.s.pool.QueryRow(ctx, "SELECT count(*) FROM pg_stat_activity WHERE datname = current_database() AND wait_event_type = 'Lock' AND query LIKE '-- name: ' || $1 || ' %'", query).Scan(&n); err != nil { + t.Fatal(err) + } + return n == 1 + } + service := sessionService(t, h.s) + revoked := make(chan error, 1) + tx, err := h.s.pool.Begin(ctx) + if err != nil { + t.Fatal(err) + } + defer tx.Rollback(context.Background()) + if first == "revocation" { + // The open transaction holds the device row, so the revocation + // starts first and the deletion queues behind it. + if _, err := tx.Exec(ctx, "SELECT 1 FROM devices WHERE id = $1 FOR NO KEY UPDATE", h.device.ID); err != nil { + t.Fatal(err) + } + go func() { revoked <- service.RevokeDevice(ctx, h.tenant, h.device.ID) }() + awaitDaemonRemoteCondition(t, ctx, 3*time.Second, "revocation waits for the device", func() bool { return waiting("RevokeDevice") }) + deleted := make(chan error, 1) + go func() { + deleted <- service.DeleteSession(ctx, sessions.DeleteSessionCommand{TenantID: h.tenant, SessionID: h.session.ID}) + }() + awaitDaemonRemoteCondition(t, ctx, 3*time.Second, "deletion ends or waits for the device", func() bool { + return len(deleted) == 1 || waiting("LockAssignmentRuntime") + }) + if err := tx.Rollback(ctx); err != nil { + t.Fatal(err) + } + if err := <-deleted; err != nil { + t.Fatal(err) + } + } else { + // The open transaction records the release, and the revocation + // runs before it commits. + tenant, _ := pgunit.ParseID(h.tenant) + session, _ := pgunit.ParseID(h.session.ID) + if err := sessionpg.BindSession(sqlc.New(tx), tenant, session).ReleaseAssignment(ctx, true); err != nil { + t.Fatal(err) + } + go func() { revoked <- service.RevokeDevice(ctx, h.tenant, h.device.ID) }() + awaitDaemonRemoteCondition(t, ctx, 3*time.Second, "revocation ends or waits for the device", func() bool { + return len(revoked) == 1 || waiting("RevokeDevice") + }) + if err := tx.Commit(ctx); err != nil { + t.Fatal(err) + } + } + if err := <-revoked; err != nil { + t.Fatal(err) + } + var epoch, applied int64 + if err := h.s.pool.QueryRow(ctx, "SELECT epoch, applied_epoch FROM session_runtime_assignments WHERE session_id=$1", h.session.ID).Scan(&epoch, &applied); err != nil { + t.Fatal(err) + } + if epoch != 2 || applied != 2 { + t.Fatalf("epoch %d, applied %d: the release racing the revocation stayed pending", epoch, applied) + } + }) + } +}