From d58fd3bc3b884aec6951324d567dcd8d7fe4cc06 Mon Sep 17 00:00:00 2001 From: Shanmukha Pasumarthy Date: Tue, 21 Jul 2026 12:51:03 +0530 Subject: [PATCH] [agentserver-core] Add FoundryStateStore durable KV storage layer Add azure.ai.agentserver.core.storage: protocol-neutral FoundryStorageClient (transport, endpoint, pipeline policies, error hierarchy) plus FoundryStateStore, a durable key-value store bound to one explicit store. get_or_create(name, ...) is the primary entry point; store-level get/update/delete act on the bound store, and explicit create_item/set_item/get_item/delete_item/list_keys act on items. Addresses PR review feedback: - Share a single _UNSET sentinel between _state and _state_serializer so one-field update() calls no longer leak the sentinel into the payload. - Close the store (pipeline + owned credential) on every get_or_create failure, including cancellation, so a failed resolution cannot leak resources. - Bound-store get() returns StateStore and raises FoundryStorageNotFoundError instead of returning None; docs and tests updated to match. - Regenerate api.md / api.metadata.yml to include the storage surface. - Add regression tests for one-field updates and get_or_create cleanup. Fix CI checks: - mypy: pass list_keys query params to _request via a typed query: Mapping[str, str] | None parameter instead of **query splat. - pylint C4758: document keyword-only args of create_item, set_item and list_keys with :keyword:/:paramtype: entries. - cspell: rename test_unparseable_body_uses_fallback_message to test_non_json_body_uses_fallback_message. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 2c9dca80-df28-4b1e-8c61-3a69bf7e9a8a --- .../azure-ai-agentserver-core/CHANGELOG.md | 5 +- .../azure-ai-agentserver-core/README.md | 33 + .../azure-ai-agentserver-core/api.md | 421 +++- .../api.metadata.yml | 4 +- .../azure/ai/agentserver/core/__init__.py | 23 +- .../ai/agentserver/core/storage/__init__.py | 51 + .../ai/agentserver/core/storage/_client.py | 111 ++ .../ai/agentserver/core/storage/_endpoint.py | 69 + .../ai/agentserver/core/storage/_errors.py | 123 ++ .../core/storage/_generated/__init__.py | 45 + .../core/storage/_generated/_enums.py | 24 + .../core/storage/_generated/_models.py | 683 +++++++ .../storage/_generated/_utils/__init__.py | 6 + .../storage/_generated/_utils/model_base.py | 1770 +++++++++++++++++ .../ai/agentserver/core/storage/_json.py | 15 + .../ai/agentserver/core/storage/_policies.py | 147 ++ .../ai/agentserver/core/storage/_state.py | 471 +++++ .../core/storage/_state_serializer.py | 176 ++ .../docs/state-store-guide.md | 340 ++++ .../azure-ai-agentserver-core/mypy.ini | 8 + .../azure-ai-agentserver-core/pyproject.toml | 2 + .../samples/state_store_sample.py | 116 ++ .../tests/test_foundry_state_store.py | 624 ++++++ .../tests/test_state_store_sample_usage.py | 246 +++ .../tests/test_storage_endpoint.py | 57 + .../tests/test_storage_errors.py | 88 + .../tests/test_storage_policies.py | 48 + 27 files changed, 5629 insertions(+), 77 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_endpoint.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_errors.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_enums.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_models.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/model_base.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_json.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_policies.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state_serializer.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md create mode 100644 sdk/agentserver/azure-ai-agentserver-core/mypy.ini create mode 100644 sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_state_store_sample_usage.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_endpoint.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_errors.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index a300e4f5517b..06b321ceb6b7 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -2,6 +2,10 @@ ## 2.0.0b8 (2026-07-22) +### Features Added + +- Added `azure.ai.agentserver.core.storage`, the protocol-neutral Foundry durable state-store layer. `FoundryStateStore.get_or_create(name, ...)` is the primary entry point -- an async classmethod that resolves (or creates, on first use) the store in one call. Store-level operations act on the bound store: `get()` (its descriptor), `update(...)` (its mutable `description` / `tags`), and `delete()` (the whole store, cascade-deleting every item). Item operations are explicit and consistently named: `create_item`, `set_item`, `get_item`, `delete_item`, `list_keys`. Store names are path-encoded with base64url on the wire, store-level `item_ttl_seconds` is configured once at create, optional `user_isolation` is declared at store create, and trusted callers may delegate end-user partitioning with `user_id` (`x-ms-user-id`). Response bodies (`StateStore`, `StateStoreItem`, `StateStoreItemRef`, `DeletedStateStore`, `DeletedStateStoreItem`, `StateStoreItemKey`) are typed model classes generated from a formal TypeSpec contract, not hand-written dataclasses. See the [Durable State Store Guide](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md). + ### Bugs Fixed - Fixed span attribute enrichment under `opentelemetry-sdk` >= 1.43.0, where `span._attributes` became a `BoundedAttributes` (backed by `._dict`) that no longer supports item assignment. Enrichment now resolves the backing store so agent identity attributes are written on both older and newer OpenTelemetry SDKs. @@ -24,7 +28,6 @@ - Populated agent metadata when operation IDs are zeroed so agent metadata remains available for telemetry and downstream processing. - Suppressed noisy observability/exporter INFO logs by default in tracing setup while preserving DEBUG visibility when explicitly enabled. - ## 2.0.0b5 (2026-05-25) ### Bugs Fixed diff --git a/sdk/agentserver/azure-ai-agentserver-core/README.md b/sdk/agentserver/azure-ai-agentserver-core/README.md index 2a9d56c42346..9dc0a8da465d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/README.md +++ b/sdk/agentserver/azure-ai-agentserver-core/README.md @@ -126,6 +126,39 @@ async def on_shutdown(): pass ``` +### Durable state storage + +`FoundryStateStore` is a durable, server-backed key-value store for agent state +— session memory, per-user preferences, counters, and checkpoints — bound to +one explicit, caller-named store, with single-item optimistic concurrency, +tag-filtered key listing, and store-level TTL. + +```python +from azure.ai.agentserver.core.storage import FoundryStateStore + +# Endpoint and credential resolve from FOUNDRY_PROJECT_ENDPOINT + DefaultAzureCredential. +# get_or_create() resolves (or creates, on first use) the store in one call. +store = await FoundryStateStore.get_or_create("checkpoints/thread-abc", user_isolation=True) +async with store: + await store.set_item("step-1", {"done": False}) + item = await store.get_item("step-1") + print(item.value) # {"done": False} +``` + +> The default credential path uses `DefaultAzureCredential`, which requires the +> optional `azure-identity` package (`pip install azure-identity`). Alternatively, +> pass any `azure.core.credentials_async.AsyncTokenCredential` explicitly to +> `get_or_create()` and `azure-identity` is not needed. + +Reads return typed `StateStoreItem` values; writes return typed item metadata and use +single-item `If-Match` concurrency. Session/conversation scoping is expressed in +the store name itself, and item expiry is controlled by the store's +`item_ttl_seconds` setting. See the [Durable State Store Guide](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md) +for the full API, the store lifecycle, and common gotchas, and +[state_store_sample.py](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py) +for a runnable end-to-end example. + + ### Configuring tracing Tracing is enabled automatically when an Application Insights connection string is available: diff --git a/sdk/agentserver/azure-ai-agentserver-core/api.md b/sdk/agentserver/azure-ai-agentserver-core/api.md index d3aafdb8a0bd..78ea68641238 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/api.md +++ b/sdk/agentserver/azure-ai-agentserver-core/api.md @@ -77,13 +77,6 @@ namespace azure.ai.agentserver.core class azure.ai.agentserver.core.AgentServerHost(Starlette): property routes: list[BaseRoute] # Read-only - async def __call__( - self, - scope: Scope, - receive: Receive, - send: Send - ) -> None: ... - def __init__( self, *, @@ -97,44 +90,6 @@ namespace azure.ai.agentserver.core **kwargs: Any ) -> None: ... - def add_exception_handler( - self, - exc_class_or_status_code: int | type[Exception], - handler: ExceptionHandler - ) -> None: ... - - def add_middleware( - self, - middleware_class: _MiddlewareFactory[P], - *args: args, - **kwargs: kwargs - ) -> None: ... - - def add_route( - self, - path: str, - route: Callable[[Request], Awaitable[Response] | Response], - methods: list[str] | None = None, - name: str | None = None, - include_in_schema: bool = True - ) -> None: ... - - def build_middleware_stack(self) -> ASGIApp: ... - - def host( - self, - host: str, - app: ASGIApp, - name: str | None = None - ) -> None: ... - - def mount( - self, - path: str, - app: ASGIApp, - name: str | None = None - ) -> None: ... - def register_server_version(self, version_segment: str) -> None: ... def run( @@ -154,13 +109,6 @@ namespace azure.ai.agentserver.core @staticmethod async def sse_keepalive_stream(iterator: AsyncIterable[_Content], interval: int) -> AsyncIterator[_Content]: ... - def url_path_for( - self, - name: str, - /, - **path_params: Any - ) -> URLPath: ... - class azure.ai.agentserver.core.FoundryAgentRequestContext: call_id: str | None @@ -202,4 +150,373 @@ namespace azure.ai.agentserver.core def __init__(self, app: ASGIApp) -> None: ... +namespace azure.ai.agentserver.core.storage + + class azure.ai.agentserver.core.storage.DeletedStateStore(_Model): + deleted: bool + id: Optional[str] + name: str + object: Literal[StateStoreObjectType.STATE_STORE] + + @overload + def __init__( + self, + *, + deleted: bool, + id: Optional[str] = ..., + name: str, + object: Literal[StateStoreObjectType.STATE_STORE] + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: ... + + + class azure.ai.agentserver.core.storage.DeletedStateStoreItem(_Model): + deleted: bool + id: Optional[str] + key: str + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] + + @overload + def __init__( + self, + *, + deleted: bool, + id: Optional[str] = ..., + key: str, + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: ... + + + class azure.ai.agentserver.core.storage.FoundryStateStore(FoundryStorageClient): implements AsyncContextManager + property name: str # Read-only + + def __init__( + self, + name: str, + credential: AsyncTokenCredential | None = None, + endpoint: FoundryStorageEndpoint | str | None = None, + *, + api_version: str = "v1", + description: str | None = ..., + item_ttl_seconds: int = DEFAULT_ITEM_TTL_SECONDS, + tags: Mapping[str, str] | None = ..., + user_id: str | None = ..., + user_isolation: bool = False, + **kwargs: Any + ) -> None: ... + + @classmethod + async def get_or_create( + cls, + name: str, + credential: AsyncTokenCredential | None = None, + endpoint: FoundryStorageEndpoint | str | None = None, + *, + api_version: str = "v1", + description: str | None = ..., + item_ttl_seconds: int = DEFAULT_ITEM_TTL_SECONDS, + tags: Mapping[str, str] | None = ..., + user_id: str | None = ..., + user_isolation: bool = False, + **kwargs: Any + ) -> FoundryStateStore: ... + + async def aclose(self) -> None: ... + + async def create_item( + self, + key: str, + value: JSONObject, + *, + tags: Mapping[str, str] | None = ... + ) -> StateStoreItemRef: ... + + async def delete(self) -> DeletedStateStore: ... + + async def delete_item( + self, + key: str, + *, + if_match: str | None = ... + ) -> DeletedStateStoreItem: ... + + async def get(self) -> StateStore: ... + + async def get_item(self, key: str) -> StateStoreItem | None: ... + + async def list_keys( + self, + *, + after: str | None = ..., + before: str | None = ..., + limit: int | None = ..., + order: Order = "desc", + tags: Mapping[str, str] | None = ... + ) -> StateStoreItemKeyPage: ... + + async def set_item( + self, + key: str, + value: JSONObject, + *, + if_match: str | None = ..., + require_exists: bool = False, + tags: Mapping[str, str] | None = ... + ) -> StateStoreItemRef: ... + + async def update( + self, + *, + description: str | None | object = _UNSET, + tags: Mapping[str, str] | None | object = _UNSET + ) -> StateStore: ... + + + class azure.ai.agentserver.core.storage.FoundryStorageApiError(FoundryStorageError): + + def __init__( + self, + message: str, + *, + response_body: dict[str, Any] | None = ..., + status_code: int | None = ... + ) -> None: ... + + + class azure.ai.agentserver.core.storage.FoundryStorageBadRequestError(FoundryStorageError): + + def __init__( + self, + message: str, + *, + param: str | None = ..., + response_body: dict[str, Any] | None = ..., + status_code: int | None = ... + ) -> None: ... + + + class azure.ai.agentserver.core.storage.FoundryStorageClient: implements AsyncContextManager + + def __init__( + self, + credential: AsyncTokenCredential, + endpoint: FoundryStorageEndpoint, + *, + get_server_version: Callable[[], str] | None = ..., + sdk_moniker: str | None = ..., + **kwargs: Any + ) -> None: ... + + async def aclose(self) -> None: ... + + + class azure.ai.agentserver.core.storage.FoundryStorageConflictError(FoundryStorageBadRequestError): + + def __init__( + self, + message: str, + *, + param: str | None = ..., + response_body: dict[str, Any] | None = ..., + status_code: int | None = ... + ) -> None: ... + + + class azure.ai.agentserver.core.storage.FoundryStorageEndpoint: + + def __init__( + self, + *, + api_version: str = _DEFAULT_API_VERSION, + storage_base_url: str + ) -> None: ... + + @classmethod + def from_endpoint( + cls, + endpoint: str, + *, + api_version: str = _DEFAULT_API_VERSION + ) -> FoundryStorageEndpoint: ... + + @classmethod + def from_env( + cls, + *, + api_version: str = _DEFAULT_API_VERSION + ) -> FoundryStorageEndpoint: ... + + def build_url( + self, + path: str, + **extra_params: str + ) -> str: ... + + + class azure.ai.agentserver.core.storage.FoundryStorageError(Exception): + + def __init__( + self, + message: str, + *, + response_body: dict[str, Any] | None = ..., + status_code: int | None = ... + ) -> None: ... + + + class azure.ai.agentserver.core.storage.FoundryStorageNotFoundError(FoundryStorageError): + + def __init__( + self, + message: str, + *, + response_body: dict[str, Any] | None = ..., + status_code: int | None = ... + ) -> None: ... + + + class azure.ai.agentserver.core.storage.FoundryStoragePreconditionError(FoundryStorageError): + + def __init__( + self, + message: str, + *, + current_etag: str | None = ..., + response_body: dict[str, Any] | None = ..., + status_code: int | None = ... + ) -> None: ... + + + class azure.ai.agentserver.core.storage.StateStore(_Model): + created_at: int + description: Optional[str] + id: str + item_ttl_seconds: int + name: str + object: Literal[StateStoreObjectType.STATE_STORE] + tags: Optional[dict[str, str]] + updated_at: int + user_isolation: bool + + @overload + def __init__( + self, + *, + created_at: int, + description: Optional[str] = ..., + id: str, + item_ttl_seconds: int, + name: str, + object: Literal[StateStoreObjectType.STATE_STORE], + tags: Optional[dict[str, str]] = ..., + updated_at: int, + user_isolation: bool + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: ... + + + class azure.ai.agentserver.core.storage.StateStoreItem(_Model): + created_at: int + etag: str + id: str + key: str + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] + tags: Optional[dict[str, str]] + updated_at: int + value: dict[str, Any] + + @overload + def __init__( + self, + *, + created_at: int, + etag: str, + id: str, + key: str, + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM], + tags: Optional[dict[str, str]] = ..., + updated_at: int, + value: dict[str, Any] + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: ... + + + class azure.ai.agentserver.core.storage.StateStoreItemKey(_Model): + created_at: int + etag: str + id: str + key: str + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] + tags: Optional[dict[str, str]] + updated_at: int + + @overload + def __init__( + self, + *, + created_at: int, + etag: str, + id: str, + key: str, + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM], + tags: Optional[dict[str, str]] = ..., + updated_at: int + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: ... + + + @dataclass(eq = True, frozen = False, init = True, kw_only = False, match_args = True, order = False, repr = True, slots = False, unsafe_hash = False, weakref_slot = False) + class azure.ai.agentserver.core.storage.StateStoreItemKeyPage: + first_id: Optional[str] + has_more: bool = field(compare = True, default = False, hash = None, init = True, kw_only = False, metadata = {}, name = "has_more", repr = True, type = "bool") + keys: list[StateStoreItemKey] + last_id: Optional[str] + + def __eq__() -> None: ... + + def __init__( + keys: list, + first_id: str | None, + last_id: str | None, + has_more: bool + ): ... + + def __repr__() -> None: ... + + + class azure.ai.agentserver.core.storage.StateStoreItemRef(_Model): + created_at: int + etag: str + id: str + key: str + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] + updated_at: int + + @overload + def __init__( + self, + *, + created_at: int, + etag: str, + id: str, + key: str, + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM], + updated_at: int + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: ... + + ``` \ No newline at end of file diff --git a/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml b/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml index b6ab0272d06d..12626e5b0ed6 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml +++ b/sdk/agentserver/azure-ai-agentserver-core/api.metadata.yml @@ -1,3 +1,3 @@ -apiMdSha256: cbbfb5e77c255601f97cc31505d7e0962ac3eeb665c690659064b2bd4940071e +apiMdSha256: a7869b1e01d125984ee825b37ed557471186b73ebfe4a0cced95a67bb063b04c parserVersion: 0.3.30 -pythonVersion: 3.11.9 +pythonVersion: 3.12.10 diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py index 1c08b10f13c6..a5fc24835fd1 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py @@ -1,28 +1,7 @@ # --------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -"""Azure AI AgentServerHost core framework. - -Provides the :class:`AgentServerHost` base class and shared utilities for -building Azure AI Hosted Agent containers. - -Public API:: - - from azure.ai.agentserver.core import ( - AgentConfig, - AgentServerHost, - FoundryAgentRequestContext, - configure_observability, - create_error_response, - detach_context, - end_span, - flush_spans, - get_request_context, - record_error, - set_current_span, - trace_stream, - ) -""" +"""Public API surface for the Azure AI Agent Server core framework.""" __path__ = __import__("pkgutil").extend_path(__path__, __name__) from ._base import AgentServerHost diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/__init__.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/__init__.py new file mode 100644 index 000000000000..9c57f54a3676 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/__init__.py @@ -0,0 +1,51 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Durable state-store client for AgentServer.""" + +from ._client import FOUNDRY_TOKEN_SCOPE, FoundryStorageClient +from ._endpoint import FoundryStorageEndpoint +from ._errors import ( + FoundryStorageApiError, + FoundryStorageBadRequestError, + FoundryStorageConflictError, + FoundryStorageError, + FoundryStorageNotFoundError, + FoundryStoragePreconditionError, +) +from ._state import DEFAULT_ITEM_TTL_SECONDS, FoundryStateStore +from ._state_serializer import ( + DeletedStateStore, + DeletedStateStoreItem, + JSONObject, + JSONValue, + StateStoreItemKeyPage, + Order, + StateStore, + StateStoreItem, + StateStoreItemRef, + StateStoreItemKey, +) + +__all__ = [ + "DEFAULT_ITEM_TTL_SECONDS", + "DeletedStateStore", + "DeletedStateStoreItem", + "FOUNDRY_TOKEN_SCOPE", + "FoundryStateStore", + "FoundryStorageApiError", + "FoundryStorageBadRequestError", + "FoundryStorageConflictError", + "FoundryStorageClient", + "FoundryStorageEndpoint", + "FoundryStorageError", + "FoundryStorageNotFoundError", + "FoundryStoragePreconditionError", + "JSONObject", + "JSONValue", + "StateStoreItemKeyPage", + "Order", + "StateStore", + "StateStoreItem", + "StateStoreItemRef", + "StateStoreItemKey", +] diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py new file mode 100644 index 000000000000..bf33c482e752 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py @@ -0,0 +1,111 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Protocol-neutral HTTP transport for Foundry-backed storage clients. + +:class:`FoundryStorageClient` owns the :class:`~azure.core.AsyncPipelineClient`, +its policy chain (retry, auth, logging, tracing), and the shared +``_send_storage_request`` helper. Protocol packages subclass it and add their +own resource methods and payload (de)serialization on top. +""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype +# pylint: disable=client-accepts-api-version-keyword + +from __future__ import annotations + +from typing import Any, Callable + +from azure.core import AsyncPipelineClient +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.exceptions import ServiceRequestError, ServiceResponseError +from azure.core.pipeline import policies +from azure.core.rest import AsyncHttpResponse, HttpRequest + +from azure.ai.agentserver.core._platform_headers import PLATFORM_ERROR_TAG +from azure.ai.agentserver.core._version import VERSION + +from ._endpoint import FoundryStorageEndpoint +from ._errors import raise_for_storage_error +from ._policies import FoundryStorageLoggingPolicy, ServerVersionUserAgentPolicy + +#: OAuth scope used to acquire bearer tokens for the Foundry storage API. +FOUNDRY_TOKEN_SCOPE = "https://ai.azure.com/.default" +#: Content type for JSON request bodies sent to the Foundry storage API. +JSON_CONTENT_TYPE = "application/json; charset=utf-8" + + +class FoundryStorageClient: + """Base HTTP client for the Foundry storage API. + + :param credential: An async credential used to obtain bearer tokens. + :type credential: AsyncTokenCredential + :param endpoint: Resolved Foundry storage endpoint configuration. + :type endpoint: FoundryStorageEndpoint + :param get_server_version: Optional callable returning the server version + string to use as the ``User-Agent`` header, evaluated lazily on each + request. When ``None``, Azure Core's default ``UserAgentPolicy`` is used. + :type get_server_version: Callable[[], str] | None + :param sdk_moniker: SDK moniker for the default ``UserAgentPolicy``. Ignored + when *get_server_version* is supplied. + :type sdk_moniker: str | None + """ + + def __init__( + self, + credential: AsyncTokenCredential, + endpoint: FoundryStorageEndpoint, + *, + get_server_version: Callable[[], str] | None = None, + sdk_moniker: str | None = None, + **kwargs: Any, + ) -> None: + self._endpoint = endpoint + + ua_policy: policies.UserAgentPolicy | ServerVersionUserAgentPolicy + if get_server_version is not None: + ua_policy = ServerVersionUserAgentPolicy(get_server_version) + else: + ua_policy = policies.UserAgentPolicy(sdk_moniker=sdk_moniker or f"ai-agentserver-core/{VERSION}") + + self._client: AsyncPipelineClient = AsyncPipelineClient( + base_url=endpoint.storage_base_url, + policies=[ + policies.RequestIdPolicy(), + policies.HeadersPolicy(), + ua_policy, + policies.AsyncRetryPolicy(), + policies.AsyncBearerTokenCredentialPolicy(credential, FOUNDRY_TOKEN_SCOPE), + FoundryStorageLoggingPolicy(), + # NOTE: ``ContentDecodePolicy`` is intentionally NOT included. It + # eagerly decodes every response body as JSON and crashes with + # ``UnicodeDecodeError`` on non-UTF-8 bodies (gzip payloads, HTML + # error pages, transport-corrupted bodies). Serializers call + # ``http_resp.text()`` directly with defensive error handling. + policies.DistributedTracingPolicy(), + ], + **kwargs, + ) + + async def aclose(self) -> None: + """Close the underlying HTTP pipeline client.""" + await self._client.close() + + async def __aenter__(self) -> "FoundryStorageClient": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.aclose() + + async def _send_storage_request(self, request: HttpRequest) -> AsyncHttpResponse: + """Send an HTTP request to the Foundry storage API. + + Transport-level exceptions are tagged as platform errors before being + re-raised; non-2xx responses raise a :class:`FoundryStorageError`. + """ + try: + http_resp: AsyncHttpResponse = await self._client.send_request(request) + except (ServiceRequestError, ServiceResponseError, OSError) as exc: + setattr(exc, PLATFORM_ERROR_TAG, True) + raise + raise_for_storage_error(http_resp) + return http_resp diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_endpoint.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_endpoint.py new file mode 100644 index 000000000000..8d32245a7c6c --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_endpoint.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Endpoint configuration shared by Foundry storage clients. + +:class:`FoundryStorageEndpoint` is protocol-neutral: it resolves the Foundry +project endpoint, derives the ``/storage/`` base URL, and builds versioned +request URLs for any storage resource path (responses, activity state, ...). +""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype +# pylint: disable=docstring-keyword-should-match-keyword-only + +from __future__ import annotations + +from urllib.parse import quote as _url_quote + +from azure.ai.agentserver.core._config import AgentConfig + +_DEFAULT_API_VERSION = "v1" + + +def _encode(value: str) -> str: + return _url_quote(value, safe="") + + +class FoundryStorageEndpoint: + """Immutable Foundry storage endpoint configuration. + + :param storage_base_url: Fully-qualified ``.../storage/`` base URL. + :type storage_base_url: str + :param api_version: Storage API version appended to every request URL. + :type api_version: str + """ + + def __init__(self, *, storage_base_url: str, api_version: str = _DEFAULT_API_VERSION) -> None: + self.storage_base_url = storage_base_url + self.api_version = api_version + + @classmethod + def from_env(cls, *, api_version: str = _DEFAULT_API_VERSION) -> "FoundryStorageEndpoint": + """Create an endpoint by reading the ``FOUNDRY_PROJECT_ENDPOINT`` environment variable.""" + config = AgentConfig.from_env() + if not config.project_endpoint: + raise EnvironmentError( + "The 'FOUNDRY_PROJECT_ENDPOINT' environment variable is required. " + "In hosted environments, the Azure AI Foundry platform must set this variable." + ) + return cls.from_endpoint(config.project_endpoint, api_version=api_version) + + @classmethod + def from_endpoint(cls, endpoint: str, *, api_version: str = _DEFAULT_API_VERSION) -> "FoundryStorageEndpoint": + """Create an endpoint from an explicit Foundry project endpoint URL.""" + if not endpoint: + raise ValueError("endpoint must be a non-empty string") + if not (endpoint.startswith("http://") or endpoint.startswith("https://")): + raise ValueError(f"endpoint must be a valid absolute URL, got: {endpoint!r}") + normalized = endpoint.rstrip("/") + if normalized.endswith("/storage"): + storage_base_url = normalized + "/" + else: + storage_base_url = normalized + "/storage/" + return cls(storage_base_url=storage_base_url, api_version=api_version) + + def build_url(self, path: str, **extra_params: str) -> str: + """Build a full storage API URL for *path* with ``api-version`` appended.""" + url = f"{self.storage_base_url}{path}?api-version={_encode(self.api_version)}" + for key, value in extra_params.items(): + url += f"&{_encode(key)}={_encode(value)}" + return url diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_errors.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_errors.py new file mode 100644 index 000000000000..8901cc3624b3 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_errors.py @@ -0,0 +1,123 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Exception hierarchy for Foundry storage API errors.""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, Union + +from azure.ai.agentserver.core._platform_headers import PLATFORM_ERROR_TAG + +if TYPE_CHECKING: + from azure.core.rest import AsyncHttpResponse, HttpResponse + + _AnyHttpResponse = Union[HttpResponse, AsyncHttpResponse] + + +class FoundryStorageError(Exception): + """Base class for errors returned by the Foundry storage API.""" + + def __init__( + self, + message: str, + *, + response_body: dict[str, Any] | None = None, + status_code: int | None = None, + ) -> None: + super().__init__(message) + self.message = message + self.response_body = response_body + self.status_code = status_code + + +class FoundryStorageNotFoundError(FoundryStorageError): + """Raised when the requested resource does not exist (HTTP 404).""" + + +class FoundryStorageBadRequestError(FoundryStorageError): + """Raised for invalid-request errors (HTTP 400).""" + + def __init__( + self, + message: str, + *, + response_body: dict[str, Any] | None = None, + status_code: int | None = None, + param: str | None = None, + ) -> None: + super().__init__(message, response_body=response_body, status_code=status_code) + self.param = param + + +class FoundryStorageConflictError(FoundryStorageBadRequestError): + """Raised when the requested create/update conflicts with an existing resource (HTTP 409).""" + + +class FoundryStoragePreconditionError(FoundryStorageError): + """Raised when an ``If-Match`` precondition fails on a single-item write.""" + + def __init__( + self, + message: str, + *, + response_body: dict[str, Any] | None = None, + status_code: int | None = None, + current_etag: str | None = None, + ) -> None: + super().__init__(message, response_body=response_body, status_code=status_code) + self.current_etag = current_etag + + +class FoundryStorageApiError(FoundryStorageError): + """Raised for all other non-success HTTP responses.""" + + +def raise_for_storage_error(response: "_AnyHttpResponse") -> None: + """Raise an appropriate :class:`FoundryStorageError` subclass for non-2xx responses.""" + status = response.status_code + if 200 <= status < 300: + return + + message, body = _extract_error_message(response, status) + + if status == 404: + raise FoundryStorageNotFoundError(message, response_body=body, status_code=status) + if status == 412: + raise FoundryStoragePreconditionError( + message, + response_body=body, + status_code=status, + current_etag=response.headers.get("ETag"), + ) + if status == 400: + param = None + if isinstance(body, dict): + error = body.get("error") + if isinstance(error, dict): + raw_param = error.get("param") + param = str(raw_param) if raw_param is not None else None + raise FoundryStorageBadRequestError(message, response_body=body, status_code=status, param=param) + if status == 409: + raise FoundryStorageConflictError(message, response_body=body, status_code=status) + exc = FoundryStorageApiError(message, response_body=body, status_code=status) + setattr(exc, PLATFORM_ERROR_TAG, True) + raise exc + + +def _extract_error_message(response: "_AnyHttpResponse", status: int) -> tuple[str, dict[str, Any] | None]: + """Extract a Foundry error-envelope message from *response* when possible.""" + try: + body = response.text() + if body: + data = json.loads(body) + error = data.get("error") + if isinstance(error, dict): + message = error.get("message") + if message: + return str(message), data + except Exception: # pylint: disable=broad-exception-caught + pass + return f"Foundry storage request failed with HTTP {status}.", None diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/__init__.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/__init__.py new file mode 100644 index 000000000000..c951ab34f2f9 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/__init__.py @@ -0,0 +1,45 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Request/response model classes generated from the Foundry storage TypeSpec contract. + +The contract is published in ``Azure/azure-rest-api-specs`` (specification for +Foundry storage); until it is merged there, these models are generated locally +and vendored into this package. Do not import from this package directly outside +``storage/``; the public, documented names are re-exported (with SDK-friendly +aliases matching earlier releases where they differ from the generated/spec +names) from ``azure.ai.agentserver.core.storage``. +""" + +from ._models import ( + ApiError, + ApiErrorResponse, + CreateItemRequest, + CreateStateStoreRequest, + DeletedStateStore, + DeletedStateStoreItem, + ListResponseStateStore, + ListResponseStateStoreItemKey, + PutItemRequest, + StateStore, + StateStoreItem, + StateStoreItemRef, + StateStoreItemKey, + UpdateStateStoreRequest, +) + +__all__ = [ + "ApiError", + "ApiErrorResponse", + "CreateItemRequest", + "CreateStateStoreRequest", + "DeletedStateStore", + "DeletedStateStoreItem", + "ListResponseStateStore", + "ListResponseStateStoreItemKey", + "PutItemRequest", + "StateStore", + "StateStoreItem", + "StateStoreItemRef", + "StateStoreItemKey", + "UpdateStateStoreRequest", +] diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_enums.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_enums.py new file mode 100644 index 000000000000..9568d7539498 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_enums.py @@ -0,0 +1,24 @@ +# coding=utf-8 +# -------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# Code generated by Microsoft (R) Python Code Generator. +# Changes may cause incorrect behavior and will be lost if the code is regenerated. +# -------------------------------------------------------------------------- + +from enum import Enum +from azure.core import CaseInsensitiveEnumMeta + + +class StateStoreItemObjectType(str, Enum, metaclass=CaseInsensitiveEnumMeta): + """The ``object`` discriminator for an item resource.""" + + STATE_STORE_ITEM = "state_store.item" + """STATE_STORE_ITEM.""" + + +class StateStoreObjectType(str, Enum, metaclass=CaseInsensitiveEnumMeta): + """The ``object`` discriminator for a state-store resource.""" + + STATE_STORE = "state_store" + """STATE_STORE.""" diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_models.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_models.py new file mode 100644 index 000000000000..3df1d9c3c068 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_models.py @@ -0,0 +1,683 @@ +# coding=utf-8 +# -------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# Code generated by Microsoft (R) Python Code Generator. +# Changes may cause incorrect behavior and will be lost if the code is regenerated. +# -------------------------------------------------------------------------- +# pylint: disable=useless-super-delegation + +from typing import Any, Literal, Mapping, Optional, TYPE_CHECKING, overload + +from ._utils.model_base import Model as _Model, rest_field +from ._enums import StateStoreItemObjectType, StateStoreObjectType + +if TYPE_CHECKING: + from . import _models as _models + + +class ApiError(_Model): + """ApiError. + + :ivar code: Machine-readable error code. Required. + :vartype code: str + :ivar message: Human-readable error message. Required. + :vartype message: str + :ivar param: For 400 responses, the name of the offending request field. + :vartype param: str + """ + + code: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Machine-readable error code. Required.""" + message: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Human-readable error message. Required.""" + param: Optional[str] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """For 400 responses, the name of the offending request field.""" + + @overload + def __init__( + self, + *, + code: str, + message: str, + param: Optional[str] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class ApiErrorResponse(_Model): + """Standard Foundry error envelope (parent spec §8.1). + + :ivar error: Required. + :vartype error: ~foundry.storage.statestore.models.ApiError + """ + + error: "_models.ApiError" = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + + @overload + def __init__( + self, + *, + error: "_models.ApiError", + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class CreateItemRequest(_Model): + """Request body for ``createItem`` (spec §4.2.1) — fails with ``409`` on a duplicate key. + + :ivar key: Required. + :vartype key: str + :ivar value: Required. + :vartype value: dict[str, any] + :ivar tags: + :vartype tags: dict[str, str] + """ + + key: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + value: dict[str, Any] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + tags: Optional[dict[str, str]] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + + @overload + def __init__( + self, + *, + key: str, + value: dict[str, Any], + tags: Optional[dict[str, str]] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class CreateStateStoreRequest(_Model): + """Request body for ``createStateStore`` (spec §4.1.1). + + :ivar name: Required. Unique within [project_id, agent_guid]. Required. + :vartype name: str + :ivar user_isolation: Optional. Omitted -> ``false`` (agent-level, shared). Fixed at create. + :vartype user_isolation: bool + :ivar item_ttl_seconds: Optional. Default 30 days (2592000). ``-1`` = never expire. Fixed at + create. + :vartype item_ttl_seconds: int + :ivar description: + :vartype description: str + :ivar tags: + :vartype tags: dict[str, str] + """ + + name: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required. Unique within [project_id, agent_guid]. Required.""" + user_isolation: Optional[bool] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Optional. Omitted -> ``false`` (agent-level, shared). Fixed at create.""" + item_ttl_seconds: Optional[int] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Optional. Default 30 days (2592000). ``-1`` = never expire. Fixed at create.""" + description: Optional[str] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + tags: Optional[dict[str, str]] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + + @overload + def __init__( + self, + *, + name: str, + user_isolation: Optional[bool] = None, + item_ttl_seconds: Optional[int] = None, + description: Optional[str] = None, + tags: Optional[dict[str, str]] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class DeletedStateStore(_Model): + """Response for ``deleteStateStore`` (spec §4.1.5). Idempotent. + + :ivar id: Present when the store existed; absent when the delete was a no-op (already gone). + :vartype id: str + :ivar object: Always ``"state_store"``. Required. STATE_STORE. + :vartype object: str or ~foundry.storage.statestore.models.STATE_STORE + :ivar name: Required. + :vartype name: str + :ivar deleted: Always ``true``. Required. + :vartype deleted: bool + """ + + id: Optional[str] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Present when the store existed; absent when the delete was a no-op (already gone).""" + object: Literal[StateStoreObjectType.STATE_STORE] = rest_field( + visibility=["read", "create", "update", "delete", "query"] + ) + """Always ``\"state_store\"``. Required. STATE_STORE.""" + name: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + deleted: bool = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Always ``true``. Required.""" + + @overload + def __init__( + self, + *, + object: Literal[StateStoreObjectType.STATE_STORE], + name: str, + deleted: bool, + id: Optional[str] = None, # pylint: disable=redefined-builtin + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class DeletedStateStoreItem(_Model): + """Response for ``deleteItem`` (spec §4.2.4). Idempotent. + + :ivar id: Present when the item existed; absent when the delete was a no-op (already gone). + :vartype id: str + :ivar object: Always ``"state_store.item"``. Required. STATE_STORE_ITEM. + :vartype object: str or ~foundry.storage.statestore.models.STATE_STORE_ITEM + :ivar key: Required. + :vartype key: str + :ivar deleted: Always ``true``. Required. + :vartype deleted: bool + """ + + id: Optional[str] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Present when the item existed; absent when the delete was a no-op (already gone).""" + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] = rest_field( + visibility=["read", "create", "update", "delete", "query"] + ) + """Always ``\"state_store.item\"``. Required. STATE_STORE_ITEM.""" + key: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + deleted: bool = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Always ``true``. Required.""" + + @overload + def __init__( + self, + *, + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM], + key: str, + deleted: bool, + id: Optional[str] = None, # pylint: disable=redefined-builtin + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class ListResponseStateStore(_Model): + """A page of results, ordered by ``created_at`` (parent spec §7). + + :ivar object: Always ``"list"``. Required. Default value is "list". + :vartype object: str + :ivar data: The page of results. Required. + :vartype data: list[~foundry.storage.statestore.models.StateStore] + :ivar first_id: The ``id`` of the first item in ``data``, or ``null`` if ``data`` is empty. + Required. + :vartype first_id: str + :ivar last_id: The ``id`` of the last item in ``data``, or ``null`` if ``data`` is empty. + Required. + :vartype last_id: str + :ivar has_more: Whether another page is available after ``last_id``. Required. + :vartype has_more: bool + """ + + object: Literal["list"] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Always ``\"list\"``. Required. Default value is \"list\".""" + data: list["_models.StateStore"] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """The page of results. Required.""" + first_id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """The ``id`` of the first item in ``data``, or ``null`` if ``data`` is empty. Required.""" + last_id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """The ``id`` of the last item in ``data``, or ``null`` if ``data`` is empty. Required.""" + has_more: bool = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Whether another page is available after ``last_id``. Required.""" + + @overload + def __init__( + self, + *, + data: list["_models.StateStore"], + first_id: str, + last_id: str, + has_more: bool, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.object: Literal["list"] = "list" + + +class ListResponseStateStoreItemKey(_Model): + """A page of results, ordered by ``created_at`` (parent spec §7). + + :ivar object: Always ``"list"``. Required. Default value is "list". + :vartype object: str + :ivar data: The page of results. Required. + :vartype data: list[~foundry.storage.statestore.models.StateStoreItemKey] + :ivar first_id: The ``id`` of the first item in ``data``, or ``null`` if ``data`` is empty. + Required. + :vartype first_id: str + :ivar last_id: The ``id`` of the last item in ``data``, or ``null`` if ``data`` is empty. + Required. + :vartype last_id: str + :ivar has_more: Whether another page is available after ``last_id``. Required. + :vartype has_more: bool + """ + + object: Literal["list"] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Always ``\"list\"``. Required. Default value is \"list\".""" + data: list["_models.StateStoreItemKey"] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """The page of results. Required.""" + first_id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """The ``id`` of the first item in ``data``, or ``null`` if ``data`` is empty. Required.""" + last_id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """The ``id`` of the last item in ``data``, or ``null`` if ``data`` is empty. Required.""" + has_more: bool = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Whether another page is available after ``last_id``. Required.""" + + @overload + def __init__( + self, + *, + data: list["_models.StateStoreItemKey"], + first_id: str, + last_id: str, + has_more: bool, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.object: Literal["list"] = "list" + + +class PutItemRequest(_Model): + """Request body for ``putItem`` (spec §4.2.2) — idempotent create-or-replace by key. + + :ivar value: Required. + :vartype value: dict[str, any] + :ivar tags: + :vartype tags: dict[str, str] + """ + + value: dict[str, Any] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + tags: Optional[dict[str, str]] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + + @overload + def __init__( + self, + *, + value: dict[str, Any], + tags: Optional[dict[str, str]] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class StateStore(_Model): + """A first-class, explicitly-created durable key-value store (spec §2, §4.1). + + :ivar id: Server-assigned, read-only. Not used for addressing (the ``name`` is). Required. + :vartype id: str + :ivar object: Always ``"state_store"``. Required. STATE_STORE. + :vartype object: str or ~foundry.storage.statestore.models.STATE_STORE + :ivar name: Caller-chosen logical name. Unique within [project_id, agent_guid]. May contain + ``/`` as a hierarchy separator. Immutable. Required. + :vartype name: str + :ivar user_isolation: Whether item operations are partitioned per resolved user (spec §3). + Omitted at create -> ``false``. Immutable. Required. + :vartype user_isolation: bool + :ivar item_ttl_seconds: Store-wide default TTL (seconds) inherited by every item, write-sliding + per item. ``-1`` = never expire. Immutable. Required. + :vartype item_ttl_seconds: int + :ivar description: Optional free-form description. Mutable via ``updateStateStore``. + :vartype description: str + :ivar tags: Optional string-to-string labels. Mutable (replaced wholesale) via + ``updateStateStore``. + :vartype tags: dict[str, str] + :ivar created_at: Creation time. Required. + :vartype created_at: int + :ivar updated_at: Last metadata-update time. Required. + :vartype updated_at: int + """ + + id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Server-assigned, read-only. Not used for addressing (the ``name`` is). Required.""" + object: Literal[StateStoreObjectType.STATE_STORE] = rest_field( + visibility=["read", "create", "update", "delete", "query"] + ) + """Always ``\"state_store\"``. Required. STATE_STORE.""" + name: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Caller-chosen logical name. Unique within [project_id, agent_guid]. May contain ``/`` as a + hierarchy separator. Immutable. Required.""" + user_isolation: bool = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Whether item operations are partitioned per resolved user (spec §3). Omitted at create -> + ``false``. Immutable. Required.""" + item_ttl_seconds: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Store-wide default TTL (seconds) inherited by every item, write-sliding per item. ``-1`` = + never expire. Immutable. Required.""" + description: Optional[str] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Optional free-form description. Mutable via ``updateStateStore``.""" + tags: Optional[dict[str, str]] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Optional string-to-string labels. Mutable (replaced wholesale) via ``updateStateStore``.""" + created_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Creation time. Required.""" + updated_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Last metadata-update time. Required.""" + + @overload + def __init__( + self, + *, + id: str, # pylint: disable=redefined-builtin + object: Literal[StateStoreObjectType.STATE_STORE], + name: str, + user_isolation: bool, + item_ttl_seconds: int, + created_at: int, + updated_at: int, + description: Optional[str] = None, + tags: Optional[dict[str, str]] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class StateStoreItem(_Model): + """A single ``key -> value`` record within a state store (spec §2, §4.2). + + :ivar id: Server-assigned, read-only. Not used for addressing (the ``key`` is.). Required. + :vartype id: str + :ivar object: Always ``"state_store.item"``. Required. STATE_STORE_ITEM. + :vartype object: str or ~foundry.storage.statestore.models.STATE_STORE_ITEM + :ivar key: Logical key, unique within the store. May contain ``/``. Required. + :vartype key: str + :ivar value: Opaque JSON object; stored and returned verbatim. <= 1 MB serialized inline. + Required. + :vartype value: dict[str, any] + :ivar tags: Optional labels used to filter ``listKeys``. Replaced wholesale on write. + :vartype tags: dict[str, str] + :ivar etag: Optimistic-concurrency token for this item; also echoed in the ``ETag`` response + header. Required. + :vartype etag: str + :ivar created_at: Required. + :vartype created_at: int + :ivar updated_at: Required. + :vartype updated_at: int + """ + + id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Server-assigned, read-only. Not used for addressing (the ``key`` is.). Required.""" + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] = rest_field( + visibility=["read", "create", "update", "delete", "query"] + ) + """Always ``\"state_store.item\"``. Required. STATE_STORE_ITEM.""" + key: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Logical key, unique within the store. May contain ``/``. Required.""" + value: dict[str, Any] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Opaque JSON object; stored and returned verbatim. <= 1 MB serialized inline. Required.""" + tags: Optional[dict[str, str]] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Optional labels used to filter ``listKeys``. Replaced wholesale on write.""" + etag: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Optimistic-concurrency token for this item; also echoed in the ``ETag`` response header. + Required.""" + created_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + updated_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + + @overload + def __init__( + self, + *, + id: str, # pylint: disable=redefined-builtin + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM], + key: str, + value: dict[str, Any], + etag: str, + created_at: int, + updated_at: int, + tags: Optional[dict[str, str]] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class StateStoreItemKey(_Model): + """A keys-only projection entry returned by ``listKeys`` (spec §4.2.5) -- never includes + ``value``. + + :ivar id: Required. + :vartype id: str + :ivar object: Always ``"state_store.item"``. Required. STATE_STORE_ITEM. + :vartype object: str or ~foundry.storage.statestore.models.STATE_STORE_ITEM + :ivar key: Required. + :vartype key: str + :ivar tags: + :vartype tags: dict[str, str] + :ivar etag: Required. + :vartype etag: str + :ivar created_at: Required. + :vartype created_at: int + :ivar updated_at: Required. + :vartype updated_at: int + """ + + id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] = rest_field( + visibility=["read", "create", "update", "delete", "query"] + ) + """Always ``\"state_store.item\"``. Required. STATE_STORE_ITEM.""" + key: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + tags: Optional[dict[str, str]] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + etag: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + created_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + updated_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + + @overload + def __init__( + self, + *, + id: str, # pylint: disable=redefined-builtin + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM], + key: str, + etag: str, + created_at: int, + updated_at: int, + tags: Optional[dict[str, str]] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class StateStoreItemRef(_Model): + """Response for ``createItem`` / ``putItem`` (spec §4.2.1/§4.2.2) — server metadata only, no + ``value`` echoed back. + + :ivar id: Required. + :vartype id: str + :ivar object: Always ``"state_store.item"``. Required. STATE_STORE_ITEM. + :vartype object: str or ~foundry.storage.statestore.models.STATE_STORE_ITEM + :ivar key: Required. + :vartype key: str + :ivar etag: Required. + :vartype etag: str + :ivar created_at: Required. + :vartype created_at: int + :ivar updated_at: Required. + :vartype updated_at: int + """ + + id: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM] = rest_field( + visibility=["read", "create", "update", "delete", "query"] + ) + """Always ``\"state_store.item\"``. Required. STATE_STORE_ITEM.""" + key: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + etag: str = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + created_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + updated_at: int = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + + @overload + def __init__( + self, + *, + id: str, # pylint: disable=redefined-builtin + object: Literal[StateStoreItemObjectType.STATE_STORE_ITEM], + key: str, + etag: str, + created_at: int, + updated_at: int, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + + +class UpdateStateStoreRequest(_Model): + """Request body for ``updateStateStore`` (spec §4.1.3). Both fields optional; a present field + replaces the stored value (tags replaced wholesale, not merged); an omitted field is left + unchanged. ``name``, ``user_isolation``, and ``item_ttl_seconds`` are immutable -- supplying + them is rejected with ``400``. + + :ivar description: + :vartype description: str + :ivar tags: + :vartype tags: dict[str, str] + """ + + description: Optional[str] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + tags: Optional[dict[str, str]] = rest_field(visibility=["read", "create", "update", "delete", "query"]) + + @overload + def __init__( + self, + *, + description: Optional[str] = None, + tags: Optional[dict[str, str]] = None, + ) -> None: ... + + @overload + def __init__(self, mapping: Mapping[str, Any]) -> None: + """ + :param mapping: raw JSON to initialize the model. + :type mapping: Mapping[str, Any] + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/__init__.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/__init__.py new file mode 100644 index 000000000000..8026245c2abc --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/__init__.py @@ -0,0 +1,6 @@ +# -------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# Code generated by Microsoft (R) Python Code Generator. +# Changes may cause incorrect behavior and will be lost if the code is regenerated. +# -------------------------------------------------------------------------- diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/model_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/model_base.py new file mode 100644 index 000000000000..b93f5120d517 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_utils/model_base.py @@ -0,0 +1,1770 @@ +# pylint: disable=line-too-long,useless-suppression,too-many-lines +# coding=utf-8 +# -------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# Code generated by Microsoft (R) Python Code Generator. +# Changes may cause incorrect behavior and will be lost if the code is regenerated. +# -------------------------------------------------------------------------- +# pylint: disable=protected-access, broad-except + +import copy +import calendar +import decimal +import functools +import sys +import logging +import base64 +import re +import typing +import enum +import email.utils +from datetime import datetime, date, time, timedelta, timezone +from json import JSONEncoder +import xml.etree.ElementTree as ET +from collections.abc import MutableMapping +import isodate +from azure.core.exceptions import DeserializationError +from azure.core import CaseInsensitiveEnumMeta +from azure.core.pipeline import PipelineResponse +from azure.core.serialization import _Null + +from azure.core.rest import HttpResponse + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +_LOGGER = logging.getLogger(__name__) + +__all__ = ["SdkJSONEncoder", "Model", "rest_field", "rest_discriminator"] + +TZ_UTC = timezone.utc +_T = typing.TypeVar("_T") +_NONE_TYPE = type(None) + + +def _timedelta_as_isostr(td: timedelta) -> str: + """Converts a datetime.timedelta object into an ISO 8601 formatted string, e.g. 'P4DT12H30M05S' + + Function adapted from the Tin Can Python project: https://github.com/RusticiSoftware/TinCanPython + + :param timedelta td: The timedelta to convert + :rtype: str + :return: ISO8601 version of this timedelta + """ + + # Split seconds to larger units + seconds = td.total_seconds() + minutes, seconds = divmod(seconds, 60) + hours, minutes = divmod(minutes, 60) + days, hours = divmod(hours, 24) + + days, hours, minutes = list(map(int, (days, hours, minutes))) + seconds = round(seconds, 6) + + # Build date + date_str = "" + if days: + date_str = "%sD" % days + + if hours or minutes or seconds: + # Build time + time_str = "T" + + # Hours + bigger_exists = date_str or hours + if bigger_exists: + time_str += "{:02}H".format(hours) + + # Minutes + bigger_exists = bigger_exists or minutes + if bigger_exists: + time_str += "{:02}M".format(minutes) + + # Seconds + try: + if seconds.is_integer(): + seconds_string = "{:02}".format(int(seconds)) + else: + # 9 chars long w/ leading 0, 6 digits after decimal + seconds_string = "%09.6f" % seconds + # Remove trailing zeros + seconds_string = seconds_string.rstrip("0") + except AttributeError: # int.is_integer() raises + seconds_string = "{:02}".format(seconds) + + time_str += "{}S".format(seconds_string) + else: + time_str = "" + + return "P" + date_str + time_str + + +def _serialize_bytes(o, format: typing.Optional[str] = None) -> str: + encoded = base64.b64encode(o).decode() + if format == "base64url": + return encoded.strip("=").replace("+", "-").replace("/", "_") + return encoded + + +def _serialize_duration(td: timedelta, format: typing.Optional[str] = None): + """Serialize a timedelta to its wire representation. + + For the ``seconds``/``milliseconds`` encodings the value is converted to a + numeric value, otherwise it falls back to an ISO 8601 duration string. + + :param timedelta td: The timedelta to serialize. + :param str format: The duration encoding format. + :rtype: int or float or str + :return: serialized duration + """ + seconds = td.total_seconds() + if format == "duration-seconds-int": + return int(seconds) + if format == "duration-seconds-float": + return seconds + if format == "duration-milliseconds-int": + return int(seconds * 1000) + if format == "duration-milliseconds-float": + return seconds * 1000 + return _timedelta_as_isostr(td) + + +def _serialize_datetime(o, format: typing.Optional[str] = None): + if hasattr(o, "year") and hasattr(o, "hour"): + if format == "rfc7231": + return email.utils.format_datetime(o, usegmt=True) + if format == "unix-timestamp": + return int(calendar.timegm(o.utctimetuple())) + + # astimezone() fails for naive times in Python 2.7, so make make sure o is aware (tzinfo is set) + if not o.tzinfo: + iso_formatted = o.replace(tzinfo=TZ_UTC).isoformat() + else: + iso_formatted = o.astimezone(TZ_UTC).isoformat() + # Replace the trailing "+00:00" UTC offset with "Z" (RFC 3339: https://www.ietf.org/rfc/rfc3339.txt) + return iso_formatted.replace("+00:00", "Z") + # Next try datetime.date or datetime.time + return o.isoformat() + + +def _is_readonly(p): + try: + return p._visibility == ["read"] + except AttributeError: + return False + + +class SdkJSONEncoder(JSONEncoder): + """A JSON encoder that's capable of serializing datetime objects and bytes.""" + + def __init__(self, *args, exclude_readonly: bool = False, format: typing.Optional[str] = None, **kwargs): + super().__init__(*args, **kwargs) + self.exclude_readonly = exclude_readonly + self.format = format + + def default(self, o): # pylint: disable=too-many-return-statements + if _is_model(o): + if self.exclude_readonly: + readonly_props = [p._rest_name for p in o._attr_to_rest_field.values() if _is_readonly(p)] + return {k: v for k, v in o.items() if k not in readonly_props} + return dict(o.items()) + try: + return super(SdkJSONEncoder, self).default(o) + except TypeError: + if isinstance(o, _Null): + return None + if isinstance(o, decimal.Decimal): + return float(o) + if isinstance(o, (bytes, bytearray)): + return _serialize_bytes(o, self.format) + try: + # First try datetime.datetime + return _serialize_datetime(o, self.format) + except AttributeError: + pass + # Last, try datetime.timedelta + try: + return _timedelta_as_isostr(o) + except AttributeError: + # This will be raised when it hits value.total_seconds in the method above + pass + return super(SdkJSONEncoder, self).default(o) + + +_VALID_DATE = re.compile(r"\d{4}[-]\d{2}[-]\d{2}T\d{2}:\d{2}:\d{2}" + r"\.?\d*Z?[-+]?[\d{2}]?:?[\d{2}]?") +_VALID_RFC7231 = re.compile( + r"(Mon|Tue|Wed|Thu|Fri|Sat|Sun),\s\d{2}\s" + r"(Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\s\d{4}\s\d{2}:\d{2}:\d{2}\sGMT" +) + +_ARRAY_ENCODE_MAPPING = { + "pipeDelimited": "|", + "spaceDelimited": " ", + "commaDelimited": ",", + "newlineDelimited": "\n", +} + + +def _deserialize_array_encoded(delimit: str, attr): + if isinstance(attr, str): + if attr == "": + return [] + return attr.split(delimit) + return attr + + +def _deserialize_datetime(attr: typing.Union[str, datetime]) -> datetime: + """Deserialize ISO-8601 formatted string into Datetime object. + + :param str attr: response string to be deserialized. + :rtype: ~datetime.datetime + :returns: The datetime object from that input + """ + if isinstance(attr, datetime): + # i'm already deserialized + return attr + attr = attr.upper() + match = _VALID_DATE.match(attr) + if not match: + raise ValueError("Invalid datetime string: " + attr) + + check_decimal = attr.split(".") + if len(check_decimal) > 1: + decimal_str = "" + for digit in check_decimal[1]: + if digit.isdigit(): + decimal_str += digit + else: + break + if len(decimal_str) > 6: + attr = attr.replace(decimal_str, decimal_str[0:6]) + + date_obj = isodate.parse_datetime(attr) + test_utc = date_obj.utctimetuple() + if test_utc.tm_year > 9999 or test_utc.tm_year < 1: + raise OverflowError("Hit max or min date") + return date_obj # type: ignore[no-any-return] + + +def _deserialize_datetime_rfc7231(attr: typing.Union[str, datetime]) -> datetime: + """Deserialize RFC7231 formatted string into Datetime object. + + :param str attr: response string to be deserialized. + :rtype: ~datetime.datetime + :returns: The datetime object from that input + """ + if isinstance(attr, datetime): + # i'm already deserialized + return attr + match = _VALID_RFC7231.match(attr) + if not match: + raise ValueError("Invalid datetime string: " + attr) + + return email.utils.parsedate_to_datetime(attr) + + +def _deserialize_datetime_unix_timestamp(attr: typing.Union[float, datetime]) -> datetime: + """Deserialize unix timestamp into Datetime object. + + :param str attr: response string to be deserialized. + :rtype: ~datetime.datetime + :returns: The datetime object from that input + """ + if isinstance(attr, datetime): + # i'm already deserialized + return attr + return datetime.fromtimestamp(attr, TZ_UTC) + + +def _deserialize_date(attr: typing.Union[str, date]) -> date: + """Deserialize ISO-8601 formatted string into Date object. + :param str attr: response string to be deserialized. + :rtype: date + :returns: The date object from that input + """ + # This must NOT use defaultmonth/defaultday. Using None ensure this raises an exception. + if isinstance(attr, date): + return attr + return isodate.parse_date(attr, defaultmonth=None, defaultday=None) # type: ignore + + +def _deserialize_time(attr: typing.Union[str, time]) -> time: + """Deserialize ISO-8601 formatted string into time object. + + :param str attr: response string to be deserialized. + :rtype: datetime.time + :returns: The time object from that input + """ + if isinstance(attr, time): + return attr + return isodate.parse_time(attr) # type: ignore[no-any-return] + + +def _deserialize_bytes(attr): + if isinstance(attr, (bytes, bytearray)): + return attr + return bytes(base64.b64decode(attr)) + + +def _deserialize_bytes_base64(attr): + if isinstance(attr, (bytes, bytearray)): + return attr + padding = "=" * (3 - (len(attr) + 3) % 4) # type: ignore + attr = attr + padding # type: ignore + encoded = attr.replace("-", "+").replace("_", "/") + return bytes(base64.b64decode(encoded)) + + +def _deserialize_duration(attr): + if isinstance(attr, timedelta): + return attr + return isodate.parse_duration(attr) + + +def _deserialize_duration_numeric(attr, unit): + if isinstance(attr, timedelta): + return attr + return timedelta(**{unit: float(attr)}) + + +def _deserialize_decimal(attr): + if isinstance(attr, decimal.Decimal): + return attr + return decimal.Decimal(str(attr)) + + +def _deserialize_int_as_str(attr): + if isinstance(attr, int): + return attr + return int(attr) + + +_DESERIALIZE_MAPPING = { + datetime: _deserialize_datetime, + date: _deserialize_date, + time: _deserialize_time, + bytes: _deserialize_bytes, + bytearray: _deserialize_bytes, + timedelta: _deserialize_duration, + typing.Any: lambda x: x, + decimal.Decimal: _deserialize_decimal, +} + +_DESERIALIZE_MAPPING_WITHFORMAT = { + "rfc3339": _deserialize_datetime, + "rfc7231": _deserialize_datetime_rfc7231, + "unix-timestamp": _deserialize_datetime_unix_timestamp, + "base64": _deserialize_bytes, + "base64url": _deserialize_bytes_base64, + "duration-seconds-int": functools.partial(_deserialize_duration_numeric, unit="seconds"), + "duration-seconds-float": functools.partial(_deserialize_duration_numeric, unit="seconds"), + "duration-milliseconds-int": functools.partial(_deserialize_duration_numeric, unit="milliseconds"), + "duration-milliseconds-float": functools.partial(_deserialize_duration_numeric, unit="milliseconds"), +} + + +def get_deserializer(annotation: typing.Any, rf: typing.Optional["_RestField"] = None): + if annotation is int and rf and rf._format == "str": + return _deserialize_int_as_str + if annotation is str and rf and rf._format in _ARRAY_ENCODE_MAPPING: + return functools.partial(_deserialize_array_encoded, _ARRAY_ENCODE_MAPPING[rf._format]) + if rf and rf._format: + return _DESERIALIZE_MAPPING_WITHFORMAT.get(rf._format) + return _DESERIALIZE_MAPPING.get(annotation) # pyright: ignore + + +def _get_type_alias_type(module_name: str, alias_name: str): + types = { + k: v + for k, v in sys.modules[module_name].__dict__.items() + if isinstance(v, typing._GenericAlias) # type: ignore + } + if alias_name not in types: + return alias_name + return types[alias_name] + + +def _get_model(module_name: str, model_name: str): + models = {k: v for k, v in sys.modules[module_name].__dict__.items() if isinstance(v, type)} + module_end = module_name.rsplit(".", 1)[0] + models.update({k: v for k, v in sys.modules[module_end].__dict__.items() if isinstance(v, type)}) + if isinstance(model_name, str): + model_name = model_name.split(".")[-1] + if model_name not in models: + return model_name + return models[model_name] + + +_UNSET = object() + + +class _MyMutableMapping(MutableMapping[str, typing.Any]): + def __init__(self, data: dict[str, typing.Any]) -> None: + self._data = data + + def __contains__(self, key: typing.Any) -> bool: + return key in self._data + + def __getitem__(self, key: str) -> typing.Any: + # If this key has been deserialized (for mutable types), we need to handle serialization + if hasattr(self, "_attr_to_rest_field"): + cache_attr = f"_deserialized_{key}" + if hasattr(self, cache_attr): + rf = _get_rest_field(getattr(self, "_attr_to_rest_field"), key) + if rf: + value = self._data.get(key) + if isinstance(value, (dict, list, set)): + # For mutable types, serialize and return + # But also update _data with serialized form and clear flag + # so mutations via this returned value affect _data + serialized = _serialize(value, rf._format) + # If serialized form is same type (no transformation needed), + # return _data directly so mutations work + if isinstance(serialized, type(value)) and serialized == value: + return self._data.get(key) + # Otherwise return serialized copy and clear flag + try: + object.__delattr__(self, cache_attr) + except AttributeError: + pass + # Store serialized form back + self._data[key] = serialized + return serialized + return self._data.__getitem__(key) + + def __setitem__(self, key: str, value: typing.Any) -> None: + # Clear any cached deserialized value when setting through dictionary access + cache_attr = f"_deserialized_{key}" + try: + object.__delattr__(self, cache_attr) + except AttributeError: + pass + self._data.__setitem__(key, value) + + def __delitem__(self, key: str) -> None: + self._data.__delitem__(key) + + def __iter__(self) -> typing.Iterator[typing.Any]: + return self._data.__iter__() + + def __len__(self) -> int: + return self._data.__len__() + + def __ne__(self, other: typing.Any) -> bool: + return not self.__eq__(other) + + def keys(self) -> typing.KeysView[str]: + """ + :returns: a set-like object providing a view on D's keys + :rtype: ~typing.KeysView + """ + return self._data.keys() + + def values(self) -> typing.ValuesView[typing.Any]: + """ + :returns: an object providing a view on D's values + :rtype: ~typing.ValuesView + """ + return self._data.values() + + def items(self) -> typing.ItemsView[str, typing.Any]: + """ + :returns: set-like object providing a view on D's items + :rtype: ~typing.ItemsView + """ + return self._data.items() + + def get(self, key: str, default: typing.Any = None) -> typing.Any: + """ + Get the value for key if key is in the dictionary, else default. + :param str key: The key to look up. + :param any default: The value to return if key is not in the dictionary. Defaults to None + :returns: D[k] if k in D, else d. + :rtype: any + """ + try: + return self[key] + except KeyError: + return default + + @typing.overload + def pop(self, key: str) -> typing.Any: ... # pylint: disable=arguments-differ + + @typing.overload + def pop(self, key: str, default: _T) -> _T: ... # pylint: disable=signature-differs + + @typing.overload + def pop(self, key: str, default: typing.Any) -> typing.Any: ... # pylint: disable=signature-differs + + def pop(self, key: str, default: typing.Any = _UNSET) -> typing.Any: + """ + Removes specified key and return the corresponding value. + :param str key: The key to pop. + :param any default: The value to return if key is not in the dictionary + :returns: The value corresponding to the key. + :rtype: any + :raises KeyError: If key is not found and default is not given. + """ + if default is _UNSET: + return self._data.pop(key) + return self._data.pop(key, default) + + def popitem(self) -> tuple[str, typing.Any]: + """ + Removes and returns some (key, value) pair + :returns: The (key, value) pair. + :rtype: tuple + :raises KeyError: if D is empty. + """ + return self._data.popitem() + + def clear(self) -> None: + """ + Remove all items from D. + """ + self._data.clear() + + def update(self, *args: typing.Any, **kwargs: typing.Any) -> None: # pylint: disable=arguments-differ + """ + Updates D from mapping/iterable E and F. + :param any args: Either a mapping object or an iterable of key-value pairs. + """ + self._data.update(*args, **kwargs) + + @typing.overload + def setdefault(self, key: str, default: None = None) -> None: ... + + @typing.overload + def setdefault(self, key: str, default: typing.Any) -> typing.Any: ... # pylint: disable=signature-differs + + def setdefault(self, key: str, default: typing.Any = _UNSET) -> typing.Any: + """ + Same as calling D.get(k, d), and setting D[k]=d if k not found + :param str key: The key to look up. + :param any default: The value to set if key is not in the dictionary + :returns: D[k] if k in D, else d. + :rtype: any + """ + if default is _UNSET: + return self._data.setdefault(key) + return self._data.setdefault(key, default) + + def __eq__(self, other: typing.Any) -> bool: + if isinstance(other, _MyMutableMapping): + return self._data == other._data + try: + other_model = self.__class__(other) + except Exception: + return False + return self._data == other_model._data + + def __repr__(self) -> str: + return str(self._data) + + +def _is_model(obj: typing.Any) -> bool: + return getattr(obj, "_is_model", False) + + +def _serialize(o, format: typing.Optional[str] = None): # pylint: disable=too-many-return-statements + if isinstance(o, list): + if format in _ARRAY_ENCODE_MAPPING and all(isinstance(x, str) for x in o): + return _ARRAY_ENCODE_MAPPING[format].join(o) + return [_serialize(x, format) for x in o] + if isinstance(o, dict): + return {k: _serialize(v, format) for k, v in o.items()} + if isinstance(o, set): + return {_serialize(x, format) for x in o} + if isinstance(o, tuple): + return tuple(_serialize(x, format) for x in o) + if isinstance(o, (bytes, bytearray)): + return _serialize_bytes(o, format) + if isinstance(o, decimal.Decimal): + return float(o) + if isinstance(o, enum.Enum): + return o.value + if isinstance(o, int): + if format == "str": + return str(o) + return o + try: + # First try datetime.datetime + return _serialize_datetime(o, format) + except AttributeError: + pass + # Last, try datetime.timedelta + try: + return _serialize_duration(o, format) + except AttributeError: + # This will be raised when it hits value.total_seconds in the method above + pass + return o + + +def _get_rest_field(attr_to_rest_field: dict[str, "_RestField"], rest_name: str) -> typing.Optional["_RestField"]: + try: + return next(rf for rf in attr_to_rest_field.values() if rf._rest_name == rest_name) + except StopIteration: + return None + + +def _create_value(rf: typing.Optional["_RestField"], value: typing.Any) -> typing.Any: + if not rf: + return _serialize(value, None) + if rf._is_multipart_file_input: + return value + if rf._is_model: + return _deserialize(rf._type, value) + if isinstance(value, ET.Element): + value = _deserialize(rf._type, value) + return _serialize(value, rf._format) + + +# ============================================================================ +# Fast-path scalar deserializer functions for rest_field(deserializer=...) +# These are referenced from rest_field declarations to bypass the generic +# _deserialize -> _deserialize_with_callable chain. +# Only simple/primitive types — no models or container types. +# ============================================================================ + + +def _xml_deser_str(value): + if isinstance(value, ET.Element): + return value.text or "" + return str(value) if value is not None else None + + +def _xml_deser_int(value): + if isinstance(value, ET.Element): + return int(value.text) if value.text else None + return int(value) if value is not None else None + + +def _xml_deser_float(value): + if isinstance(value, ET.Element): + return float(value.text) if value.text else None + return float(value) if value is not None else None + + +def _xml_deser_bool(value): + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + if text in (True, False): + return text + return text.lower() == "true" + + +# pylint: disable=docstring-missing-param +def _xml_deser_bytes(value): + """Deserialize bytes from XML (base64).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_bytes(text) + + +def _xml_deser_bytes_base64url(value): + """Deserialize bytes from XML (base64url).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_bytes_base64(text) + + +def _xml_deser_datetime(value): + """Deserialize a datetime from XML (ISO 8601 / rfc3339).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_datetime(text) + + +def _xml_deser_datetime_rfc7231(value): + """Deserialize a datetime from XML (RFC7231 format).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_datetime_rfc7231(text) + + +def _xml_deser_datetime_unix_timestamp(value): + """Deserialize a datetime from XML (Unix timestamp).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_datetime_unix_timestamp(float(text)) + + +def _xml_deser_date(value): + """Deserialize a date from XML (ISO 8601).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_date(text) + + +def _xml_deser_time(value): + """Deserialize a time from XML (ISO 8601).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_time(text) + + +def _xml_deser_duration(value): + """Deserialize a timedelta from XML (ISO 8601 duration).""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_duration(text) + + +def _xml_deser_decimal(value): + """Deserialize a Decimal from XML.""" + if isinstance(value, ET.Element): + text = value.text + else: + text = value + if text is None: + return None + return _deserialize_decimal(text) + + +def _xml_deser_enum_or_str(enum_cls, value): + """Deserialize a Union[EnumType, str] from XML.""" + text = value.text if isinstance(value, ET.Element) else value + if text is None: + return None + try: + return enum_cls(text) + except ValueError: + return text + + +def _extract_xml_model_type(rf_type): + """Extract the concrete Model class from a resolved rf._type partial chain. + + Unwraps ``Optional[Model]`` and ``_deserialize_model(Model, ...)`` + wrappers. Only handles Model and Optional[Model] — other composite + types (List, Dict, Union, etc.) return None and fall through to the + generic ``_deserialize`` path at runtime. + """ + if rf_type is None: + return None + if isinstance(rf_type, type) and _is_model(rf_type): + return rf_type + if not isinstance(rf_type, functools.partial): + return None + func = rf_type.func + args = rf_type.args + if func is _deserialize_with_optional and args: + return _extract_xml_model_type(args[0]) + if func is _deserialize_model and args: + cls = args[0] + return cls if isinstance(cls, type) and _is_model(cls) else None + return None + + +def _build_xml_field_plan( # pylint: disable=docstring-missing-return, docstring-missing-rtype, unused-variable + cls, attr_to_rest_field: dict +) -> list: + """Build a precomputed XML field plan for fast _init_from_xml iteration. + + Called once per model class in __new__. Returns a list of tuples: + (rest_name, xml_name, kind, deser, rf_type, is_optional, items_name) + + kind: 0=wrapped, 1=attribute, 2=unwrapped, 3=text + + For Model and Optional[Model] fields that lack a scalar + ``_deserializer``, this function precomputes the Model class as the + deserializer so ``_init_from_xml`` can call ``ModelClass(element)`` + directly instead of going through the expensive + ``_get_deserialize_callable_from_annotation`` chain at runtime. + """ + model_meta = getattr(cls, "_xml", {}) + model_ns = model_meta.get("ns") or model_meta.get("namespace") + plan = [] + + for rf in attr_to_rest_field.values(): + prop_meta = getattr(rf, "_xml", {}) + deser = rf._deserializer + + xml_name = prop_meta.get("name", rf._rest_name) + xml_ns = _resolve_xml_ns(prop_meta, model_meta) + if xml_ns: + xml_name = "{" + xml_ns + "}" + xml_name + + is_optional = rf._is_optional + + # For Model / Optional[Model] fields without a scalar deserializer, + # precompute the Model class as the deserializer. + if deser is None and rf._type is not None: + model_cls = _extract_xml_model_type(rf._type) + if model_cls is not None: + deser = model_cls + + if prop_meta.get("attribute", False): + plan.append((rf._rest_name, xml_name, 1, deser, rf._type, is_optional, None)) + elif prop_meta.get("unwrapped", False): + items_name = prop_meta.get("itemsName") + if items_name: + items_ns = prop_meta.get("itemsNs") + if items_ns is not None: + xml_ns = items_ns + if xml_ns: + items_name = "{" + xml_ns + "}" + items_name + else: + items_name = xml_name + plan.append((rf._rest_name, xml_name, 2, deser, rf._type, is_optional, items_name)) + elif prop_meta.get("text", False): + plan.append((rf._rest_name, xml_name, 3, deser, rf._type, is_optional, None)) + else: + plan.append((rf._rest_name, xml_name, 0, deser, rf._type, is_optional, None)) + + return plan + + +# pylint: enable=docstring-missing-param +class Model(_MyMutableMapping): + _is_model = True + # label whether current class's _attr_to_rest_field has been calculated + # could not see _attr_to_rest_field directly because subclass inherits it from parent class + _calculated: set[str] = set() + + def __init__(self, *args: typing.Any, **kwargs: typing.Any) -> None: + class_name = self.__class__.__name__ + if len(args) > 1: + raise TypeError(f"{class_name}.__init__() takes 2 positional arguments but {len(args) + 1} were given") + dict_to_pass: dict[str, typing.Any] = {} + if args: + if isinstance(args[0], ET.Element): + dict_to_pass.update(self._init_from_xml(args[0])) + else: + dict_to_pass.update( + {k: _create_value(_get_rest_field(self._attr_to_rest_field, k), v) for k, v in args[0].items()} + ) + else: + non_attr_kwargs = [k for k in kwargs if k not in self._attr_to_rest_field] + if non_attr_kwargs: + # actual type errors only throw the first wrong keyword arg they see, so following that. + raise TypeError(f"{class_name}.__init__() got an unexpected keyword argument '{non_attr_kwargs[0]}'") + dict_to_pass.update( + { + self._attr_to_rest_field[k]._rest_name: _create_value(self._attr_to_rest_field[k], v) + for k, v in kwargs.items() + if v is not None + } + ) + # Apply client default values for fields the caller didn't set so that + # defaults are part of `_data` and therefore included during serialization. + for rf in self._attr_to_rest_field.values(): + if rf._default is _UNSET: + continue + if rf._rest_name in dict_to_pass: + continue + dict_to_pass[rf._rest_name] = _create_value(rf, rf._default) + super().__init__(dict_to_pass) + + def _init_from_xml( # pylint: disable=too-many-branches, too-many-statements + self, element: ET.Element + ) -> dict[str, typing.Any]: + """Deserialize an XML element into a dict mapping rest field names to values. + + :param ET.Element element: The XML element to deserialize from. + :returns: A dictionary of rest_name to deserialized value pairs. + :rtype: dict + """ + result: dict[str, typing.Any] = {} + existed_attr_keys: list[str] = [] + + field_plan = getattr(self, "_xml_field_plan", None) + if field_plan: + for rest_name, xml_name, kind, deser, rf_type, is_optional, items_name in field_plan: + if kind == 0: # wrapped element (most common) + item = element.find(xml_name) + if item is not None: + existed_attr_keys.append(xml_name) + if deser: + result[rest_name] = deser(item) + else: + result[rest_name] = _deserialize(rf_type, item) + elif kind == 1: # attribute + attr_val = element.get(xml_name) + if attr_val is not None: + existed_attr_keys.append(xml_name) + if deser: + result[rest_name] = deser(attr_val) + else: + result[rest_name] = attr_val + elif kind == 2: # unwrapped array + items = element.findall(items_name) # pyright: ignore + if len(items) > 0: + existed_attr_keys.append(items_name) + if deser: + result[rest_name] = deser(items) + else: + result[rest_name] = _deserialize(rf_type, items) + elif not is_optional: + existed_attr_keys.append(items_name) + result[rest_name] = [] + elif kind == 3: # text + if element.text is not None: + if deser: + result[rest_name] = deser(element.text) + else: + result[rest_name] = element.text + else: + model_meta = getattr(self, "_xml", {}) + for rf in self._attr_to_rest_field.values(): + prop_meta = getattr(rf, "_xml", {}) + xml_name = prop_meta.get("name", rf._rest_name) + xml_ns = _resolve_xml_ns(prop_meta, model_meta) + if xml_ns: + xml_name = "{" + xml_ns + "}" + xml_name + + # attribute + if prop_meta.get("attribute", False) and element.get(xml_name) is not None: + existed_attr_keys.append(xml_name) + result[rf._rest_name] = _deserialize(rf._type, element.get(xml_name)) + continue + + # unwrapped element is array + if prop_meta.get("unwrapped", False): + _items_name = prop_meta.get("itemsName") + if _items_name: + xml_name = _items_name + _items_ns = prop_meta.get("itemsNs") + if _items_ns is not None: + xml_ns = _items_ns + if xml_ns: + xml_name = "{" + xml_ns + "}" + xml_name + items = element.findall(xml_name) # pyright: ignore + if len(items) > 0: + existed_attr_keys.append(xml_name) + result[rf._rest_name] = _deserialize(rf._type, items) + elif not rf._is_optional: + existed_attr_keys.append(xml_name) + result[rf._rest_name] = [] + continue + + # text element is primitive type + if prop_meta.get("text", False): + if element.text is not None: + result[rf._rest_name] = _deserialize(rf._type, element.text) + continue + + # wrapped element could be normal property or array + item = element.find(xml_name) + if item is not None: + existed_attr_keys.append(xml_name) + result[rf._rest_name] = _deserialize(rf._type, item) + + # rest thing is additional properties + for e in element: + if e.tag not in existed_attr_keys: + result[e.tag] = _convert_element(e) + + return result + + def copy(self) -> "Model": + return Model(self.__dict__) + + def __new__(cls, *args: typing.Any, **kwargs: typing.Any) -> Self: + if f"{cls.__module__}.{cls.__qualname__}" not in cls._calculated: + # we know the last nine classes in mro are going to be 'Model', '_MyMutableMapping', 'MutableMapping', + # 'Mapping', 'Collection', 'Sized', 'Iterable', 'Container' and 'object' + mros = cls.__mro__[:-9][::-1] # ignore parents, and reverse the mro order + attr_to_rest_field: dict[str, _RestField] = { # map attribute name to rest_field property + k: v for mro_class in mros for k, v in mro_class.__dict__.items() if k[0] != "_" and hasattr(v, "_type") + } + annotations = { + k: v + for mro_class in mros + if hasattr(mro_class, "__annotations__") + for k, v in mro_class.__annotations__.items() + } + for attr, rf in attr_to_rest_field.items(): + rf._module = cls.__module__ + if not rf._type: + rf._type = rf._get_deserialize_callable_from_annotation(annotations.get(attr, None)) + if not rf._rest_name_input: + rf._rest_name_input = attr + cls._attr_to_rest_field: dict[str, _RestField] = dict(attr_to_rest_field.items()) + # Build XML field plan for fast _init_from_xml (only for XML models) + if getattr(cls, "_xml", None): + cls._xml_field_plan = _build_xml_field_plan(cls, attr_to_rest_field) + cls._calculated.add(f"{cls.__module__}.{cls.__qualname__}") + + return super().__new__(cls) + + def __init_subclass__(cls, discriminator: typing.Optional[str] = None) -> None: + for base in cls.__bases__: + if hasattr(base, "__mapping__"): + base.__mapping__[discriminator or cls.__name__] = cls # type: ignore + + @classmethod + def _get_discriminator(cls, exist_discriminators) -> typing.Optional["_RestField"]: + for v in cls.__dict__.values(): + if isinstance(v, _RestField) and v._is_discriminator and v._rest_name not in exist_discriminators: + return v + return None + + @classmethod + def _deserialize(cls, data, exist_discriminators): + if not hasattr(cls, "__mapping__"): + return cls(data) + discriminator = cls._get_discriminator(exist_discriminators) + if discriminator is None: + return cls(data) + exist_discriminators.append(discriminator._rest_name) + if isinstance(data, ET.Element): + model_meta = getattr(cls, "_xml", {}) + prop_meta = getattr(discriminator, "_xml", {}) + xml_name = prop_meta.get("name", discriminator._rest_name) + xml_ns = _resolve_xml_ns(prop_meta, model_meta) + if xml_ns: + xml_name = "{" + xml_ns + "}" + xml_name + + if data.get(xml_name) is not None: + discriminator_value = data.get(xml_name) + else: + discriminator_value = data.find(xml_name).text # pyright: ignore + else: + discriminator_value = data.get(discriminator._rest_name) + mapped_cls = cls.__mapping__.get(discriminator_value, cls) # pyright: ignore # pylint: disable=no-member + return mapped_cls._deserialize(data, exist_discriminators) + + def as_dict(self, *, exclude_readonly: bool = False) -> dict[str, typing.Any]: + """Return a dict that can be turned into json using json.dump. + + :keyword bool exclude_readonly: Whether to remove the readonly properties. + :returns: A dict JSON compatible object + :rtype: dict + """ + + result = {} + readonly_props = [] + if exclude_readonly: + readonly_props = [p._rest_name for p in self._attr_to_rest_field.values() if _is_readonly(p)] + for k, v in self.items(): + if exclude_readonly and k in readonly_props: # pyright: ignore + continue + is_multipart_file_input = False + try: + is_multipart_file_input = next( + rf for rf in self._attr_to_rest_field.values() if rf._rest_name == k + )._is_multipart_file_input + except StopIteration: + pass + result[k] = v if is_multipart_file_input else Model._as_dict_value(v, exclude_readonly=exclude_readonly) + return result + + @staticmethod + def _as_dict_value(v: typing.Any, exclude_readonly: bool = False) -> typing.Any: + if v is None or isinstance(v, _Null): + return None + if isinstance(v, (list, tuple, set)): + return type(v)(Model._as_dict_value(x, exclude_readonly=exclude_readonly) for x in v) + if isinstance(v, dict): + return {dk: Model._as_dict_value(dv, exclude_readonly=exclude_readonly) for dk, dv in v.items()} + return v.as_dict(exclude_readonly=exclude_readonly) if hasattr(v, "as_dict") else v + + +def _deserialize_model(model_deserializer: typing.Optional[typing.Callable], obj): + if _is_model(obj): + return obj + return _deserialize(model_deserializer, obj) + + +def _deserialize_with_optional(if_obj_deserializer: typing.Optional[typing.Callable], obj): + if obj is None: + return obj + return _deserialize_with_callable(if_obj_deserializer, obj) + + +def _deserialize_with_union(deserializers, obj): + for deserializer in deserializers: + try: + return _deserialize(deserializer, obj) + except DeserializationError: + pass + raise DeserializationError() + + +def _deserialize_dict( + value_deserializer: typing.Optional[typing.Callable], + module: typing.Optional[str], + obj: dict[typing.Any, typing.Any], +): + if obj is None: + return obj + if isinstance(obj, ET.Element): + obj = {child.tag: child for child in obj} + return {k: _deserialize(value_deserializer, v, module) for k, v in obj.items()} + + +def _deserialize_multiple_sequence( + entry_deserializers: list[typing.Optional[typing.Callable]], + module: typing.Optional[str], + obj, +): + if obj is None: + return obj + return type(obj)(_deserialize(deserializer, entry, module) for entry, deserializer in zip(obj, entry_deserializers)) + + +def _is_array_encoded_deserializer(deserializer: functools.partial) -> bool: + return ( + isinstance(deserializer, functools.partial) + and isinstance(deserializer.args[0], functools.partial) + and deserializer.args[0].func == _deserialize_array_encoded # pylint: disable=comparison-with-callable + ) + + +def _deserialize_sequence( + deserializer: typing.Optional[typing.Callable], + module: typing.Optional[str], + obj, +): + if obj is None: + return obj + if isinstance(obj, ET.Element): + obj = list(obj) + + # encoded string may be deserialized to sequence + if isinstance(obj, str) and isinstance(deserializer, functools.partial): + # for list[str] + if _is_array_encoded_deserializer(deserializer): + return deserializer(obj) + + # for list[Union[...]] + if isinstance(deserializer.args[0], list): + for sub_deserializer in deserializer.args[0]: + if _is_array_encoded_deserializer(sub_deserializer): + return sub_deserializer(obj) + + return type(obj)(_deserialize(deserializer, entry, module) for entry in obj) + + +def _sorted_annotations(types: list[typing.Any]) -> list[typing.Any]: + return sorted( + types, + key=lambda x: hasattr(x, "__name__") and x.__name__.lower() in ("str", "float", "int", "bool"), + ) + + +def _get_deserialize_callable_from_annotation( # pylint: disable=too-many-return-statements, too-many-statements, too-many-branches + annotation: typing.Any, + module: typing.Optional[str], + rf: typing.Optional["_RestField"] = None, +) -> typing.Optional[typing.Callable[[typing.Any], typing.Any]]: + if not annotation: + return None + + # is it a type alias? + if isinstance(annotation, str): + if module is not None: + annotation = _get_type_alias_type(module, annotation) + + # is it a forward ref / in quotes? + if isinstance(annotation, (str, typing.ForwardRef)): + try: + model_name = annotation.__forward_arg__ # type: ignore + except AttributeError: + model_name = annotation + if module is not None: + annotation = _get_model(module, model_name) # type: ignore + + try: + if module and _is_model(annotation): + if rf: + rf._is_model = True + + return functools.partial(_deserialize_model, annotation) # pyright: ignore + except Exception: + pass + + # is it a literal? + try: + if annotation.__origin__ is typing.Literal: # pyright: ignore + return None + except AttributeError: + pass + + # is it optional? + try: + if any(a is _NONE_TYPE for a in annotation.__args__): # pyright: ignore + if rf: + rf._is_optional = True + if len(annotation.__args__) <= 2: # pyright: ignore + if_obj_deserializer = _get_deserialize_callable_from_annotation( + next(a for a in annotation.__args__ if a is not _NONE_TYPE), module, rf # pyright: ignore + ) + + return functools.partial(_deserialize_with_optional, if_obj_deserializer) + # the type is Optional[Union[...]], we need to remove the None type from the Union + annotation_copy = copy.copy(annotation) + annotation_copy.__args__ = [a for a in annotation_copy.__args__ if a is not _NONE_TYPE] # pyright: ignore + return _get_deserialize_callable_from_annotation(annotation_copy, module, rf) + except AttributeError: + pass + + # is it union? + if getattr(annotation, "__origin__", None) is typing.Union: + # initial ordering is we make `string` the last deserialization option, because it is often them most generic + deserializers = [ + _get_deserialize_callable_from_annotation(arg, module, rf) + for arg in _sorted_annotations(annotation.__args__) # pyright: ignore + ] + + return functools.partial(_deserialize_with_union, deserializers) + + try: + annotation_name = ( + annotation.__name__ if hasattr(annotation, "__name__") else annotation._name # pyright: ignore + ) + if annotation_name.lower() == "dict": + value_deserializer = _get_deserialize_callable_from_annotation( + annotation.__args__[1], module, rf # pyright: ignore + ) + + return functools.partial( + _deserialize_dict, + value_deserializer, + module, + ) + except (AttributeError, IndexError): + pass + try: + annotation_name = ( + annotation.__name__ if hasattr(annotation, "__name__") else annotation._name # pyright: ignore + ) + if annotation_name.lower() in ["list", "set", "tuple", "sequence"]: + if len(annotation.__args__) > 1: # pyright: ignore + entry_deserializers = [ + _get_deserialize_callable_from_annotation(dt, module, rf) + for dt in annotation.__args__ # pyright: ignore + ] + return functools.partial(_deserialize_multiple_sequence, entry_deserializers, module) + deserializer = _get_deserialize_callable_from_annotation( + annotation.__args__[0], module, rf # pyright: ignore + ) + + return functools.partial(_deserialize_sequence, deserializer, module) + except (TypeError, IndexError, AttributeError, SyntaxError): + pass + + def _deserialize_default( + deserializer, + obj, + ): + if obj is None: + return obj + try: + return _deserialize_with_callable(deserializer, obj) + except Exception: + pass + return obj + + if get_deserializer(annotation, rf): + return functools.partial(_deserialize_default, get_deserializer(annotation, rf)) + + return functools.partial(_deserialize_default, annotation) + + +def _deserialize_with_callable( + deserializer: typing.Optional[typing.Callable[[typing.Any], typing.Any]], + value: typing.Any, +): # pylint: disable=too-many-return-statements + try: + if value is None or isinstance(value, _Null): + return None + if isinstance(value, ET.Element): + if deserializer is str: + return value.text or "" + if deserializer is int: + return int(value.text) if value.text else None + if deserializer is float: + return float(value.text) if value.text else None + if deserializer is bool: + return value.text == "true" if value.text else None + if deserializer and deserializer in _DESERIALIZE_MAPPING.values(): + return deserializer(value.text) if value.text else None + if deserializer and deserializer in _DESERIALIZE_MAPPING_WITHFORMAT.values(): + return deserializer(value.text) if value.text else None + if deserializer is None: + return value + if deserializer in [int, float, bool]: + return deserializer(value) + if isinstance(deserializer, CaseInsensitiveEnumMeta): + try: + return deserializer(value.text if isinstance(value, ET.Element) else value) + except ValueError: + # for unknown value, return raw value + return value.text if isinstance(value, ET.Element) else value + if isinstance(deserializer, type) and issubclass(deserializer, Model): + return deserializer._deserialize(value, []) + return typing.cast(typing.Callable[[typing.Any], typing.Any], deserializer)(value) + except Exception as e: + raise DeserializationError() from e + + +def _deserialize( + deserializer: typing.Any, + value: typing.Any, + module: typing.Optional[str] = None, + rf: typing.Optional["_RestField"] = None, + format: typing.Optional[str] = None, +) -> typing.Any: + if isinstance(value, PipelineResponse): + value = value.http_response.json() + if rf is None and format: + rf = _RestField(format=format) + if not isinstance(deserializer, functools.partial): + deserializer = _get_deserialize_callable_from_annotation(deserializer, module, rf) + return _deserialize_with_callable(deserializer, value) + + +def _failsafe_deserialize( + deserializer: typing.Any, + response: HttpResponse, + module: typing.Optional[str] = None, + rf: typing.Optional["_RestField"] = None, + format: typing.Optional[str] = None, +) -> typing.Any: + try: + return _deserialize(deserializer, response.json(), module, rf, format) + except Exception: # pylint: disable=broad-except + _LOGGER.warning( + "Ran into a deserialization error. Ignoring since this is failsafe deserialization", exc_info=True + ) + return None + + +def _failsafe_deserialize_xml( + deserializer: typing.Any, + response: HttpResponse, +) -> typing.Any: + try: + return _deserialize_xml(deserializer, response.text()) + except Exception: # pylint: disable=broad-except + _LOGGER.warning( + "Ran into a deserialization error. Ignoring since this is failsafe deserialization", exc_info=True + ) + return None + + +# pylint: disable=too-many-instance-attributes +class _RestField: + def __init__( + self, + *, + name: typing.Optional[str] = None, + type: typing.Optional[typing.Callable] = None, # pylint: disable=redefined-builtin + is_discriminator: bool = False, + visibility: typing.Optional[list[str]] = None, + default: typing.Any = _UNSET, + format: typing.Optional[str] = None, + is_multipart_file_input: bool = False, + xml: typing.Optional[dict[str, typing.Any]] = None, + deserializer: typing.Optional[typing.Callable] = None, + ): + self._type = type + self._rest_name_input = name + self._module: typing.Optional[str] = None + self._is_discriminator = is_discriminator + self._visibility = visibility + self._is_model = False + self._is_optional = False + self._default = default + self._format = format + self._is_multipart_file_input = is_multipart_file_input + self._xml = xml if xml is not None else {} + self._deserializer = deserializer + + @property + def _class_type(self) -> typing.Any: + result = getattr(self._type, "args", [None])[0] + # type may be wrapped by nested functools.partial so we need to check for that + if isinstance(result, functools.partial): + return getattr(result, "args", [None])[0] + return result + + @property + def _rest_name(self) -> str: + if self._rest_name_input is None: + raise ValueError("Rest name was never set") + return self._rest_name_input + + def __get__(self, obj: Model, type=None): # pylint: disable=redefined-builtin + # by this point, type and rest_name will have a value bc we default + # them in __new__ of the Model class + # Use _data.get() directly to avoid triggering __getitem__ which clears the cache + item = obj._data.get(self._rest_name, _UNSET) + if item is _UNSET: + # Field not set by user; return the client default if one exists, otherwise None + return self._default if self._default is not _UNSET else None + if item is None: + return item + if self._is_model: + return item + + # For mutable types, we want mutations to directly affect _data + # Check if we've already deserialized this value + cache_attr = f"_deserialized_{self._rest_name}" + if hasattr(obj, cache_attr): + # Return the value from _data directly (it's been deserialized in place) + return obj._data.get(self._rest_name) + + # Fast path: use _deserializer directly (avoids _serialize/_deserialize chain) + if self._deserializer: + deserialized = self._deserializer(item) + else: + deserialized = _deserialize(self._type, _serialize(item, self._format), rf=self) + + # For mutable types, store the deserialized value back in _data + # so mutations directly affect _data + if isinstance(deserialized, (dict, list, set)): + obj._data[self._rest_name] = deserialized + object.__setattr__(obj, cache_attr, True) # Mark as deserialized + return deserialized + + return deserialized + + def __set__(self, obj: Model, value) -> None: + # Clear the cached deserialized object when setting a new value + cache_attr = f"_deserialized_{self._rest_name}" + if hasattr(obj, cache_attr): + object.__delattr__(obj, cache_attr) + + if value is None: + # we want to wipe out entries if users set attr to None + try: + obj.__delitem__(self._rest_name) + except KeyError: + pass + return + if self._is_model: + if not _is_model(value): + value = _deserialize(self._type, value) + obj.__setitem__(self._rest_name, value) + return + obj.__setitem__(self._rest_name, _serialize(value, self._format)) + + def _get_deserialize_callable_from_annotation( + self, annotation: typing.Any + ) -> typing.Optional[typing.Callable[[typing.Any], typing.Any]]: + return _get_deserialize_callable_from_annotation(annotation, self._module, self) + + +def rest_field( + *, + name: typing.Optional[str] = None, + type: typing.Optional[typing.Callable] = None, # pylint: disable=redefined-builtin + visibility: typing.Optional[list[str]] = None, + default: typing.Any = _UNSET, + format: typing.Optional[str] = None, + is_multipart_file_input: bool = False, + xml: typing.Optional[dict[str, typing.Any]] = None, + deserializer: typing.Optional[typing.Callable] = None, +) -> typing.Any: + return _RestField( + name=name, + type=type, + visibility=visibility, + default=default, + format=format, + is_multipart_file_input=is_multipart_file_input, + xml=xml, + deserializer=deserializer, + ) + + +def rest_discriminator( + *, + name: typing.Optional[str] = None, + type: typing.Optional[typing.Callable] = None, # pylint: disable=redefined-builtin + visibility: typing.Optional[list[str]] = None, + xml: typing.Optional[dict[str, typing.Any]] = None, +) -> typing.Any: + return _RestField(name=name, type=type, is_discriminator=True, visibility=visibility, xml=xml) + + +def serialize_xml(model: Model, exclude_readonly: bool = False) -> str: + """Serialize a model to XML. + + :param Model model: The model to serialize. + :param bool exclude_readonly: Whether to exclude readonly properties. + :returns: The XML representation of the model. + :rtype: str + """ + return ET.tostring(_get_element(model, exclude_readonly), encoding="unicode") # type: ignore + + +def _get_xml_ns(meta: dict[str, typing.Any]) -> typing.Optional[str]: + """Return the XML namespace from a metadata dict, checking both 'ns' (old-style) and 'namespace' (DPG) keys. + + :param dict meta: The metadata dictionary to extract namespace from. + :returns: The namespace string if 'ns' or 'namespace' key is present, None otherwise. + :rtype: str or None + """ + ns = meta.get("ns") + if ns is None: + ns = meta.get("namespace") + return ns + + +def _resolve_xml_ns( + prop_meta: dict[str, typing.Any], model_meta: typing.Optional[dict[str, typing.Any]] = None +) -> typing.Optional[str]: + """Resolve XML namespace for a property, falling back to model namespace when appropriate. + + Checks the property metadata first; if no namespace is found and the model does not declare + an explicit prefix, falls back to the model-level namespace. + + :param dict prop_meta: The property metadata dictionary. + :param dict model_meta: The model metadata dictionary, used as fallback. + :returns: The resolved namespace string, or None. + :rtype: str or None + """ + ns = _get_xml_ns(prop_meta) + if ns is None and model_meta is not None and not model_meta.get("prefix"): + ns = _get_xml_ns(model_meta) + return ns + + +def _set_xml_attribute(element: ET.Element, name: str, value: typing.Any, prop_meta: dict[str, typing.Any]) -> None: + """Set an XML attribute on an element, handling namespace prefix registration. + + :param ET.Element element: The element to set the attribute on. + :param str name: The default attribute name (wire name). + :param any value: The attribute value. + :param dict prop_meta: The property metadata dictionary. + """ + xml_name = prop_meta.get("name", name) + _attr_ns = _get_xml_ns(prop_meta) + if _attr_ns: + _attr_prefix = prop_meta.get("prefix") + if _attr_prefix: + _safe_register_namespace(_attr_prefix, _attr_ns) + xml_name = "{" + _attr_ns + "}" + xml_name + element.set(xml_name, _get_primitive_type_value(value)) + + +def _get_element( + o: typing.Any, + exclude_readonly: bool = False, + parent_meta: typing.Optional[dict[str, typing.Any]] = None, + wrapped_element: typing.Optional[ET.Element] = None, +) -> typing.Union[ET.Element, list[ET.Element]]: + if _is_model(o): + model_meta = getattr(o, "_xml", {}) + + # if prop is a model, then use the prop element directly, else generate a wrapper of model + if wrapped_element is None: + # When serializing as an array item (parent_meta is set), check if the parent has an + # explicit itemsName. This ensures correct element names for unwrapped arrays (where + # the element tag is the property/items name, not the model type name). + _items_name = parent_meta.get("itemsName") if parent_meta is not None else None + element_name = _items_name if _items_name else (model_meta.get("name") or o.__class__.__name__) + _model_ns = _get_xml_ns(model_meta) + wrapped_element = _create_xml_element( + element_name, + model_meta.get("prefix"), + _model_ns, + ) + + readonly_props = [] + if exclude_readonly: + readonly_props = [p._rest_name for p in o._attr_to_rest_field.values() if _is_readonly(p)] + + for k, v in o.items(): + # do not serialize readonly properties + if exclude_readonly and k in readonly_props: + continue + + prop_rest_field = _get_rest_field(o._attr_to_rest_field, k) + if prop_rest_field: + prop_meta = getattr(prop_rest_field, "_xml").copy() + # use the wire name as xml name if no specific name is set + if prop_meta.get("name") is None: + prop_meta["name"] = k + else: + # additional properties will not have rest field, use the wire name as xml name + prop_meta = {"name": k} + + # Propagate model namespace to properties only for old-style "ns"-keyed models. + # DPG-generated models use the "namespace" key and explicitly declare namespace on + # each property that needs it, so propagation is intentionally skipped for them. + if prop_meta.get("ns") is None and model_meta.get("ns"): + prop_meta["ns"] = model_meta.get("ns") + prop_meta["prefix"] = model_meta.get("prefix") + + if prop_meta.get("unwrapped", False): + # unwrapped could only set on array + wrapped_element.extend(_get_element(v, exclude_readonly, prop_meta)) + elif prop_meta.get("text", False): + # text could only set on primitive type + wrapped_element.text = _get_primitive_type_value(v) + elif prop_meta.get("attribute", False): + _set_xml_attribute(wrapped_element, k, v, prop_meta) + else: + # other wrapped prop element + wrapped_element.append(_get_wrapped_element(v, exclude_readonly, prop_meta)) + return wrapped_element + if isinstance(o, list): + return [_get_element(x, exclude_readonly, parent_meta) for x in o] # type: ignore + if isinstance(o, dict): + result = [] + _dict_ns = _get_xml_ns(parent_meta) if parent_meta else None + for k, v in o.items(): + result.append( + _get_wrapped_element( + v, + exclude_readonly, + { + "name": k, + "ns": _dict_ns, + "prefix": parent_meta.get("prefix") if parent_meta else None, + }, + ) + ) + return result + + # primitive case need to create element based on parent_meta + if parent_meta: + _items_ns = parent_meta.get("itemsNs") + if _items_ns is None: + _items_ns = _get_xml_ns(parent_meta) + return _get_wrapped_element( + o, + exclude_readonly, + { + "name": parent_meta.get("itemsName", parent_meta.get("name")), + "prefix": parent_meta.get("itemsPrefix", parent_meta.get("prefix")), + "ns": _items_ns, + }, + ) + + raise ValueError("Could not serialize value into xml: " + o) + + +def _get_wrapped_element( + v: typing.Any, + exclude_readonly: bool, + meta: typing.Optional[dict[str, typing.Any]], +) -> ET.Element: + _meta_ns = _get_xml_ns(meta) if meta else None + wrapped_element = _create_xml_element( + meta.get("name") if meta else None, meta.get("prefix") if meta else None, _meta_ns + ) + if isinstance(v, (dict, list)): + wrapped_element.extend(_get_element(v, exclude_readonly, meta)) + elif _is_model(v): + _get_element(v, exclude_readonly, meta, wrapped_element) + else: + wrapped_element.text = _get_primitive_type_value(v) + return wrapped_element # type: ignore[no-any-return] + + +def _get_primitive_type_value(v) -> str: + if v is True: + return "true" + if v is False: + return "false" + if isinstance(v, _Null): + return "" + return str(v) + + +def _safe_register_namespace(prefix: str, ns: str) -> None: + """Register an XML namespace prefix, handling reserved prefix patterns. + + Some prefixes (e.g. 'ns2') match Python's reserved 'ns\\d+' pattern used for + auto-generated prefixes, causing register_namespace to raise ValueError. + Falls back to directly registering in the internal namespace map. + + :param str prefix: The namespace prefix to register. + :param str ns: The namespace URI. + """ + try: + ET.register_namespace(prefix, ns) + except ValueError: + _ns_map = getattr(ET, "_namespace_map", None) + if _ns_map is not None: + _ns_map[ns] = prefix + + +def _create_xml_element( + tag: typing.Any, prefix: typing.Optional[str] = None, ns: typing.Optional[str] = None +) -> ET.Element: + if prefix and ns: + _safe_register_namespace(prefix, ns) + if ns: + return ET.Element("{" + ns + "}" + tag) + return ET.Element(tag) + + +def _deserialize_xml( + deserializer: typing.Any, + value: str, +) -> typing.Any: + element = ET.fromstring(value) # nosec + if _is_model(deserializer): + return deserializer._deserialize(element, []) + return _deserialize(deserializer, element) + + +def _convert_element(e: ET.Element): + # dict case + if len(e.attrib) > 0 or len({child.tag for child in e}) > 1: + dict_result: dict[str, typing.Any] = {} + for child in e: + if dict_result.get(child.tag) is not None: + if isinstance(dict_result[child.tag], list): + dict_result[child.tag].append(_convert_element(child)) + else: + dict_result[child.tag] = [dict_result[child.tag], _convert_element(child)] + else: + dict_result[child.tag] = _convert_element(child) + dict_result.update(e.attrib) + return dict_result + # array case + if len(e) > 0: + array_result: list[typing.Any] = [] + for child in e: + array_result.append(_convert_element(child)) + return array_result + # primitive case + return e.text diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_json.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_json.py new file mode 100644 index 000000000000..335e9a0c58c4 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_json.py @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Small JSON helpers shared by Foundry storage serializers.""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype + +from __future__ import annotations + +import json +from typing import Any + + +def load_json(body: str | None) -> Any: + """Parse a JSON response body, treating empty/None bodies as an empty object.""" + return json.loads(body or "{}") diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_policies.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_policies.py new file mode 100644 index 000000000000..f592933183a4 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_policies.py @@ -0,0 +1,147 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Azure Core pipeline policies shared by Foundry storage clients. + +Both the request User-Agent policy and the per-retry logging policy are +protocol-neutral and reused by every AgentServer storage client. +""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype + +from __future__ import annotations + +import logging +import time +import urllib.parse +from typing import Callable + +from azure.ai.agentserver.core._platform_headers import ( + APIM_REQUEST_ID, + CLIENT_REQUEST_ID, + REQUEST_ID, + TRACEPARENT, +) +from azure.core.pipeline import PipelineRequest, PipelineResponse +from azure.core.pipeline.policies import AsyncHTTPPolicy, SansIOHTTPPolicy +from azure.core.rest import AsyncHttpResponse, HttpRequest, HttpResponse + +logger = logging.getLogger("azure.ai.agentserver.core") + +_LOG_FMT_START = "Foundry storage %s %s starting (x-ms-client-request-id=%s, traceparent=%s)" +_LOG_FMT_TRANSPORT_FAILURE = ( + "Foundry storage %s %s transport failure after %.1fms (x-ms-client-request-id=%s, traceparent=%s)" +) +_LOG_FMT_RESPONSE = ( + "Foundry storage %s %s -> %d (%.1fms, " + "x-ms-client-request-id=%s, traceparent=%s, x-request-id=%s, apim-request-id=%s)" +) + + +class ServerVersionUserAgentPolicy(SansIOHTTPPolicy[HttpRequest, HttpResponse]): + """Pipeline policy that sets the ``User-Agent`` header lazily from a callback. + + Unlike :class:`~azure.core.pipeline.policies.UserAgentPolicy` which captures + the value at construction time, this policy evaluates the callback on each + request so it reflects segments registered after the pipeline was built + (e.g. by cooperative ``__init__`` in a multi-protocol host). + """ + + def __init__(self, get_server_version: Callable[[], str]) -> None: + super().__init__() + self._get_server_version = get_server_version + + def on_request(self, request: PipelineRequest[HttpRequest]) -> None: + """Set the ``User-Agent`` header before the request is sent.""" + request.http_request.headers["User-Agent"] = self._get_server_version() + + +def _mask_storage_url(url: str) -> str: + """Mask the sensitive portions of a Foundry storage URL.""" + try: + if not url: + return "(redacted)" + parsed = urllib.parse.urlparse(url) + path = parsed.path or "" + idx = path.find("/storage") + if idx < 0: + return "(redacted)" + storage_path = path[idx:] + masked = f"***{_redact_storage_path(storage_path)}" + qs = urllib.parse.parse_qs(parsed.query) + api_version = qs.get("api-version") + if api_version: + masked += f"?api-version={api_version[0]}" + return masked + except Exception: # pylint: disable=broad-exception-caught + return "(redacted)" + + +def _redact_storage_path(path: str) -> str: + if not path.startswith("/storage/"): + return "/storage" + segments = path.split("/") + redacted: list[str] = [] + for index, segment in enumerate(segments): + if index < 3 or segment in {"items", "items:keys", "state_stores"} or not segment: + redacted.append(segment) + else: + redacted.append("*") + return "/".join(redacted) + + +class FoundryStorageLoggingPolicy(AsyncHTTPPolicy[HttpRequest, AsyncHttpResponse]): + """Azure Core per-retry pipeline policy that logs Foundry storage calls.""" + + async def send(self, request: PipelineRequest[HttpRequest]) -> PipelineResponse[HttpRequest, AsyncHttpResponse]: + """Send the request and log the operation details.""" + http_request = request.http_request + method = http_request.method + url = _mask_storage_url(str(http_request.url)) + client_request_id = http_request.headers.get(CLIENT_REQUEST_ID, "") + traceparent = http_request.headers.get(TRACEPARENT, "") + + logger.debug( + _LOG_FMT_START, + method, + url, + client_request_id, + traceparent, + ) + + start = time.monotonic() + try: + response = await self.next.send(request) + except Exception: + elapsed_ms = (time.monotonic() - start) * 1000 + logger.error( + _LOG_FMT_TRANSPORT_FAILURE, + method, + url, + elapsed_ms, + client_request_id, + traceparent, + exc_info=True, + ) + raise + + elapsed_ms = (time.monotonic() - start) * 1000 + http_response = response.http_response + status_code = http_response.status_code + x_request_id = http_response.headers.get(REQUEST_ID, "") + apim_request_id = http_response.headers.get(APIM_REQUEST_ID, "") + + log_level = logging.INFO if 200 <= status_code < 400 else logging.WARNING + logger.log( + log_level, + _LOG_FMT_RESPONSE, + method, + url, + status_code, + elapsed_ms, + client_request_id, + traceparent, + x_request_id, + apim_request_id, + ) + + return response diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state.py new file mode 100644 index 000000000000..9f9ee11522de --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state.py @@ -0,0 +1,471 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Developer-facing Foundry state-store client bound to one explicit store.""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype +# pylint: disable=client-accepts-api-version-keyword + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from typing import Any + +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.rest import HttpRequest + +from .._request_context import get_request_context +from ._client import JSON_CONTENT_TYPE, FoundryStorageClient +from ._endpoint import FoundryStorageEndpoint +from ._errors import FoundryStorageConflictError, FoundryStorageNotFoundError +from ._state_serializer import ( + _UNSET, + DeletedStateStore, + DeletedStateStoreItem, + JSONObject, + StateStoreItemKeyPage, + Order, + StateStore, + StateStoreItem, + StateStoreItemRef, + deserialize_deleted_state_item, + deserialize_deleted_state_store, + deserialize_list_keys_response, + deserialize_state_item, + deserialize_state_item_ref, + deserialize_state_store, + serialize_item_create_request, + serialize_item_put_request, + serialize_store_create_request, + serialize_store_update_request, +) + +DEFAULT_ITEM_TTL_SECONDS = 30 * 24 * 60 * 60 +DELEGATED_USER_ID_HEADER = "x-ms-user-id" + + +def _encode_segment(value: str) -> str: + encoded = base64.urlsafe_b64encode(value.encode("utf-8")).decode("ascii") + return encoded.rstrip("=") + + +def _validate_key(key: str) -> None: + if not key: + raise ValueError("key must be a non-empty string") + + +class FoundryStateStore(FoundryStorageClient): + """Developer-facing client for one explicit Foundry state store. + + The instance is bound to a single caller-chosen store ``name``. Session or + conversation scoping is expressed by encoding that identity into the store + name itself (for example ``checkpoints/``), matching spec + PR 247's removal of built-in session isolation. + + Prefer :meth:`get_or_create` over the constructor: it resolves (or creates) + the store's server-side resource in one call, so you never need a separate + lifecycle step before reading or writing items:: + + store = await FoundryStateStore.get_or_create("checkpoints/thread-abc") + await store.set_item("step-1", {"done": False}) + + Store-level operations (:meth:`get`, :meth:`update`, :meth:`delete`) act on + the bound store itself. The explicit ``*_item`` methods + (:meth:`create_item`, :meth:`get_item`, :meth:`set_item`, + :meth:`delete_item`, :meth:`list_keys`) act on individual items within it. + """ + + def __init__( + self, + name: str, + credential: AsyncTokenCredential | None = None, + endpoint: FoundryStorageEndpoint | str | None = None, + *, + user_isolation: bool = False, + item_ttl_seconds: int = DEFAULT_ITEM_TTL_SECONDS, + description: str | None = None, + tags: Mapping[str, str] | None = None, + user_id: str | None = None, + api_version: str = "v1", + **kwargs: Any, + ) -> None: + """Create a store-bound durable state-store client. + + Prefer :meth:`get_or_create`, which also resolves the server-side store + resource; use the constructor directly only when you already know the + store exists and want to skip that round trip. + + :param name: The logical state-store name. Encode conversation/thread + identity into this name when you need that scope. + :type name: str + :param credential: Async token credential. Defaults to + ``DefaultAzureCredential`` when omitted. + :type credential: AsyncTokenCredential | None + :param endpoint: Foundry storage endpoint or project endpoint URL. + :type endpoint: FoundryStorageEndpoint | str | None + :keyword user_isolation: Whether item operations should be partitioned + per resolved user. Fixed at store creation; ignored if the store + already exists. + :paramtype user_isolation: bool + :keyword item_ttl_seconds: Store-level default TTL inherited by every + item. Fixed at store creation; ignored if the store already exists. + :paramtype item_ttl_seconds: int + :keyword description: Optional mutable store description, set at + creation. Change it later with :meth:`update`. + :paramtype description: str or None + :keyword tags: Optional mutable store metadata tags, set at creation. + Change them later with :meth:`update`. + :paramtype tags: ~collections.abc.Mapping[str, str] or None + :keyword user_id: Delegated end-user identity for trusted callers. + :paramtype user_id: str or None + :keyword api_version: Storage API version. + :paramtype api_version: str + :raises ValueError: If ``name`` is empty. + """ + if not name: + raise ValueError("name must be a non-empty string") + self._owns_credential = False + if credential is None: + try: + from azure.identity.aio import DefaultAzureCredential + except ImportError as exc: # pragma: no cover + raise ImportError( + "FoundryStateStore requires azure-identity when no credential is supplied. " + "Install azure-identity or pass an async credential." + ) from exc + credential = DefaultAzureCredential() + self._owns_credential = True + self._credential = credential + + if isinstance(endpoint, FoundryStorageEndpoint): + resolved = endpoint + elif isinstance(endpoint, str): + resolved = FoundryStorageEndpoint.from_endpoint(endpoint, api_version=api_version) + else: + resolved = FoundryStorageEndpoint.from_env(api_version=api_version) + + self._name = name + self._user_isolation = user_isolation + self._item_ttl_seconds = item_ttl_seconds + self._description = description + self._tags = {} if tags is None else dict(tags) + self._user_id = user_id + super().__init__(credential, resolved, **kwargs) + + async def aclose(self) -> None: + """Close the pipeline client and any owned default credential.""" + try: + await super().aclose() + finally: + if self._owns_credential and hasattr(self._credential, "close"): + await self._credential.close() + + @property + def name(self) -> str: + """Return the logical store name bound to this client.""" + return self._name + + def _store_path(self) -> str: + return f"state_stores/{_encode_segment(self._name)}" + + def _item_path(self, key: str) -> str: + _validate_key(key) + return f"{self._store_path()}/items/{_encode_segment(key)}" + + def _request( + self, + method: str, + path: str, + *, + content: bytes | None = None, + include_user_id: bool = False, + if_match: str | None = None, + query: Mapping[str, str] | None = None, + ) -> HttpRequest: + headers: dict[str, str] = get_request_context().platform_headers() + if content is not None: + headers["Content-Type"] = JSON_CONTENT_TYPE + if include_user_id and self._user_id is not None: + headers[DELEGATED_USER_ID_HEADER] = self._user_id + if if_match is not None: + headers["If-Match"] = if_match + return HttpRequest(method, self._endpoint.build_url(path, **(query or {})), content=content, headers=headers) + + @classmethod + async def get_or_create( + cls, + name: str, + credential: AsyncTokenCredential | None = None, + endpoint: FoundryStorageEndpoint | str | None = None, + *, + user_isolation: bool = False, + item_ttl_seconds: int = DEFAULT_ITEM_TTL_SECONDS, + description: str | None = None, + tags: Mapping[str, str] | None = None, + user_id: str | None = None, + api_version: str = "v1", + **kwargs: Any, + ) -> "FoundryStateStore": + """Return a client bound to *name*, creating the store if needed. + + This is the recommended entry point: it resolves the server-side store + resource in one call (fetch, or create on first use) so callers never + need a separate lifecycle step before reading or writing items. + + :param name: The logical state-store name. See the constructor. + :type name: str + :param credential: Async token credential. See the constructor. + :type credential: AsyncTokenCredential | None + :param endpoint: Foundry storage endpoint or project endpoint URL. + :type endpoint: FoundryStorageEndpoint | str | None + :keyword user_isolation: See the constructor. Only applied if the store + does not already exist. + :paramtype user_isolation: bool + :keyword item_ttl_seconds: See the constructor. Only applied if the + store does not already exist. + :paramtype item_ttl_seconds: int + :keyword description: See the constructor. Only applied if the store + does not already exist. + :paramtype description: str or None + :keyword tags: See the constructor. Only applied if the store does not + already exist. + :paramtype tags: ~collections.abc.Mapping[str, str] or None + :keyword user_id: Delegated end-user identity for trusted callers. + :paramtype user_id: str or None + :keyword api_version: Storage API version. + :paramtype api_version: str + :return: The bound, ready-to-use store client. + :rtype: FoundryStateStore + """ + store = cls( + name, + credential, + endpoint, + user_isolation=user_isolation, + item_ttl_seconds=item_ttl_seconds, + description=description, + tags=tags, + user_id=user_id, + api_version=api_version, + **kwargs, + ) + try: + try: + await store._fetch_properties() + except FoundryStorageNotFoundError: + try: + await store._create_properties() + except FoundryStorageConflictError: + await store._fetch_properties() + except BaseException: + # Never leak the pipeline (or an owned DefaultAzureCredential) when + # resolution fails: the caller never receives ``store`` and so has + # no handle to close it. ``BaseException`` covers task cancellation. + await store.aclose() + raise + return store + + async def _create_properties(self) -> StateStore: + body = serialize_store_create_request( + self._name, + user_isolation=self._user_isolation, + item_ttl_seconds=self._item_ttl_seconds, + description=self._description, + tags=self._tags, + ) + response = await self._send_storage_request(self._request("POST", "state_stores", content=body)) + return deserialize_state_store(response.text()) + + async def _fetch_properties(self) -> StateStore: + response = await self._send_storage_request(self._request("GET", self._store_path())) + return deserialize_state_store(response.text()) + + async def update( + self, + *, + description: str | None | object = _UNSET, + tags: Mapping[str, str] | None | object = _UNSET, + ) -> StateStore: + """Update the bound store's mutable metadata (``description`` / ``tags``). + + :keyword description: The new description, or ``None`` to clear it. + Omit to leave the description unchanged. + :paramtype description: str or None + :keyword tags: The new tags (replaces the existing set wholesale), or + ``None`` to clear them. Omit to leave the tags unchanged. + :paramtype tags: ~collections.abc.Mapping[str, str] or None + :return: The updated store descriptor. + :rtype: ~azure.ai.agentserver.core.storage.StateStore + """ + body = serialize_store_update_request(description, tags) + response = await self._send_storage_request(self._request("PATCH", self._store_path(), content=body)) + if description is not _UNSET: + self._description = ( + description if isinstance(description, str) or description is None else self._description + ) + if tags is not _UNSET: + self._tags = dict(tags) if isinstance(tags, Mapping) else {} + return deserialize_state_store(response.text()) + + async def create_item( + self, + key: str, + value: JSONObject, + *, + tags: Mapping[str, str] | None = None, + ) -> StateStoreItemRef: + """Create a new item and fail on duplicate keys. + + :param key: The item key to create. + :type key: str + :param value: The item payload as a JSON object. + :type value: dict[str, any] + :keyword tags: Optional string labels stored with the item and used to + filter :meth:`list_keys`. + :paramtype tags: ~collections.abc.Mapping[str, str] or None + :return: A reference to the created item. + :rtype: ~azure.ai.agentserver.core.storage.StateStoreItemRef + """ + body = serialize_item_create_request(key, value, tags) + response = await self._send_storage_request( + self._request("POST", f"{self._store_path()}/items", content=body, include_user_id=True) + ) + return deserialize_state_item_ref(response.text()) + + async def set_item( + self, + key: str, + value: JSONObject, + *, + tags: Mapping[str, str] | None = None, + if_match: str | None = None, + require_exists: bool = False, + ) -> StateStoreItemRef: + """Create or replace one item by key. + + :param key: The item key to create or replace. + :type key: str + :param value: The item payload as a JSON object. + :type value: dict[str, any] + :keyword tags: Optional string labels stored with the item and used to + filter :meth:`list_keys`. + :paramtype tags: ~collections.abc.Mapping[str, str] or None + :keyword if_match: Optional concurrency token; the write only succeeds + when the current item ETag matches. Mutually exclusive with + ``require_exists``. + :paramtype if_match: str or None + :keyword require_exists: When ``True``, only replace an existing item and + fail if the key is absent. Mutually exclusive with ``if_match``. + :paramtype require_exists: bool + :return: A reference to the created or replaced item. + :rtype: ~azure.ai.agentserver.core.storage.StateStoreItemRef + """ + if if_match is not None and require_exists: + raise ValueError("if_match and require_exists are mutually exclusive") + body = serialize_item_put_request(value, tags) + header = "*" if require_exists else if_match + response = await self._send_storage_request( + self._request( + "PUT", + self._item_path(key), + content=body, + include_user_id=True, + if_match=header, + ) + ) + return deserialize_state_item_ref(response.text()) + + async def get(self) -> StateStore: + """Fetch the bound store's own descriptor. + + :return: The store descriptor. + :rtype: ~azure.ai.agentserver.core.storage.StateStore + :raises ~azure.ai.agentserver.core.storage.FoundryStorageNotFoundError: + If the store does not exist. + """ + return await self._fetch_properties() + + async def get_item(self, key: str) -> StateStoreItem | None: + """Fetch one item by key. + + :param key: The item key to fetch. + :type key: str + :return: The item, or ``None`` if it does not exist. + :rtype: ~azure.ai.agentserver.core.storage.StateStoreItem or None + """ + try: + response = await self._send_storage_request( + self._request("GET", self._item_path(key), include_user_id=True) + ) + except FoundryStorageNotFoundError: + return None + return deserialize_state_item(response.text()) + + async def delete(self) -> DeletedStateStore: + """Delete the bound store, cascading to every item under it. + + :return: The deleted-store marker. + :rtype: ~azure.ai.agentserver.core.storage.DeletedStateStore + """ + response = await self._send_storage_request(self._request("DELETE", self._store_path())) + return deserialize_deleted_state_store(response.text()) + + async def delete_item(self, key: str, *, if_match: str | None = None) -> DeletedStateStoreItem: + """Delete one item by key. + + :param key: The item key to delete. + :type key: str + :keyword if_match: Optional concurrency token. + :paramtype if_match: str or None + :return: The deleted-item marker. + :rtype: ~azure.ai.agentserver.core.storage.DeletedStateStoreItem + """ + response = await self._send_storage_request( + self._request("DELETE", self._item_path(key), include_user_id=True, if_match=if_match) + ) + return deserialize_deleted_state_item(response.text()) + + async def list_keys( + self, + *, + tags: Mapping[str, str] | None = None, + limit: int | None = None, + after: str | None = None, + before: str | None = None, + order: Order = "desc", + ) -> StateStoreItemKeyPage: + """List keys within the bound store. + + :keyword tags: Optional tag filter; only keys whose tags match every + provided pair are returned. + :paramtype tags: ~collections.abc.Mapping[str, str] or None + :keyword limit: Optional maximum number of keys to return in the page. + :paramtype limit: int or None + :keyword after: Return keys ordered after this cursor. Mutually exclusive + with ``before``. + :paramtype after: str or None + :keyword before: Return keys ordered before this cursor. Mutually + exclusive with ``after``. + :paramtype before: str or None + :keyword order: Sort order, ``"asc"`` or ``"desc"`` (default ``"desc"``). + :paramtype order: str + :return: A page of item keys. + :rtype: ~azure.ai.agentserver.core.storage.StateStoreItemKeyPage + """ + if after is not None and before is not None: + raise ValueError("after and before are mutually exclusive") + query: dict[str, str] = {} + if tags is not None: + for key, value in tags.items(): + query[f"tags.{key}"] = value + if limit is not None: + query["limit"] = str(limit) + if after is not None: + query["after"] = after + if before is not None: + query["before"] = before + query["order"] = order + response = await self._send_storage_request( + self._request("GET", f"{self._store_path()}/items:keys", include_user_id=True, query=query) + ) + return deserialize_list_keys_response(response.text()) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state_serializer.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state_serializer.py new file mode 100644 index 000000000000..aa9a2feb3bfd --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state_serializer.py @@ -0,0 +1,176 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Serialization helpers and value types for the Foundry state-store client. + +Request and response bodies are typed models generated from a formal +TypeSpec contract (see ``azure.ai.agentserver.core.storage._generated``) +instead of hand-written dataclasses + ad hoc JSON (de)serialization. The +resource/response model names below are re-exported unchanged from that +generated package; this module only adds the thin ``serialize_*`` / +``deserialize_*`` helpers that translate between those models and the raw +JSON bytes/text the HTTP layer sends and receives, plus a couple of +hand-written convenience types (``JSONObject``, ``Order``, ``StateStoreItemKeyPage``) that +have no wire representation of their own. +""" + +# Internal serialize/deserialize helpers below intentionally omit per-param docs. +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Literal, Union + +try: + from typing import TypeAlias +except ImportError: # pragma: no cover + from typing_extensions import TypeAlias # type: ignore[assignment] # Python <3.10 fallback + +from ._generated import ( + CreateItemRequest, + CreateStateStoreRequest, + DeletedStateStore, + DeletedStateStoreItem, + ListResponseStateStoreItemKey, + PutItemRequest, + StateStore, + StateStoreItem, + StateStoreItemRef, + StateStoreItemKey, + UpdateStateStoreRequest, +) +from ._json import load_json + +JSONValue: TypeAlias = Union[ + str, + int, + float, + bool, + None, + list["JSONValue"], + dict[str, "JSONValue"], +] +JSONObject: TypeAlias = dict[str, JSONValue] +Order: TypeAlias = Literal["asc", "desc"] + +__all__ = [ + "JSONObject", + "JSONValue", + "Order", + "StateStore", + "DeletedStateStore", + "StateStoreItem", + "StateStoreItemRef", + "DeletedStateStoreItem", + "StateStoreItemKey", + "StateStoreItemKeyPage", + "serialize_store_create_request", + "serialize_store_update_request", + "serialize_item_create_request", + "serialize_item_put_request", + "deserialize_state_store", + "deserialize_deleted_state_store", + "deserialize_state_item_ref", + "deserialize_state_item", + "deserialize_deleted_state_item", + "deserialize_list_keys_response", +] + + +@dataclass +class StateStoreItemKeyPage: + """A page of keys returned by ``list_keys``. + + A hand-written convenience wrapper (not itself part of the wire contract) + over the generated ``ListResponseStateStoreItemKey`` envelope, exposing its + ``data`` array as ``keys`` to match this client's naming. + """ + + keys: list[StateStoreItemKey] + first_id: str | None = None + last_id: str | None = None + has_more: bool = False + + +_UNSET = object() + + +def _to_wire_tags(tags: Mapping[str, str] | None) -> dict[str, str] | None: + return dict(tags) if tags else None + + +def serialize_store_create_request( + name: str, + *, + user_isolation: bool, + item_ttl_seconds: int, + description: str | None, + tags: Mapping[str, str], +) -> bytes: + request = CreateStateStoreRequest( + name=name, + user_isolation=user_isolation, + item_ttl_seconds=item_ttl_seconds, + description=description, + tags=_to_wire_tags(tags), + ) + return json.dumps(dict(request)).encode("utf-8") + + +def serialize_store_update_request(description: str | None | object, tags: Mapping[str, str] | None | object) -> bytes: + # description/tags use the module-level _UNSET sentinel to distinguish + # "leave unchanged" (key omitted) from "clear" (key present with a null + # value). The generated model's *mapping* constructor preserves that + # distinction (unlike its kwargs constructor, which drops an explicit + # ``None`` as if the field had been omitted) -- see model_base.Model + # .__init__: the kwargs branch filters ``v is not None``, the mapping + # branch does not. + payload: dict[str, Any] = {} + if description is not _UNSET: + payload["description"] = description + if tags is not _UNSET: + payload["tags"] = dict(tags) if isinstance(tags, Mapping) else {} + request = UpdateStateStoreRequest(payload) + return json.dumps(dict(request)).encode("utf-8") + + +def serialize_item_create_request(key: str, value: JSONObject, tags: Mapping[str, str] | None) -> bytes: + request = CreateItemRequest(key=key, value=value, tags=_to_wire_tags(tags)) + return json.dumps(dict(request)).encode("utf-8") + + +def serialize_item_put_request(value: JSONObject, tags: Mapping[str, str] | None) -> bytes: + request = PutItemRequest(value=value, tags=_to_wire_tags(tags)) + return json.dumps(dict(request)).encode("utf-8") + + +def deserialize_state_store(body: str) -> StateStore: + return StateStore(load_json(body)) + + +def deserialize_deleted_state_store(body: str) -> DeletedStateStore: + return DeletedStateStore(load_json(body)) + + +def deserialize_state_item_ref(body: str) -> StateStoreItemRef: + return StateStoreItemRef(load_json(body)) + + +def deserialize_state_item(body: str) -> StateStoreItem: + return StateStoreItem(load_json(body)) + + +def deserialize_deleted_state_item(body: str) -> DeletedStateStoreItem: + return DeletedStateStoreItem(load_json(body)) + + +def deserialize_list_keys_response(body: str) -> StateStoreItemKeyPage: + page = ListResponseStateStoreItemKey(load_json(body)) + return StateStoreItemKeyPage( + keys=list(page.data or []), + first_id=page.first_id, + last_id=page.last_id, + has_more=bool(page.has_more), + ) diff --git a/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md b/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md new file mode 100644 index 000000000000..59f59bcc9185 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md @@ -0,0 +1,340 @@ +# Durable State Store Guide + +`FoundryStateStore` is a durable, server-backed store for agent state. Each +store holds JSON items that you read, write, and list by key. + +## Overview + +`FoundryStateStore` is **bound to one caller-chosen store name**. That store +name is the main scoping tool for your data: + +- Use one store per conversation/thread when you need conversation isolation. +- Use `user_isolation=True` when the store name is shared across many users and + the platform should partition items per user. +- Set `item_ttl_seconds` once at store creation when you want idle items to + age out automatically. + +## Typed Models + +Every request and response is a typed model class, so you get IDE completion +and type-checker support. Import them from +`azure.ai.agentserver.core.storage`: + +| Returned by | Model | +|---|---| +| `get_or_create()` | `FoundryStateStore` (bound client) | +| `get()`, `update()` | `StateStore` | +| `delete()` | `DeletedStateStore` | +| `create_item()`, `set_item()` | `StateStoreItemRef` | +| `get_item(key)` | `StateStoreItem` | +| `delete_item(key)` | `DeletedStateStoreItem` | +| `list_keys()` | `StateStoreItemKeyPage` (of `StateStoreItemKey`) | + +Item values (`StateStoreItem.value`) are your own application JSON. Pass a +plain `dict` when writing and read one back on `get_item()`; the SDK stores and +returns the value as-is. + +```python +from azure.ai.agentserver.core.storage import StateStore, StateStoreItem + +store_info: StateStore = await store.get() +item: StateStoreItem | None = await store.get_item("step-1") +``` + +## Getting Started + +`get_or_create()` is the recommended entry point: it fetches the store, or +creates it if it does not exist, in a single call -- so you can read and write +items right away. + +```python +from azure.ai.agentserver.core.storage import FoundryStateStore, StateStoreItem + +store = await FoundryStateStore.get_or_create( + "checkpoints/thread-abc", + user_isolation=True, + item_ttl_seconds=3600, + description="Checkpoint store for thread abc", +) +async with store: + await store.set_item("step-1", {"done": False}) + + item: StateStoreItem | None = await store.get_item("step-1") + assert item is not None + print(item.value) # {"done": False} + print(item.etag) +``` + +By default, the client resolves: + +- `FOUNDRY_PROJECT_ENDPOINT` for the project endpoint +- `DefaultAzureCredential` for authentication (requires `azure-identity`) + +## Store Name = Scope + +To scope data to a conversation or thread, encode it directly into the store +name: + +```python +await FoundryStateStore.get_or_create("checkpoints/thread-abc") +await FoundryStateStore.get_or_create("workflow-state/run-42") +await FoundryStateStore.get_or_create("user-prefs/defaults", user_isolation=True) +``` + +Because the store name is its identity, choose a stable naming scheme up +front. Names may contain `/`, so you can use it as a hierarchy separator. + +## Store Lifecycle + +`get_or_create()` is the only lifecycle call you need for the common case: + +```python +store = await FoundryStateStore.get_or_create( + "checkpoints/thread-abc", + user_isolation=True, + item_ttl_seconds=3600, +) +print(store.name) +``` + +Store-level operations (`get()`, `update()`, `delete()`) act on the bound +store itself; the explicit `*_item` methods act on individual items within it. + +```python +info: StateStore = await store.get() # the store's metadata; raises if absent +info = await store.update( + description="Checkpoint store for prod traffic", + tags={"env": "prod", "team": "agents"}, +) + +deleted = await store.delete() # deletes the store, cascading to every item +assert deleted.deleted is True +``` + +### Key points + +- `get_or_create()` fetches the store first, or creates it when it is absent + (falling back to a fetch if another caller created it in the meantime). It + does not update `user_isolation`, `item_ttl_seconds`, `description`, or + `tags` on a store that already exists -- those are only applied on first + creation. +- `update()` only changes `description` and `tags`. +- `user_isolation` and `item_ttl_seconds` are fixed at create time. +- `delete()` cascades to every item under that store name. + +## User Isolation and Delegated User IDs + +Set `user_isolation=True` when the same store name should fan out per user. + +```python +store = await FoundryStateStore.get_or_create( + "user-prefs/defaults", + user_isolation=True, + user_id="aad-user-42", +) +``` + +- For direct callers, the platform can derive user identity from the token. +- For trusted callers acting on behalf of an end user, pass `user_id` so the SDK + sends the delegated `x-ms-user-id` header on item operations. +- Store-management calls (`get_or_create`, `get()`, `update()`, `delete()`) + stay store-scoped and do not send the delegated user + header. + +## Values, Tags, and TTL + +Each item value is a JSON object -- pass a `dict`. + +```python +await store.create_item( + "step-1", + {"done": False, "attempt": 1}, + tags={"kind": "checkpoint"}, +) +``` + +Tags are simple string labels used only for filtering `list_keys()`. + +TTL is **store-level**, not per-item: + +```python +store = await FoundryStateStore.get_or_create("otp/user-42", item_ttl_seconds=300) +``` + +- Default: `30 days` +- `-1`: never expire +- Any item write renews the TTL window for that item +- Reads do **not** renew the TTL window + +## Single-Item Operations + +### Create a new item + +```python +created = await store.create_item( + "step-1", + {"done": False}, + tags={"kind": "checkpoint"}, +) +print(created.etag) +``` + +Use `create_item()` when duplicate keys should fail with `409`. + +### Create-or-replace + +```python +updated = await store.set_item( + "step-1", + {"done": True}, + tags={"kind": "checkpoint"}, +) +print(updated.etag) +``` + +`set_item()` creates the item, or replaces it if the key already exists. + +### Fetch one item + +```python +item: StateStoreItem | None = await store.get_item("step-1") +if item is not None: + print(item.id, item.key, item.value, item.tags, item.etag) +``` + +`get_item(key)` returns `None` when the item is missing; `get()` fetches the +store's metadata instead, raising `FoundryStorageNotFoundError` if the store is +absent. + +### Delete one item + +```python +deleted = await store.delete_item("step-1") +assert deleted.deleted is True +``` + +Deletes are idempotent. + +## Optimistic Concurrency + +Use `if_match` when you want a guarded update or delete: + +```python +from azure.ai.agentserver.core.storage import FoundryStoragePreconditionError + +item = await store.get_item("counter") +assert item is not None + +try: + await store.set_item("counter", {"value": item.value["value"] + 1}, if_match=item.etag) +except FoundryStoragePreconditionError as err: + print("Current etag:", err.current_etag) +``` + +If you want a strict update that only succeeds when the item already exists, use +`require_exists=True`: + +```python +await store.set_item("counter", {"value": 2}, require_exists=True) +``` + +## Listing Keys + +`list_keys()` returns a keys-only page within the bound store. + +```python +from azure.ai.agentserver.core.storage import StateStoreItemKeyPage + +page: StateStoreItemKeyPage = await store.list_keys(tags={"kind": "checkpoint"}, limit=50, order="asc") +for key in page.keys: + print(key.id, key.key, key.tags, key.etag) + +while page.has_more and page.last_id is not None: + page = await store.list_keys( + tags={"kind": "checkpoint"}, + after=page.last_id, + limit=50, + order="asc", + ) +``` + +Use: + +- `tags={...}` for AND-filtered tag matching +- `limit` for page size +- `after` / `before` for cursor paging by item id +- `order="desc"` (default) or `"asc"` + +## Error Handling + +All storage errors derive from `FoundryStorageError`. + +| Exception | HTTP | Meaning | +|---|---|---| +| `FoundryStoragePreconditionError` | 412 | `If-Match` failed; `current_etag` may be populated. | +| `FoundryStorageNotFoundError` | 404 | Store or resource path not found. | +| `FoundryStorageConflictError` | 409 | A `create`/`create_item` duplicated an existing name/key. | +| `FoundryStorageBadRequestError` | 400 | Invalid request; `param` names the offending field. | +| `FoundryStorageApiError` | other 4xx/5xx | Server-side failure. | + +```python +from azure.ai.agentserver.core.storage import ( + FoundryStorageError, + FoundryStoragePreconditionError, +) + +try: + await store.set_item("step-1", {"done": True}, if_match='"stale"') +except FoundryStoragePreconditionError as err: + print(err.current_etag) +except FoundryStorageError as err: + print(err.message, err.response_body) +``` + +## Limits + +All request-body and query fields are bounded server-side; a violating request +is rejected with `400 Bad Request` (`error.param` names the offending field). + +**Store (`get_or_create`, `update`):** + +| Field | Constraints | Mutability | +|-------|-------------|------------| +| `name` | 1-128 chars. Unicode; may contain `/` as a hierarchy separator. Unique within the project + agent. | Immutable | +| `user_isolation` | Boolean. Omitted -> `false` (agent-level, shared). | Immutable (fixed at first creation) | +| `item_ttl_seconds` | Default `2592000` (30 days); `-1` = never expire; else `1`-`2147483647`. Write-sliding per item (renews on write, not read). | Immutable (fixed at first creation) | +| `description` | <= 1024 chars. Free-form. | Mutable via `update()` | +| `tags` | <= 16 entries. Key: 1-64 chars, `[a-zA-Z0-9_.-]`. Value: <= 256 chars. Replaced wholesale. | Mutable via `update()` | + +**Item (`create_item`, `set_item`):** + +| Field | Constraints | Mutability | +|-------|-------------|------------| +| `key` | 1-128 chars. Unicode; may contain `/`. Unique within the store. | Immutable | +| `value` | Opaque JSON object, <= 1 MB serialized inline. | Mutable via `set_item()` (replace) | +| `tags` | <= 16 entries, same shape as store tags. | Mutable via `set_item()` (replace) | + +Items carry no TTL of their own -- expiry is inherited from the store's +`item_ttl_seconds`. + +**Query parameters (`list_keys`):** + +| Parameter | Constraints | +|-----------|-------------| +| `limit` | 1-100. Default `20`. | +| `order` | `"asc"` or `"desc"`. Default `"desc"`. | +| `after` / `before` | Opaque cursor; mutually exclusive. | + +## Best Practices + +1. **Prefer `get_or_create()`.** It is the only lifecycle call you need for the + common case; do not assume item writes will create the store for you. +2. **Encode conversation scope in the store name.** There is no separate + session-isolation knob. +3. **Use `user_isolation=True` only when needed.** Prefer a stable store naming + scheme first, then add per-user partitioning when the store name is shared. +4. **Use `if_match` for read-modify-write flows.** Counters and checkpoints are + race-prone without it. +5. **Keep values as JSON objects.** Serialize your own models explicitly. +6. **Reuse the client.** It owns an HTTP pipeline; construct it once (via + `get_or_create()`) and close it with `async with` or `await store.aclose()`. diff --git a/sdk/agentserver/azure-ai-agentserver-core/mypy.ini b/sdk/agentserver/azure-ai-agentserver-core/mypy.ini new file mode 100644 index 000000000000..850cfd736a67 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/mypy.ini @@ -0,0 +1,8 @@ +[mypy] +explicit_package_bases = True + +[mypy-samples.*] +ignore_errors = true + +[mypy-azure.ai.agentserver.core.storage._generated.*] +ignore_errors = true diff --git a/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml index cbe9448f849f..989f85540252 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml @@ -21,6 +21,8 @@ classifiers = [ keywords = ["azure", "azure sdk", "agent", "agentserver", "core"] dependencies = [ + "azure-core>=1.37.0", + "isodate>=0.6.1", "starlette>=0.45.0", "hypercorn>=0.17.0", "opentelemetry-api>=1.43.0", diff --git a/sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py b/sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py new file mode 100644 index 000000000000..7e9de6bce757 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py @@ -0,0 +1,116 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +FILE: state_store_sample.py + +DESCRIPTION: + Demonstrates the explicit Foundry state-store client (`FoundryStateStore`): + resolving a store with `get_or_create`, item create/update/get/delete, + metadata updates, optimistic concurrency, and tag-filtered key listing. + +USAGE: + python state_store_sample.py + + Set these environment variables before running: + 1) FOUNDRY_PROJECT_ENDPOINT - the Azure AI Foundry project endpoint. + + The sample authenticates with DefaultAzureCredential, so sign in with the + Azure CLI (`az login`) or configure another supported credential source. + Requires the `azure-identity` package. +""" + +import asyncio +from uuid import uuid4 + +from azure.ai.agentserver.core.storage import ( + FoundryStateStore, + FoundryStoragePreconditionError, +) + + +async def create_and_get_item(store: FoundryStateStore) -> None: + """Create an item, then fetch it back.""" + created = await store.create_item( + "step-1", + {"done": False, "attempt": 1}, + tags={"kind": "checkpoint"}, + ) + print(f"created step-1 with etag={created.etag}") + + item = await store.get_item("step-1") + assert item is not None + print(f"get step-1 -> value={item.value!r}, tags={item.tags}, etag={item.etag}") + + +async def update_metadata(store: FoundryStateStore) -> None: + """Update mutable store metadata.""" + info = await store.update( + description="Sample checkpoint store", + tags={"scenario": "state-store-sample", "env": "dev"}, + ) + print(f"updated store metadata -> description={info.description!r}, tags={info.tags}") + + +async def optimistic_concurrency(store: FoundryStateStore) -> None: + """Guard a read-modify-write with if_match and handle the conflict.""" + item = await store.get_item("step-1") + assert item is not None + + await store.set_item("step-1", {"done": True, "attempt": 2}, tags={"kind": "checkpoint"}) + + try: + await store.set_item( + "step-1", + {"done": True, "attempt": 3}, + tags={"kind": "checkpoint"}, + if_match=item.etag, + ) + except FoundryStoragePreconditionError as err: + print(f"lost the race as expected; current_etag={err.current_etag}") + + +async def list_with_tags(store: FoundryStateStore) -> None: + """List keys filtered by tag and page using last_id.""" + await store.create_item("step-2", {"done": False, "attempt": 1}, tags={"kind": "checkpoint"}) + await store.create_item("audit-1", {"event": "created"}, tags={"kind": "audit"}) + + page = await store.list_keys(tags={"kind": "checkpoint"}, limit=1, order="asc") + while True: + print(f"checkpoint keys: {[entry.key for entry in page.keys]}") + if not page.has_more or page.last_id is None: + break + page = await store.list_keys(tags={"kind": "checkpoint"}, after=page.last_id, limit=1, order="asc") + + +async def delete_item_and_store(store: FoundryStateStore) -> None: + """Delete one item, then cascade-delete the whole store.""" + deleted_item = await store.delete_item("audit-1") + print(f"deleted item -> key={deleted_item.key}, deleted={deleted_item.deleted}") + + deleted_store = await store.delete() + print(f"deleted store -> name={deleted_store.name}, deleted={deleted_store.deleted}") + + +async def main() -> None: + store_name = f"checkpoints/sample-thread-{uuid4().hex}" + # get_or_create resolves (or creates, on first use) the server-side store + # resource in one call, so there is no separate lifecycle step before + # reading or writing items. + store = await FoundryStateStore.get_or_create( + store_name, + user_isolation=True, + item_ttl_seconds=3600, + description="Sample state store", + ) + async with store: + print(f"using store name={store.name}") + await create_and_get_item(store) + await update_metadata(store) + await optimistic_concurrency(store) + await list_with_tags(store) + await delete_item_and_store(store) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py new file mode 100644 index 000000000000..db945bea5125 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py @@ -0,0 +1,624 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Unit tests for FoundryStateStore request construction and response handling.""" + +from __future__ import annotations + +import base64 +import json +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from azure.ai.agentserver.core.storage import ( + DeletedStateStore, + DeletedStateStoreItem, + FoundryStateStore, + FoundryStorageConflictError, + FoundryStorageEndpoint, + FoundryStorageNotFoundError, + StateStoreItemKeyPage, + StateStore, + StateStoreItem, + StateStoreItemRef, + StateStoreItemKey, +) + +_BASE_URL = "https://foundry.example.com/storage/" +_ENDPOINT = FoundryStorageEndpoint(storage_base_url=_BASE_URL) + + +def _encode_segment(value: str) -> str: + encoded = base64.urlsafe_b64encode(value.encode("utf-8")).decode("ascii") + return encoded.rstrip("=") + + +def _make_response(status_code: int, body: Any, *, headers: dict[str, str] | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = status_code + resp.headers = headers or {} + resp.text = MagicMock(return_value=json.dumps(body)) + return resp + + +def _make_store( + response: MagicMock, + *, + name: str = "langGraphCheckpoints/thread-abc", + user_isolation: bool = False, + item_ttl_seconds: int = 2592000, + description: str | None = None, + tags: dict[str, str] | None = None, + user_id: str | None = None, +) -> FoundryStateStore: + store = FoundryStateStore.__new__(FoundryStateStore) + store._endpoint = _ENDPOINT + store._owns_credential = False + store._name = name + store._user_isolation = user_isolation + store._item_ttl_seconds = item_ttl_seconds + store._description = description + store._tags = {} if tags is None else dict(tags) + store._user_id = user_id + mock_pipeline = AsyncMock() + mock_pipeline.send_request = AsyncMock(return_value=response) + mock_pipeline.close = AsyncMock() + store._client = mock_pipeline + return store + + +def _make_store_with_responses( + *responses: MagicMock, name: str = "langGraphCheckpoints/thread-abc" +) -> FoundryStateStore: + store = _make_store(responses[0], name=name) + store._client.send_request = AsyncMock(side_effect=list(responses)) + return store + + +def _sent_request(store: FoundryStateStore) -> Any: + return store._client.send_request.call_args[0][0] + + +def _state_store_body(**overrides: Any) -> dict[str, Any]: + body: dict[str, Any] = { + "id": "ss_1", + "object": "state_store", + "name": "checkpoints", + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": None, + "tags": {}, + "created_at": 1, + "updated_at": 1, + } + body.update(overrides) + return body + + +def _state_store(**overrides: Any) -> StateStore: + return StateStore(_state_store_body(**overrides)) + + +def _item_metadata_body(**overrides: Any) -> dict[str, Any]: + body: dict[str, Any] = { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + } + body.update(overrides) + return body + + +def _item_metadata(**overrides: Any) -> StateStoreItemRef: + return StateStoreItemRef(_item_metadata_body(**overrides)) + + +def _item_body(**overrides: Any) -> dict[str, Any]: + body: dict[str, Any] = { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "value": {}, + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + } + body.update(overrides) + return body + + +def _item(**overrides: Any) -> StateStoreItem: + return StateStoreItem(_item_body(**overrides)) + + +def _key_body(**overrides: Any) -> dict[str, Any]: + body: dict[str, Any] = { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + } + body.update(overrides) + return body + + +def _key(**overrides: Any) -> StateStoreItemKey: + return StateStoreItemKey(_key_body(**overrides)) + + +# --------------------------------------------------------------------------- +# get_or_create (classmethod orchestration: fetch / create / conflict-refetch) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_or_create_returns_existing_store_when_present(monkeypatch: pytest.MonkeyPatch) -> None: + info = _state_store() + fetch = AsyncMock(return_value=info) + create = AsyncMock() + monkeypatch.setattr(FoundryStateStore, "_fetch_properties", fetch) + monkeypatch.setattr(FoundryStateStore, "_create_properties", create) + + store = await FoundryStateStore.get_or_create("checkpoints", credential=MagicMock(), endpoint=_ENDPOINT) + try: + fetch.assert_awaited_once() + create.assert_not_awaited() + assert store.name == "checkpoints" + finally: + await store.aclose() + + +@pytest.mark.asyncio +async def test_get_or_create_creates_store_when_absent(monkeypatch: pytest.MonkeyPatch) -> None: + info = _state_store() + fetch = AsyncMock(side_effect=FoundryStorageNotFoundError("not found")) + create = AsyncMock(return_value=info) + monkeypatch.setattr(FoundryStateStore, "_fetch_properties", fetch) + monkeypatch.setattr(FoundryStateStore, "_create_properties", create) + + store = await FoundryStateStore.get_or_create("checkpoints", credential=MagicMock(), endpoint=_ENDPOINT) + try: + fetch.assert_awaited_once() + create.assert_awaited_once() + assert store.name == "checkpoints" + finally: + await store.aclose() + + +@pytest.mark.asyncio +async def test_get_or_create_refetches_when_create_races_with_another_caller(monkeypatch: pytest.MonkeyPatch) -> None: + created_elsewhere = _state_store() + fetch = AsyncMock(side_effect=[FoundryStorageNotFoundError("not found"), created_elsewhere]) + create = AsyncMock(side_effect=FoundryStorageConflictError("duplicate store")) + monkeypatch.setattr(FoundryStateStore, "_fetch_properties", fetch) + monkeypatch.setattr(FoundryStateStore, "_create_properties", create) + + store = await FoundryStateStore.get_or_create("checkpoints", credential=MagicMock(), endpoint=_ENDPOINT) + try: + assert fetch.await_count == 2 + create.assert_awaited_once() + assert store.name == "checkpoints" + finally: + await store.aclose() + + +@pytest.mark.asyncio +async def test_get_or_create_closes_store_when_fetch_fails(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(FoundryStateStore, "_fetch_properties", AsyncMock(side_effect=RuntimeError("boom"))) + closed = AsyncMock() + monkeypatch.setattr(FoundryStateStore, "aclose", closed) + + with pytest.raises(RuntimeError, match="boom"): + await FoundryStateStore.get_or_create("checkpoints", credential=MagicMock(), endpoint=_ENDPOINT) + + closed.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_or_create_closes_store_when_create_fails(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + FoundryStateStore, "_fetch_properties", AsyncMock(side_effect=FoundryStorageNotFoundError("not found")) + ) + monkeypatch.setattr( + FoundryStateStore, "_create_properties", AsyncMock(side_effect=RuntimeError("create-fail")) + ) + closed = AsyncMock() + monkeypatch.setattr(FoundryStateStore, "aclose", closed) + + with pytest.raises(RuntimeError, match="create-fail"): + await FoundryStateStore.get_or_create("checkpoints", credential=MagicMock(), endpoint=_ENDPOINT) + + closed.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_or_create_forwards_creation_options_to_create() -> None: + info = _state_store( + user_isolation=True, + item_ttl_seconds=600, + description="checkpoint store", + tags={"team": "agents"}, + ) + store = _make_store( + _make_response( + 201, + _state_store_body( + user_isolation=True, + item_ttl_seconds=600, + description="checkpoint store", + tags={"team": "agents"}, + ), + ), + name="checkpoints", + user_isolation=True, + item_ttl_seconds=600, + description="checkpoint store", + tags={"team": "agents"}, + ) + + result = await store._create_properties() + + request = _sent_request(store) + assert json.loads(request.content.decode("utf-8")) == { + "name": "checkpoints", + "user_isolation": True, + "item_ttl_seconds": 600, + "description": "checkpoint store", + "tags": {"team": "agents"}, + } + assert result == info + + +# --------------------------------------------------------------------------- +# get() -- overloaded on key: None fetches the store descriptor, else an item +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_get_with_no_key_returns_the_store_descriptor() -> None: + store_name = "langGraphCheckpoints/thread-abc" + store = _make_store( + _make_response( + 200, + { + "id": "ss_1", + "object": "state_store", + "name": store_name, + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": None, + "tags": {}, + "created_at": 1, + "updated_at": 2, + }, + ), + name=store_name, + ) + + result = await store.get() + + request = _sent_request(store) + assert request.method == "GET" + assert request.url == f"{_BASE_URL}state_stores/{_encode_segment(store_name)}?api-version=v1" + assert "x-ms-user-id" not in request.headers # store-level ops never send the delegated user header + assert result is not None + assert result.name == store_name + assert result.id == "ss_1" + + +@pytest.mark.asyncio +async def test_get_with_no_key_raises_when_store_is_absent() -> None: + store = _make_store(_make_response(404, {"error": {"message": "not found"}}), name="checkpoints") + + with pytest.raises(FoundryStorageNotFoundError): + await store.get() + + +@pytest.mark.asyncio +async def test_get_with_key_returns_state_item_with_value_and_metadata() -> None: + store = _make_store( + _make_response( + 200, + { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "value": {"done": True}, + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + }, + ), + name="checkpoints", + user_id="user-42", + ) + + result = await store.get_item("step/1") + + request = _sent_request(store) + assert request.method == "GET" + assert request.headers["x-ms-user-id"] == "user-42" + assert request.url == ( + f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}/items/{_encode_segment('step/1')}?api-version=v1" + ) + assert result == _item( + value={"done": True}, + tags={"kind": "checkpoint"}, + ) + + +@pytest.mark.asyncio +async def test_get_with_key_returns_none_when_item_is_absent() -> None: + store = _make_store(_make_response(404, {"error": {"message": "not found"}}), name="checkpoints") + assert await store.get_item("missing") is None + + +# --------------------------------------------------------------------------- +# update() -- store mutable metadata (was update_metadata) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_update_sends_only_present_fields() -> None: + store = _make_store( + _make_response( + 200, + { + "id": "ss_1", + "object": "state_store", + "name": "prefs", + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": "updated", + "tags": {"env": "prod"}, + "created_at": 1, + "updated_at": 3, + }, + ), + name="prefs", + ) + + result = await store.update(description="updated", tags={"env": "prod"}) + + request = _sent_request(store) + assert request.method == "PATCH" + assert request.url == f"{_BASE_URL}state_stores/{_encode_segment('prefs')}?api-version=v1" + assert json.loads(request.content.decode("utf-8")) == {"description": "updated", "tags": {"env": "prod"}} + assert result.updated_at == 3 + + +@pytest.mark.asyncio +async def test_update_with_only_description_omits_tags() -> None: + store = _make_store( + _make_response(200, _state_store_body(name="prefs", description="only-desc")), + name="prefs", + ) + + await store.update(description="only-desc") + + request = _sent_request(store) + assert json.loads(request.content.decode("utf-8")) == {"description": "only-desc"} + + +@pytest.mark.asyncio +async def test_update_with_only_tags_omits_description() -> None: + store = _make_store( + _make_response(200, _state_store_body(name="prefs", tags={"env": "prod"})), + name="prefs", + ) + + await store.update(tags={"env": "prod"}) + + request = _sent_request(store) + assert json.loads(request.content.decode("utf-8")) == {"tags": {"env": "prod"}} + + +@pytest.mark.asyncio +async def test_update_with_no_arguments_sends_empty_body() -> None: + store = _make_store( + _make_response(200, _state_store_body(name="prefs")), + name="prefs", + ) + + await store.update() + + request = _sent_request(store) + assert json.loads(request.content.decode("utf-8")) == {} + + +# --------------------------------------------------------------------------- +# delete() -- overloaded on key: None deletes the store, else one item +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_delete_with_no_key_deletes_the_store() -> None: + store = _make_store( + _make_response(200, {"id": "ss_1", "object": "state_store", "name": "prefs", "deleted": True}), + name="prefs", + ) + + result = await store.delete() + + request = _sent_request(store) + assert request.method == "DELETE" + assert request.url == f"{_BASE_URL}state_stores/{_encode_segment('prefs')}?api-version=v1" + assert "x-ms-user-id" not in request.headers + assert result == DeletedStateStore({"id": "ss_1", "object": "state_store", "name": "prefs", "deleted": True}) + + +@pytest.mark.asyncio +async def test_delete_with_key_returns_deleted_item_marker() -> None: + store = _make_store( + _make_response(200, {"id": "it_1", "object": "state_store.item", "key": "step/1", "deleted": True}), + name="checkpoints", + user_id="user-42", + ) + + result = await store.delete_item("step/1", if_match='"0x8DD"') + + request = _sent_request(store) + assert request.method == "DELETE" + assert request.headers["If-Match"] == '"0x8DD"' + assert request.headers["x-ms-user-id"] == "user-42" + assert result == DeletedStateStoreItem( + {"id": "it_1", "object": "state_store.item", "key": "step/1", "deleted": True} + ) + + +# --------------------------------------------------------------------------- +# Item operations unaffected by the store-admin refactor +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_create_item_posts_key_value_and_tags() -> None: + store = _make_store( + _make_response( + 201, + { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "etag": '"0x8DC"', + "created_at": 10, + "updated_at": 10, + }, + ), + name="checkpoints", + ) + + result = await store.create_item("step/1", {"done": False}, tags={"kind": "checkpoint"}) + + request = _sent_request(store) + assert request.method == "POST" + assert request.url == f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}/items?api-version=v1" + assert json.loads(request.content.decode("utf-8")) == { + "key": "step/1", + "value": {"done": False}, + "tags": {"kind": "checkpoint"}, + } + assert "If-Match" not in request.headers + assert result == _item_metadata(etag='"0x8DC"', created_at=10, updated_at=10) + + +@pytest.mark.asyncio +async def test_set_puts_value_and_if_match_header() -> None: + store = _make_store( + _make_response( + 200, + { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + }, + headers={"ETag": '"0x8DD"'}, + ), + name="checkpoints", + ) + + result = await store.set_item("step/1", {"done": True}, tags={"kind": "checkpoint"}, if_match='"0x8DC"') + + request = _sent_request(store) + assert request.method == "PUT" + assert request.url == ( + f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}/items/{_encode_segment('step/1')}?api-version=v1" + ) + assert request.headers["If-Match"] == '"0x8DC"' + assert json.loads(request.content.decode("utf-8")) == {"value": {"done": True}, "tags": {"kind": "checkpoint"}} + assert result.etag == '"0x8DD"' + + +@pytest.mark.asyncio +async def test_set_require_exists_uses_wildcard_if_match() -> None: + store = _make_store( + _make_response( + 200, + { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + }, + ), + name="checkpoints", + ) + + await store.set_item("step/1", {"done": True}, require_exists=True) + + request = _sent_request(store) + assert request.headers["If-Match"] == "*" + + +@pytest.mark.asyncio +async def test_list_keys_uses_query_parameters_and_returns_page() -> None: + store = _make_store( + _make_response( + 200, + { + "object": "list", + "data": [ + { + "id": "it_1", + "object": "state_store.item", + "key": "step/1", + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + } + ], + "first_id": "it_1", + "last_id": "it_1", + "has_more": False, + }, + ), + name="checkpoints", + user_id="user-42", + ) + + page = await store.list_keys(tags={"kind": "checkpoint", "phase": "run"}, limit=10, after="it_0", order="asc") + + request = _sent_request(store) + assert request.method == "GET" + assert request.headers["x-ms-user-id"] == "user-42" + assert request.url == ( + f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}/items:keys" + "?api-version=v1&tags.kind=checkpoint&tags.phase=run&limit=10&after=it_0&order=asc" + ) + assert page == StateStoreItemKeyPage( + keys=[_key(tags={"kind": "checkpoint"})], + first_id="it_1", + last_id="it_1", + has_more=False, + ) + + +@pytest.mark.asyncio +async def test_list_keys_defaults_to_desc_order() -> None: + store = _make_store(_make_response(200, {"object": "list", "data": [], "has_more": False}), name="checkpoints") + + await store.list_keys() + + request = _sent_request(store) + assert ( + request.url == f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}/items:keys?api-version=v1&order=desc" + ) + + +def test_empty_key_is_rejected() -> None: + store = _make_store(_make_response(200, {}), name="checkpoints") + + with pytest.raises(ValueError): + store._item_path("") diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_state_store_sample_usage.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_state_store_sample_usage.py new file mode 100644 index 000000000000..1d93abdceb61 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_state_store_sample_usage.py @@ -0,0 +1,246 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Programmatic sample-flow coverage for the Foundry state-store sample.""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest +from azure.ai.agentserver.core.storage import ( + FoundryStateStore, + FoundryStorageEndpoint, + FoundryStoragePreconditionError, +) + +_BASE_URL = "https://foundry.example.com/storage/" +_ENDPOINT = FoundryStorageEndpoint(storage_base_url=_BASE_URL) + + +def _make_response(status_code: int, body: object, *, headers: dict[str, str] | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = status_code + resp.headers = {} if headers is None else headers + resp.text = MagicMock(return_value=json.dumps(body)) + return resp + + +def _make_store_with_responses(*responses: MagicMock) -> FoundryStateStore: + store = FoundryStateStore.__new__(FoundryStateStore) + store._endpoint = _ENDPOINT + store._owns_credential = False + store._name = "checkpoints/thread-abc" + store._user_isolation = True + store._item_ttl_seconds = 3600 + store._description = "Sample state store" + store._tags = {} + store._user_id = "user-42" + mock_pipeline = AsyncMock() + mock_pipeline.send_request = AsyncMock(side_effect=list(responses)) + mock_pipeline.close = AsyncMock() + store._client = mock_pipeline + return store + + +@pytest.mark.asyncio +async def test_state_store_sample_flow() -> None: + store = _make_store_with_responses( + _make_response( + 200, + { + "id": "ss_1", + "object": "state_store", + "name": "checkpoints/thread-abc", + "user_isolation": True, + "item_ttl_seconds": 3600, + "description": "Sample state store", + "tags": {}, + "created_at": 1, + "updated_at": 1, + }, + ), + _make_response( + 201, + { + "id": "it_1", + "object": "state_store.item", + "key": "step-1", + "etag": '"0x8DA"', + "created_at": 2, + "updated_at": 2, + }, + ), + _make_response( + 200, + { + "id": "it_1", + "object": "state_store.item", + "key": "step-1", + "value": {"done": False, "attempt": 1}, + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DA"', + "created_at": 2, + "updated_at": 2, + }, + ), + _make_response( + 200, + { + "id": "ss_1", + "object": "state_store", + "name": "checkpoints/thread-abc", + "user_isolation": True, + "item_ttl_seconds": 3600, + "description": "Sample checkpoint store", + "tags": {"scenario": "state-store-sample", "env": "dev"}, + "created_at": 1, + "updated_at": 3, + }, + ), + _make_response( + 200, + { + "id": "it_1", + "object": "state_store.item", + "key": "step-1", + "value": {"done": False, "attempt": 1}, + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DA"', + "created_at": 2, + "updated_at": 2, + }, + ), + _make_response( + 200, + { + "id": "it_1", + "object": "state_store.item", + "key": "step-1", + "etag": '"0x8DB"', + "created_at": 2, + "updated_at": 4, + }, + ), + _make_response( + 412, + {"error": {"message": "etag mismatch"}}, + headers={"ETag": '"0x8DB"'}, + ), + _make_response( + 201, + { + "id": "it_2", + "object": "state_store.item", + "key": "step-2", + "etag": '"0x8DC"', + "created_at": 5, + "updated_at": 5, + }, + ), + _make_response( + 201, + { + "id": "it_3", + "object": "state_store.item", + "key": "audit-1", + "etag": '"0x8DD"', + "created_at": 6, + "updated_at": 6, + }, + ), + _make_response( + 200, + { + "object": "list", + "data": [ + { + "id": "it_1", + "object": "state_store.item", + "key": "step-1", + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DB"', + "created_at": 2, + "updated_at": 4, + } + ], + "first_id": "it_1", + "last_id": "it_1", + "has_more": True, + }, + ), + _make_response( + 200, + { + "object": "list", + "data": [ + { + "id": "it_2", + "object": "state_store.item", + "key": "step-2", + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DC"', + "created_at": 5, + "updated_at": 5, + } + ], + "first_id": "it_2", + "last_id": "it_2", + "has_more": False, + }, + ), + _make_response( + 200, + {"id": "it_3", "object": "state_store.item", "key": "audit-1", "deleted": True}, + ), + _make_response( + 200, + {"id": "ss_1", "object": "state_store", "name": "checkpoints/thread-abc", "deleted": True}, + ), + ) + + # This test drives a hand-built store double directly (bypassing the real + # constructor/classmethod), so it exercises the store-resolution wire call + # via the private helper get_or_create() delegates to, rather than the + # classmethod itself (which constructs its own instance). + store_info = await store._fetch_properties() + assert store_info.id == "ss_1" + + created = await store.create_item("step-1", {"done": False, "attempt": 1}, tags={"kind": "checkpoint"}) + assert created.etag == '"0x8DA"' + + item = await store.get_item("step-1") + assert item is not None + assert item.value["attempt"] == 1 + + updated_store = await store.update( + description="Sample checkpoint store", + tags={"scenario": "state-store-sample", "env": "dev"}, + ) + assert updated_store.tags is not None + assert updated_store.tags["env"] == "dev" + + stale_item = await store.get_item("step-1") + assert stale_item is not None + await store.set_item("step-1", {"done": True, "attempt": 2}, tags={"kind": "checkpoint"}) + + with pytest.raises(FoundryStoragePreconditionError) as exc: + await store.set_item("step-1", {"done": True, "attempt": 3}, if_match=stale_item.etag) + assert exc.value.current_etag == '"0x8DB"' + + await store.create_item("step-2", {"done": False, "attempt": 1}, tags={"kind": "checkpoint"}) + await store.create_item("audit-1", {"event": "created"}, tags={"kind": "audit"}) + + first_page = await store.list_keys(tags={"kind": "checkpoint"}, limit=1, order="asc") + assert [entry.key for entry in first_page.keys] == ["step-1"] + assert first_page.has_more is True + + second_page = await store.list_keys(tags={"kind": "checkpoint"}, after=first_page.last_id, limit=1, order="asc") + assert [entry.key for entry in second_page.keys] == ["step-2"] + assert second_page.has_more is False + + deleted_item = await store.delete_item("audit-1") + assert deleted_item.deleted is True + + deleted_store = await store.delete() + assert deleted_store.deleted is True diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_endpoint.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_endpoint.py new file mode 100644 index 000000000000..fc41756d4fb8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_endpoint.py @@ -0,0 +1,57 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Unit tests for FoundryStorageEndpoint resolution and URL building.""" + +from __future__ import annotations + +import pytest +from azure.ai.agentserver.core.storage import FoundryStorageEndpoint + + +def test_from_endpoint_derives_storage_base_url() -> None: + ep = FoundryStorageEndpoint.from_endpoint("https://proj.example.com/api/projects/p") + assert ep.storage_base_url == "https://proj.example.com/api/projects/p/storage/" + assert ep.api_version == "v1" + + +def test_from_endpoint_strips_trailing_slash() -> None: + ep = FoundryStorageEndpoint.from_endpoint("https://proj.example.com/") + assert ep.storage_base_url == "https://proj.example.com/storage/" + + +def test_from_endpoint_accepts_storage_base_url() -> None: + ep = FoundryStorageEndpoint.from_endpoint("https://proj.example.com/storage/") + assert ep.storage_base_url == "https://proj.example.com/storage/" + + +def test_from_endpoint_rejects_empty() -> None: + with pytest.raises(ValueError): + FoundryStorageEndpoint.from_endpoint("") + + +def test_from_endpoint_rejects_non_absolute() -> None: + with pytest.raises(ValueError): + FoundryStorageEndpoint.from_endpoint("proj.example.com") + + +def test_from_env_reads_project_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FOUNDRY_PROJECT_ENDPOINT", "https://proj.example.com") + ep = FoundryStorageEndpoint.from_env() + assert ep.storage_base_url == "https://proj.example.com/storage/" + + +def test_from_env_requires_variable(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("FOUNDRY_PROJECT_ENDPOINT", raising=False) + with pytest.raises(EnvironmentError): + FoundryStorageEndpoint.from_env() + + +def test_build_url_appends_api_version() -> None: + ep = FoundryStorageEndpoint(storage_base_url="https://x/storage/", api_version="v1") + assert ep.build_url("state_stores") == "https://x/storage/state_stores?api-version=v1" + + +def test_build_url_appends_extra_params_encoded() -> None: + ep = FoundryStorageEndpoint(storage_base_url="https://x/storage/", api_version="v1") + url = ep.build_url("state_stores", after="it 1/2") + assert url == "https://x/storage/state_stores?api-version=v1&after=it%201%2F2" diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_errors.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_errors.py new file mode 100644 index 000000000000..4071729c8a3f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_errors.py @@ -0,0 +1,88 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Unit tests for storage error mapping (raise_for_storage_error).""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest +from azure.ai.agentserver.core._platform_headers import PLATFORM_ERROR_TAG +from azure.ai.agentserver.core.storage import ( + FoundryStorageApiError, + FoundryStorageBadRequestError, + FoundryStorageConflictError, + FoundryStorageError, + FoundryStorageNotFoundError, + FoundryStoragePreconditionError, +) +from azure.ai.agentserver.core.storage._errors import raise_for_storage_error + + +class _FakeResponse: + def __init__(self, status_code: int, body: Any, *, headers: dict[str, str] | None = None) -> None: + self.status_code = status_code + self.headers: dict[str, str] = {} if headers is None else headers + self._body = "" if body is None else json.dumps(body) + + def text(self) -> str: + return self._body + + +def test_2xx_does_not_raise() -> None: + raise_for_storage_error(_FakeResponse(204, None)) + + +def test_404_maps_to_not_found() -> None: + with pytest.raises(FoundryStorageNotFoundError) as exc: + raise_for_storage_error(_FakeResponse(404, {"error": {"message": "nope"}})) + assert exc.value.message == "nope" + assert exc.value.status_code == 404 + + +def test_400_maps_to_bad_request_and_param() -> None: + with pytest.raises(FoundryStorageBadRequestError) as exc: + raise_for_storage_error(_FakeResponse(400, {"error": {"message": "bad", "param": "item_ttl_seconds"}})) + assert exc.value.param == "item_ttl_seconds" + + +def test_409_maps_to_bad_request() -> None: + with pytest.raises(FoundryStorageBadRequestError) as exc: + raise_for_storage_error(_FakeResponse(409, {"error": {"message": "duplicate"}})) + assert isinstance(exc.value, FoundryStorageConflictError) + assert exc.value.status_code == 409 + + +def test_412_maps_to_precondition_with_current_etag() -> None: + with pytest.raises(FoundryStoragePreconditionError) as exc: + raise_for_storage_error( + _FakeResponse(412, {"error": {"message": "etag"}}, headers={"ETag": '"0x8DD"'}) + ) + assert exc.value.current_etag == '"0x8DD"' + + +def test_412_without_current_etag_defaults_to_none() -> None: + with pytest.raises(FoundryStoragePreconditionError) as exc: + raise_for_storage_error(_FakeResponse(412, {"error": {"message": "etag"}})) + assert exc.value.current_etag is None + + +def test_500_maps_to_api_error_and_tagged_platform() -> None: + with pytest.raises(FoundryStorageApiError) as exc: + raise_for_storage_error(_FakeResponse(500, {"error": {"message": "boom"}})) + assert getattr(exc.value, PLATFORM_ERROR_TAG) is True + + +def test_non_json_body_uses_fallback_message() -> None: + resp = _FakeResponse(503, None) + resp._body = "not json" + with pytest.raises(FoundryStorageApiError) as exc: + raise_for_storage_error(resp) + assert "HTTP 503" in exc.value.message + assert exc.value.response_body is None + + +def test_error_hierarchy_is_catchable_as_base() -> None: + with pytest.raises(FoundryStorageError): + raise_for_storage_error(_FakeResponse(404, {"error": {"message": "x"}})) diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py new file mode 100644 index 000000000000..db466422ca93 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py @@ -0,0 +1,48 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Unit tests for storage pipeline policies (URL masking, UA policy).""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +from azure.ai.agentserver.core.storage._policies import ( + ServerVersionUserAgentPolicy, + _mask_storage_url, +) + + +def test_mask_keeps_only_storage_path_and_api_version() -> None: + url = "https://proj.example.com/api/projects/secret/storage/state_stores/abc/items:keys?api-version=v1&after=it_1" + masked = _mask_storage_url(url) + assert masked == "***/storage/state_stores/*/items:keys?api-version=v1" + assert "secret" not in masked + assert "after=it_1" not in masked + assert "abc" not in masked + + +def test_mask_without_storage_segment_is_fully_redacted() -> None: + assert _mask_storage_url("https://proj.example.com/other/path") == "(redacted)" + + +def test_mask_empty_url_is_redacted() -> None: + assert _mask_storage_url("") == "(redacted)" + + +def test_mask_malformed_url_is_redacted() -> None: + assert _mask_storage_url(None) == "(redacted)" # type: ignore[arg-type] # defensive branch + + +def test_user_agent_policy_evaluates_callback_per_request() -> None: + versions = iter(["ua-1", "ua-2"]) + policy = ServerVersionUserAgentPolicy(lambda: next(versions)) + + req1 = MagicMock() + req1.http_request.headers = {} + policy.on_request(req1) + assert req1.http_request.headers["User-Agent"] == "ua-1" + + req2 = MagicMock() + req2.http_request.headers = {} + policy.on_request(req2) + assert req2.http_request.headers["User-Agent"] == "ua-2"