Repository: nats-io/nats.go

Issue:
Confusing doc for MessageContext.Next
https://github.com/nats-io/nats.go/blob/3d824b07c65b981268de904c18cb44a0ed4dc303/jetstream/pull.go#L35-L38

The doc for `MessageContext.Next()` is a bit confusing. In "If the context is canceled" I originally thought that meant the context.Context used to create the underlying `Consumer` because I didn't see any "cancel" associated with `MessageContext`. Now I realize it's talking about calling `Stop()` or `Drain()` on `MessageContext`. 


Hints from issue comments:
@piotrpio: Hello @WillAbides, thanks for submitting the issue. I agree, this was confusing - I created [a PR](https://github.com/nats-io/nats.go/pull/1951) improving and clarifying the doc - please let me know if that looks better to you@WillAbides: That looks better to me. I'm pretty sure it would have saved me a bit of confusion (though at this point it's impossible to reread without knowing what I learned)

Solution patch:
diff --git a/jetstream/object.go b/jetstream/object.go
index 3ccd397..ccc6f86 100644
--- a/jetstream/object.go
+++ b/jetstream/object.go
@@ -471,6 +471,7 @@ type (
 		r      io.ReadCloser
 		err    error
 		ctx    context.Context
+		cancel context.CancelFunc
 		digest hash.Hash
 	}
 )
@@ -867,7 +868,7 @@ func (obs *obs) Get(ctx context.Context, name string, opts ...GetObjectOpt) (Obj
 		return lobs.Get(ctx, info.ObjectMeta.Opts.Link.Name)
 	}
 
-	result := &objResult{info: info, ctx: ctx}
+	result := &objResult{info: info, ctx: ctx, cancel: cancel}
 	if info.Size == 0 {
 		return result, nil
 	}
@@ -934,15 +935,10 @@ func (obs *obs) Get(ctx context.Context, name string, opts ...GetObjectOpt) (Obj
 		nats.Context(ctx),
 		nats.BindStream(streamName),
 	}
-	sub, err := obs.pushJS.Subscribe(chunkSubj, processChunk, subscribeOpts...)
+	_, err = obs.pushJS.Subscribe(chunkSubj, processChunk, subscribeOpts...)
 	if err != nil {
 		return nil, err
 	}
-	sub.SetClosedHandler(func(subject string) {
-		if cancel != nil {
-			cancel()
-		}
-	})
 
 	return result, nil
 }
