l-ai-chat-client-go/subscription_test.go

package chatclient

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"testing"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/grpc/test/bufconn"

	pb "git.be.kaien.ru/l-ai/l-ai-chat-client-go/proto/l-ai-chat-proto"
)

type fakeChatServer struct {
	pb.UnimplementedChatServiceServer

	mu             sync.Mutex
	authByMethod   map[string][]string
	postReqs       []*pb.AgentPostMessageRequest
	genericPosts   int
	ackCalls       []ackCall
	ackErrors      map[int64][]error
	ackFunc        func(context.Context, *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error)
	subscribeCount int
	subscribeFunc  func(*pb.SubscribeRequest, pb.ChatService_SubscribeServer) error
	historyPages   []*pb.HistoryResponse
	historyReqs    []*pb.HistoryRequest
}

type ackCall struct {
	topic string
	subID string
	seq   int64
}

func newFake() *fakeChatServer {
	return &fakeChatServer{
		authByMethod: make(map[string][]string),
		ackErrors:    make(map[int64][]error),
	}
}

func (f *fakeChatServer) record(ctx context.Context, method string) {
	md, _ := metadata.FromIncomingContext(ctx)
	f.mu.Lock()
	defer f.mu.Unlock()
	f.authByMethod[method] = append([]string(nil), md.Get("authorization")...)
}

func (f *fakeChatServer) ListTopics(ctx context.Context, _ *pb.ListTopicsRequest) (*pb.ListTopicsResponse, error) {
	f.record(ctx, "ListTopics")
	return &pb.ListTopicsResponse{Topics: []*pb.Topic{{Id: "topic-1", Slug: "general"}}}, nil
}

func (f *fakeChatServer) GetTopic(ctx context.Context, _ *pb.GetTopicRequest) (*pb.Topic, error) {
	f.record(ctx, "GetTopic")
	return &pb.Topic{Id: "topic-1", Slug: "general"}, nil
}

func (f *fakeChatServer) GetMessage(ctx context.Context, req *pb.GetMessageRequest) (*pb.Message, error) {
	f.record(ctx, "GetMessage")
	return &pb.Message{Id: req.GetMessageId(), Sequence: 1}, nil
}

func (f *fakeChatServer) PostMessage(context.Context, *pb.PostMessageRequest) (*pb.PostMessageResponse, error) {
	f.mu.Lock()
	f.genericPosts++
	f.mu.Unlock()
	return nil, status.Error(codes.Unimplemented, "generic post is fail-closed")
}

func (f *fakeChatServer) AgentPostMessage(ctx context.Context, req *pb.AgentPostMessageRequest) (*pb.PostMessageResponse, error) {
	f.record(ctx, "AgentPostMessage")
	f.mu.Lock()
	f.postReqs = append(f.postReqs, req)
	f.mu.Unlock()
	return &pb.PostMessageResponse{
		Message: &pb.Message{Id: "posted", TopicId: req.GetTopicId(), Content: req.GetContent(), Sequence: 10},
	}, nil
}

func (f *fakeChatServer) History(ctx context.Context, req *pb.HistoryRequest) (*pb.HistoryResponse, error) {
	f.record(ctx, "History")
	f.mu.Lock()
	defer f.mu.Unlock()
	f.historyReqs = append(f.historyReqs, req)
	if len(f.historyPages) == 0 {
		return &pb.HistoryResponse{}, nil
	}
	page := f.historyPages[0]
	f.historyPages = f.historyPages[1:]
	return page, nil
}

func (f *fakeChatServer) Subscribe(req *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
	f.record(stream.Context(), "Subscribe")
	f.mu.Lock()
	f.subscribeCount++
	fn := f.subscribeFunc
	f.mu.Unlock()
	if fn == nil {
		return status.Error(codes.Unavailable, "no subscribe script")
	}
	return fn(req, stream)
}

