Advance durable Agent session generations on restart

internal/rpc/server.go

@@ -294,6 +294,9 @@ func (r *SessionRegistry) Activate(binding *agentpb.AgentBinding, activationOper
 			return cloneSession(existing.session), true, nil
 		}
 		if binding.SessionGeneration == 0 {
+			if existing.binding.SessionGeneration == ^uint64(0) {
+				return nil, false, status.Error(codes.Aborted, "session generation is exhausted")
+			}
 			binding = proto.Clone(binding).(*agentpb.AgentBinding)
 			binding.SessionGeneration = existing.binding.SessionGeneration + 1
 		}
@@ -304,6 +307,12 @@ func (r *SessionRegistry) Activate(binding *agentpb.AgentBinding, activationOper
 	if binding.SessionGeneration == 0 {
 		binding = proto.Clone(binding).(*agentpb.AgentBinding)
 		binding.SessionGeneration = 1
+		if previous, ok := r.generations[binding.AgentId]; ok {
+			if previous == ^uint64(0) {
+				return nil, false, status.Error(codes.Aborted, "session generation is exhausted")
+			}
+			binding.SessionGeneration = previous + 1
+		}
 	}
 	if previous, ok := r.generations[binding.AgentId]; ok && binding.SessionGeneration <= previous {
 		return nil, false, status.Error(codes.Aborted, "persisted session generation is fenced")

internal/rpc/session_test.go

@@ -1,6 +1,7 @@
 package rpc
 
 import (
+	"strings"
 	"testing"
 	"time"
 
@@ -25,3 +26,45 @@ func TestSessionRegistryPersistsGenerationAcrossRestart(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestSessionRegistryZeroGenerationAdvancesDurableHighwaterAfterRestart(t *testing.T) {
+	path := t.TempDir() + "/rpc-session.json"
+	first := NewSessionRegistry(path)
+	now := time.Unix(100, 0)
+	if _, _, err := first.Activate(&agentpb.AgentBinding{AgentId: "agent-1", CellId: "cell-1", ExpectedBootId: "boot-1", DispatcherEpoch: "epoch-1", SessionGeneration: 7}, "activate-1", "digest-1", now); err != nil {
+		t.Fatal(err)
+	}
+	second := NewSessionRegistry(path)
+	session, replay, err := second.Activate(&agentpb.AgentBinding{AgentId: "agent-1", CellId: "cell-1", ExpectedBootId: "boot-2", DispatcherEpoch: "epoch-2"}, "activate-2", "digest-2", now)
+	if err != nil || replay || session.GetSessionGeneration() != 8 {
+		t.Fatalf("Agent could not advance persisted fencing after restart: session=%v replay=%t err=%v", session, replay, err)
+	}
+	third := NewSessionRegistry(path)
+	if _, _, err := third.Activate(&agentpb.AgentBinding{AgentId: "agent-1", CellId: "cell-1", ExpectedBootId: "boot-3", DispatcherEpoch: "epoch-3", SessionGeneration: 7}, "activate-3", "digest-3", now); status.Code(err) != codes.Aborted {
+		t.Fatalf("explicit stale generation bypassed the durable fence: %v", err)
+	}
+}
+
+func TestSessionRegistryZeroGenerationDoesNotWrapAtExhaustion(t *testing.T) {
+	path := t.TempDir() + "/rpc-session.json"
+	first := NewSessionRegistry(path)
+	now := time.Unix(100, 0)
+	if _, _, err := first.Activate(&agentpb.AgentBinding{AgentId: "agent-1", CellId: "cell-1", ExpectedBootId: "boot-1", DispatcherEpoch: "epoch-1", SessionGeneration: ^uint64(0)}, "activate-1", "digest-1", now); err != nil {
+		t.Fatal(err)
+	}
+	for _, tc := range []struct {
+		name     string
+		registry *SessionRegistry
+		boot     string
+	}{
+		{"active", first, "boot-1"},
+		{"after restart", NewSessionRegistry(path), "boot-2"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			_, _, err := tc.registry.Activate(&agentpb.AgentBinding{AgentId: "agent-1", CellId: "cell-1", ExpectedBootId: tc.boot, DispatcherEpoch: "epoch-2"}, "activate-2", "digest-2", now)
+			if status.Code(err) != codes.Aborted || !strings.Contains(err.Error(), "exhausted") {
+				t.Fatalf("exhausted generation wrapped or was hidden: %v", err)
+			}
+		})
+	}
+}