[Key Vault] Update continuation tokens

sdk/keyvault/azure-keyvault-administration/CHANGELOG.md

@@ -6,10 +6,15 @@
 
 ### Breaking Changes
 
-### Bugs Fixed
+- Changed the continuation token format. Continuation tokens generated by previous versions of
+  `azure-keyvault-administration` are not compatible with this version. Similarly, continuation tokens generated by
+  previous versions of this library are not compatible with versions of `azure-core>=1.38.0`.
 
+### Bugs Fixed
 ### Other Changes
 
+- Updated minimum `azure-core` version to 1.38.0
+
 ## 4.6.0 (2025-06-16)
 
 ### Features Added

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_backup_client.py

@@ -4,12 +4,13 @@
 # ------------------------------------
 import base64
 import functools
-import pickle
+import json
 from typing import Any, Callable, Optional, overload
 from urllib.parse import urlparse
 
 from typing_extensions import Literal
 
+from azure.core.pipeline import PipelineResponse
 from azure.core.polling import LROPoller
 from azure.core.tracing.decorator import distributed_trace
 
@@ -19,12 +20,61 @@
 from ._internal.polling import KeyVaultBackupClientPolling, KeyVaultBackupClientPollingMethod
 
 
-def _parse_status_url(url):
+def _parse_status_url(url: str) -> str:
     parsed = urlparse(url)
     job_id = parsed.path.split("/")[2]
     return job_id
 
 
+def _get_continuation_token(pipeline_response: PipelineResponse) -> str:
+    """Returns an opaque token which can be used by the user to rehydrate/restart the LRO.
+
+    Saves the state of the LRO based on a status response so that polling can be resumed from that context. Because
+    the service has different operations for backup/restore starting vs. status checking, the caller is expected to
+    first use the status URL from the initial response to make a status request and then pass that status response to
+    this function to be serialized into a continuation token.
+
+    :param pipeline_response: The pipeline response of the operation status request.
+    :type pipeline_response: ~azure.core.pipeline.PipelineResponse
+    :returns: An opaque continuation token that can be provided to Core to rehydrate the LRO.
+    :rtype: str
+    """
+    # Headers needed for LRO rehydration - use an allowlist approach for security
+    lro_headers = {"azure-asyncoperation", "operation-location", "location", "content-type", "retry-after"}
+    response = pipeline_response.http_response
+    filtered_headers = {k: v for k, v in response.headers.items() if k.lower() in lro_headers}
+
+    request = response.request
+    # Serialize the essential parts of the PipelineResponse to JSON.
+    if request:
+        request_headers = {}
+        # Preserve x-ms-client-request-id for request correlation
+        if "x-ms-client-request-id" in request.headers:
+            request_headers["x-ms-client-request-id"] = request.headers["x-ms-client-request-id"]
+        request_state = {
+            "method": request.method,
+            "url": request.url,
+            "headers": request_headers,
+        }
+    else:
+        request_state = None
+
+    # Use a versioned token schema: {"version": <int>, "data": <your state>}
+    # This allows for future compatibility checking when deserializing
+    token = {
+        "version": 1,
+        "data": {
+            "request": request_state,
+            "response": {
+                "status_code": response.status_code,
+                "headers": filtered_headers,
+                "content": base64.b64encode(response.content).decode("ascii"),
+            },
+        },
+    }
+    return base64.b64encode(json.dumps(token).encode("utf-8")).decode("ascii")
+
+
 class KeyVaultBackupClient(KeyVaultClientBase):
     """Performs Key Vault backup and restore operations.
 
@@ -56,7 +106,7 @@ def _use_continuation_token(self, continuation_token: str, status_method: Callab
         )
         if "azure-asyncoperation" not in pipeline_response.http_response.headers:
             pipeline_response.http_response.headers["azure-asyncoperation"] = status_url
-        return base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii")
+        return _get_continuation_token(pipeline_response)
 
     @overload
     def begin_backup(

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/_internal/polling.py

@@ -15,7 +15,7 @@
 class KeyVaultBackupClientPolling(OperationResourcePolling):
     def __init__(self) -> None:
         self._polling_url = ""
-        super(KeyVaultBackupClientPolling, self).__init__(operation_location_header="azure-asyncoperation")
+        super().__init__(operation_location_header="azure-asyncoperation")
 
     def get_polling_url(self) -> str:
         return self._polling_url
@@ -34,4 +34,13 @@ def set_initial_status(self, pipeline_response: "PipelineResponse") -> str:
 
 class KeyVaultBackupClientPollingMethod(LROBasePolling):
     def get_continuation_token(self) -> str:
+        """
+        Get a continuation token to resume the polling later.
+
+        :return: A continuation token.
+        :rtype: str
+        """
+        # Because of the operation structure, we need to use a "continuation token" that is just the status URL.
+        # This URL can then be used to fetch the status of the operation when resuming, at which point a genuine
+        # continuation token will be created from the response and provided to Core.
         return base64.b64encode(self._operation.get_polling_url().encode()).decode("ascii")

sdk/keyvault/azure-keyvault-administration/azure/keyvault/administration/aio/_backup_client.py

@@ -4,7 +4,6 @@
 # ------------------------------------
 import base64
 import functools
-import pickle
 from typing import Any, Callable, Optional, overload
 
 from typing_extensions import Literal
@@ -13,7 +12,7 @@
 from azure.core.tracing.decorator_async import distributed_trace_async
 
 from .._generated.models import PreBackupOperationParameters, PreRestoreOperationParameters, SASTokenParameter
-from .._backup_client import _parse_status_url
+from .._backup_client import _get_continuation_token, _parse_status_url
 from .._internal import AsyncKeyVaultClientBase, parse_folder_url
 from .._internal.async_polling import KeyVaultAsyncBackupClientPollingMethod
 from .._internal.polling import KeyVaultBackupClientPolling
@@ -51,7 +50,7 @@ async def _use_continuation_token(self, continuation_token: str, status_method:
         )
         if "azure-asyncoperation" not in pipeline_response.http_response.headers:
             pipeline_response.http_response.headers["azure-asyncoperation"] = status_url
-        return base64.b64encode(pickle.dumps(pipeline_response)).decode("ascii")
+        return _get_continuation_token(pipeline_response)
 
     @overload
     async def begin_backup(

sdk/keyvault/azure-keyvault-administration/setup.py

@@ -62,7 +62,7 @@
     ),
     install_requires=[
         "isodate>=0.6.1",
-        "azure-core>=1.31.0",
+        "azure-core>=1.38.0",
         "typing-extensions>=4.6.0",
     ],
     python_requires=">=3.9",