@@ -1500,6 +1496,9 @@ func (o *objResult) Read(p []byte) (n int, err error) {
 func (o *objResult) Close() error {
 	o.Lock()
 	defer o.Unlock()
+	if o.cancel != nil {
+		o.cancel()
+	}
 	if o.r == nil {
 		return nil
 	}
diff --git a/jetstream/pull.go b/jetstream/pull.go
index d2b2778..b65c012 100644
--- a/jetstream/pull.go
+++ b/jetstream/pull.go
@@ -33,12 +33,11 @@ type (
 	// MessagesContext supports iterating over a messages on a stream.
 	// It is returned by [Consumer.Messages] method.
 	MessagesContext interface {
-		// Next retrieves next message on a stream. It will block until the next
-		// message is available. If the context is canceled, Next will return
-		// ErrMsgIteratorClosed error. An optional timeout or context can be
-		// provided using NextOpt options. If none are provided, Next will block
-		// indefinitely until a message is available, iterator is closed or a
-		// heartbeat error occurs.
+		// Next retrieves next message on a stream. If MessagesContext is closed
+		// (either stopped or drained), Next will return ErrMsgIteratorClosed
+		// error. An optional timeout or context can be provided using NextOpt
+		// options. If none are provided, Next will block indefinitely until a
+		// message is available, iterator is closed or a heartbeat error occurs.
 		Next(opts ...NextOpt) (Msg, error)
 
 		// Stop unsubscribes from the stream and cancels subscription. Calling
@@ -137,6 +136,7 @@ type (
 		consumer          *pullConsumer
 		subscription      *nats.Subscription
 		msgs              chan *nats.Msg
+		msgsClosed        atomic.Uint32
 		errs              chan error
 		pending           pendingMsgs
 		hbMonitor         *hbMonitor
@@ -552,7 +552,7 @@ func (p *pullConsumer) Messages(opts ...PullMessagesOpt) (MessagesContext, error
 				// in Next
 				p.subs.Delete(sid)
 			}
-			close(msgs)
+			sub.closeMsgs()
 		}
 	}(sub.id))
 
@@ -588,9 +588,11 @@ var (
 	errDisconnected = errors.New("disconnected")
 )
 
-// Next retrieves next message on a stream. It will block until the next
-// message is available. If the context is canceled, Next will return
-// ErrMsgIteratorClosed error.
+// Next retrieves next message on a stream. If MessagesContext is closed
+// (either stopped or drained), Next will return ErrMsgIteratorClosed
+// error. An optional timeout or context can be provided using NextOpt
+// options. If none are provided, Next will block indefinitely until a
+// message is available, iterator is closed or a heartbeat error occurs.
 func (s *pullSubscription) Next(opts ...NextOpt) (Msg, error) {
 	var nextOpts nextOpts
 	for _, opt := range opts {
@@ -1057,6 +1059,12 @@ func (s *pullSubscription) pullMessages(subject string) {
 	}
 }
 
+func (s *pullSubscription) closeMsgs() {
+	if s.msgsClosed.CompareAndSwap(0, 1) {
+		close(s.msgs)
+	}
+}
+
 func (s *pullSubscription) scheduleHeartbeatCheck(dur time.Duration) *hbMonitor {
 	if dur == 0 {
 		return nil
diff --git a/nats.go b/nats.go
index 7e8c11b..cba133a 100644
--- a/nats.go
+++ b/nats.go
@@ -151,6 +151,7 @@ var (
 	ErrMaxAccountConnectionsExceeded = errors.New("nats: maximum account active connections exceeded")
 	ErrConnectionNotTLS              = errors.New("nats: connection is not tls")
 	ErrMaxSubscriptionsExceeded      = errors.New("nats: server maximum subscriptions exceeded")
+	ErrWebSocketHeadersAlreadySet    = errors.New("nats: websocket connection headers already set")
 )
 
 // GetDefaultOptions returns default configuration options for the client.
@@ -250,6 +251,9 @@ type UserInfoCB func() (string, string)
 // whole list of URLs and failed to reconnect.
 type ReconnectDelayHandler func(attempts int) time.Duration
 
+// WebSocketHeadersHandler is an optional callback handler for generating token used for WebSocket connections.
+type WebSocketHeadersHandler func() (http.Header, error)
+
 // asyncCB is used to preserve order for async callbacks.
 type asyncCB struct {
 	f    func()
@@ -524,6 +528,12 @@ type Options struct {
 	// from SubscribeSync if the server returns a permissions error for a subscription.
 	// Defaults to false.
 	PermissionErrOnSubscribe bool
+
+	// WebSocketConnectionHeaders is an optional http request headers to be sent with the WebSocket request.
+	WebSocketConnectionHeaders http.Header
+
+	// WebSocketConnectionHeadersHandler is an optional callback handler for generating token  used for WebSocket connections.
+	WebSocketConnectionHeadersHandler WebSocketHeadersHandler
 }
 
 const (
@@ -1472,6 +1482,36 @@ func TLSHandshakeFirst() Option {
 	}
 }
 
+// WebSocketConnectionHeaders sets a fixed set of HTTP headers that will be
+// sent during the WebSocket connection handshake.
+// This option is mutually exclusive with WebSocketConnectionHeadersHandler;
+// if a headers handler has already been configured, it returns
+// ErrWebSocketHeadersAlreadySet.
+func WebSocketConnectionHeaders(headers http.Header) Option {
+	return func(o *Options) error {
+		if o.WebSocketConnectionHeadersHandler != nil {
+			return ErrWebSocketHeadersAlreadySet
+		}
+		o.WebSocketConnectionHeaders = headers
+		return nil
+	}
+}
+
+// WebSocketConnectionHeadersHandler registers a callback used to supply HTTP
+// headers for the WebSocket connection handshake.
+// This option is mutually exclusive with WebSocketConnectionHeaders; if
+// non-empty static headers have already been configured, it returns
+// ErrWebSocketHeadersAlreadySet.
+func WebSocketConnectionHeadersHandler(cb WebSocketHeadersHandler) Option {
+	return func(o *Options) error {
+		if len(o.WebSocketConnectionHeaders) != 0 {
+			return ErrWebSocketHeadersAlreadySet
+		}
+		o.WebSocketConnectionHeadersHandler = cb
+		return nil
+	}
+}
+
 // Handler processing
 
 // SetDisconnectHandler will set the disconnect event handler.
@@ -1671,8 +1711,10 @@ func (o Options) Connect() (*Conn, error) {
 		return nil, err
 	}
 
-	if connectionEstablished && nc.Opts.ConnectedCB != nil {
-		nc.ach.push(func() { nc.Opts.ConnectedCB(nc) })
+	if connectionEstablished {
+		if connectedCB := nc.Opts.ConnectedCB; connectedCB != nil {
+			nc.ach.push(func() { connectedCB(nc) })
+		}
 	}
 
 	return nc, nil
@@ -2747,8 +2789,10 @@ func (nc *Conn) sendConnect() error {
 	// Construct the CONNECT protocol string
 	cProto, err := nc.connectProto()
 	if err != nil {
-		if !nc.initc && nc.Opts.AsyncErrorCB != nil {
-			nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, nil, err) })
+		if !nc.initc {
+			if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+				nc.ach.push(func() { asyncErrorCB(nc, nil, err) })
+			}
 		}
 		return err
 	}
@@ -2764,8 +2808,10 @@ func (nc *Conn) sendConnect() error {
 	// reading byte-by-byte here is ok.
 	proto, err := nc.readProto()
 	if err != nil {
-		if !nc.initc && nc.Opts.AsyncErrorCB != nil {
-			nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, nil, err) })
+		if !nc.initc {
+			if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+				nc.ach.push(func() { asyncErrorCB(nc, nil, err) })
+			}
 		}
 		return err
 	}
@@ -2775,8 +2821,10 @@ func (nc *Conn) sendConnect() error {
 		// Read the rest now...
 		proto, err = nc.readProto()
 		if err != nil {
-			if !nc.initc && nc.Opts.AsyncErrorCB != nil {
-				nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, nil, err) })
+			if !nc.initc {
+				if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+					nc.ach.push(func() { asyncErrorCB(nc, nil, err) })
+				}
 			}
 			return err
 		}
@@ -2884,10 +2932,10 @@ func (nc *Conn) doReconnect(err error, forceReconnect bool) {
 	// Perform appropriate callback if needed for a disconnect.
 	// DisconnectedErrCB has priority over deprecated DisconnectedCB
 	if !nc.initc {
-		if nc.Opts.DisconnectedErrCB != nil {
-			nc.ach.push(func() { nc.Opts.DisconnectedErrCB(nc, err) })
-		} else if nc.Opts.DisconnectedCB != nil {
-			nc.ach.push(func() { nc.Opts.DisconnectedCB(nc) })
+		if disconnectedErrCB := nc.Opts.DisconnectedErrCB; disconnectedErrCB != nil {
+			nc.ach.push(func() { disconnectedErrCB(nc, err) })
+		} else if disconnectedCB := nc.Opts.DisconnectedCB; disconnectedCB != nil {
+			nc.ach.push(func() { disconnectedCB(nc) })
 		}
 	} else if nc.Opts.RetryOnFailedConnect && nc.initc && err != nil {
 		// For initial connection failure with RetryOnFailedConnect,
@@ -2996,8 +3044,8 @@ func (nc *Conn) doReconnect(err error, forceReconnect bool) {
 		// Continue to hold the lock
 		if err != nil {
 			// Perform appropriate callback for a failed connection attempt.
-			if nc.Opts.ReconnectErrCB != nil {
-				nc.ach.push(func() { nc.Opts.ReconnectErrCB(nc, err) })
+			if reconnectErrCB := nc.Opts.ReconnectErrCB; reconnectErrCB != nil {
+				nc.ach.push(func() { reconnectErrCB(nc, err) })
 			}
 			nc.err = nil
 			continue
@@ -3047,10 +3095,10 @@ func (nc *Conn) doReconnect(err error, forceReconnect bool) {
 		// Queue up the correct callback. If we are in initial connect state
 		// (using retry on failed connect), we will call the ConnectedCB,
 		// otherwise the ReconnectedCB.
-		if nc.Opts.ReconnectedCB != nil && !nc.initc {
-			nc.ach.push(func() { nc.Opts.ReconnectedCB(nc) })
-		} else if nc.Opts.ConnectedCB != nil && nc.initc {
-			nc.ach.push(func() { nc.Opts.ConnectedCB(nc) })
+		if reconnectedCB := nc.Opts.ReconnectedCB; reconnectedCB != nil && !nc.initc {
+			nc.ach.push(func() { reconnectedCB(nc) })
+		} else if connectedCB := nc.Opts.ConnectedCB; connectedCB != nil && nc.initc {
+			nc.ach.push(func() { connectedCB(nc) })
 		}
 
 		// If we are here with a retry on failed connect, indicate that the
@@ -3364,8 +3412,8 @@ func (nc *Conn) processMsg(data []byte) {
 			// We will pass the message through but send async error.
 			nc.mu.Lock()
 			nc.err = ErrBadHeaderMsg
-			if nc.Opts.AsyncErrorCB != nil {
-				nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, sub, ErrBadHeaderMsg) })
+			if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+				nc.ach.push(func() { asyncErrorCB(nc, sub, ErrBadHeaderMsg) })
 			}
 			nc.mu.Unlock()
 		}
@@ -3542,8 +3590,8 @@ slowConsumer:
 		// is already experiencing client-side slow consumer situation.
 		nc.mu.Lock()
 		nc.err = ErrSlowConsumer
-		if nc.Opts.AsyncErrorCB != nil {
-			nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, sub, ErrSlowConsumer) })
+		if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+			nc.ach.push(func() { asyncErrorCB(nc, sub, ErrSlowConsumer) })
 		}
 		nc.mu.Unlock()
 	} else {
@@ -3586,8 +3634,8 @@ func (nc *Conn) processTransientError(err error) {
 			}
 		}
 	}
-	if nc.Opts.AsyncErrorCB != nil {
-		nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, nil, err) })
+	if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+		nc.ach.push(func() { asyncErrorCB(nc, nil, err) })
 	}
 	nc.mu.Unlock()
 }