func (f *fakeChatServer) Acknowledge(ctx context.Context, req *pb.AcknowledgeRequest) (*pb.AcknowledgeResponse, error) {
	f.record(ctx, "Acknowledge")
	f.mu.Lock()
	fn := f.ackFunc
	f.mu.Unlock()
	if fn != nil {
		return fn(ctx, req)
	}
	f.mu.Lock()
	f.ackCalls = append(f.ackCalls, ackCall{topic: req.GetTopicId(), subID: req.GetSubscriptionId(), seq: req.GetLastAckSequence()})
	errs := f.ackErrors[req.GetLastAckSequence()]
	if len(errs) > 0 {
		err := errs[0]
		f.ackErrors[req.GetLastAckSequence()] = errs[1:]
		f.mu.Unlock()
		return nil, err
	}
	f.mu.Unlock()
	return &pb.AcknowledgeResponse{LastAckSequence: req.GetLastAckSequence()}, nil
}

func newClientForTest(t *testing.T, fake *fakeChatServer, opts ...Option) *Client {
	t.Helper()
	lis := bufconn.Listen(1 << 20)
	srv := grpc.NewServer()
	pb.RegisterChatServiceServer(srv, fake)
	go func() { _ = srv.Serve(lis) }()
	t.Cleanup(func() {
		srv.Stop()
		_ = lis.Close()
	})

	all := []Option{
		WithStaticToken("test-token"),
		WithInsecureTransport(),
		WithGRPCDialOptions(grpc.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
			return lis.DialContext(ctx)
		})),
		withJitterForTest(func(d time.Duration) time.Duration { return d }),
		withSleeperForTest(func(ctx context.Context, d time.Duration) error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			default:
				return nil
			}
		}),
	}
	all = append(all, opts...)
	c, err := New(context.Background(), "passthrough:///bufnet", all...)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = c.Close() })
	return c
}

func msg(seq int64) *pb.DeliveryMessage {
	return &pb.DeliveryMessage{
		Message:        &pb.Message{Id: fmt.Sprintf("m-%d", seq), TopicId: "topic", Sequence: seq},
		SubscriptionId: "sub-1",
	}
}

func recvDelivery(t *testing.T, sub *Subscription) Delivery {
	t.Helper()
	select {
	case d := <-sub.Deliveries():
		return d
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for delivery")
		return Delivery{}
	}
}

func waitDone(t *testing.T, sub *Subscription) {
	t.Helper()
	select {
	case <-sub.Done():
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for subscription done")
	}
}

func waitACKs(t *testing.T, fake *fakeChatServer, want ...int64) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		fake.mu.Lock()
		got := make([]int64, len(fake.ackCalls))
		for i, c := range fake.ackCalls {
			got[i] = c.seq
		}
		fake.mu.Unlock()
		if len(got) >= len(want) {
			ok := true
			for i := range want {
				if got[i] != want[i] {
					ok = false
				}
			}
			if ok {
				return
			}
		}
		time.Sleep(10 * time.Millisecond)
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	t.Fatalf("ack calls = %+v, want prefix %v", fake.ackCalls, want)
}

func TestAuthMetadataUnarySubscribeAcknowledgeAndEmptyTokenFailsClosed(t *testing.T) {
	fake := newFake()
	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		if err := stream.Send(&pb.DeliveryMessage{Ready: true}); err != nil {
			return err
		}
		if err := stream.Send(msg(100)); err != nil {
			return err
		}
		<-stream.Context().Done()
		return stream.Context().Err()
	}
	c := newClientForTest(t, fake)
	if _, err := c.ListTopics(context.Background()); err != nil {
		t.Fatal(err)
	}
	sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
	if err != nil {
		t.Fatal(err)
	}
	d := recvDelivery(t, sub)
	if err := sub.MarkProcessed(context.Background(), d); err != nil {
		t.Fatal(err)
	}
	waitACKs(t, fake, 100)

	fake.mu.Lock()
	defer fake.mu.Unlock()
	for _, method := range []string{"ListTopics", "Subscribe", "Acknowledge"} {
		got := fake.authByMethod[method]
		if len(got) != 1 || got[0] != "Bearer test-token" {
			t.Fatalf("%s auth metadata = %v", method, got)
		}
	}

	empty := newClientForTest(t, newFake(), WithStaticToken(""))
	_, err = empty.ListTopics(context.Background())
	if status.Code(err) != codes.Unauthenticated {
		t.Fatalf("empty token: got %v", err)
	}
	if err != nil && strings.Contains(err.Error(), "test-token") {
		t.Fatalf("error leaked token: %v", err)
	}
}

