Harden Stage 5 subscription safety checks

Fence stale ACK results by generation, make duplicate MarkProcessed harmless after confirmed ACK, and cover forged delivery, mutation, in-flight ACK, initial subscribe, and replacement edge cases.

Co-authored-by: Cursor <cursoragent@cursor.com>

README.md

@@ -106,3 +106,19 @@ The SDK wraps implemented read/agent RPCs:
 Generic `PostMessage`, `CreateTopic`, and `UpdateTopic` are intentionally not
 exposed as ergonomic helpers because the current server returns
 `Unimplemented` for those fail-closed gRPC paths.
+
+## Lifecycle Notes
+
+`Client` owns and closes the gRPC connection it creates in `New`. If callers
+provide an existing connection with `WithClientConn`, the SDK uses it but does
+not close it; callers keep ownership of that connection.
+
+`Subscription.Run` is sequential: it calls the handler for one delivery at a
+time and calls `MarkProcessed` only after the handler returns nil. Handler
+errors stop `Run` and do not ACK that delivery. Automatic reconnect cancels the
+subscription stream generation, but it cannot undo external handler side
+effects; handlers should honor their context and remain idempotent.
+
+Calling `MarkProcessed` more than once for the same delivery is harmless. The
+first call may advance the safe ACK boundary; later duplicate calls do not add
+tracker entries, move the cursor backward, or ACK a caller-mutated sequence.

ack_tracker.go