@@ -3599,8 +3647,10 @@ func (nc *Conn) processTransientError(err error) {
 // Connection lock is held on entry
 func (nc *Conn) processAuthError(err error) bool {
 	nc.err = err
-	if !nc.initc && nc.Opts.AsyncErrorCB != nil {
-		nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, nil, err) })
+	if !nc.initc {
+		if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+			nc.ach.push(func() { asyncErrorCB(nc, nil, err) })
+		}
 	}
 	// We should give up if we tried twice on this server and got the
 	// same error. This behavior can be modified using IgnoreAuthErrorAbort.
@@ -3645,8 +3695,8 @@ func (nc *Conn) flusher() {
 				if nc.err == nil {
 					nc.err = err
 				}
-				if nc.Opts.AsyncErrorCB != nil {
-					nc.ach.push(func() { nc.Opts.AsyncErrorCB(nc, nil, err) })
+				if asyncErrorCB := nc.Opts.AsyncErrorCB; asyncErrorCB != nil {
+					nc.ach.push(func() { asyncErrorCB(nc, nil, err) })
 				}
 			}
 		}
@@ -3760,12 +3810,16 @@ func (nc *Conn) processInfo(info string) error {
 		if !nc.Opts.NoRandomize {
 			nc.shufflePool(1)
 		}
-		if !nc.initc && nc.Opts.DiscoveredServersCB != nil {
-			nc.ach.push(func() { nc.Opts.DiscoveredServersCB(nc) })
+		if !nc.initc {
+			if discoveredServersCB := nc.Opts.DiscoveredServersCB; discoveredServersCB != nil {
+				nc.ach.push(func() { discoveredServersCB(nc) })
+			}
 		}
 	}
-	if !nc.initc && ncInfo.LameDuckMode && nc.Opts.LameDuckModeHandler != nil {
-		nc.ach.push(func() { nc.Opts.LameDuckModeHandler(nc) })
+	if !nc.initc && ncInfo.LameDuckMode {
+		if lameDuckModeHandler := nc.Opts.LameDuckModeHandler; lameDuckModeHandler != nil {
+			nc.ach.push(func() { lameDuckModeHandler(nc) })
+		}
 	}
 	return nil
 }
@@ -5607,8 +5661,8 @@ func (nc *Conn) close(status Status, doCBs bool, err error) {
 				nc.ach.push(func() { disconnectedCB(nc) })
 			}
 		}