func TestPostMessageUsesAgentPostMessageAndPreservesIdempotencyKey(t *testing.T) {
	fake := newFake()
	c := newClientForTest(t, fake)
	_, err := c.PostMessage(context.Background(), PostMessageRequest{
		TopicID:        "topic",
		Content:        "payload",
		IdempotencyKey: "logical-key",
	})
	if err != nil {
		t.Fatal(err)
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	if fake.genericPosts != 0 {
		t.Fatal("generic PostMessage must not be called")
	}
	if len(fake.postReqs) != 1 || fake.postReqs[0].GetIdempotencyKey() != "logical-key" {
		t.Fatalf("AgentPostMessage reqs = %+v", fake.postReqs)
	}
}

func TestOrderedCumulativeACKWithSequenceGaps(t *testing.T) {
	fake := newFake()
	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		for _, seq := range []int64{100, 104, 111} {
			_ = stream.Send(msg(seq))
		}
		<-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)
	}
	d100, d104, d111 := recvDelivery(t, sub), recvDelivery(t, sub), recvDelivery(t, sub)
	if err := sub.MarkProcessed(context.Background(), d104); err != nil {
		t.Fatal(err)
	}
	time.Sleep(30 * time.Millisecond)
	fake.mu.Lock()
	if len(fake.ackCalls) != 0 {
		t.Fatalf("ACK advanced over unprocessed 100: %+v", fake.ackCalls)
	}
	fake.mu.Unlock()
	if err := sub.MarkProcessed(context.Background(), d100); err != nil {
		t.Fatalf("mark 100: %v", err)
	}
	waitACKs(t, fake, 104)
	if err := sub.MarkProcessed(context.Background(), d111); err != nil {
		t.Fatalf("mark 111: %v", err)
	}
	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 {
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		for _, seq := range []int64{100, 104, 111, 118, 130} {
			_ = stream.Send(msg(seq))
		}
		<-stream.Context().Done()
		return stream.Context().Err()
	}
	c := newClientForTest(t, fake)
	sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
	deliveries := map[int64]Delivery{}
	for i := 0; i < 5; i++ {
		d := recvDelivery(t, sub)
		deliveries[d.Message.GetSequence()] = d
	}
	for _, seq := range []int64{118, 104} {
		_ = sub.MarkProcessed(context.Background(), deliveries[seq])
	}
	time.Sleep(30 * time.Millisecond)
	fake.mu.Lock()
	if len(fake.ackCalls) != 0 {
		t.Fatalf("ACK skipped earliest pending: %+v", fake.ackCalls)
	}
	fake.mu.Unlock()
	if err := sub.MarkProcessed(context.Background(), deliveries[100]); err != nil {
		t.Fatalf("mark 100: %v", err)
	}
	waitACKs(t, fake, 104)
	if err := sub.MarkProcessed(context.Background(), deliveries[130]); err != nil {
		t.Fatalf("mark 130: %v", err)
	}
	time.Sleep(30 * time.Millisecond)
	fake.mu.Lock()
	if len(fake.ackCalls) != 1 {
		t.Fatalf("ACK skipped 111: %+v", fake.ackCalls)
	}
	fake.mu.Unlock()
	if err := sub.MarkProcessed(context.Background(), deliveries[111]); err != nil {
		t.Fatalf("mark 111: %v", err)
	}
	waitACKs(t, fake, 104, 130)
}

func TestACKTransientFailureIsRetriedWithoutLosingBoundary(t *testing.T) {
	fake := newFake()
	fake.ackErrors[104] = []error{status.Error(codes.Unavailable, "try again")}
	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()
	}
	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.Fatalf("mark 100: %v", err)
	}
	if err := sub.MarkProcessed(context.Background(), d104); err != nil {
		t.Fatalf("mark 104: %v", err)
	}
	waitACKs(t, fake, 104, 104)
}

