Improve gateway reliability and client e2e coverage

clients/go/mxgateway/client_session_test.go

@@ -77,6 +77,42 @@ func TestStreamEventsAttachesAuthMetadataAndClosesOnCancellation(t *testing.T) {
 	}
 }
 
+func TestEventSubscriptionCloseStopsStream(t *testing.T) {
+	fake := &fakeGatewayServer{
+		streamStarted: make(chan struct{}),
+		streamDone:    make(chan struct{}),
+	}
+	client, cleanup := newBufconnClient(t, fake)
+	defer cleanup()
+	session := NewSessionForID(client, "session-1")
+
+	subscription, err := session.SubscribeEvents(context.Background())
+	if err != nil {
+		t.Fatalf("SubscribeEvents() error = %v", err)
+	}
+	<-fake.streamStarted
+	first := <-subscription.Events()
+	if first.Err != nil {
+		t.Fatalf("first event error = %v", first.Err)
+	}
+
+	subscription.Close()
+
+	select {
+	case <-fake.streamDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("event stream did not stop after subscription close")
+	}
+	select {
+	case _, ok := <-subscription.Events():
+		if ok {
+			t.Fatal("subscription channel remained open after close")
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("subscription channel did not close")
+	}
+}
+
 func TestSessionHelpersBuildCommandsAndExposeRawReply(t *testing.T) {
 	fake := &fakeGatewayServer{
 		invokeReply: &pb.MxCommandReply{
@@ -235,6 +271,7 @@ type fakeGatewayServer struct {
 	openAuth      string
 	streamAuth    string
 	streamStarted chan struct{}
+	streamDone    chan struct{}
 	invokeReply   *pb.MxCommandReply
 	invokeRequest *pb.MxCommandRequest
 }
@@ -277,6 +314,9 @@ func (s *fakeGatewayServer) Invoke(ctx context.Context, req *pb.MxCommandRequest
 
 func (s *fakeGatewayServer) StreamEvents(req *pb.StreamEventsRequest, stream grpc.ServerStreamingServer[pb.MxEvent]) error {
 	s.streamAuth = authorizationFromContext(stream.Context())
+	if s.streamDone != nil {
+		defer close(s.streamDone)
+	}
 	if s.streamStarted != nil {
 		close(s.streamStarted)
 	}

clients/go/mxgateway/session.go

@@ -22,6 +22,30 @@ type EventResult struct {
 	Err   error
 }
 
+// EventSubscription owns a running gateway event stream.
+type EventSubscription struct {
+	results <-chan EventResult
+	cancel  context.CancelFunc
+	done    <-chan struct{}
+	once    sync.Once
+}
+
+// Events returns the stream results channel.
+func (s *EventSubscription) Events() <-chan EventResult {
+	return s.results
+}
+
+// Close cancels the stream and waits for the receive goroutine to stop.
+func (s *EventSubscription) Close() {
+	if s == nil {
+		return
+	}
+	s.once.Do(func() {
+		s.cancel()
+		<-s.done
+	})
+}
+
 // Session represents one gateway-backed MXAccess session.
 type Session struct {
 	client     *Client
@@ -394,34 +418,56 @@ func (s *Session) Events(ctx context.Context) (<-chan EventResult, error) {
 
 // EventsAfter streams ordered session events after the given worker sequence.
 func (s *Session) EventsAfter(ctx context.Context, afterWorkerSequence uint64) (<-chan EventResult, error) {
-	stream, err := s.client.StreamEventsRaw(ctx, &pb.StreamEventsRequest{
+	subscription, err := s.SubscribeEventsAfter(ctx, afterWorkerSequence)
+	if err != nil {
+		return nil, err
+	}
+	return subscription.Events(), nil
+}
+
+// SubscribeEvents starts an owned event subscription.
+func (s *Session) SubscribeEvents(ctx context.Context) (*EventSubscription, error) {
+	return s.SubscribeEventsAfter(ctx, 0)
+}
+
+// SubscribeEventsAfter starts an owned event subscription after the given worker sequence.
+func (s *Session) SubscribeEventsAfter(ctx context.Context, afterWorkerSequence uint64) (*EventSubscription, error) {
+	streamCtx, cancel := context.WithCancel(ctx)
+	stream, err := s.client.StreamEventsRaw(streamCtx, &pb.StreamEventsRequest{
 		SessionId:           s.ID(),
 		AfterWorkerSequence: afterWorkerSequence,
 	})
 	if err != nil {
+		cancel()
 		return nil, err
 	}
 
 	results := make(chan EventResult, 16)
+	done := make(chan struct{})
 	go func() {
 		defer close(results)
+		defer close(done)
 		for {
 			event, err := stream.Recv()
 			if err == nil {
-				if !sendEventResult(ctx, results, EventResult{Event: event}) {
+				if !sendEventResult(streamCtx, results, EventResult{Event: event}) {
 					return
 				}
 				continue
 			}
-			if err == io.EOF || status.Code(err) == codes.Canceled || ctx.Err() != nil {
+			if err == io.EOF || status.Code(err) == codes.Canceled || streamCtx.Err() != nil {
 				return
 			}
-			sendEventResult(ctx, results, EventResult{Err: &GatewayError{Op: "stream events", Err: err}})
+			sendEventResult(streamCtx, results, EventResult{Err: &GatewayError{Op: "stream events", Err: err}})
 			return
 		}
 	}()
 
-	return results, nil
+	return &EventSubscription{
+		results: results,
+		cancel:  cancel,
+		done:    done,
+	}, nil
 }
 
 func ensureBulkSize(name string, length int) error {