fix: Save lock-token from receive_deferred_messages AMQP payload, into ServiceBusReceivedMessage (#42471)

* Save lock-token from `receive_deferred_message` AMQP payload, so it can be used as the message's lock_token property. This allows the deferred message to be correctly renewed or settled (abandoned, completed, deferred again).

* Use `kwargs.get()` instead of `kwargs.pop()` to get the optional received lock token in the constructor

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* test: add regression tests and changelog entry for deferred message lock-token (#42454)

Adds unit tests in test_message.py verifying parse_received_message populates lock_token for deferred messages and leaves it None for peeked messages, plus a CHANGELOG bug-fix entry.

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Pranjal Patel <pranjalpatel@microsoft.com>
Co-authored-by: Pranjal Patel <patelpranzp25@gmail.com>

Author: 4 people

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Bugs Fixed
 
+- Fixed a bug where messages returned by `receive_deferred_messages` had a `lock_token` of `None`, which prevented settling (completing, abandoning, dead-lettering, deferring) or renewing the lock on a deferred message in `PEEK_LOCK` mode. The lock token is now read from the `lock-token` field of the management-link response for deferred messages. ([#42454](https://github.com/Azure/azure-sdk-for-python/issues/42454))
 - Read `com.microsoft:max-message-batch-size` vendor property from the AMQP sender link to correctly limit batch size on Premium large-message entities, where `max-message-size` can be up to 100 MB but the batch limit is 1 MB.
 
 ## 7.14.3 (2025-11-11)

sdk/servicebus/azure-servicebus/azure/servicebus/_common/message.py

@@ -787,6 +787,7 @@ def __init__(
         self._settled = receive_mode == ServiceBusReceiveMode.RECEIVE_AND_DELETE
         self._delivery_tag = self._amqp_transport.get_message_delivery_tag(message, frame)
         self._delivery_id = self._amqp_transport.get_message_delivery_id(message, frame)  # only used by pyamqp
+        self._received_lock_token = kwargs.get("lock_token", None)
         self._received_timestamp_utc = utc_now()
         self._is_deferred_message = kwargs.get("is_deferred_message", False)
         self._is_peeked_message = kwargs.get("is_peeked_message", False)
@@ -1090,6 +1091,10 @@ def lock_token(self) -> Optional[Union[uuid.UUID, str]]:
         if self._settled:
             return None
 
+        if self._received_lock_token:
+            # this will already be a uuid.UUID
+            return self._received_lock_token
+
         if self._delivery_tag:
             return uuid.UUID(bytes_le=self._delivery_tag)
 

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_pyamqp_transport.py

@@ -840,15 +840,21 @@ def parse_received_message( # pylint:disable=docstring-keyword-should-match-keyw
         :keyword ~azure.servicebus.ServiceBusReceiver receiver: Required.
         :keyword bool is_peeked_message: Optional. For peeked messages.
         :keyword bool is_deferred_message: Optional. For deferred messages.
+        :keyword uuid.UUID lock_token: Optional. Lock token, if it is given by the message receiver.
         :keyword ~azure.servicebus.ServiceBusReceiveMode receive_mode: Optional.
         :return: List of service bus received messages.
         :rtype: list[~azure.servicebus.ServiceBusReceivedMessage]
         """
+        is_deferred_message = kwargs.get("is_deferred_message", False)
         parsed = []
         if message.value:
             for m in message.value[b"messages"]:
                 wrapped = decode_payload(memoryview(m[b"message"]))
-                parsed.append(message_type(wrapped, **kwargs))
+                if is_deferred_message and b"lock-token" in m:
+                    lock_token = m[b"lock-token"]
+                else:
+                    lock_token = kwargs.pop("lock_token", None)
+                parsed.append(message_type(wrapped, lock_token=lock_token, **kwargs))
         return parsed
 
     @staticmethod

sdk/servicebus/azure-servicebus/azure/servicebus/_transport/_uamqp_transport.py

@@ -883,14 +883,20 @@ def parse_received_message( # pylint:disable=docstring-keyword-should-match-keyw
             :keyword ~azure.servicebus.ServiceBusReceiver receiver: Required.
             :keyword bool is_peeked_message: Optional. For peeked messages.
             :keyword bool is_deferred_message: Optional. For deferred messages.
+            :keyword uuid.UUID lock_token: Optional. Lock token, if it is given by the message receiver.
             :keyword ~azure.servicebus.ServiceBusReceiveMode receive_mode: Optional.
             :return: List of ServiceBusReceivedMessage.
             :rtype: list[~azure.servicebus.ServiceBusReceivedMessage]
             """
+            is_deferred_message = kwargs.get("is_deferred_message", False)
             parsed = []
             for m in message.get_data()[b"messages"]:
                 wrapped = Message.decode_from_bytes(bytearray(m[b"message"]))
-                parsed.append(message_type(wrapped, **kwargs))
+                if is_deferred_message and b"lock-token" in m:
+                    lock_token = m[b"lock-token"]
+                else:
+                    lock_token = kwargs.pop("lock_token", None)
+                parsed.append(message_type(wrapped, lock_token=lock_token, **kwargs))
             return parsed
 
         @staticmethod

sdk/servicebus/azure-servicebus/tests/test_message.py

@@ -1,5 +1,7 @@
 import os
+import uuid
 from decimal import Decimal
+from unittest import mock
 import pytest
 
 try:
@@ -28,6 +30,7 @@
 from azure.servicebus.amqp import AmqpAnnotatedMessage, AmqpMessageBodyType, AmqpMessageProperties, AmqpMessageHeader
 from azure.servicebus._pyamqp.message import Message
 from azure.servicebus._pyamqp._message_backcompat import LegacyBatchMessage
+from azure.servicebus._pyamqp._encode import encode_payload
 from azure.servicebus._transport._pyamqp_transport import PyamqpTransport
 
 from devtools_testutils import AzureMgmtRecordedTestCase
@@ -360,6 +363,50 @@ def test_servicebus_message_time_to_live():
     assert message.time_to_live == timedelta(days=1)
 
 
+def _pyamqp_encoded_inner_message(body=b"data"):
+    # Encode a minimal AMQP message payload, as the service wraps each
+    # deferred/peeked message inside the management-link response.
+    output = bytearray()
+    encode_payload(output, Message(data=[body]))
+    return bytes(output)
+
+
+def test_servicebus_received_message_deferred_has_lock_token():
+    # Regression test for #42454: receive_deferred_messages returned messages
+    # with lock_token=None. For deferred messages the service returns the
+    # lock-token alongside the message in the management-link response (not as a
+    # delivery-tag), so parse_received_message must read m[b"lock-token"].
+    lock_token = uuid.uuid4()
+    mgmt_response = Message(
+        value={b"messages": [{b"message": _pyamqp_encoded_inner_message(), b"lock-token": lock_token}]}
+    )
+    parsed = PyamqpTransport.parse_received_message(
+        mgmt_response,
+        ServiceBusReceivedMessage,
+        receiver=mock.Mock(),
+        is_deferred_message=True,
+        receive_mode=ServiceBusReceiveMode.PEEK_LOCK,
+    )
+    assert len(parsed) == 1
+    assert parsed[0].lock_token == lock_token
+
+
+def test_servicebus_received_message_peeked_has_no_lock_token():
+    # Peeked messages are not locked and carry no lock-token in the response,
+    # so lock_token must remain None (the deferred-only extraction must not leak
+    # into the peek path).
+    mgmt_response = Message(value={b"messages": [{b"message": _pyamqp_encoded_inner_message()}]})
+    parsed = PyamqpTransport.parse_received_message(
+        mgmt_response,
+        ServiceBusReceivedMessage,
+        receiver=mock.Mock(),
+        is_peeked_message=True,
+        receive_mode=ServiceBusReceiveMode.PEEK_LOCK,
+    )
+    assert len(parsed) == 1
+    assert parsed[0].lock_token is None
+
+
 class TestServiceBusMessageBackcompat(AzureMgmtRecordedTestCase):
 
     @pytest.mark.liveTest