diff --git a/sdk/agentserver/azure-ai-agentserver-activity/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-activity/CHANGELOG.md index 66d72b72187c..fa40ad019f9f 100644 --- a/sdk/agentserver/azure-ai-agentserver-activity/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-activity/CHANGELOG.md @@ -2,6 +2,18 @@ ## 1.0.0b4 (Unreleased) +### Features Added + +- Added `FoundryStorage`, an M365 Agents SDK `Storage` adapter backed by + `azure.ai.agentserver.core.storage.FoundryStateStore`. Each M365 storage key + (already a scope identifier, e.g. `f"{channel_id}/conversations/{conversation_id}"` + or `f"{channel_id}/users/{user_id}"`) gets its own lazily-created, lazily-cached + `FoundryStateStore` — "store name = scope", not a single shared namespace. + Keys matching the M365 `UserState` shape automatically get + `user_isolation=True` (override via `is_user_scoped=`). Conversation, user, and + proactive-reference state written through the M365 bridge persists across + restarts and scale-out. Requires `azure-ai-agentserver-core>=2.0.0b8`. + ### Other Changes - Internal refactor only; no public API changes. Consolidated the default storage diff --git a/sdk/agentserver/azure-ai-agentserver-activity/README.md b/sdk/agentserver/azure-ai-agentserver-activity/README.md index 564e7bea1786..d0ed16aaa401 100644 --- a/sdk/agentserver/azure-ai-agentserver-activity/README.md +++ b/sdk/agentserver/azure-ai-agentserver-activity/README.md @@ -54,6 +54,16 @@ from azure.ai.agentserver.activity import ActivityAgentServerHost app = ActivityAgentServerHost(storage=MemoryStorage()) ``` +**Foundry durable storage** — drop-in durable state for the M365 bridge, backed by +Foundry-managed Activity state storage instead of the in-memory default: + +```python +from azure.ai.agentserver.activity import ActivityAgentServerHost, FoundryStorage + +storage = FoundryStorage() +app = ActivityAgentServerHost(storage=storage) +``` + **Inject a pre-built `AgentApplication`** — host an M365 `AgentApplication` you built yourself: ```python @@ -108,6 +118,14 @@ app.run() the `03-self-hosted-app` sample. - `ActivityAgentServerHost.adapter` — the channel adapter for the underlying `AgentApplication`. +- `FoundryStorage` — platform-managed durable storage for M365 conversation, user, + and proactive state. Implements the M365 Agents SDK `Storage` interface. Each + M365 storage key is already a scope identifier (for example + `f"{channel_id}/conversations/{conversation_id}"` or + `f"{channel_id}/users/{user_id}"`), so `FoundryStorage` backs each one with its + own lazily-created `azure.ai.agentserver.core.storage.FoundryStateStore` + ("store name = scope" — no shared namespace); user-shaped keys automatically + get `user_isolation=True` (override via `is_user_scoped=`). ## Examples @@ -118,6 +136,9 @@ See the [samples directory](https://github.com/Azure/azure-sdk-for-python/tree/m - `03-self-hosted-app` — build the M365 `AgentApplication` yourself and host it with `from_agent_application`. - `04-custom-handler` — own the request pipeline with `from_request_handler` (the M365 SDK is not initialized). - `05-multi-protocol` — compose the Activity and Invocations protocols on a single server. +- `06-foundry-storage-state` — durable conversation and user state with `FoundryStorage`. +- `07-foundry-storage-proactive` — durable proactive conversation references with `FoundryStorage`. +- `08-foundry-storage-history` — persist the full conversation transcript with `FoundryStorage` (`/history` and `/clear` commands). ## Troubleshooting diff --git a/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/__init__.py b/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/__init__.py index f043c5984021..4047efae86b9 100644 --- a/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/__init__.py +++ b/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/__init__.py @@ -59,7 +59,8 @@ async def handle(request): from ._activity import ActivityAgentServerHost from ._config import get_hosted_agent_env +from ._foundry_storage import FoundryStorage from ._version import VERSION -__all__ = ["ActivityAgentServerHost", "get_hosted_agent_env"] +__all__ = ["ActivityAgentServerHost", "FoundryStorage", "get_hosted_agent_env"] __version__ = VERSION diff --git a/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/_foundry_storage.py b/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/_foundry_storage.py new file mode 100644 index 000000000000..7263c1e016ca --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/_foundry_storage.py @@ -0,0 +1,286 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""M365 Agents SDK storage adapter backed by per-key Foundry state stores.""" + +from __future__ import annotations + +import asyncio +import re +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Callable, TypeVar + +from azure.ai.agentserver.core.storage import FoundryStateStore, FoundryStorageEndpoint, FoundryStorageNotFoundError +from azure.core.credentials_async import AsyncTokenCredential + +if TYPE_CHECKING: + # Only needed to give the TypeVar bound below a concrete type for static + # analysis; never imported at runtime so the package stays importable + # without the optional M365 Agents SDK. + # pylint: disable=import-error,no-name-in-module + from microsoft_agents.hosting.core.storage import StoreItem + +try: + # pylint: disable=import-error,no-name-in-module + from microsoft_agents.hosting.core.storage import AsyncStorageBase +except ImportError: # pragma: no cover - keeps package importable without optional M365 SDK bits. + + class AsyncStorageBase: # type: ignore[no-redef] # pylint: disable=too-few-public-methods + """Fallback base class used only when the M365 Agents SDK is not installed.""" + + +StoreItemT = TypeVar("StoreItemT", bound="StoreItem") + +#: Regex matching the M365 storage keys that identify a single user, so their +#: backing store gets ``user_isolation=True``. Two shapes are recognized: +#: +#: * ``"{channel}/users/{user_id}"`` -- the M365 ``UserState`` key shape (see +#: ``microsoft_agents.hosting.core.state.user_state.UserState.get_storage_key``). +#: * ``"auth:_SignInState:{channel}:{user_id}"`` -- the M365 ``Authorization`` +#: sign-in state key shape. +_USER_SCOPED_KEY_RE = re.compile(r"/users/|^auth:_SignInState:[^:]+:[^:]+$") + +#: Upper bound on the number of per-key ``FoundryStateStore`` clients the adapter +#: keeps in its in-memory LRU cache. Once exceeded, the least-recently-used store +#: is evicted and closed; its server-side state is untouched and gets a fresh +#: client on next access. Chosen large enough that the coldest (evicted) store is +#: never one being actively fanned out to in a single batch read/write/delete. +_MAX_CACHED_STORES = 1024 + + +def _default_is_user_scoped(key: str) -> bool: + """Match the per-user M365 key shapes (``UserState`` and ``Authorization`` sign-in state).""" + return _USER_SCOPED_KEY_RE.search(key) is not None + + +class FoundryStorage(AsyncStorageBase): + """Durable M365 Agents SDK storage adapter for Foundry-hosted Activity agents. + + Backed by :class:`~azure.ai.agentserver.core.storage.FoundryStateStore`, whose + protocol binds every client to one explicit, caller-named store -- "store + name = scope". Each M365 storage key (already a scope identifier, for + example ``f"{channel_id}/conversations/{conversation_id}"`` or + ``f"{channel_id}/users/{user_id}"``) gets its own lazily-created, + lazily-cached :class:`FoundryStateStore`, with the key doubling as both the + store name and the single item key stored in it. This mirrors the M365 SDK + itself, which never batches keys from different scopes in one call (every + ``AgentState`` / ``Authorization`` / ``Proactive`` call operates on exactly + one key). + + Subclasses :class:`~microsoft_agents.hosting.core.storage.AsyncStorageBase`, + which implements the batch :meth:`read` / :meth:`write` / :meth:`delete` + (validation + concurrent fan-out) in terms of the single-item + ``_read_item`` / ``_write_item`` / ``_delete_item`` hooks below. + + :keyword credential: Async token credential shared by every per-key store. + Defaults to ``DefaultAzureCredential`` (requires ``azure-identity``). + :keyword endpoint: Foundry storage endpoint or project endpoint URL override. + :keyword item_ttl_seconds: Store-level TTL applied to every per-key store + created by this adapter. Defaults to ``FoundryStateStore``'s own default + (30 days) when omitted. + :keyword is_user_scoped: Predicate deciding which keys get + ``user_isolation=True`` on their backing store. Defaults to matching the + per-user M365 key shapes ``"{channel}/users/{user_id}"`` (``UserState``) + and ``"auth:_SignInState:{channel}:{user_id}"`` (``Authorization``). + """ + + def __init__( + self, + *, + credential: AsyncTokenCredential | None = None, + endpoint: FoundryStorageEndpoint | str | None = None, + item_ttl_seconds: int | None = None, + is_user_scoped: Callable[[str], bool] = _default_is_user_scoped, + ) -> None: + """Create the adapter. No backing stores are created until first use. + + :keyword credential: Async token credential shared by every per-key + store, or ``None`` to build an owned ``DefaultAzureCredential``. + :paramtype credential: ~azure.core.credentials_async.AsyncTokenCredential | None + :keyword endpoint: Foundry storage endpoint or project endpoint URL + override, or ``None`` to resolve from the environment. + :paramtype endpoint: ~azure.ai.agentserver.core.storage.FoundryStorageEndpoint | str | None + :keyword item_ttl_seconds: Store-level TTL applied to every per-key + store, or ``None`` to use ``FoundryStateStore``'s own default. + :paramtype item_ttl_seconds: int | None + :keyword is_user_scoped: Predicate deciding which keys get + ``user_isolation=True`` on their backing store. + :paramtype is_user_scoped: ~collections.abc.Callable[[str], bool] + :raises ImportError: If no ``credential`` is supplied and + ``azure-identity`` is not installed. + :return: None. + :rtype: None + """ + super().__init__() + self._owns_credential = credential is None + if credential is None: + try: + from azure.identity.aio import DefaultAzureCredential + except ImportError as exc: # pragma: no cover + raise ImportError( + "FoundryStorage requires azure-identity when no credential is supplied. " + "Install azure-identity or pass an async credential." + ) from exc + credential = DefaultAzureCredential() + self._credential = credential + self._endpoint = endpoint + self._item_ttl_seconds = item_ttl_seconds + self._is_user_scoped = is_user_scoped + self._max_cached_stores = _MAX_CACHED_STORES + + # LRU-ordered cache: most-recently-used key at the end, coldest at the front. + self._stores: "OrderedDict[str, FoundryStateStore]" = OrderedDict() + self._ensured_keys: set[str] = set() + self._creation_lock = asyncio.Lock() + + def _store_kwargs(self, key: str) -> dict[str, Any]: + """Build the keyword arguments for the per-key backing store. + + :param key: The M365 storage key the store will be bound to. + :type key: str + :return: Keyword arguments for ``FoundryStateStore(key, **kwargs)`` / + ``FoundryStateStore.get_or_create(key, **kwargs)``. + :rtype: dict[str, ~typing.Any] + """ + kwargs: dict[str, Any] = { + "credential": self._credential, + "endpoint": self._endpoint, + "user_isolation": self._is_user_scoped(key), + } + if self._item_ttl_seconds is not None: + kwargs["item_ttl_seconds"] = self._item_ttl_seconds + return kwargs + + async def _get_store(self, key: str, *, ensure_exists: bool) -> FoundryStateStore: + """Return the cached per-key store, creating it on first use. + + Reads/deletes get a plain (unconfirmed) client -- ``FoundryStateStore`` + gracefully treats a not-yet-created store as a missing key/item, so no + network round trip is needed just to look something up. Writes need + the server-side store to actually exist, so the first write for a key + upgrades the cache entry via :meth:`~FoundryStateStore.get_or_create`. + + :param key: The M365 storage key to resolve a store for. + :type key: str + :keyword ensure_exists: When ``True``, ensures the server-side store + resource exists (once per key) before returning. + :paramtype ensure_exists: bool + :return: The cached (or newly created) per-key store. + :rtype: ~azure.ai.agentserver.core.storage.FoundryStateStore + """ + store = self._stores.get(key) + if store is not None and (not ensure_exists or key in self._ensured_keys): + self._stores.move_to_end(key) # mark most-recently-used + return store + async with self._creation_lock: + if ensure_exists and key not in self._ensured_keys: + store = await FoundryStateStore.get_or_create(key, **self._store_kwargs(key)) + self._stores[key] = store + self._ensured_keys.add(key) + else: + store = self._stores.get(key) + if store is None: + store = FoundryStateStore(key, **self._store_kwargs(key)) + self._stores[key] = store + self._stores.move_to_end(key) # mark most-recently-used + await self._evict_if_needed() + return store + + async def _evict_if_needed(self) -> None: + """Evict and close least-recently-used stores until the cache is within bounds. + + Must be called while holding ``_creation_lock``. Only ever evicts the + coldest entries (front of the LRU order), which -- given ``max_cached_stores`` + exceeds the fan-out width of any single batch -- are never stores currently + being read from or written to. Evicting a key drops it from + ``_ensured_keys`` too; a later access simply recreates its client (and + re-runs ``get_or_create`` on the next write), the server-side state being + untouched. + + :return: None. + :rtype: None + """ + while len(self._stores) > self._max_cached_stores: + evicted_key, evicted_store = self._stores.popitem(last=False) + self._ensured_keys.discard(evicted_key) + await evicted_store.aclose() + + async def aclose(self) -> None: + """Close every cached per-key store and an owned default credential. + + :return: None. + :rtype: None + """ + for store in list(self._stores.values()): + await store.aclose() + self._stores.clear() + self._ensured_keys.clear() + if self._owns_credential and hasattr(self._credential, "close"): + await self._credential.close() + + async def __aenter__(self) -> "FoundryStorage": + """Enter the async context manager. + + :return: This instance. + :rtype: ~azure.ai.agentserver.activity.FoundryStorage + """ + return self + + async def __aexit__(self, *args: Any) -> None: + """Exit the async context manager, closing all cached stores. + + :param args: The exception type, value, and traceback (unused). + :type args: object + :return: None. + :rtype: None + """ + await self.aclose() + + async def _read_item( + self, + key: str, + *, + target_cls: type[StoreItemT] | None = None, + **kwargs: Any, + ) -> tuple[str | None, StoreItemT | None]: + """Fetch one item. A store that does not exist yet is just a missing key. + + :param key: The M365 storage key to read (also the store name). + :type key: str + :keyword target_cls: The ``StoreItem`` subclass to deserialize into. + :paramtype target_cls: type[StoreItemT] | None + :return: ``(key, item)`` if found, otherwise ``(None, None)``. + :rtype: tuple[str | None, StoreItemT | None] + """ + _ = kwargs + store = await self._get_store(key, ensure_exists=False) + item = await store.get(key) + if item is None or target_cls is None: + return None, None + return key, target_cls.from_json_to_store_item(item.value) + + async def _write_item(self, key: str, value: StoreItemT) -> None: + """Create-or-replace one item, creating its backing store on first write. + + :param key: The M365 storage key to write (also the store name). + :type key: str + :param value: The ``StoreItem`` to persist. + :type value: StoreItemT + :return: None. + :rtype: None + """ + store = await self._get_store(key, ensure_exists=True) + await store.set(key, value.store_item_to_json()) + + async def _delete_item(self, key: str) -> None: + """Delete one item. Missing keys (or stores) are ignored. + + :param key: The M365 storage key to delete (also the store name). + :type key: str + :return: None. + :rtype: None + """ + store = await self._get_store(key, ensure_exists=False) + try: + await store.delete(key) + except FoundryStorageNotFoundError: + pass diff --git a/sdk/agentserver/azure-ai-agentserver-activity/dev_requirements.txt b/sdk/agentserver/azure-ai-agentserver-activity/dev_requirements.txt index 25c128efb198..04231583735d 100644 --- a/sdk/agentserver/azure-ai-agentserver-activity/dev_requirements.txt +++ b/sdk/agentserver/azure-ai-agentserver-activity/dev_requirements.txt @@ -1,6 +1,7 @@ # keep in sync with pyproject.toml#dependency-groups.dev -e ../../../eng/tools/azure-sdk-tools -e ../azure-ai-agentserver-core +azure-identity>=1.17.0 pytest httpx pytest-asyncio diff --git a/sdk/agentserver/azure-ai-agentserver-activity/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-activity/pyproject.toml index 37c14c54a891..2f024f3c2780 100644 --- a/sdk/agentserver/azure-ai-agentserver-activity/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-activity/pyproject.toml @@ -21,7 +21,8 @@ classifiers = [ keywords = ["azure", "azure sdk", "agent", "agentserver", "activity"] dependencies = [ - "azure-ai-agentserver-core>=2.0.0b7", + "azure-ai-agentserver-core>=2.0.0b8", + "azure-identity>=1.17.0", "opentelemetry-api>=1.40.0", "aiohttp>=3.10.0,<4.0.0a0", ] diff --git a/sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/main.py b/sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/main.py new file mode 100644 index 000000000000..a4c500f68d4d --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/main.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Foundry Storage State Agent — Activity Protocol with durable state. + +Demonstrates the zero-config decorator pattern with platform-managed durable +storage. Conversation and user counters are persisted by the M365 Agents SDK +through ``FoundryStorage`` after each turn. + +Usage:: + + python foundry_storage_state_agent.py +""" + +import logging +import sys +import traceback + +from azure.ai.agentserver.activity import ActivityAgentServerHost, FoundryStorage + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s | %(message)s", +) + +storage = FoundryStorage() +app = ActivityAgentServerHost(storage=storage) + + +@app.activity("conversationUpdate") +async def on_members_added(context, _state): + """Welcome new members.""" + for member in context.activity.members_added or []: + if member.id != context.activity.recipient.id: + await context.send_activity( + "Hello! I persist conversation and user state with FoundryStorage.\n\n" + "Send any message to increment the durable counters." + ) + + +@app.activity("message") +async def on_message(context, state): + """Increment durable conversation and user counters.""" + conversation_count = state.conversation.get_value("message_count", lambda: 0) + user_count = state.user.get_value("message_count", lambda: 0) + + conversation_count += 1 + user_count += 1 + state.conversation.set_value("message_count", conversation_count) + state.user.set_value("message_count", user_count) + + await context.send_activity( + "FoundryStorage persisted this turn.\n\n" + f"- Conversation messages: **{conversation_count}**\n" + f"- Messages from you: **{user_count}**" + ) + + +@app.error +async def on_error(context, error): + """Handle unhandled errors.""" + print(f"\n [on_turn_error] unhandled error: {error}", file=sys.stderr) + traceback.print_exc() + await context.send_activity("The agent encountered an error or bug.") + + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/requirements.txt b/sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/requirements.txt new file mode 100644 index 000000000000..060932fccb2f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/requirements.txt @@ -0,0 +1,5 @@ +azure-ai-agentserver-activity +microsoft-agents-hosting-core +microsoft-agents-authentication-msal +microsoft-agents-activity +azure-identity diff --git a/sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/main.py b/sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/main.py new file mode 100644 index 000000000000..5cda149e7b44 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/main.py @@ -0,0 +1,100 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Foundry Storage Proactive Agent — durable proactive conversation references. + +Builds the M365 ``AgentApplication`` with ``ProactiveOptions`` (self-hosted-app +pattern) and injects it into the host via ``agent_app=``. Users send +``/subscribe`` to store the current conversation reference in ``FoundryStorage``. +An external caller can then POST ``/notify/{conversation_id}`` (registered via +the host's ``routes=`` override) to resume that stored conversation and send a +proactive notification. + +Usage:: + + python main.py +""" + +import sys +import traceback + +from starlette.responses import JSONResponse, Response +from starlette.routing import Route + +from azure.ai.agentserver.activity import ActivityAgentServerHost, FoundryStorage, get_hosted_agent_env +from microsoft_agents.activity import load_configuration_from_env +from microsoft_agents.authentication.msal import MsalConnectionManager +from microsoft_agents.hosting.core import ( + AgentApplication, + Authorization, + HttpAdapterBase, + RestChannelServiceClientFactory, + TurnContext, + TurnState, +) +from microsoft_agents.hosting.core.app.proactive.proactive_options import ProactiveOptions + +# get_hosted_agent_env() resolves the CONNECTIONS__* settings WITHOUT mutating +# the process environment (see 03-self-hosted-app). +config = load_configuration_from_env(get_hosted_agent_env()) +storage = FoundryStorage() +connection_manager = MsalConnectionManager(**config) +client_factory = RestChannelServiceClientFactory(connection_manager) +adapter = HttpAdapterBase(channel_service_client_factory=client_factory) +authorization = Authorization(storage, connection_manager, **config) +agent_app = AgentApplication[TurnState]( + storage=storage, + adapter=adapter, + authorization=authorization, + proactive=ProactiveOptions(storage=storage), + **config, +) + + +@agent_app.message("/subscribe") +async def on_subscribe(context: TurnContext, _state: TurnState): + """Persist the current conversation reference for future proactive sends.""" + await agent_app.proactive.store_conversation(context) + conversation_id = context.activity.conversation.id + await context.send_activity( + "Stored this conversation in FoundryStorage.\n\n" + f"POST `/notify/{conversation_id}` to send a proactive notification." + ) + + +@agent_app.activity("message") +async def on_message(context: TurnContext, _state: TurnState): + """Default message handler.""" + await context.send_activity("Send **/subscribe** to store this conversation for proactive notifications.") + + +@agent_app.error +async def on_error(context: TurnContext, error: Exception): + """Handle unhandled errors.""" + print(f"\n [on_turn_error] unhandled error: {error}", file=sys.stderr) + traceback.print_exc() + await context.send_activity("The agent encountered an error or bug.") + + +async def notify(request) -> Response: + """Resume a stored conversation reference and send a proactive message.""" + conversation_id = request.path_params["conversation_id"] + + async def send_notification(context: TurnContext, _state: TurnState): + await context.send_activity("Proactive notification sent from a conversation reference in FoundryStorage.") + + try: + await agent_app.proactive.continue_conversation(adapter, conversation_id, send_notification) + except KeyError: + return JSONResponse(status_code=404, content={"error": "Conversation reference not found."}) + return JSONResponse({"sent": True, "conversation_id": conversation_id}) + + +# routes= extends the Foundry activity routes with the custom /notify endpoint; +# agent_app= hosts the pre-built AgentApplication (adapter taken from agent_app.adapter). +app = ActivityAgentServerHost( + agent_app=agent_app, + routes=[Route("/notify/{conversation_id}", notify, methods=["POST"])], +) + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/requirements.txt b/sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/requirements.txt new file mode 100644 index 000000000000..060932fccb2f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/requirements.txt @@ -0,0 +1,5 @@ +azure-ai-agentserver-activity +microsoft-agents-hosting-core +microsoft-agents-authentication-msal +microsoft-agents-activity +azure-identity diff --git a/sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/main.py b/sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/main.py new file mode 100644 index 000000000000..4a57dc270614 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/main.py @@ -0,0 +1,88 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Foundry Storage History Agent — Activity Protocol with durable history. + +The simplest durable-storage sample: it persists the full conversation +transcript with ``FoundryStorage`` so the history survives restarts and +scale-out. Each turn is appended to a list held in conversation state; +the agent echoes the running transcript back. + +Commands: + + /history show the stored transcript + /clear forget the stored transcript + +Usage:: + + python foundry_storage_history_agent.py +""" + +import logging +import sys +import traceback + +from azure.ai.agentserver.activity import ActivityAgentServerHost, FoundryStorage + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s | %(message)s", +) + +# Platform-managed storage — no Cosmos account or connection string to manage. +storage = FoundryStorage() +app = ActivityAgentServerHost(storage=storage) + + +@app.activity("conversationUpdate") +async def on_members_added(context, _state): + """Welcome new members.""" + for member in context.activity.members_added or []: + if member.id != context.activity.recipient.id: + await context.send_activity( + "Hello! I remember our conversation with FoundryStorage.\n\n" + "Send any message and I'll append it to the durable transcript. " + "Use `/history` to see it or `/clear` to forget it." + ) + + +@app.activity("message") +async def on_message(context, state): + """Persist the turn in conversation history and echo the transcript back.""" + user_text = (context.activity.text or "").strip() + if not user_text: + return + + history = state.conversation.get_value("history", lambda: []) + + if user_text == "/clear": + state.conversation.set_value("history", []) + await context.send_activity("Transcript cleared.") + return + + if user_text == "/history": + if not history: + await context.send_activity("No messages stored yet.") + else: + transcript = "\n".join(f"{i}. {line}" for i, line in enumerate(history, 1)) + await context.send_activity(f"**Stored transcript ({len(history)}):**\n\n{transcript}") + return + + history.append(f"You: {user_text}") + state.conversation.set_value("history", history) + + await context.send_activity( + f"Saved. I've persisted **{len(history)}** message(s) this conversation. " + "Send `/history` to see them all." + ) + + +@app.error +async def on_error(context, error): + """Handle unhandled errors.""" + print(f"\n [on_turn_error] unhandled error: {error}", file=sys.stderr) + traceback.print_exc() + await context.send_activity("The agent encountered an error or bug.") + + +if __name__ == "__main__": + app.run() diff --git a/sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/requirements.txt b/sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/requirements.txt new file mode 100644 index 000000000000..060932fccb2f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/requirements.txt @@ -0,0 +1,5 @@ +azure-ai-agentserver-activity +microsoft-agents-hosting-core +microsoft-agents-authentication-msal +microsoft-agents-activity +azure-identity diff --git a/sdk/agentserver/azure-ai-agentserver-activity/tests/test_foundry_storage.py b/sdk/agentserver/azure-ai-agentserver-activity/tests/test_foundry_storage.py new file mode 100644 index 000000000000..9ec84d2be88f --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/tests/test_foundry_storage.py @@ -0,0 +1,202 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Unit tests for the M365 FoundryStorage adapter (backed by FoundryStateStore).""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from azure.ai.agentserver.activity import FoundryStorage +import azure.ai.agentserver.activity._foundry_storage as module +from azure.ai.agentserver.core.storage import FoundryStorageNotFoundError, StateStoreItem + + +class _TestStoreItem: + def __init__(self, value: dict[str, Any]) -> None: + self.value = value + + def store_item_to_json(self) -> dict[str, Any]: + return self.value + + @staticmethod + def from_json_to_store_item(json_data: dict[str, Any]) -> "_TestStoreItem": + return _TestStoreItem(json_data) + + +def _fake_store() -> MagicMock: + store = MagicMock() + store.get = AsyncMock(return_value=None) + store.set = AsyncMock() + store.delete = AsyncMock() + store.aclose = AsyncMock() + return store + + +def _patch_stores(monkeypatch: pytest.MonkeyPatch, stores_by_key: dict[str, MagicMock]) -> MagicMock: + """Patch ``FoundryStateStore`` so both the plain constructor (reads/deletes) + and the ``get_or_create`` classmethod (writes) resolve to the given fakes.""" + + def factory(key: str, **kwargs: Any) -> MagicMock: + return stores_by_key[key] + + mock_cls = MagicMock(side_effect=factory) + mock_cls.get_or_create = AsyncMock(side_effect=lambda key, **kwargs: stores_by_key[key]) + monkeypatch.setattr(module, "FoundryStateStore", mock_cls) + return mock_cls + + +@pytest.mark.asyncio +async def test_read_missing_key_does_not_create_a_store(monkeypatch: pytest.MonkeyPatch) -> None: + store = _fake_store() + mock_cls = _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + result = await storage.read(["k"], target_cls=_TestStoreItem) + + assert result == {} + mock_cls.assert_called_once() # plain client constructed to issue the GET ... + mock_cls.get_or_create.assert_not_awaited() # ... but the store resource is never created for a read + + +@pytest.mark.asyncio +async def test_read_deserializes_existing_item(monkeypatch: pytest.MonkeyPatch) -> None: + store = _fake_store() + store.get = AsyncMock( + return_value=StateStoreItem( + { + "id": "i1", + "object": "state_store.item", + "key": "k", + "value": {"count": 3}, + "etag": "e1", + "created_at": 0, + "updated_at": 0, + } + ) + ) + _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + result = await storage.read(["k"], target_cls=_TestStoreItem) + + assert result["k"].value == {"count": 3} + store.get.assert_awaited_once_with("k") + + +@pytest.mark.asyncio +async def test_read_treats_missing_item_as_none(monkeypatch: pytest.MonkeyPatch) -> None: + """FoundryStateStore.get() already returns None for a missing store/item.""" + store = _fake_store() + store.get = AsyncMock(return_value=None) + _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + result = await storage.read(["k"], target_cls=_TestStoreItem) + + assert result == {} + + +@pytest.mark.asyncio +async def test_write_creates_the_store_then_sets_the_item(monkeypatch: pytest.MonkeyPatch) -> None: + store = _fake_store() + mock_cls = _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + await storage.write({"k": _TestStoreItem({"turn": 4})}) + + mock_cls.get_or_create.assert_awaited_once() + store.set.assert_awaited_once_with("k", {"turn": 4}) + + +@pytest.mark.asyncio +async def test_write_only_ensures_the_store_exists_once(monkeypatch: pytest.MonkeyPatch) -> None: + store = _fake_store() + mock_cls = _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + await storage.write({"k": _TestStoreItem({"turn": 1})}) + await storage.write({"k": _TestStoreItem({"turn": 2})}) + + mock_cls.get_or_create.assert_awaited_once() + assert store.set.await_count == 2 + + +@pytest.mark.asyncio +async def test_delete_forwards_the_key_and_ignores_missing(monkeypatch: pytest.MonkeyPatch) -> None: + store = _fake_store() + store.delete = AsyncMock(side_effect=FoundryStorageNotFoundError("not found")) + _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + await storage.delete(["k"]) # must not raise + + store.delete.assert_awaited_once_with("k") + + +@pytest.mark.asyncio +async def test_user_scoped_keys_get_user_isolation(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def factory(key: str, **kwargs: Any) -> MagicMock: + captured[key] = kwargs + return _fake_store() + + monkeypatch.setattr(module, "FoundryStateStore", MagicMock(side_effect=factory)) + storage = FoundryStorage() + + await storage.read(["teams/conversations/abc"], target_cls=_TestStoreItem) + await storage.read(["teams/users/user-42"], target_cls=_TestStoreItem) + + assert captured["teams/conversations/abc"]["user_isolation"] is False + assert captured["teams/users/user-42"]["user_isolation"] is True + + +@pytest.mark.asyncio +async def test_custom_is_user_scoped_predicate_is_honored(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def factory(key: str, **kwargs: Any) -> MagicMock: + captured[key] = kwargs + return _fake_store() + + monkeypatch.setattr(module, "FoundryStateStore", MagicMock(side_effect=factory)) + storage = FoundryStorage(is_user_scoped=lambda key: key.startswith("private/")) + + await storage.read(["private/whatever"], target_cls=_TestStoreItem) + + assert captured["private/whatever"]["user_isolation"] is True + + +@pytest.mark.asyncio +async def test_validates_like_m365_storage(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_stores(monkeypatch, {}) + storage = FoundryStorage() + + with pytest.raises(ValueError, match="Keys are required"): + await storage.read([], target_cls=_TestStoreItem) + with pytest.raises(ValueError, match="target_cls cannot be None"): + await storage.read(["k"]) + with pytest.raises(ValueError, match="Changes are required"): + await storage.write({}) + with pytest.raises(ValueError, match="Keys are required"): + await storage.delete([]) + + +@pytest.mark.asyncio +async def test_aclose_closes_every_cached_store_but_not_a_supplied_credential(monkeypatch: pytest.MonkeyPatch) -> None: + store_a, store_b = _fake_store(), _fake_store() + _patch_stores(monkeypatch, {"a": store_a, "b": store_b}) + credential = MagicMock() + credential.close = AsyncMock() + storage = FoundryStorage(credential=credential) + + await storage.read(["a"], target_cls=_TestStoreItem) + await storage.read(["b"], target_cls=_TestStoreItem) + await storage.aclose() + + store_a.aclose.assert_awaited_once() + store_b.aclose.assert_awaited_once() + credential.close.assert_not_called() # caller-supplied credential is not owned diff --git a/sdk/agentserver/azure-ai-agentserver-activity/tests/test_storage_samples_e2e.py b/sdk/agentserver/azure-ai-agentserver-activity/tests/test_storage_samples_e2e.py new file mode 100644 index 000000000000..e0dd0674e50c --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/tests/test_storage_samples_e2e.py @@ -0,0 +1,366 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""End-to-end tests for the FoundryStorage samples (06-08). + +Per the repo's sample-testing convention (see +``azure-ai-agentserver-responses/tests/e2e/test_sample_e2e.py``), each sample's +handler logic is replicated here *inline* (not imported from the sample +module) and driven through the real M365 Agents SDK turn pipeline +(``AgentApplication`` + ``HttpAdapterBase.process_activity``) for the full +in-process lifecycle: state load -> handler dispatch -> state save -> outbound +send. ``MemoryStorage`` stands in for ``FoundryStorage`` (both implement the +same M365 ``Storage`` contract; see ``test_foundry_storage.py`` for the adapter +unit tests), and a fake ``ChannelServiceClientFactory`` captures outbound +activities in place of a real Bot Connector call, so these tests exercise the +samples' actual behavior end-to-end without any network access. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from microsoft_agents.activity import Activity, ChannelAccount, ConversationAccount, ResourceResponse +from microsoft_agents.hosting.core import AgentApplication, ClaimsIdentity, HttpAdapterBase, MemoryStorage, TurnState +from microsoft_agents.hosting.core.app.proactive.proactive_options import ProactiveOptions + +# --------------------------------------------------------------------------- +# Shared network-free turn-driving harness +# --------------------------------------------------------------------------- + + +class _FakeConversations: + """Captures outbound sends instead of calling the real Bot Connector API.""" + + def __init__(self) -> None: + self.sent: list[Activity] = [] + + async def reply_to_activity(self, conversation_id: str, activity_id: str, activity: Activity) -> ResourceResponse: + _ = conversation_id, activity_id + self.sent.append(activity) + return ResourceResponse(id="r1") + + async def send_to_conversation(self, conversation_id: str, activity: Activity) -> ResourceResponse: + _ = conversation_id + self.sent.append(activity) + return ResourceResponse(id="r1") + + +class _FakeConnectorClient: + def __init__(self) -> None: + self.conversations = _FakeConversations() + + async def close(self) -> None: + pass + + +class _FakeUserTokenClient: + async def close(self) -> None: + pass + + +class _FakeChannelServiceClientFactory: + """Network-free stand-in for ``RestChannelServiceClientFactory``.""" + + def __init__(self) -> None: + self.client = _FakeConnectorClient() + + async def create_connector_client( + self, + context: Any, + claims_identity: Any, + service_url: str, + audience: str, + scopes: list[str] | None = None, + use_anonymous: bool = False, + ) -> _FakeConnectorClient: + _ = context, claims_identity, service_url, audience, scopes, use_anonymous + return self.client + + async def create_user_token_client( + self, context: Any, claims_identity: Any, use_anonymous: bool = False + ) -> _FakeUserTokenClient: + _ = context, claims_identity, use_anonymous + return _FakeUserTokenClient() + + +def _test_adapter() -> tuple[HttpAdapterBase, _FakeChannelServiceClientFactory]: + factory = _FakeChannelServiceClientFactory() + return HttpAdapterBase(channel_service_client_factory=factory), factory + + +def _anonymous_claims() -> ClaimsIdentity: + return ClaimsIdentity({}, is_authenticated=False, authentication_type="anonymous") + + +def _message(text: str, *, conversation_id: str = "c1", user_id: str = "u1") -> Activity: + return Activity( + type="message", + text=text, + channel_id="test", + conversation=ConversationAccount(id=conversation_id), + from_property=ChannelAccount(id=user_id), + recipient=ChannelAccount(id="bot"), + service_url="https://test.example", + ) + + +def _conversation_update(*, conversation_id: str = "c1", member_id: str = "u1") -> Activity: + return Activity( + type="conversationUpdate", + channel_id="test", + conversation=ConversationAccount(id=conversation_id), + from_property=ChannelAccount(id=member_id), + recipient=ChannelAccount(id="bot"), + service_url="https://test.example", + members_added=[ChannelAccount(id=member_id)], + ) + + +# --------------------------------------------------------------------------- +# Sample 06: durable conversation/user state (samples/06-foundry-storage-state) +# --------------------------------------------------------------------------- + + +def _build_state_sample_app(storage: MemoryStorage, adapter: HttpAdapterBase) -> AgentApplication: + """Replicates samples/06-foundry-storage-state/main.py's handlers.""" + app = AgentApplication[TurnState](storage=storage, adapter=adapter) + + @app.activity("conversationUpdate") + async def on_members_added(context, _state): + for member in context.activity.members_added or []: + if member.id != context.activity.recipient.id: + await context.send_activity( + "Hello! I persist conversation and user state with FoundryStorage.\n\n" + "Send any message to increment the durable counters." + ) + + @app.activity("message") + async def on_message(context, state): + conversation_count = state.conversation.get_value("message_count", lambda: 0) + user_count = state.user.get_value("message_count", lambda: 0) + conversation_count += 1 + user_count += 1 + state.conversation.set_value("message_count", conversation_count) + state.user.set_value("message_count", user_count) + await context.send_activity( + "FoundryStorage persisted this turn.\n\n" + f"- Conversation messages: **{conversation_count}**\n" + f"- Messages from you: **{user_count}**" + ) + + return app + + +@pytest.mark.asyncio +async def test_sample06_welcomes_new_members() -> None: + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_state_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _conversation_update(), app.on_turn) + + assert len(factory.client.conversations.sent) == 1 + assert "persist conversation and user state" in factory.client.conversations.sent[0].text + + +@pytest.mark.asyncio +async def test_sample06_counters_persist_across_turns() -> None: + """Conversation and user counters survive across separate turns via storage.""" + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_state_sample_app(storage, adapter) + + for _ in range(3): + await adapter.process_activity(_anonymous_claims(), _message("hi"), app.on_turn) + + replies = [a.text for a in factory.client.conversations.sent] + assert "Conversation messages: **1**" in replies[0] + assert "Conversation messages: **2**" in replies[1] + assert "Conversation messages: **3**" in replies[2] + assert "Messages from you: **3**" in replies[2] + + +@pytest.mark.asyncio +async def test_sample06_user_counter_is_scoped_per_user() -> None: + """The per-user counter is independent of the conversation-wide counter.""" + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_state_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _message("hi", user_id="alice"), app.on_turn) + await adapter.process_activity(_anonymous_claims(), _message("hi", user_id="bob"), app.on_turn) + + replies = [a.text for a in factory.client.conversations.sent] + # Conversation counter keeps incrementing across both users ... + assert "Conversation messages: **1**" in replies[0] + assert "Conversation messages: **2**" in replies[1] + # ... but each user's own counter starts fresh. + assert "Messages from you: **1**" in replies[0] + assert "Messages from you: **1**" in replies[1] + + +# --------------------------------------------------------------------------- +# Sample 08: durable transcript history (samples/08-foundry-storage-history) +# --------------------------------------------------------------------------- + + +def _build_history_sample_app(storage: MemoryStorage, adapter: HttpAdapterBase) -> AgentApplication: + """Replicates samples/08-foundry-storage-history/main.py's handlers.""" + app = AgentApplication[TurnState](storage=storage, adapter=adapter) + + @app.activity("message") + async def on_message(context, state): + user_text = (context.activity.text or "").strip() + if not user_text: + return + + history = state.conversation.get_value("history", lambda: []) + + if user_text == "/clear": + state.conversation.set_value("history", []) + await context.send_activity("Transcript cleared.") + return + + if user_text == "/history": + if not history: + await context.send_activity("No messages stored yet.") + else: + transcript = "\n".join(f"{i}. {line}" for i, line in enumerate(history, 1)) + await context.send_activity(f"**Stored transcript ({len(history)}):**\n\n{transcript}") + return + + history.append(f"You: {user_text}") + state.conversation.set_value("history", history) + await context.send_activity( + f"Saved. I've persisted **{len(history)}** message(s) this conversation. " + "Send `/history` to see them all." + ) + + return app + + +@pytest.mark.asyncio +async def test_sample08_history_command_reports_no_messages_initially() -> None: + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_history_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _message("/history"), app.on_turn) + + assert factory.client.conversations.sent[0].text == "No messages stored yet." + + +@pytest.mark.asyncio +async def test_sample08_transcript_persists_and_history_command_replays_it() -> None: + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_history_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _message("hello"), app.on_turn) + await adapter.process_activity(_anonymous_claims(), _message("world"), app.on_turn) + await adapter.process_activity(_anonymous_claims(), _message("/history"), app.on_turn) + + transcript_reply = factory.client.conversations.sent[-1].text + assert "**Stored transcript (2):**" in transcript_reply + assert "1. You: hello" in transcript_reply + assert "2. You: world" in transcript_reply + + +@pytest.mark.asyncio +async def test_sample08_clear_command_erases_the_transcript() -> None: + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_history_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _message("hello"), app.on_turn) + await adapter.process_activity(_anonymous_claims(), _message("/clear"), app.on_turn) + await adapter.process_activity(_anonymous_claims(), _message("/history"), app.on_turn) + + assert factory.client.conversations.sent[1].text == "Transcript cleared." + assert factory.client.conversations.sent[2].text == "No messages stored yet." + + +# --------------------------------------------------------------------------- +# Sample 07: durable proactive conversation references +# (samples/07-foundry-storage-proactive) +# --------------------------------------------------------------------------- + + +def _build_proactive_sample_app(storage: MemoryStorage, adapter: HttpAdapterBase) -> AgentApplication: + """Replicates samples/07-foundry-storage-proactive/main.py's handlers.""" + app = AgentApplication[TurnState](storage=storage, adapter=adapter, proactive=ProactiveOptions(storage=storage)) + + @app.message("/subscribe") + async def on_subscribe(context, _state): + await app.proactive.store_conversation(context) + conversation_id = context.activity.conversation.id + await context.send_activity( + "Stored this conversation in FoundryStorage.\n\n" + f"POST `/notify/{conversation_id}` to send a proactive notification." + ) + + @app.activity("message") + async def on_message(context, _state): + await context.send_activity("Send **/subscribe** to store this conversation for proactive notifications.") + + return app + + +async def _notify(app: AgentApplication, adapter: HttpAdapterBase, conversation_id: str) -> None: + """Replicates samples/07-foundry-storage-proactive/main.py's notify() route.""" + + async def send_notification(context, _state): + await context.send_activity("Proactive notification sent from a conversation reference in FoundryStorage.") + + await app.proactive.continue_conversation(adapter, conversation_id, send_notification) + + +@pytest.mark.asyncio +async def test_sample07_subscribe_stores_the_conversation_reference() -> None: + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_proactive_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _message("/subscribe"), app.on_turn) + + assert "Stored this conversation in FoundryStorage" in factory.client.conversations.sent[0].text + assert await app.proactive.get_conversation("c1") is not None + + +@pytest.mark.asyncio +async def test_sample07_notify_resumes_the_stored_conversation() -> None: + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_proactive_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _message("/subscribe"), app.on_turn) + await _notify(app, adapter, "c1") + + assert factory.client.conversations.sent[-1].text == ( + "Proactive notification sent from a conversation reference in FoundryStorage." + ) + + +@pytest.mark.asyncio +async def test_sample07_notify_unknown_conversation_raises_key_error() -> None: + """Mirrors the sample's notify() route, which maps this to a 404 response.""" + storage = MemoryStorage() + adapter, _factory = _test_adapter() + app = _build_proactive_sample_app(storage, adapter) + + with pytest.raises(KeyError): + await _notify(app, adapter, "never-subscribed") + + +@pytest.mark.asyncio +async def test_sample07_default_message_prompts_to_subscribe() -> None: + storage = MemoryStorage() + adapter, factory = _test_adapter() + app = _build_proactive_sample_app(storage, adapter) + + await adapter.process_activity(_anonymous_claims(), _message("hi"), app.on_turn) + + assert factory.client.conversations.sent[0].text == ( + "Send **/subscribe** to store this conversation for proactive notifications." + ) diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index ba3a1ab6b47b..700044b9e9ad 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -1,5 +1,11 @@ # Release History +## 2.0.0b8 (Unreleased) + +### Features Added + +- Added `azure.ai.agentserver.core.storage`: `FoundryStorageClient` (transport, endpoint, pipeline policies, error hierarchy) and `FoundryStateStore`, a developer-facing client bound to one explicit, caller-named store over the `/storage/state_stores/*` protocol. Session/conversation/user scoping is expressed in the store name itself (for example `checkpoints/`) rather than a shared namespace. `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. `get(key=None)` and `delete(key=None, ...)` are overloaded on whether `key` is supplied: with no `key` they act on the bound store itself (its descriptor, or the whole store cascade-deleted); with a `key` they act on one item. `update(...)` changes the store's mutable `description` / `tags`. Item operations (`create_item`, `set`, `list_keys`) round out the single-item surface. Store-level `item_ttl_seconds` is fixed at create time (default 30 days, `-1` never expires); optional `user_isolation` partitions items per end user; the caller identity is resolved server-side from the opaque per-request call ID (`x-agent-foundry-call-id`), which the SDK forwards automatically on every storage request from `azure.ai.agentserver.core.get_request_context()` (nothing to configure on the client). `if_match` gives single-item optimistic concurrency; `list_keys` is ordered and paged. Response bodies (`StateStore`, `StateStoreItem`, `StateStoreItemMetadata`, `DeletedStateStore`, `DeletedStateStoreItem`, `StateStoreKey`) are typed model classes generated from a formal TypeSpec contract (`type_spec/main.tsp`), not hand-written dataclasses -- see `type_spec/README.md` for the local generation workflow (this contract does not live under `Azure/azure-rest-api-specs` yet). Write methods take plain scalar keywords (`user_isolation`, `item_ttl_seconds`, `description`, `tags`) rather than typed request-body objects. 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) and `samples/state_store_sample.py`. + ## 2.0.0b7 (2026-06-28) ### Features Added diff --git a/sdk/agentserver/azure-ai-agentserver-core/Makefile b/sdk/agentserver/azure-ai-agentserver-core/Makefile new file mode 100644 index 000000000000..c8ebd8ede7f3 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/Makefile @@ -0,0 +1,83 @@ +# TypeSpec Code Generation Tooling for storage request/response models +# Targets: generate-models, clean, install-typespec-deps +# +# See type_spec/README.md for why this compiles main.tsp locally instead of +# `tsp-client sync`-ing from Azure/azure-rest-api-specs (no upstream +# state-stores directory exists there yet). + +OUTPUT_DIR ?= azure/ai/agentserver/core/storage/_generated +TYPESPEC_DIR ?= type_spec +TEMP_OUTPUT_DIR := $(TYPESPEC_DIR)/.tmp_codegen + +.PHONY: generate-models clean install-typespec-deps + +ifeq ($(OS),Windows_NT) +SHELL := cmd +.SHELLFLAGS := /c +endif + +# -------------------------------------------------------------------------- +# install-typespec-deps: Install the pinned TypeSpec toolchain (one-time) +# -------------------------------------------------------------------------- +ifeq ($(OS),Windows_NT) +install-typespec-deps: + @where npm >NUL 2>NUL || (echo Error: npm is required. Install Node.js ^(v18+^) from https://nodejs.org/ 1>&2 && exit /b 1) + cd /d $(TYPESPEC_DIR) && npm install --silent +else +install-typespec-deps: + @command -v npm >/dev/null 2>&1 || { \ + echo "Error: npm is required. Install Node.js (v18+) from https://nodejs.org/" >&2; \ + exit 1; \ + } + cd $(TYPESPEC_DIR) && npm install --silent +endif + +# -------------------------------------------------------------------------- +# generate-models: Compile type_spec/main.tsp into Python model classes +# -------------------------------------------------------------------------- +ifeq ($(OS),Windows_NT) +generate-models: + @where npm >NUL 2>NUL || (echo Error: npm is required. Install Node.js ^(v18+^) from https://nodejs.org/ 1>&2 && exit /b 1) + @if not exist "$(TYPESPEC_DIR)\node_modules" (echo Error: TypeSpec toolchain not installed. Run 'make install-typespec-deps' first. 1>&2 && exit /b 1) + @echo Compiling TypeSpec into Python models... + cd /d $(TYPESPEC_DIR) && npx tsp compile main.tsp --emit @azure-tools/typespec-python --option "@azure-tools/typespec-python.emitter-output-dir=$(abspath $(TEMP_OUTPUT_DIR))" + @if not exist "$(OUTPUT_DIR)" mkdir "$(OUTPUT_DIR)" + @if not exist "$(OUTPUT_DIR)\_utils" mkdir "$(OUTPUT_DIR)\_utils" + @for /f "delims=" %%d in ('dir /b /s /ad "$(TEMP_OUTPUT_DIR)\models"') do @copy /Y "%%d\_models.py" "$(OUTPUT_DIR)\_models.py" >NUL & copy /Y "%%d\_enums.py" "$(OUTPUT_DIR)\_enums.py" >NUL + @for /f "delims=" %%d in ('dir /b /s /ad "$(TEMP_OUTPUT_DIR)\_utils"') do @copy /Y "%%d\model_base.py" "$(OUTPUT_DIR)\_utils\model_base.py" >NUL + @powershell -NoProfile -Command "(Get-Content '$(OUTPUT_DIR)\_models.py') -replace 'from \.\.\._utils\.model_base', 'from ._utils.model_base' | Set-Content '$(OUTPUT_DIR)\_models.py'" + @echo Generated models at $(OUTPUT_DIR) + @if exist "$(TEMP_OUTPUT_DIR)" rmdir /s /q "$(TEMP_OUTPUT_DIR)" +else +generate-models: + @command -v npm >/dev/null 2>&1 || { \ + echo "Error: npm is required. Install Node.js (v18+) from https://nodejs.org/" >&2; \ + exit 1; \ + } + @test -d "$(TYPESPEC_DIR)/node_modules" || { \ + echo "Error: TypeSpec toolchain not installed. Run 'make install-typespec-deps' first." >&2; \ + exit 1; \ + } + @echo "Compiling TypeSpec into Python models..." + cd $(TYPESPEC_DIR) && npx tsp compile main.tsp --emit @azure-tools/typespec-python --option "@azure-tools/typespec-python.emitter-output-dir=$(abspath $(TEMP_OUTPUT_DIR))" + mkdir -p $(OUTPUT_DIR)/_utils + find $(TEMP_OUTPUT_DIR) -path "*/models/_models.py" -exec cp {} $(OUTPUT_DIR)/_models.py \; + find $(TEMP_OUTPUT_DIR) -path "*/models/_enums.py" -exec cp {} $(OUTPUT_DIR)/_enums.py \; + find $(TEMP_OUTPUT_DIR) -path "*/_utils/model_base.py" -exec cp {} $(OUTPUT_DIR)/_utils/model_base.py \; + sed -i 's/from \.\.\._utils\.model_base/from ._utils.model_base/' $(OUTPUT_DIR)/_models.py + @echo "Generated models at $(OUTPUT_DIR)" + rm -rf $(TEMP_OUTPUT_DIR) +endif + +# -------------------------------------------------------------------------- +# clean: Remove all previously generated Python model files +# -------------------------------------------------------------------------- +ifeq ($(OS),Windows_NT) +clean: + @if exist "$(OUTPUT_DIR)\_models.py" del /q "$(OUTPUT_DIR)\_models.py" + @if exist "$(OUTPUT_DIR)\_enums.py" del /q "$(OUTPUT_DIR)\_enums.py" + @if exist "$(OUTPUT_DIR)\_utils\model_base.py" del /q "$(OUTPUT_DIR)\_utils\model_base.py" +else +clean: + rm -f $(OUTPUT_DIR)/_models.py $(OUTPUT_DIR)/_enums.py $(OUTPUT_DIR)/_utils/model_base.py +endif diff --git a/sdk/agentserver/azure-ai-agentserver-core/README.md b/sdk/agentserver/azure-ai-agentserver-core/README.md index 2a9d56c42346..edfb7c62368b 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/README.md +++ b/sdk/agentserver/azure-ai-agentserver-core/README.md @@ -126,6 +126,35 @@ 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. +# The store name is the scope -- encode conversation/thread identity into it. +# 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("step-1", {"done": False}) + item = await store.get("step-1") + print(item.value) # {"done": False} +``` + +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/azure/ai/agentserver/core/_version.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_version.py index 369f0dcc3bea..213b705c79c6 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_version.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_version.py @@ -2,4 +2,4 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- -VERSION = "2.0.0b7" +VERSION = "2.0.0b8" 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..c023bcda3af9 --- /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, + KeyPage, + Order, + StateStore, + StateStoreItem, + StateStoreItemMetadata, + StateStoreKey, +) + +__all__ = [ + "DEFAULT_ITEM_TTL_SECONDS", + "DeletedStateStore", + "DeletedStateStoreItem", + "FOUNDRY_TOKEN_SCOPE", + "FoundryStateStore", + "FoundryStorageApiError", + "FoundryStorageBadRequestError", + "FoundryStorageConflictError", + "FoundryStorageClient", + "FoundryStorageEndpoint", + "FoundryStorageError", + "FoundryStorageNotFoundError", + "FoundryStoragePreconditionError", + "JSONObject", + "JSONValue", + "KeyPage", + "Order", + "StateStore", + "StateStoreItem", + "StateStoreItemMetadata", + "StateStoreKey", +] 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..74c3e17e0bd8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py @@ -0,0 +1,112 @@ +# 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, PlatformCallIdPolicy, 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(), + PlatformCallIdPolicy(), + 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..c46b81cc418e --- /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"&{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..2f7c5400e348 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Request/response model classes generated from ``type_spec/main.tsp``. + +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, + ListResponseStateStoreKey, + PutItemRequest, + StateStore, + StateStoreItem, + StateStoreItemMetadata, + StateStoreKey, + UpdateStateStoreRequest, +) + +__all__ = [ + "ApiError", + "ApiErrorResponse", + "CreateItemRequest", + "CreateStateStoreRequest", + "DeletedStateStore", + "DeletedStateStoreItem", + "ListResponseStateStore", + "ListResponseStateStoreKey", + "PutItemRequest", + "StateStore", + "StateStoreItem", + "StateStoreItemMetadata", + "StateStoreKey", + "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..a4afe40debfb --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_generated/_models.py @@ -0,0 +1,681 @@ +# 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=useless-super-delegation + +from typing import Any, Literal, Mapping, Optional, overload + +from ._utils.model_base import Model as _Model, rest_field +from ._enums import StateStoreItemObjectType, StateStoreObjectType + + +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: "ApiError" = rest_field(visibility=["read", "create", "update", "delete", "query"]) + """Required.""" + + @overload + def __init__( + self, + *, + error: "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["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["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 ListResponseStateStoreKey(_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.StateStoreKey] + :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["StateStoreKey"] = 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["StateStoreKey"], + 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 StateStoreItemMetadata(_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 StateStoreKey(_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 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..667d26fbfdb7 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_policies.py @@ -0,0 +1,166 @@ +# 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.ai.agentserver.core._request_context import get_request_context +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() + + +class PlatformCallIdPolicy(SansIOHTTPPolicy[HttpRequest, HttpResponse]): + """Forward the per-request Foundry call ID on every outbound storage request. + + Container protocol version ``2.0.0`` mints an opaque per-request call ID + (``x-agent-foundry-call-id``) that the container **MUST** forward on all + outbound calls to Foundry 1P services (Storage, Toolboxes/MCP proxy, A2A) so + the service resolves the caller context -- including the end-user identity + used for per-user isolation -- server-side. This policy stamps that header on + every request from the current + :func:`~azure.ai.agentserver.core.get_request_context`. It is a no-op under + protocol version ``1.0.0`` or local development, where no call ID is bound. + """ + + def on_request(self, request: PipelineRequest[HttpRequest]) -> None: + """Merge the platform call-ID header before the request is sent.""" + request.http_request.headers.update(get_request_context().platform_headers()) + + +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..e1e579e17bba --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state.py @@ -0,0 +1,398 @@ +# 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, overload + +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.rest import HttpRequest + +from ._client import JSON_CONTENT_TYPE, FoundryStorageClient +from ._endpoint import FoundryStorageEndpoint +from ._errors import FoundryStorageConflictError, FoundryStorageNotFoundError +from ._state_serializer import ( + DeletedStateStore, + DeletedStateStoreItem, + JSONObject, + KeyPage, + Order, + StateStore, + StateStoreItem, + StateStoreItemMetadata, + deserialize_deleted_state_item, + deserialize_deleted_state_store, + deserialize_list_keys_response, + deserialize_state_item, + deserialize_state_item_metadata, + 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 +_UNSET = object() + + +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/``). + + 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("step-1", {"done": False}) + + ``get`` and ``delete`` are overloaded on whether *key* is supplied: with no + *key* they act on the store itself (its descriptor, or the whole store + cascade-deleted); with a *key* they act on one item 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, + 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 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) + super().__init__(credential, resolved, **kwargs) + + @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, + 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 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, + api_version=api_version, + **kwargs, + ) + try: + await store._fetch_properties() + except FoundryStorageNotFoundError: + try: + await store._create_properties() + except FoundryStorageConflictError: + await store._fetch_properties() + return store + + async def aclose(self) -> None: + """Close the pipeline client and any owned default credential.""" + await super().aclose() + 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, + if_match: str | None = None, + **query: str, + ) -> HttpRequest: + headers: dict[str, str] = {} + if content is not None: + headers["Content-Type"] = JSON_CONTENT_TYPE + if if_match is not None: + headers["If-Match"] = if_match + return HttpRequest(method, self._endpoint.build_url(path, **query), content=content, headers=headers) + + 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, + ) -> StateStoreItemMetadata: + """Create a new item and fail on duplicate keys.""" + body = serialize_item_create_request(key, value, tags) + response = await self._send_storage_request(self._request("POST", f"{self._store_path()}/items", content=body)) + return deserialize_state_item_metadata(response.text()) + + async def set( + self, + key: str, + value: JSONObject, + *, + tags: Mapping[str, str] | None = None, + if_match: str | None = None, + require_exists: bool = False, + ) -> StateStoreItemMetadata: + """Create or replace one item by key.""" + 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, + if_match=header, + ) + ) + return deserialize_state_item_metadata(response.text()) + + @overload + async def get(self, key: None = None) -> StateStore | None: ... + @overload + async def get(self, key: str) -> StateStoreItem | None: ... + async def get(self, key: str | None = None) -> StateStoreItem | StateStore | None: + """Fetch the bound store's own descriptor, or one item by key. + + :param key: The item key to fetch, or ``None`` (the default) to fetch + the bound store's own descriptor instead. + :type key: str or None + :return: The store descriptor (``key=None``) or the item (``key=``), + or ``None`` if it does not exist. + :rtype: ~azure.ai.agentserver.core.storage.StateStore or + ~azure.ai.agentserver.core.storage.StateStoreItem or None + """ + if key is None: + try: + return await self._fetch_properties() + except FoundryStorageNotFoundError: + return None + try: + response = await self._send_storage_request(self._request("GET", self._item_path(key))) + except FoundryStorageNotFoundError: + return None + return deserialize_state_item(response.text()) + + @overload + async def delete(self, key: None = None, *, if_match: str | None = None) -> DeletedStateStore: ... + @overload + async def delete(self, key: str, *, if_match: str | None = None) -> DeletedStateStoreItem: ... + async def delete( + self, key: str | None = None, *, if_match: str | None = None + ) -> DeletedStateStoreItem | DeletedStateStore: + """Delete the bound store (cascades to every item), or one item by key. + + :param key: The item key to delete, or ``None`` (the default) to + delete the bound store itself, cascading to every item under it. + :type key: str or None + :keyword if_match: Optional concurrency token. Only meaningful for a + single-item delete (``key`` supplied); ignored for a store delete. + :paramtype if_match: str or None + :return: The deleted-store marker (``key=None``) or the deleted-item + marker (``key=``). + :rtype: ~azure.ai.agentserver.core.storage.DeletedStateStore or + ~azure.ai.agentserver.core.storage.DeletedStateStoreItem + """ + if key is None: + response = await self._send_storage_request(self._request("DELETE", self._store_path())) + return deserialize_deleted_state_store(response.text()) + response = await self._send_storage_request(self._request("DELETE", self._item_path(key), 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", + ) -> KeyPage: + """List keys within the bound store.""" + 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", **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..2611e39d77cf --- /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 +``type_spec/main.tsp`` (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``, ``KeyPage``) 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, + ListResponseStateStoreKey, + PutItemRequest, + StateStore, + StateStoreItem, + StateStoreItemMetadata, + StateStoreKey, + 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", + "StateStoreItemMetadata", + "DeletedStateStoreItem", + "StateStoreKey", + "KeyPage", + "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_metadata", + "deserialize_state_item", + "deserialize_deleted_state_item", + "deserialize_list_keys_response", +] + + +@dataclass +class KeyPage: + """A page of keys returned by ``list_keys``. + + A hand-written convenience wrapper (not itself part of the wire contract) + over the generated ``ListResponseStateStoreKey`` envelope, exposing its + ``data`` array as ``keys`` to match this client's naming. + """ + + keys: list[StateStoreKey] + 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_metadata(body: str) -> StateStoreItemMetadata: + return StateStoreItemMetadata(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) -> KeyPage: + page = ListResponseStateStoreKey(load_json(body)) + return KeyPage( + 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..0d30bf8788c0 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md @@ -0,0 +1,339 @@ +# 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()`, `get()` (no `key`), `update()` | `StateStore` | +| `delete()` (no `key`) | `DeletedStateStore` | +| `create_item()`, `set()` | `StateStoreItemMetadata` | +| `get(key)` | `StateStoreItem` | +| `delete(key)` | `DeletedStateStoreItem` | +| `list_keys()` | `KeyPage` (of `StateStoreKey`) | + +Item values (`StateStoreItem.value`) are your own application JSON. Pass a +plain `dict` when writing and read one back on `get()`; the SDK stores and +returns the value as-is. + +```python +from azure.ai.agentserver.core.storage import StateStore, StateStoreItem + +store_info: StateStore | None = await store.get() +item: StateStoreItem | None = await store.get("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("step-1", {"done": False}) + + item: StateStoreItem | None = await store.get("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) +``` + +`get()` and `delete()` are overloaded on whether you pass a `key`: with no +`key` they act on the bound store itself; with a `key` they act on one item. + +```python +info: StateStore | None = await store.get() # the store's metadata, or None 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()` with no `key` cascades to every item under that store name. + +## User Isolation + +Set `user_isolation=True` when the same store name should fan out per end user. + +```python +store = await FoundryStateStore.get_or_create("user-prefs/defaults", user_isolation=True) +``` + +- Item operations on a user-isolated store are automatically scoped to the + current end user. You never pass or configure a user identity, and a single + client instance can safely serve requests for different users over its + lifetime. +- For hosted agents, the platform mints an opaque per-request call ID that the + SDK forwards on every storage call; the service derives the end user from it. + There is nothing to wire up. +- Store-management calls (`get_or_create`, `get()` with no key, `update()`, + `delete()` with no key) stay store-scoped and are shared across every user of + the store. + +## 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( + "step-1", + {"done": True}, + tags={"kind": "checkpoint"}, +) +print(updated.etag) +``` + +`set()` creates the item, or replaces it if the key already exists. + +### Fetch one item + +```python +item: StateStoreItem | None = await store.get("step-1") +if item is not None: + print(item.id, item.key, item.value, item.tags, item.etag) +``` + +`get(key)` returns `None` when the item is missing; `get()` with no `key` +returns the store's metadata instead (or `None` if the store is absent). + +### Delete one item + +```python +deleted = await store.delete("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("counter") +assert item is not None + +try: + await store.set("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("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 KeyPage + +page: KeyPage = 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("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`):** + +| 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()` (replace) | +| `tags` | <= 16 entries, same shape as store tags. | Mutable via `set()` (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 4d43985d6cbc..c9bbda534011 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml @@ -21,12 +21,13 @@ classifiers = [ keywords = ["azure", "azure sdk", "agent", "agentserver", "core"] dependencies = [ + "azure-core>=1.30.0", + "isodate>=0.6.1", "starlette>=0.45.0", "hypercorn>=0.17.0", "opentelemetry-api>=1.40.0", "opentelemetry-sdk>=1.40.0", "microsoft-opentelemetry>=1.0.0", - "aiohttp>=3.10.0,<4.0.0a0", ] [build-system] 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..f79443b9a827 --- /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("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("step-1") + assert item is not None + + await store.set("step-1", {"done": True, "attempt": 2}, tags={"kind": "checkpoint"}) + + try: + await store.set( + "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("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..cf0cfff0eb0c --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py @@ -0,0 +1,583 @@ +# 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 import FoundryAgentRequestContext, reset_request_context, set_request_context +from azure.ai.agentserver.core.storage import ( + DeletedStateStore, + DeletedStateStoreItem, + FoundryStateStore, + FoundryStorageConflictError, + FoundryStorageEndpoint, + FoundryStorageNotFoundError, + KeyPage, + StateStore, + StateStoreItem, + StateStoreItemMetadata, + StateStoreKey, +) +from azure.ai.agentserver.core.storage._policies import PlatformCallIdPolicy +from azure.core.pipeline import PipelineContext, PipelineRequest +from azure.core.rest import HttpRequest + +_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, +) -> 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) + 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) -> StateStoreItemMetadata: + return StateStoreItemMetadata(_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) -> StateStoreKey: + return StateStoreKey(_key_body(**overrides)) + + +# --------------------------------------------------------------------------- +# PlatformCallIdPolicy -- forwards x-agent-foundry-call-id on every request +# --------------------------------------------------------------------------- + + +def _pipeline_request(method: str = "GET") -> PipelineRequest: + http_request = HttpRequest(method, f"{_BASE_URL}state_stores") + return PipelineRequest(http_request, PipelineContext(None)) + + +def test_platform_call_id_policy_forwards_call_id_when_bound() -> None: + policy = PlatformCallIdPolicy() + request = _pipeline_request() + token = set_request_context(FoundryAgentRequestContext(call_id="cid", user_id="uid")) + try: + policy.on_request(request) + finally: + reset_request_context(token) + + assert request.http_request.headers["x-agent-foundry-call-id"] == "cid" + # user_id is the sole hosted-agent derivation via the opaque call ID; it is + # never forwarded to storage as its own header. + assert "x-ms-user-id" not in request.http_request.headers + + +def test_platform_call_id_policy_is_noop_without_context() -> None: + policy = PlatformCallIdPolicy() + request = _pipeline_request() + + policy.on_request(request) + + assert "x-agent-foundry-call-id" not in request.http_request.headers + + +# --------------------------------------------------------------------------- +# 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_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 result is not None + assert result.name == store_name + assert result.id == "ss_1" + + +@pytest.mark.asyncio +async def test_get_with_no_key_returns_none_when_store_is_absent() -> None: + store = _make_store(_make_response(404, {"error": {"message": "not found"}}), name="checkpoints") + + assert await store.get() is None + + +@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", + ) + + result = await store.get("step/1") + + request = _sent_request(store) + assert request.method == "GET" + 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("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 + + +# --------------------------------------------------------------------------- +# 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 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", + ) + + result = await store.delete("step/1", if_match='"0x8DD"') + + request = _sent_request(store) + assert request.method == "DELETE" + assert request.headers["If-Match"] == '"0x8DD"' + 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("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("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", + ) + + 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.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 == KeyPage( + 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..2e34c24b0d41 --- /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 = {} + 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 -- see + # test_foundry_state_store.py's get_or_create orchestration tests for that). + 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("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("step-1") + assert stale_item is not None + await store.set("step-1", {"done": True, "attempt": 2}, tags={"kind": "checkpoint"}) + + with pytest.raises(FoundryStoragePreconditionError) as exc: + await store.set("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("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..a7a818773a8d --- /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_unparseable_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" diff --git a/sdk/agentserver/azure-ai-agentserver-core/type_spec/README.md b/sdk/agentserver/azure-ai-agentserver-core/type_spec/README.md new file mode 100644 index 000000000000..8b45f4a1b935 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/type_spec/README.md @@ -0,0 +1,41 @@ +# TypeSpec source for storage models + +`main.tsp` is the formal TypeSpec contract for the Foundry State Storage wire +API (`/storage/state_stores/*`), companion to the prose spec +`foundry-state-storage-provider-spec.md` in `coreai-microsoft/foundrysdk_specs`. + +It is the source of truth for the generated request/response model classes +under `azure/ai/agentserver/core/storage/_generated/` (`StateStore`, +`StateStoreItem`, `CreateStateStoreRequest`, etc.) — those files carry a +"Code generated ... do not edit" header and should only be changed by +re-running generation, not by hand. + +## Why this isn't synced via `tsp-client` (yet) + +Sibling packages like `azure-ai-agentserver-responses` generate models by +running `tsp-client sync` against a `tsp-location.yaml` pointing at a +directory under `Azure/azure-rest-api-specs` +(`specification/ai-foundry/data-plane/Foundry/src/...`). This contract does +not live there yet — there is no `state-stores` directory upstream. Until it +is published there, `main.tsp` is authored and compiled locally. + +Once it lands upstream, replace `main.tsp` with a `tsp-location.yaml` +pointing at the new `specification/ai-foundry/data-plane/Foundry/src/state-stores` +directory and switch `make generate-models` back to the standard +`tsp-client sync` + compile flow, matching `azure-ai-agentserver-responses`. + +## Regenerating models + +From the package root: + +``` +make install-typespec-deps # one-time: npm install pinned TypeSpec toolchain +make generate-models # compile main.tsp -> azure/ai/agentserver/core/storage/_generated/ +``` + +This runs `npx tsp compile main.tsp --emit @azure-tools/typespec-python` +(the same emitter used to generate `azure-ai-agentserver-responses`' models) +and copies the resulting `_models.py`, `_enums.py`, and `_utils/model_base.py` +into `_generated/`, discarding the generated client/operations code (this +package builds its own requests by hand in `_state.py`; only the model +classes are used). diff --git a/sdk/agentserver/azure-ai-agentserver-core/type_spec/main.tsp b/sdk/agentserver/azure-ai-agentserver-core/type_spec/main.tsp new file mode 100644 index 000000000000..8c444f1bf226 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/type_spec/main.tsp @@ -0,0 +1,482 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +/** + * Foundry State Storage — formal TypeSpec contract. + * + * Companion to the prose spec `foundry-state-storage-provider-spec.md` in + * coreai-microsoft/foundrysdk_specs: this is the same `/storage/state_stores/*` + * surface expressed as a compilable TypeSpec contract. It is the source for + * the generated request/response model classes under + * `azure/ai/agentserver/core/storage/_generated/` (see `Makefile` / + * `type_spec/README.md` in this package for the local generation workflow). + * + * Authentication, base URL / API version, and the error envelope are + * inherited from the parent `foundry-storage-protocol-spec.md` (§2, §3, §8) + * and are stubbed locally below just enough to make this file self-contained. + * + * Status: this contract does not (yet) live under Azure/azure-rest-api-specs + * the way sibling agentserver contracts (e.g. `sdk-service-agentserver-contracts`, + * consumed by `azure-ai-agentserver-responses`) do — there is no + * `state-stores` directory there yet. Until it lands upstream, models are + * generated locally from this file directly (`make generate-models`) rather + * than synced via `tsp-client`. Once the contract is published under + * `specification/ai-foundry/data-plane/Foundry/src/state-stores/`, this file + * should be replaced by a `tsp-location.yaml` pointing there, matching the + * `azure-ai-agentserver-responses` pattern. + * + * Not public-facing (for now): this is an internal platform<->container wire + * API. Agent code reaches it only through an SDK state-provider layer (see + * `azure.ai.agentserver.core.storage.FoundryStateStore`). + */ + +import "@typespec/http"; + +using TypeSpec.Http; + +@service(#{ title: "Foundry State Storage" }) +namespace Foundry.Storage.StateStore; + +// --------------------------------------------------------------------------- +// Shared / inherited-from-parent stubs (kept minimal; see parent spec §2/§8). +// --------------------------------------------------------------------------- + +/** Unix epoch seconds, as returned in `created_at` / `updated_at`. */ +scalar FoundryTimestamp extends int64; + +/** Standard Foundry error envelope (parent spec §8.1). */ +@error +model ApiErrorResponse { + error: ApiError; +} + +model ApiError { + @doc("Machine-readable error code.") + code: string; + + @doc("Human-readable error message.") + message: string; + + @doc("For 400 responses, the name of the offending request field.") + param?: string; +} + +/** A generic paged-list envelope, shared by store list and item list-keys. */ +@doc("A page of results, ordered by `created_at` (parent spec §7).") +model ListResponse { + @doc("Always `\"list\"`.") + object: "list"; + + @doc("The page of results.") + data: T[]; + + @doc("The `id` of the first item in `data`, or `null` if `data` is empty.") + first_id: string | null; + + @doc("The `id` of the last item in `data`, or `null` if `data` is empty.") + last_id: string | null; + + @doc("Whether another page is available after `last_id`.") + has_more: boolean; +} + +/** Shared list query parameters (parent spec §7): limit, order, and cursor paging. */ +model ListQueryParams { + @query + @doc("Maximum number of results to return.") + @minValue(1) + @maxValue(100) + limit?: int32 = 20; + + @query + @doc("Sort order over `created_at`.") + order?: "asc" | "desc" = "desc"; + + @query + @doc("Cursor: return results after this id. Mutually exclusive with `before`.") + after?: string; + + @query + @doc("Cursor: return results before this id. Mutually exclusive with `after`.") + before?: string; +} + +/** + * Trusted-only delegation header: supplies the resolved `user_id` for a + * user-isolated store's item operations (spec §3). Honored only for trusted + * (S2S/managed-identity) callers; ignored otherwise (falls back to + * `hash(oid, tid)`). + */ +model DelegatedUserIdHeader { + @header("x-ms-user-id") + xMsUserId?: string; +} + +/** Optional single-item optimistic-concurrency header (spec §4.2). */ +model IfMatchHeader { + @doc("`\"\"` for a conditional write/delete, or `\"*\"` to require the item already exists. Omit for unconditional last-write-wins.") + @header("If-Match") + ifMatch?: string; +} + +/** The `ETag` response header carrying an item's current concurrency token. */ +model ETagHeader { + @header("ETag") + etag?: string; +} + +// --------------------------------------------------------------------------- +// 4.1 State Store Resource +// --------------------------------------------------------------------------- + +/** The `object` discriminator for a state-store resource. */ +union StateStoreObjectType { + stateStore: "state_store", +} + +/** A first-class, explicitly-created durable key-value store (spec §2, §4.1). */ +model StateStore { + @doc("Server-assigned, read-only. Not used for addressing (the `name` is).") + id: string; + + @doc("Always `\"state_store\"`.") + object: StateStoreObjectType.stateStore; + + @doc("Caller-chosen logical name. Unique within [project_id, agent_guid]. May contain `/` as a hierarchy separator. Immutable.") + @minLength(1) + @maxLength(128) + name: string; + + @doc("Whether item operations are partitioned per resolved user (spec §3). Omitted at create -> `false`. Immutable.") + user_isolation: boolean = false; + + @doc("Store-wide default TTL (seconds) inherited by every item, write-sliding per item. `-1` = never expire. Immutable.") + @minValue(-1) + @maxValue(2147483647) + item_ttl_seconds: int32 = 2592000; + + @doc("Optional free-form description. Mutable via `updateStateStore`.") + @maxLength(1024) + description?: string; + + @doc("Optional string-to-string labels. Mutable (replaced wholesale) via `updateStateStore`.") + tags?: Record; + + @doc("Creation time.") + created_at: FoundryTimestamp; + + @doc("Last metadata-update time.") + updated_at: FoundryTimestamp; +} + +/** Request body for `createStateStore` (spec §4.1.1). */ +model CreateStateStoreRequest { + @doc("Required. Unique within [project_id, agent_guid].") + @minLength(1) + @maxLength(128) + name: string; + + @doc("Optional. Omitted -> `false` (agent-level, shared). Fixed at create.") + user_isolation?: boolean = false; + + @doc("Optional. Default 30 days (2592000). `-1` = never expire. Fixed at create.") + @minValue(-1) + @maxValue(2147483647) + item_ttl_seconds?: int32 = 2592000; + + @maxLength(1024) + description?: string; + + tags?: Record; +} + +/** + * 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`. + */ +model UpdateStateStoreRequest { + @maxLength(1024) + description?: string; + + tags?: Record; +} + +/** Response for `deleteStateStore` (spec §4.1.5). Idempotent. */ +model DeletedStateStore { + @doc("Present when the store existed; absent when the delete was a no-op (already gone).") + id?: string; + + @doc("Always `\"state_store\"`.") + object: StateStoreObjectType.stateStore; + + name: string; + + @doc("Always `true`.") + deleted: boolean; +} + +// --------------------------------------------------------------------------- +// 4.2 Item Resource +// --------------------------------------------------------------------------- + +/** The `object` discriminator for an item resource. */ +union StateStoreItemObjectType { + stateStoreItem: "state_store.item", +} + +/** A single `key -> value` record within a state store (spec §2, §4.2). */ +model StateStoreItem { + @doc("Server-assigned, read-only. Not used for addressing (the `key` is.)") + id: string; + + @doc("Always `\"state_store.item\"`.") + object: StateStoreItemObjectType.stateStoreItem; + + @doc("Logical key, unique within the store. May contain `/`.") + @minLength(1) + @maxLength(128) + key: string; + + @doc("Opaque JSON object; stored and returned verbatim. <= 1 MB serialized inline.") + value: Record; + + @doc("Optional labels used to filter `listKeys`. Replaced wholesale on write.") + tags?: Record; + + @doc("Optimistic-concurrency token for this item; also echoed in the `ETag` response header.") + etag: string; + + created_at: FoundryTimestamp; + updated_at: FoundryTimestamp; +} + +/** Request body for `createItem` (spec §4.2.1) — fails with `409` on a duplicate key. */ +model CreateItemRequest { + @minLength(1) + @maxLength(128) + key: string; + + value: Record; + + tags?: Record; +} + +/** Request body for `putItem` (spec §4.2.2) — idempotent create-or-replace by key. */ +model PutItemRequest { + value: Record; + + tags?: Record; +} + +/** + * Response for `createItem` / `putItem` (spec §4.2.1/§4.2.2) — server + * metadata only, no `value` echoed back. + */ +model StateStoreItemMetadata { + id: string; + + @doc("Always `\"state_store.item\"`.") + object: StateStoreItemObjectType.stateStoreItem; + + key: string; + + etag: string; + + created_at: FoundryTimestamp; + updated_at: FoundryTimestamp; +} + +/** Response for `deleteItem` (spec §4.2.4). Idempotent. */ +model DeletedStateStoreItem { + @doc("Present when the item existed; absent when the delete was a no-op (already gone).") + id?: string; + + @doc("Always `\"state_store.item\"`.") + object: StateStoreItemObjectType.stateStoreItem; + + key: string; + + @doc("Always `true`.") + deleted: boolean; +} + +/** A keys-only projection entry returned by `listKeys` (spec §4.2.5) -- never includes `value`. */ +model StateStoreKey { + id: string; + + @doc("Always `\"state_store.item\"`.") + object: StateStoreItemObjectType.stateStoreItem; + + key: string; + + tags?: Record; + + etag: string; + + created_at: FoundryTimestamp; + updated_at: FoundryTimestamp; +} + +// --------------------------------------------------------------------------- +// 4. Endpoints +// --------------------------------------------------------------------------- + +@route("/storage/state_stores") +interface StateStores { + /** Create — spec §4.1.1. `201` on success, `409` if `name` already exists. */ + @post + create(@body body: CreateStateStoreRequest): { + @statusCode statusCode: 201; + @body store: StateStore; + } | { + @statusCode statusCode: 409; + @body error: ApiErrorResponse; + }; + + /** List — spec §4.1.4. Single-partition query over the store catalog. */ + @get + list(...ListQueryParams): { + @statusCode statusCode: 200; + @body page: ListResponse; + }; + + /** Read — spec §4.1.2. `404` if no such store exists. */ + @get + @route("{state_store}") + read(@path state_store: string): { + @statusCode statusCode: 200; + @body store: StateStore; + } | { + @statusCode statusCode: 404; + @body error: ApiErrorResponse; + }; + + /** + * Update — spec §4.1.3. Mutable metadata only (`description` / `tags`); + * `400` if the caller tries to change an immutable field. + */ + @patch + @route("{state_store}") + update(@path state_store: string, @body body: UpdateStateStoreRequest): { + @statusCode statusCode: 200; + @body store: StateStore; + } | { + @statusCode statusCode: 404; + @body error: ApiErrorResponse; + } | { + @statusCode statusCode: 400; + @body error: ApiErrorResponse; + }; + + /** + * Delete (cascade) — spec §4.1.5. Removes the store descriptor and every + * item under it. Idempotent: deleting an absent store still returns `200`. + */ + @delete + @route("{state_store}") + delete(@path state_store: string): { + @statusCode statusCode: 200; + @body deleted: DeletedStateStore; + }; +} + +@route("/storage/state_stores/{state_store}/items") +interface StateStoreItems { + /** + * Create — spec §4.2.1. Fails with `409` on a duplicate `key`; `404` if the + * store does not exist (there is no implicit store creation). + */ + @post + create(@path state_store: string, @body body: CreateItemRequest, ...DelegatedUserIdHeader): { + @statusCode statusCode: 201; + @body item: StateStoreItemMetadata; + } | { + @statusCode statusCode: 409; + @body error: ApiErrorResponse; + } | { + @statusCode statusCode: 404; + @body error: ApiErrorResponse; + }; + + /** + * List keys — spec §4.2.5 (`items:keys` custom method). Never returns + * `value`; ordered by write time within the single partition. + */ + @get + @route("{state_store}/items:keys") + listKeys( + @path state_store: string, + ...ListQueryParams, + @doc("Optional, repeatable tag-equality filter, e.g. `tags.version=1`. AND-combined; omit to list every key.") + @query + tags?: Record, + ...DelegatedUserIdHeader + ): { + @statusCode statusCode: 200; + @body page: ListResponse; + }; + + /** + * Update (create-or-replace) — spec §4.2.2. Idempotent `PUT` by key; + * `If-Match` gives optimistic concurrency (`412` on mismatch); each write + * re-slides the item's inherited TTL. `404` if the store does not exist. + */ + @put + @route("{state_store}/items/{key}") + put( + @path state_store: string, + @path key: string, + @body body: PutItemRequest, + ...IfMatchHeader, + ...DelegatedUserIdHeader + ): { + @statusCode statusCode: 201; + ...ETagHeader; + @body item: StateStoreItemMetadata; + } | { + @statusCode statusCode: 200; + ...ETagHeader; + @body item: StateStoreItemMetadata; + } | { + @statusCode statusCode: 412; + ...ETagHeader; + @body error: ApiErrorResponse; + } | { + @statusCode statusCode: 404; + @body error: ApiErrorResponse; + }; + + /** Get — spec §4.2.3. `404` if the key has never been written (a normal empty result). */ + @get + @route("{state_store}/items/{key}") + get(@path state_store: string, @path key: string, ...DelegatedUserIdHeader): { + @statusCode statusCode: 200; + ...ETagHeader; + @body item: StateStoreItem; + } | { + @statusCode statusCode: 404; + @body error: ApiErrorResponse; + }; + + /** + * Delete — spec §4.2.4. `If-Match` gives a conditional delete (`412` on + * mismatch); idempotent -- deleting a missing key still returns `200`. + */ + @delete + @route("{state_store}/items/{key}") + delete( + @path state_store: string, + @path key: string, + ...IfMatchHeader, + ...DelegatedUserIdHeader + ): { + @statusCode statusCode: 200; + @body deleted: DeletedStateStoreItem; + } | { + @statusCode statusCode: 412; + ...ETagHeader; + @body error: ApiErrorResponse; + }; +} diff --git a/sdk/agentserver/azure-ai-agentserver-core/type_spec/package.json b/sdk/agentserver/azure-ai-agentserver-core/type_spec/package.json new file mode 100644 index 000000000000..10c2eddf8fd8 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/type_spec/package.json @@ -0,0 +1,16 @@ +{ + "name": "foundry-state-store-typespec", + "private": true, + "description": "Pinned TypeSpec toolchain for generating azure-ai-agentserver-core's storage request/response models from main.tsp. Not published; used only by `make generate-models`.", + "devDependencies": { + "@typespec/compiler": "1.13.0", + "@typespec/http": "1.13.0", + "@typespec/rest": "0.83.0", + "@typespec/versioning": "0.83.0", + "@typespec/openapi": "1.13.0", + "@typespec/openapi3": "1.13.0", + "@azure-tools/typespec-azure-core": "0.69.0", + "@azure-tools/typespec-client-generator-core": "0.69.2", + "@azure-tools/typespec-python": "0.63.2" + } +}