-		if nc.Opts.ClosedCB != nil {
-			nc.ach.push(func() { nc.Opts.ClosedCB(nc) })
+		if closedCB := nc.Opts.ClosedCB; closedCB != nil {
+			nc.ach.push(func() { closedCB(nc) })
 		}
 	}
 	// If this is terminal, then we have to notify the asyncCB handler that
diff --git a/ws.go b/ws.go
index 46d8300..3642917 100644
--- a/ws.go
+++ b/ws.go
@@ -610,6 +610,9 @@ func (nc *Conn) wsInitHandshake(u *url.URL) error {
 	if compress {
 		req.Header.Add("Sec-WebSocket-Extensions", wsPMCReqHeaderValue)
 	}
+	if err := nc.wsUpdateConnectionHeaders(req); err != nil {
+		return err
+	}
 	if err := req.Write(nc.conn); err != nil {
 		return err
 	}
@@ -728,6 +731,25 @@ func (nc *Conn) wsEnqueueControlMsg(needsLock bool, frameType wsOpCode, payload
 	nc.bw.flush()
 }
 
+func (nc *Conn) wsUpdateConnectionHeaders(req *http.Request) error {
+	var headers http.Header
+	var err error
+	if nc.Opts.WebSocketConnectionHeadersHandler != nil {
+		headers, err = nc.Opts.WebSocketConnectionHeadersHandler()
+		if err != nil {
+			return err
+		}
+	} else {
+		headers = nc.Opts.WebSocketConnectionHeaders
+	}
+	for key, values := range headers {
+		for _, val := range values {
+			req.Header.Add(key, val)
+		}
+	}
+	return nil
+}
+
 func wsPMCExtensionSupport(header http.Header) (bool, bool) {
 	for _, extensionList := range header["Sec-Websocket-Extensions"] {
 		extensions := strings.Split(extensionList, ",")



Test patch:
diff --git a/jetstream/test/consumer_test.go b/jetstream/test/consumer_test.go
index 9449eec..6411382 100644
--- a/jetstream/test/consumer_test.go
+++ b/jetstream/test/consumer_test.go
@@ -1522,3 +1522,79 @@ func TestConsumerCachedInfo(t *testing.T) {
 	}
 
 }
+
+func TestJetStreamPauseAndResumeConsumer(t *testing.T) {
+	srv := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, srv)
+
+	nc, err := nats.Connect(srv.ClientURL())
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer nc.Close()
+
+	ctx := context.Background()
+	js, err := jetstream.New(nc)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	_, err = js.CreateStream(ctx, jetstream.StreamConfig{
+		Name:     "TEST",
+		Subjects: []string{"foo"},
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	_, err = js.CreateConsumer(ctx, "TEST", jetstream.ConsumerConfig{
+		Durable:   "cons",
+		AckPolicy: jetstream.AckExplicitPolicy,
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	t.Run("pause consumer", func(t *testing.T) {
+		pauseUntil := time.Now().Add(1 * time.Minute)
+		resp, err := js.PauseConsumer(ctx, "TEST", "cons", pauseUntil)
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if !resp.Paused {
+			t.Fatal("Consumer should be paused")
+		}
+		if resp.PauseRemaining <= 0 {
+			t.Fatal("PauseRemaining should be greater than 0")
+		}
+	})
+
+	t.Run("resume consumer", func(t *testing.T) {
+		resp, err := js.ResumeConsumer(ctx, "TEST", "cons")
+		if err != nil {
+			t.Fatalf("Unexpected error: %v", err)
+		}
+
+		if resp.Paused {
+			t.Fatal("Consumer should not be paused")
+		}
+		if resp.PauseRemaining != 0 {
+			t.Fatal("PauseRemaining should be 0")
+		}
+	})
+
+	t.Run("pause consumer with invalid stream name", func(t *testing.T) {
+		_, err := js.PauseConsumer(ctx, "", "cons", time.Now())
+		if err == nil {
+			t.Fatal("Expected error for empty stream name")
+		}
+	})
+
+	t.Run("resume consumer with invalid stream name", func(t *testing.T) {
+		_, err := js.ResumeConsumer(ctx, "", "cons")
+		if err == nil {
+			t.Fatal("Expected error for empty stream name")
+		}
+	})
+}
diff --git a/jetstream/test/jetstream_test.go b/jetstream/test/jetstream_test.go
index b541d01..69ea7a6 100644
--- a/jetstream/test/jetstream_test.go
+++ b/jetstream/test/jetstream_test.go
@@ -2478,3 +2478,24 @@ func TestPromoteMirrorToStream(t *testing.T) {
 		t.Fatalf("Expected subject to be 'bar', got %q", info.Config.Subjects[0])
 	}
 }
+
+func TestJetStreamConn(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, err := nats.Connect(s.ClientURL())
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer nc.Close()
+
+	js, err := jetstream.New(nc)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	conn := js.Conn()
+	if conn != nc {
+		t.Fatal("Expected Conn() to return the same connection")
+	}
+}
diff --git a/jetstream/test/kv_test.go b/jetstream/test/kv_test.go
index 9b7e1b7..62e5ff8 100644
--- a/jetstream/test/kv_test.go
+++ b/jetstream/test/kv_test.go
@@ -2145,3 +2145,153 @@ func TestKeyValueListKeysDuplicates(t *testing.T) {
 		})
 	}
 }
+
+func TestKeyValueWithSources(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	kv1, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+		Bucket:  "SOURCE1",
+		History: 5,
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	kv2, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+		Bucket:  "SOURCE2",
+		History: 5,
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	kv3, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+		Bucket:  "SOURCED",
+		History: 5,
+		Sources: []*jetstream.StreamSource{
+			{Name: "SOURCE1"},
+			{Name: "SOURCE2"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if _, err := kv1.Put(ctx, "key1", []byte("value1")); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if _, err := kv2.Put(ctx, "key2", []byte("value2")); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	time.Sleep(200 * time.Millisecond)
+
+	val, err := kv3.Get(ctx, "key1")
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if string(val.Value()) != "value1" {
+		t.Fatalf("Expected value1, got %s", string(val.Value()))
+	}
+
+	val, err = kv3.Get(ctx, "key2")
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	if string(val.Value()) != "value2" {
+		t.Fatalf("Expected value2, got %s", string(val.Value()))
+	}
+
+	stream, err := js.Stream(ctx, "KV_SOURCED")
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	info, err := stream.Info(ctx)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if len(info.Config.Sources) != 2 {
+		t.Fatalf("Expected 2 sources, got %d", len(info.Config.Sources))
+	}
+
+	for _, src := range info.Config.Sources {
+		if len(src.SubjectTransforms) != 1 {
+			t.Fatalf("Expected 1 subject transform, got %d", len(src.SubjectTransforms))
+		}
+	}
+}
+
+func TestKeyValueGetRevision(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	kv, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+		Bucket:  "TEST",
+		History: 5,
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if _, err := kv.Put(ctx, "key", []byte("value1")); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if _, err := kv.Put(ctx, "key", []byte("value2")); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	entry, err := kv.GetRevision(ctx, "key", 1)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if string(entry.Value()) != "value1" {
+		t.Fatalf("Expected value1, got %s", string(entry.Value()))
+	}
+
+	if entry.Revision() != 1 {
+		t.Fatalf("Expected revision 1, got %d", entry.Revision())
+	}
+}
+
+func TestKeyValueStatusBytes(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	kv, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{
+		Bucket: "TEST",
+	})
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if _, err := kv.Put(ctx, "key", []byte("value")); err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	status, err := kv.Status(ctx)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if status.Bytes() == 0 {
+		t.Fatal("Expected Bytes() to return non-zero value")
+	}
+}
diff --git a/jetstream/test/object_test.go b/jetstream/test/object_test.go
index d7d3100..30ccc85 100644
--- a/jetstream/test/object_test.go
+++ b/jetstream/test/object_test.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"crypto/rand"
 	"crypto/sha256"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -1298,3 +1299,20 @@ func TestObjectStoreMirror(t *testing.T) {
 		}
 	}
 }