@@ -13,10 +13,14 @@ type ackEntry struct {
 type ackTracker struct {
 	entries []ackEntry
 	index   map[*deliveryToken]int
+	retired map[*deliveryToken]struct{}
 }
 
 func newACKTracker() ackTracker {
-	return ackTracker{index: make(map[*deliveryToken]int)}
+	return ackTracker{
+		index:   make(map[*deliveryToken]int),
+		retired: make(map[*deliveryToken]struct{}),
+	}
 }
 
 func (t *ackTracker) add(seq int64, token *deliveryToken) {
@@ -27,6 +31,9 @@ func (t *ackTracker) add(seq int64, token *deliveryToken) {
 func (t *ackTracker) mark(token *deliveryToken) (int64, bool, bool) {
 	i, ok := t.index[token]
 	if !ok {
+		if _, retired := t.retired[token]; retired {
+			return t.safeBoundary(), true, true
+		}
 		return 0, false, false
 	}
 	if t.entries[i].done {
@@ -53,6 +60,7 @@ func (t *ackTracker) removeConfirmed(lastAck int64) {
 		if !t.entries[n].done || t.entries[n].seq > lastAck {
 			break
 		}
+		t.retired[t.entries[n].token] = struct{}{}
 		delete(t.index, t.entries[n].token)
 		n++
 	}

subscription.go

@@ -303,8 +303,9 @@ func (s *Subscription) runACK() {
 				SubscriptionId:  subID,
 			})
 			if err != nil {
-				if s.ctx.Err() != nil || !isRetryableACKError(err) || !s.sameGeneration(gen) {
+				current := s.sameGeneration(gen)
-					if s.ctx.Err() == nil && !isRetryableACKError(err) {
+				if s.ctx.Err() != nil || !isRetryableACKError(err) || !current {
+					if s.ctx.Err() == nil && current && !isRetryableACKError(err) {
 						s.setErr(err)
 						s.cancel()
 					}

subscription_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"strings"
 	"sync"
@@ -325,6 +326,101 @@ func TestOrderedCumulativeACKWithSequenceGaps(t *testing.T) {
 	waitACKs(t, fake, 104, 111)
 }
 
+func TestForgedDeliveryRejectedWithoutACK(t *testing.T) {
+	fake := newFake()
+	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
+		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	c := newClientForTest(t, fake)
+	sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = sub.MarkProcessed(context.Background(), Delivery{
+		Message: &pb.Message{Sequence: 999999},
+		Replay:  true,
+	})
+	if !errors.Is(err, ErrUnknownDelivery) {
+		t.Fatalf("forged delivery MarkProcessed = %v", err)
+	}
+	fake.mu.Lock()
+	defer fake.mu.Unlock()
+	if len(fake.ackCalls) != 0 {
+		t.Fatalf("forged delivery produced ACKs: %+v", fake.ackCalls)
+	}
+}
+
+func TestCrossSubscriptionDeliveryRejectedWithoutACK(t *testing.T) {
+	fake := newFake()
+	fake.subscribeFunc = func(req *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
+		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
+		dm := msg(100)
+		dm.Message.TopicId = req.GetTopicId()
+		dm.SubscriptionId = "sub-" + req.GetTopicId()
+		_ = stream.Send(dm)
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	c := newClientForTest(t, fake)
+	subA, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "a"})
+	subB, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "b"})
+	dA := recvDelivery(t, subA)
+	_ = recvDelivery(t, subB)
+	if err := subB.MarkProcessed(context.Background(), dA); !errors.Is(err, ErrUnknownDelivery) {
+		t.Fatalf("cross-subscription MarkProcessed = %v", err)
+	}
+	fake.mu.Lock()
+	defer fake.mu.Unlock()
+	if len(fake.ackCalls) != 0 {
+		t.Fatalf("cross-subscription delivery produced ACKs: %+v", fake.ackCalls)
+	}
+}
+
+func TestMutatedDeliveryMessageCannotAlterACKSequence(t *testing.T) {
+	fake := newFake()
+	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
+		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
+		_ = stream.Send(msg(100))
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	c := newClientForTest(t, fake)
+	sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+	d := recvDelivery(t, sub)
+	d.Message.Sequence = 111
+	if err := sub.MarkProcessed(context.Background(), d); err != nil {
+		t.Fatal(err)
+	}
+	waitACKs(t, fake, 100)
+}
+
+func TestDuplicateMarkProcessedIsHarmless(t *testing.T) {
+	fake := newFake()
+	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
+		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
+		_ = stream.Send(msg(100))
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	c := newClientForTest(t, fake)
+	sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+	d := recvDelivery(t, sub)
+	if err := sub.MarkProcessed(context.Background(), d); err != nil {
+		t.Fatal(err)
+	}
+	waitACKs(t, fake, 100)
+	if err := sub.MarkProcessed(context.Background(), d); err != nil {
+		t.Fatalf("duplicate MarkProcessed = %v", err)
+	}
+	fake.mu.Lock()
+	defer fake.mu.Unlock()
+	if len(fake.ackCalls) != 1 || fake.ackCalls[0].seq != 100 {
+		t.Fatalf("duplicate MarkProcessed changed ACKs: %+v", fake.ackCalls)
+	}
+}
+
 func TestLargerOutOfOrderCompletionDoesNotSkipEarliestPending(t *testing.T) {
 	fake := newFake()
 	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
@@ -447,6 +543,102 @@ func TestACKResponseLostAmbiguityReconnectUsesServerCursor(t *testing.T) {
 	waitACKs(t, fake, 104, 104, 111)
 }
 
+func TestACKPrefixGrowthWhileACKInFlightAdvancesAfterSuccess(t *testing.T) {
+	fake := newFake()
+	ack100Started := make(chan struct{})
+	release100 := make(chan struct{})
+	var once sync.Once
+	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
+		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
+		_ = stream.Send(msg(100))
+		_ = stream.Send(msg(104))
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	fake.ackFunc = func(_ context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) {
+		fake.mu.Lock()
+		fake.ackCalls = append(fake.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()})
+		fake.mu.Unlock()
+		if req.GetLastAckSequence() == 100 {
+			once.Do(func() { close(ack100Started) })
+			<-release100
+		}
+		return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil
+	}
+	c := newClientForTest(t, fake)
+	sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+	d100, d104 := recvDelivery(t, sub), recvDelivery(t, sub)
+	if err := sub.MarkProcessed(context.Background(), d100); err != nil {
+		t.Fatal(err)
+	}
+	<-ack100Started
+	if err := sub.MarkProcessed(context.Background(), d104); err != nil {
+		t.Fatal(err)
+	}
+	close(release100)
+	waitACKs(t, fake, 100, 104)
+}
+
+func TestOldACKResultCannotCorruptReconnectedGeneration(t *testing.T) {
+	fake := newFake()
+	dropFirst := make(chan struct{})
+	oldACKStarted := make(chan struct{})
+	releaseOldACK := make(chan struct{})
+	var oldOnce sync.Once
+	var calls int
+	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
+		calls++
+		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
+		if calls == 1 {
+			_ = stream.Send(msg(100))
+			<-dropFirst
+			return status.Error(codes.Unavailable, "generation A failed")
+		}
+		dm := msg(100)
+		dm.Replay = true
+		dm.SubscriptionId = "sub-2"
+		_ = stream.Send(dm)
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	fake.ackFunc = func(_ context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) {
+		fake.mu.Lock()
+		fake.ackCalls = append(fake.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()})
+		fake.mu.Unlock()
+		if req.GetSubscriptionId() == "sub-1" {
+			oldOnce.Do(func() { close(oldACKStarted) })
+			<-releaseOldACK
+			return nil, status.Error(codes.Aborted, "old subscription_id rejected")
+		}
+		return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil
+	}
+	c := newClientForTest(t, fake,
+		WithReconnectBackoff(10*time.Millisecond, 10*time.Millisecond),
+		withSleeperForTest(sleepContext),
+	)
+	sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+	old100 := recvDelivery(t, sub)
+	if err := sub.MarkProcessed(context.Background(), old100); err != nil {
+		t.Fatal(err)
+	}
+	<-oldACKStarted
+	close(dropFirst)
+	new100 := recvDelivery(t, sub)
+	close(releaseOldACK)
+	if err := sub.MarkProcessed(context.Background(), new100); err != nil {
+		t.Fatal(err)
+	}
+	waitACKs(t, fake, 100, 100)
+	if err := sub.Err(); err != nil {
+		t.Fatalf("old ACK result corrupted generation B: %v", err)
+	}
+	fake.mu.Lock()
+	defer fake.mu.Unlock()
+	if fake.ackCalls[0].subID != "sub-1" || fake.ackCalls[1].subID != "sub-2" {
+		t.Fatalf("ACKs used wrong subscription ids: %+v", fake.ackCalls)
+	}
+}
+
 func TestReconnectReplayAndStaleGenerationFencing(t *testing.T) {
 	fake := newFake()
 	var calls int
@@ -507,6 +699,48 @@ func TestStreamReplacementAbortedIsTerminalNoReconnect(t *testing.T) {
 	}
 }
 
+func TestReplacementBetweenTwoSDKSubscriptionsDoesNotPingPong(t *testing.T) {
+	fake := newFake()
+	replaceFirst := make(chan struct{})
+	secondReady := make(chan struct{})
+	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
+		fake.mu.Lock()
+		call := fake.subscribeCount
+		fake.mu.Unlock()
+		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
+		if call == 1 {
+			<-replaceFirst
+			return status.Error(codes.Aborted, "subscription replaced")
+		}
+		close(secondReady)
+		<-stream.Context().Done()
+		return stream.Context().Err()
+	}
+	c := newClientForTest(t, fake)
+	first, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	second, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	<-secondReady
+	close(replaceFirst)
+	waitDone(t, first)
+	if !errors.Is(first.Err(), ErrSubscriptionReplaced) {
+		t.Fatalf("first Err = %v", first.Err())
+	}
+	if err := second.Close(); err != nil {
+		t.Fatal(err)
+	}
+	fake.mu.Lock()
+	defer fake.mu.Unlock()
+	if fake.subscribeCount != 2 {
+		t.Fatalf("replacement caused reconnect ping-pong: %d subscribes", fake.subscribeCount)
+	}
+}
+
 func TestAuthTerminalErrorsDoNotEndlesslyReconnect(t *testing.T) {
 	for _, code := range []codes.Code{codes.Unauthenticated, codes.PermissionDenied} {
 		t.Run(code.String(), func(t *testing.T) {
@@ -527,6 +761,40 @@ func TestAuthTerminalErrorsDoNotEndlesslyReconnect(t *testing.T) {
 	}
 }
 
+func TestInitialSubscribeFailuresBeforeReadyReturnImmediately(t *testing.T) {
+	cases := []struct {
+		name string
+		err  error
+		code codes.Code
+	}{
+		{"Unavailable", status.Error(codes.Unavailable, "transient outage"), codes.Unavailable},
+		{"EOF", io.EOF, codes.Unknown},
+		{"DeadlineExceeded", status.Error(codes.DeadlineExceeded, "deadline"), codes.DeadlineExceeded},
+		{"Unauthenticated", status.Error(codes.Unauthenticated, "bad token"), codes.Unauthenticated},
+		{"PermissionDenied", status.Error(codes.PermissionDenied, "denied"), codes.PermissionDenied},
+		{"NotFound", status.Error(codes.NotFound, "missing topic"), codes.NotFound},
+		{"InvalidArgument", status.Error(codes.InvalidArgument, "bad topic"), codes.InvalidArgument},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			fake := newFake()
+			fake.subscribeFunc = func(_ *pb.SubscribeRequest, _ pb.ChatService_SubscribeServer) error {
+				return tc.err
+			}
+			c := newClientForTest(t, fake)
+			_, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
+			if status.Code(err) != tc.code {
+				t.Fatalf("Subscribe error = %v, code %v want %v", err, status.Code(err), tc.code)
+			}
+			fake.mu.Lock()
+			defer fake.mu.Unlock()
+			if fake.subscribeCount != 1 {
+				t.Fatalf("initial failure retried before READY: %d subscribes", fake.subscribeCount)
+			}
+		})
+	}
+}
+
 func TestTransientReconnect(t *testing.T) {
 	fake := newFake()
 	var calls int