From a30fdfc11f0a46a45b69609beba281d71c7aafd7 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 10:08:07 -0400 Subject: [PATCH 01/17] test(runner): event-gate the stale-sweep Sessions start (RIG-4193) TestRunStartsSessionsBeforeStaleSweepDeadline waited a fixed sweep+1s for Sessions to start. Gate on the blocked Remove unwinding with DeadlineExceeded, then on Sessions, so load cannot trip a wall-clock bound. Join Run in cleanup and rename to match the asserted order. Co-authored-by: Matt Wilkinson --- go/internal/runner/run_sweep_test.go | 62 ++++++++++++++++------------ 1 file changed, 36 insertions(+), 26 deletions(-) diff --git a/go/internal/runner/run_sweep_test.go b/go/internal/runner/run_sweep_test.go index 36d9a786f..fc0493013 100644 --- a/go/internal/runner/run_sweep_test.go +++ b/go/internal/runner/run_sweep_test.go @@ -293,9 +293,10 @@ func TestRunReturnsNilWhenSweepCancelsContext(t *testing.T) { type blockingRemoveTestEngine struct { *pipeRuntime - listed chan struct{} - entered chan struct{} - removed chan struct{} + listed chan struct{} + entered chan struct{} + removed chan struct{} + removeErr error // the ctx error Remove unwound on; read after removed closes } func (e *blockingRemoveTestEngine) ListByOwner(context.Context, string, string) ([]runtime.WorkloadID, error) { @@ -306,8 +307,9 @@ func (e *blockingRemoveTestEngine) ListByOwner(context.Context, string, string) func (e *blockingRemoveTestEngine) Remove(ctx context.Context, _ runtime.WorkloadID) error { close(e.entered) <-ctx.Done() + e.removeErr = ctx.Err() close(e.removed) - return ctx.Err() + return e.removeErr } type sessionsStartedTestHandler struct { @@ -320,9 +322,9 @@ func (h *sessionsStartedTestHandler) Sessions(context.Context, *connect.BidiStre return nil } -func TestRunStartsSessionsBeforeStaleSweepDeadline(t *testing.T) { +func TestRunStartsSessionsWhenStaleSweepTimesOut(t *testing.T) { prev := staleContainerSweepTimeout - staleContainerSweepTimeout = 2 * time.Second + staleContainerSweepTimeout = 100 * time.Millisecond t.Cleanup(func() { staleContainerSweepTimeout = prev }) handler := &sessionsStartedTestHandler{sessions: make(chan struct{})} path, service := compassv1internalconnect.NewRunnerServiceHandler(handler) @@ -340,11 +342,13 @@ func TestRunStartsSessionsBeforeStaleSweepDeadline(t *testing.T) { removed: make(chan struct{}), } ctx, cancel := context.WithCancel(t.Context()) - done := make(chan error, 1) + runDone := make(chan struct{}) + var runErr error runtimeDir := shortRuntimeDir(t) httpClient := h2cHTTPClient(t) go func() { - done <- Run(ctx, RunnerConfig{ + defer close(runDone) + runErr = Run(ctx, RunnerConfig{ RunnerID: "runner-1", ServerAddr: server.URL, Token: "tok", @@ -353,27 +357,24 @@ func TestRunStartsSessionsBeforeStaleSweepDeadline(t *testing.T) { HTTPClient: httpClient, }, nil, discardLoggerRunner()) }() - - select { - case <-handler.sessions: - case <-time.After(staleContainerSweepTimeout + time.Second): + // Registered after the timeout restore so it runs first: Run is joined + // before the var and runtime dir it reads are torn down. + t.Cleanup(func() { cancel() select { - case <-engine.removed: + case <-runDone: case <-time.After(testTimeout): - t.Fatal("stale cleanup did not unwind after cancellation") + t.Error("Run did not return after cancellation") } - select { - case <-done: - case <-time.After(testTimeout): - t.Fatal("Run did not return after cancellation") - } - t.Fatalf("Sessions did not start within %s while stale cleanup was blocked", staleContainerSweepTimeout+time.Second) - } + }) + + // Event-gated, no wall-clock bound asserted: the blocked Remove unwinds on the + // sweep's own deadline, then Sessions must start without any cancel. select { case <-engine.entered: + case <-runDone: + t.Fatalf("Run returned %v before the stale Remove started", runErr) case <-time.After(testTimeout): - cancel() t.Fatal("stale Remove did not start") } select { @@ -381,15 +382,24 @@ func TestRunStartsSessionsBeforeStaleSweepDeadline(t *testing.T) { case <-time.After(testTimeout): t.Fatal("stale cleanup did not stop at its deadline") } + if !errors.Is(engine.removeErr, context.DeadlineExceeded) { + t.Fatalf("stale Remove unwound on %v, want the sweep deadline", engine.removeErr) + } + select { + case <-handler.sessions: + case <-runDone: + t.Fatalf("Run returned %v before Sessions started", runErr) + case <-time.After(testTimeout): + t.Fatal("Sessions did not start after the stale sweep deadline") + } select { - case err := <-done: - if err != nil { - t.Fatalf("Run after bounded startup sweep = %v, want nil", err) + case <-runDone: + if runErr != nil { + t.Fatalf("Run after bounded startup sweep = %v, want nil", runErr) } case <-time.After(testTimeout): t.Fatal("Run did not return after Sessions completed") } - cancel() } // The sweep reaches container backends only through a type assertion. From dd7bfae0e71dabf3546130b7292d2bcb92652fb4 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 00:35:25 -0400 Subject: [PATCH 02/17] fix(runnerhub): ignore lifecycle frames that arrive after ERRORED (RIG-4452) When the Runner's shared-stream ERRORED send stalls, it falls back to a one-shot stream. Frames still buffered on the cancelled shared stream can then be delivered after ERRORED and republish the session as live. The hub now marks a session ERRORED and drops its later lifecycle frames. A lifecycle lock serializes the guard with the status publish, so an in-flight frame cannot publish after ERRORED. A recovery command (resume Start, Reload) clears the mark under the same lock as the command's queue admission, so a command that never reached the Runner leaves it set. A re-enroll clears all marks. Trace frames still relay. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/commands.go | 47 ++- go/internal/runnerhub/hub.go | 21 +- go/internal/runnerhub/lost_session_test.go | 336 +++++++++++++++++++++ go/internal/runnerhub/resume_start.go | 2 +- go/internal/runnerhub/router.go | 14 +- 5 files changed, 410 insertions(+), 10 deletions(-) diff --git a/go/internal/runnerhub/commands.go b/go/internal/runnerhub/commands.go index fd5840260..efc8af668 100644 --- a/go/internal/runnerhub/commands.go +++ b/go/internal/runnerhub/commands.go @@ -75,7 +75,8 @@ func (h *Hub) Start(ctx context.Context, requestID string, req *compassv1.StartA return nil, fmt.Errorf("minting fresh session id: %w", err) } } - result, _, err := h.relay(ctx, req.GetContainerName(), &compassv1internal.SessionsResponse{ + // A fresh Start has no resume id, so relayRecovery clears no guard for it. + result, err := h.relayRecovery(ctx, req.GetContainerName(), req.GetResumeSessionId(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), FreshSessionId: freshID, Command: &compassv1internal.SessionsResponse_Start{Start: req}, @@ -133,7 +134,7 @@ func (h *Hub) Remove(ctx context.Context, requestID string, req *compassv1.Remov // Reload relays a ReloadAgentSession command to the owning Runner. func (h *Hub) Reload(ctx context.Context, requestID string, req *compassv1.ReloadAgentSessionRequest) (*compassv1.ReloadAgentSessionResponse, error) { - result, _, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ + result, err := h.relayRecovery(ctx, req.GetSessionId(), req.GetSessionId(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Reload{Reload: req}, }) @@ -217,14 +218,48 @@ func (h *Hub) relay(ctx context.Context, sessionKey string, cmd *compassv1intern if err != nil { return nil, "", connect.NewError(connect.CodeUnavailable, err) } - result, err := router.dispatch(ctx, cmd) + result, err := relayResult(router.dispatch(ctx, cmd)) if err != nil { - return nil, "", connect.NewError(connect.CodeUnavailable, err) + return nil, "", err + } + return result, runnerID, nil +} + +// relayRecovery is relay for a command that recovers sessionID. It clears the +// ERRORED guard atomically with admission, before the new lifetime can emit a frame. +func (h *Hub) relayRecovery(ctx context.Context, sessionKey, sessionID string, cmd *compassv1internal.SessionsResponse) (*compassv1internal.SessionsRequest, error) { + if sessionID == "" { + result, _, err := h.relay(ctx, sessionKey, cmd) + return result, err + } + if err := CheckClientRequestID(cmd.GetRequestId()); err != nil { + return nil, err + } + router, _, err := h.routerFor(sessionKey) + if err != nil { + return nil, connect.NewError(connect.CodeUnavailable, err) + } + h.lifecycleMu.Lock() + call, err := router.admit(cmd) + if err == nil { + delete(h.erroredSessions, sessionID) + } + h.lifecycleMu.Unlock() + if err != nil { + return nil, connect.NewError(connect.CodeUnavailable, err) + } + return relayResult(waitCall(ctx, call)) +} + +// relayResult maps a dispatched command's outcome onto relay's error contract. +func relayResult(result *compassv1internal.SessionsRequest, err error) (*compassv1internal.SessionsRequest, error) { + if err != nil { + return nil, connect.NewError(connect.CodeUnavailable, err) } if runnerErr := result.GetError(); runnerErr != nil { - return nil, "", runnerErrorToConnect(runnerErr) + return nil, runnerErrorToConnect(runnerErr) } - return result, runnerID, nil + return result, nil } // runnerErrorToConnect maps a RunnerError to the Connect status the client sees. diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 02c00f19d..7d82d3aad 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -429,7 +429,9 @@ type Hub struct { // bindingWriteMu serializes whole enrolls (map-clear through reap) with promotion // writes and cache updates. Lock it before mu; never hold mu across a store call. bindingWriteMu sync.Mutex - mu sync.Mutex + // lifecycleMu serializes lifecycle guards and their sink publication order. + lifecycleMu sync.Mutex + mu sync.Mutex // runner is the single attached Runner (single-Runner MVP, OQ6). A second // enrollment re-attaches rather than registering a second entry. @@ -444,6 +446,9 @@ type Hub struct { // Stop removes, a Runner reconnect drops ALL, so a re-minted id fails closed // (CodeNotFound) not inheriting a stale account (OQ-2). sessionAccounts map[string]sessionBinding + // erroredSessions is lifecycleMu-guarded; entries live until recovery or re-enroll. + // Its size is bounded by the sessions errored during one Runner lifetime. + erroredSessions map[string]struct{} // accountSessions is the REVERSE of sessionAccounts (account -> live session_id), // maintained wherever sessionAccounts is so the two never drift. The delivery // consumer (RIG-1569 T3) resolves a subscribed account to its live session to @@ -503,6 +508,7 @@ func NewHub(lifecycle LifecycleSink, tail SessionTailSink, comms CommsCaller, lo containerAccounts: make(map[string]sessionBinding), sessionAccounts: make(map[string]sessionBinding), accountSessions: make(map[store.AccountID]string), + erroredSessions: make(map[string]struct{}), reapStale: make(map[string]uint64), runnerEpoch: make(map[string]uint64), } @@ -805,6 +811,15 @@ func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, sf if !lifecycle { return } + h.lifecycleMu.Lock() + if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + h.erroredSessions[sessionID] = struct{}{} + } else if _, errored := h.erroredSessions[sessionID]; errored { + h.lifecycleMu.Unlock() + h.log.Debug("ignored stale lifecycle frame after ERRORED", + slog.String("runner_id", runnerID), slog.String("session_id", sessionID), slog.String("state", state.String())) + return + } // Resolve the session's agent account and stamp it onto the published status — the // DL-167 attribution join. A status published after a Runner reconnect cleared the // maps carries none (the residual gap). runnerRuntimeIdentity reads tier/posture in @@ -836,6 +851,7 @@ func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, sf if presence != nil && hasAccount { presence.OnSessionLifecycle(account, sessionID, state) } + h.lifecycleMu.Unlock() // The Runner can see the exit before any deliver is refused, so ERRORED is a loss too. if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED && hasAccount { h.dropLostSessionDetached(ctx, runnerID, sessionID, true) @@ -1001,6 +1017,9 @@ type promotedPair struct { // so none of its pre-enroll sessions live. A failed durable reap still runs the // in-RAM fallback, then returns the error so the Runner retries enrollment. func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier compassv1.RuntimeTier, egressPosture compassv1.EgressPosture) (reattached bool, err error) { + h.lifecycleMu.Lock() + clear(h.erroredSessions) + h.lifecycleMu.Unlock() // Held from the map-clear through the reap, so no promotion lands in between. h.bindingWriteMu.Lock() h.mu.Lock() diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 8fe38f17d..fb7e47321 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -5,11 +5,13 @@ package runnerhub import ( "context" "errors" + "sync" "testing" "testing/synctest" "time" compassv1 "github.com/RigelBuild/compass/go/gen/compass/v1" + compassv1internal "github.com/RigelBuild/compass/go/internal/gen/compass/v1" "github.com/RigelBuild/compass/go/internal/store" ) @@ -122,6 +124,340 @@ func TestErroredSessionDropsBindingAndReportsLoss(t *testing.T) { } } +func TestStaleStateAfterErroredIsIgnored(t *testing.T) { + ctx := t.Context() + hub, lifecycle, tail := newHub() + bindings := newFakeBindingStore() + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }); err != nil { + t.Fatalf("Deliver(ERRORED) = %v, want nil", err) + } + if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { + t.Fatalf("lost report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) + } + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }); err != nil { + t.Fatalf("Deliver(WORKING) = %v, want nil", err) + } + statuses := lifecycle.snapshot() + if len(statuses) != 1 || statuses[0].GetSessionId() != "sess-1" || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("published statuses = %+v, want only ERRORED for sess-1", statuses) + } + frames := tail.snapshot() + if len(frames) != 2 || frames[0].sessionID != "sess-1" || frames[0].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED || frames[1].sessionID != "sess-1" || frames[1].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { + t.Fatalf("relayed session frames = %+v, want ERRORED then WORKING", frames) + } +} + +func TestRecoveryCommandClearsErroredGuard(t *testing.T) { + ctx := t.Context() + cases := []struct { + name string + run func(*Hub) error + result *compassv1internal.SessionsRequest + }{ + { + name: "Start", + run: func(hub *Hub) error { + _, err := hub.Start(ctx, "recover-start", &compassv1.StartAgentSessionRequest{ + ContainerName: "cont-1", ResumeSessionId: "sess-1", + }) + return err + }, + result: &compassv1internal.SessionsRequest{ + Result: &compassv1internal.SessionsRequest_Start{ + Start: &compassv1.StartAgentSessionResponse{SessionId: "sess-recovered"}, + }, + }, + }, + { + name: "StartResume", + run: func(hub *Hub) error { + _, err := hub.StartResume(ctx, "recover-resume", &compassv1.StartAgentSessionRequest{ + ContainerName: "cont-1", ResumeSessionId: "sess-1", + }, nil) + return err + }, + result: &compassv1internal.SessionsRequest{ + Result: &compassv1internal.SessionsRequest_Start{ + Start: &compassv1.StartAgentSessionResponse{SessionId: "sess-recovered"}, + }, + }, + }, + { + name: "Reload", + run: func(hub *Hub) error { + _, err := hub.Reload(ctx, "recover-reload", &compassv1.ReloadAgentSessionRequest{SessionId: "sess-1"}) + return err + }, + result: &compassv1internal.SessionsRequest{ + Result: &compassv1internal.SessionsRequest_Reload{Reload: &compassv1.ReloadAgentSessionResponse{}}, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + hub, lifecycle, _ := newHub() + bindings := newFakeBindingStore() + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }); err != nil { + t.Fatalf("Deliver(ERRORED) = %v, want nil", err) + } + if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { + t.Fatalf("lost report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) + } + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }); err != nil { + t.Fatalf("Deliver(stale WORKING) = %v, want nil", err) + } + statuses := lifecycle.snapshot() + if len(statuses) != 1 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("published statuses before recovery = %+v, want only ERRORED", statuses) + } + + router, _, err := hub.routerFor("sess-1") + if err != nil { + t.Fatalf("routerFor(sess-1) = %v, want router", err) + } + router.attach(func(cmd *compassv1internal.SessionsResponse) error { + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 3, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_READY), + }); err != nil { + t.Errorf("Deliver(READY) during %s = %v, want nil", tc.name, err) + } + go router.complete(&compassv1internal.SessionsRequest{ + RequestId: cmd.GetRequestId(), Result: tc.result.GetResult(), + }) + return nil + }) + if err := tc.run(hub); err != nil { + t.Fatalf("recovery command = %v, want nil", err) + } + + statuses = lifecycle.snapshot() + if len(statuses) != 2 || statuses[1].GetSessionId() != "sess-1" || statuses[1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_READY { + t.Fatalf("published statuses after recovery = %+v, want ERRORED then READY for sess-1", statuses) + } + }) + } +} + +func TestConcurrentErroredPublishesAfterInFlightLifecycle(t *testing.T) { + ctx := t.Context() + lifecycle := &blockingLifecycleSink{workingEntered: make(chan struct{}), releaseWorking: make(chan struct{})} + tail := &blockingTailSink{errored: make(chan struct{})} + var releaseOnce sync.Once + releaseWorking := func() { releaseOnce.Do(func() { close(lifecycle.releaseWorking) }) } + t.Cleanup(releaseWorking) + + hub := NewHub(lifecycle, tail, nil, discardLogger()) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + workingDone := make(chan error, 1) + go func() { + workingDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }) + }() + select { + case <-lifecycle.workingEntered: + case <-time.After(10 * time.Second): + t.Fatal("WORKING lifecycle publish did not block") + } + if hub.lifecycleMu.TryLock() { + hub.lifecycleMu.Unlock() + t.Fatal("lifecycle lock released while WORKING is still being published") + } + + erroredDone := make(chan error, 1) + go func() { + erroredDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }) + }() + select { + case <-tail.errored: + case <-time.After(10 * time.Second): + t.Fatal("ERRORED frame did not reach the tail sink") + } + if hub.lifecycleMu.TryLock() { + hub.lifecycleMu.Unlock() + t.Fatal("lifecycle lock released before ERRORED delivery completed") + } + releaseWorking() + select { + case err := <-workingDone: + if err != nil { + t.Fatalf("Deliver(WORKING) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("Deliver(WORKING) did not complete after release") + } + select { + case err := <-erroredDone: + if err != nil { + t.Fatalf("Deliver(ERRORED) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("Deliver(ERRORED) did not complete after WORKING release") + } + + statuses := lifecycle.snapshot() + if len(statuses) != 2 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING || statuses[1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("published lifecycle order = %+v, want WORKING then ERRORED", statuses) + } +} + +func TestRecoveryPreSendFailureKeepsErroredGuard(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }); err != nil { + t.Fatalf("Deliver(ERRORED) = %v, want nil", err) + } + if _, err := hub.Reload(ctx, "recover-no-sender", &compassv1.ReloadAgentSessionRequest{SessionId: "sess-1"}); err == nil { + t.Fatal("Reload with no attached Runner = nil, want pre-send error") + } + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }); err != nil { + t.Fatalf("Deliver(stale WORKING) = %v, want nil", err) + } + statuses := lifecycle.snapshot() + if len(statuses) != 1 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("published statuses after failed recovery = %+v, want only ERRORED", statuses) + } +} + +func TestRecoveryQueueFullFailureKeepsErroredGuard(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }); err != nil { + t.Fatalf("Deliver(ERRORED) = %v, want nil", err) + } + router, _, err := hub.routerFor("sess-1") + if err != nil { + t.Fatalf("routerFor(sess-1) = %v, want router", err) + } + sendEntered := make(chan struct{}) + releaseSend := make(chan struct{}) + var releaseOnce sync.Once + var enteredOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseSend) }) } + router.attach(func(*compassv1internal.SessionsResponse) error { + enteredOnce.Do(func() { close(sendEntered) }) + <-releaseSend + return nil + }) + if err := router.push(&compassv1internal.SessionsResponse{ + Command: &compassv1internal.SessionsResponse_SecretsVersion{ + SecretsVersion: &compassv1internal.SecretsVersion{SessionId: "sess-1"}, + }, + }); err != nil { + t.Fatalf("queue blocking signal = %v, want nil", err) + } + t.Cleanup(func() { + release() + router.detach(errStreamClosed) + }) + select { + case <-sendEntered: + case <-time.After(10 * time.Second): + t.Fatal("Runner send did not block") + } + for range sendQueueCap { + if err := router.push(&compassv1internal.SessionsResponse{ + Command: &compassv1internal.SessionsResponse_SecretsVersion{ + SecretsVersion: &compassv1internal.SecretsVersion{SessionId: "sess-1"}, + }, + }); err != nil { + t.Fatalf("fill command queue = %v, want nil", err) + } + } + if _, err := hub.Reload(ctx, "recover-full-queue", &compassv1.ReloadAgentSessionRequest{SessionId: "sess-1"}); err == nil { + t.Fatal("Reload with a full Runner queue = nil, want pre-send error") + } + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }); err != nil { + t.Fatalf("Deliver(stale WORKING) = %v, want nil", err) + } + statuses := lifecycle.snapshot() + if len(statuses) != 1 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("published statuses after full-queue recovery failure = %+v, want only ERRORED", statuses) + } +} + +type blockingLifecycleSink struct { + mu sync.Mutex + statuses []*compassv1.AgentSessionStatus + workingEntered chan struct{} + releaseWorking chan struct{} +} + +func (s *blockingLifecycleSink) PublishSessionStatus(status *compassv1.AgentSessionStatus) { + if status.GetState() == compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { + close(s.workingEntered) + <-s.releaseWorking + } + s.mu.Lock() + s.statuses = append(s.statuses, status) + s.mu.Unlock() +} + +func (s *blockingLifecycleSink) snapshot() []*compassv1.AgentSessionStatus { + s.mu.Lock() + defer s.mu.Unlock() + return append([]*compassv1.AgentSessionStatus(nil), s.statuses...) +} + +type blockingTailSink struct { + errored chan struct{} +} + +func (s *blockingTailSink) RelaySessionFrame(_ string, frame *compassv1internal.SessionFrame) { + if frame.GetState() == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + close(s.errored) + } +} + type chanEndSink chan string func (c chanEndSink) OnSessionEnded(_ context.Context, sessionID string) { c <- sessionID } diff --git a/go/internal/runnerhub/resume_start.go b/go/internal/runnerhub/resume_start.go index 918ac4c29..3fd719730 100644 --- a/go/internal/runnerhub/resume_start.go +++ b/go/internal/runnerhub/resume_start.go @@ -27,7 +27,7 @@ import ( // session id exactly as Start does, so a resumed session's comms calls resolve // the same way a fresh one's do. func (h *Hub) StartResume(ctx context.Context, requestID string, req *compassv1.StartAgentSessionRequest, resumeBody []byte) (*compassv1.StartAgentSessionResponse, error) { - result, _, err := h.relay(ctx, req.GetContainerName(), &compassv1internal.SessionsResponse{ + result, err := h.relayRecovery(ctx, req.GetContainerName(), req.GetResumeSessionId(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Start{Start: req}, ResumeBody: &compassv1internal.ResumeBody{SessionBody: string(resumeBody)}, diff --git a/go/internal/runnerhub/router.go b/go/internal/runnerhub/router.go index 4688fccba..b6bae4a9f 100644 --- a/go/internal/runnerhub/router.go +++ b/go/internal/runnerhub/router.go @@ -260,6 +260,16 @@ func (r *commandRouter) logger() *slog.Logger { // result (OQ6 idempotency). The command variant is set by the caller; dispatch // only stamps correlation + waits. func (r *commandRouter) dispatch(ctx context.Context, cmd *compassv1internal.SessionsResponse) (*compassv1internal.SessionsRequest, error) { + call, err := r.admit(cmd) + if err != nil { + return nil, err + } + return waitCall(ctx, call) +} + +// admit registers and enqueues cmd without waiting, so a caller can act atomically +// with the point the command becomes deliverable. +func (r *commandRouter) admit(cmd *compassv1internal.SessionsResponse) (*pendingCall, error) { id := cmd.GetRequestId() if id == "" { return nil, errors.New("session command requires a request id") @@ -270,7 +280,7 @@ func (r *commandRouter) dispatch(ctx context.Context, cmd *compassv1internal.Ses // than issuing a second command to the Runner. if existing, ok := r.inflight[id]; ok { r.mu.Unlock() - return waitCall(ctx, existing) + return existing, nil } if r.sender == nil { r.mu.Unlock() @@ -288,7 +298,7 @@ func (r *commandRouter) dispatch(ctx context.Context, cmd *compassv1internal.Ses return nil, fmt.Errorf("runner send queue full for command %q", id) } r.mu.Unlock() - return waitCall(ctx, call) + return call, nil } // push enqueues cmd onto the outbound queue WITHOUT registering a pendingCall or From 2e8c9630e6922e5fbd20b19da3280825d0e649a9 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 14:02:56 -0400 Subject: [PATCH 03/17] fix(runnerhub): key the ERRORED guard on RunnerSeq, not recovery admission (RIG-4673) A recovery command clearing the guard at admission left a window before the Runner ran it, in which a dead-lifetime frame could still publish. The hub now records ERRORED's RunnerSeq per session and drops lifecycle frames at or below it. RunnerSeq is Runner-wide and monotonic, so new-lifetime frames pass on their own; re-enroll resets the counter and clears the map. The recovery- admission path (relayRecovery, router admit) is removed. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/commands.go | 47 +--- go/internal/runnerhub/hub.go | 22 +- go/internal/runnerhub/lost_session_test.go | 260 +++++---------------- go/internal/runnerhub/resume_start.go | 2 +- go/internal/runnerhub/router.go | 14 +- 5 files changed, 75 insertions(+), 270 deletions(-) diff --git a/go/internal/runnerhub/commands.go b/go/internal/runnerhub/commands.go index efc8af668..fd5840260 100644 --- a/go/internal/runnerhub/commands.go +++ b/go/internal/runnerhub/commands.go @@ -75,8 +75,7 @@ func (h *Hub) Start(ctx context.Context, requestID string, req *compassv1.StartA return nil, fmt.Errorf("minting fresh session id: %w", err) } } - // A fresh Start has no resume id, so relayRecovery clears no guard for it. - result, err := h.relayRecovery(ctx, req.GetContainerName(), req.GetResumeSessionId(), &compassv1internal.SessionsResponse{ + result, _, err := h.relay(ctx, req.GetContainerName(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), FreshSessionId: freshID, Command: &compassv1internal.SessionsResponse_Start{Start: req}, @@ -134,7 +133,7 @@ func (h *Hub) Remove(ctx context.Context, requestID string, req *compassv1.Remov // Reload relays a ReloadAgentSession command to the owning Runner. func (h *Hub) Reload(ctx context.Context, requestID string, req *compassv1.ReloadAgentSessionRequest) (*compassv1.ReloadAgentSessionResponse, error) { - result, err := h.relayRecovery(ctx, req.GetSessionId(), req.GetSessionId(), &compassv1internal.SessionsResponse{ + result, _, err := h.relay(ctx, req.GetSessionId(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Reload{Reload: req}, }) @@ -218,48 +217,14 @@ func (h *Hub) relay(ctx context.Context, sessionKey string, cmd *compassv1intern if err != nil { return nil, "", connect.NewError(connect.CodeUnavailable, err) } - result, err := relayResult(router.dispatch(ctx, cmd)) + result, err := router.dispatch(ctx, cmd) if err != nil { - return nil, "", err - } - return result, runnerID, nil -} - -// relayRecovery is relay for a command that recovers sessionID. It clears the -// ERRORED guard atomically with admission, before the new lifetime can emit a frame. -func (h *Hub) relayRecovery(ctx context.Context, sessionKey, sessionID string, cmd *compassv1internal.SessionsResponse) (*compassv1internal.SessionsRequest, error) { - if sessionID == "" { - result, _, err := h.relay(ctx, sessionKey, cmd) - return result, err - } - if err := CheckClientRequestID(cmd.GetRequestId()); err != nil { - return nil, err - } - router, _, err := h.routerFor(sessionKey) - if err != nil { - return nil, connect.NewError(connect.CodeUnavailable, err) - } - h.lifecycleMu.Lock() - call, err := router.admit(cmd) - if err == nil { - delete(h.erroredSessions, sessionID) - } - h.lifecycleMu.Unlock() - if err != nil { - return nil, connect.NewError(connect.CodeUnavailable, err) - } - return relayResult(waitCall(ctx, call)) -} - -// relayResult maps a dispatched command's outcome onto relay's error contract. -func relayResult(result *compassv1internal.SessionsRequest, err error) (*compassv1internal.SessionsRequest, error) { - if err != nil { - return nil, connect.NewError(connect.CodeUnavailable, err) + return nil, "", connect.NewError(connect.CodeUnavailable, err) } if runnerErr := result.GetError(); runnerErr != nil { - return nil, runnerErrorToConnect(runnerErr) + return nil, "", runnerErrorToConnect(runnerErr) } - return result, nil + return result, runnerID, nil } // runnerErrorToConnect maps a RunnerError to the Connect status the client sees. diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 7d82d3aad..8ea1fcf0b 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -446,9 +446,10 @@ type Hub struct { // Stop removes, a Runner reconnect drops ALL, so a re-minted id fails closed // (CodeNotFound) not inheriting a stale account (OQ-2). sessionAccounts map[string]sessionBinding - // erroredSessions is lifecycleMu-guarded; entries live until recovery or re-enroll. - // Its size is bounded by the sessions errored during one Runner lifetime. - erroredSessions map[string]struct{} + // erroredSessions maps a session to the RunnerSeq of its latest ERRORED. Lifecycle + // frames at or below it are from the dead lifetime; the counter resets only with + // re-enroll, which clears the map. lifecycleMu-guarded. + erroredSessions map[string]uint64 // accountSessions is the REVERSE of sessionAccounts (account -> live session_id), // maintained wherever sessionAccounts is so the two never drift. The delivery // consumer (RIG-1569 T3) resolves a subscribed account to its live session to @@ -508,7 +509,7 @@ func NewHub(lifecycle LifecycleSink, tail SessionTailSink, comms CommsCaller, lo containerAccounts: make(map[string]sessionBinding), sessionAccounts: make(map[string]sessionBinding), accountSessions: make(map[store.AccountID]string), - erroredSessions: make(map[string]struct{}), + erroredSessions: make(map[string]uint64), reapStale: make(map[string]uint64), runnerEpoch: make(map[string]uint64), } @@ -694,7 +695,7 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { } switch f := oneof.(type) { case *compassv1internal.AgentFrame_Session: - h.deliverSession(ctx, ev.RunnerID, ev.SessionID, f.Session) + h.deliverSession(ctx, ev.RunnerID, ev.SessionID, ev.RunnerSeq, f.Session) return nil case *compassv1internal.AgentFrame_DeliveryAck: h.deliverAck(ctx, ev, f.DeliveryAck) @@ -796,7 +797,7 @@ func (h *Hub) fireRunnerReady() { // deliverSession routes session frames to the observation pane, publishes lifecycle // transitions, and retires an owned session when its Runner reports ERRORED. // UNSPECIFIED means "trace only, no transition". -func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, sf *compassv1internal.SessionFrame) { +func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, seq uint64, sf *compassv1internal.SessionFrame) { state := sf.GetState() lifecycle := state != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED // A frame the publishing Runner may not speak for is dropped whole: its trace, @@ -813,11 +814,12 @@ func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, sf } h.lifecycleMu.Lock() if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - h.erroredSessions[sessionID] = struct{}{} - } else if _, errored := h.erroredSessions[sessionID]; errored { + h.erroredSessions[sessionID] = max(h.erroredSessions[sessionID], seq) + } else if erroredSeq, errored := h.erroredSessions[sessionID]; errored && seq <= erroredSeq { h.lifecycleMu.Unlock() - h.log.Debug("ignored stale lifecycle frame after ERRORED", - slog.String("runner_id", runnerID), slog.String("session_id", sessionID), slog.String("state", state.String())) + h.log.Debug("ignored stale lifecycle frame from before ERRORED", + slog.String("runner_id", runnerID), slog.String("session_id", sessionID), + slog.String("state", state.String()), slog.Uint64("runner_seq", seq), slog.Uint64("errored_seq", erroredSeq)) return } // Resolve the session's agent account and stamp it onto the published status — the diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index fb7e47321..cd95d07fd 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -5,6 +5,7 @@ package runnerhub import ( "context" "errors" + "slices" "sync" "testing" "testing/synctest" @@ -135,133 +136,70 @@ func TestStaleStateAfterErroredIsIgnored(t *testing.T) { hub.bindContainer("cont-1", testAgentAccount, "runner-1") hub.promoteSession(ctx, "cont-1", "sess-1") - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), - }); err != nil { - t.Fatalf("Deliver(ERRORED) = %v, want nil", err) - } + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { t.Fatalf("lost report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) } - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), - }); err != nil { - t.Fatalf("Deliver(WORKING) = %v, want nil", err) - } - statuses := lifecycle.snapshot() - if len(statuses) != 1 || statuses[0].GetSessionId() != "sess-1" || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - t.Fatalf("published statuses = %+v, want only ERRORED for sess-1", statuses) - } + // The fallback ERRORED overtook a frame buffered on the cancelled shared stream. + deliverState(t, hub, 4, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) frames := tail.snapshot() - if len(frames) != 2 || frames[0].sessionID != "sess-1" || frames[0].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED || frames[1].sessionID != "sess-1" || frames[1].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { + if len(frames) != 2 || frames[0].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED || frames[1].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { t.Fatalf("relayed session frames = %+v, want ERRORED then WORKING", frames) } } -func TestRecoveryCommandClearsErroredGuard(t *testing.T) { +func TestNewLifetimeStateAfterErroredPublishes(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + deliverState(t, hub, 6, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + // A dead-lifetime frame that arrives after recovery began must still not publish. + deliverState(t, hub, 4, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + deliverState(t, hub, 7, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + + assertPublished(t, lifecycle, + compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED, + compassv1.AgentSessionState_AGENT_SESSION_STATE_READY, + compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) +} + +func TestReenrollClearsErroredBoundary(t *testing.T) { ctx := t.Context() - cases := []struct { - name string - run func(*Hub) error - result *compassv1internal.SessionsRequest - }{ - { - name: "Start", - run: func(hub *Hub) error { - _, err := hub.Start(ctx, "recover-start", &compassv1.StartAgentSessionRequest{ - ContainerName: "cont-1", ResumeSessionId: "sess-1", - }) - return err - }, - result: &compassv1internal.SessionsRequest{ - Result: &compassv1internal.SessionsRequest_Start{ - Start: &compassv1.StartAgentSessionResponse{SessionId: "sess-recovered"}, - }, - }, - }, - { - name: "StartResume", - run: func(hub *Hub) error { - _, err := hub.StartResume(ctx, "recover-resume", &compassv1.StartAgentSessionRequest{ - ContainerName: "cont-1", ResumeSessionId: "sess-1", - }, nil) - return err - }, - result: &compassv1internal.SessionsRequest{ - Result: &compassv1internal.SessionsRequest_Start{ - Start: &compassv1.StartAgentSessionResponse{SessionId: "sess-recovered"}, - }, - }, - }, - { - name: "Reload", - run: func(hub *Hub) error { - _, err := hub.Reload(ctx, "recover-reload", &compassv1.ReloadAgentSessionRequest{SessionId: "sess-1"}) - return err - }, - result: &compassv1internal.SessionsRequest{ - Result: &compassv1internal.SessionsRequest_Reload{Reload: &compassv1.ReloadAgentSessionResponse{}}, - }, - }, + hub, lifecycle, _ := newHub() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + + // A restarted Runner's counter starts over, so its low seqs are a new lifetime. + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + deliverState(t, hub, 1, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + + assertPublished(t, lifecycle, + compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED, + compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) +} + +func deliverState(t *testing.T, hub *Hub, seq uint64, state compassv1.AgentSessionState) { + t.Helper() + if err := hub.Deliver(t.Context(), RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: seq, SessionID: "sess-1", Frame: sessionStateFrame(state), + }); err != nil { + t.Fatalf("Deliver(%s, seq %d) = %v, want nil", state, seq, err) } +} - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - hub, lifecycle, _ := newHub() - bindings := newFakeBindingStore() - hub.SetSessionBindingStore(bindings) - lost := newRecordingLostSink() - hub.SetSessionLostSink(lost) - hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) - hub.bindContainer("cont-1", testAgentAccount, "runner-1") - hub.promoteSession(ctx, "cont-1", "sess-1") - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), - }); err != nil { - t.Fatalf("Deliver(ERRORED) = %v, want nil", err) - } - if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { - t.Fatalf("lost report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) - } - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), - }); err != nil { - t.Fatalf("Deliver(stale WORKING) = %v, want nil", err) - } - statuses := lifecycle.snapshot() - if len(statuses) != 1 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - t.Fatalf("published statuses before recovery = %+v, want only ERRORED", statuses) - } - - router, _, err := hub.routerFor("sess-1") - if err != nil { - t.Fatalf("routerFor(sess-1) = %v, want router", err) - } - router.attach(func(cmd *compassv1internal.SessionsResponse) error { - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 3, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_READY), - }); err != nil { - t.Errorf("Deliver(READY) during %s = %v, want nil", tc.name, err) - } - go router.complete(&compassv1internal.SessionsRequest{ - RequestId: cmd.GetRequestId(), Result: tc.result.GetResult(), - }) - return nil - }) - if err := tc.run(hub); err != nil { - t.Fatalf("recovery command = %v, want nil", err) - } - - statuses = lifecycle.snapshot() - if len(statuses) != 2 || statuses[1].GetSessionId() != "sess-1" || statuses[1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_READY { - t.Fatalf("published statuses after recovery = %+v, want ERRORED then READY for sess-1", statuses) - } - }) +func assertPublished(t *testing.T, lifecycle *fakeLifecycleSink, want ...compassv1.AgentSessionState) { + t.Helper() + statuses := lifecycle.snapshot() + got := make([]compassv1.AgentSessionState, 0, len(statuses)) + for _, s := range statuses { + got = append(got, s.GetState()) + } + if !slices.Equal(got, want) { + t.Fatalf("published states = %v, want %v", got, want) } } @@ -335,96 +273,6 @@ func TestConcurrentErroredPublishesAfterInFlightLifecycle(t *testing.T) { } } -func TestRecoveryPreSendFailureKeepsErroredGuard(t *testing.T) { - ctx := t.Context() - hub, lifecycle, _ := newHub() - hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) - - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), - }); err != nil { - t.Fatalf("Deliver(ERRORED) = %v, want nil", err) - } - if _, err := hub.Reload(ctx, "recover-no-sender", &compassv1.ReloadAgentSessionRequest{SessionId: "sess-1"}); err == nil { - t.Fatal("Reload with no attached Runner = nil, want pre-send error") - } - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), - }); err != nil { - t.Fatalf("Deliver(stale WORKING) = %v, want nil", err) - } - statuses := lifecycle.snapshot() - if len(statuses) != 1 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - t.Fatalf("published statuses after failed recovery = %+v, want only ERRORED", statuses) - } -} - -func TestRecoveryQueueFullFailureKeepsErroredGuard(t *testing.T) { - ctx := t.Context() - hub, lifecycle, _ := newHub() - hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), - }); err != nil { - t.Fatalf("Deliver(ERRORED) = %v, want nil", err) - } - router, _, err := hub.routerFor("sess-1") - if err != nil { - t.Fatalf("routerFor(sess-1) = %v, want router", err) - } - sendEntered := make(chan struct{}) - releaseSend := make(chan struct{}) - var releaseOnce sync.Once - var enteredOnce sync.Once - release := func() { releaseOnce.Do(func() { close(releaseSend) }) } - router.attach(func(*compassv1internal.SessionsResponse) error { - enteredOnce.Do(func() { close(sendEntered) }) - <-releaseSend - return nil - }) - if err := router.push(&compassv1internal.SessionsResponse{ - Command: &compassv1internal.SessionsResponse_SecretsVersion{ - SecretsVersion: &compassv1internal.SecretsVersion{SessionId: "sess-1"}, - }, - }); err != nil { - t.Fatalf("queue blocking signal = %v, want nil", err) - } - t.Cleanup(func() { - release() - router.detach(errStreamClosed) - }) - select { - case <-sendEntered: - case <-time.After(10 * time.Second): - t.Fatal("Runner send did not block") - } - for range sendQueueCap { - if err := router.push(&compassv1internal.SessionsResponse{ - Command: &compassv1internal.SessionsResponse_SecretsVersion{ - SecretsVersion: &compassv1internal.SecretsVersion{SessionId: "sess-1"}, - }, - }); err != nil { - t.Fatalf("fill command queue = %v, want nil", err) - } - } - if _, err := hub.Reload(ctx, "recover-full-queue", &compassv1.ReloadAgentSessionRequest{SessionId: "sess-1"}); err == nil { - t.Fatal("Reload with a full Runner queue = nil, want pre-send error") - } - if err := hub.Deliver(ctx, RunnerEvent{ - RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-1", - Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), - }); err != nil { - t.Fatalf("Deliver(stale WORKING) = %v, want nil", err) - } - statuses := lifecycle.snapshot() - if len(statuses) != 1 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - t.Fatalf("published statuses after full-queue recovery failure = %+v, want only ERRORED", statuses) - } -} - type blockingLifecycleSink struct { mu sync.Mutex statuses []*compassv1.AgentSessionStatus diff --git a/go/internal/runnerhub/resume_start.go b/go/internal/runnerhub/resume_start.go index 3fd719730..918ac4c29 100644 --- a/go/internal/runnerhub/resume_start.go +++ b/go/internal/runnerhub/resume_start.go @@ -27,7 +27,7 @@ import ( // session id exactly as Start does, so a resumed session's comms calls resolve // the same way a fresh one's do. func (h *Hub) StartResume(ctx context.Context, requestID string, req *compassv1.StartAgentSessionRequest, resumeBody []byte) (*compassv1.StartAgentSessionResponse, error) { - result, err := h.relayRecovery(ctx, req.GetContainerName(), req.GetResumeSessionId(), &compassv1internal.SessionsResponse{ + result, _, err := h.relay(ctx, req.GetContainerName(), &compassv1internal.SessionsResponse{ RequestId: orNewRequestID(requestID), Command: &compassv1internal.SessionsResponse_Start{Start: req}, ResumeBody: &compassv1internal.ResumeBody{SessionBody: string(resumeBody)}, diff --git a/go/internal/runnerhub/router.go b/go/internal/runnerhub/router.go index b6bae4a9f..4688fccba 100644 --- a/go/internal/runnerhub/router.go +++ b/go/internal/runnerhub/router.go @@ -260,16 +260,6 @@ func (r *commandRouter) logger() *slog.Logger { // result (OQ6 idempotency). The command variant is set by the caller; dispatch // only stamps correlation + waits. func (r *commandRouter) dispatch(ctx context.Context, cmd *compassv1internal.SessionsResponse) (*compassv1internal.SessionsRequest, error) { - call, err := r.admit(cmd) - if err != nil { - return nil, err - } - return waitCall(ctx, call) -} - -// admit registers and enqueues cmd without waiting, so a caller can act atomically -// with the point the command becomes deliverable. -func (r *commandRouter) admit(cmd *compassv1internal.SessionsResponse) (*pendingCall, error) { id := cmd.GetRequestId() if id == "" { return nil, errors.New("session command requires a request id") @@ -280,7 +270,7 @@ func (r *commandRouter) admit(cmd *compassv1internal.SessionsResponse) (*pending // than issuing a second command to the Runner. if existing, ok := r.inflight[id]; ok { r.mu.Unlock() - return existing, nil + return waitCall(ctx, existing) } if r.sender == nil { r.mu.Unlock() @@ -298,7 +288,7 @@ func (r *commandRouter) admit(cmd *compassv1internal.SessionsResponse) (*pending return nil, fmt.Errorf("runner send queue full for command %q", id) } r.mu.Unlock() - return call, nil + return waitCall(ctx, call) } // push enqueues cmd onto the outbound queue WITHOUT registering a pendingCall or From f2db0e3d63b412ee7973bd80680a152f75ccabfc Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 14:45:56 -0400 Subject: [PATCH 04/17] fix(runner): share RunnerSeq across Gateways and fence ERRORED by enrollment (RIG-4673) Each container socket built its own Gateway counter, so a resumed session's new socket restarted RunnerSeq at 1 and fell under the hub's ERRORED boundary. agentHost now passes one SeqCounter to every Gateway, matching the proto's per-Runner contract. The hub also stamps an enrollment generation, so an ERRORED delivery paused across a re-enroll cannot reinstall a boundary the re-enroll cleared. The stale-frame test now covers settle and presence edges. Co-authored-by: Matt Wilkinson --- go/internal/runner/gateway/gateway.go | 16 +++-- go/internal/runner/gateway/publisher.go | 25 ++++--- go/internal/runner/gateway/session_state.go | 2 +- .../runner/gateway/telemetry_ingest_test.go | 24 ++++++- go/internal/runner/host.go | 7 +- go/internal/runnerhub/hub.go | 9 ++- go/internal/runnerhub/lost_session_test.go | 70 +++++++++++++++++++ 7 files changed, 130 insertions(+), 23 deletions(-) diff --git a/go/internal/runner/gateway/gateway.go b/go/internal/runner/gateway/gateway.go index 24282882a..d22df4464 100644 --- a/go/internal/runner/gateway/gateway.go +++ b/go/internal/runner/gateway/gateway.go @@ -174,13 +174,12 @@ type Gateway struct { // across the whole event stream and the hub only flags seq > lastSeq+1, so // replayed low seqs are accepted and loss in that range stops being detectable. - // Scope is per-Runner-link: relay.go's eventPublisher owns a SECOND counter and - // both feed the hub's one high-water mark, so gap detection is meaningful only - // while exactly one is live. This path replaced the stdout relay for gateway - // traffic, so that holds today; unifying the two is T9. + // In production every Gateway of one Runner link shares deps.Seq, so RunnerSeq + // stays monotonic across container sockets. The hub's stale-ERRORED guard + // relies on that: a resumed session's new socket must not restart at 1. pubMu sync.Mutex pub *sessionPublisher - seq seqCounter + seq *SeqCounter // publishMu fences telemetry admission while a lifecycle terminal state is sent. publishMu sync.Mutex @@ -222,6 +221,8 @@ type Deps struct { Events EventRelay // Committer forwards a durable conversation frame to the Server for commit (CommitConversationFrame). Committer ConversationCommitter + // Seq is the Runner link's RunnerSeq allocator; nil gives this Gateway its own. + Seq *SeqCounter } // NewGateway builds the AgentGateway handler for the container's socket: @@ -240,6 +241,10 @@ type Deps struct { // the stream outlives any one agent request. A caller with no distinct socket // scope (a hermetic test) passes context.Background(). func NewGateway(baseCtx context.Context, containerName string, deps Deps) *Gateway { + seq := deps.Seq + if seq == nil { + seq = &SeqCounter{} + } return &Gateway{ baseCtx: baseCtx, containerName: containerName, @@ -250,6 +255,7 @@ func NewGateway(baseCtx context.Context, containerName string, deps Deps) *Gatew board: deps.Board, events: deps.Events, committer: deps.Committer, + seq: seq, control: noopControlRouter{}, // ttl=0: no expiry, a pure size-bounded LRU (committedKeysMax). The cache // is advisory, so eviction is safe — it never drops the store's boundary. diff --git a/go/internal/runner/gateway/publisher.go b/go/internal/runner/gateway/publisher.go index de94adf12..9f94cd3a1 100644 --- a/go/internal/runner/gateway/publisher.go +++ b/go/internal/runner/gateway/publisher.go @@ -11,10 +11,9 @@ package gateway // emission order, else the hub's gap detector records a false gap. So a publisher // holds its stream mutex across BOTH the seq allocation AND the Send. -// Scope: the counter is Gateway-scoped (per socket / per Runner link) and survives -// a publisher replacement, since a per-publisher counter would restart the -// sequence on a swap. relay.go's eventPublisher owns a SECOND counter feeding the -// same high-water mark, so gap detection holds only while one is live; unifying is T9. +// Scope: the counter is Runner-link-wide (shared by every container's Gateway) and +// survives a publisher replacement, since a per-publisher counter would restart the +// sequence on a swap. import ( "context" @@ -32,9 +31,9 @@ type EventRelay interface { PublishEvents(ctx context.Context) *connect.ClientStreamForClient[compassv1internal.PublishEventsRequest, compassv1internal.PublishEventsResponse] } -// seqCounter is the RunnerSeq allocator shared by every publisher a Gateway -// builds for its socket. It carries ONLY the counter and the lock that guards -// the counter — deliberately not the publishers' stream lock. +// SeqCounter is the RunnerSeq allocator shared by every publisher of every +// Gateway on one Runner link. It carries ONLY the counter and the lock that +// guards the counter — deliberately not the publishers' stream lock. // // Sharing the counter is required: a publisher is replaceable within one session, // and a per-publisher counter restarts the sequence on that swap. Sharing the @@ -45,14 +44,14 @@ type EventRelay interface { // distinct streams need no mutual ordering — the hub keeps one global high-water // mark and cannot observe an interleaving between them — so the coupling would // buy nothing and cost unbounded liveness. -type seqCounter struct { +type SeqCounter struct { mu sync.Mutex n uint64 } // next allocates and returns the next sequence. Called with the publisher's own // stream lock held, so allocation order still equals emission order. -func (c *seqCounter) next() uint64 { +func (c *SeqCounter) next() uint64 { c.mu.Lock() defer c.mu.Unlock() c.n++ @@ -66,7 +65,7 @@ func (c *seqCounter) next() uint64 { // expected behaviour — it must not make the Server report a loss that did not // happen. Safe because the caller holds its stream lock across allocate-and-send, // so no other goroutine can have sent past this value. -func (c *seqCounter) rollback(seq uint64) { +func (c *SeqCounter) rollback(seq uint64) { c.mu.Lock() defer c.mu.Unlock() if c.n == seq { @@ -99,7 +98,7 @@ type sessionPublisher struct { // seq is the Gateway's shared allocator, carried across publishers so the // sequence survives a replacement. Only the counter is shared; its lock is // held just long enough to allocate. - seq *seqCounter + seq *SeqCounter // admit, when set, allocates under the Gateway's publish gate so a sealed // session refuses frames instead of sequencing them after its ERRORED report. admit func() (uint64, error) @@ -122,7 +121,7 @@ type sessionPublisher struct { // life: the Publish handler uses the socket-lifetime context, while lifecycle // reports use a bounded caller ctx. seq is the Gateway counter, carried across // publishers so the sequence never restarts. -func newSessionPublisher(ctx context.Context, relay EventRelay, sessionID string, seq *seqCounter) *sessionPublisher { +func newSessionPublisher(ctx context.Context, relay EventRelay, sessionID string, seq *SeqCounter) *sessionPublisher { ctx, cancel := context.WithCancel(ctx) return &sessionPublisher{ sessionID: sessionID, @@ -230,7 +229,7 @@ func (g *Gateway) acquirePublisher(sessionID string) *sessionPublisher { g.pub = nil } if g.pub == nil { - g.pub = newSessionPublisher(g.baseCtx, g.events, sessionID, &g.seq) + g.pub = newSessionPublisher(g.baseCtx, g.events, sessionID, g.seq) g.pub.admit = func() (uint64, error) { return g.admitFrame(sessionID) } g.pub.afterAdmit = g.afterAdmit } diff --git a/go/internal/runner/gateway/session_state.go b/go/internal/runner/gateway/session_state.go index 7febb5853..c1e4eaeba 100644 --- a/go/internal/runner/gateway/session_state.go +++ b/go/internal/runner/gateway/session_state.go @@ -52,7 +52,7 @@ func (l *SocketListener) PublishSessionState(ctx context.Context, sessionID stri } sendCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), stateSendTimeout) defer cancel() - pub := newSessionPublisher(sendCtx, g.events, sessionID, &g.seq) + pub := newSessionPublisher(sendCtx, g.events, sessionID, g.seq) forwardErr := pub.forward(frame) if forwardErr == nil { sharedErr = nil diff --git a/go/internal/runner/gateway/telemetry_ingest_test.go b/go/internal/runner/gateway/telemetry_ingest_test.go index d8c6636be..56cbf7392 100644 --- a/go/internal/runner/gateway/telemetry_ingest_test.go +++ b/go/internal/runner/gateway/telemetry_ingest_test.go @@ -1164,7 +1164,7 @@ func (s *seqSink) seqs() []uint64 { // a variant defect that restarts at a value which happens not to collide would // pass a duplicates-only check. [1 2] pins the contract. // -// RED: give newSessionPublisher its own &seqCounter{} instead of the Gateway's +// RED: give newSessionPublisher its own &SeqCounter{} instead of the Gateway's // -> seqs = [1 1], and this fails every run. func TestSequenceSurvivesPublisherReplacement(t *testing.T) { sink := &seqSink{} @@ -1199,6 +1199,28 @@ func TestSequenceSurvivesPublisherReplacement(t *testing.T) { } } +// A resumed session gets a new container socket, so a per-Gateway counter would +// restart at 1 and fall under the hub's stale-ERRORED boundary. +// +// RED: ignore deps.Seq in NewGateway -> seqs = [1 1]. +func TestSequenceSharedAcrossGateways(t *testing.T) { + sink := &seqSink{} + events := newRunnerServiceServer(t, sink) + shared := &SeqCounter{} + for _, name := range []string{"cont-1", "cont-2"} { + g := NewGateway(context.Background(), name, Deps{Sessions: boundSessions(), Events: events, Seq: shared}) + if err := g.acquirePublisher("sess-1").forward(traceFrame(name)); err != nil { + t.Fatalf("forward on %s = %v, want success", name, err) + } + if err := releaseCurrentPublisher(g); err != nil { + t.Fatalf("release on %s = %v", name, err) + } + } + if got, want := sink.seqs(), []uint64{1, 2}; !slices.Equal(got, want) { + t.Fatalf("RunnerSeq sequence = %v, want %v across Gateways sharing one counter", got, want) + } +} + // A failed forward must not burn a sequence number. The counter is // socket-lifetime, so an allocated-but-unsent number is a permanent hole, and // the hub reads a skipped number as in-transit loss (runnerhub/hub.go:230) — so diff --git a/go/internal/runner/host.go b/go/internal/runner/host.go index 6fb95e744..b226c7425 100644 --- a/go/internal/runner/host.go +++ b/go/internal/runner/host.go @@ -97,6 +97,9 @@ type agentHost struct { mu sync.Mutex sessions map[string]*liveSession sockets map[string]*gateway.SocketListener + // runnerSeq is shared by every container Gateway so RunnerSeq never restarts + // within one enrollment; the hub drops lifecycle frames at or below ERRORED's. + runnerSeq gateway.SeqCounter // afterExitCheck is a test seam between detecting exit and acquiring the // container lock; its returned func runs when retireOnExit returns. Nil in production. afterExitCheck func() func() @@ -871,7 +874,7 @@ func (h *agentHost) provisionVsockGateway(ctx context.Context, spec runtime.Agen h.teardownContainer(ctx, name) return "", fmt.Errorf("resolving vsock gateway endpoint for container %q: backend reports no session", name) } - deps := gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client} + deps := gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client, Seq: &h.runnerSeq} h.log.InfoContext(ctx, "serving agent gateway over vsock path", slog.String("container", name), slog.String("path", endpoint)) listener, err := gateway.Serve(ctx, endpoint, name, deps) @@ -1250,7 +1253,7 @@ func (h *agentHost) serveSocket(ctx context.Context, containerName string) (*gat // serveSocketAt is serveSocket with an explicit socket path: container tiers pass the // fixed RuntimeDir socket, the host tier a path in the handle's state dir (no mount). func (h *agentHost) serveSocketAt(ctx context.Context, containerName, path string) (*gateway.SocketListener, error) { - listener, err := gateway.Serve(ctx, path, containerName, gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client}) + listener, err := gateway.Serve(ctx, path, containerName, gateway.Deps{Sessions: h, Relay: h.link.client, Lifecycle: h.link.client, Events: h.link.client, Committer: h.link.client, Forge: h.link.client, Board: h.link.client, Seq: &h.runnerSeq}) if err != nil { return nil, fmt.Errorf("serving agent socket for container %q: %w", containerName, err) } diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 8ea1fcf0b..f0ca7dd40 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -450,6 +450,9 @@ type Hub struct { // frames at or below it are from the dead lifetime; the counter resets only with // re-enroll, which clears the map. lifecycleMu-guarded. erroredSessions map[string]uint64 + // lifecycleGen counts enrollments; a delivery that began under an older one must + // not reinstall a boundary the re-enroll cleared. Written under lifecycleMu. + lifecycleGen atomic.Uint64 // accountSessions is the REVERSE of sessionAccounts (account -> live session_id), // maintained wherever sessionAccounts is so the two never drift. The delivery // consumer (RIG-1569 T3) resolves a subscribed account to its live session to @@ -798,6 +801,7 @@ func (h *Hub) fireRunnerReady() { // transitions, and retires an owned session when its Runner reports ERRORED. // UNSPECIFIED means "trace only, no transition". func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, seq uint64, sf *compassv1internal.SessionFrame) { + gen := h.lifecycleGen.Load() state := sf.GetState() lifecycle := state != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED // A frame the publishing Runner may not speak for is dropped whole: its trace, @@ -814,7 +818,9 @@ func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, se } h.lifecycleMu.Lock() if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - h.erroredSessions[sessionID] = max(h.erroredSessions[sessionID], seq) + if gen == h.lifecycleGen.Load() { + h.erroredSessions[sessionID] = max(h.erroredSessions[sessionID], seq) + } } else if erroredSeq, errored := h.erroredSessions[sessionID]; errored && seq <= erroredSeq { h.lifecycleMu.Unlock() h.log.Debug("ignored stale lifecycle frame from before ERRORED", @@ -1021,6 +1027,7 @@ type promotedPair struct { func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier compassv1.RuntimeTier, egressPosture compassv1.EgressPosture) (reattached bool, err error) { h.lifecycleMu.Lock() clear(h.erroredSessions) + h.lifecycleGen.Add(1) h.lifecycleMu.Unlock() // Held from the map-clear through the reap, so no promotion lands in between. h.bindingWriteMu.Lock() diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index cd95d07fd..68a7340e7 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -132,6 +132,10 @@ func TestStaleStateAfterErroredIsIgnored(t *testing.T) { hub.SetSessionBindingStore(bindings) lost := newRecordingLostSink() hub.SetSessionLostSink(lost) + settle := &fakeSettleSink{} + hub.SetSettleSink(settle) + presence := &fakePresenceSink{} + hub.SetPresenceSink(presence) hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) hub.bindContainer("cont-1", testAgentAccount, "runner-1") hub.promoteSession(ctx, "cont-1", "sess-1") @@ -148,6 +152,15 @@ func TestStaleStateAfterErroredIsIgnored(t *testing.T) { if len(frames) != 2 || frames[0].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED || frames[1].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { t.Fatalf("relayed session frames = %+v, want ERRORED then WORKING", frames) } + if got := settle.snapshot(); len(got) != 1 || got[0].state != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("settle edges = %+v, want only ERRORED", got) + } + // The ERRORED loss path adds its own DISCONNECTED edge; only the stale WORKING is wrong. + for _, rec := range presence.lifecycleSnapshot() { + if rec.state == compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { + t.Fatalf("presence edges = %+v, want no edge for the stale frame", presence.lifecycleSnapshot()) + } + } } func TestNewLifetimeStateAfterErroredPublishes(t *testing.T) { @@ -182,6 +195,63 @@ func TestReenrollClearsErroredBoundary(t *testing.T) { compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) } +func TestErroredFromBeforeReenrollDoesNotSurviveIt(t *testing.T) { + ctx := t.Context() + lifecycle := &fakeLifecycleSink{} + tail := &pausingTailSink{entered: make(chan struct{}), release: make(chan struct{})} + hub := NewHub(lifecycle, tail, nil, discardLogger()) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(tail.release) }) } + t.Cleanup(release) + + oldDone := make(chan error, 1) + go func() { + oldDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 5, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }) + }() + select { + case <-tail.entered: + case <-time.After(10 * time.Second): + t.Fatal("old ERRORED did not reach the tail sink") + } + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + release() + select { + case err := <-oldDone: + if err != nil { + t.Fatalf("Deliver(old ERRORED) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("old ERRORED did not complete after release") + } + + deliverState(t, hub, 1, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + statuses := lifecycle.snapshot() + if n := len(statuses); n == 0 || statuses[n-1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_READY { + t.Fatalf("published statuses = %+v, want the new enrollment's READY last", statuses) + } +} + +// pausingTailSink holds the first ERRORED frame in RelaySessionFrame until released. +type pausingTailSink struct { + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (s *pausingTailSink) RelaySessionFrame(_ string, frame *compassv1internal.SessionFrame) { + if frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + return + } + s.once.Do(func() { + close(s.entered) + <-s.release + }) +} + func deliverState(t *testing.T, hub *Hub, seq uint64, state compassv1.AgentSessionState) { t.Helper() if err := hub.Deliver(t.Context(), RunnerEvent{ From 8a2ac888cec0de113b190bb2fdee1fb2c42c10b5 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 15:12:24 -0400 Subject: [PATCH 05/17] fix(runnerhub): fence session frames by stream enrollment; let late seqs close gaps (RIG-4673) PublishEvents now stamps the hub's enrollment generation when the stream opens. A session frame from an older stream is dropped before the tail relay, lifecycle edges, and ERRORED's lost-session cleanup, so a resumed session cannot be unbound by its dead process's late ERRORED. Container Gateways share one RunnerSeq counter on separate streams, so a lower seq can arrive after a higher one. The hub now tracks skipped seqs (bounded) and SeenGap reports only those still unseen. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/handler.go | 2 + go/internal/runnerhub/hub.go | 94 +++++++++++++++++----- go/internal/runnerhub/hub_test.go | 15 ++++ go/internal/runnerhub/lost_session_test.go | 35 ++++++++ 4 files changed, 124 insertions(+), 22 deletions(-) diff --git a/go/internal/runnerhub/handler.go b/go/internal/runnerhub/handler.go index 4df7d08e4..5f980fc4f 100644 --- a/go/internal/runnerhub/handler.go +++ b/go/internal/runnerhub/handler.go @@ -169,6 +169,7 @@ func (h *Handler) PublishEvents(ctx context.Context, stream *connect.ClientStrea if !ok { return nil, errUnauthenticated } + enrollGen := h.hub.EnrollGeneration() for stream.Receive() { msg := stream.Msg() if err := h.hub.Deliver(ctx, RunnerEvent{ @@ -176,6 +177,7 @@ func (h *Handler) PublishEvents(ctx context.Context, stream *connect.ClientStrea SessionID: msg.GetSessionId(), RunnerID: subj.ID, Frame: msg.GetFrame(), + EnrollGen: enrollGen, }); err != nil { return nil, connect.NewError(connect.CodeInternal, err) } diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index f0ca7dd40..97eb41d40 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -43,6 +43,9 @@ type RunnerEvent struct { RunnerID string // Frame is the relayed agent stdout frame, verbatim. Frame *compassv1internal.AgentFrame + // EnrollGen is the hub's enrollment generation when the frame's stream opened; + // a session frame from an older one is dropped. Zero means no enrollment fence. + EnrollGen uint64 } // LifecycleSink publishes an extracted agent-session lifecycle transition onto @@ -450,8 +453,8 @@ type Hub struct { // frames at or below it are from the dead lifetime; the counter resets only with // re-enroll, which clears the map. lifecycleMu-guarded. erroredSessions map[string]uint64 - // lifecycleGen counts enrollments; a delivery that began under an older one must - // not reinstall a boundary the re-enroll cleared. Written under lifecycleMu. + // lifecycleGen counts enrollments, so frames from a stream opened before a + // re-enroll cannot act on its sessions. Written under lifecycleMu. lifecycleGen atomic.Uint64 // accountSessions is the REVERSE of sessionAccounts (account -> live session_id), // maintained wherever sessionAccounts is so the two never drift. The delivery @@ -460,9 +463,12 @@ type Hub struct { accountSessions map[store.AccountID]string // lastSeq is the highest RunnerSeq Deliver has accepted, for gap detection. lastSeq uint64 - // seenGap records whether a sequence gap was ever observed (in-transit - // loss), surfaced for the board/diagnostics. - seenGap bool + // missingSeqs holds skipped RunnerSeqs not yet seen. Container Gateways share + // one counter but send on separate streams, so a skipped seq may arrive late; + // a gap is loss only while it stays open. Capped at maxMissingSeqs. + missingSeqs map[uint64]struct{} + // gapOverflow is set once more seqs went missing than the cap can track. + gapOverflow bool // unknownFrames counts frames whose oneof variant was unset or unrecognized // — logged and counted, never silently dropped (agent.proto:38-39). unknownFrames uint64 @@ -513,6 +519,7 @@ func NewHub(lifecycle LifecycleSink, tail SessionTailSink, comms CommsCaller, lo sessionAccounts: make(map[string]sessionBinding), accountSessions: make(map[store.AccountID]string), erroredSessions: make(map[string]uint64), + missingSeqs: make(map[uint64]struct{}), reapStale: make(map[string]uint64), runnerEpoch: make(map[string]uint64), } @@ -698,7 +705,7 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { } switch f := oneof.(type) { case *compassv1internal.AgentFrame_Session: - h.deliverSession(ctx, ev.RunnerID, ev.SessionID, ev.RunnerSeq, f.Session) + h.deliverSession(ctx, ev, f.Session) return nil case *compassv1internal.AgentFrame_DeliveryAck: h.deliverAck(ctx, ev, f.DeliveryAck) @@ -715,12 +722,18 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { } } -// SeenGap reports whether Deliver ever observed a Runner-sequence gap. For the +// EnrollGeneration returns the current enrollment generation, for a stream to +// stamp on its RunnerEvents. +func (h *Hub) EnrollGeneration() uint64 { + return h.lifecycleGen.Load() +} + +// SeenGap reports whether a skipped Runner sequence is still unseen. For the // board/diagnostics; a gap means in-transit loss the Client bus resync recovers. func (h *Hub) SeenGap() bool { h.mu.Lock() defer h.mu.Unlock() - return h.seenGap + return h.seenGapLocked() } // UnknownFrames reports the count of unset/unrecognized frames Deliver has seen. @@ -746,7 +759,7 @@ func (h *Hub) DroppedAcks() uint64 { // one would interleave with Deliver and could report a gap without the drop that // accompanied it. type FrameDiagnostics struct { - // SeenGap is true once a Runner-sequence gap was observed (in-transit loss + // SeenGap is true while a skipped Runner sequence is unseen (in-transit loss // the Client bus resync recovers). SeenGap bool // UnknownFrames counts frames whose oneof variant was unset or unrecognized. @@ -765,7 +778,7 @@ func (h *Hub) FrameDiagnostics() FrameDiagnostics { h.mu.Lock() defer h.mu.Unlock() return FrameDiagnostics{ - SeenGap: h.seenGap, + SeenGap: h.seenGapLocked(), UnknownFrames: h.unknownFrames, DroppedAcks: h.droppedAcks, } @@ -800,8 +813,15 @@ func (h *Hub) fireRunnerReady() { // deliverSession routes session frames to the observation pane, publishes lifecycle // transitions, and retires an owned session when its Runner reports ERRORED. // UNSPECIFIED means "trace only, no transition". -func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, seq uint64, sf *compassv1internal.SessionFrame) { - gen := h.lifecycleGen.Load() +func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1internal.SessionFrame) { + runnerID, sessionID, seq := ev.RunnerID, ev.SessionID, ev.RunnerSeq + gen := ev.EnrollGen + if gen == 0 { + gen = h.lifecycleGen.Load() + } + if h.staleEnrollment(runnerID, sessionID, gen) { + return + } state := sf.GetState() lifecycle := state != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED // A frame the publishing Runner may not speak for is dropped whole: its trace, @@ -817,10 +837,13 @@ func (h *Hub) deliverSession(ctx context.Context, runnerID, sessionID string, se return } h.lifecycleMu.Lock() + // Re-checked under the lock: enroll may have run since the first check. + if h.staleEnrollment(runnerID, sessionID, gen) { + h.lifecycleMu.Unlock() + return + } if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - if gen == h.lifecycleGen.Load() { - h.erroredSessions[sessionID] = max(h.erroredSessions[sessionID], seq) - } + h.erroredSessions[sessionID] = max(h.erroredSessions[sessionID], seq) } else if erroredSeq, errored := h.erroredSessions[sessionID]; errored && seq <= erroredSeq { h.lifecycleMu.Unlock() h.log.Debug("ignored stale lifecycle frame from before ERRORED", @@ -964,20 +987,47 @@ func (h *Hub) forgeNotificationAck(ctx context.Context, ev RunnerEvent, ack *com } } -// recordSeq advances the accepted-sequence high-water mark and flags a gap when -// the observed seq is not exactly one past the last (in-transit loss). The first -// event (lastSeq == 0) establishes the baseline without flagging. +// staleEnrollment reports, and logs, a session frame whose stream predates the +// current enrollment. +func (h *Hub) staleEnrollment(runnerID, sessionID string, gen uint64) bool { + if gen == h.lifecycleGen.Load() { + return false + } + h.log.Debug("dropped session frame from a stream opened before re-enroll", + slog.String("runner_id", runnerID), slog.String("session_id", sessionID)) + return true +} + +// maxMissingSeqs bounds missingSeqs against a corrupt or hostile seq jump. +const maxMissingSeqs = 4096 + +// recordSeq advances the accepted-sequence high-water mark and records the seqs +// a jump skipped; a late arrival fills its gap. The first event (lastSeq == 0) +// establishes the baseline without flagging. func (h *Hub) recordSeq(seq uint64) { h.mu.Lock() defer h.mu.Unlock() + if seq <= h.lastSeq { + delete(h.missingSeqs, seq) + return + } if h.lastSeq != 0 && seq > h.lastSeq+1 { - h.seenGap = true h.log.Warn("runner event sequence gap", slog.Uint64("expected", h.lastSeq+1), slog.Uint64("got", seq)) + for missing := h.lastSeq + 1; missing < seq; missing++ { + if len(h.missingSeqs) >= maxMissingSeqs { + h.gapOverflow = true + break + } + h.missingSeqs[missing] = struct{}{} + } } - if seq > h.lastSeq { - h.lastSeq = seq - } + h.lastSeq = seq +} + +// seenGapLocked reports whether any skipped seq is still unseen. Caller holds mu. +func (h *Hub) seenGapLocked() bool { + return h.gapOverflow || len(h.missingSeqs) > 0 } // countUnknown records and logs an unknown frame. diff --git a/go/internal/runnerhub/hub_test.go b/go/internal/runnerhub/hub_test.go index 044cd530a..92149b06d 100644 --- a/go/internal/runnerhub/hub_test.go +++ b/go/internal/runnerhub/hub_test.go @@ -149,6 +149,21 @@ func TestDeliverSequenceGapDetection(t *testing.T) { t.Fatal("SeenGap true on a contiguous 1,2,3; a gap-free stream must not flag") } }) + + // Container Gateways share one counter on separate streams, so 2 can trail 3. + t.Run("late arrival closes the gap", func(t *testing.T) { + hub := newHubOnly() + deliverSeq(t, hub, 1) + deliverSeq(t, hub, 3) + deliverSeq(t, hub, 4) + if !hub.SeenGap() { + t.Fatal("SeenGap false while seq 2 is still missing") + } + deliverSeq(t, hub, 2) + if hub.SeenGap() { + t.Fatal("SeenGap true after the late seq 2 arrived; nothing was lost") + } + }) } // deliverSeq delivers one trace frame at seq n (a frame that touches no sink of diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 68a7340e7..cd023dfff 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -252,6 +252,41 @@ func (s *pausingTailSink) RelaySessionFrame(_ string, frame *compassv1internal.S }) } +func TestFramesFromStreamBeforeReenrollAreDropped(t *testing.T) { + ctx := t.Context() + hub, lifecycle, tail := newHub() + settle := &fakeSettleSink{} + hub.SetSettleSink(settle) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + oldGen := hub.EnrollGeneration() + + // The Runner restarts and resumes the same session id before the old stream drains. + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + if err := hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 9, SessionID: "sess-1", EnrollGen: oldGen, + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }); err != nil { + t.Fatalf("Deliver(old-stream ERRORED) = %v, want nil", err) + } + + if got := lifecycle.snapshot(); len(got) != 0 { + t.Fatalf("published statuses = %+v, want none from the old stream", got) + } + if got := tail.snapshot(); len(got) != 0 { + t.Fatalf("relayed frames = %+v, want none from the old stream", got) + } + if got := settle.snapshot(); len(got) != 0 { + t.Fatalf("settle edges = %+v, want none from the old stream", got) + } + lost.none(t, "an old-stream ERRORED must not unbind the resumed session") + deliverState(t, hub, 1, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) +} + func deliverState(t *testing.T, hub *Hub, seq uint64, state compassv1.AgentSessionState) { t.Helper() if err := hub.Deliver(t.Context(), RunnerEvent{ From 02d5f5f9f6c9a519d91298012d9046beeabf1bd9 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 16:01:34 -0400 Subject: [PATCH 06/17] fix(runnerhub): serialize session delivery with enroll; scope gap tracking per enrollment (RIG-4673) An RWMutex now fences session-frame delivery (read) against enroll (write), so the generation check, tail relay, and lifecycle edges cannot straddle a re-enroll. ERRORED's detached cleanup carries its enrollment generation and skips if a re-enroll has happened, so it cannot unbind a re-bound session. Re-enroll resets the RunnerSeq gap tracker, and stale-stream events no longer feed it. Co-authored-by: Matt Wilkinson --- go/internal/runner/gateway/publisher.go | 9 ++-- go/internal/runnerhub/handler.go | 3 +- go/internal/runnerhub/hub.go | 59 +++++++++++++--------- go/internal/runnerhub/hub_test.go | 20 ++++++++ go/internal/runnerhub/lost_session_test.go | 40 +++++++++++++-- go/internal/runnerhub/relay_comms.go | 9 +++- 6 files changed, 106 insertions(+), 34 deletions(-) diff --git a/go/internal/runner/gateway/publisher.go b/go/internal/runner/gateway/publisher.go index 9f94cd3a1..a77f09046 100644 --- a/go/internal/runner/gateway/publisher.go +++ b/go/internal/runner/gateway/publisher.go @@ -41,9 +41,8 @@ type EventRelay interface { // a replacement and then closes the stale publisher outside pubMu, so a // CloseAndReceive round-trip against an unresponsive-but-connected Server would // block every forward on the live replacement's separate upstream stream. Two -// distinct streams need no mutual ordering — the hub keeps one global high-water -// mark and cannot observe an interleaving between them — so the coupling would -// buy nothing and cost unbounded liveness. +// distinct streams may therefore deliver out of seq order; the hub tolerates a +// late lower seq, so the coupling would buy nothing and cost unbounded liveness. type SeqCounter struct { mu sync.Mutex n uint64 @@ -63,8 +62,8 @@ func (c *SeqCounter) next() uint64 { // permanent hole, and the hub flags a skipped number as in-transit loss // (runnerhub/hub.go:230). A durable frame erring back to the agent is correct, // expected behaviour — it must not make the Server report a loss that did not -// happen. Safe because the caller holds its stream lock across allocate-and-send, -// so no other goroutine can have sent past this value. +// happen. Only the latest seq is reclaimed: if another Gateway allocated since, +// the number stays burned and the hub's gap diagnostic reports it. func (c *SeqCounter) rollback(seq uint64) { c.mu.Lock() defer c.mu.Unlock() diff --git a/go/internal/runnerhub/handler.go b/go/internal/runnerhub/handler.go index 5f980fc4f..90de84335 100644 --- a/go/internal/runnerhub/handler.go +++ b/go/internal/runnerhub/handler.go @@ -128,8 +128,9 @@ func (h *Handler) Sessions(ctx context.Context, stream *connect.BidiStream[compa return connect.NewError(connect.CodeFailedPrecondition, err) } + enrollGen := h.hub.EnrollGeneration() router.setSessionUnknown(func(sessionID string) { - h.hub.dropLostSessionDetached(ctx, subj.ID, sessionID, false) + h.hub.dropLostSessionDetached(ctx, enrollGen, subj.ID, sessionID, false) }) router.attach(stream.Send) defer router.detach(errStreamClosed) diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 97eb41d40..772704b5b 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -453,9 +453,11 @@ type Hub struct { // frames at or below it are from the dead lifetime; the counter resets only with // re-enroll, which clears the map. lifecycleMu-guarded. erroredSessions map[string]uint64 - // lifecycleGen counts enrollments, so frames from a stream opened before a - // re-enroll cannot act on its sessions. Written under lifecycleMu. - lifecycleGen atomic.Uint64 + // enrollMu fences session-frame delivery (read) against enroll (write), and + // guards enrollGen, which counts enrollments so a stream opened before a + // re-enroll cannot act on its sessions. Lock order: enrollMu, lifecycleMu, mu. + enrollMu sync.RWMutex + enrollGen uint64 // accountSessions is the REVERSE of sessionAccounts (account -> live session_id), // maintained wherever sessionAccounts is so the two never drift. The delivery // consumer (RIG-1569 T3) resolves a subscribed account to its live session to @@ -467,7 +469,8 @@ type Hub struct { // one counter but send on separate streams, so a skipped seq may arrive late; // a gap is loss only while it stays open. Capped at maxMissingSeqs. missingSeqs map[uint64]struct{} - // gapOverflow is set once more seqs went missing than the cap can track. + // gapOverflow is set once more seqs went missing than the cap can track; it + // stays set for the enrollment, since untracked seqs can never be proven seen. gapOverflow bool // unknownFrames counts frames whose oneof variant was unset or unrecognized // — logged and counted, never silently dropped (agent.proto:38-39). @@ -693,7 +696,12 @@ func (h *Hub) SetBoardCaller(c BoardCaller) { // whose variant is unset or unrecognized is logged and counted, never silently // dropped (design.md:1427-1434, agent.proto:38-39). func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { - h.recordSeq(ev.RunnerSeq) + h.enrollMu.RLock() + stale := h.staleEnrollmentLocked(ev) + h.enrollMu.RUnlock() + if !stale { + h.recordSeq(ev.RunnerSeq) + } frame := ev.Frame oneof := frame.GetFrame() @@ -725,7 +733,9 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { // EnrollGeneration returns the current enrollment generation, for a stream to // stamp on its RunnerEvents. func (h *Hub) EnrollGeneration() uint64 { - return h.lifecycleGen.Load() + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + return h.enrollGen } // SeenGap reports whether a skipped Runner sequence is still unseen. For the @@ -815,11 +825,11 @@ func (h *Hub) fireRunnerReady() { // UNSPECIFIED means "trace only, no transition". func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1internal.SessionFrame) { runnerID, sessionID, seq := ev.RunnerID, ev.SessionID, ev.RunnerSeq - gen := ev.EnrollGen - if gen == 0 { - gen = h.lifecycleGen.Load() - } - if h.staleEnrollment(runnerID, sessionID, gen) { + // Held through the tail relay and lifecycle edges so enroll cannot slip between + // the generation check and them. + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + if h.staleEnrollmentLocked(ev) { return } state := sf.GetState() @@ -837,11 +847,6 @@ func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1i return } h.lifecycleMu.Lock() - // Re-checked under the lock: enroll may have run since the first check. - if h.staleEnrollment(runnerID, sessionID, gen) { - h.lifecycleMu.Unlock() - return - } if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { h.erroredSessions[sessionID] = max(h.erroredSessions[sessionID], seq) } else if erroredSeq, errored := h.erroredSessions[sessionID]; errored && seq <= erroredSeq { @@ -885,7 +890,7 @@ func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1i h.lifecycleMu.Unlock() // The Runner can see the exit before any deliver is refused, so ERRORED is a loss too. if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED && hasAccount { - h.dropLostSessionDetached(ctx, runnerID, sessionID, true) + h.dropLostSessionDetached(ctx, h.enrollGen, runnerID, sessionID, true) } } @@ -987,14 +992,14 @@ func (h *Hub) forgeNotificationAck(ctx context.Context, ev RunnerEvent, ack *com } } -// staleEnrollment reports, and logs, a session frame whose stream predates the -// current enrollment. -func (h *Hub) staleEnrollment(runnerID, sessionID string, gen uint64) bool { - if gen == h.lifecycleGen.Load() { +// staleEnrollmentLocked reports, and logs, an event whose stream predates the +// current enrollment. Zero EnrollGen means unfenced. Caller holds enrollMu. +func (h *Hub) staleEnrollmentLocked(ev RunnerEvent) bool { + if ev.EnrollGen == 0 || ev.EnrollGen == h.enrollGen { return false } - h.log.Debug("dropped session frame from a stream opened before re-enroll", - slog.String("runner_id", runnerID), slog.String("session_id", sessionID)) + h.log.Debug("dropped event from a stream opened before re-enroll", + slog.String("runner_id", ev.RunnerID), slog.String("session_id", ev.SessionID)) return true } @@ -1075,10 +1080,12 @@ type promotedPair struct { // so none of its pre-enroll sessions live. A failed durable reap still runs the // in-RAM fallback, then returns the error so the Runner retries enrollment. func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier compassv1.RuntimeTier, egressPosture compassv1.EgressPosture) (reattached bool, err error) { + h.enrollMu.Lock() h.lifecycleMu.Lock() clear(h.erroredSessions) - h.lifecycleGen.Add(1) + h.enrollGen++ h.lifecycleMu.Unlock() + h.enrollMu.Unlock() // Held from the map-clear through the reap, so no promotion lands in between. h.bindingWriteMu.Lock() h.mu.Lock() @@ -1108,6 +1115,10 @@ func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier clear(h.containerAccounts) clear(h.sessionAccounts) clear(h.accountSessions) + // A new Runner process restarts RunnerSeq, so gap tracking restarts with it. + h.lastSeq = 0 + clear(h.missingSeqs) + h.gapOverflow = false // Refuse read-through for this Runner from the instant the maps are cleared: // a concurrent lookup could otherwise resurrect rows while the reap is in flight. h.bindingEpoch++ diff --git a/go/internal/runnerhub/hub_test.go b/go/internal/runnerhub/hub_test.go index 92149b06d..f5b8d9952 100644 --- a/go/internal/runnerhub/hub_test.go +++ b/go/internal/runnerhub/hub_test.go @@ -150,6 +150,26 @@ func TestDeliverSequenceGapDetection(t *testing.T) { } }) + // A restarted Runner's counter starts over, so the tracker must start over too. + t.Run("re-enroll restarts tracking", func(t *testing.T) { + hub := newHubOnly() + enroll := func() { + hub.enroll(context.Background(), testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + } + enroll() + deliverSeq(t, hub, 10) + deliverSeq(t, hub, 12) + enroll() + deliverSeq(t, hub, 1) + if hub.SeenGap() { + t.Fatal("SeenGap true after re-enroll; the old process's missing 11 must not carry over") + } + deliverSeq(t, hub, 3) + if !hub.SeenGap() { + t.Fatal("SeenGap false after the new process's 1 then 3") + } + }) + // Container Gateways share one counter on separate streams, so 2 can trail 3. t.Run("late arrival closes the gap", func(t *testing.T) { hub := newHubOnly() diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index cd023dfff..d41545d96 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -5,6 +5,7 @@ package runnerhub import ( "context" "errors" + "runtime" "slices" "sync" "testing" @@ -195,12 +196,16 @@ func TestReenrollClearsErroredBoundary(t *testing.T) { compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) } -func TestErroredFromBeforeReenrollDoesNotSurviveIt(t *testing.T) { +func TestReenrollWaitsForInFlightErroredAndFencesItsCleanup(t *testing.T) { ctx := t.Context() lifecycle := &fakeLifecycleSink{} tail := &pausingTailSink{entered: make(chan struct{}), release: make(chan struct{})} hub := NewHub(lifecycle, tail, nil, discardLogger()) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") var releaseOnce sync.Once release := func() { releaseOnce.Do(func() { close(tail.release) }) } t.Cleanup(release) @@ -217,7 +222,25 @@ func TestErroredFromBeforeReenrollDoesNotSurviveIt(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("old ERRORED did not reach the tail sink") } - hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + if hub.enrollMu.TryLock() { + hub.enrollMu.Unlock() + t.Fatal("enroll lock free while a session frame is mid-delivery") + } + enrolled := make(chan struct{}) + go func() { + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + close(enrolled) + }() + // A pending writer refuses new readers, so the ERRORED's cleanup queues behind enroll. + for deadline := time.Now().Add(10 * time.Second); hub.enrollMu.TryRLock(); { + hub.enrollMu.RUnlock() + if time.Now().After(deadline) { + t.Fatal("re-enroll never queued behind the in-flight delivery") + } + runtime.Gosched() + } release() select { case err := <-oldDone: @@ -227,7 +250,18 @@ func TestErroredFromBeforeReenrollDoesNotSurviveIt(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("old ERRORED did not complete after release") } - + select { + case <-enrolled: + case <-time.After(10 * time.Second): + t.Fatal("re-enroll did not complete") + } + // Readers queued behind enroll, the cleanup among them, finish before this Lock. + hub.enrollMu.Lock() + lost.none(t, "the old ERRORED's cleanup must not unbind the re-bound session") + hub.enrollMu.Unlock() + if _, _, err := hub.routerFor("sess-1"); err != nil { + t.Fatalf("routerFor(sess-1) after re-bind = %v, want the binding kept", err) + } deliverState(t, hub, 1, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) statuses := lifecycle.snapshot() if n := len(statuses); n == 0 || statuses[n-1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_READY { diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index 382d9e3d6..5345895aa 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -914,10 +914,17 @@ func (h *Hub) dropLostSession(ctx context.Context, runnerID, sessionID string, e // dropLostSessionDetached runs dropLostSession off the caller's receive loop: it // does store work, and the stream ctx dies with the stream. -func (h *Hub) dropLostSessionDetached(ctx context.Context, runnerID, sessionID string, errored bool) { +// gen is the enrollment the loss was observed under. +func (h *Hub) dropLostSessionDetached(ctx context.Context, gen uint64, runnerID, sessionID string, errored bool) { go func() { dctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), lostSessionTimeout) defer cancel() + // A re-enroll since the loss may have re-bound sessionID; hold it off until done. + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + if h.enrollGen != gen { + return + } h.dropLostSession(dctx, runnerID, sessionID, errored) }() } From 5a2733b4b672453bb8021ecf5f770cd7d92f8114 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 16:42:56 -0400 Subject: [PATCH 07/17] fix(runnerhub): drop any lifecycle frame older than the session's latest (RIG-4673) Separate streams can deliver one session's frames out of seq order, so an ERRORED at seq 5 could land after the resumed READY at 6 and retire it. The hub now keeps the highest accepted lifecycle seq per session and drops any older lifecycle frame, ERRORED included, before the tail relay. Deliver holds the enrollment read lock for the whole event, so a re-enroll cannot reset sequence state between the generation check and its use. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/hub.go | 42 +++++++++++----------- go/internal/runnerhub/hub_test.go | 22 ++++++++++++ go/internal/runnerhub/lost_session_test.go | 35 ++++++++++++------ 3 files changed, 67 insertions(+), 32 deletions(-) diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 772704b5b..a16bf6e0c 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -449,10 +449,11 @@ type Hub struct { // Stop removes, a Runner reconnect drops ALL, so a re-minted id fails closed // (CodeNotFound) not inheriting a stale account (OQ-2). sessionAccounts map[string]sessionBinding - // erroredSessions maps a session to the RunnerSeq of its latest ERRORED. Lifecycle - // frames at or below it are from the dead lifetime; the counter resets only with - // re-enroll, which clears the map. lifecycleMu-guarded. - erroredSessions map[string]uint64 + // lifecycleSeqs maps a session to the highest RunnerSeq of an accepted lifecycle + // frame. One lifetime's frames arrive in seq order, so a lower one is from a + // cancelled stream or a dead lifetime. Kept past unbind so the boundary outlives + // ERRORED; cleared on re-enroll, when RunnerSeq restarts. lifecycleMu-guarded. + lifecycleSeqs map[string]uint64 // enrollMu fences session-frame delivery (read) against enroll (write), and // guards enrollGen, which counts enrollments so a stream opened before a // re-enroll cannot act on its sessions. Lock order: enrollMu, lifecycleMu, mu. @@ -521,7 +522,7 @@ func NewHub(lifecycle LifecycleSink, tail SessionTailSink, comms CommsCaller, lo containerAccounts: make(map[string]sessionBinding), sessionAccounts: make(map[string]sessionBinding), accountSessions: make(map[store.AccountID]string), - erroredSessions: make(map[string]uint64), + lifecycleSeqs: make(map[string]uint64), missingSeqs: make(map[uint64]struct{}), reapStale: make(map[string]uint64), runnerEpoch: make(map[string]uint64), @@ -696,9 +697,12 @@ func (h *Hub) SetBoardCaller(c BoardCaller) { // whose variant is unset or unrecognized is logged and counted, never silently // dropped (design.md:1427-1434, agent.proto:38-39). func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { + // Held for the whole delivery so enroll cannot reset sequence or lifecycle state + // between the generation check and their use. Never re-acquired below: a nested + // RLock deadlocks behind a pending enroll. h.enrollMu.RLock() + defer h.enrollMu.RUnlock() stale := h.staleEnrollmentLocked(ev) - h.enrollMu.RUnlock() if !stale { h.recordSeq(ev.RunnerSeq) } @@ -713,7 +717,10 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { } switch f := oneof.(type) { case *compassv1internal.AgentFrame_Session: - h.deliverSession(ctx, ev, f.Session) + // Acks stay unfenced: they are idempotent and record work already done. + if !stale { + h.deliverSession(ctx, ev, f.Session) + } return nil case *compassv1internal.AgentFrame_DeliveryAck: h.deliverAck(ctx, ev, f.DeliveryAck) @@ -825,13 +832,6 @@ func (h *Hub) fireRunnerReady() { // UNSPECIFIED means "trace only, no transition". func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1internal.SessionFrame) { runnerID, sessionID, seq := ev.RunnerID, ev.SessionID, ev.RunnerSeq - // Held through the tail relay and lifecycle edges so enroll cannot slip between - // the generation check and them. - h.enrollMu.RLock() - defer h.enrollMu.RUnlock() - if h.staleEnrollmentLocked(ev) { - return - } state := sf.GetState() lifecycle := state != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED // A frame the publishing Runner may not speak for is dropped whole: its trace, @@ -842,20 +842,20 @@ func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1i slog.String("runner_id", runnerID), slog.String("session_id", sessionID)) return } - h.tail.RelaySessionFrame(sessionID, sf) if !lifecycle { + h.tail.RelaySessionFrame(sessionID, sf) return } h.lifecycleMu.Lock() - if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { - h.erroredSessions[sessionID] = max(h.erroredSessions[sessionID], seq) - } else if erroredSeq, errored := h.erroredSessions[sessionID]; errored && seq <= erroredSeq { + if last, seen := h.lifecycleSeqs[sessionID]; seen && seq <= last { h.lifecycleMu.Unlock() - h.log.Debug("ignored stale lifecycle frame from before ERRORED", + h.log.Debug("ignored lifecycle frame older than the session's latest", slog.String("runner_id", runnerID), slog.String("session_id", sessionID), - slog.String("state", state.String()), slog.Uint64("runner_seq", seq), slog.Uint64("errored_seq", erroredSeq)) + slog.String("state", state.String()), slog.Uint64("runner_seq", seq), slog.Uint64("latest_seq", last)) return } + h.lifecycleSeqs[sessionID] = seq + h.tail.RelaySessionFrame(sessionID, sf) // Resolve the session's agent account and stamp it onto the published status — the // DL-167 attribution join. A status published after a Runner reconnect cleared the // maps carries none (the residual gap). runnerRuntimeIdentity reads tier/posture in @@ -1082,7 +1082,7 @@ type promotedPair struct { func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier compassv1.RuntimeTier, egressPosture compassv1.EgressPosture) (reattached bool, err error) { h.enrollMu.Lock() h.lifecycleMu.Lock() - clear(h.erroredSessions) + clear(h.lifecycleSeqs) h.enrollGen++ h.lifecycleMu.Unlock() h.enrollMu.Unlock() diff --git a/go/internal/runnerhub/hub_test.go b/go/internal/runnerhub/hub_test.go index f5b8d9952..0f125aaf8 100644 --- a/go/internal/runnerhub/hub_test.go +++ b/go/internal/runnerhub/hub_test.go @@ -170,6 +170,28 @@ func TestDeliverSequenceGapDetection(t *testing.T) { } }) + // Untracked seqs can never be proven seen, so overflow holds for the enrollment. + t.Run("overflow stays a gap until re-enroll", func(t *testing.T) { + hub := newHubOnly() + enroll := func() { + hub.enroll(context.Background(), testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + } + enroll() + deliverSeq(t, hub, 1) + deliverSeq(t, hub, maxMissingSeqs+3) + for seq := uint64(2); seq < maxMissingSeqs+3; seq++ { + deliverSeq(t, hub, seq) + } + if !hub.SeenGap() { + t.Fatal("SeenGap false after an overflowed jump; untracked seqs cannot be proven seen") + } + enroll() + deliverSeq(t, hub, 1) + if hub.SeenGap() { + t.Fatal("SeenGap true after re-enroll; overflow must reset with the enrollment") + } + }) + // Container Gateways share one counter on separate streams, so 2 can trail 3. t.Run("late arrival closes the gap", func(t *testing.T) { hub := newHubOnly() diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index d41545d96..4dbe92874 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -150,8 +150,8 @@ func TestStaleStateAfterErroredIsIgnored(t *testing.T) { assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) frames := tail.snapshot() - if len(frames) != 2 || frames[0].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED || frames[1].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING { - t.Fatalf("relayed session frames = %+v, want ERRORED then WORKING", frames) + if len(frames) != 1 || frames[0].frame.GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { + t.Fatalf("relayed session frames = %+v, want only ERRORED; a stale state must not reach the tail", frames) } if got := settle.snapshot(); len(got) != 1 || got[0].state != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { t.Fatalf("settle edges = %+v, want only ERRORED", got) @@ -181,6 +181,23 @@ func TestNewLifetimeStateAfterErroredPublishes(t *testing.T) { compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) } +func TestErroredOlderThanNewLifetimeIsIgnored(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + // Separate streams can deliver out of seq order: the resumed READY overtakes ERRORED. + deliverState(t, hub, 6, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + deliverState(t, hub, 5, compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED) + + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) + lost.none(t, "an older ERRORED must not retire the resumed lifetime") +} + func TestReenrollClearsErroredBoundary(t *testing.T) { ctx := t.Context() hub, lifecycle, _ := newHub() @@ -379,15 +396,6 @@ func TestConcurrentErroredPublishesAfterInFlightLifecycle(t *testing.T) { Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), }) }() - select { - case <-tail.errored: - case <-time.After(10 * time.Second): - t.Fatal("ERRORED frame did not reach the tail sink") - } - if hub.lifecycleMu.TryLock() { - hub.lifecycleMu.Unlock() - t.Fatal("lifecycle lock released before ERRORED delivery completed") - } releaseWorking() select { case err := <-workingDone: @@ -405,6 +413,11 @@ func TestConcurrentErroredPublishesAfterInFlightLifecycle(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("Deliver(ERRORED) did not complete after WORKING release") } + select { + case <-tail.errored: + default: + t.Fatal("ERRORED never reached the tail sink") + } statuses := lifecycle.snapshot() if len(statuses) != 2 || statuses[0].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING || statuses[1].GetState() != compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED { From 1ae2220db894c758972ba9377f6a3cdea4dd3f4a Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 16:46:50 -0400 Subject: [PATCH 08/17] docs(gateway): state when a burned RunnerSeq shows as a gap (RIG-4673) Co-authored-by: Matt Wilkinson --- go/internal/runner/gateway/publisher.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/go/internal/runner/gateway/publisher.go b/go/internal/runner/gateway/publisher.go index a77f09046..9a00c9dcb 100644 --- a/go/internal/runner/gateway/publisher.go +++ b/go/internal/runner/gateway/publisher.go @@ -63,7 +63,7 @@ func (c *SeqCounter) next() uint64 { // (runnerhub/hub.go:230). A durable frame erring back to the agent is correct, // expected behaviour — it must not make the Server report a loss that did not // happen. Only the latest seq is reclaimed: if another Gateway allocated since, -// the number stays burned and the hub's gap diagnostic reports it. +// the number stays burned and shows as a gap once the hub has a baseline. func (c *SeqCounter) rollback(seq uint64) { c.mu.Lock() defer c.mu.Unlock() From 76b67cbd7f86e4e694e2c6d58dae8a797ce512ed Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 17:53:04 -0400 Subject: [PATCH 09/17] fix(runnerhub): admit lifecycle frames in seq order; pair router and generation (RIG-4673) A lifecycle frame now takes lifecycleMu before its seq is recorded, so a lower seq cannot overtake a higher one that is still mid-delivery. Sessions reads its router and enrollment generation as one pair, and enroll holds the write lock until the new router is installed. Generations start at 1, so a PublishEvents stream opened before the first Enroll is fenced too. lifecycleSeqs becomes a bounded LRU. ERRORED cleanup's generation check is a synchronous helper with its own test. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/handler.go | 3 +- go/internal/runnerhub/hub.go | 57 ++++++++++--- go/internal/runnerhub/lost_session_test.go | 93 ++++++++++++++++++++-- go/internal/runnerhub/relay_comms.go | 19 +++-- go/internal/runnerhub/seam_test.go | 43 ++++++++++ 5 files changed, 190 insertions(+), 25 deletions(-) diff --git a/go/internal/runnerhub/handler.go b/go/internal/runnerhub/handler.go index 90de84335..cdc931844 100644 --- a/go/internal/runnerhub/handler.go +++ b/go/internal/runnerhub/handler.go @@ -121,14 +121,13 @@ func (h *Handler) Sessions(ctx context.Context, stream *connect.BidiStream[compa if !ok { return errUnauthenticated } - router, _, err := h.hub.routerFor(subj.ID) + router, enrollGen, err := h.hub.routerForStream(subj.ID) if err != nil { // A Sessions stream with no enrolled Runner — the Runner must Enroll // before opening Sessions. return connect.NewError(connect.CodeFailedPrecondition, err) } - enrollGen := h.hub.EnrollGeneration() router.setSessionUnknown(func(sessionID string) { h.hub.dropLostSessionDetached(ctx, enrollGen, subj.ID, sessionID, false) }) diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index a16bf6e0c..35357eba7 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -21,6 +21,8 @@ import ( "sync" "sync/atomic" + lru "github.com/hashicorp/golang-lru/v2" + compassv1 "github.com/RigelBuild/compass/go/gen/compass/v1" "github.com/RigelBuild/compass/go/internal/fabric" compassv1internal "github.com/RigelBuild/compass/go/internal/gen/compass/v1" @@ -451,9 +453,9 @@ type Hub struct { sessionAccounts map[string]sessionBinding // lifecycleSeqs maps a session to the highest RunnerSeq of an accepted lifecycle // frame. One lifetime's frames arrive in seq order, so a lower one is from a - // cancelled stream or a dead lifetime. Kept past unbind so the boundary outlives - // ERRORED; cleared on re-enroll, when RunnerSeq restarts. lifecycleMu-guarded. - lifecycleSeqs map[string]uint64 + // cancelled stream or a dead lifetime. An LRU, so a boundary outlives unbind but + // not maxLifecycleSeqs newer sessions; purged on re-enroll. lifecycleMu-guarded. + lifecycleSeqs *lru.Cache[string, uint64] // enrollMu fences session-frame delivery (read) against enroll (write), and // guards enrollGen, which counts enrollments so a stream opened before a // re-enroll cannot act on its sessions. Lock order: enrollMu, lifecycleMu, mu. @@ -522,7 +524,8 @@ func NewHub(lifecycle LifecycleSink, tail SessionTailSink, comms CommsCaller, lo containerAccounts: make(map[string]sessionBinding), sessionAccounts: make(map[string]sessionBinding), accountSessions: make(map[store.AccountID]string), - lifecycleSeqs: make(map[string]uint64), + lifecycleSeqs: newLifecycleSeqs(), + enrollGen: 1, missingSeqs: make(map[uint64]struct{}), reapStale: make(map[string]uint64), runnerEpoch: make(map[string]uint64), @@ -703,6 +706,12 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { h.enrollMu.RLock() defer h.enrollMu.RUnlock() stale := h.staleEnrollmentLocked(ev) + // A lifecycle frame is admitted under lifecycleMu from before its seq is + // recorded, so a higher seq the hub has already seen is never overtaken. + if !stale && sessionState(ev.Frame) != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED { + h.lifecycleMu.Lock() + defer h.lifecycleMu.Unlock() + } if !stale { h.recordSeq(ev.RunnerSeq) } @@ -846,15 +855,14 @@ func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1i h.tail.RelaySessionFrame(sessionID, sf) return } - h.lifecycleMu.Lock() - if last, seen := h.lifecycleSeqs[sessionID]; seen && seq <= last { - h.lifecycleMu.Unlock() + // Deliver holds lifecycleMu for a lifecycle frame. + if last, seen := h.lifecycleSeqs.Get(sessionID); seen && seq <= last { h.log.Debug("ignored lifecycle frame older than the session's latest", slog.String("runner_id", runnerID), slog.String("session_id", sessionID), slog.String("state", state.String()), slog.Uint64("runner_seq", seq), slog.Uint64("latest_seq", last)) return } - h.lifecycleSeqs[sessionID] = seq + h.lifecycleSeqs.Add(sessionID, seq) h.tail.RelaySessionFrame(sessionID, sf) // Resolve the session's agent account and stamp it onto the published status — the // DL-167 attribution join. A status published after a Runner reconnect cleared the @@ -887,7 +895,6 @@ func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1i if presence != nil && hasAccount { presence.OnSessionLifecycle(account, sessionID, state) } - h.lifecycleMu.Unlock() // The Runner can see the exit before any deliver is refused, so ERRORED is a loss too. if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED && hasAccount { h.dropLostSessionDetached(ctx, h.enrollGen, runnerID, sessionID, true) @@ -992,6 +999,23 @@ func (h *Hub) forgeNotificationAck(ctx context.Context, ev RunnerEvent, ack *com } } +// maxLifecycleSeqs bounds lifecycleSeqs. A buffered stale frame lands within +// seconds, long before this many newer sessions could evict its boundary. +const maxLifecycleSeqs = 4096 + +func newLifecycleSeqs() *lru.Cache[string, uint64] { + c, err := lru.New[string, uint64](maxLifecycleSeqs) + if err != nil { + panic(err) // only a non-positive size errors + } + return c +} + +// sessionState is a session frame's lifecycle state, UNSPECIFIED for any other frame. +func sessionState(frame *compassv1internal.AgentFrame) compassv1.AgentSessionState { + return frame.GetSession().GetState() +} + // staleEnrollmentLocked reports, and logs, an event whose stream predates the // current enrollment. Zero EnrollGen means unfenced. Caller holds enrollMu. func (h *Hub) staleEnrollmentLocked(ev RunnerEvent) bool { @@ -1080,12 +1104,13 @@ type promotedPair struct { // so none of its pre-enroll sessions live. A failed durable reap still runs the // in-RAM fallback, then returns the error so the Runner retries enrollment. func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier compassv1.RuntimeTier, egressPosture compassv1.EgressPosture) (reattached bool, err error) { + // Held until the runner and router are replaced, so a stream that snapshots + // both under it sees one enrollment's pair. h.enrollMu.Lock() h.lifecycleMu.Lock() - clear(h.lifecycleSeqs) + h.lifecycleSeqs.Purge() h.enrollGen++ h.lifecycleMu.Unlock() - h.enrollMu.Unlock() // Held from the map-clear through the reap, so no promotion lands in between. h.bindingWriteMu.Lock() h.mu.Lock() @@ -1128,6 +1153,7 @@ func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier h.reapStale[id] = epoch } h.mu.Unlock() + h.enrollMu.Unlock() // Choose the reap set. With a durable store the ROWS are reaped and the edges // driven from what was removed (the authoritative set); without one, the in-RAM @@ -1202,6 +1228,15 @@ func (h *Hub) enroll(ctx context.Context, id string, subject store.Subject, tier // no Runner is enrolled. The id travels out with the router so a caller that must // attribute the call to a Runner names the one that served it, rather than re-reading // the registry and racing a re-enroll onto the wrong id. +// routerForStream returns the Runner's router and the enrollment generation it +// belongs to, read as one pair. +func (h *Hub) routerForStream(sessionID string) (*commandRouter, uint64, error) { + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + router, _, err := h.routerFor(sessionID) + return router, h.enrollGen, err +} + func (h *Hub) routerFor(sessionID string) (*commandRouter, string, error) { h.mu.Lock() defer h.mu.Unlock() diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 4dbe92874..707f3cc2b 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -198,6 +198,93 @@ func TestErroredOlderThanNewLifetimeIsIgnored(t *testing.T) { lost.none(t, "an older ERRORED must not retire the resumed lifetime") } +func TestLostSessionCleanupFromOldEnrollmentKeepsRebinding(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + oldGen := hub.EnrollGeneration() + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + + hub.dropLostSessionIfCurrent(ctx, oldGen, "runner-1", "sess-1", true) + lost.none(t, "cleanup from the old enrollment must not unbind the re-bound session") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after stale cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } + + hub.dropLostSessionIfCurrent(ctx, hub.EnrollGeneration(), "runner-1", "sess-1", true) + if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { + t.Fatalf("current-enrollment cleanup report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) + } +} + +func TestLowerSeqCannotOvertakeHigherSeqMidDelivery(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + store := &pausingResolveBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(store) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(store.release) }) } + t.Cleanup(release) + + // READY seq 6 has recorded its seq and pauses resolving the (uncached) binding. + readyDone := make(chan error, 1) + go func() { + readyDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 6, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_READY), + }) + }() + select { + case <-store.entered: + case <-time.After(10 * time.Second): + t.Fatal("READY did not reach the binding lookup") + } + if hub.lifecycleMu.TryLock() { + hub.lifecycleMu.Unlock() + t.Fatal("lifecycle lock free while READY seq 6 is mid-delivery; seq 5 could overtake it") + } + erroredDone := make(chan error, 1) + go func() { + erroredDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 5, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED), + }) + }() + release() + for name, done := range map[string]chan error{"READY": readyDone, "ERRORED": erroredDone} { + select { + case err := <-done: + if err != nil { + t.Fatalf("Deliver(%s) = %v, want nil", name, err) + } + case <-time.After(10 * time.Second): + t.Fatalf("Deliver(%s) did not complete", name) + } + } + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) +} + +// pausingResolveBindingStore holds the first ResolveSessionBinding until released. +type pausingResolveBindingStore struct { + *fakeBindingStore + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (b *pausingResolveBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) { + b.once.Do(func() { + close(b.entered) + <-b.release + }) + return b.fakeBindingStore.ResolveSessionBinding(ctx, sessionID) +} + func TestReenrollClearsErroredBoundary(t *testing.T) { ctx := t.Context() hub, lifecycle, _ := newHub() @@ -213,7 +300,7 @@ func TestReenrollClearsErroredBoundary(t *testing.T) { compassv1.AgentSessionState_AGENT_SESSION_STATE_READY) } -func TestReenrollWaitsForInFlightErroredAndFencesItsCleanup(t *testing.T) { +func TestReenrollWaitsForInFlightErrored(t *testing.T) { ctx := t.Context() lifecycle := &fakeLifecycleSink{} tail := &pausingTailSink{entered: make(chan struct{}), release: make(chan struct{})} @@ -272,10 +359,6 @@ func TestReenrollWaitsForInFlightErroredAndFencesItsCleanup(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("re-enroll did not complete") } - // Readers queued behind enroll, the cleanup among them, finish before this Lock. - hub.enrollMu.Lock() - lost.none(t, "the old ERRORED's cleanup must not unbind the re-bound session") - hub.enrollMu.Unlock() if _, _, err := hub.routerFor("sess-1"); err != nil { t.Fatalf("routerFor(sess-1) after re-bind = %v, want the binding kept", err) } diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index 5345895aa..b0d8c5b4d 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -919,16 +919,21 @@ func (h *Hub) dropLostSessionDetached(ctx context.Context, gen uint64, runnerID, go func() { dctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), lostSessionTimeout) defer cancel() - // A re-enroll since the loss may have re-bound sessionID; hold it off until done. - h.enrollMu.RLock() - defer h.enrollMu.RUnlock() - if h.enrollGen != gen { - return - } - h.dropLostSession(dctx, runnerID, sessionID, errored) + h.dropLostSessionIfCurrent(dctx, gen, runnerID, sessionID, errored) }() } +// dropLostSessionIfCurrent runs dropLostSession only while enrollment gen is +// current: a re-enroll since the loss may have re-bound sessionID. +func (h *Hub) dropLostSessionIfCurrent(ctx context.Context, gen uint64, runnerID, sessionID string, errored bool) { + h.enrollMu.RLock() + defer h.enrollMu.RUnlock() + if h.enrollGen != gen { + return + } + h.dropLostSession(ctx, runnerID, sessionID, errored) +} + // runnerSessionCtx scopes a Runner-originated ctx to the tenant that binds // sessionID to runnerID. Runners are shared across tenants, so the door sets none; // without this, binding reads and deletes run under the bootstrap tenant. The diff --git a/go/internal/runnerhub/seam_test.go b/go/internal/runnerhub/seam_test.go index ac8154d71..908f26034 100644 --- a/go/internal/runnerhub/seam_test.go +++ b/go/internal/runnerhub/seam_test.go @@ -23,7 +23,9 @@ import ( "context" "errors" "io" + "runtime" "testing" + "time" "connectrpc.com/connect" @@ -121,6 +123,47 @@ func TestSeamEnrollSessionsRoundTripAndPublishEvents(t *testing.T) { } } +func TestSeamPublishEventsOpenedBeforeEnrollIsFenced(t *testing.T) { + hub, _, tail := newHub() + resolver := &fakeResolver{tokens: map[string]resolverEntry{ + "runner-tok": {subj: store.Subject{Kind: store.SubjectRunner, ID: "runner-1"}}, + }} + url := newMountedH2CServer(t, hub, resolver.resolve) + client := newRawRunnerClient(t, url, "runner-tok") + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + pub := client.PublishEvents(ctx) + if err := pub.Send(&compassv1internal.PublishEventsRequest{RunnerSeq: 1, SessionId: "sess-wire", Frame: sessionTraceFrame("before enroll")}); err != nil { + t.Fatalf("PublishEvents.Send = %v", err) + } + // Gate on the server handler having taken the frame, so the stream opened pre-enroll. + for deadline := time.Now().Add(10 * time.Second); ; runtime.Gosched() { + hub.mu.Lock() + seen := hub.lastSeq == 1 + hub.mu.Unlock() + if seen { + break + } + if time.Now().After(deadline) { + t.Fatal("pre-enroll frame never reached Deliver") + } + } + if _, err := client.Enroll(ctx, connect.NewRequest(&compassv1internal.EnrollRequest{RunnerId: "runner-1"})); err != nil { + t.Fatalf("Enroll = %v, want success", err) + } + bindSession(hub, "sess-wire") + if err := pub.Send(&compassv1internal.PublishEventsRequest{RunnerSeq: 2, SessionId: "sess-wire", Frame: sessionTraceFrame("after enroll")}); err != nil { + t.Fatalf("PublishEvents.Send = %v", err) + } + if _, err := pub.CloseAndReceive(); err != nil { + t.Fatalf("PublishEvents.CloseAndReceive = %v", err) + } + if calls := tail.snapshot(); len(calls) != 0 { + t.Fatalf("tail sink saw %d frames, want none from a stream opened before enroll", len(calls)) + } +} + // startRoundTrip is the outcome of Hub.Start driven over the wire. type startRoundTrip struct { resp *compassv1.StartAgentSessionResponse From ca21ef317f4e1c76cd5df3556dd3b54346674d8a Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 18:15:02 -0400 Subject: [PATCH 10/17] fix(runnerhub): serialize lifecycle frames per session, not hub-wide (RIG-4673) A cold binding lookup ran under the hub-wide lifecycleMu, so one slow store read stalled lifecycle delivery for every session. Each session now has its own refcounted lock held from seq record through publication. lifecycleMu only guards the seq LRU and the lock map. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/hub.go | 58 ++++++++++++++--- go/internal/runnerhub/lost_session_test.go | 72 ++++++++++++++++++++-- 2 files changed, 116 insertions(+), 14 deletions(-) diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 35357eba7..13a820827 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -434,9 +434,14 @@ type Hub struct { // bindingWriteMu serializes whole enrolls (map-clear through reap) with promotion // writes and cache updates. Lock it before mu; never hold mu across a store call. bindingWriteMu sync.Mutex - // lifecycleMu serializes lifecycle guards and their sink publication order. + // lifecycleMu guards lifecycleSeqs and sessionLocks. Held only for map access; + // a session's ordering is held by its own sessionLocks entry. lifecycleMu sync.Mutex - mu sync.Mutex + // sessionLocks serialize one session's lifecycle frames from seq record through + // publication, so a slow binding lookup stalls only that session. Refcounted and + // removed when idle. lifecycleMu-guarded. + sessionLocks map[string]*sessionLock + mu sync.Mutex // runner is the single attached Runner (single-Runner MVP, OQ6). A second // enrollment re-attaches rather than registering a second entry. @@ -458,7 +463,8 @@ type Hub struct { lifecycleSeqs *lru.Cache[string, uint64] // enrollMu fences session-frame delivery (read) against enroll (write), and // guards enrollGen, which counts enrollments so a stream opened before a - // re-enroll cannot act on its sessions. Lock order: enrollMu, lifecycleMu, mu. + // re-enroll cannot act on its sessions. Lock order: enrollMu, bindingWriteMu, a + // session lock, lifecycleMu, mu. enrollMu sync.RWMutex enrollGen uint64 // accountSessions is the REVERSE of sessionAccounts (account -> live session_id), @@ -525,6 +531,7 @@ func NewHub(lifecycle LifecycleSink, tail SessionTailSink, comms CommsCaller, lo sessionAccounts: make(map[string]sessionBinding), accountSessions: make(map[store.AccountID]string), lifecycleSeqs: newLifecycleSeqs(), + sessionLocks: make(map[string]*sessionLock), enrollGen: 1, missingSeqs: make(map[uint64]struct{}), reapStale: make(map[string]uint64), @@ -706,11 +713,10 @@ func (h *Hub) Deliver(ctx context.Context, ev RunnerEvent) error { h.enrollMu.RLock() defer h.enrollMu.RUnlock() stale := h.staleEnrollmentLocked(ev) - // A lifecycle frame is admitted under lifecycleMu from before its seq is - // recorded, so a higher seq the hub has already seen is never overtaken. + // A lifecycle frame holds its session's lock from before its seq is recorded, + // so a higher seq the hub has already seen is never overtaken. if !stale && sessionState(ev.Frame) != compassv1.AgentSessionState_AGENT_SESSION_STATE_UNSPECIFIED { - h.lifecycleMu.Lock() - defer h.lifecycleMu.Unlock() + defer h.lockSession(ev.SessionID)() } if !stale { h.recordSeq(ev.RunnerSeq) @@ -855,14 +861,19 @@ func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1i h.tail.RelaySessionFrame(sessionID, sf) return } - // Deliver holds lifecycleMu for a lifecycle frame. - if last, seen := h.lifecycleSeqs.Get(sessionID); seen && seq <= last { + // Deliver holds the session's lock for a lifecycle frame. + h.lifecycleMu.Lock() + last, seen := h.lifecycleSeqs.Get(sessionID) + if !seen || seq > last { + h.lifecycleSeqs.Add(sessionID, seq) + } + h.lifecycleMu.Unlock() + if seen && seq <= last { h.log.Debug("ignored lifecycle frame older than the session's latest", slog.String("runner_id", runnerID), slog.String("session_id", sessionID), slog.String("state", state.String()), slog.Uint64("runner_seq", seq), slog.Uint64("latest_seq", last)) return } - h.lifecycleSeqs.Add(sessionID, seq) h.tail.RelaySessionFrame(sessionID, sf) // Resolve the session's agent account and stamp it onto the published status — the // DL-167 attribution join. A status published after a Runner reconnect cleared the @@ -1011,6 +1022,33 @@ func newLifecycleSeqs() *lru.Cache[string, uint64] { return c } +// sessionLock is one session's lifecycle mutex; refs counts holders and waiters. +type sessionLock struct { + mu sync.Mutex + refs int +} + +// lockSession locks sessionID's lifecycle mutex and returns its unlock. +func (h *Hub) lockSession(sessionID string) func() { + h.lifecycleMu.Lock() + l := h.sessionLocks[sessionID] + if l == nil { + l = &sessionLock{} + h.sessionLocks[sessionID] = l + } + l.refs++ + h.lifecycleMu.Unlock() + l.mu.Lock() + return func() { + l.mu.Unlock() + h.lifecycleMu.Lock() + if l.refs--; l.refs == 0 { + delete(h.sessionLocks, sessionID) + } + h.lifecycleMu.Unlock() + } +} + // sessionState is a session frame's lifecycle state, UNSPECIFIED for any other frame. func sessionState(frame *compassv1internal.AgentFrame) compassv1.AgentSessionState { return frame.GetSession().GetState() diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 707f3cc2b..b4ee702b2 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -244,8 +244,7 @@ func TestLowerSeqCannotOvertakeHigherSeqMidDelivery(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("READY did not reach the binding lookup") } - if hub.lifecycleMu.TryLock() { - hub.lifecycleMu.Unlock() + if !sessionLockHeld(hub, "sess-1") { t.Fatal("lifecycle lock free while READY seq 6 is mid-delivery; seq 5 could overtake it") } erroredDone := make(chan error, 1) @@ -285,6 +284,72 @@ func (b *pausingResolveBindingStore) ResolveSessionBinding(ctx context.Context, return b.fakeBindingStore.ResolveSessionBinding(ctx, sessionID) } +// sessionLockHeld reports whether sessionID's lifecycle lock is held right now. +func sessionLockHeld(hub *Hub, sessionID string) bool { + hub.lifecycleMu.Lock() + l := hub.sessionLocks[sessionID] + hub.lifecycleMu.Unlock() + if l == nil { + return false + } + if l.mu.TryLock() { + l.mu.Unlock() + return false + } + return true +} + +func TestSlowBindingLookupStallsOnlyItsSession(t *testing.T) { + ctx := t.Context() + hub, lifecycle, _ := newHub() + store := &pausingResolveBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(store) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(store.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-2", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-2", "sess-2") + + slowDone := make(chan error, 1) + go func() { + slowDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 1, SessionID: "sess-1", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_READY), + }) + }() + select { + case <-store.entered: + case <-time.After(10 * time.Second): + t.Fatal("sess-1 did not reach the binding lookup") + } + fastDone := make(chan error, 1) + go func() { + fastDone <- hub.Deliver(ctx, RunnerEvent{ + RunnerID: "runner-1", RunnerSeq: 2, SessionID: "sess-2", + Frame: sessionStateFrame(compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING), + }) + }() + select { + case err := <-fastDone: + if err != nil { + t.Fatalf("Deliver(sess-2) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("sess-2 lifecycle delivery stalled behind sess-1's binding lookup") + } + assertPublished(t, lifecycle, compassv1.AgentSessionState_AGENT_SESSION_STATE_WORKING) + release() + select { + case err := <-slowDone: + if err != nil { + t.Fatalf("Deliver(sess-1) = %v, want nil", err) + } + case <-time.After(10 * time.Second): + t.Fatal("sess-1 did not complete") + } +} + func TestReenrollClearsErroredBoundary(t *testing.T) { ctx := t.Context() hub, lifecycle, _ := newHub() @@ -467,8 +532,7 @@ func TestConcurrentErroredPublishesAfterInFlightLifecycle(t *testing.T) { case <-time.After(10 * time.Second): t.Fatal("WORKING lifecycle publish did not block") } - if hub.lifecycleMu.TryLock() { - hub.lifecycleMu.Unlock() + if !sessionLockHeld(hub, "sess-1") { t.Fatal("lifecycle lock released while WORKING is still being published") } From 73c1eb4ca4cf79a68db3aa15c52701eb6cb30a0b Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 19:16:12 -0400 Subject: [PATCH 11/17] fix(runnerhub): ERRORED cleanup releases only the binding it saw (RIG-4742) The detached ERRORED cleanup unbound whatever row held the session id when it ran, so a resume that re-bound the same id in the same enrollment could be unbound and archived by the old lifetime's cleanup. A binding now carries the durable row version (xmin) and a cache lifetime. The cleanup passes the binding it saw; dropLostSession skips a different one, and releaseSession re-checks the cache around a delete conditioned on that version. DeleteSessionBinding takes a version and reports a removal, so a peer's re-bind is also left alone. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/binding_cache_test.go | 52 +++-- go/internal/runnerhub/handler.go | 2 +- go/internal/runnerhub/hub.go | 19 +- go/internal/runnerhub/lost_session_test.go | 188 +++++++++++++++++- go/internal/runnerhub/relay_comms.go | 87 ++++++-- .../runnerhub/runner_tenant_pgtest_test.go | 8 +- .../store/compute_usage_pgtest_test.go | 12 +- go/internal/store/db/querier.go | 9 +- go/internal/store/db/session_bindings.sql.go | 58 ++++-- .../store/queries/session_bindings.sql | 38 ++-- go/internal/store/session_bindings.go | 62 +++--- .../store/session_bindings_pgtest_test.go | 99 ++++++--- go/server/lifecycle_wake_pgtest_test.go | 2 +- 13 files changed, 481 insertions(+), 155 deletions(-) diff --git a/go/internal/runnerhub/binding_cache_test.go b/go/internal/runnerhub/binding_cache_test.go index 4c1e17683..057bf9481 100644 --- a/go/internal/runnerhub/binding_cache_test.go +++ b/go/internal/runnerhub/binding_cache_test.go @@ -54,6 +54,8 @@ type fakeBindingStore struct { mu sync.Mutex tenant store.TenantID bindings map[string]store.SessionBinding // session id -> binding + versions map[string]string // session id -> row version, new per write + written int resolveCalled bool resolveCtxSystemRole bool @@ -66,14 +68,14 @@ type fakeBindingStore struct { } func newFakeBindingStore() *fakeBindingStore { - return &fakeBindingStore{tenant: "tenant-a", bindings: map[string]store.SessionBinding{}} + return &fakeBindingStore{tenant: "tenant-a", bindings: map[string]store.SessionBinding{}, versions: map[string]string{}} } -func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, error) { +func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, string, error) { f.mu.Lock() defer f.mu.Unlock() if f.recordErr != nil { - return "", f.recordErr + return "", "", f.recordErr } // session_bindings_session_key: this session id already belongs to a // DIFFERENT account. The real store raises ErrConflict rather than @@ -81,7 +83,7 @@ func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID str // single-valued; a fake that silently overwrote would hide the whole // class (a Runner restart re-mints "sess-1", so id reuse is routine). if b, ok := f.bindings[sessionID]; ok && b.AccountID != accountID { - return "", fmt.Errorf("%w: session %q is already bound to a different agent", store.ErrConflict, sessionID) + return "", "", fmt.Errorf("%w: session %q is already bound to a different agent", store.ErrConflict, sessionID) } var displaced string // Account-keyed UPSERT: a prior binding for this account is displaced. @@ -89,26 +91,26 @@ func (f *fakeBindingStore) RecordSessionBinding(_ context.Context, sessionID str if b.AccountID == accountID && sid != sessionID { displaced = sid delete(f.bindings, sid) + delete(f.versions, sid) break } } - f.bindings[sessionID] = store.SessionBinding{TenantID: f.tenant, SessionID: sessionID, AccountID: accountID, RunnerID: runnerID} - return displaced, nil + return displaced, f.putLocked(sessionID, accountID, runnerID), nil } -func (f *fakeBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) { +func (f *fakeBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, string, error) { f.mu.Lock() defer f.mu.Unlock() f.resolveCalled = true f.resolveCtxSystemRole = store.IsSystemRole(ctx) if f.resolveErr != nil { - return "", "", f.resolveErr + return "", "", "", f.resolveErr } b, ok := f.bindings[sessionID] if !ok { - return "", "", store.ErrNotFound + return "", "", "", store.ErrNotFound } - return b.AccountID, b.RunnerID, nil + return b.AccountID, b.RunnerID, f.versions[sessionID], nil } func (f *fakeBindingStore) SessionForAccount(_ context.Context, accountID store.AccountID) (string, string, error) { @@ -125,11 +127,15 @@ func (f *fakeBindingStore) SessionForAccount(_ context.Context, accountID store. return "", "", store.ErrNotFound } -func (f *fakeBindingStore) DeleteSessionBinding(_ context.Context, sessionID string) error { +func (f *fakeBindingStore) DeleteSessionBinding(_ context.Context, sessionID, version string) (bool, error) { f.mu.Lock() defer f.mu.Unlock() + if _, ok := f.bindings[sessionID]; !ok || (version != "" && f.versions[sessionID] != version) { + return false, nil + } delete(f.bindings, sessionID) - return nil + delete(f.versions, sessionID) + return true, nil } func (f *fakeBindingStore) DeleteSessionBindingsForRunner(_ context.Context, runnerID string) ([]store.SessionBinding, error) { @@ -143,6 +149,7 @@ func (f *fakeBindingStore) DeleteSessionBindingsForRunner(_ context.Context, run if b.RunnerID == runnerID { removed = append(removed, b) delete(f.bindings, sid) + delete(f.versions, sid) } } return removed, nil @@ -173,6 +180,15 @@ func (f *fakeBindingStore) SessionBindingTenant(_ context.Context, sessionID, ru return "", store.ErrNotFound } +// putLocked writes a binding row with a fresh version, as every upsert does. +func (f *fakeBindingStore) putLocked(sessionID string, accountID store.AccountID, runnerID string) string { + f.written++ + version := fmt.Sprintf("v%d", f.written) + f.bindings[sessionID] = store.SessionBinding{TenantID: f.tenant, SessionID: sessionID, AccountID: accountID, RunnerID: runnerID} + f.versions[sessionID] = version + return version +} + func (f *fakeBindingStore) seed(sessionID string) { f.seedBinding(sessionID, testAgentAccount, testRunnerID) } @@ -180,7 +196,7 @@ func (f *fakeBindingStore) seed(sessionID string) { func (f *fakeBindingStore) seedBinding(sessionID string, accountID store.AccountID, runnerID string) { f.mu.Lock() defer f.mu.Unlock() - f.bindings[sessionID] = store.SessionBinding{TenantID: f.tenant, SessionID: sessionID, AccountID: accountID, RunnerID: runnerID} + f.putLocked(sessionID, accountID, runnerID) } // fakeRoutingFabric is an in-memory RoutingFabric double: PublishBindingChange @@ -843,11 +859,11 @@ func (b *reapAndReadBlockingBindingStore) DeleteSessionBindingsForRunner(ctx con return b.fakeBindingStore.DeleteSessionBindingsForRunner(ctx, runnerID) } -func (b *reapAndReadBlockingBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) { - account, runnerID, err := b.fakeBindingStore.ResolveSessionBinding(ctx, sessionID) +func (b *reapAndReadBlockingBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, string, error) { + account, runnerID, version, err := b.fakeBindingStore.ResolveSessionBinding(ctx, sessionID) close(b.readEntered) <-b.readRelease - return account, runnerID, err + return account, runnerID, version, err } func (b *reapAndReadBlockingBindingStore) SessionForAccount(ctx context.Context, accountID store.AccountID) (string, string, error) { @@ -985,7 +1001,7 @@ func (b *overlapDetectingBindingStore) DeleteSessionBindingsForRunner(ctx contex return b.fakeBindingStore.DeleteSessionBindingsForRunner(ctx, runnerID) } -func (b *overlapDetectingBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, error) { +func (b *overlapDetectingBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, string, error) { b.mu.Lock() if b.parked { b.overlap = true @@ -1134,7 +1150,7 @@ type blockingRecordBindingStore struct { recordRelease chan struct{} } -func (b *blockingRecordBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, error) { +func (b *blockingRecordBindingStore) RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (string, string, error) { close(b.recordEntered) <-b.recordRelease return b.fakeBindingStore.RecordSessionBinding(ctx, sessionID, accountID, runnerID) diff --git a/go/internal/runnerhub/handler.go b/go/internal/runnerhub/handler.go index cdc931844..fcf5672d5 100644 --- a/go/internal/runnerhub/handler.go +++ b/go/internal/runnerhub/handler.go @@ -129,7 +129,7 @@ func (h *Handler) Sessions(ctx context.Context, stream *connect.BidiStream[compa } router.setSessionUnknown(func(sessionID string) { - h.hub.dropLostSessionDetached(ctx, enrollGen, subj.ID, sessionID, false) + h.hub.dropLostSessionDetached(ctx, enrollGen, subj.ID, sessionID, nil, false) }) router.attach(stream.Send) defer router.detach(errStreamClosed) diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index 13a820827..d68c5ccbb 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -213,19 +213,21 @@ type SessionBindingStore interface { // returns the session id it DISPLACED (empty when the account held none) — // the prior session the hub must evict from both maps. It runs on the // request ctx (tenant-scoped), so the write lands under the acting tenant. - RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (displaced string, err error) + // version names the written row; a release conditioned on it cannot remove a re-bind. + RecordSessionBinding(ctx context.Context, sessionID string, accountID store.AccountID, runnerID string) (displaced, version string, err error) // ResolveSessionBinding resolves the agent account and owning Runner a live // session speaks for — the cache-miss read behind accountForSession. // store.ErrNotFound is the fail-closed miss (mapped to ok=false), - // byte-identical to today's CodeNotFound. - ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) + // byte-identical to today's CodeNotFound. version names the row read. + ResolveSessionBinding(ctx context.Context, sessionID string) (account store.AccountID, runnerID, version string, err error) // SessionForAccount resolves the live session bound to an account — the // cache-miss read behind SessionForAccount (the reverse direction). Same // fail-closed store.ErrNotFound contract. Returns the owning Runner id too. SessionForAccount(ctx context.Context, accountID store.AccountID) (sessionID, runnerID string, err error) // DeleteSessionBinding releases one session's binding — the unbind write. // Idempotent: releasing an already-released session is a no-op success. - DeleteSessionBinding(ctx context.Context, sessionID string) error + // A non-empty version releases only that row; removed reports a deletion. + DeleteSessionBinding(ctx context.Context, sessionID, version string) (removed bool, err error) // DeleteSessionBindingsForRunner is the enroll sweep: it releases every // binding attached to runnerID and RETURNS the rows it removed, driving the // enroll reap (offline edges + held-deliver reap) from durable truth @@ -411,6 +413,8 @@ type Hub struct { reapStale map[string]uint64 // bindingEpoch fences read-throughs across reap start and completion. bindingEpoch uint64 + // lastLifetime numbers sessionAccounts inserts (sessionBinding.lifetime). mu-guarded. + lastLifetime uint64 // runnerEpoch entries are bounded by the number of distinct Runner ids. runnerEpoch map[string]uint64 // lifecycleCaller is the spawn/despawn execution seam RelayLifecycleCall delegates @@ -506,9 +510,14 @@ type attachedRunner struct { } // sessionBinding is one live session's principal and the Runner that owns it. +// version and lifetime name one binding of the session id, so a release from an +// older lifetime can skip a re-bind: version is the durable row ("" when none was +// written), lifetime is the cache entry, unique per insert and never zero. type sessionBinding struct { account store.AccountID runnerID string + version string + lifetime uint64 } // NewHub constructs a hub over the two write-through sinks and the agent-comms @@ -908,7 +917,7 @@ func (h *Hub) deliverSession(ctx context.Context, ev RunnerEvent, sf *compassv1i } // The Runner can see the exit before any deliver is refused, so ERRORED is a loss too. if state == compassv1.AgentSessionState_AGENT_SESSION_STATE_ERRORED && hasAccount { - h.dropLostSessionDetached(ctx, h.enrollGen, runnerID, sessionID, true) + h.dropLostSessionDetached(ctx, h.enrollGen, runnerID, sessionID, &binding, true) } } diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index b4ee702b2..23e5b5680 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -67,13 +67,13 @@ func TestDropLostSessionOnlyForOwningRunner(t *testing.T) { hub.bindContainer("cont-1", testAgentAccount, "runner-1") hub.promoteSession(ctx, "cont-1", "sess-1") - hub.dropLostSession(ctx, "runner-other", "sess-1", false) + hub.dropLostSession(ctx, "runner-other", "sess-1", nil, false) if _, ok := hub.accountForSession(ctx, "sess-1"); !ok { t.Fatal("foreign refusal unbound sess-1") } sink.none(t, "foreign refusal") - hub.dropLostSession(ctx, "runner-1", "sess-1", false) + hub.dropLostSession(ctx, "runner-1", "sess-1", nil, false) if _, ok := hub.accountForSession(ctx, "sess-1"); ok { t.Fatal("owning refusal left sess-1 bound") } @@ -198,6 +198,180 @@ func TestErroredOlderThanNewLifetimeIsIgnored(t *testing.T) { lost.none(t, "an older ERRORED must not retire the resumed lifetime") } +func TestErroredCleanupKeepsSameEnrollmentRebind(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingTenantBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old, ok := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + if !ok { + t.Fatal("sess-1 not bound after promote") + } + + // The old lifetime's cleanup pauses resolving the tenant; a resume re-binds + // the same session in the same enrollment before it continues. + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("ERRORED cleanup did not start") + } + hub.promoteSession(ctx, "cont-1", "sess-1") + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("ERRORED cleanup did not finish") + } + + lost.none(t, "the old lifetime's cleanup must not report the resumed session lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } + if account, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil || account != testAgentAccount { + t.Fatalf("durable binding after old cleanup = (%s, %v), want (%s, nil)", account, err, testAgentAccount) + } +} + +// pausingTenantBindingStore holds the first SessionBindingTenant until released. +type pausingTenantBindingStore struct { + *fakeBindingStore + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (b *pausingTenantBindingStore) SessionBindingTenant(ctx context.Context, sessionID, runnerID string) (store.TenantID, error) { + b.once.Do(func() { + close(b.entered) + <-b.release + }) + return b.fakeBindingStore.SessionBindingTenant(ctx, sessionID, runnerID) +} + +func TestErroredCleanupKeepsRebindDuringDurableDelete(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingDeleteBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + // The cleanup has matched the old binding and pauses in its durable delete; + // a resume re-binds the session before the delete runs. + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not reach the durable delete") + } + hub.promoteSession(ctx, "cont-1", "sess-1") + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not finish") + } + + lost.none(t, "a cleanup whose row was re-bound must not report the session lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } + if account, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil || account != testAgentAccount { + t.Fatalf("durable binding after old cleanup = (%s, %v), want (%s, nil)", account, err, testAgentAccount) + } +} + +func TestErroredCleanupKeepsPeerRebind(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingDeleteBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not reach the durable delete") + } + // Another Server re-binds the row; this hub's cache still holds the old binding. + if _, _, err := bindings.fakeBindingStore.RecordSessionBinding(ctx, "sess-1", testAgentAccount, "runner-1"); err != nil { + t.Fatalf("peer RecordSessionBinding: %v", err) + } + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not finish") + } + + lost.none(t, "a cleanup whose row a peer re-bound must not report the session lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) + } +} + +// pausingDeleteBindingStore holds the first DeleteSessionBinding until released. +type pausingDeleteBindingStore struct { + *fakeBindingStore + once sync.Once + entered chan struct{} + release chan struct{} +} + +func (b *pausingDeleteBindingStore) DeleteSessionBinding(ctx context.Context, sessionID, version string) (bool, error) { + b.once.Do(func() { + close(b.entered) + <-b.release + }) + return b.fakeBindingStore.DeleteSessionBinding(ctx, sessionID, version) +} + func TestLostSessionCleanupFromOldEnrollmentKeepsRebinding(t *testing.T) { ctx := t.Context() hub, _, _ := newHub() @@ -209,13 +383,13 @@ func TestLostSessionCleanupFromOldEnrollmentKeepsRebinding(t *testing.T) { hub.bindContainer("cont-1", testAgentAccount, "runner-1") hub.promoteSession(ctx, "cont-1", "sess-1") - hub.dropLostSessionIfCurrent(ctx, oldGen, "runner-1", "sess-1", true) + hub.dropLostSessionIfCurrent(ctx, oldGen, "runner-1", "sess-1", nil, true) lost.none(t, "cleanup from the old enrollment must not unbind the re-bound session") if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { t.Fatalf("binding after stale cleanup = (%s, %v), want (%s, true)", account, ok, testAgentAccount) } - hub.dropLostSessionIfCurrent(ctx, hub.EnrollGeneration(), "runner-1", "sess-1", true) + hub.dropLostSessionIfCurrent(ctx, hub.EnrollGeneration(), "runner-1", "sess-1", nil, true) if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { t.Fatalf("current-enrollment cleanup report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) } @@ -276,7 +450,7 @@ type pausingResolveBindingStore struct { release chan struct{} } -func (b *pausingResolveBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, error) { +func (b *pausingResolveBindingStore) ResolveSessionBinding(ctx context.Context, sessionID string) (store.AccountID, string, string, error) { b.once.Do(func() { close(b.entered) <-b.release @@ -706,8 +880,8 @@ func TestSessionEndedWithoutStopIsArchived(t *testing.T) { hub.bindContainer("cont-1", testAgentAccount, "runner-1") hub.promoteSession(ctx, "cont-1", "sess-lost") - hub.dropLostSession(ctx, "runner-other", "sess-lost", false) - hub.dropLostSession(ctx, "runner-1", "sess-lost", false) + hub.dropLostSession(ctx, "runner-other", "sess-lost", nil, false) + hub.dropLostSession(ctx, "runner-1", "sess-lost", nil, false) if got := recvEnded(t, ended); got != "sess-lost" { t.Fatalf("archived %q after the lost drop, want sess-lost", got) } diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index b0d8c5b4d..882ce1782 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -84,10 +84,10 @@ func (h *Hub) promoteSession(ctx context.Context, containerName, sessionID strin // Write the durable binding first (h.mu released). The store upsert is keyed // on the account, so it lands whether or not the account held a prior // session, and it returns the displaced session id. - var displaced string + var displaced, version string tenant := "" if bindings != nil && runnerID != "" { - d, err := bindings.RecordSessionBinding(ctx, sessionID, account, runnerID) + d, v, err := bindings.RecordSessionBinding(ctx, sessionID, account, runnerID) if err != nil { // A durable-write fault must not fail the Start that already succeeded // on the Runner: log it and fall back to the in-RAM cache so the session @@ -96,7 +96,7 @@ func (h *Hub) promoteSession(ctx context.Context, containerName, sessionID strin h.log.Error("record session binding failed; falling back to in-RAM cache", "session_id", sessionID, "account", string(account), "error", err) } else { - displaced = d + displaced, version = d, v tenant = string(bindings.EffectiveTenant(ctx)) } } @@ -105,7 +105,7 @@ func (h *Hub) promoteSession(ctx context.Context, containerName, sessionID strin h.mu.Lock() // Keep the container->account entry: a resume Start on this container fetches // secrets by container_name before exec. Remove and re-enroll clear it. - h.sessionAccounts[sessionID] = sessionBinding{account: account, runnerID: runnerID} + h.sessionAccounts[sessionID] = h.newBindingLocked(account, runnerID, version) h.accountSessions[account] = sessionID // Evict the displaced session from the forward map: the account moved off it, // so it now resolves nowhere. Guard displaced != sessionID for the rebind @@ -180,6 +180,13 @@ func (h *Hub) publishBindingChange(ctx context.Context, routing RoutingFabric, t // lock-then-release-then-fire discipline promoteSession uses, so the sink (which // enqueues into the presence loop) never runs under h.mu. func (h *Hub) unbindSession(ctx context.Context, sessionID string) { + h.releaseSession(ctx, sessionID, nil) +} + +// releaseSession is unbindSession limited, when only is non-nil, to that one +// binding of sessionID, so a release from an older lifetime leaves a re-bind +// alone. It reports whether the binding was released. +func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessionBinding) bool { // The maps are a cache, so the durable row is deleted FIRST (on the request ctx // so it stays tenant-scoped), then the maps are evicted under h.mu. Delete is // by session id and idempotent — a session already displaced has no row, so a @@ -187,17 +194,32 @@ func (h *Hub) unbindSession(ctx context.Context, sessionID string) { h.mu.Lock() bindings := h.bindings routing := h.routing + if only != nil { + if live, ok := h.sessionAccounts[sessionID]; ok && !sameBinding(live, *only) { + h.mu.Unlock() + return false + } + } h.mu.Unlock() tenant := "" if bindings != nil { - if err := bindings.DeleteSessionBinding(ctx, sessionID); err != nil { + version := "" + if only != nil { + version = only.version + } + removed, err := bindings.DeleteSessionBinding(ctx, sessionID, version) + switch { + case err != nil: // A durable-delete fault must not fail the Stop that already // succeeded on the Runner: log and continue to evict the cache. The // next re-enroll sweep retires any surviving row. h.log.Error("delete session binding failed; evicting cache anyway", "session_id", sessionID, "error", err) - } else { + case only != nil && version != "" && !removed: + // The row was re-bound since only was read: it is not ours to release. + return false + default: tenant = string(bindings.EffectiveTenant(ctx)) } } @@ -207,7 +229,13 @@ func (h *Hub) unbindSession(ctx context.Context, sessionID string) { account store.AccountID wentOffline bool ) - if binding, ok := h.sessionAccounts[sessionID]; ok { + binding, ok := h.sessionAccounts[sessionID] + if ok && only != nil && !sameBinding(binding, *only) { + // Re-bound in the cache while the durable delete ran. + h.mu.Unlock() + return false + } + if ok { // Drop the reverse entry only if it still points at THIS session — a // promoteSession for the account onto a newer session would have already // repointed it, and a stale delete would then unbind the live one. @@ -238,6 +266,13 @@ func (h *Hub) unbindSession(ctx context.Context, sessionID string) { if wentOffline && presence != nil { presence.OnSessionLifecycle(account, sessionID, compassv1.AgentSessionState_AGENT_SESSION_STATE_DISCONNECTED) } + return true +} + +// newBindingLocked builds a sessionAccounts entry with a fresh lifetime. Caller holds mu. +func (h *Hub) newBindingLocked(account store.AccountID, runnerID, version string) sessionBinding { + h.lastLifetime++ + return sessionBinding{account: account, runnerID: runnerID, version: version, lifetime: h.lastLifetime} } // unbindContainer drops a container's provisioned account binding. Remove is the @@ -357,7 +392,7 @@ func (h *Hub) lookupSessionBinding(ctx context.Context, sessionID string) (sessi if !h.readThroughAllowed(ctx, bindings, enrolled) { return sessionBinding{}, bindingUnverifiable } - account, runnerID, err := bindings.ResolveSessionBinding(ctx, sessionID) + account, runnerID, version, err := bindings.ResolveSessionBinding(ctx, sessionID) if errors.Is(err, store.ErrNotFound) { return sessionBinding{}, bindingNotFound } @@ -368,7 +403,6 @@ func (h *Hub) lookupSessionBinding(ctx context.Context, sessionID string) (sessi // session hits without a table round-trip. Re-check under the lock: a // concurrent promote/unbind may have run, so a live map entry wins over the // row just read (avoids clobbering a fresher binding with a staler one). - resolved := sessionBinding{account: account, runnerID: runnerID} h.mu.Lock() if live, ok := h.sessionAccounts[sessionID]; ok { h.mu.Unlock() @@ -380,6 +414,7 @@ func (h *Hub) lookupSessionBinding(ctx context.Context, sessionID string) (sessi h.mu.Unlock() return sessionBinding{}, bindingUnverifiable } + resolved := h.newBindingLocked(account, runnerID, version) h.sessionAccounts[sessionID] = resolved h.mu.Unlock() return resolved, bindingFound @@ -890,48 +925,60 @@ func commsCallError(err error) *compassv1internal.CommsCallError { } // dropLostSession unbinds, archives and reports a session its Runner lost. Only -// the owning Runner may unbind it. -func (h *Hub) dropLostSession(ctx context.Context, runnerID, sessionID string, errored bool) { +// the owning Runner may unbind it. A non-nil seen limits it to that binding: a +// resume may have re-bound the session id since the loss was observed. +func (h *Hub) dropLostSession(ctx context.Context, runnerID, sessionID string, seen *sessionBinding, errored bool) { ctx, scoped := h.runnerSessionCtx(ctx, runnerID, sessionID) if !scoped { return } - account, ok := h.accountForRunnerSession(ctx, runnerID, sessionID) - if !ok { + binding, ok := h.resolveSessionBinding(ctx, sessionID) + if !ok || runnerID == "" || binding.runnerID != runnerID || (seen != nil && !sameBinding(binding, *seen)) { + return + } + if !h.releaseSession(ctx, sessionID, &binding) { return } - h.unbindSession(ctx, sessionID) h.mu.Lock() lost := h.lost h.mu.Unlock() h.archiveEnded(ctx, sessionID) h.log.Warn("runner reports bound session lost; released binding", - "session_id", sessionID, "agent_account_id", account, "errored", errored) + "session_id", sessionID, "agent_account_id", binding.account, "errored", errored) if lost != nil { - lost.OnSessionLost(sessionID, account, errored) + lost.OnSessionLost(sessionID, binding.account, errored) } } // dropLostSessionDetached runs dropLostSession off the caller's receive loop: it // does store work, and the stream ctx dies with the stream. // gen is the enrollment the loss was observed under. -func (h *Hub) dropLostSessionDetached(ctx context.Context, gen uint64, runnerID, sessionID string, errored bool) { +func (h *Hub) dropLostSessionDetached(ctx context.Context, gen uint64, runnerID, sessionID string, seen *sessionBinding, errored bool) { go func() { dctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), lostSessionTimeout) defer cancel() - h.dropLostSessionIfCurrent(dctx, gen, runnerID, sessionID, errored) + h.dropLostSessionIfCurrent(dctx, gen, runnerID, sessionID, seen, errored) }() } // dropLostSessionIfCurrent runs dropLostSession only while enrollment gen is // current: a re-enroll since the loss may have re-bound sessionID. -func (h *Hub) dropLostSessionIfCurrent(ctx context.Context, gen uint64, runnerID, sessionID string, errored bool) { +func (h *Hub) dropLostSessionIfCurrent(ctx context.Context, gen uint64, runnerID, sessionID string, seen *sessionBinding, errored bool) { h.enrollMu.RLock() defer h.enrollMu.RUnlock() if h.enrollGen != gen { return } - h.dropLostSession(ctx, runnerID, sessionID, errored) + h.dropLostSession(ctx, runnerID, sessionID, seen, errored) +} + +// sameBinding reports whether a and b are one binding of a session id: the same +// durable row when either has one, else the same cache entry. +func sameBinding(a, b sessionBinding) bool { + if a.version != "" || b.version != "" { + return a.version == b.version + } + return a.lifetime == b.lifetime } // runnerSessionCtx scopes a Runner-originated ctx to the tenant that binds diff --git a/go/internal/runnerhub/runner_tenant_pgtest_test.go b/go/internal/runnerhub/runner_tenant_pgtest_test.go index ba00b6e7a..6025cd49e 100644 --- a/go/internal/runnerhub/runner_tenant_pgtest_test.go +++ b/go/internal/runnerhub/runner_tenant_pgtest_test.go @@ -28,16 +28,16 @@ func TestDropLostSessionScopesToTheSessionTenant(t *testing.T) { hub.SetSessionLostSink(sink) hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) // Enroll reaps every binding on this Runner; re-record so the drop reads a live row cold. - if _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { + if _, _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { t.Fatalf("RecordSessionBinding after enroll: %v", err) } - hub.dropLostSession(ctx, "runner-1", "sess-b", false) + hub.dropLostSession(ctx, "runner-1", "sess-b", nil, false) if lost, _ := sink.waitOne(t); lost != agent.ID { t.Fatalf("lost = %s, want %s: the tenant-B session was not resolved, so no wake", lost, agent.ID) } - if _, _, err := st.ResolveSessionBinding(ctxB, "sess-b"); err == nil { + if _, _, _, err := st.ResolveSessionBinding(ctxB, "sess-b"); err == nil { t.Fatal("tenant B's durable binding survived the drop; a cache miss would resurrect it") } } @@ -101,7 +101,7 @@ func openTenantBSession(t *testing.T, ctx context.Context) (*store.Store, store. if err != nil { t.Fatalf("CreateAgent: %v", err) } - if _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { + if _, _, err := st.RecordSessionBinding(ctxB, "sess-b", agent.ID, "runner-1"); err != nil { t.Fatalf("RecordSessionBinding: %v", err) } return st, agent, ctxB diff --git a/go/internal/store/compute_usage_pgtest_test.go b/go/internal/store/compute_usage_pgtest_test.go index 4ef91ab89..eae70168a 100644 --- a/go/internal/store/compute_usage_pgtest_test.go +++ b/go/internal/store/compute_usage_pgtest_test.go @@ -71,7 +71,7 @@ func TestComputeUsageBindingLifecycle(t *testing.T) { t.Fatalf("start event = %+v, want the tenant, owner, agent, session, runner, and non-estimated start", start) } - if err := s.DeleteSessionBinding(ctx, "compute-session"); err != nil { + if _, err := s.DeleteSessionBinding(ctx, "compute-session", ""); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } events = computeEvents(t, s, tenant, agent.ID) @@ -88,7 +88,7 @@ func TestComputeUsageBindingLifecycle(t *testing.T) { end.OwnerUserID != string(owner.ID) || end.SessionID != start.SessionID || end.RunnerID != start.RunnerID || end.Estimated { t.Fatalf("end event = %+v, want same interval metadata and a non-estimated end", end) } - if err := s.DeleteSessionBinding(ctx, "compute-session"); err != nil { + if _, err := s.DeleteSessionBinding(ctx, "compute-session", ""); err != nil { t.Fatalf("second DeleteSessionBinding: %v", err) } if got := computeEvents(t, s, tenant, agent.ID); len(got) != 2 { @@ -159,7 +159,7 @@ func TestComputeUsageEventsStampedAfterLockWait(t *testing.T) { } done := make(chan error, 1) go func() { - _, err := s.RecordSessionBinding(ctx, "wait-after", agent.ID, "runner-1") + _, _, err := s.RecordSessionBinding(ctx, "wait-after", agent.ID, "runner-1") done <- err }() deadline := time.After(5 * time.Second) @@ -223,7 +223,7 @@ func TestComputeUsageConflictRollsBackIntervalEvents(t *testing.T) { mustBind(t, ctx, s, "sess-2", accountY.ID, "runner-1") before := computeEvents(t, s, tenant, accountX.ID) - if _, err := s.RecordSessionBinding(ctx, "sess-2", accountX.ID, "runner-1"); !errors.Is(err, ErrConflict) { + if _, _, err := s.RecordSessionBinding(ctx, "sess-2", accountX.ID, "runner-1"); !errors.Is(err, ErrConflict) { t.Fatalf("RecordSessionBinding(X, sess-2) error = %v, want ErrConflict", err) } after := computeEvents(t, s, tenant, accountX.ID) @@ -259,7 +259,7 @@ func TestComputeUsageSameSessionRunnerRebindKeepsInterval(t *testing.T) { func TestComputeUsageLegacyBindingGetsEstimatedStart(t *testing.T) { for name, release := range map[string]func(*testing.T, context.Context, *Store){ "single release": func(t *testing.T, ctx context.Context, s *Store) { - if err := s.DeleteSessionBinding(ctx, "legacy-session"); err != nil { + if _, err := s.DeleteSessionBinding(ctx, "legacy-session", ""); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } }, @@ -300,7 +300,7 @@ func TestComputeUsageLegacyBindingReleaseLogsInterval(t *testing.T) { tenant := s.EffectiveTenant(ctx) execAsSystem(t, s, "INSERT INTO session_bindings (tenant_id, agent_account_id, session_id, runner_id) VALUES ($1, $2, 'legacy-only', 'runner-old')", string(tenant), string(agent.ID)) - if err := s.DeleteSessionBinding(ctx, "legacy-only"); err != nil { + if _, err := s.DeleteSessionBinding(ctx, "legacy-only", ""); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } events := computeEvents(t, s, tenant, agent.ID) diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 2c921b45f..6b5f69f27 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -127,7 +127,10 @@ type Querier interface { DeleteServerSecret(ctx context.Context, name string) (int64, error) // Both deletes also write an estimated start for a binding an older server made // without one; ON CONFLICT keeps any real start. - DeleteSessionBinding(ctx context.Context, sessionID string) error + // An empty version releases whatever row binds the session; a non-empty one + // releases only that version, so a stale release cannot remove a re-bind. + // Returns the rows removed. Postgres runs every data-modifying CTE to completion. + DeleteSessionBinding(ctx context.Context, arg DeleteSessionBindingParams) (int64, error) // The reconnect sweep, run by Hub.enroll under the system role because a Runner // is shared across tenants. :many with RETURNING: each removed row drives a // presence DISCONNECTED edge, a held-deliver reap, and a tenant-scoped archive. @@ -492,7 +495,9 @@ type Querier interface { RecordOwedMention(ctx context.Context, arg RecordOwedMentionParams) error // What it DISPLACED comes from SessionBindingForUpdate above, not from a // RETURNING here. The binding update and event writes share the Store tx. - RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error + // xmin is the row version: every upsert writes a new tuple, so a re-bind of the + // same session id still gets a new version a stale release can be fenced by. + RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) // Forge state-transition memo queries (compass-forge-state-transition §Actor // attribution). The write chokepoint upserts one memo per forge coordinate // AFTER a successful agent-driven transition; the notify lane consumes it on diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go index 5608f5704..8a449b8ff 100644 --- a/go/internal/store/db/session_bindings.sql.go +++ b/go/internal/store/db/session_bindings.sql.go @@ -11,10 +11,11 @@ import ( "github.com/jackc/pgx/v5/pgtype" ) -const deleteSessionBinding = `-- name: DeleteSessionBinding :exec +const deleteSessionBinding = `-- name: DeleteSessionBinding :one WITH d AS ( DELETE FROM session_bindings AS b WHERE b.session_id = $1 + AND ($2::text = '' OR b.xmin::text = $2::text) RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, b.session_id, b.runner_id, b.created_at ), starts AS ( @@ -28,23 +29,36 @@ WITH d AS ( JOIN agent_accounts AS a ON a.account_id = d.agent_account_id ON CONFLICT DO NOTHING RETURNING 1 +), ends AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 ) -INSERT INTO compute_usage_events ( - tenant_id, id, interval_id, kind, occurred_at, agent_account_id, - owner_user_id, session_id, runner_id -) -SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), - d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id - FROM d - JOIN agent_accounts AS a ON a.account_id = d.agent_account_id -ON CONFLICT DO NOTHING +SELECT count(*) FROM d ` +type DeleteSessionBindingParams struct { + SessionID string + Version string +} + // Both deletes also write an estimated start for a binding an older server made // without one; ON CONFLICT keeps any real start. -func (q *Queries) DeleteSessionBinding(ctx context.Context, sessionID string) error { - _, err := q.db.Exec(ctx, deleteSessionBinding, sessionID) - return err +// An empty version releases whatever row binds the session; a non-empty one +// releases only that version, so a stale release cannot remove a re-bind. +// Returns the rows removed. Postgres runs every data-modifying CTE to completion. +func (q *Queries) DeleteSessionBinding(ctx context.Context, arg DeleteSessionBindingParams) (int64, error) { + row := q.db.QueryRow(ctx, deleteSessionBinding, arg.SessionID, arg.Version) + var count int64 + err := row.Scan(&count) + return count, err } const deleteSessionBindingsForRunner = `-- name: DeleteSessionBindingsForRunner :many @@ -223,13 +237,14 @@ func (q *Queries) LockSessionBindingAccount(ctx context.Context, arg LockSession return err } -const recordSessionBinding = `-- name: RecordSessionBinding :exec +const recordSessionBinding = `-- name: RecordSessionBinding :one INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id) VALUES ($1, $2, $3, $4) ON CONFLICT (tenant_id, agent_account_id) DO UPDATE SET session_id = EXCLUDED.session_id, runner_id = EXCLUDED.runner_id, usage_interval_id = EXCLUDED.usage_interval_id +RETURNING xmin::text AS version ` type RecordSessionBindingParams struct { @@ -241,29 +256,34 @@ type RecordSessionBindingParams struct { // What it DISPLACED comes from SessionBindingForUpdate above, not from a // RETURNING here. The binding update and event writes share the Store tx. -func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error { - _, err := q.db.Exec(ctx, recordSessionBinding, +// xmin is the row version: every upsert writes a new tuple, so a re-bind of the +// same session id still gets a new version a stale release can be fenced by. +func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) { + row := q.db.QueryRow(ctx, recordSessionBinding, arg.AgentAccountID, arg.SessionID, arg.RunnerID, arg.UsageIntervalID, ) - return err + var version string + err := row.Scan(&version) + return version, err } const sessionBinding = `-- name: SessionBinding :one -SELECT agent_account_id, runner_id FROM session_bindings WHERE session_id = $1 +SELECT agent_account_id, runner_id, xmin::text AS version FROM session_bindings WHERE session_id = $1 ` type SessionBindingRow struct { AgentAccountID string RunnerID string + Version string } func (q *Queries) SessionBinding(ctx context.Context, sessionID string) (SessionBindingRow, error) { row := q.db.QueryRow(ctx, sessionBinding, sessionID) var i SessionBindingRow - err := row.Scan(&i.AgentAccountID, &i.RunnerID) + err := row.Scan(&i.AgentAccountID, &i.RunnerID, &i.Version) return i, err } diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql index 8b850a916..b89808028 100644 --- a/go/internal/store/queries/session_bindings.sql +++ b/go/internal/store/queries/session_bindings.sql @@ -114,26 +114,33 @@ ON CONFLICT (tenant_id, interval_id, kind) DO NOTHING; -- What it DISPLACED comes from SessionBindingForUpdate above, not from a -- RETURNING here. The binding update and event writes share the Store tx. --- name: RecordSessionBinding :exec +-- xmin is the row version: every upsert writes a new tuple, so a re-bind of the +-- same session id still gets a new version a stale release can be fenced by. +-- name: RecordSessionBinding :one INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id) VALUES ($1, $2, $3, $4) ON CONFLICT (tenant_id, agent_account_id) DO UPDATE SET session_id = EXCLUDED.session_id, runner_id = EXCLUDED.runner_id, - usage_interval_id = EXCLUDED.usage_interval_id; + usage_interval_id = EXCLUDED.usage_interval_id +RETURNING xmin::text AS version; -- name: SessionBinding :one -SELECT agent_account_id, runner_id FROM session_bindings WHERE session_id = $1; +SELECT agent_account_id, runner_id, xmin::text AS version FROM session_bindings WHERE session_id = $1; -- name: SessionBindingForAccount :one SELECT session_id, runner_id FROM session_bindings WHERE agent_account_id = $1; -- Both deletes also write an estimated start for a binding an older server made -- without one; ON CONFLICT keeps any real start. --- name: DeleteSessionBinding :exec +-- An empty version releases whatever row binds the session; a non-empty one +-- releases only that version, so a stale release cannot remove a re-bind. +-- Returns the rows removed. Postgres runs every data-modifying CTE to completion. +-- name: DeleteSessionBinding :one WITH d AS ( DELETE FROM session_bindings AS b - WHERE b.session_id = $1 + WHERE b.session_id = sqlc.arg(session_id) + AND (sqlc.arg(version)::text = '' OR b.xmin::text = sqlc.arg(version)::text) RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, b.session_id, b.runner_id, b.created_at ), starts AS ( @@ -147,16 +154,19 @@ WITH d AS ( JOIN agent_accounts AS a ON a.account_id = d.agent_account_id ON CONFLICT DO NOTHING RETURNING 1 +), ends AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 ) -INSERT INTO compute_usage_events ( - tenant_id, id, interval_id, kind, occurred_at, agent_account_id, - owner_user_id, session_id, runner_id -) -SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), - d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id - FROM d - JOIN agent_accounts AS a ON a.account_id = d.agent_account_id -ON CONFLICT DO NOTHING; +SELECT count(*) FROM d; -- The reconnect sweep, run by Hub.enroll under the system role because a Runner -- is shared across tenants. :many with RETURNING: each removed row drives a diff --git a/go/internal/store/session_bindings.go b/go/internal/store/session_bindings.go index ea236788f..4f00936aa 100644 --- a/go/internal/store/session_bindings.go +++ b/go/internal/store/session_bindings.go @@ -94,19 +94,19 @@ type SessionBinding struct { // resolves the principal a comms call runs under. Any OTHER unique violation is // unexpected, so it falls through to the generic wrap with its own message // intact rather than being relabelled as a session collision. -func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) (string, error) { +func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) (displaced, version string, err error) { if sessionID == "" { - return "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) + return "", "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) } if accountID == "" { - return "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + return "", "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) } // Unlike agent_placements.runner_id, '' is NOT an accepted unknown sentinel: // a placement outlives its Runner, but a binding exists only while a Runner // is attached, and runner_id is the sweep key that retires it. A '' binding // no real Runner's re-enroll could sweep would linger as a stale session. if runnerID == "" { - return "", fmt.Errorf("%w: runner id is required", ErrInvalidArgument) + return "", "", fmt.Errorf("%w: runner id is required", ErrInvalidArgument) } // beginTenantTx, not s.pool.Begin: it arms SET LOCAL ROLE + the @@ -115,7 +115,7 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco // would run as the owner with no GUC and silently disable tenant isolation. tx, err := s.beginTenantTx(ctx) if err != nil { - return "", fmt.Errorf("store: begin record session binding: %w", err) + return "", "", fmt.Errorf("store: begin record session binding: %w", err) } // No-op after a successful commit; the rollback that matters is on every // error path below, where there is nothing further to report about it. @@ -129,23 +129,23 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco Column1: pgtype.Text{String: string(s.resolveTenant(ctx)), Valid: true}, Column2: pgtype.Text{String: string(accountID), Valid: true}, }); err != nil { - return "", fmt.Errorf("store: lock session binding account: %w", err) + return "", "", fmt.Errorf("store: lock session binding account: %w", err) } prior, err := qtx.SessionBindingForUpdate(ctx, string(accountID)) if err != nil && !noRows(err) { - return "", fmt.Errorf("store: read prior session binding: %w", err) + return "", "", fmt.Errorf("store: read prior session binding: %w", err) } if noRows(err) { prior = db.SessionBindingForUpdateRow{} } - displaced := prior.SessionID + displaced = prior.SessionID intervalID := prior.UsageIntervalID if intervalID != "" { // An older server may have written the prior row without a start event. if err := qtx.EnsureComputeUsageIntervalStart(ctx, string(accountID)); err != nil { - return "", fmt.Errorf("store: ensure compute usage interval start: %w", err) + return "", "", fmt.Errorf("store: ensure compute usage interval start: %w", err) } } if prior.SessionID != sessionID { @@ -156,7 +156,7 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco SessionID: prior.SessionID, RunnerID: prior.RunnerID, }); err != nil { - return "", fmt.Errorf("store: end displaced compute usage interval: %w", err) + return "", "", fmt.Errorf("store: end displaced compute usage interval: %w", err) } } intervalID = uuid.NewString() @@ -166,28 +166,29 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco SessionID: sessionID, RunnerID: runnerID, }); err != nil { - return "", fmt.Errorf("store: start compute usage interval: %w", err) + return "", "", fmt.Errorf("store: start compute usage interval: %w", err) } } - if err := qtx.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ + version, err = qtx.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ SessionID: sessionID, AgentAccountID: string(accountID), RunnerID: runnerID, UsageIntervalID: intervalID, - }); err != nil { + }) + if err != nil { if pgErrIs(err, pgForeignKeyViolation) { - return "", fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) + return "", "", fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) } if pgErrIs(err, pgUniqueViolation) && pgConstraintName(err) == "session_bindings_session_key" { - return "", fmt.Errorf("%w: session %q is already bound to a different agent", ErrConflict, sessionID) + return "", "", fmt.Errorf("%w: session %q is already bound to a different agent", ErrConflict, sessionID) } - return "", fmt.Errorf("store: record session binding: %w", err) + return "", "", fmt.Errorf("store: record session binding: %w", err) } if err := tx.Commit(ctx); err != nil { - return "", fmt.Errorf("store: commit record session binding: %w", err) + return "", "", fmt.Errorf("store: commit record session binding: %w", err) } - return displaced, nil + return displaced, version, nil } // ResolveSessionBinding resolves the account and Runner a live session speaks @@ -197,18 +198,18 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco // resolves the scope a call runs under, so a miss must never surface as an // empty AccountID with a nil error. A zero-value account id would flow onward as // a real (wrong) principal instead of stopping the call. -func (s *Store) ResolveSessionBinding(ctx context.Context, sessionID string) (AccountID, string, error) { +func (s *Store) ResolveSessionBinding(ctx context.Context, sessionID string) (account AccountID, runnerID, version string, err error) { if sessionID == "" { - return "", "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) + return "", "", "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) } row, err := s.q.SessionBinding(ctx, sessionID) if err != nil { if noRows(err) { - return "", "", fmt.Errorf("%w: session %q is not bound", ErrNotFound, sessionID) + return "", "", "", fmt.Errorf("%w: session %q is not bound", ErrNotFound, sessionID) } - return "", "", fmt.Errorf("store: resolve session binding: %w", err) + return "", "", "", fmt.Errorf("store: resolve session binding: %w", err) } - return AccountID(row.AgentAccountID), row.RunnerID, nil + return AccountID(row.AgentAccountID), row.RunnerID, row.Version, nil } // SessionForAccount resolves the live session bound to an agent account — the @@ -253,14 +254,19 @@ func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (str // call on a session RecordSessionBinding has already displaced: the row now // names the newer session, so the stale release matches nothing and leaves the // live binding alone. -func (s *Store) DeleteSessionBinding(ctx context.Context, sessionID string) error { +// +// A non-empty version releases only that row version, as returned by Record or +// Resolve: a stale release cannot remove a re-bind of the same session id. +// removed reports whether a row was deleted. +func (s *Store) DeleteSessionBinding(ctx context.Context, sessionID, version string) (removed bool, err error) { if sessionID == "" { - return fmt.Errorf("%w: session id is required", ErrInvalidArgument) + return false, fmt.Errorf("%w: session id is required", ErrInvalidArgument) } - if err := s.q.DeleteSessionBinding(ctx, sessionID); err != nil { - return fmt.Errorf("store: delete session binding: %w", err) + n, err := s.q.DeleteSessionBinding(ctx, db.DeleteSessionBindingParams{SessionID: sessionID, Version: version}) + if err != nil { + return false, fmt.Errorf("store: delete session binding: %w", err) } - return nil + return n > 0, nil } // DeleteSessionBindingsForRunner is the enroll sweep: it releases every binding diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index 65613e9b6..e91b02a82 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -35,7 +35,7 @@ import ( // ctx is the SECOND parameter, after t, matching mustTopic (messages_test.go). func mustBind(t *testing.T, ctx context.Context, s *Store, sessionID string, accountID AccountID, runnerID string) string { t.Helper() - displaced, err := s.RecordSessionBinding(ctx, sessionID, accountID, runnerID) + displaced, _, err := s.RecordSessionBinding(ctx, sessionID, accountID, runnerID) if err != nil { t.Fatalf("RecordSessionBinding(%q, %q, %q): %v", sessionID, accountID, runnerID, err) } @@ -91,7 +91,7 @@ func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - displaced, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") + displaced, _, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") if err != nil { t.Fatalf("RecordSessionBinding: %v", err) } @@ -99,7 +99,7 @@ func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { t.Fatalf("first bind displaced %q, want \"\" — the account held no prior session, and a caller reaping a phantom id would clear a live registry entry", displaced) } - gotAccount, gotRunner, err := s.ResolveSessionBinding(ctx, "sess-1") + gotAccount, gotRunner, _, err := s.ResolveSessionBinding(ctx, "sess-1") if err != nil { t.Fatalf("ResolveSessionBinding: %v", err) } @@ -128,7 +128,7 @@ func TestSessionBindingLookupsFailClosed(t *testing.T) { ctx := context.Background() s := newTestStore(t) - account, runner, err := s.ResolveSessionBinding(ctx, "never-bound") + account, runner, _, err := s.ResolveSessionBinding(ctx, "never-bound") if !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionBinding(never-bound) err = %v, want errors.Is(_, ErrNotFound)", err) } @@ -167,7 +167,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { mustBind(t, ctx, s, "sess-old", agent.ID, "runner-1") - displaced, err := s.RecordSessionBinding(ctx, "sess-new", agent.ID, "runner-2") + displaced, _, err := s.RecordSessionBinding(ctx, "sess-new", agent.ID, "runner-2") if err != nil { t.Fatalf("re-point onto sess-new: %v (a re-point must never be refused — promoteSession has nowhere to put an error)", err) } @@ -192,7 +192,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { // The displaced session id no longer resolves: it named the same row, which // now carries the new session. - if _, _, err := s.ResolveSessionBinding(ctx, "sess-old"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-old"); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionBinding(sess-old) err = %v, want ErrNotFound — the displaced session must stop resolving", err) } @@ -260,7 +260,7 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T mustBind(t, ctx, s, "sess-1", agentA.ID, "runner-1") - displaced, err := s.RecordSessionBinding(ctx, "sess-1", agentB.ID, "runner-1") + displaced, _, err := s.RecordSessionBinding(ctx, "sess-1", agentB.ID, "runner-1") sentinelIs(t, err, ErrConflict, "a session id already bound to a different agent") if displaced != "" { t.Fatalf("the refused bind returned displaced = %q, want \"\" — a failed write displaced nothing", displaced) @@ -268,7 +268,7 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T // The refused write changed nothing: the session still speaks for agent A, // and agent B still has no live session. - gotAccount, _, err := s.ResolveSessionBinding(ctx, "sess-1") + gotAccount, _, _, err := s.ResolveSessionBinding(ctx, "sess-1") if err != nil { t.Fatalf("ResolveSessionAccount after the refused bind: %v", err) } @@ -287,7 +287,7 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T func TestRecordSessionBindingUnknownAgentIsInvalidArgument(t *testing.T) { s := newTestStore(t) - _, err := s.RecordSessionBinding(context.Background(), "sess-1", "no-such-agent", "runner-1") + _, _, err := s.RecordSessionBinding(context.Background(), "sess-1", "no-such-agent", "runner-1") sentinelIs(t, err, ErrInvalidArgument, "binding for an unknown agent") } @@ -303,11 +303,11 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { agent := mustAgent(t, s, owner.ID, "agent") mustBind(t, ctx, s, "sess-1", agent.ID, "runner-1") - if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { + if _, err := s.DeleteSessionBinding(ctx, "sess-1", ""); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } - if _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionAccount after delete err = %v, want ErrNotFound", err) } if _, _, err := s.SessionForAccount(ctx, agent.ID); !errors.Is(err, ErrNotFound) { @@ -322,14 +322,53 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { // A second delete of an already-released session is a no-op, and so is // deleting one that never existed. - if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { + if _, err := s.DeleteSessionBinding(ctx, "sess-1", ""); err != nil { t.Fatalf("second DeleteSessionBinding(sess-1) = %v, want nil (idempotent)", err) } - if err := s.DeleteSessionBinding(ctx, "never-bound"); err != nil { + if _, err := s.DeleteSessionBinding(ctx, "never-bound", ""); err != nil { t.Fatalf("DeleteSessionBinding(never-bound) = %v, want nil (idempotent)", err) } } +// A release conditioned on a version removes only that row version: a re-bind of +// the same session id writes a new version, so the stale release leaves it bound. +func TestDeleteSessionBindingByVersionSkipsRebind(t *testing.T) { + ctx := context.Background() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + _, oldVersion, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") + if err != nil { + t.Fatalf("RecordSessionBinding: %v", err) + } + if _, _, read, err := s.ResolveSessionBinding(ctx, "sess-1"); err != nil || read != oldVersion { + t.Fatalf("ResolveSessionBinding version = (%q, %v), want (%q, nil)", read, err, oldVersion) + } + _, newVersion, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") + if err != nil { + t.Fatalf("re-bind RecordSessionBinding: %v", err) + } + if newVersion == oldVersion { + t.Fatalf("re-bind kept version %q, want a new one", oldVersion) + } + + removed, err := s.DeleteSessionBinding(ctx, "sess-1", oldVersion) + if err != nil || removed { + t.Fatalf("stale DeleteSessionBinding = (%v, %v), want (false, nil)", removed, err) + } + if got, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); err != nil || got != agent.ID { + t.Fatalf("binding after stale delete = (%q, %v), want (%q, nil)", got, err, agent.ID) + } + removed, err = s.DeleteSessionBinding(ctx, "sess-1", newVersion) + if err != nil || !removed { + t.Fatalf("current DeleteSessionBinding = (%v, %v), want (true, nil)", removed, err) + } + if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { + t.Fatalf("ResolveSessionBinding after delete err = %v, want ErrNotFound", err) + } +} + // TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding is the reconnect // sweep itself, and the test that would catch a `:exec` regression. The returned // rows are NOT diagnostic: each drives a presence DISCONNECTED edge for its @@ -390,14 +429,14 @@ func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { // The swept bindings are actually gone in both directions. for _, sessionID := range []string{"sess-a", "sess-b", "sess-z"} { - if _, _, err := s.ResolveSessionBinding(ctx, sessionID); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, sessionID); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionBinding(%s) after the sweep err = %v, want ErrNotFound", sessionID, err) } } // The other Runner's binding SURVIVES: a re-enroll must not drop live // sessions on a Runner that never reconnected. - gotAccount, _, err := s.ResolveSessionBinding(ctx, "sess-elsewhere") + gotAccount, _, _, err := s.ResolveSessionBinding(ctx, "sess-elsewhere") if err != nil { t.Fatalf("runner-2's binding did not survive the runner-1 sweep: %v", err) } @@ -470,7 +509,7 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { // Tenant B resolves A's session id: the row is not in B's view, so this must // fail closed rather than hand B tenant A's account. - gotAccount, _, err := s.ResolveSessionBinding(ctxB, "sess-a") + gotAccount, _, _, err := s.ResolveSessionBinding(ctxB, "sess-a") if !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B ResolveSessionBinding(sess-a) err = %v, want ErrNotFound — cross-tenant read leak", err) } @@ -495,7 +534,7 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { // Control: tenant A still sees its OWN binding, proving the policy is not // simply hiding everything. - if gotAccount, _, err := s.ResolveSessionBinding(ctxA, "sess-a"); err != nil { + if gotAccount, _, _, err := s.ResolveSessionBinding(ctxA, "sess-a"); err != nil { t.Fatalf("tenant A cannot see its OWN binding — policy over-blocks: %v", err) } else if gotAccount != agentA.ID { t.Fatalf("tenant A's own binding resolves to %q, want %q", gotAccount, agentA.ID) @@ -554,14 +593,14 @@ func TestSessionBindingSameSessionIDInTwoTenantsCoexist(t *testing.T) { // Each tenant resolves its OWN account from the shared session id. A single // surviving row would make one of these two answer with the other tenant's // account — the relay resolving a foreign principal. - gotA, _, err := s.ResolveSessionBinding(ctxA, "sess-shared") + gotA, _, _, err := s.ResolveSessionBinding(ctxA, "sess-shared") if err != nil { t.Fatalf("tenant A ResolveSessionBinding(sess-shared) after B's bind: %v (B's write clobbered A's binding)", err) } if gotA != agentA.ID { t.Fatalf("tenant A resolves sess-shared to %q, want its own agent %q", gotA, agentA.ID) } - gotB, _, err := s.ResolveSessionBinding(ctxB, "sess-shared") + gotB, _, _, err := s.ResolveSessionBinding(ctxB, "sess-shared") if err != nil { t.Fatalf("tenant B ResolveSessionBinding(sess-shared): %v", err) } @@ -625,10 +664,10 @@ func TestSessionBindingSameAccountIDInTwoTenantsCoexist(t *testing.T) { } // Each tenant's session resolves only in its own tenant. - if _, _, err := s.ResolveSessionBinding(ctxA, "sess-b"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxA, "sess-b"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant A resolved B's session err = %v, want ErrNotFound", err) } - if _, _, err := s.ResolveSessionBinding(ctxB, "sess-a"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxB, "sess-a"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B resolved A's session err = %v, want ErrNotFound", err) } } @@ -700,7 +739,7 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { // pooled connection it draws, so this asserts the DISJUNCTION rather than // pretending either branch is the contract. Both are defects; PR3 owns the // fix, and this test is what will fail when it lands one. - sysDisplaced, sysErr := s.RecordSessionBinding(WithSystemRole(context.Background()), "sess-sys", shared.ID, "runner-sys") + sysDisplaced, _, sysErr := s.RecordSessionBinding(WithSystemRole(context.Background()), "sess-sys", shared.ID, "runner-sys") switch { case sysErr != nil: t.Logf("system-role RecordSessionBinding failed (connection had no prior armed statement, so the GUC was undefined and the DEFAULT was NULL): %v", sysErr) @@ -716,10 +755,10 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { t.Fatalf("system-role write left %d rows stamped tenant_id = '', want 1 — the observed failure mode is an ORPHAN row, so if it is now stamped with a real tenant the write path grew scoping and this test must be updated deliberately", orphans) } // And it is invisible to every tenant: no policy matches ''. - if _, _, err := s.ResolveSessionBinding(ctxA, "sess-sys"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxA, "sess-sys"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant A sees the untenanted binding err = %v, want ErrNotFound", err) } - if _, _, err := s.ResolveSessionBinding(ctxB, "sess-sys"); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctxB, "sess-sys"); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B sees the untenanted binding err = %v, want ErrNotFound", err) } } @@ -786,7 +825,7 @@ func TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced(t *testin results := make(chan outcome, 2) launch := func(sessionID, runnerID string) { go func() { - d, err := s.RecordSessionBinding(ctx, sessionID, agent.ID, runnerID) + d, _, err := s.RecordSessionBinding(ctx, sessionID, agent.ID, runnerID) results <- outcome{sessionID, d, err} }() } @@ -848,7 +887,7 @@ func TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced(t *testin // Both displaced sessions are gone, and between them they were reported to // exactly the two callers: nothing was destroyed unreported. for _, dead := range []string{"sess-A", winner} { - if _, _, err := s.ResolveSessionBinding(ctx, dead); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, dead); !errors.Is(err, ErrNotFound) { t.Fatalf("displaced session %q still resolves (err = %v), want ErrNotFound", dead, err) } } @@ -914,7 +953,7 @@ func TestRecordSessionBindingConcurrentFirstBindsReportTheDestroyedSession(t *te {"sess-2", "runner-2"}, } { go func() { - d, err := s.RecordSessionBinding(ctx, b.session, agent.ID, b.runner) + d, _, err := s.RecordSessionBinding(ctx, b.session, agent.ID, b.runner) results <- outcome{b.session, d, err} }() } @@ -957,7 +996,7 @@ func TestRecordSessionBindingConcurrentFirstBindsReportTheDestroyedSession(t *te t.Fatalf("bindings after two concurrent first binds = %d, want 1 — the PK must fold them into one row", n) } // The winner's session is gone, and it WAS reported, so it can be reaped. - if _, _, err := s.ResolveSessionBinding(ctx, winner); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, winner); !errors.Is(err, ErrNotFound) { t.Fatalf("the overwritten session %q still resolves (err = %v), want ErrNotFound", winner, err) } } @@ -1046,7 +1085,7 @@ func TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep(t *t go func() { ready.Done() <-start - displaced, err := s.RecordSessionBinding(ctx, fresh, agent.ID, "runner-2") + displaced, _, err := s.RecordSessionBinding(ctx, fresh, agent.ID, "runner-2") binds <- bindResult{displaced, err} }() go func() { @@ -1107,7 +1146,7 @@ func TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep(t *t } // And the displaced session is gone in every ordering, which is what // makes "reported to nobody" a genuine strand rather than a deferral. - if _, _, err := s.ResolveSessionBinding(ctx, old); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.ResolveSessionBinding(ctx, old); !errors.Is(err, ErrNotFound) { t.Fatalf("iteration %d: displaced session %q still resolves (err = %v), want ErrNotFound", i, old, err) } } diff --git a/go/server/lifecycle_wake_pgtest_test.go b/go/server/lifecycle_wake_pgtest_test.go index fe8511882..3dc3c283e 100644 --- a/go/server/lifecycle_wake_pgtest_test.go +++ b/go/server/lifecycle_wake_pgtest_test.go @@ -134,7 +134,7 @@ func TestWakeAgentStaleBindingRowStillResumes(t *testing.T) { if err := f.store.AppendTranscriptEntry(ctx, logical, 1, true, `{"header":true}`, "k1"); err != nil { t.Fatalf("append checkpoint: %v", err) } - if _, err := f.store.RecordSessionBinding(ctx, "sess-dead", f.agentID, fakeRunnerID); err != nil { + if _, _, err := f.store.RecordSessionBinding(ctx, "sess-dead", f.agentID, fakeRunnerID); err != nil { t.Fatalf("RecordSessionBinding: %v", err) } // Checked on the store, not the hub: a hub read-through would warm the cache. From 6a71170292784a415f702280645a26001c1f11e4 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 19:52:10 -0400 Subject: [PATCH 12/17] fix(store): version session bindings with a UUID column, not xmin (RIG-4742) xmin is a 32-bit transaction id and repeats after wraparound, so an old cached version could match a later re-bind. binding_version is a fresh UUID per upsert. The conditional release is a separate DeleteSessionBindingVersion, and DeleteSessionBinding is back to its unconditional form for Stop. A cached binding whose durable write failed has an empty version, which matches no row, so its cleanup cannot delete a peer's row. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/binding_cache_test.go | 12 ++- go/internal/runnerhub/hub.go | 10 ++- go/internal/runnerhub/lost_session_test.go | 39 ++++++++- go/internal/runnerhub/relay_comms.go | 12 ++- .../store/compute_usage_pgtest_test.go | 8 +- go/internal/store/db/models.go | 1 + go/internal/store/db/querier.go | 13 ++- go/internal/store/db/session_bindings.sql.go | 85 +++++++++++++------ .../0003_session_binding_version.sql | 5 ++ .../store/queries/session_bindings.sql | 53 +++++++++--- go/internal/store/session_bindings.go | 32 ++++--- .../store/session_bindings_pgtest_test.go | 12 +-- 12 files changed, 200 insertions(+), 82 deletions(-) create mode 100644 go/internal/store/migrations/0003_session_binding_version.sql diff --git a/go/internal/runnerhub/binding_cache_test.go b/go/internal/runnerhub/binding_cache_test.go index 057bf9481..93827c9aa 100644 --- a/go/internal/runnerhub/binding_cache_test.go +++ b/go/internal/runnerhub/binding_cache_test.go @@ -127,10 +127,18 @@ func (f *fakeBindingStore) SessionForAccount(_ context.Context, accountID store. return "", "", store.ErrNotFound } -func (f *fakeBindingStore) DeleteSessionBinding(_ context.Context, sessionID, version string) (bool, error) { +func (f *fakeBindingStore) DeleteSessionBinding(_ context.Context, sessionID string) error { f.mu.Lock() defer f.mu.Unlock() - if _, ok := f.bindings[sessionID]; !ok || (version != "" && f.versions[sessionID] != version) { + delete(f.bindings, sessionID) + delete(f.versions, sessionID) + return nil +} + +func (f *fakeBindingStore) DeleteSessionBindingVersion(_ context.Context, sessionID, version string) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + if _, ok := f.bindings[sessionID]; !ok || f.versions[sessionID] != version { return false, nil } delete(f.bindings, sessionID) diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index d68c5ccbb..e027513d9 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -226,8 +226,10 @@ type SessionBindingStore interface { SessionForAccount(ctx context.Context, accountID store.AccountID) (sessionID, runnerID string, err error) // DeleteSessionBinding releases one session's binding — the unbind write. // Idempotent: releasing an already-released session is a no-op success. - // A non-empty version releases only that row; removed reports a deletion. - DeleteSessionBinding(ctx context.Context, sessionID, version string) (removed bool, err error) + DeleteSessionBinding(ctx context.Context, sessionID string) error + // DeleteSessionBindingVersion releases sessionID only while its row is still + // the write that returned version; removed reports whether it was. + DeleteSessionBindingVersion(ctx context.Context, sessionID, version string) (removed bool, err error) // DeleteSessionBindingsForRunner is the enroll sweep: it releases every // binding attached to runnerID and RETURNS the rows it removed, driving the // enroll reap (offline edges + held-deliver reap) from durable truth @@ -511,8 +513,8 @@ type attachedRunner struct { // sessionBinding is one live session's principal and the Runner that owns it. // version and lifetime name one binding of the session id, so a release from an -// older lifetime can skip a re-bind: version is the durable row ("" when none was -// written), lifetime is the cache entry, unique per insert and never zero. +// older lifetime can skip a re-bind: version is the durable write ("" when none +// was made), lifetime is the cache entry, unique per insert and never zero. type sessionBinding struct { account store.AccountID runnerID string diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 23e5b5680..f2c433a0a 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -356,7 +356,40 @@ func TestErroredCleanupKeepsPeerRebind(t *testing.T) { } } -// pausingDeleteBindingStore holds the first DeleteSessionBinding until released. +func TestErroredCleanupWithoutDurableWriteKeepsPeerRow(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := newFakeBindingStore() + bindings.recordErr = errors.New("store down") + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + // The durable write fails, so this hub caches a binding with no version. + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + if old.version != "" { + t.Fatalf("cached version after failed write = %q, want empty", old.version) + } + // A peer Server then writes the session's row. + bindings.mu.Lock() + bindings.recordErr = nil + bindings.mu.Unlock() + if _, _, err := bindings.RecordSessionBinding(ctx, "sess-1", testAgentAccount, "runner-1"); err != nil { + t.Fatalf("peer RecordSessionBinding: %v", err) + } + + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + if _, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("peer's durable binding after old cleanup: %v, want it kept", err) + } +} + +// pausingDeleteBindingStore holds the first DeleteSessionBindingVersion until released. type pausingDeleteBindingStore struct { *fakeBindingStore once sync.Once @@ -364,12 +397,12 @@ type pausingDeleteBindingStore struct { release chan struct{} } -func (b *pausingDeleteBindingStore) DeleteSessionBinding(ctx context.Context, sessionID, version string) (bool, error) { +func (b *pausingDeleteBindingStore) DeleteSessionBindingVersion(ctx context.Context, sessionID, version string) (bool, error) { b.once.Do(func() { close(b.entered) <-b.release }) - return b.fakeBindingStore.DeleteSessionBinding(ctx, sessionID, version) + return b.fakeBindingStore.DeleteSessionBindingVersion(ctx, sessionID, version) } func TestLostSessionCleanupFromOldEnrollmentKeepsRebinding(t *testing.T) { diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index 882ce1782..dacafba2b 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -204,11 +204,15 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio tenant := "" if bindings != nil { - version := "" + // A limited release deletes only its own write, matching even "": a + // version-less cache entry must not remove a row someone else wrote. + removed := true + var err error if only != nil { - version = only.version + removed, err = bindings.DeleteSessionBindingVersion(ctx, sessionID, only.version) + } else { + err = bindings.DeleteSessionBinding(ctx, sessionID) } - removed, err := bindings.DeleteSessionBinding(ctx, sessionID, version) switch { case err != nil: // A durable-delete fault must not fail the Stop that already @@ -216,7 +220,7 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio // next re-enroll sweep retires any surviving row. h.log.Error("delete session binding failed; evicting cache anyway", "session_id", sessionID, "error", err) - case only != nil && version != "" && !removed: + case !removed && only.version != "": // The row was re-bound since only was read: it is not ours to release. return false default: diff --git a/go/internal/store/compute_usage_pgtest_test.go b/go/internal/store/compute_usage_pgtest_test.go index eae70168a..78f4e0585 100644 --- a/go/internal/store/compute_usage_pgtest_test.go +++ b/go/internal/store/compute_usage_pgtest_test.go @@ -71,7 +71,7 @@ func TestComputeUsageBindingLifecycle(t *testing.T) { t.Fatalf("start event = %+v, want the tenant, owner, agent, session, runner, and non-estimated start", start) } - if _, err := s.DeleteSessionBinding(ctx, "compute-session", ""); err != nil { + if err := s.DeleteSessionBinding(ctx, "compute-session"); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } events = computeEvents(t, s, tenant, agent.ID) @@ -88,7 +88,7 @@ func TestComputeUsageBindingLifecycle(t *testing.T) { end.OwnerUserID != string(owner.ID) || end.SessionID != start.SessionID || end.RunnerID != start.RunnerID || end.Estimated { t.Fatalf("end event = %+v, want same interval metadata and a non-estimated end", end) } - if _, err := s.DeleteSessionBinding(ctx, "compute-session", ""); err != nil { + if err := s.DeleteSessionBinding(ctx, "compute-session"); err != nil { t.Fatalf("second DeleteSessionBinding: %v", err) } if got := computeEvents(t, s, tenant, agent.ID); len(got) != 2 { @@ -259,7 +259,7 @@ func TestComputeUsageSameSessionRunnerRebindKeepsInterval(t *testing.T) { func TestComputeUsageLegacyBindingGetsEstimatedStart(t *testing.T) { for name, release := range map[string]func(*testing.T, context.Context, *Store){ "single release": func(t *testing.T, ctx context.Context, s *Store) { - if _, err := s.DeleteSessionBinding(ctx, "legacy-session", ""); err != nil { + if err := s.DeleteSessionBinding(ctx, "legacy-session"); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } }, @@ -300,7 +300,7 @@ func TestComputeUsageLegacyBindingReleaseLogsInterval(t *testing.T) { tenant := s.EffectiveTenant(ctx) execAsSystem(t, s, "INSERT INTO session_bindings (tenant_id, agent_account_id, session_id, runner_id) VALUES ($1, $2, 'legacy-only', 'runner-old')", string(tenant), string(agent.ID)) - if _, err := s.DeleteSessionBinding(ctx, "legacy-only", ""); err != nil { + if err := s.DeleteSessionBinding(ctx, "legacy-only"); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } events := computeEvents(t, s, tenant, agent.ID) diff --git a/go/internal/store/db/models.go b/go/internal/store/db/models.go index 45cf884d5..b3abbe31a 100644 --- a/go/internal/store/db/models.go +++ b/go/internal/store/db/models.go @@ -342,6 +342,7 @@ type SessionBinding struct { CreatedAt pgtype.Timestamptz UpdatedAt pgtype.Timestamptz UsageIntervalID string + BindingVersion string } type SystemAccount struct { diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 6b5f69f27..c336b53ac 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -127,10 +127,11 @@ type Querier interface { DeleteServerSecret(ctx context.Context, name string) (int64, error) // Both deletes also write an estimated start for a binding an older server made // without one; ON CONFLICT keeps any real start. - // An empty version releases whatever row binds the session; a non-empty one - // releases only that version, so a stale release cannot remove a re-bind. - // Returns the rows removed. Postgres runs every data-modifying CTE to completion. - DeleteSessionBinding(ctx context.Context, arg DeleteSessionBindingParams) (int64, error) + DeleteSessionBinding(ctx context.Context, sessionID string) error + // DeleteSessionBinding limited to one write of the row: a re-bind since that + // write set a new binding_version, so it is left alone. Returns rows removed; + // Postgres runs every data-modifying CTE to completion. + DeleteSessionBindingVersion(ctx context.Context, arg DeleteSessionBindingVersionParams) (int64, error) // The reconnect sweep, run by Hub.enroll under the system role because a Runner // is shared across tenants. :many with RETURNING: each removed row drives a // presence DISCONNECTED edge, a held-deliver reap, and a tenant-scoped archive. @@ -495,9 +496,7 @@ type Querier interface { RecordOwedMention(ctx context.Context, arg RecordOwedMentionParams) error // What it DISPLACED comes from SessionBindingForUpdate above, not from a // RETURNING here. The binding update and event writes share the Store tx. - // xmin is the row version: every upsert writes a new tuple, so a re-bind of the - // same session id still gets a new version a stale release can be fenced by. - RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) + RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error // Forge state-transition memo queries (compass-forge-state-transition §Actor // attribution). The write chokepoint upserts one memo per forge coordinate // AFTER a successful agent-driven transition; the notify lane consumes it on diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go index 8a449b8ff..c57770338 100644 --- a/go/internal/store/db/session_bindings.sql.go +++ b/go/internal/store/db/session_bindings.sql.go @@ -11,11 +11,46 @@ import ( "github.com/jackc/pgx/v5/pgtype" ) -const deleteSessionBinding = `-- name: DeleteSessionBinding :one +const deleteSessionBinding = `-- name: DeleteSessionBinding :exec WITH d AS ( DELETE FROM session_bindings AS b WHERE b.session_id = $1 - AND ($2::text = '' OR b.xmin::text = $2::text) + RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, + b.session_id, b.runner_id, b.created_at +), starts AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id, estimated + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'start', d.created_at, + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id, TRUE + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 +) +INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id +) +SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id +ON CONFLICT DO NOTHING +` + +// Both deletes also write an estimated start for a binding an older server made +// without one; ON CONFLICT keeps any real start. +func (q *Queries) DeleteSessionBinding(ctx context.Context, sessionID string) error { + _, err := q.db.Exec(ctx, deleteSessionBinding, sessionID) + return err +} + +const deleteSessionBindingVersion = `-- name: DeleteSessionBindingVersion :one +WITH d AS ( + DELETE FROM session_bindings AS b + WHERE b.session_id = $1 AND b.binding_version = $2 RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, b.session_id, b.runner_id, b.created_at ), starts AS ( @@ -44,18 +79,16 @@ WITH d AS ( SELECT count(*) FROM d ` -type DeleteSessionBindingParams struct { - SessionID string - Version string +type DeleteSessionBindingVersionParams struct { + SessionID string + BindingVersion string } -// Both deletes also write an estimated start for a binding an older server made -// without one; ON CONFLICT keeps any real start. -// An empty version releases whatever row binds the session; a non-empty one -// releases only that version, so a stale release cannot remove a re-bind. -// Returns the rows removed. Postgres runs every data-modifying CTE to completion. -func (q *Queries) DeleteSessionBinding(ctx context.Context, arg DeleteSessionBindingParams) (int64, error) { - row := q.db.QueryRow(ctx, deleteSessionBinding, arg.SessionID, arg.Version) +// DeleteSessionBinding limited to one write of the row: a re-bind since that +// write set a new binding_version, so it is left alone. Returns rows removed; +// Postgres runs every data-modifying CTE to completion. +func (q *Queries) DeleteSessionBindingVersion(ctx context.Context, arg DeleteSessionBindingVersionParams) (int64, error) { + row := q.db.QueryRow(ctx, deleteSessionBindingVersion, arg.SessionID, arg.BindingVersion) var count int64 err := row.Scan(&count) return count, err @@ -237,14 +270,14 @@ func (q *Queries) LockSessionBindingAccount(ctx context.Context, arg LockSession return err } -const recordSessionBinding = `-- name: RecordSessionBinding :one -INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id) -VALUES ($1, $2, $3, $4) +const recordSessionBinding = `-- name: RecordSessionBinding :exec +INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id, binding_version) +VALUES ($1, $2, $3, $4, $5) ON CONFLICT (tenant_id, agent_account_id) DO UPDATE SET session_id = EXCLUDED.session_id, runner_id = EXCLUDED.runner_id, - usage_interval_id = EXCLUDED.usage_interval_id -RETURNING xmin::text AS version + usage_interval_id = EXCLUDED.usage_interval_id, + binding_version = EXCLUDED.binding_version ` type RecordSessionBindingParams struct { @@ -252,38 +285,36 @@ type RecordSessionBindingParams struct { SessionID string RunnerID string UsageIntervalID string + BindingVersion string } // What it DISPLACED comes from SessionBindingForUpdate above, not from a // RETURNING here. The binding update and event writes share the Store tx. -// xmin is the row version: every upsert writes a new tuple, so a re-bind of the -// same session id still gets a new version a stale release can be fenced by. -func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) { - row := q.db.QueryRow(ctx, recordSessionBinding, +func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error { + _, err := q.db.Exec(ctx, recordSessionBinding, arg.AgentAccountID, arg.SessionID, arg.RunnerID, arg.UsageIntervalID, + arg.BindingVersion, ) - var version string - err := row.Scan(&version) - return version, err + return err } const sessionBinding = `-- name: SessionBinding :one -SELECT agent_account_id, runner_id, xmin::text AS version FROM session_bindings WHERE session_id = $1 +SELECT agent_account_id, runner_id, binding_version FROM session_bindings WHERE session_id = $1 ` type SessionBindingRow struct { AgentAccountID string RunnerID string - Version string + BindingVersion string } func (q *Queries) SessionBinding(ctx context.Context, sessionID string) (SessionBindingRow, error) { row := q.db.QueryRow(ctx, sessionBinding, sessionID) var i SessionBindingRow - err := row.Scan(&i.AgentAccountID, &i.RunnerID, &i.Version) + err := row.Scan(&i.AgentAccountID, &i.RunnerID, &i.BindingVersion) return i, err } diff --git a/go/internal/store/migrations/0003_session_binding_version.sql b/go/internal/store/migrations/0003_session_binding_version.sql new file mode 100644 index 000000000..cc8d7b5f4 --- /dev/null +++ b/go/internal/store/migrations/0003_session_binding_version.sql @@ -0,0 +1,5 @@ +-- binding_version names one write of a session_bindings row. Every upsert sets a +-- fresh UUID, so a release conditioned on it cannot remove a later re-bind of the +-- same session id. Not xmin: transaction ids wrap around and repeat. +ALTER TABLE session_bindings + ADD COLUMN binding_version TEXT NOT NULL DEFAULT ''; diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql index b89808028..d8653b1a1 100644 --- a/go/internal/store/queries/session_bindings.sql +++ b/go/internal/store/queries/session_bindings.sql @@ -114,33 +114,58 @@ ON CONFLICT (tenant_id, interval_id, kind) DO NOTHING; -- What it DISPLACED comes from SessionBindingForUpdate above, not from a -- RETURNING here. The binding update and event writes share the Store tx. --- xmin is the row version: every upsert writes a new tuple, so a re-bind of the --- same session id still gets a new version a stale release can be fenced by. --- name: RecordSessionBinding :one -INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id) -VALUES ($1, $2, $3, $4) +-- name: RecordSessionBinding :exec +INSERT INTO session_bindings (agent_account_id, session_id, runner_id, usage_interval_id, binding_version) +VALUES ($1, $2, $3, $4, $5) ON CONFLICT (tenant_id, agent_account_id) DO UPDATE SET session_id = EXCLUDED.session_id, runner_id = EXCLUDED.runner_id, - usage_interval_id = EXCLUDED.usage_interval_id -RETURNING xmin::text AS version; + usage_interval_id = EXCLUDED.usage_interval_id, + binding_version = EXCLUDED.binding_version; -- name: SessionBinding :one -SELECT agent_account_id, runner_id, xmin::text AS version FROM session_bindings WHERE session_id = $1; +SELECT agent_account_id, runner_id, binding_version FROM session_bindings WHERE session_id = $1; -- name: SessionBindingForAccount :one SELECT session_id, runner_id FROM session_bindings WHERE agent_account_id = $1; -- Both deletes also write an estimated start for a binding an older server made -- without one; ON CONFLICT keeps any real start. --- An empty version releases whatever row binds the session; a non-empty one --- releases only that version, so a stale release cannot remove a re-bind. --- Returns the rows removed. Postgres runs every data-modifying CTE to completion. --- name: DeleteSessionBinding :one +-- name: DeleteSessionBinding :exec WITH d AS ( DELETE FROM session_bindings AS b - WHERE b.session_id = sqlc.arg(session_id) - AND (sqlc.arg(version)::text = '' OR b.xmin::text = sqlc.arg(version)::text) + WHERE b.session_id = $1 + RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, + b.session_id, b.runner_id, b.created_at +), starts AS ( + INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id, estimated + ) + SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'start', d.created_at, + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id, TRUE + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id + ON CONFLICT DO NOTHING + RETURNING 1 +) +INSERT INTO compute_usage_events ( + tenant_id, id, interval_id, kind, occurred_at, agent_account_id, + owner_user_id, session_id, runner_id +) +SELECT d.tenant_id, gen_random_uuid()::text, d.usage_interval_id, 'end', clock_timestamp(), + d.agent_account_id, a.owner_user_id, d.session_id, d.runner_id + FROM d + JOIN agent_accounts AS a ON a.account_id = d.agent_account_id +ON CONFLICT DO NOTHING; + +-- DeleteSessionBinding limited to one write of the row: a re-bind since that +-- write set a new binding_version, so it is left alone. Returns rows removed; +-- Postgres runs every data-modifying CTE to completion. +-- name: DeleteSessionBindingVersion :one +WITH d AS ( + DELETE FROM session_bindings AS b + WHERE b.session_id = $1 AND b.binding_version = $2 RETURNING b.tenant_id, b.usage_interval_id, b.agent_account_id, b.session_id, b.runner_id, b.created_at ), starts AS ( diff --git a/go/internal/store/session_bindings.go b/go/internal/store/session_bindings.go index 4f00936aa..168fc7f3b 100644 --- a/go/internal/store/session_bindings.go +++ b/go/internal/store/session_bindings.go @@ -170,13 +170,14 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco } } - version, err = qtx.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ + version = uuid.NewString() + if err := qtx.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ SessionID: sessionID, AgentAccountID: string(accountID), RunnerID: runnerID, UsageIntervalID: intervalID, - }) - if err != nil { + BindingVersion: version, + }); err != nil { if pgErrIs(err, pgForeignKeyViolation) { return "", "", fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) } @@ -209,7 +210,7 @@ func (s *Store) ResolveSessionBinding(ctx context.Context, sessionID string) (ac } return "", "", "", fmt.Errorf("store: resolve session binding: %w", err) } - return AccountID(row.AgentAccountID), row.RunnerID, row.Version, nil + return AccountID(row.AgentAccountID), row.RunnerID, row.BindingVersion, nil } // SessionForAccount resolves the live session bound to an agent account — the @@ -254,17 +255,26 @@ func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (str // call on a session RecordSessionBinding has already displaced: the row now // names the newer session, so the stale release matches nothing and leaves the // live binding alone. -// -// A non-empty version releases only that row version, as returned by Record or -// Resolve: a stale release cannot remove a re-bind of the same session id. -// removed reports whether a row was deleted. -func (s *Store) DeleteSessionBinding(ctx context.Context, sessionID, version string) (removed bool, err error) { +func (s *Store) DeleteSessionBinding(ctx context.Context, sessionID string) error { + if sessionID == "" { + return fmt.Errorf("%w: session id is required", ErrInvalidArgument) + } + if err := s.q.DeleteSessionBinding(ctx, sessionID); err != nil { + return fmt.Errorf("store: delete session binding: %w", err) + } + return nil +} + +// DeleteSessionBindingVersion is DeleteSessionBinding limited to the write that +// returned version (from Record or Resolve): a re-bind of the same session id +// since then has a new version and stays. removed reports whether a row went. +func (s *Store) DeleteSessionBindingVersion(ctx context.Context, sessionID, version string) (removed bool, err error) { if sessionID == "" { return false, fmt.Errorf("%w: session id is required", ErrInvalidArgument) } - n, err := s.q.DeleteSessionBinding(ctx, db.DeleteSessionBindingParams{SessionID: sessionID, Version: version}) + n, err := s.q.DeleteSessionBindingVersion(ctx, db.DeleteSessionBindingVersionParams{SessionID: sessionID, BindingVersion: version}) if err != nil { - return false, fmt.Errorf("store: delete session binding: %w", err) + return false, fmt.Errorf("store: delete session binding version: %w", err) } return n > 0, nil } diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index e91b02a82..2a97e07ce 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -303,7 +303,7 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { agent := mustAgent(t, s, owner.ID, "agent") mustBind(t, ctx, s, "sess-1", agent.ID, "runner-1") - if _, err := s.DeleteSessionBinding(ctx, "sess-1", ""); err != nil { + if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } @@ -322,10 +322,10 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { // A second delete of an already-released session is a no-op, and so is // deleting one that never existed. - if _, err := s.DeleteSessionBinding(ctx, "sess-1", ""); err != nil { + if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { t.Fatalf("second DeleteSessionBinding(sess-1) = %v, want nil (idempotent)", err) } - if _, err := s.DeleteSessionBinding(ctx, "never-bound", ""); err != nil { + if err := s.DeleteSessionBinding(ctx, "never-bound"); err != nil { t.Fatalf("DeleteSessionBinding(never-bound) = %v, want nil (idempotent)", err) } } @@ -349,18 +349,18 @@ func TestDeleteSessionBindingByVersionSkipsRebind(t *testing.T) { if err != nil { t.Fatalf("re-bind RecordSessionBinding: %v", err) } - if newVersion == oldVersion { + if newVersion == oldVersion || oldVersion == "" { t.Fatalf("re-bind kept version %q, want a new one", oldVersion) } - removed, err := s.DeleteSessionBinding(ctx, "sess-1", oldVersion) + removed, err := s.DeleteSessionBindingVersion(ctx, "sess-1", oldVersion) if err != nil || removed { t.Fatalf("stale DeleteSessionBinding = (%v, %v), want (false, nil)", removed, err) } if got, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); err != nil || got != agent.ID { t.Fatalf("binding after stale delete = (%q, %v), want (%q, nil)", got, err, agent.ID) } - removed, err = s.DeleteSessionBinding(ctx, "sess-1", newVersion) + removed, err = s.DeleteSessionBindingVersion(ctx, "sess-1", newVersion) if err != nil || !removed { t.Fatalf("current DeleteSessionBinding = (%v, %v), want (true, nil)", removed, err) } From ce396aea37de5850004e310e02a3acb6425d99b4 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 19:59:19 -0400 Subject: [PATCH 13/17] fix(runnerhub): a version-less cleanup deletes no durable row (RIG-4742) A cached binding whose durable write failed has an empty version. A legacy row also has the column default '', so a versioned delete with '' could remove it. The limited release now skips the store for an empty version. The removed check no longer reads only, which nilaway could not prove non-nil. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/lost_session_test.go | 25 ++++++++++++++++++++++ go/internal/runnerhub/relay_comms.go | 16 ++++++++------ 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index f2c433a0a..62464df8b 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -389,6 +389,31 @@ func TestErroredCleanupWithoutDurableWriteKeepsPeerRow(t *testing.T) { } } +func TestErroredCleanupWithoutDurableWriteKeepsLegacyRow(t *testing.T) { + ctx := t.Context() + hub, _, _ := newHub() + bindings := newFakeBindingStore() + bindings.recordErr = errors.New("store down") + hub.SetSessionBindingStore(bindings) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + // A row written before binding_version existed carries the column default "". + bindings.mu.Lock() + bindings.bindings["sess-1"] = store.SessionBinding{TenantID: bindings.tenant, SessionID: "sess-1", AccountID: testAgentAccount, RunnerID: "runner-1"} + bindings.versions["sess-1"] = "" + bindings.mu.Unlock() + + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + if _, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("legacy durable binding after version-less cleanup: %v, want it kept", err) + } +} + // pausingDeleteBindingStore holds the first DeleteSessionBindingVersion until released. type pausingDeleteBindingStore struct { *fakeBindingStore diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index dacafba2b..0d58426ef 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -204,14 +204,15 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio tenant := "" if bindings != nil { - // A limited release deletes only its own write, matching even "": a - // version-less cache entry must not remove a row someone else wrote. + // A limited release deletes only its own write. A version-less entry + // wrote no row, so it deletes none: a legacy or peer row is not its own. removed := true var err error - if only != nil { - removed, err = bindings.DeleteSessionBindingVersion(ctx, sessionID, only.version) - } else { + switch { + case only == nil: err = bindings.DeleteSessionBinding(ctx, sessionID) + case only.version != "": + removed, err = bindings.DeleteSessionBindingVersion(ctx, sessionID, only.version) } switch { case err != nil: @@ -220,8 +221,9 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio // next re-enroll sweep retires any surviving row. h.log.Error("delete session binding failed; evicting cache anyway", "session_id", sessionID, "error", err) - case !removed && only.version != "": - // The row was re-bound since only was read: it is not ours to release. + case !removed: + // Only a versioned release reports false: the row was re-bound since + // only was read, so it is not ours to release. return false default: tenant = string(bindings.EffectiveTenant(ctx)) From a0d3c6a1a8957b81ad5254dccd0e88dbf7b71f47 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 20:42:35 -0400 Subject: [PATCH 14/17] fix(runnerhub): a re-bound version-less cleanup is not a loss; backfill versions (RIG-4742) A cached binding with no durable version now checks for a competing row before releasing. If one exists, the session was re-bound elsewhere: the stale cache entry is dropped and no loss is reported. The migration also gives pre-existing rows a real version, so no row keeps the '' default. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/lost_session_test.go | 4 ++++ go/internal/runnerhub/relay_comms.go | 21 +++++++++++++++++-- .../0003_session_binding_version.sql | 4 ++++ 3 files changed, 27 insertions(+), 2 deletions(-) diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 62464df8b..5d036da5b 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -387,6 +387,10 @@ func TestErroredCleanupWithoutDurableWriteKeepsPeerRow(t *testing.T) { if _, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); err != nil { t.Fatalf("peer's durable binding after old cleanup: %v, want it kept", err) } + lost.none(t, "a session the peer re-bound is not lost") + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != testAgentAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want the peer's (%s, true)", account, ok, testAgentAccount) + } } func TestErroredCleanupWithoutDurableWriteKeepsLegacyRow(t *testing.T) { diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index 0d58426ef..39bbf3392 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -213,6 +213,10 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio err = bindings.DeleteSessionBinding(ctx, sessionID) case only.version != "": removed, err = bindings.DeleteSessionBindingVersion(ctx, sessionID, only.version) + default: + // Any row is someone else's write, so the session is still bound there. + _, _, _, rerr := bindings.ResolveSessionBinding(ctx, sessionID) + removed = errors.Is(rerr, store.ErrNotFound) } switch { case err != nil: @@ -222,8 +226,8 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio h.log.Error("delete session binding failed; evicting cache anyway", "session_id", sessionID, "error", err) case !removed: - // Only a versioned release reports false: the row was re-bound since - // only was read, so it is not ours to release. + // The row was re-bound since only was read: the session is not lost. + h.evictStaleBinding(sessionID, only) return false default: tenant = string(bindings.EffectiveTenant(ctx)) @@ -275,6 +279,19 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio return true } +// evictStaleBinding drops sessionID's cache entry while it is still stale, so the +// next read resolves the row that replaced it. +func (h *Hub) evictStaleBinding(sessionID string, stale *sessionBinding) { + if stale == nil { + return + } + h.mu.Lock() + defer h.mu.Unlock() + if live, ok := h.sessionAccounts[sessionID]; ok && sameBinding(live, *stale) { + delete(h.sessionAccounts, sessionID) + } +} + // newBindingLocked builds a sessionAccounts entry with a fresh lifetime. Caller holds mu. func (h *Hub) newBindingLocked(account store.AccountID, runnerID, version string) sessionBinding { h.lastLifetime++ diff --git a/go/internal/store/migrations/0003_session_binding_version.sql b/go/internal/store/migrations/0003_session_binding_version.sql index cc8d7b5f4..f9574b32a 100644 --- a/go/internal/store/migrations/0003_session_binding_version.sql +++ b/go/internal/store/migrations/0003_session_binding_version.sql @@ -3,3 +3,7 @@ -- same session id. Not xmin: transaction ids wrap around and repeat. ALTER TABLE session_bindings ADD COLUMN binding_version TEXT NOT NULL DEFAULT ''; + +-- Give rows written before this column a real version too. +UPDATE session_bindings SET binding_version = gen_random_uuid()::TEXT +WHERE binding_version = ''; From c1e204a03564eebf4daf2f02c7cda6967b75dce8 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 21:01:34 -0400 Subject: [PATCH 15/17] fix(runnerhub): drop the stale account entry when a peer took the session (RIG-4742) A cleanup that found its row re-bound evicted only the session's cache entry. The reverse account entry still named the session, so delivery for the old account could reach the new owner's session. Evict both. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/lost_session_test.go | 50 ++++++++++++++++++++++ go/internal/runnerhub/relay_comms.go | 3 ++ 2 files changed, 53 insertions(+) diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index 5d036da5b..d2185ae9a 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -356,6 +356,56 @@ func TestErroredCleanupKeepsPeerRebind(t *testing.T) { } } +func TestErroredCleanupDropsStaleAccountForPeerRebind(t *testing.T) { + const otherAccount store.AccountID = "acct-other" + ctx := t.Context() + hub, _, _ := newHub() + bindings := &pausingDeleteBindingStore{fakeBindingStore: newFakeBindingStore(), entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + hub.SetSessionLostSink(newRecordingLostSink()) + hub.enroll(ctx, "runner-1", runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(bindings.release) }) } + t.Cleanup(release) + hub.bindContainer("cont-1", testAgentAccount, "runner-1") + hub.promoteSession(ctx, "cont-1", "sess-1") + gen := hub.EnrollGeneration() + hub.mu.Lock() + old := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + done := make(chan struct{}) + go func() { + defer close(done) + hub.dropLostSessionIfCurrent(ctx, gen, "runner-1", "sess-1", &old, true) + }() + select { + case <-bindings.entered: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not reach the durable delete") + } + // A peer re-binds the session id to another account. + if err := bindings.fakeBindingStore.DeleteSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("peer DeleteSessionBinding: %v", err) + } + if _, _, err := bindings.fakeBindingStore.RecordSessionBinding(ctx, "sess-1", otherAccount, "runner-1"); err != nil { + t.Fatalf("peer RecordSessionBinding: %v", err) + } + release() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatal("cleanup did not finish") + } + + if sess, ok := hub.SessionForAccount(ctx, testAgentAccount); ok { + t.Fatalf("old account resolves to %q after the peer took the session, want a miss", sess) + } + if account, ok := hub.accountForSession(ctx, "sess-1"); !ok || account != otherAccount { + t.Fatalf("binding after old cleanup = (%s, %v), want (%s, true)", account, ok, otherAccount) + } +} + func TestErroredCleanupWithoutDurableWriteKeepsPeerRow(t *testing.T) { ctx := t.Context() hub, _, _ := newHub() diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index 39bbf3392..d1caa5351 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -289,6 +289,9 @@ func (h *Hub) evictStaleBinding(sessionID string, stale *sessionBinding) { defer h.mu.Unlock() if live, ok := h.sessionAccounts[sessionID]; ok && sameBinding(live, *stale) { delete(h.sessionAccounts, sessionID) + if h.accountSessions[stale.account] == sessionID { + delete(h.accountSessions, stale.account) + } } } From 5822a457385619239eac68a0cedc494f84dbf4c5 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 21:48:01 -0400 Subject: [PATCH 16/17] test(store): check usage events on a versioned binding delete (RIG-4742) A stale-version delete records nothing; the current-version delete ends its open interval. Co-authored-by: Matt Wilkinson --- go/internal/store/session_bindings_pgtest_test.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index 2a97e07ce..59a1ee02e 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -360,6 +360,10 @@ func TestDeleteSessionBindingByVersionSkipsRebind(t *testing.T) { if got, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); err != nil || got != agent.ID { t.Fatalf("binding after stale delete = (%q, %v), want (%q, nil)", got, err, agent.ID) } + tenant := s.EffectiveTenant(ctx) + if events := computeEvents(t, s, tenant, agent.ID); len(events) != 1 || events[0].Kind != "start" { + t.Fatalf("usage events after stale delete = %+v, want only the open start", events) + } removed, err = s.DeleteSessionBindingVersion(ctx, "sess-1", newVersion) if err != nil || !removed { t.Fatalf("current DeleteSessionBinding = (%v, %v), want (true, nil)", removed, err) @@ -367,6 +371,10 @@ func TestDeleteSessionBindingByVersionSkipsRebind(t *testing.T) { if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionBinding after delete err = %v, want ErrNotFound", err) } + events := computeEvents(t, s, tenant, agent.ID) + if len(events) != 2 || events[0].Kind != "start" || events[1].Kind != "end" || events[0].IntervalID != events[1].IntervalID { + t.Fatalf("usage events after current delete = %+v, want one start and its end", events) + } } // TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding is the reconnect From 260e55e414d9bca06e6717e4faccaa34fd748d33 Mon Sep 17 00:00:00 2001 From: mintaka Date: Tue, 6 Oct 2026 22:54:54 -0400 Subject: [PATCH 17/17] fix(runnerhub): version reverse read-through bindings; order Stop with enroll (RIG-4742) A binding learned through SessionForAccount had no version or lifetime, so its ERRORED cleanup took the version-less path, found its own row, and never released or reported it. The reverse read now returns the version. A Stop's release now takes bindingWriteMu, as promotion does since the rebase onto the enroll serialization, so its delete cannot land between an enroll's map-clear and its reap and hide the session from the sweep. Co-authored-by: Matt Wilkinson --- go/internal/runnerhub/binding_cache_test.go | 60 ++++++++++++++++--- go/internal/runnerhub/hub.go | 2 +- go/internal/runnerhub/lost_session_test.go | 27 +++++++++ go/internal/runnerhub/relay_comms.go | 16 ++++- .../store/compute_usage_pgtest_test.go | 2 +- go/internal/store/db/session_bindings.sql.go | 9 +-- .../store/queries/session_bindings.sql | 2 +- go/internal/store/session_bindings.go | 13 ++-- .../store/session_bindings_pgtest_test.go | 30 +++++----- go/server/lifecycle_wake_pgtest_test.go | 2 +- 10 files changed, 123 insertions(+), 40 deletions(-) diff --git a/go/internal/runnerhub/binding_cache_test.go b/go/internal/runnerhub/binding_cache_test.go index 93827c9aa..c1cac768e 100644 --- a/go/internal/runnerhub/binding_cache_test.go +++ b/go/internal/runnerhub/binding_cache_test.go @@ -113,18 +113,18 @@ func (f *fakeBindingStore) ResolveSessionBinding(ctx context.Context, sessionID return b.AccountID, b.RunnerID, f.versions[sessionID], nil } -func (f *fakeBindingStore) SessionForAccount(_ context.Context, accountID store.AccountID) (string, string, error) { +func (f *fakeBindingStore) SessionForAccount(_ context.Context, accountID store.AccountID) (string, string, string, error) { f.mu.Lock() defer f.mu.Unlock() if f.reverseErr != nil { - return "", "", f.reverseErr + return "", "", "", f.reverseErr } for sid, b := range f.bindings { if b.AccountID == accountID { - return sid, b.RunnerID, nil + return sid, b.RunnerID, f.versions[sid], nil } } - return "", "", store.ErrNotFound + return "", "", "", store.ErrNotFound } func (f *fakeBindingStore) DeleteSessionBinding(_ context.Context, sessionID string) error { @@ -874,11 +874,11 @@ func (b *reapAndReadBlockingBindingStore) ResolveSessionBinding(ctx context.Cont return account, runnerID, version, err } -func (b *reapAndReadBlockingBindingStore) SessionForAccount(ctx context.Context, accountID store.AccountID) (string, string, error) { - sessionID, runnerID, err := b.fakeBindingStore.SessionForAccount(ctx, accountID) +func (b *reapAndReadBlockingBindingStore) SessionForAccount(ctx context.Context, accountID store.AccountID) (string, string, string, error) { + sessionID, runnerID, version, err := b.fakeBindingStore.SessionForAccount(ctx, accountID) close(b.readEntered) <-b.readRelease - return sessionID, runnerID, err + return sessionID, runnerID, version, err } // TestConcurrentResolveDuringAFaultingReapCannotResurrect fences the window @@ -975,7 +975,7 @@ func TestOlderEnrollReapCannotDeleteNewerPromotion(t *testing.T) { t.Fatal("promotion wrote its binding while the older reap's delete was in flight") } // Ordered after the reap, the live session's binding survives in the store and cache. - if sessionID, runnerID, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); err != nil || sessionID != "sess-new" || runnerID != testRunnerID { + if sessionID, runnerID, _, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); err != nil || sessionID != "sess-new" || runnerID != testRunnerID { t.Fatalf("durable SessionForAccount = (%q, %q, %v), want (sess-new, %s, nil)", sessionID, runnerID, err, testRunnerID) } if sessionID, ok := hub.CachedSessionForAccount(testAgentAccount); !ok || sessionID != "sess-new" { @@ -983,6 +983,48 @@ func TestOlderEnrollReapCannotDeleteNewerPromotion(t *testing.T) { } } +// A Stop's unbind racing an in-flight reap waits for the whole enroll, so the reap +// still finds the session and drives its OFFLINE edge. +func TestStopUnbindWaitsForEnrollReap(t *testing.T) { + hub := newHubOnly() + pres := &fakePresenceSink{} + hub.SetPresenceSink(pres) + plain := newFakeBindingStore() + hub.SetSessionBindingStore(plain) + ctx := store.WithTenant(context.Background(), "tenant-a") + hub.enroll(ctx, testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + hub.bindContainer("cont-1", testAgentAccount, testRunnerID) + hub.promoteSession(ctx, "cont-1", "sess-1") + + bindings := &blockingBindingStore{fakeBindingStore: plain, entered: make(chan struct{}), release: make(chan struct{})} + hub.SetSessionBindingStore(bindings) + reapDone := make(chan struct{}) + go func() { + defer close(reapDone) + hub.enroll(ctx, testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + }() + <-bindings.entered + unbindDone := make(chan struct{}) + go func() { + defer close(unbindDone) + hub.unbindSession(ctx, "sess-1") + }() + waitParkedOrDone(t, "(*Hub).releaseSession", unbindDone) + close(bindings.release) + <-reapDone + <-unbindDone + + offline := 0 + for _, rec := range pres.lifecycleSnapshot() { + if rec.sessionID == "sess-1" && rec.state == compassv1.AgentSessionState_AGENT_SESSION_STATE_DISCONNECTED { + offline++ + } + } + if offline == 0 { + t.Fatal("neither the reap nor the Stop drove sess-1 OFFLINE") + } +} + // overlapDetectingBindingStore parks the first reap and records whether a binding // write arrived while it was parked. type overlapDetectingBindingStore struct { @@ -1093,7 +1135,7 @@ func TestEnrollDuringPromotionCannotResurrectBinding(t *testing.T) { if forwardCached || reverseCached { t.Fatalf("promotion repopulated cache after enroll: forward=%v reverse=%v", forwardCached, reverseCached) } - if sessionID, _, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); !errors.Is(err, store.ErrNotFound) { + if sessionID, _, _, err := plain.SessionForAccount(store.WithTenant(context.Background(), "tenant-a"), testAgentAccount); !errors.Is(err, store.ErrNotFound) { t.Fatalf("durable SessionForAccount = (%q, %v), want ErrNotFound after enroll", sessionID, err) } } diff --git a/go/internal/runnerhub/hub.go b/go/internal/runnerhub/hub.go index e027513d9..01ca98d8c 100644 --- a/go/internal/runnerhub/hub.go +++ b/go/internal/runnerhub/hub.go @@ -223,7 +223,7 @@ type SessionBindingStore interface { // SessionForAccount resolves the live session bound to an account — the // cache-miss read behind SessionForAccount (the reverse direction). Same // fail-closed store.ErrNotFound contract. Returns the owning Runner id too. - SessionForAccount(ctx context.Context, accountID store.AccountID) (sessionID, runnerID string, err error) + SessionForAccount(ctx context.Context, accountID store.AccountID) (sessionID, runnerID, version string, err error) // DeleteSessionBinding releases one session's binding — the unbind write. // Idempotent: releasing an already-released session is a no-op success. DeleteSessionBinding(ctx context.Context, sessionID string) error diff --git a/go/internal/runnerhub/lost_session_test.go b/go/internal/runnerhub/lost_session_test.go index d2185ae9a..4a54409bd 100644 --- a/go/internal/runnerhub/lost_session_test.go +++ b/go/internal/runnerhub/lost_session_test.go @@ -468,6 +468,33 @@ func TestErroredCleanupWithoutDurableWriteKeepsLegacyRow(t *testing.T) { } } +func TestErroredCleanupReleasesReverseReadThroughBinding(t *testing.T) { + ctx := store.WithTenant(t.Context(), "tenant-a") + hub, _, _ := newHub() + bindings := newFakeBindingStore() + hub.SetSessionBindingStore(bindings) + lost := newRecordingLostSink() + hub.SetSessionLostSink(lost) + hub.enroll(ctx, testRunnerID, runnerSubject(), compassv1.RuntimeTier_RUNTIME_TIER_UNSPECIFIED, compassv1.EgressPosture_EGRESS_POSTURE_UNSPECIFIED) + // Another instance promoted the session; this hub learns it from a delivery's reverse read. + bindings.seed("sess-1") + if sess, ok := hub.SessionForAccount(ctx, testAgentAccount); !ok || sess != "sess-1" { + t.Fatalf("SessionForAccount = (%q, %v), want (sess-1, true)", sess, ok) + } + gen := hub.EnrollGeneration() + hub.mu.Lock() + seen := hub.sessionAccounts["sess-1"] + hub.mu.Unlock() + + hub.dropLostSessionIfCurrent(ctx, gen, testRunnerID, "sess-1", &seen, true) + if account, errored := lost.waitOne(t); account != testAgentAccount || !errored { + t.Fatalf("loss report = (%s, %v), want (%s, true)", account, errored, testAgentAccount) + } + if _, _, _, err := bindings.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, store.ErrNotFound) { + t.Fatalf("durable binding after ERRORED cleanup: %v, want ErrNotFound", err) + } +} + // pausingDeleteBindingStore holds the first DeleteSessionBindingVersion until released. type pausingDeleteBindingStore struct { *fakeBindingStore diff --git a/go/internal/runnerhub/relay_comms.go b/go/internal/runnerhub/relay_comms.go index d1caa5351..c921e1595 100644 --- a/go/internal/runnerhub/relay_comms.go +++ b/go/internal/runnerhub/relay_comms.go @@ -191,12 +191,21 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio // so it stays tenant-scoped), then the maps are evicted under h.mu. Delete is // by session id and idempotent — a session already displaced has no row, so a // stale release matches nothing and leaves the live binding alone. + // A Stop's release takes bindingWriteMu so it orders wholly before or after an + // enroll's map-clear through reap. A limited release skips it: mid-reap no binding + // is cached (promotion waits here, read-through is refused), so it releases nothing. + unlockWrite := func() {} + if only == nil { + h.bindingWriteMu.Lock() + unlockWrite = h.bindingWriteMu.Unlock + } h.mu.Lock() bindings := h.bindings routing := h.routing if only != nil { if live, ok := h.sessionAccounts[sessionID]; ok && !sameBinding(live, *only) { h.mu.Unlock() + unlockWrite() return false } } @@ -228,6 +237,7 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio case !removed: // The row was re-bound since only was read: the session is not lost. h.evictStaleBinding(sessionID, only) + unlockWrite() return false default: tenant = string(bindings.EffectiveTenant(ctx)) @@ -243,6 +253,7 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio if ok && only != nil && !sameBinding(binding, *only) { // Re-bound in the cache while the durable delete ran. h.mu.Unlock() + unlockWrite() return false } if ok { @@ -258,6 +269,7 @@ func (h *Hub) releaseSession(ctx context.Context, sessionID string, only *sessio delete(h.sessionAccounts, sessionID) presence := h.presence h.mu.Unlock() + unlockWrite() // Invalidate peer instances' caches (h.mu released, nil-safe, best-effort): // this session has no binding any more (BindingUnbound), so a peer can drop @@ -491,7 +503,7 @@ func (h *Hub) SessionForAccount(ctx context.Context, account store.AccountID) (s if !h.readThroughAllowed(ctx, bindings, enrolled) { return "", false } - sessionID, ownerID, err := bindings.SessionForAccount(ctx, account) + sessionID, ownerID, version, err := bindings.SessionForAccount(ctx, account) if err != nil || ownerID != runnerID { return "", false } @@ -509,7 +521,7 @@ func (h *Hub) SessionForAccount(ctx context.Context, account store.AccountID) (s return "", false } h.accountSessions[account] = sessionID - h.sessionAccounts[sessionID] = sessionBinding{account: account, runnerID: ownerID} + h.sessionAccounts[sessionID] = h.newBindingLocked(account, ownerID, version) h.mu.Unlock() return sessionID, true } diff --git a/go/internal/store/compute_usage_pgtest_test.go b/go/internal/store/compute_usage_pgtest_test.go index 78f4e0585..ded0b40d2 100644 --- a/go/internal/store/compute_usage_pgtest_test.go +++ b/go/internal/store/compute_usage_pgtest_test.go @@ -231,7 +231,7 @@ func TestComputeUsageConflictRollsBackIntervalEvents(t *testing.T) { after[0].IntervalID != before[0].IntervalID || !after[0].OccurredAt.Equal(before[0].OccurredAt) { t.Fatalf("account X events before=%+v after=%+v, want its unchanged open start only", before, after) } - if sessionID, runnerID, err := s.SessionForAccount(ctx, accountX.ID); err != nil || sessionID != "sess-1" || runnerID != "runner-1" { + if sessionID, runnerID, _, err := s.SessionForAccount(ctx, accountX.ID); err != nil || sessionID != "sess-1" || runnerID != "runner-1" { t.Fatalf("account X resolves to (%q, %q, %v), want (sess-1, runner-1)", sessionID, runnerID, err) } } diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go index c57770338..ff14b5827 100644 --- a/go/internal/store/db/session_bindings.sql.go +++ b/go/internal/store/db/session_bindings.sql.go @@ -319,18 +319,19 @@ func (q *Queries) SessionBinding(ctx context.Context, sessionID string) (Session } const sessionBindingForAccount = `-- name: SessionBindingForAccount :one -SELECT session_id, runner_id FROM session_bindings WHERE agent_account_id = $1 +SELECT session_id, runner_id, binding_version FROM session_bindings WHERE agent_account_id = $1 ` type SessionBindingForAccountRow struct { - SessionID string - RunnerID string + SessionID string + RunnerID string + BindingVersion string } func (q *Queries) SessionBindingForAccount(ctx context.Context, agentAccountID string) (SessionBindingForAccountRow, error) { row := q.db.QueryRow(ctx, sessionBindingForAccount, agentAccountID) var i SessionBindingForAccountRow - err := row.Scan(&i.SessionID, &i.RunnerID) + err := row.Scan(&i.SessionID, &i.RunnerID, &i.BindingVersion) return i, err } diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql index d8653b1a1..df2a915f7 100644 --- a/go/internal/store/queries/session_bindings.sql +++ b/go/internal/store/queries/session_bindings.sql @@ -127,7 +127,7 @@ ON CONFLICT (tenant_id, agent_account_id) DO UPDATE SELECT agent_account_id, runner_id, binding_version FROM session_bindings WHERE session_id = $1; -- name: SessionBindingForAccount :one -SELECT session_id, runner_id FROM session_bindings WHERE agent_account_id = $1; +SELECT session_id, runner_id, binding_version FROM session_bindings WHERE agent_account_id = $1; -- Both deletes also write an estimated start for a binding an older server made -- without one; ON CONFLICT keeps any real start. diff --git a/go/internal/store/session_bindings.go b/go/internal/store/session_bindings.go index 168fc7f3b..625a1aaad 100644 --- a/go/internal/store/session_bindings.go +++ b/go/internal/store/session_bindings.go @@ -230,19 +230,20 @@ func (s *Store) ResolveSessionBinding(ctx context.Context, sessionID string) (ac // dropped on a Runner reconnect. Fail-closed for the same reason as above: an // empty session id with a nil error would be dispatched to as if it were a live // session. The consumer's own contract turns this into "push nothing now, let -// the cursor sweep deliver on the recipient's next start". -func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (string, string, error) { +// the cursor sweep deliver on the recipient's next start". version names the +// row's write, as ResolveSessionBinding returns it. +func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (sessionID, runnerID, version string, err error) { if accountID == "" { - return "", "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + return "", "", "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) } row, err := s.q.SessionBindingForAccount(ctx, string(accountID)) if err != nil { if noRows(err) { - return "", "", fmt.Errorf("%w: agent %q has no live session", ErrNotFound, accountID) + return "", "", "", fmt.Errorf("%w: agent %q has no live session", ErrNotFound, accountID) } - return "", "", fmt.Errorf("store: resolve session for account: %w", err) + return "", "", "", fmt.Errorf("store: resolve session for account: %w", err) } - return row.SessionID, row.RunnerID, nil + return row.SessionID, row.RunnerID, row.BindingVersion, nil } // DeleteSessionBinding releases the binding for sessionID — the single-session diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index 59a1ee02e..6c26fc997 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -110,7 +110,7 @@ func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { t.Fatalf("ResolveSessionBinding runner = %q, want runner-1", gotRunner) } - gotSession, gotRunner, err := s.SessionForAccount(ctx, agent.ID) + gotSession, gotRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("SessionForAccount: %v", err) } @@ -140,7 +140,7 @@ func TestSessionBindingLookupsFailClosed(t *testing.T) { // started a session resolves nothing to dispatch to. owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - session, runner, err := s.SessionForAccount(ctx, agent.ID) + session, runner, _, err := s.SessionForAccount(ctx, agent.ID) if !errors.Is(err, ErrNotFound) { t.Fatalf("SessionForAccount(unbound agent) err = %v, want errors.Is(_, ErrNotFound)", err) } @@ -176,7 +176,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { } // The account now resolves to the NEW session (the session_id assignment). - gotSession, gotRunner, err := s.SessionForAccount(ctx, agent.ID) + gotSession, gotRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("SessionForAccount after the re-point: %v", err) } @@ -275,7 +275,7 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T if gotAccount != agentA.ID { t.Fatalf("sess-1 resolves to %q, want the original owner %q", gotAccount, agentA.ID) } - if _, _, err := s.SessionForAccount(ctx, agentB.ID); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.SessionForAccount(ctx, agentB.ID); !errors.Is(err, ErrNotFound) { t.Fatalf("SessionForAccount(agent-b) err = %v, want ErrNotFound — the refused bind must not have landed", err) } } @@ -310,7 +310,7 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { if _, _, _, err := s.ResolveSessionBinding(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionAccount after delete err = %v, want ErrNotFound", err) } - if _, _, err := s.SessionForAccount(ctx, agent.ID); !errors.Is(err, ErrNotFound) { + if _, _, _, err := s.SessionForAccount(ctx, agent.ID); !errors.Is(err, ErrNotFound) { t.Fatalf("SessionForAccount after delete err = %v, want ErrNotFound", err) } @@ -526,7 +526,7 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { } // The reverse direction leaks nothing either. - if gotSession, _, err := s.SessionForAccount(ctxB, agentA.ID); !errors.Is(err, ErrNotFound) { + if gotSession, _, _, err := s.SessionForAccount(ctxB, agentA.ID); !errors.Is(err, ErrNotFound) { t.Fatalf("tenant B SessionForAccount(A's agent) = (%q, %v), want ErrNotFound — cross-tenant read leak", gotSession, err) } @@ -617,10 +617,10 @@ func TestSessionBindingSameSessionIDInTwoTenantsCoexist(t *testing.T) { } // And the reverse direction, per tenant. - if got, runner, err := s.SessionForAccount(ctxA, agentA.ID); err != nil || got != "sess-shared" || runner != "runner-a" { + if got, runner, _, err := s.SessionForAccount(ctxA, agentA.ID); err != nil || got != "sess-shared" || runner != "runner-a" { t.Fatalf("tenant A SessionForAccount = (%q, %q, %v), want (sess-shared, runner-a, nil)", got, runner, err) } - if got, runner, err := s.SessionForAccount(ctxB, agentB.ID); err != nil || got != "sess-shared" || runner != "runner-b" { + if got, runner, _, err := s.SessionForAccount(ctxB, agentB.ID); err != nil || got != "sess-shared" || runner != "runner-b" { t.Fatalf("tenant B SessionForAccount = (%q, %q, %v), want (sess-shared, runner-b, nil)", got, runner, err) } } @@ -655,10 +655,10 @@ func TestSessionBindingSameAccountIDInTwoTenantsCoexist(t *testing.T) { } // Both rows survive, one per tenant, each resolving its own session. - if got, runner, err := s.SessionForAccount(ctxA, shared.ID); err != nil || got != "sess-a" || runner != "runner-a" { + if got, runner, _, err := s.SessionForAccount(ctxA, shared.ID); err != nil || got != "sess-a" || runner != "runner-a" { t.Fatalf("tenant A SessionForAccount = (%q, %q, %v), want (sess-a, runner-a, nil) — B's write reached A's row", got, runner, err) } - if got, runner, err := s.SessionForAccount(ctxB, shared.ID); err != nil || got != "sess-b" || runner != "runner-b" { + if got, runner, _, err := s.SessionForAccount(ctxB, shared.ID); err != nil || got != "sess-b" || runner != "runner-b" { t.Fatalf("tenant B SessionForAccount = (%q, %q, %v), want (sess-b, runner-b, nil)", got, runner, err) } @@ -710,7 +710,7 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { // Under the system role BOTH rows are visible, so the :one read is ambiguous. // It does NOT error — that is the hazard. It returns one of the two, and // which one is not something the caller can control or detect. - got, gotRunner, err := s.SessionForAccount(WithSystemRole(context.Background()), shared.ID) + got, gotRunner, _, err := s.SessionForAccount(WithSystemRole(context.Background()), shared.ID) if err != nil { t.Fatalf("SessionForAccount under the system role: %v — the current behaviour is a SILENT pick, not an error; if this now errors, PR3 changed the contract and this test must be updated deliberately", err) } @@ -721,10 +721,10 @@ func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { // The same read on the REQUEST path is exact in both tenants: the hazard is // the system role's missing scoping, NOT anything about the data. - if v, runner, err := s.SessionForAccount(ctxA, shared.ID); err != nil || v != "sess-a" || runner != "runner-a" { + if v, runner, _, err := s.SessionForAccount(ctxA, shared.ID); err != nil || v != "sess-a" || runner != "runner-a" { t.Fatalf("tenant A request-path SessionForAccount = (%q, %q, %v), want (sess-a, runner-a, nil)", v, runner, err) } - if v, runner, err := s.SessionForAccount(ctxB, shared.ID); err != nil || v != "sess-b" || runner != "runner-b" { + if v, runner, _, err := s.SessionForAccount(ctxB, shared.ID); err != nil || v != "sess-b" || runner != "runner-b" { t.Fatalf("tenant B request-path SessionForAccount = (%q, %q, %v), want (sess-b, runner-b, nil)", v, runner, err) } @@ -883,7 +883,7 @@ func TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced(t *testin if n := countBindings(t, ctx, s, agent.ID); n != 1 { t.Fatalf("bindings after two concurrent re-points = %d, want 1", n) } - live, liveRunner, err := s.SessionForAccount(ctx, agent.ID) + live, liveRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("SessionForAccount after the race: %v", err) } @@ -1145,7 +1145,7 @@ func TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep(t *t if n := countBindings(t, ctx, s, agent.ID); n != 1 { t.Fatalf("iteration %d: bindings after the race = %d, want 1 — the bind must land whether or not the sweep removed the prior row", i, n) } - live, liveRunner, err := s.SessionForAccount(ctx, agent.ID) + live, liveRunner, _, err := s.SessionForAccount(ctx, agent.ID) if err != nil { t.Fatalf("iteration %d: SessionForAccount after the race: %v", i, err) } diff --git a/go/server/lifecycle_wake_pgtest_test.go b/go/server/lifecycle_wake_pgtest_test.go index 3dc3c283e..db1b20da7 100644 --- a/go/server/lifecycle_wake_pgtest_test.go +++ b/go/server/lifecycle_wake_pgtest_test.go @@ -138,7 +138,7 @@ func TestWakeAgentStaleBindingRowStillResumes(t *testing.T) { t.Fatalf("RecordSessionBinding: %v", err) } // Checked on the store, not the hub: a hub read-through would warm the cache. - if got, runnerID, err := f.store.SessionForAccount(ctx, f.agentID); err != nil || got != "sess-dead" || runnerID != fakeRunnerID { + if got, runnerID, _, err := f.store.SessionForAccount(ctx, f.agentID); err != nil || got != "sess-dead" || runnerID != fakeRunnerID { t.Fatalf("precondition: durable row = (%q, %q, %v), want (sess-dead, %s, nil)", got, runnerID, err, fakeRunnerID) } f.runner.forget()