+
+func TestObjectStoreDeleteNonExistent(t *testing.T) {
+	s := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, s)
+
+	nc, js := jsClient(t, s)
+	defer nc.Close()
+	ctx := context.Background()
+
+	err := js.DeleteObjectStore(ctx, "NONEXISTENT")
+	if err == nil {
+		t.Fatal("Expected error when deleting non-existent object store")
+	}
+	if !errors.Is(err, jetstream.ErrBucketNotFound) {
+		t.Fatalf("Expected ErrBucketNotFound, got: %v", err)
+	}
+}
diff --git a/jetstream/test/pull_test.go b/jetstream/test/pull_test.go
index d4fb8b9..d688d4d 100644
--- a/jetstream/test/pull_test.go
+++ b/jetstream/test/pull_test.go
@@ -554,6 +554,70 @@ func TestPullConsumerFetch(t *testing.T) {
 	})
 }
 
+func TestPullConsumerMessagesConcurrentStopAndDrain(t *testing.T) {
+	srv := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, srv)
+
+	nc, err := nats.Connect(srv.ClientURL())
+	if err != nil {
+		t.Fatalf("connect: %v", err)
+	}
+	t.Cleanup(nc.Close)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	js, err := jetstream.New(nc)
+	if err != nil {
+		t.Fatalf("jetstream: %v", err)
+	}
+
+	_, err = js.CreateStream(ctx, jetstream.StreamConfig{
+		Name:     "FOO",
+		Subjects: []string{"FOO.>"},
+	})
+	if err != nil {
+		t.Fatalf("create stream: %v", err)
+	}
+
+	consumer, err := js.CreateOrUpdateConsumer(ctx, "FOO", jetstream.ConsumerConfig{
+		Durable:   "race-consumer",
+		AckPolicy: jetstream.AckExplicitPolicy,
+	})
+	if err != nil {
+		t.Fatalf("create consumer: %v", err)
+	}
+
+	messages, err := consumer.Messages()
+	if err != nil {
+		t.Fatalf("messages: %v", err)
+	}
+
+	start := make(chan struct{})
+	var ready, done sync.WaitGroup
+	ready.Add(2)
+	done.Add(2)
+
+	go func() {
+		defer done.Done()
+		ready.Done()
+		<-start
+		messages.Stop()
+	}()
+
+	go func() {
+		defer done.Done()
+		ready.Done()
+		<-start
+		nc.Drain()
+	}()
+
+	ready.Wait()
+	time.Sleep(2 * time.Millisecond)
+	close(start)
+	done.Wait()
+}
+
 func TestPullConsumerFetchRace(t *testing.T) {
 	srv := RunBasicJetStreamServer()
 	defer shutdownJSServerAndRemoveStorage(t, srv)
diff --git a/nats_test.go b/nats_test.go
index d481948..a0d1215 100644
--- a/nats_test.go
+++ b/nats_test.go
@@ -237,11 +237,11 @@ func TestSimplifiedURLs(t *testing.T) {
 				"host4:1234",
 				"host5:",
 				"host6",
-				"nats://[1:2:3:4]:1234",
-				"nats://[5:6:7:8]:",
-				"nats://[9:10:11:12]",
-				"[13:14:15:16]:",
-				"[17:18:19:20]:1234",
+				"nats://[::1:2:3:4]:1234",
+				"nats://[::5:6:7:8]:",
+				"nats://[::9:10:11:12]",
+				"[::13:14:15:16]:",
+				"[::17:18:19:20]:1234",
 			},
 			[]string{
 				"nats://host1:1234/",
@@ -251,11 +251,11 @@ func TestSimplifiedURLs(t *testing.T) {
 				"nats://host4:1234",
 				"nats://host5:4222",
 				"nats://host6:4222",
-				"nats://[1:2:3:4]:1234",
-				"nats://[5:6:7:8]:4222",
-				"nats://[9:10:11:12]:4222",
-				"nats://[13:14:15:16]:4222",
-				"nats://[17:18:19:20]:1234",
+				"nats://[::1:2:3:4]:1234",
+				"nats://[::5:6:7:8]:4222",
+				"nats://[::9:10:11:12]:4222",
+				"nats://[::13:14:15:16]:4222",
+				"nats://[::17:18:19:20]:1234",
 			},
 		},
 		{
@@ -264,17 +264,17 @@ func TestSimplifiedURLs(t *testing.T) {
 				"ws://host1:1234",
 				"ws://host2:",
 				"ws://host3",
-				"ws://[1:2:3:4]:1234",
-				"ws://[5:6:7:8]:",
-				"ws://[9:10:11:12]",
+				"ws://[::1:2:3:4]:1234",
+				"ws://[::5:6:7:8]:",
+				"ws://[::9:10:11:12]",
 			},
 			[]string{
 				"ws://host1:1234",
 				"ws://host2:80",
 				"ws://host3:80",
-				"ws://[1:2:3:4]:1234",
-				"ws://[5:6:7:8]:80",
-				"ws://[9:10:11:12]:80",
+				"ws://[::1:2:3:4]:1234",
+				"ws://[::5:6:7:8]:80",
+				"ws://[::9:10:11:12]:80",
 			},
 		},
 		{
@@ -283,17 +283,17 @@ func TestSimplifiedURLs(t *testing.T) {
 				"wss://host1:1234",
 				"wss://host2:",
 				"wss://host3",
-				"wss://[1:2:3:4]:1234",
-				"wss://[5:6:7:8]:",
-				"wss://[9:10:11:12]",
+				"wss://[::1:2:3:4]:1234",
+				"wss://[::5:6:7:8]:",
+				"wss://[::9:10:11:12]",
 			},
 			[]string{
 				"wss://host1:1234",
 				"wss://host2:443",
 				"wss://host3:443",
-				"wss://[1:2:3:4]:1234",
-				"wss://[5:6:7:8]:443",
-				"wss://[9:10:11:12]:443",
+				"wss://[::1:2:3:4]:1234",
+				"wss://[::5:6:7:8]:443",
+				"wss://[::9:10:11:12]:443",
 			},
 		},
 	} {
diff --git a/test/conn_test.go b/test/conn_test.go
index 7a2d02e..6b3a93f 100644
--- a/test/conn_test.go
+++ b/test/conn_test.go
@@ -299,7 +299,13 @@ func TestClientTLSConfig(t *testing.T) {
 	// Should fail because of missing certificate
 	_, err = nats.Connect(secureURL,
 		nats.ClientTLSConfig(nil, caCB))
-	if !strings.Contains(err.Error(), "bad certificate") && !strings.Contains(err.Error(), "certificate required") {
+	if err == nil {
+		t.Fatal("Expected connection to fail due to missing certificate")
+	}
+	errStr := err.Error()
+	if !strings.Contains(errStr, "bad certificate") &&
+		!strings.Contains(errStr, "certificate required") &&
+		!strings.Contains(errStr, "connection reset by peer") {
 		t.Fatalf("Expected missing certificate error; got: %s", err)
 	}
 
diff --git a/test/context_test.go b/test/context_test.go
index 0502eab..81c1050 100644
--- a/test/context_test.go
+++ b/test/context_test.go
@@ -731,6 +731,19 @@ func TestFlushWithContext(t *testing.T) {
 	if err := nc.FlushWithContext(dctx); err != context.Canceled {
 		t.Fatalf("Expected '%v', got '%v'", context.Canceled, err)
 	}
+
+	t.Run("flush on closed connection", func(t *testing.T) {
+		nc2 := NewDefaultConnection(t)
+		nc2.Close()
+
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		defer cancel()
+
+		err := nc2.FlushWithContext(ctx)
+		if err != nats.ErrConnectionClosed {
+			t.Fatalf("Expected '%v', got '%v'", nats.ErrConnectionClosed, err)
+		}
+	})
 }
 
 func TestUnsubscribeAndNextMsgWithContext(t *testing.T) {
diff --git a/test/drain_test.go b/test/drain_test.go
index da07c89..f3861f0 100644
--- a/test/drain_test.go
+++ b/test/drain_test.go
@@ -494,3 +494,23 @@ func TestDrainConnDuringReconnect(t *testing.T) {
 		t.Fatalf("Timeout waiting for closed state for connection")
 	}
 }
+
+func TestDrainClosedHandlerRace(t *testing.T) {
+	s := RunDefaultServer()
+	defer s.Shutdown()
+
+	nc, err := nats.Connect(s.ClientURL(), nats.ClosedHandler(func(_ *nats.Conn) {}))
+
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer nc.Close()
+
+	if err := nc.Drain(); err != nil {
+		t.Fatalf("Unexpected error on drain: %v", err)
+	}
+	go func() {
+		nc.SetClosedHandler(nil)
+	}()
+	time.Sleep(500 * time.Millisecond)
+}
diff --git a/test/ws_test.go b/test/ws_test.go
index ee13b33..b0b44ef 100644
--- a/test/ws_test.go
+++ b/test/ws_test.go
@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"math/rand"
 	"net"
+	"net/http"
 	"runtime"
 	"strings"
 	"sync"
@@ -611,3 +612,88 @@ func TestWSNoDeadlockOnAuthFailure(t *testing.T) {
 
 	tm.Stop()
 }
+
+func TestWsWithCustomHeaders(t *testing.T) {
+	sopts := testWSGetDefaultOptions(t, false)
+	s := RunServerWithOptions(sopts)
+	defer s.Shutdown()
+
+	staticHeader := make(http.Header, 0)
+	staticHeader.Set("Authorization", "Bearer Random Token")
+	headerProvider := func() (http.Header, error) {
+		return staticHeader, nil
+	}
+
+	for _, test := range []struct {
+		name              string
+		connectionOptions []nats.Option
+		wantErr           bool
+	}{
+		{
+			name: "Failure 1: Both headers and handler present",
+			connectionOptions: []nats.Option{
+				nats.WebSocketConnectionHeadersHandler(headerProvider),
+				nats.WebSocketConnectionHeaders(staticHeader),
+			},
+			wantErr: true,
+		},
+		{
+			name: "Success 1: Headers present as static headers",
+			connectionOptions: []nats.Option{
+				nats.WebSocketConnectionHeaders(staticHeader),
+			},
+			wantErr: false,
+		},
+		{
+			name: "Success 2: Header supplied through handler",
+			connectionOptions: []nats.Option{
+				nats.WebSocketConnectionHeadersHandler(headerProvider),
+			},
+			wantErr: false,
+		},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			url := fmt.Sprintf("ws://127.0.0.1:%d", sopts.Websocket.Port)
+			nc, err := nats.Connect(url, test.connectionOptions...)
+			if err != nil && test.wantErr {
+				return
+			}
+			if err != nil && !test.wantErr {
+				t.Fatalf("Did not expect error, found error: %v", err)
+			}
+			defer nc.Close()
+			sub, err := nc.SubscribeSync("foo")
+			if err != nil {
+				t.Fatalf("Error on subscribe: %v", err)
+			}
+
+			msgs := make([][]byte, 100)
+			for i := 0; i < len(msgs); i++ {
+				msg := make([]byte, 100)
+				for j := 0; j < len(msg); j++ {
+					msg[j] = 'A'
+				}
+				msgs[i] = msg
+			}
+			for i, msg := range msgs {
+				if err := nc.Publish("foo", msg); err != nil {
+					t.Fatalf("Error on publish: %v", err)
+				}
+				// Make sure that compression/masking does not touch user data
+				if !bytes.Equal(msgs[i], msg) {
+					t.Fatalf("User content has been changed: %v, got %v", msgs[i], msg)
+				}
+			}
+
+			for i := 0; i < len(msgs); i++ {
+				msg, err := sub.NextMsg(time.Second)
+				if err != nil {
+					t.Fatalf("Error getting next message (%d): %v", i+1, err)
+				}
+				if !bytes.Equal(msgs[i], msg.Data) {
+					t.Fatalf("Expected message (%d): %v, got %v", i+1, msgs[i], msg)
+				}
+			}
+		})
+	}
+}
diff --git a/ws_test.go b/ws_test.go
index 2227c82..4c2a15e 100644
--- a/ws_test.go
+++ b/ws_test.go
@@ -14,6 +14,7 @@
 package nats
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"fmt"
@@ -608,3 +609,123 @@ func TestWSProxyPath(t *testing.T) {
 		})
 	}
 }
