Add Storage V2 support across hosting providers - #569
Add Storage V2 support across hosting providers#569Cecilia Avila (ceciliaavila) wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
🔵 Needs a closer look
It introduces a new storage contract and propagates it through multiple core runtime paths (state/OAuth/proactive) plus external providers, making regression risk non-trivial despite strong test coverage.
Pull request overview
Adds Storage V2 (per-key operation results + optimistic concurrency tokens) to the Python SDK while keeping Storage V1 as the default, and wires V2 support through core hosting flows (state, OAuth, proactive messaging, and conversation ID storage) plus Blob/Cosmos providers.
Changes:
- Introduces Storage V2 types (statuses, versions, options/modes) and V1↔V2 compatibility adapters/helpers.
- Extends Memory/Blob/Cosmos storage providers to support V2 behaviors (read/write/delete results, concurrency control) with V1 preserved by default.
- Expands docs and tests to cover V2 behaviors, config defaults, and concurrency batching.
File summaries
| File | Description |
|---|---|
| tests/storage_cosmos/test_cosmos_db_storage.py | Adds V2 option rejection checks for V1 and concurrency batching tests for Cosmos. |
| tests/storage_cosmos/test_cosmos_db_config.py | Verifies Cosmos config defaults to V1 and can select V2. |
| tests/storage_blob/test_blob_storage.py | Adds Blob config version tests plus V1 option rejection and V2 concurrency batching tests. |
| tests/hosting_core/storage/test_storage_compatibility.py | Adds tests for V1↔V2 adapters and helper utilities. |
| tests/hosting_core/storage/test_memory_storage.py | Adds MemoryStorage V2 behavior tests (results, modes, versions, immutability) and V1 validation checks. |
| tests/hosting_core/state/test_agent_state.py | Adds integration coverage for state management using MemoryStorage V2. |
| test_samples/app_style/README.md | Documents env-based selection of provider + storage version and /v2-demo usage. |
| libraries/microsoft-agents-storage-cosmos/readme.md | Documents Cosmos storage version selection and what V2 provides. |
| libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py | Implements V2 read/write/delete paths (results + optimistic concurrency) alongside V1 behavior. |
| libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage_config.py | Adds storage_version configuration and typing updates. |
| libraries/microsoft-agents-storage-blob/readme.md | Documents Blob config storage_version option. |
| libraries/microsoft-agents-storage-blob/microsoft_agents/storage/blob/blob_storage.py | Implements V2 read/write/delete (results + ETag-based concurrency) alongside V1 behavior. |
| libraries/microsoft-agents-storage-blob/microsoft_agents/storage/blob/blob_storage_config.py | Adds storage_version to Blob config with typing updates. |
| libraries/microsoft-agents-hosting-core/setup.py | Adds typing-extensions dependency to support TypeVar defaults. |
| libraries/microsoft-agents-hosting-core/readme.md | Adds Storage V2 usage documentation (section heading hierarchy needs a small fix). |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/storage/storage.py | Introduces Storage V2 contracts, result/status types, options, and version typing. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/storage/storage_compatibility.py | Adds adapter layer and validation/helpers for bridging V1 and V2 semantics. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/storage/memory_storage.py | Adds selectable V1/V2 behavior, V2 result/mode/version support, and deep-copy semantics for V2. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/storage/init.py | Re-exports new storage V2 symbols from the package. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/state/user_state.py | Broadens storage type to StorageProvider for V1/V2 compatibility. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/state/agent_state.py | Routes state load/save/delete through V2 compatibility seam to preserve V1 behavior. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/client/conversation_id_factory.py | Updates conversation ID storage to use V2 helpers and fixes AgentsModel StoreItem method binding. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/state/turn_state.py | Updates TurnState APIs to accept StorageProvider. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/state/conversation_state.py | Updates ConversationState to accept StorageProvider. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/proactive/proactive.py | Updates proactive conversation persistence to use V2 helpers and preserve V1 semantics. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/proactive/proactive_options.py | Updates proactive options storage type to StorageProvider. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/authorization.py | Updates OAuth flow state reads/writes/deletes through V2 helpers for compatibility. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/_handlers/connector_user_authorization.py | Updates handler storage type to StorageProvider. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/_handlers/agentic_user_authorization.py | Updates handler storage type to StorageProvider. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/oauth/_handlers/_authorization_handler.py | Updates base handler storage type to StorageProvider. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/app/app_options.py | Updates app storage option type to StorageProvider. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/_oauth/_flow_storage_client.py | Updates two-tier OAuth cache/backing storage client to use V2 helpers. |
| libraries/microsoft-agents-hosting-core/microsoft_agents/hosting/core/init.py | Re-exports new storage V2 symbols at the package root. |
Review details
- Files reviewed: 33/33 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Cecilia Avila (ceciliaavila)
changed the title
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
CosmosDBStorage V2 replace writes appear to omit the required partition key for replace_item, which can cause runtime failures against partitioned Cosmos containers.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 33/33 changed files
- Comments generated: 1
- Review effort level: Lite
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Storage V2 REPLACE mode in BlobStorage and CosmosDBStorage is currently conditional even without expected_version, which breaks the intended “replace vs conditional replace” semantics under concurrency.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
libraries/microsoft-agents-storage-blob/microsoft_agents/storage/blob/blob_storage.py:364
- In V2 write(), REPLACE mode is currently treated as conditional even when expected_version is not provided (it sets the upload condition to the blob’s current ETag). That makes plain “replace” fail with CONDITION_NOT_MET under concurrent updates, which conflicts with the intended split between replace vs conditional replace (expected_version). Consider making REPLACE require existence but only enforce ETag matching when expected_version is set.
- Files reviewed: 33/33 changed files
- Comments generated: 1
- Review effort level: Lite
| elif write_options.mode == StorageWriteMode.REPLACE: | ||
| response = await self._container.replace_item( | ||
| escaped_key, | ||
| document, | ||
| etag=current_version, | ||
| match_condition=MatchConditions.IfNotModified, | ||
| partition_key=self._get_partition_key(escaped_key), | ||
| ) |
There was a problem hiding this comment.
🟡 Changes recommended
Conditional Cosmos writes and unconditional provider deletes have race-related correctness issues, and the documented sample functionality is absent.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (4)
libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py:396
- This conditional UPSERT is vulnerable to recreating an item that is deleted after the pre-read: Cosmos can treat the operation as a create and not enforce
IfNotModifiedfor the now-missing document. Usereplace_itemwith the expected ETag so the expected-version write can never create data during that race.
elif write_options.expected_version is not None:
response = await self._container.upsert_item(
body=document,
etag=write_options.expected_version,
match_condition=MatchConditions.IfNotModified,
libraries/microsoft-agents-storage-blob/microsoft_agents/storage/blob/blob_storage.py:420
- When
expected_versionis omitted, this pre-read ETag is still passed todelete_blob, so an unconditional delete becomes conditional and can returnconditionNotMetif a concurrent writer updates the blob between the two calls. Only read and send an ETag when the caller explicitly supplied an expected version; otherwise delete directly and update the concurrency test double to synchronizedelete_blobitself.
current_version = await self._get_current_version(blob_client)
if current_version is None:
libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py:450
- Even without
expected_version, this reads the current ETag and always performs an If-Match delete. A concurrent update between these calls therefore makes an unconditional delete fail withconditionNotMet; only pre-read and attach the condition when the caller requested conditional deletion, and make the test barrier synchronizedelete_itemdirectly.
current = await self._try_read_document(key)
if current is None:
return StorageDeleteResult(
key=key, status=StorageOperationStatus.NOT_FOUND
)
libraries/microsoft-agents-storage-cosmos/microsoft_agents/storage/cosmos/cosmos_db_storage.py:419
- For an UPSERT carrying
expected_version, a 404 from the atomic replace path means the expected version ceased to match because the item disappeared; the other providers report that asconditionNotMet. ReturningnotFoundhere makes the same conditional race provider-dependent, so map this case toCONDITION_NOT_METwhile retainingNOT_FOUNDfor explicit replace-without-version.
if status_code == 404:
return StorageWriteResult(
key=key, status=StorageOperationStatus.NOT_FOUND
)
- Files reviewed: 33/33 changed files
- Comments generated: 3
- Review effort level: Balanced
|
|
||
| async def write_one(key: str, value: StoreItem) -> StorageWriteResult: | ||
| blob_client = self._container_client.get_blob_client(key) | ||
| current_version = await self._get_current_version(blob_client) |
| current = await self._try_read_document(key) | ||
| current_version = current.get("_etag") if current else None |
|
|
||
| ## Echo storage V2 test | ||
|
|
||
| `echo_proactive_agent.py` selects storage with these `.env` values: |
| def get_storage_read_value( | ||
| results: StorageReadResults[StoreItemT] | None, key: str | ||
| ) -> StoreItemT | None: | ||
| """Return a successful V2 value, map not-found to ``None``, or raise.""" | ||
| result = results.get(key) if results is not None else None | ||
| if result is not None and result.status == StorageOperationStatus.NOT_FOUND: | ||
| return None | ||
| if result is not None and result.status == StorageOperationStatus.SUCCEEDED: | ||
| return result.value | ||
| _raise_result_error("read", key, result.status if result else None) |
There was a problem hiding this comment.
I would rather we make StorageReadResults a class and have this be a method of it. This API isn't very user friendly as the user has to import helper functions in order to deal with the result of the new storage operations.
There was a problem hiding this comment.
I think you could override the bool method of StorageReadResults to return False if the dictionary is empty. That would be useful and would handle None results.
| def assert_storage_write_succeeded( | ||
| results: StorageWriteResults | None, keys: list[str] | ||
| ) -> None: | ||
| """Raise unless every V2 write result succeeded.""" | ||
| _assert_results("write", results, keys, {StorageOperationStatus.SUCCEEDED}) |
There was a problem hiding this comment.
Same here, and same applies to the other methods below.
| async def read(self, keys, *, target_cls, **kwargs): | ||
| validate_storage_v2_keys(keys) | ||
| if not keys: | ||
| return {} | ||
| items = await self._storage.read(keys, target_cls=target_cls, **kwargs) | ||
| return { |
There was a problem hiding this comment.
Let's annotate the type of the arguments. We trust in our type annotations and the linter to catch things for us. In this SDK, we often don't perform checks on preconditions like this unless the incoming value is typed as Any or object.
There was a problem hiding this comment.
Also, annotate the return type.
| class _StorageToStorageV2Adapter(StorageV2): | ||
| """Adapt a legacy provider where V2 behavior is safely available.""" | ||
|
|
||
| storage_version = StorageVersion.V2 |
There was a problem hiding this comment.
Can we make this a private field with a public getter instead?
There was a problem hiding this comment.
Or perhaps just a get:
Property
def storage_version(...
| def validate_write_mode(mode: StorageWriteMode) -> None: | ||
| """Validate a V2 write mode.""" | ||
| if not isinstance(mode, StorageWriteMode): | ||
| raise ValueError(f'Storage V2 write mode "{mode}" is not supported.') |
There was a problem hiding this comment.
Similar, we trust our typing annotations. In this SDK we don't do a lot of defensive programming because of that.
| def store_item_to_json(instance): | ||
| return instance.model_dump(mode="json", exclude_none=True) | ||
|
|
||
| @classmethod | ||
| def from_json_to_store_item(cls, data): | ||
| return cls.model_validate(data) |
There was a problem hiding this comment.
add type annotations
|
|
||
|
|
||
| class BlobStorageConfig: | ||
| class BlobStorageConfig(Generic[StorageVersionT]): |
There was a problem hiding this comment.
Why is the storage version a generic type argument to the class? Would it be sufficient to remove this and keep the storage_version field only?
|
|
||
|
|
||
| class BlobStorage(AsyncStorageBase): | ||
| class BlobStorage(AsyncStorageBase, StorageV2, Generic[StorageVersionT]): |
There was a problem hiding this comment.
Can we consolidate read_one and _read_item (and same idea for delete and write) into one method? Make the v1 logic a wrapper around the v2 code.
|
|
||
|
|
||
| class CosmosDBStorageConfig: | ||
| class CosmosDBStorageConfig(Generic[StorageVersionT]): |
There was a problem hiding this comment.
Same feedback I gave for BlobStorage
|
|
||
|
|
||
| class CosmosDBStorage(AsyncStorageBase): | ||
| class CosmosDBStorage(AsyncStorageBase, StorageV2, Generic[StorageVersionT]): |
There was a problem hiding this comment.
Why do we need to support both versions in the same class? I think it would be cleaner to provider a CosmosDBStorageV2 than to support both paths in the same class. I don't like the idea of the signature of the class' API being affected by the generic argument to CosmosDBStorageConfig. It seems backwards. As a user, I would rather select the API/model I will use and then configure it after the fact.
Fixes #559
Important
This PR ports Storage V2 from #1272 (JS PR) using microsoft/Agents-for-net#883 as reference.
Description
Add Storage V2 to Python. Storage V1 remains the default and remains compatible.
Changes
Storage API
read()return types from the storage constructor version.Storage providers
Sample and docs
STORAGE_PROVIDERandSTORAGE_VERSIONconfiguration./v2-demoto validate V2 statuses and optimistic concurrency./deleteflow to validate V1 and V2 state behavior.Tests
Testing
pytest tests/hosting_core/storage tests/hosting_core/state tests/storage_blob tests/storage_cosmos -qThe following images show the V1 and V2 state checks for Memory, Blob, and Cosmos DB: counter progression,
/delete, and reset to[0].The V2 capability checks passed for Memory, Blob, and Cosmos DB:
notFound.conflict.conditionNotMet.notFound.