Fixes:

- Correctly import blobs at aio
- default capabilities to empty list
- Revert async deferred init
- Add async tests
- More aggressive checks for externalized payload

Author: andystaples

durabletask/extensions/azure_blob_payloads/blob_payload_store.py

@@ -8,12 +8,10 @@
 import gzip
 import logging
 import uuid
-from typing import TYPE_CHECKING, Optional
+from typing import Optional
 
 from azure.storage.blob import BlobServiceClient
-
-if TYPE_CHECKING:
-    from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
+from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
 
 from durabletask.extensions.azure_blob_payloads.options import BlobPayloadStoreOptions
 from durabletask.payload.store import PayloadStore
@@ -69,35 +67,25 @@ def __init__(self, options: BlobPayloadStoreOptions):
                 **extra_kwargs,
             )
 
-        # Async client is built lazily to avoid importing
-        # azure.storage.blob.aio when only sync methods are used.
-        self._async_blob_service_client: AsyncBlobServiceClient | None = None
-        self._extra_kwargs = extra_kwargs
+        # Build async client
+        if options.connection_string:
+            self._async_blob_service_client = AsyncBlobServiceClient.from_connection_string(
+                options.connection_string, **extra_kwargs,
+            )
+        else:
+            assert options.account_url is not None  # guaranteed by validation above
+            self._async_blob_service_client = AsyncBlobServiceClient(
+                account_url=options.account_url,
+                credential=options.credential,
+                **extra_kwargs,
+            )
 
         self._ensure_container_created = False
 
     @property
     def options(self) -> BlobPayloadStoreOptions:
         return self._options
 
-    def _get_async_blob_service_client(self) -> AsyncBlobServiceClient:
-        """Lazily create the async blob service client."""
-        if self._async_blob_service_client is None:
-            from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
-
-            if self._options.connection_string:
-                self._async_blob_service_client = AsyncBlobServiceClient.from_connection_string(
-                    self._options.connection_string, **self._extra_kwargs,
-                )
-            else:
-                assert self._options.account_url is not None
-                self._async_blob_service_client = AsyncBlobServiceClient(
-                    account_url=self._options.account_url,
-                    credential=self._options.credential,
-                    **self._extra_kwargs,
-                )
-        return self._async_blob_service_client
-
     # ------------------------------------------------------------------
     # Sync operations
     # ------------------------------------------------------------------
@@ -138,8 +126,7 @@ async def upload_async(self, data: bytes, *, instance_id: Optional[str] = None)
             data = gzip.compress(data)
 
         blob_name = self._make_blob_name(instance_id)
-        client = self._get_async_blob_service_client()
-        container_client = client.get_container_client(self._container_name)
+        container_client = self._async_blob_service_client.get_container_client(self._container_name)
         await container_client.upload_blob(name=blob_name, data=data, overwrite=True)
 
         token = f"{_TOKEN_PREFIX}{self._container_name}:{blob_name}"
@@ -148,8 +135,7 @@ async def upload_async(self, data: bytes, *, instance_id: Optional[str] = None)
 
     async def download_async(self, token: str) -> bytes:
         container, blob_name = self._parse_token(token)
-        client = self._get_async_blob_service_client()
-        container_client = client.get_container_client(container)
+        container_client = self._async_blob_service_client.get_container_client(container)
         stream = await container_client.download_blob(blob_name)
         blob_data = await stream.readall()
 
@@ -164,7 +150,11 @@ async def download_async(self, token: str) -> bytes:
     # ------------------------------------------------------------------
 
     def is_known_token(self, value: str) -> bool:
-        return value.startswith(_TOKEN_PREFIX)
+        try:
+            self._parse_token(value)
+            return True
+        except ValueError:
+            return False
 
     @staticmethod
     def _parse_token(token: str) -> tuple[str, str]:
@@ -203,8 +193,7 @@ def _ensure_container_sync(self) -> None:
     async def _ensure_container_async(self) -> None:
         if self._ensure_container_created:
             return
-        client = self._get_async_blob_service_client()
-        container_client = client.get_container_client(self._container_name)
+        container_client = self._async_blob_service_client.get_container_client(self._container_name)
         try:
             await container_client.create_container()
         except Exception:

durabletask/worker.py

