From 57c1afe39a3a4cafc8ab428349a2a2581ab00f30 Mon Sep 17 00:00:00 2001 From: Charles-Edouard Cady Date: Tue, 11 Aug 2026 10:29:21 +0200 Subject: [PATCH] Add optional parameter to force payload offloading When sending a gRPC payload to Temporal, if the payload is too large, we offload some of it to Amazon S3: this commit allows you to do this even if the payload is small. This is useful when uploading several things at once (eg registering multiple child workflows), where any single payload is small but their concatenation is too big and no offload happens. --- .../extra/tests/test_workflow_encoding.py | 80 +++++++++++++++++++ .../workflows/encoding/payload_encoder.py | 12 ++- 2 files changed, 89 insertions(+), 3 deletions(-) diff --git a/src/mistralai/extra/tests/test_workflow_encoding.py b/src/mistralai/extra/tests/test_workflow_encoding.py index ad7ff192..9f3c74ef 100644 --- a/src/mistralai/extra/tests/test_workflow_encoding.py +++ b/src/mistralai/extra/tests/test_workflow_encoding.py @@ -760,3 +760,83 @@ async def test_workflow_encoding_hook_handles_gzipped_response(): assert isinstance(result, httpx.Response) response_body = json.loads(result.content) assert response_body["result"] == original_data + + +@pytest.mark.asyncio +async def test_encode_payload_content_does_not_offload_below_threshold(monkeypatch): + storage = InMemoryBlobStorage() + monkeypatch.setattr( + "mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage", + lambda _: storage, + ) + config = WorkflowEncodingConfig( + payload_offloading=PayloadOffloadingConfig( + min_size_bytes=1024, + storage_config=BlobStorageConfig( + storage_provider=StorageProvider.S3, + bucket_name="test-bucket", + ), + ), + ) + encoder = PayloadEncoder(encoding_config=config) + context = WorkflowContext(namespace="test", execution_id="exec") + + data, options = await encoder.encode_payload_content(b"tiny", context) + + assert options == [] + assert storage.blobs == {} + assert data == b"tiny" + + +@pytest.mark.asyncio +async def test_encode_payload_content_force_offload_bypasses_threshold(monkeypatch): + storage = InMemoryBlobStorage() + monkeypatch.setattr( + "mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage", + lambda _: storage, + ) + config = WorkflowEncodingConfig( + payload_offloading=PayloadOffloadingConfig( + min_size_bytes=1024, + storage_config=BlobStorageConfig( + storage_provider=StorageProvider.S3, + bucket_name="test-bucket", + ), + ), + ) + encoder = PayloadEncoder(encoding_config=config) + context = WorkflowContext(namespace="test", execution_id="exec") + + data, options = await encoder.encode_payload_content( + b"tiny", context, force_offload=True + ) + + assert EncodedPayloadOptions.OFFLOADED in options + assert len(storage.blobs) == 1 + ref = json.loads(data) + assert storage.blobs[ref["key"]] == b"tiny" + + +@pytest.mark.asyncio +async def test_encode_payload_content_force_offload_is_idempotent(monkeypatch): + storage = InMemoryBlobStorage() + monkeypatch.setattr( + "mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage", + lambda _: storage, + ) + config = WorkflowEncodingConfig( + payload_offloading=PayloadOffloadingConfig( + min_size_bytes=1024, + storage_config=BlobStorageConfig( + storage_provider=StorageProvider.S3, + bucket_name="test-bucket", + ), + ), + ) + encoder = PayloadEncoder(encoding_config=config) + context = WorkflowContext(namespace="test", execution_id="exec") + + await encoder.encode_payload_content(b"same", context, force_offload=True) + await encoder.encode_payload_content(b"same", context, force_offload=True) + + assert len(storage.blobs) == 1 diff --git a/src/mistralai/extra/workflows/encoding/payload_encoder.py b/src/mistralai/extra/workflows/encoding/payload_encoder.py index a90cb8aa..6699a7e4 100644 --- a/src/mistralai/extra/workflows/encoding/payload_encoder.py +++ b/src/mistralai/extra/workflows/encoding/payload_encoder.py @@ -213,7 +213,10 @@ def _decrypt(self, data: bytes) -> bytes: ) from main_exc async def _handle_offloading( - self, data: bytes, context: Optional[WorkflowContext] + self, + data: bytes, + context: Optional[WorkflowContext], + force: bool = False, ) -> tuple[bytes, bool]: if ( self.offloading_config is None @@ -223,7 +226,7 @@ async def _handle_offloading( "You must configure payload offloading storage" ) - if len(data) < self.offloading_config.min_size_bytes: + if not force and len(data) < self.offloading_config.min_size_bytes: return data, False if not context: @@ -325,6 +328,7 @@ async def encode_payload_content( context: Optional[WorkflowContext] = None, *, allow_offloading: bool = True, + force_offload: bool = False, ) -> tuple[bytes, list[EncodedPayloadOptions]]: """Handle payload encoding. @@ -353,7 +357,9 @@ async def encode_payload_content( encoding_options.append(EncodedPayloadOptions.COMPRESSED) if allow_offloading and self.offloading_config is not None: - data, offloaded = await self._handle_offloading(data, context) + data, offloaded = await self._handle_offloading( + data, context, force=force_offload + ) if offloaded: encoding_options.append(EncodedPayloadOptions.OFFLOADED)