func TestACKResponseLostAmbiguityReconnectUsesServerCursor(t *testing.T) {
	fake := newFake()
	dropFirst := make(chan struct{})
	var dropOnce 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))
			_ = stream.Send(msg(104))
			<-dropFirst
			return status.Error(codes.Unavailable, "ack response was lost with the stream")
		}
		dm := msg(111)
		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.GetLastAckSequence() == 104 {
			// Simulate "server applied ACK, client lost response"; the next
			// stream script starts after that durable cursor and does not
			// replay 100/104.
			dropOnce.Do(func() { close(dropFirst) })
			return nil, status.Error(codes.Unavailable, "response lost")
		}
		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"})
	d100, d104 := recvDelivery(t, sub), recvDelivery(t, sub)
	if err := sub.MarkProcessed(context.Background(), d104); err != nil {
		t.Fatal(err)
	}
	if err := sub.MarkProcessed(context.Background(), d100); err != nil {
		t.Fatal(err)
	}
	d111 := recvDelivery(t, sub)
	if d111.Message.GetSequence() != 111 {
		t.Fatalf("reconnect must follow server cursor, got seq %d", d111.Message.GetSequence())
	}
	if err := sub.MarkProcessed(context.Background(), d111); err != nil {
		t.Fatal(err)
	}
	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
	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		calls++
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		if calls == 1 {
			_ = stream.Send(msg(100))
			_ = stream.Send(msg(104))
			return status.Error(codes.Unavailable, "drop")
		}
		dm := msg(104)
		dm.Replay = true
		dm.SubscriptionId = "sub-2"
		_ = stream.Send(dm)
		<-stream.Context().Done()
		return stream.Context().Err()
	}
	c := newClientForTest(t, fake)
	sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
	d100, old104 := recvDelivery(t, sub), recvDelivery(t, sub)
	if err := sub.MarkProcessed(context.Background(), d100); err != nil {
		t.Fatalf("mark 100: %v", err)
	}
	waitACKs(t, fake, 100)
	new104 := recvDelivery(t, sub)
	if !new104.Replay {
		t.Fatal("reconnected delivery must expose replay=true")
	}
	if err := sub.MarkProcessed(context.Background(), old104); !errors.Is(err, ErrStaleDelivery) {
		t.Fatalf("old generation MarkProcessed = %v", err)
	}
	if err := sub.MarkProcessed(context.Background(), new104); err != nil {
		t.Fatal(err)
	}
	waitACKs(t, fake, 100, 104)
}

func TestStreamReplacementAbortedIsTerminalNoReconnect(t *testing.T) {
	fake := newFake()
	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		return status.Error(codes.Aborted, "subscription replaced")
	}
	c := newClientForTest(t, fake)
	sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
	if err != nil {
		t.Fatal(err)
	}
	waitDone(t, sub)
	if !errors.Is(sub.Err(), ErrSubscriptionReplaced) {
		t.Fatalf("Err = %v", sub.Err())
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	if fake.subscribeCount != 1 {
		t.Fatalf("replacement triggered reconnect storm: %d subscribes", fake.subscribeCount)
	}
}

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) {
			fake := newFake()
			fake.subscribeFunc = func(_ *pb.SubscribeRequest, _ pb.ChatService_SubscribeServer) error {
				return status.Error(code, "terminal")
			}
			c := newClientForTest(t, fake)
			if _, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"}); status.Code(err) != code {
				t.Fatalf("Subscribe error = %v", err)
			}
			fake.mu.Lock()
			defer fake.mu.Unlock()
			if fake.subscribeCount != 1 {
				t.Fatalf("terminal auth reconnects: %d", fake.subscribeCount)
			}
		})
	}
}

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
	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		calls++
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		if calls == 1 {
			return status.Error(codes.Unavailable, "transient")
		}
		_ = stream.Send(msg(200))
		<-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)
	}
	d := recvDelivery(t, sub)
	if d.Message.GetSequence() != 200 {
		t.Fatalf("delivery seq = %d", d.Message.GetSequence())
	}
}