@@ -508,7 +508,7 @@ def should_invalidate_connection(rpc_error):
                 get_work_items_request = pb.GetWorkItemsRequest(
                     maxConcurrentOrchestrationWorkItems=self._concurrency_options.maximum_concurrent_orchestration_work_items,
                     maxConcurrentActivityWorkItems=self._concurrency_options.maximum_concurrent_activity_work_items,
-                    capabilities=capabilities if capabilities else None,
+                    capabilities=capabilities,
                 )
                 self._response_stream = stub.GetWorkItems(get_work_items_request)
                 self._logger.info(

pyproject.toml

@@ -37,7 +37,7 @@ opentelemetry = [
     "opentelemetry-sdk>=1.0.0"
 ]
 azure-blob-payloads = [
-    "azure-storage-blob>=12.0.0"
+    "azure-storage-blob[aio]>=12.0.0"
 ]
 
 [project.urls]

tests/durabletask/test_large_payload.py

@@ -3,7 +3,6 @@
 
 """Tests for large-payload externalization and de-externalization."""
 
-import asyncio
 from typing import Optional
 from unittest.mock import MagicMock
 
@@ -317,7 +316,8 @@ def test_externalize_then_deexternalize(self):
 
 
 class TestAsyncPayloadHelpers:
-    def test_async_externalize_and_deexternalize(self):
+    @pytest.mark.asyncio
+    async def test_async_externalize_and_deexternalize(self):
         """Async versions should work identically to sync."""
         store = FakePayloadStore(threshold_bytes=10)
         original = "async round trip " * 20
@@ -328,14 +328,10 @@ def test_async_externalize_and_deexternalize(self):
             input=sv(original),
         )
 
-        asyncio.get_event_loop().run_until_complete(
-            externalize_payloads_async(req, store, instance_id="async-1")
-        )
+        await externalize_payloads_async(req, store, instance_id="async-1")
         assert req.input.value.startswith(FakePayloadStore.TOKEN_PREFIX)
 
-        asyncio.get_event_loop().run_until_complete(
-            deexternalize_payloads_async(req, store)
-        )
+        await deexternalize_payloads_async(req, store)
         assert req.input.value == original
 
 
@@ -402,10 +398,13 @@ def test_is_known_token(self):
 
         store = MagicMock(spec=BlobPayloadStore)
         store.is_known_token = BlobPayloadStore.is_known_token.__get__(store)
+        store._parse_token = BlobPayloadStore._parse_token
 
         assert store.is_known_token("blob:v1:c:b") is True
         assert store.is_known_token("not-a-token") is False
         assert store.is_known_token("") is False
+        assert store.is_known_token("blob:v1:") is False
+        assert store.is_known_token("blob:v1:container:") is False
 
 
 # ------------------------------------------------------------------

tests/durabletask/test_large_payload_e2e.py

@@ -357,3 +357,59 @@ def echo(ctx: task.OrchestrationContext, inp: str):
                 svc.delete_container(fresh_container)
             except Exception:
                 pass
+
+
+class TestAsyncBlobClient:
+    """Test upload_async / download_async against Azurite."""
+
+    @pytest.fixture()
+    def async_store(self):
+        """Per-test BlobPayloadStore with a unique container."""
+        container = f"async-test-{uuid.uuid4().hex[:8]}"
+        store = BlobPayloadStore(BlobPayloadStoreOptions(
+            connection_string=AZURITE_CONN_STR,
+            container_name=container,
+            threshold_bytes=THRESHOLD_BYTES,
+            enable_compression=True,
+            api_version=AZURITE_API_VERSION,
+        ))
+        yield store
+        try:
+            svc = azure_blob.BlobServiceClient.from_connection_string(
+                AZURITE_CONN_STR, api_version=AZURITE_API_VERSION,
+            )
+            svc.delete_container(container)
+        except Exception:
+            pass
+
+    @pytest.mark.asyncio
+    async def test_async_upload_and_download_round_trip(self, async_store):
+        """upload_async stores data that download_async can retrieve."""
+        payload = b"async round-trip payload " * 200
+        token = await async_store.upload_async(payload, instance_id="async-1")
+
+        assert async_store.is_known_token(token)
+        result = await async_store.download_async(token)
+        assert result == payload
+
+    @pytest.mark.asyncio
+    async def test_async_upload_with_compression(self, async_store):
+        """Compressed upload should still decompress on download."""
+        payload = b"Z" * 5000
+        token = await async_store.upload_async(payload)
+
+        downloaded = await async_store.download_async(token)
+        assert downloaded == payload
+
+    @pytest.mark.asyncio
+    async def test_async_upload_instance_id_scopes_blob(self, async_store):
+        """Blobs uploaded with instance_id are scoped under that prefix."""
+        payload = b"scoped payload"
+        token = await async_store.upload_async(payload, instance_id="inst-42")
+
+        # Token format: blob:v1:<container>:<instance_id>/<uuid>
+        _, blob_name = BlobPayloadStore._parse_token(token)
+        assert blob_name.startswith("inst-42/")
+
+        downloaded = await async_store.download_async(token)
+        assert downloaded == payload