+
+// --- helpers ---
+
+func startHeaderCatcher(t *testing.T) (addr string, got chan []string, closer func()) {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	got = make(chan []string, 1)
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			// surface nothing; test will timeout
+			return
+		}
+		defer conn.Close()
+		r := bufio.NewReader(conn)
+		var lines []string
+		for {
+			s, err := r.ReadString('\n')
+			if err != nil {
+				break
+			}
+			s = strings.TrimRight(s, "\r\n")
+			if s == "" { // end of HTTP headers
+				break
+			}
+			lines = append(lines, s)
+		}
+		got <- lines
+	}()
+
+	return ln.Addr().String(), got, func() { _ = ln.Close() }
+}
+
+func hasHeaderValue(headers []string, name, want string) bool {
+	prefix := strings.ToLower(name) + ":"
+	for _, h := range headers {
+		if !strings.HasPrefix(strings.ToLower(h), prefix) {
+			continue
+		}
+		val := strings.TrimSpace(strings.SplitN(h, ":", 2)[1])
+		for _, part := range strings.Split(val, ",") {
+			if strings.EqualFold(strings.TrimSpace(part), want) {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func TestWSHeaders_StaticAppliedOnHandshake(t *testing.T) {
+	addr, got, closeLn := startHeaderCatcher(t)
+	defer closeLn()
+
+	static := make(http.Header)
+	static.Set("Authorization", "Bearer Random Token")
+	static.Add("X-Multi", "v1")
+	static.Add("X-Multi", "v2")
+
+	// Intentionally connect to our fake server; it won't complete the upgrade.
+	opts := GetDefaultOptions()
+	opts.WebSocketConnectionHeaders = static
+	opts.Url = "ws://" + addr
+	_, err := opts.Connect()
+	if err == nil {
+		t.Fatalf("expected connect to fail because server does not reply")
+	}
+
+	var headers []string
+	select {
+	case headers = <-got:
+	case <-time.After(2 * time.Second):
+		t.Fatal("did not capture headers in time")
+	}
+
+	if !hasHeaderValue(headers, "Authorization", "Bearer Random Token") {
+		t.Fatalf("Authorization header missing: %v", headers)
+	}
+	if !hasHeaderValue(headers, "X-Multi", "v1") || !hasHeaderValue(headers, "X-Multi", "v2") {
+		t.Fatalf("X-Multi headers missing/combined incorrectly: %v", headers)
+	}
+}
+
+func TestWSHeaders_HandlerAppliedOnHandshake(t *testing.T) {
+	addr, got, closeLn := startHeaderCatcher(t)
+	defer closeLn()
+
+	provider := func() (http.Header, error) {
+		h := make(http.Header)
+		h.Set("Authorization", "Bearer FromHandler")
+		h.Add("X-Multi", "h1")
+		h.Add("X-Multi", "h2")
+		return h, nil
+	}
+
+	opts := GetDefaultOptions()
+	opts.WebSocketConnectionHeadersHandler = provider
+	opts.Url = "ws://" + addr
+	_, err := opts.Connect()
+	if err == nil {
+		t.Fatalf("expected connect to fail because server does not reply")
+	}
+
+	var headers []string
+	select {
+	case headers = <-got:
+	case <-time.After(2 * time.Second):
+		t.Fatal("did not capture headers in time")
+	}
+
+	if !hasHeaderValue(headers, "Authorization", "Bearer FromHandler") {
+		t.Fatalf("Authorization header missing: %v", headers)
+	}
+	if !hasHeaderValue(headers, "X-Multi", "h1") || !hasHeaderValue(headers, "X-Multi", "h2") {
+		t.Fatalf("X-Multi headers missing/combined incorrectly: %v", headers)
+	}
+}