func TestBoundedBufferShutdownDoesNotDropOrPanic(t *testing.T) {
	fake := newFake()
	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		for _, seq := range []int64{1, 2, 3} {
			_ = stream.Send(msg(seq))
		}
		<-stream.Context().Done()
		return stream.Context().Err()
	}
	c := newClientForTest(t, fake, WithSubscriptionBufferSize(1))
	sub, err := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
	if err != nil {
		t.Fatal(err)
	}
	if err := sub.Close(); err != nil {
		t.Fatal(err)
	}
}

func TestMultipleTopicSubscriptionsAreIndependent(t *testing.T) {
	fake := newFake()
	fake.subscribeFunc = func(req *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		dm := msg(1)
		dm.Message.TopicId = req.GetTopicId()
		dm.SubscriptionId = "sub-" + req.GetTopicId()
		_ = stream.Send(dm)
		if req.GetTopicId() == "general" {
			return status.Error(codes.Unavailable, "one topic failed")
		}
		<-stream.Context().Done()
		return stream.Context().Err()
	}
	c := newClientForTest(t, fake)
	general, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "general"})
	security, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "security"})
	if recvDelivery(t, general).Message.GetTopicId() != "general" {
		t.Fatal("general delivery mismatch")
	}
	ds := recvDelivery(t, security)
	if ds.Message.GetTopicId() != "security" {
		t.Fatal("security delivery mismatch")
	}
	if err := security.MarkProcessed(context.Background(), ds); err != nil {
		t.Fatal(err)
	}
	waitACKs(t, fake, 1)
}

func TestCloseRacesWithReconnectRecvAckAndMarkProcessed(t *testing.T) {
	fake := newFake()
	ackEntered := make(chan struct{})
	releaseACK := make(chan struct{})
	fake.subscribeFunc = func(_ *pb.SubscribeRequest, stream pb.ChatService_SubscribeServer) error {
		_ = stream.Send(&pb.DeliveryMessage{Ready: true})
		_ = stream.Send(msg(1))
		<-stream.Context().Done()
		return stream.Context().Err()
	}
	fake.ackErrors[1] = []error{status.Error(codes.Unavailable, "retry")}
	c := newClientForTest(t, fake, withSleeperForTest(func(ctx context.Context, d time.Duration) error {
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
			return nil
		}
	}))
	origAck := fake.Acknowledge
	_ = origAck
	sub, _ := c.Subscribe(context.Background(), SubscribeRequest{TopicID: "topic"})
	d := recvDelivery(t, sub)
	go func() { _ = sub.MarkProcessed(context.Background(), d) }()
	close(ackEntered)
	<-ackEntered
	close(releaseACK)
	go func() { _ = sub.MarkProcessed(context.Background(), d) }()
	if err := sub.Close(); err != nil {
		t.Fatal(err)
	}
}

func TestHistoryPagerUsesReturnedSequenceAcrossGaps(t *testing.T) {
	fake := newFake()
	fake.historyPages = []*pb.HistoryResponse{
		{Messages: []*pb.Message{{Sequence: 100}, {Sequence: 104}, {Sequence: 111}}, HasMore: true},
		{Messages: []*pb.Message{{Sequence: 118}}, HasMore: false},
	}
	c := newClientForTest(t, fake)
	pager, err := c.NewHistoryPager(HistoryRequest{TopicID: "topic", AfterSequence: 99, Limit: 3})
	if err != nil {
		t.Fatal(err)
	}
	first, err := pager.Next(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if first.Messages[len(first.Messages)-1].GetSequence() != 111 {
		t.Fatal("first page setup broken")
	}
	second, err := pager.Next(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if len(second.Messages) != 1 || second.Messages[0].GetSequence() != 118 {
		t.Fatalf("second page = %+v", second.Messages)
	}
	fake.mu.Lock()
	defer fake.mu.Unlock()
	if len(fake.historyReqs) != 2 || fake.historyReqs[1].GetAfterSequence() != 111 {
		t.Fatalf("history reqs = %+v", fake.historyReqs)
	}
}