From 3c62b0f5ad6927a4267897896f8a7ab2205e22e5 Mon Sep 17 00:00:00 2001 From: Shanmukha Pasumarthy Date: Thu, 9 Jul 2026 20:57:01 +0530 Subject: [PATCH 1/2] feat(agentserver): FoundryStorage - M365 Storage adapter on Foundry durable state (Activity Protocol) Adds FoundryStorage, an implementation of the M365 Agents SDK Storage interface for azure-ai-agentserver-activity, backed by the durable FoundryStateStore state-store client from azure-ai-agentserver-core (see PR #47763). Design: - Store name = scope: every M365 storage key is already a scope identifier (e.g. "{channel}/conversations/{conversation_id}", "{channel}/users/{user_id}", "proactive/conversations/{conversation_id}"), so FoundryStorage lazily creates and caches one FoundryStateStore per distinct key instead of a shared namespace. The key doubles as both the store name and the item key. - Keys shaped like the M365 UserState key ("/users/" segment) automatically get user_isolation=True on their backing store; override via is_user_scoped=. - Subclasses the M365 SDK's AsyncStorageBase and implements _read_item/_write_item/_delete_item; batching, validation, and concurrent fan-out come from the base class. - The backing store is only created (get_or_create) lazily on first write for a given key; read/delete treat a not-yet-created store as a missing key. - ActivityAgentServerHost gains a storage= override wired through the M365 bridge (falls back to MemoryStorage). Also renumbers/fixes the new FoundryStorage samples (06-08) to match the package's 01-05 sample convention and current host constructor API, and pulls in azure-ai-agentserver-core's storage/ subtree from PR #47763 as-is (no other core changes) since FoundryStorage depends on FoundryStateStore's statestores-protocol shape. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../CHANGELOG.md | 12 + .../azure-ai-agentserver-activity/README.md | 21 + .../azure/ai/agentserver/activity/__init__.py | 3 +- .../agentserver/activity/_foundry_storage.py | 168 +++++ .../dev_requirements.txt | 1 + .../pyproject.toml | 3 +- .../samples/06-foundry-storage-state/main.py | 67 ++ .../06-foundry-storage-state/requirements.txt | 5 + .../07-foundry-storage-proactive/main.py | 100 +++ .../requirements.txt | 5 + .../08-foundry-storage-history/main.py | 88 +++ .../requirements.txt | 5 + .../tests/test_foundry_storage.py | 183 ++++++ .../azure-ai-agentserver-core/CHANGELOG.md | 6 + .../azure-ai-agentserver-core/README.md | 23 + .../azure/ai/agentserver/core/_version.py | 2 +- .../ai/agentserver/core/storage/__init__.py | 51 ++ .../ai/agentserver/core/storage/_client.py | 111 ++++ .../ai/agentserver/core/storage/_endpoint.py | 69 +++ .../ai/agentserver/core/storage/_errors.py | 123 ++++ .../ai/agentserver/core/storage/_json.py | 15 + .../ai/agentserver/core/storage/_policies.py | 147 +++++ .../ai/agentserver/core/storage/_state.py | 300 +++++++++ .../core/storage/_state_serializer.py | 253 ++++++++ .../docs/state-store-guide.md | 272 +++++++++ .../azure-ai-agentserver-core/pyproject.toml | 1 + .../samples/state_store_sample.py | 118 ++++ .../tests/test_foundry_state_store.py | 576 ++++++++++++++++++ .../tests/test_state_store_sample_usage.py | 241 ++++++++ .../tests/test_storage_endpoint.py | 57 ++ .../tests/test_storage_errors.py | 88 +++ .../tests/test_storage_policies.py | 48 ++ 32 files changed, 3159 insertions(+), 3 deletions(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/_foundry_storage.py create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/main.py create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/samples/06-foundry-storage-state/requirements.txt create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/main.py create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/samples/07-foundry-storage-proactive/requirements.txt create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/main.py create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/samples/08-foundry-storage-history/requirements.txt create mode 100644 sdk/agentserver/azure-ai-agentserver-activity/tests/test_foundry_storage.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/__init__.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_endpoint.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_errors.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_json.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_policies.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state_serializer.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md create mode 100644 sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_state_store_sample_usage.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_endpoint.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_errors.py create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py diff --git a/sdk/agentserver/azure-ai-agentserver-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..cff47ed91e52 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/azure/ai/agentserver/activity/_foundry_storage.py @@ -0,0 +1,168 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""M365 Agents SDK storage adapter backed by per-key Foundry state stores.""" +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype +# pylint: disable=docstring-keyword-should-match-keyword-only,import-error,no-name-in-module + +from __future__ import annotations + +import asyncio +from typing import Any, Callable, Optional, Tuple, Type, TypeVar + +from azure.ai.agentserver.core.storage import FoundryStateStore, FoundryStorageEndpoint, FoundryStorageNotFoundError +from azure.core.credentials_async import AsyncTokenCredential + +try: + from microsoft_agents.hosting.core.storage import AsyncStorageBase, Storage +except ImportError: # pragma: no cover - keeps package importable without optional M365 SDK bits. + class Storage: # type: ignore[no-redef] + """Fallback base class used only when the M365 Agents SDK is not installed.""" + + class AsyncStorageBase(Storage): # type: ignore[no-redef] + """Fallback base class used only when the M365 Agents SDK is not installed.""" + + +StoreItemT = TypeVar("StoreItemT") + +#: Segment the M365 Agents SDK ``UserState`` uses to build per-user storage keys +#: (``f"{channel_id}/users/{user_id}"`` -- see +#: ``microsoft_agents.hosting.core.state.user_state.UserState.get_storage_key``). +#: Keys containing this segment get ``user_isolation=True`` on their backing +#: store by default. +_USER_SCOPE_SEGMENT = "/users/" + + +def _default_is_user_scoped(key: str) -> bool: + """Match the M365 ``UserState`` key shape ``"{channel_id}/users/{user_id}"``.""" + return _USER_SCOPE_SEGMENT in key + + +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 + M365 ``UserState`` key shape (``"{channel_id}/users/{user_id}"``). + """ + + 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: + 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._stores: dict[str, FoundryStateStore] = {} + self._ensured_keys: set[str] = set() + self._creation_lock = asyncio.Lock() + + def _new_store(self, key: str) -> FoundryStateStore: + 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 FoundryStateStore(key, **kwargs) + + async def _get_store(self, key: str, *, ensure_exists: bool) -> FoundryStateStore: + """Return the cached per-key store, creating the client (and, when + ``ensure_exists``, the server-side store resource) on first use.""" + store = self._stores.get(key) + if store is None: + async with self._creation_lock: + store = self._stores.get(key) + if store is None: + store = self._new_store(key) + self._stores[key] = store + if ensure_exists and key not in self._ensured_keys: + async with self._creation_lock: + if key not in self._ensured_keys: + await store.get_or_create() + self._ensured_keys.add(key) + return store + + async def aclose(self) -> None: + """Close every cached per-key store and an owned default credential.""" + 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": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.aclose() + + async def _read_item( + self, + key: str, + *, + target_cls: Type[StoreItemT] | None = None, + **kwargs: Any, + ) -> Tuple[Optional[str], Optional[StoreItemT]]: + """Fetch one item. A store that does not exist yet is just a missing key.""" + _ = kwargs + store = await self._get_store(key, ensure_exists=False) + try: + item = await store.get(key) + except FoundryStorageNotFoundError: + item = None + if item is None: + return None, None + return key, target_cls.from_json_to_store_item(item.value) # type: ignore[attr-defined] + + async def _write_item(self, key: str, value: StoreItemT) -> None: + """Create-or-replace one item, creating its backing store on first write.""" + store = await self._get_store(key, ensure_exists=True) + await store.set(key, value.store_item_to_json()) # type: ignore[attr-defined] + + async def _delete_item(self, key: str) -> None: + """Delete one item. Missing keys (or stores) are ignored.""" + 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..269943868e60 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-activity/tests/test_foundry_storage.py @@ -0,0 +1,183 @@ +# 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, StateItem + + +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_or_create = AsyncMock() + 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: + factory = MagicMock(side_effect=lambda key, **kwargs: stores_by_key[key]) + monkeypatch.setattr(module, "FoundryStateStore", factory) + return factory + + +@pytest.mark.asyncio +async def test_read_missing_key_does_not_create_a_store(monkeypatch: pytest.MonkeyPatch) -> None: + store = _fake_store() + factory = _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + result = await storage.read(["k"], target_cls=_TestStoreItem) + + assert result == {} + factory.assert_called_once() # client constructed to issue the GET ... + store.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=StateItem(id="i1", key="k", value={"count": 3}, etag="e1")) + _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_store_not_found_as_missing(monkeypatch: pytest.MonkeyPatch) -> None: + store = _fake_store() + store.get = AsyncMock(side_effect=FoundryStorageNotFoundError("not found")) + _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() + _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + await storage.write({"k": _TestStoreItem({"turn": 4})}) + + store.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() + _patch_stores(monkeypatch, {"k": store}) + storage = FoundryStorage() + + await storage.write({"k": _TestStoreItem({"turn": 1})}) + await storage.write({"k": _TestStoreItem({"turn": 2})}) + + store.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-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index ba3a1ab6b47b..f66efd712630 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` package providing the protocol-neutral Foundry storage layer: `FoundryStorageClient` (transport, endpoint, pipeline policies, error hierarchy) and `FoundryStateStore`, a generic durable key-value store. Every operation is scoped to exactly one caller-supplied `namespace`, carried as a percent-encoded URL path segment (`POST /storage/namespaces/{namespace}/state:read|:write|:listKeys`); server-side isolation is selected by the trusted `x-ms-internal-state-session-isolation` / `x-ms-internal-state-user-isolation` headers (`isolate_sessions` / `isolate_users` knobs; disabling both is rejected). Batch `write` is atomic and reports per-key `WriteResult` metadata (etag + `created_at` / `updated_at`); `read` and `list_keys` return the same server-managed timestamps. Supports optional `if_match` optimistic concurrency, optional per-item `ttl_seconds` (surfaced as `StateItem.expires_at`), and ordered, paged `list_keys` (keys + metadata only). `FoundryStateStore.for_namespace(...)` returns a namespace-bound `NamespaceStateStore` handle. Writes use the strongly-typed `Upsert` / `Delete` change objects (`WriteChange`) and stored values are typed as `JSONValue`. Protocol packages can build resource-specific clients on top of `FoundryStorageClient`. 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/README.md b/sdk/agentserver/azure-ai-agentserver-core/README.md index 2a9d56c42346..b4ee137d8551 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/README.md +++ b/sdk/agentserver/azure-ai-agentserver-core/README.md @@ -126,6 +126,29 @@ 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 — with +optimistic concurrency, tag filtering, key listing, and optional per-item TTL. + +```python +from azure.ai.agentserver.core.storage import FoundryStateStore + +# Endpoint and credential resolve from FOUNDRY_PROJECT_ENDPOINT + DefaultAzureCredential. +async with FoundryStateStore() as store: + etag = await store.set("counters", "page-views", 1) + item = await store.get("counters", "page-views") + print(item.value) # 1 +``` + +Reads return typed `StateItem` values; writes are expressed as typed `Upsert` / +`Delete` changes. 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 concurrency model, 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..fa07dc12f2cc --- /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 ( + DeletedStateItem, + DeletedStateStore, + JSONObject, + JSONValue, + KeyPage, + Order, + StateItem, + StateItemMetadata, + StateKey, + StateStoreInfo, +) + +__all__ = [ + "DEFAULT_ITEM_TTL_SECONDS", + "DeletedStateItem", + "DeletedStateStore", + "FOUNDRY_TOKEN_SCOPE", + "FoundryStateStore", + "FoundryStorageApiError", + "FoundryStorageBadRequestError", + "FoundryStorageConflictError", + "FoundryStorageClient", + "FoundryStorageEndpoint", + "FoundryStorageError", + "FoundryStorageNotFoundError", + "FoundryStoragePreconditionError", + "JSONObject", + "JSONValue", + "KeyPage", + "Order", + "StateItem", + "StateItemMetadata", + "StateKey", + "StateStoreInfo", +] diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py new file mode 100644 index 000000000000..bf33c482e752 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_client.py @@ -0,0 +1,111 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Protocol-neutral HTTP transport for Foundry-backed storage clients. + +:class:`FoundryStorageClient` owns the :class:`~azure.core.AsyncPipelineClient`, +its policy chain (retry, auth, logging, tracing), and the shared +``_send_storage_request`` helper. Protocol packages subclass it and add their +own resource methods and payload (de)serialization on top. +""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype +# pylint: disable=client-accepts-api-version-keyword + +from __future__ import annotations + +from typing import Any, Callable + +from azure.core import AsyncPipelineClient +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.exceptions import ServiceRequestError, ServiceResponseError +from azure.core.pipeline import policies +from azure.core.rest import AsyncHttpResponse, HttpRequest + +from azure.ai.agentserver.core._platform_headers import PLATFORM_ERROR_TAG +from azure.ai.agentserver.core._version import VERSION + +from ._endpoint import FoundryStorageEndpoint +from ._errors import raise_for_storage_error +from ._policies import FoundryStorageLoggingPolicy, ServerVersionUserAgentPolicy + +#: OAuth scope used to acquire bearer tokens for the Foundry storage API. +FOUNDRY_TOKEN_SCOPE = "https://ai.azure.com/.default" +#: Content type for JSON request bodies sent to the Foundry storage API. +JSON_CONTENT_TYPE = "application/json; charset=utf-8" + + +class FoundryStorageClient: + """Base HTTP client for the Foundry storage API. + + :param credential: An async credential used to obtain bearer tokens. + :type credential: AsyncTokenCredential + :param endpoint: Resolved Foundry storage endpoint configuration. + :type endpoint: FoundryStorageEndpoint + :param get_server_version: Optional callable returning the server version + string to use as the ``User-Agent`` header, evaluated lazily on each + request. When ``None``, Azure Core's default ``UserAgentPolicy`` is used. + :type get_server_version: Callable[[], str] | None + :param sdk_moniker: SDK moniker for the default ``UserAgentPolicy``. Ignored + when *get_server_version* is supplied. + :type sdk_moniker: str | None + """ + + def __init__( + self, + credential: AsyncTokenCredential, + endpoint: FoundryStorageEndpoint, + *, + get_server_version: Callable[[], str] | None = None, + sdk_moniker: str | None = None, + **kwargs: Any, + ) -> None: + self._endpoint = endpoint + + ua_policy: policies.UserAgentPolicy | ServerVersionUserAgentPolicy + if get_server_version is not None: + ua_policy = ServerVersionUserAgentPolicy(get_server_version) + else: + ua_policy = policies.UserAgentPolicy(sdk_moniker=sdk_moniker or f"ai-agentserver-core/{VERSION}") + + self._client: AsyncPipelineClient = AsyncPipelineClient( + base_url=endpoint.storage_base_url, + policies=[ + policies.RequestIdPolicy(), + policies.HeadersPolicy(), + ua_policy, + policies.AsyncRetryPolicy(), + policies.AsyncBearerTokenCredentialPolicy(credential, FOUNDRY_TOKEN_SCOPE), + FoundryStorageLoggingPolicy(), + # NOTE: ``ContentDecodePolicy`` is intentionally NOT included. It + # eagerly decodes every response body as JSON and crashes with + # ``UnicodeDecodeError`` on non-UTF-8 bodies (gzip payloads, HTML + # error pages, transport-corrupted bodies). Serializers call + # ``http_resp.text()`` directly with defensive error handling. + policies.DistributedTracingPolicy(), + ], + **kwargs, + ) + + async def aclose(self) -> None: + """Close the underlying HTTP pipeline client.""" + await self._client.close() + + async def __aenter__(self) -> "FoundryStorageClient": + return self + + async def __aexit__(self, *args: Any) -> None: + await self.aclose() + + async def _send_storage_request(self, request: HttpRequest) -> AsyncHttpResponse: + """Send an HTTP request to the Foundry storage API. + + Transport-level exceptions are tagged as platform errors before being + re-raised; non-2xx responses raise a :class:`FoundryStorageError`. + """ + try: + http_resp: AsyncHttpResponse = await self._client.send_request(request) + except (ServiceRequestError, ServiceResponseError, OSError) as exc: + setattr(exc, PLATFORM_ERROR_TAG, True) + raise + raise_for_storage_error(http_resp) + return http_resp diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_endpoint.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_endpoint.py new file mode 100644 index 000000000000..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/_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..030c79bf71cb --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_policies.py @@ -0,0 +1,147 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Azure Core pipeline policies shared by Foundry storage clients. + +Both the request User-Agent policy and the per-retry logging policy are +protocol-neutral and reused by every AgentServer storage client. +""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype + +from __future__ import annotations + +import logging +import time +import urllib.parse +from typing import Callable + +from azure.ai.agentserver.core._platform_headers import ( + APIM_REQUEST_ID, + CLIENT_REQUEST_ID, + REQUEST_ID, + TRACEPARENT, +) +from azure.core.pipeline import PipelineRequest, PipelineResponse +from azure.core.pipeline.policies import AsyncHTTPPolicy, SansIOHTTPPolicy +from azure.core.rest import AsyncHttpResponse, HttpRequest, HttpResponse + +logger = logging.getLogger("azure.ai.agentserver.core") + +_LOG_FMT_START = "Foundry storage %s %s starting (x-ms-client-request-id=%s, traceparent=%s)" +_LOG_FMT_TRANSPORT_FAILURE = ( + "Foundry storage %s %s transport failure after %.1fms (x-ms-client-request-id=%s, traceparent=%s)" +) +_LOG_FMT_RESPONSE = ( + "Foundry storage %s %s -> %d (%.1fms, " + "x-ms-client-request-id=%s, traceparent=%s, x-request-id=%s, apim-request-id=%s)" +) + + +class ServerVersionUserAgentPolicy(SansIOHTTPPolicy[HttpRequest, HttpResponse]): + """Pipeline policy that sets the ``User-Agent`` header lazily from a callback. + + Unlike :class:`~azure.core.pipeline.policies.UserAgentPolicy` which captures + the value at construction time, this policy evaluates the callback on each + request so it reflects segments registered after the pipeline was built + (e.g. by cooperative ``__init__`` in a multi-protocol host). + """ + + def __init__(self, get_server_version: Callable[[], str]) -> None: + super().__init__() + self._get_server_version = get_server_version + + def on_request(self, request: PipelineRequest[HttpRequest]) -> None: + """Set the ``User-Agent`` header before the request is sent.""" + request.http_request.headers["User-Agent"] = self._get_server_version() + + +def _mask_storage_url(url: str) -> str: + """Mask the sensitive portions of a Foundry storage URL.""" + try: + if not url: + return "(redacted)" + parsed = urllib.parse.urlparse(url) + path = parsed.path or "" + idx = path.find("/storage") + if idx < 0: + return "(redacted)" + storage_path = path[idx:] + masked = f"***{_redact_storage_path(storage_path)}" + qs = urllib.parse.parse_qs(parsed.query) + api_version = qs.get("api-version") + if api_version: + masked += f"?api-version={api_version[0]}" + return masked + except Exception: # pylint: disable=broad-exception-caught + return "(redacted)" + + +def _redact_storage_path(path: str) -> str: + if not path.startswith("/storage/"): + return "/storage" + segments = path.split("/") + redacted: list[str] = [] + for index, segment in enumerate(segments): + if index < 3 or segment in {"items", "items:keys", "statestores"} 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..a56fe4a5e558 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state.py @@ -0,0 +1,300 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Developer-facing Foundry state-store client bound to one explicit store.""" + +# pylint: disable=docstring-missing-param,docstring-missing-return,docstring-missing-rtype +# pylint: disable=client-accepts-api-version-keyword + +from __future__ import annotations + +import base64 +from collections.abc import Mapping +from typing import Any + +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.rest import HttpRequest + +from ._client import JSON_CONTENT_TYPE, FoundryStorageClient +from ._endpoint import FoundryStorageEndpoint +from ._errors import FoundryStorageConflictError, FoundryStorageNotFoundError +from ._state_serializer import ( + JSONObject, + KeyPage, + Order, + StateItem, + StateItemMetadata, + StateStoreInfo, + DeletedStateItem, + DeletedStateStore, + 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 +DELEGATED_USER_ID_HEADER = "x-ms-user-id" +_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/``), matching spec + PR 247's removal of built-in session isolation. + """ + + def __init__( + self, + name: str, + credential: AsyncTokenCredential | None = None, + endpoint: FoundryStorageEndpoint | str | None = None, + *, + user_isolation: bool = False, + item_ttl_seconds: int = DEFAULT_ITEM_TTL_SECONDS, + description: str | None = None, + tags: Mapping[str, str] | None = None, + user_id: str | None = None, + api_version: str = "v1", + **kwargs: Any, + ) -> None: + """Create a store-bound durable state-store client. + + :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. + :paramtype user_isolation: bool + :keyword item_ttl_seconds: Store-level default TTL inherited by every + item. + :paramtype item_ttl_seconds: int + :keyword description: Optional mutable store description. + :paramtype description: str or None + :keyword tags: Optional mutable store metadata tags. + :paramtype tags: ~collections.abc.Mapping[str, str] or None + :keyword user_id: Delegated end-user identity for trusted callers. + :paramtype user_id: str or None + :keyword api_version: Storage API version. + :paramtype api_version: str + :raises ValueError: If ``name`` is empty. + """ + if not name: + raise ValueError("name must be a non-empty string") + self._owns_credential = False + if credential is None: + try: + from azure.identity.aio import DefaultAzureCredential + except ImportError as exc: # pragma: no cover + raise ImportError( + "FoundryStateStore requires azure-identity when no credential is supplied. " + "Install azure-identity or pass an async credential." + ) from exc + credential = DefaultAzureCredential() + self._owns_credential = True + self._credential = credential + + if isinstance(endpoint, FoundryStorageEndpoint): + resolved = endpoint + elif isinstance(endpoint, str): + resolved = FoundryStorageEndpoint.from_endpoint(endpoint, api_version=api_version) + else: + resolved = FoundryStorageEndpoint.from_env(api_version=api_version) + + self._name = name + self._user_isolation = user_isolation + self._item_ttl_seconds = item_ttl_seconds + self._description = description + self._tags = {} if tags is None else dict(tags) + self._user_id = user_id + super().__init__(credential, resolved, **kwargs) + + async def aclose(self) -> None: + """Close the pipeline client and any owned default credential.""" + 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"statestores/{_encode_segment(self._name)}" + + def _item_path(self, key: str) -> str: + _validate_key(key) + return f"{self._store_path()}/items/{_encode_segment(key)}" + + def _request( + self, + method: str, + path: str, + *, + content: bytes | None = None, + include_user_id: bool = False, + if_match: str | None = None, + **query: str, + ) -> HttpRequest: + headers: dict[str, str] = {} + if content is not None: + headers["Content-Type"] = JSON_CONTENT_TYPE + if include_user_id and self._user_id is not None: + headers[DELEGATED_USER_ID_HEADER] = self._user_id + if if_match is not None: + headers["If-Match"] = if_match + return HttpRequest(method, self._endpoint.build_url(path, **query), content=content, headers=headers) + + async def create(self) -> StateStoreInfo: + """Create the bound store resource.""" + 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", "statestores", content=body)) + return deserialize_state_store(response.text()) + + async def create_or_get(self) -> StateStoreInfo: + """Create the bound store resource, or fetch it when it already exists.""" + try: + return await self.create() + except FoundryStorageConflictError: + return await self.get_properties() + + async def get_or_create(self) -> StateStoreInfo: + """Fetch the bound store resource, or create it when it does not exist.""" + try: + return await self.get_properties() + except FoundryStorageNotFoundError: + try: + return await self.create() + except FoundryStorageConflictError: + return await self.get_properties() + + async def get_properties(self) -> StateStoreInfo: + """Fetch the bound store descriptor.""" + response = await self._send_storage_request(self._request("GET", self._store_path())) + return deserialize_state_store(response.text()) + + async def update_metadata( + self, + *, + description: str | None | object = _UNSET, + tags: Mapping[str, str] | None | object = _UNSET, + ) -> StateStoreInfo: + """Update mutable store metadata on the bound store.""" + 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 = {} if tags is None else dict(tags) + return deserialize_state_store(response.text()) + + async def delete_store(self) -> DeletedStateStore: + """Delete the bound store and cascade-delete its items.""" + response = await self._send_storage_request(self._request("DELETE", self._store_path())) + return deserialize_deleted_state_store(response.text()) + + async def create_item(self, key: str, value: JSONObject, *, tags: Mapping[str, str] | None = None) -> StateItemMetadata: + """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, include_user_id=True) + ) + 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, + ) -> StateItemMetadata: + """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, + include_user_id=True, + if_match=header, + ) + ) + return deserialize_state_item_metadata(response.text()) + + async def get(self, key: str) -> StateItem | None: + """Fetch one item by key, returning ``None`` when it is absent.""" + try: + response = await self._send_storage_request(self._request("GET", self._item_path(key), include_user_id=True)) + except FoundryStorageNotFoundError: + return None + return deserialize_state_item(response.text()) + + async def delete(self, key: str, *, if_match: str | None = None) -> DeletedStateItem: + """Delete one item by key.""" + response = await self._send_storage_request( + self._request("DELETE", self._item_path(key), include_user_id=True, if_match=if_match) + ) + return deserialize_deleted_state_item(response.text()) + + async def list_keys( + self, + *, + tags: Mapping[str, str] | None = None, + limit: int | None = None, + after: str | None = None, + before: str | None = None, + order: Order = "desc", + ) -> 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", include_user_id=True, **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..59422f5ed320 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/storage/_state_serializer.py @@ -0,0 +1,253 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Serialization helpers and value types for the Foundry state-store client.""" + +# 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, field +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 ._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"] + + +@dataclass +class StateStoreInfo: + """Descriptor for a state store resource.""" + + id: str | None + name: str + user_isolation: bool + item_ttl_seconds: int | None + description: str | None = None + tags: dict[str, str] = field(default_factory=dict) + created_at: int | None = None + updated_at: int | None = None + + +@dataclass +class DeletedStateStore: + """Deleted-store marker returned by store delete operations.""" + + id: str | None + name: str + deleted: bool + + +@dataclass +class StateItemMetadata: + """Metadata returned by create/update item operations.""" + + id: str | None + key: str + etag: str | None = None + created_at: int | None = None + updated_at: int | None = None + + +@dataclass +class StateItem: + """A single stored item returned by ``get``.""" + + id: str | None + key: str + value: JSONObject + tags: dict[str, str] = field(default_factory=dict) + etag: str | None = None + created_at: int | None = None + updated_at: int | None = None + + +@dataclass +class DeletedStateItem: + """Deleted-item marker returned by item delete operations.""" + + id: str | None + key: str + deleted: bool + + +@dataclass +class StateKey: + """A key entry returned by ``list_keys``.""" + + id: str | None + key: str + tags: dict[str, str] = field(default_factory=dict) + etag: str | None = None + created_at: int | None = None + updated_at: int | None = None + + +@dataclass +class KeyPage: + """A page of keys returned by ``list_keys``.""" + + keys: list[StateKey] + first_id: str | None = None + last_id: str | None = None + has_more: bool = False + + +def serialize_store_create_request( + name: str, + *, + user_isolation: bool, + item_ttl_seconds: int, + description: str | None, + tags: Mapping[str, str], +) -> bytes: + payload: dict[str, Any] = { + "name": name, + "user_isolation": user_isolation, + "item_ttl_seconds": item_ttl_seconds, + } + if description is not None: + payload["description"] = description + if tags: + payload["tags"] = dict(tags) + return json.dumps(payload).encode("utf-8") + + +def serialize_store_update_request(description: str | None | object, tags: Mapping[str, str] | None | object) -> bytes: + payload: dict[str, Any] = {} + if description is not _UNSET: + payload["description"] = description + if tags is not _UNSET: + payload["tags"] = {} if tags is None else dict(tags) + return json.dumps(payload).encode("utf-8") + + +def serialize_item_create_request(key: str, value: JSONObject, tags: Mapping[str, str] | None) -> bytes: + payload: dict[str, Any] = {"key": key, "value": value} + if tags: + payload["tags"] = dict(tags) + return json.dumps(payload).encode("utf-8") + + +def serialize_item_put_request(value: JSONObject, tags: Mapping[str, str] | None) -> bytes: + payload: dict[str, Any] = {"value": value} + if tags: + payload["tags"] = dict(tags) + return json.dumps(payload).encode("utf-8") + + +def _as_epoch(value: Any) -> int | None: + return value if isinstance(value, int) and not isinstance(value, bool) else None + + +def _as_string_dict(value: Any) -> dict[str, str]: + if not isinstance(value, dict): + return {} + return {str(key): str(item) for key, item in value.items()} + + +def deserialize_state_store(body: str) -> StateStoreInfo: + data = load_json(body) + return StateStoreInfo( + id=_as_optional_str(data.get("id")), + name=str(data.get("name", "")), + user_isolation=bool(data.get("user_isolation", False)), + item_ttl_seconds=_as_epoch(data.get("item_ttl_seconds")), + description=_as_optional_str(data.get("description")), + tags=_as_string_dict(data.get("tags")), + created_at=_as_epoch(data.get("created_at")), + updated_at=_as_epoch(data.get("updated_at")), + ) + + +def deserialize_deleted_state_store(body: str) -> DeletedStateStore: + data = load_json(body) + return DeletedStateStore( + id=_as_optional_str(data.get("id")), + name=str(data.get("name", "")), + deleted=bool(data.get("deleted", False)), + ) + + +def deserialize_state_item_metadata(body: str) -> StateItemMetadata: + data = load_json(body) + return StateItemMetadata( + id=_as_optional_str(data.get("id")), + key=str(data.get("key", "")), + etag=_as_optional_str(data.get("etag")), + created_at=_as_epoch(data.get("created_at")), + updated_at=_as_epoch(data.get("updated_at")), + ) + + +def deserialize_state_item(body: str) -> StateItem: + data = load_json(body) + raw_value = data.get("value") + value = raw_value if isinstance(raw_value, dict) else {} + return StateItem( + id=_as_optional_str(data.get("id")), + key=str(data.get("key", "")), + value=value, + tags=_as_string_dict(data.get("tags")), + etag=_as_optional_str(data.get("etag")), + created_at=_as_epoch(data.get("created_at")), + updated_at=_as_epoch(data.get("updated_at")), + ) + + +def deserialize_deleted_state_item(body: str) -> DeletedStateItem: + data = load_json(body) + return DeletedStateItem( + id=_as_optional_str(data.get("id")), + key=str(data.get("key", "")), + deleted=bool(data.get("deleted", False)), + ) + + +def deserialize_list_keys_response(body: str) -> KeyPage: + data = load_json(body) + raw_items = data.get("data", []) + keys: list[StateKey] = [] + if isinstance(raw_items, list): + for entry in raw_items: + if isinstance(entry, dict) and entry.get("key") is not None: + keys.append( + StateKey( + id=_as_optional_str(entry.get("id")), + key=str(entry["key"]), + tags=_as_string_dict(entry.get("tags")), + etag=_as_optional_str(entry.get("etag")), + created_at=_as_epoch(entry.get("created_at")), + updated_at=_as_epoch(entry.get("updated_at")), + ) + ) + return KeyPage( + keys=keys, + first_id=_as_optional_str(data.get("first_id")), + last_id=_as_optional_str(data.get("last_id")), + has_more=bool(data.get("has_more", False)), + ) + + +def _as_optional_str(value: Any) -> str | None: + return str(value) if value is not None else None + + +_UNSET = object() 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..a82f3c80dc1b --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md @@ -0,0 +1,272 @@ +# Durable State Store Guide + +`FoundryStateStore` is a durable, server-backed state-store client for agent +state. It gives your agent an explicit store resource plus single-item +operations over that store: create the store once, then create, replace, fetch, +delete, and list items inside it. + +## 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. + +The SDK is the developer-facing layer over the internal `/storage/statestores/*` +protocol: it keeps the transport/auth pipeline in `FoundryStorageClient`, while +`FoundryStateStore` owns the ergonomic store-bound API. + +## Getting Started + +```python +from azure.ai.agentserver.core.storage import FoundryStateStore + +async with FoundryStateStore( + "checkpoints/thread-abc", + user_isolation=True, + item_ttl_seconds=3600, + description="Checkpoint store for thread abc", +) as store: + await store.get_or_create() + await store.set("step-1", {"done": False}) + + item = await store.get("step-1") + 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 + +The protocol no longer has a built-in session-isolation knob. If you want +conversation or thread scoping, encode it directly into the store name: + +```python +FoundryStateStore("checkpoints/thread-abc") +FoundryStateStore("workflow-state/run-42") +FoundryStateStore("user-prefs/defaults", user_isolation=True) +``` + +Because the store name is the logical identity, choose a stable naming scheme +up front. The raw name may contain `/`; the SDK handles the required base64url +path encoding on the wire. + +## Store Lifecycle + +Stores are **explicit resources**. Create or resolve them before writing items. + +```python +info = await store.create() +print(info.id, info.name, info.item_ttl_seconds) + +info = await store.create_or_get() +print(info.id, info.name, info.item_ttl_seconds) + +info = await store.get_or_create() +print(info.id, info.name, info.item_ttl_seconds) + +info = await store.get_properties() +info = await store.update_metadata( + description="Checkpoint store for prod traffic", + tags={"env": "prod", "team": "agents"}, +) + +deleted = await store.delete_store() +assert deleted.deleted is True +``` + +### Key points + +- `create()` raises `FoundryStorageConflictError` if the store name already exists. +- `get_or_create()` fetches the store first, or creates it when it is absent. +- `create_or_get()` creates the store first, or fetches the existing descriptor on `409`. +- Neither helper updates existing `user_isolation`, `item_ttl_seconds`, `description`, or `tags`. +- `update_metadata()` only changes `description` and `tags`. +- `user_isolation` and `item_ttl_seconds` are fixed at create time. +- `delete_store()` cascades to every item under that store name. + +## User Isolation and Delegated User IDs + +Set `user_isolation=True` when the same store name should fan out per user. + +```python +store = FoundryStateStore( + "user-prefs/defaults", + user_isolation=True, + user_id="aad-user-42", +) +``` + +- For direct callers, the platform can derive user identity from the token. +- For trusted callers acting on behalf of an end user, pass `user_id` so the SDK + sends the delegated `x-ms-user-id` header on item operations. +- Store-management calls (`create`, `get_properties`, `update_metadata`, + `delete_store`) stay store-scoped and do not send the delegated user header. + +## Values, Tags, and TTL + +Each item value is a JSON **object**. Store plain JSON, not Python objects. + +```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 = FoundryStateStore("otp/user-42", item_ttl_seconds=300) +await store.get_or_create() +``` + +- 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()` maps to the protocol's single-item `PUT`: create-or-replace by key. + +### Fetch one item + +```python +item = await store.get("step-1") +if item is not None: + print(item.id, item.key, item.value, item.tags, item.etag) +``` + +`get()` returns `None` when the item is missing. + +### 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 +page = 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. | +| `FoundryStorageBadRequestError` | 400 / 409 | Invalid request or duplicate/conflict. | +| `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) +``` + +## Best Practices + +1. **Create stores deliberately.** Do not assume item writes will create the + store for you. +2. **Encode conversation scope in the store name.** The store name replaces the + old 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 and close it + with `async with` or `await store.aclose()`. diff --git a/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml index 4d43985d6cbc..73e2e508142c 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml @@ -21,6 +21,7 @@ classifiers = [ keywords = ["azure", "azure sdk", "agent", "agentserver", "core"] dependencies = [ + "azure-core>=1.30.0", "starlette>=0.45.0", "hypercorn>=0.17.0", "opentelemetry-api>=1.40.0", diff --git a/sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py b/sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py new file mode 100644 index 000000000000..4845254494f5 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/samples/state_store_sample.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +FILE: state_store_sample.py + +DESCRIPTION: + Demonstrates the explicit Foundry state-store client (`FoundryStateStore`): + store creation, 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 ensure_store(store: FoundryStateStore) -> None: + """Create the backing store resource or reuse the existing one.""" + info = await store.get_or_create() + print(f"using store name={info.name}, id={info.id}, ttl={info.item_ttl_seconds}") + + +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_metadata( + 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_store() + print(f"deleted store -> name={deleted_store.name}, deleted={deleted_store.deleted}") + + +async def main() -> None: + store_name = f"checkpoints/sample-thread-{uuid4().hex}" + async with FoundryStateStore( + store_name, + user_isolation=True, + item_ttl_seconds=3600, + description="Sample state store", + ) as store: + await ensure_store(store) + 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..94c1a08f1200 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py @@ -0,0 +1,576 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Unit tests for FoundryStateStore request construction and response handling.""" + +from __future__ import annotations + +import base64 +import json +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from azure.ai.agentserver.core.storage import ( + DeletedStateItem, + DeletedStateStore, + FoundryStateStore, + FoundryStorageEndpoint, + KeyPage, + StateItem, + StateItemMetadata, + StateKey, + StateStoreInfo, +) + +_BASE_URL = "https://foundry.example.com/storage/" +_ENDPOINT = FoundryStorageEndpoint(storage_base_url=_BASE_URL) + + +def _encode_segment(value: str) -> str: + encoded = base64.urlsafe_b64encode(value.encode("utf-8")).decode("ascii") + return encoded.rstrip("=") + + +def _make_response(status_code: int, body: Any, *, headers: dict[str, str] | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = status_code + resp.headers = headers or {} + resp.text = MagicMock(return_value=json.dumps(body)) + return resp + + +def _make_store( + response: MagicMock, + *, + name: str = "langGraphCheckpoints/thread-abc", + user_isolation: bool = False, + item_ttl_seconds: int = 2592000, + description: str | None = None, + tags: dict[str, str] | None = None, + user_id: str | None = None, +) -> FoundryStateStore: + store = FoundryStateStore.__new__(FoundryStateStore) + store._endpoint = _ENDPOINT + store._owns_credential = False + store._name = name + store._user_isolation = user_isolation + store._item_ttl_seconds = item_ttl_seconds + store._description = description + store._tags = {} if tags is None else dict(tags) + store._user_id = user_id + mock_pipeline = AsyncMock() + mock_pipeline.send_request = AsyncMock(return_value=response) + mock_pipeline.close = AsyncMock() + store._client = mock_pipeline + return store + + +def _make_store_with_responses(*responses: MagicMock, name: str = "langGraphCheckpoints/thread-abc") -> FoundryStateStore: + store = _make_store(responses[0], name=name) + store._client.send_request = AsyncMock(side_effect=list(responses)) + return store + + +def _sent_request(store: FoundryStateStore) -> Any: + return store._client.send_request.call_args[0][0] + + +@pytest.mark.asyncio +async def test_create_posts_store_descriptor() -> None: + store = _make_store( + _make_response( + 201, + { + "id": "ss_1", + "object": "statestore", + "name": "langGraphCheckpoints/thread-abc", + "user_isolation": True, + "item_ttl_seconds": 600, + "description": "checkpoint store", + "tags": {"team": "agents"}, + "created_at": 1, + "updated_at": 1, + }, + ), + user_isolation=True, + item_ttl_seconds=600, + description="checkpoint store", + tags={"team": "agents"}, + ) + + result = await store.create() + + request = _sent_request(store) + assert request.method == "POST" + assert request.url == f"{_BASE_URL}statestores?api-version=v1" + assert request.headers["Content-Type"] == "application/json; charset=utf-8" + assert "x-ms-user-id" not in request.headers + assert json.loads(request.content.decode("utf-8")) == { + "name": "langGraphCheckpoints/thread-abc", + "user_isolation": True, + "item_ttl_seconds": 600, + "description": "checkpoint store", + "tags": {"team": "agents"}, + } + assert result == StateStoreInfo( + id="ss_1", + name="langGraphCheckpoints/thread-abc", + user_isolation=True, + item_ttl_seconds=600, + description="checkpoint store", + tags={"team": "agents"}, + created_at=1, + updated_at=1, + ) + + +@pytest.mark.asyncio +async def test_create_or_get_returns_created_store_when_absent() -> None: + store = _make_store( + _make_response( + 201, + { + "id": "ss_1", + "object": "statestore", + "name": "checkpoints", + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": None, + "tags": {}, + "created_at": 1, + "updated_at": 1, + }, + ), + name="checkpoints", + ) + + result = await store.create_or_get() + + request = _sent_request(store) + assert request.method == "POST" + assert request.url == f"{_BASE_URL}statestores?api-version=v1" + assert result.name == "checkpoints" + + +@pytest.mark.asyncio +async def test_create_or_get_fetches_existing_store_on_conflict() -> None: + store = _make_store_with_responses( + _make_response(409, {"error": {"message": "duplicate store"}}), + _make_response( + 200, + { + "id": "ss_1", + "object": "statestore", + "name": "checkpoints", + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": "existing", + "tags": {"env": "dev"}, + "created_at": 1, + "updated_at": 2, + }, + ), + name="checkpoints", + ) + + result = await store.create_or_get() + + first_request = store._client.send_request.call_args_list[0][0][0] + second_request = store._client.send_request.call_args_list[1][0][0] + assert first_request.method == "POST" + assert second_request.method == "GET" + assert second_request.url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}?api-version=v1" + assert result.description == "existing" + assert result.tags == {"env": "dev"} + + +@pytest.mark.asyncio +async def test_get_or_create_returns_existing_store_when_present() -> None: + store = _make_store( + _make_response( + 200, + { + "id": "ss_1", + "object": "statestore", + "name": "checkpoints", + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": "existing", + "tags": {"env": "dev"}, + "created_at": 1, + "updated_at": 2, + }, + ), + name="checkpoints", + ) + + result = await store.get_or_create() + + request = _sent_request(store) + assert request.method == "GET" + assert request.url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}?api-version=v1" + assert result.description == "existing" + assert result.tags == {"env": "dev"} + + +@pytest.mark.asyncio +async def test_get_or_create_creates_store_when_absent() -> None: + store = _make_store_with_responses( + _make_response(404, {"error": {"message": "not found"}}), + _make_response( + 201, + { + "id": "ss_1", + "object": "statestore", + "name": "checkpoints", + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": None, + "tags": {}, + "created_at": 1, + "updated_at": 1, + }, + ), + name="checkpoints", + ) + + result = await store.get_or_create() + + first_request = store._client.send_request.call_args_list[0][0][0] + second_request = store._client.send_request.call_args_list[1][0][0] + assert first_request.method == "GET" + assert second_request.method == "POST" + assert second_request.url == f"{_BASE_URL}statestores?api-version=v1" + assert result.name == "checkpoints" + + +@pytest.mark.asyncio +async def test_get_or_create_fetches_store_when_create_races_with_another_caller() -> None: + store = _make_store_with_responses( + _make_response(404, {"error": {"message": "not found"}}), + _make_response(409, {"error": {"message": "duplicate store"}}), + _make_response( + 200, + { + "id": "ss_1", + "object": "statestore", + "name": "checkpoints", + "user_isolation": False, + "item_ttl_seconds": 2592000, + "description": "created elsewhere", + "tags": {"env": "dev"}, + "created_at": 1, + "updated_at": 2, + }, + ), + name="checkpoints", + ) + + result = await store.get_or_create() + + requests = [call_args[0][0] for call_args in store._client.send_request.call_args_list] + assert [request.method for request in requests] == ["GET", "POST", "GET"] + assert requests[2].url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}?api-version=v1" + assert result.description == "created elsewhere" + + +@pytest.mark.asyncio +async def test_get_properties_uses_base64url_store_name() -> None: + store_name = "langGraphCheckpoints/thread-abc" + store = _make_store( + _make_response( + 200, + { + "id": "ss_1", + "object": "statestore", + "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_properties() + + request = _sent_request(store) + assert request.method == "GET" + assert request.url == f"{_BASE_URL}statestores/{_encode_segment(store_name)}?api-version=v1" + assert result.name == store_name + assert result.id == "ss_1" + + +@pytest.mark.asyncio +async def test_update_metadata_sends_only_present_fields() -> None: + store = _make_store( + _make_response( + 200, + { + "id": "ss_1", + "object": "statestore", + "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_metadata(description="updated", tags={"env": "prod"}) + + request = _sent_request(store) + assert request.method == "PATCH" + assert request.url == f"{_BASE_URL}statestores/{_encode_segment('prefs')}?api-version=v1" + assert json.loads(request.content.decode("utf-8")) == {"description": "updated", "tags": {"env": "prod"}} + assert result.updated_at == 3 + + +@pytest.mark.asyncio +async def test_delete_store_returns_deleted_marker() -> None: + store = _make_store( + _make_response( + 200, + {"id": "ss_1", "object": "statestore.deleted", "name": "prefs", "deleted": True}, + ), + name="prefs", + ) + + result = await store.delete_store() + + request = _sent_request(store) + assert request.method == "DELETE" + assert request.url == f"{_BASE_URL}statestores/{_encode_segment('prefs')}?api-version=v1" + assert result == DeletedStateStore(id="ss_1", name="prefs", deleted=True) + + +@pytest.mark.asyncio +async def test_create_item_posts_key_value_and_tags() -> None: + store = _make_store( + _make_response( + 201, + { + "id": "it_1", + "object": "statestore_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}statestores/{_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 == StateItemMetadata( + id="it_1", + key="step/1", + 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": "statestore_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}statestores/{_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": "statestore_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_get_returns_state_item_with_value_and_metadata() -> None: + store = _make_store( + _make_response( + 200, + { + "id": "it_1", + "object": "statestore_item", + "key": "step/1", + "value": {"done": True}, + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + }, + ), + name="checkpoints", + user_id="user-42", + ) + + result = await store.get("step/1") + + request = _sent_request(store) + assert request.method == "GET" + assert request.headers["x-ms-user-id"] == "user-42" + assert request.url == ( + f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}/items/{_encode_segment('step/1')}?api-version=v1" + ) + assert result == StateItem( + id="it_1", + key="step/1", + value={"done": True}, + tags={"kind": "checkpoint"}, + etag='"0x8DD"', + created_at=10, + updated_at=20, + ) + + +@pytest.mark.asyncio +async def test_get_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 + + +@pytest.mark.asyncio +async def test_delete_item_returns_deleted_marker() -> None: + store = _make_store( + _make_response( + 200, + {"id": "it_1", "object": "statestore_item.deleted", "key": "step/1", "deleted": True}, + ), + name="checkpoints", + user_id="user-42", + ) + + result = await store.delete("step/1", if_match='"0x8DD"') + + request = _sent_request(store) + assert request.method == "DELETE" + assert request.headers["If-Match"] == '"0x8DD"' + assert request.headers["x-ms-user-id"] == "user-42" + assert result == DeletedStateItem(id="it_1", key="step/1", deleted=True) + + +@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": "statestore_item", + "key": "step/1", + "tags": {"kind": "checkpoint"}, + "etag": '"0x8DD"', + "created_at": 10, + "updated_at": 20, + } + ], + "first_id": "it_1", + "last_id": "it_1", + "has_more": False, + }, + ), + name="checkpoints", + user_id="user-42", + ) + + page = await store.list_keys(tags={"kind": "checkpoint", "phase": "run"}, limit=10, after="it_0", order="asc") + + request = _sent_request(store) + assert request.method == "GET" + assert request.headers["x-ms-user-id"] == "user-42" + assert request.url == ( + f"{_BASE_URL}statestores/{_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=[ + StateKey( + id="it_1", + 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, + ) + + +@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}statestores/{_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..c6d986d209b0 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_state_store_sample_usage.py @@ -0,0 +1,241 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. +"""Programmatic sample-flow coverage for the Foundry state-store sample.""" + +from __future__ import annotations + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest +from azure.ai.agentserver.core.storage import ( + FoundryStateStore, + FoundryStorageEndpoint, + FoundryStoragePreconditionError, +) + +_BASE_URL = "https://foundry.example.com/storage/" +_ENDPOINT = FoundryStorageEndpoint(storage_base_url=_BASE_URL) + + +def _make_response(status_code: int, body: object, *, headers: dict[str, str] | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = status_code + resp.headers = {} if headers is None else headers + resp.text = MagicMock(return_value=json.dumps(body)) + return resp + + +def _make_store_with_responses(*responses: MagicMock) -> FoundryStateStore: + store = FoundryStateStore.__new__(FoundryStateStore) + store._endpoint = _ENDPOINT + store._owns_credential = False + store._name = "checkpoints/thread-abc" + store._user_isolation = True + store._item_ttl_seconds = 3600 + store._description = "Sample state store" + store._tags = {} + store._user_id = "user-42" + mock_pipeline = AsyncMock() + mock_pipeline.send_request = AsyncMock(side_effect=list(responses)) + mock_pipeline.close = AsyncMock() + store._client = mock_pipeline + return store + + +@pytest.mark.asyncio +async def test_state_store_sample_flow() -> None: + store = _make_store_with_responses( + _make_response( + 201, + { + "id": "ss_1", + "object": "statestore", + "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": "statestore_item", + "key": "step-1", + "etag": '"0x8DA"', + "created_at": 2, + "updated_at": 2, + }, + ), + _make_response( + 200, + { + "id": "it_1", + "object": "statestore_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": "statestore", + "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": "statestore_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": "statestore_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": "statestore_item", + "key": "step-2", + "etag": '"0x8DC"', + "created_at": 5, + "updated_at": 5, + }, + ), + _make_response( + 201, + { + "id": "it_3", + "object": "statestore_item", + "key": "audit-1", + "etag": '"0x8DD"', + "created_at": 6, + "updated_at": 6, + }, + ), + _make_response( + 200, + { + "object": "list", + "data": [ + { + "id": "it_1", + "object": "statestore_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": "statestore_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": "statestore_item.deleted", "key": "audit-1", "deleted": True}, + ), + _make_response( + 200, + {"id": "ss_1", "object": "statestore.deleted", "name": "checkpoints/thread-abc", "deleted": True}, + ), + ) + + store_info = await store.get_or_create() + 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_metadata( + description="Sample checkpoint store", + tags={"scenario": "state-store-sample", "env": "dev"}, + ) + 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_store() + 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..db9aa9d32cbf --- /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("statestores") == "https://x/storage/statestores?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("statestores", after="it 1/2") + assert url == "https://x/storage/statestores?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..4117d021e80b --- /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/statestores/abc/items:keys?api-version=v1&after=it_1" + masked = _mask_storage_url(url) + assert masked == "***/storage/statestores/*/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" From 9937d968a7d37cd68626a882557a341fca0a0003 Mon Sep 17 00:00:00 2001 From: Shanmukha Pasumarthy Date: Fri, 10 Jul 2026 13:21:31 +0530 Subject: [PATCH 2/2] fix(agentserver-core): align FoundryStateStore route with latest state-storage spec (foundrysdk_specs#247) Per the latest commit on coreai-microsoft/foundrysdk_specs#247 ("rename route to /storage/state_stores, add store PATCH update, align object descriptors"), the state-store REST path is /storage/state_stores/* (snake_case with underscore), not /storage/statestores/*. - _state.py: store path + create() now target state_stores. - _policies.py: masked-logging allowlist updated to the new segment name. - Updated docs/state-store-guide.md, README, and test URL assertions. - Fixed the core CHANGELOG/README, which still described the earlier namespace-based design (pre-dating the PR #47763 pull) instead of the current store-bound statestores-protocol API; also dropped an unused aiohttp dependency left over from that earlier design. No functional change beyond the route rename -- the object-type descriptor changes in the spec commit (state_store / state_store.item) are response-only fields the SDK does not parse or assert on. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../azure-ai-agentserver-core/CHANGELOG.md | 2 +- .../azure-ai-agentserver-core/README.md | 23 +++++++++------ .../ai/agentserver/core/storage/_policies.py | 2 +- .../ai/agentserver/core/storage/_state.py | 4 +-- .../docs/state-store-guide.md | 2 +- .../azure-ai-agentserver-core/pyproject.toml | 1 - .../tests/test_foundry_state_store.py | 28 +++++++++---------- .../tests/test_storage_endpoint.py | 6 ++-- .../tests/test_storage_policies.py | 4 +-- 9 files changed, 38 insertions(+), 34 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md index f66efd712630..8a2fd9186073 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md +++ b/sdk/agentserver/azure-ai-agentserver-core/CHANGELOG.md @@ -4,7 +4,7 @@ ### Features Added -- Added `azure.ai.agentserver.core.storage` package providing the protocol-neutral Foundry storage layer: `FoundryStorageClient` (transport, endpoint, pipeline policies, error hierarchy) and `FoundryStateStore`, a generic durable key-value store. Every operation is scoped to exactly one caller-supplied `namespace`, carried as a percent-encoded URL path segment (`POST /storage/namespaces/{namespace}/state:read|:write|:listKeys`); server-side isolation is selected by the trusted `x-ms-internal-state-session-isolation` / `x-ms-internal-state-user-isolation` headers (`isolate_sessions` / `isolate_users` knobs; disabling both is rejected). Batch `write` is atomic and reports per-key `WriteResult` metadata (etag + `created_at` / `updated_at`); `read` and `list_keys` return the same server-managed timestamps. Supports optional `if_match` optimistic concurrency, optional per-item `ttl_seconds` (surfaced as `StateItem.expires_at`), and ordered, paged `list_keys` (keys + metadata only). `FoundryStateStore.for_namespace(...)` returns a namespace-bound `NamespaceStateStore` handle. Writes use the strongly-typed `Upsert` / `Delete` change objects (`WriteChange`) and stored values are typed as `JSONValue`. Protocol packages can build resource-specific clients on top of `FoundryStorageClient`. 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`. +- 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. Stores are explicit resources with a lifecycle (`create`, `get_or_create`, `create_or_get`, `get_properties`, `update_metadata`, `delete_store`) plus single-item operations (`create_item`, `set`, `get`, `delete`, `list_keys`). Store-level `item_ttl_seconds` is fixed at create time (default 30 days, `-1` never expires); optional `user_isolation` partitions items per resolved (or delegated `user_id`) caller. `if_match` gives single-item optimistic concurrency; `list_keys` is ordered and paged. 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) diff --git a/sdk/agentserver/azure-ai-agentserver-core/README.md b/sdk/agentserver/azure-ai-agentserver-core/README.md index b4ee137d8551..04feb4348f88 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/README.md +++ b/sdk/agentserver/azure-ai-agentserver-core/README.md @@ -129,22 +129,27 @@ async def on_shutdown(): ### Durable state storage `FoundryStateStore` is a durable, server-backed key-value store for agent state -— session memory, per-user preferences, counters, and checkpoints — with -optimistic concurrency, tag filtering, key listing, and optional per-item TTL. +— session memory, per-user preferences, counters, and checkpoints — bound to +one explicit, caller-named store, with an explicit store lifecycle, 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. -async with FoundryStateStore() as store: - etag = await store.set("counters", "page-views", 1) - item = await store.get("counters", "page-views") - print(item.value) # 1 +# The store name is the scope -- encode conversation/thread identity into it. +async with FoundryStateStore("checkpoints/thread-abc", user_isolation=True) as store: + await store.get_or_create() + await store.set("step-1", {"done": False}) + item = await store.get("step-1") + print(item.value) # {"done": False} ``` -Reads return typed `StateItem` values; writes are expressed as typed `Upsert` / -`Delete` changes. 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 concurrency model, and common gotchas, and +Reads return typed `StateItem` 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. 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 index 030c79bf71cb..f592933183a4 100644 --- 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 @@ -82,7 +82,7 @@ def _redact_storage_path(path: str) -> str: segments = path.split("/") redacted: list[str] = [] for index, segment in enumerate(segments): - if index < 3 or segment in {"items", "items:keys", "statestores"} or not segment: + if index < 3 or segment in {"items", "items:keys", "state_stores"} or not segment: redacted.append(segment) else: redacted.append("*") 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 index a56fe4a5e558..1f6d588ef756 100644 --- 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 @@ -144,7 +144,7 @@ def name(self) -> str: return self._name def _store_path(self) -> str: - return f"statestores/{_encode_segment(self._name)}" + return f"state_stores/{_encode_segment(self._name)}" def _item_path(self, key: str) -> str: _validate_key(key) @@ -178,7 +178,7 @@ async def create(self) -> StateStoreInfo: description=self._description, tags=self._tags, ) - response = await self._send_storage_request(self._request("POST", "statestores", content=body)) + response = await self._send_storage_request(self._request("POST", "state_stores", content=body)) return deserialize_state_store(response.text()) async def create_or_get(self) -> StateStoreInfo: 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 index a82f3c80dc1b..8e793f216b96 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md +++ b/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md @@ -16,7 +16,7 @@ name is the main scoping tool for your data: - Set `item_ttl_seconds` once at store creation when you want idle items to age out automatically. -The SDK is the developer-facing layer over the internal `/storage/statestores/*` +The SDK is the developer-facing layer over the internal `/storage/state_stores/*` protocol: it keeps the transport/auth pipeline in `FoundryStorageClient`, while `FoundryStateStore` owns the ergonomic store-bound API. diff --git a/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml b/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml index 73e2e508142c..69f1df4599ee 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml +++ b/sdk/agentserver/azure-ai-agentserver-core/pyproject.toml @@ -27,7 +27,6 @@ dependencies = [ "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/tests/test_foundry_state_store.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_foundry_state_store.py index 94c1a08f1200..b0e8f3d27ea8 100644 --- 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 @@ -102,7 +102,7 @@ async def test_create_posts_store_descriptor() -> None: request = _sent_request(store) assert request.method == "POST" - assert request.url == f"{_BASE_URL}statestores?api-version=v1" + assert request.url == f"{_BASE_URL}state_stores?api-version=v1" assert request.headers["Content-Type"] == "application/json; charset=utf-8" assert "x-ms-user-id" not in request.headers assert json.loads(request.content.decode("utf-8")) == { @@ -148,7 +148,7 @@ async def test_create_or_get_returns_created_store_when_absent() -> None: request = _sent_request(store) assert request.method == "POST" - assert request.url == f"{_BASE_URL}statestores?api-version=v1" + assert request.url == f"{_BASE_URL}state_stores?api-version=v1" assert result.name == "checkpoints" @@ -179,7 +179,7 @@ async def test_create_or_get_fetches_existing_store_on_conflict() -> None: second_request = store._client.send_request.call_args_list[1][0][0] assert first_request.method == "POST" assert second_request.method == "GET" - assert second_request.url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}?api-version=v1" + assert second_request.url == f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}?api-version=v1" assert result.description == "existing" assert result.tags == {"env": "dev"} @@ -208,7 +208,7 @@ async def test_get_or_create_returns_existing_store_when_present() -> None: request = _sent_request(store) assert request.method == "GET" - assert request.url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}?api-version=v1" + assert request.url == f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}?api-version=v1" assert result.description == "existing" assert result.tags == {"env": "dev"} @@ -240,7 +240,7 @@ async def test_get_or_create_creates_store_when_absent() -> None: second_request = store._client.send_request.call_args_list[1][0][0] assert first_request.method == "GET" assert second_request.method == "POST" - assert second_request.url == f"{_BASE_URL}statestores?api-version=v1" + assert second_request.url == f"{_BASE_URL}state_stores?api-version=v1" assert result.name == "checkpoints" @@ -270,7 +270,7 @@ async def test_get_or_create_fetches_store_when_create_races_with_another_caller requests = [call_args[0][0] for call_args in store._client.send_request.call_args_list] assert [request.method for request in requests] == ["GET", "POST", "GET"] - assert requests[2].url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}?api-version=v1" + assert requests[2].url == f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}?api-version=v1" assert result.description == "created elsewhere" @@ -299,7 +299,7 @@ async def test_get_properties_uses_base64url_store_name() -> None: request = _sent_request(store) assert request.method == "GET" - assert request.url == f"{_BASE_URL}statestores/{_encode_segment(store_name)}?api-version=v1" + assert request.url == f"{_BASE_URL}state_stores/{_encode_segment(store_name)}?api-version=v1" assert result.name == store_name assert result.id == "ss_1" @@ -328,7 +328,7 @@ async def test_update_metadata_sends_only_present_fields() -> None: request = _sent_request(store) assert request.method == "PATCH" - assert request.url == f"{_BASE_URL}statestores/{_encode_segment('prefs')}?api-version=v1" + 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 @@ -347,7 +347,7 @@ async def test_delete_store_returns_deleted_marker() -> None: request = _sent_request(store) assert request.method == "DELETE" - assert request.url == f"{_BASE_URL}statestores/{_encode_segment('prefs')}?api-version=v1" + assert request.url == f"{_BASE_URL}state_stores/{_encode_segment('prefs')}?api-version=v1" assert result == DeletedStateStore(id="ss_1", name="prefs", deleted=True) @@ -372,7 +372,7 @@ async def test_create_item_posts_key_value_and_tags() -> None: request = _sent_request(store) assert request.method == "POST" - assert request.url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}/items?api-version=v1" + 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}, @@ -411,7 +411,7 @@ async def test_set_puts_value_and_if_match_header() -> None: request = _sent_request(store) assert request.method == "PUT" assert request.url == ( - f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}/items/{_encode_segment('step/1')}?api-version=v1" + 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"}} @@ -467,7 +467,7 @@ async def test_get_returns_state_item_with_value_and_metadata() -> None: assert request.method == "GET" assert request.headers["x-ms-user-id"] == "user-42" assert request.url == ( - f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}/items/{_encode_segment('step/1')}?api-version=v1" + f"{_BASE_URL}state_stores/{_encode_segment('checkpoints')}/items/{_encode_segment('step/1')}?api-version=v1" ) assert result == StateItem( id="it_1", @@ -539,7 +539,7 @@ async def test_list_keys_uses_query_parameters_and_returns_page() -> None: assert request.method == "GET" assert request.headers["x-ms-user-id"] == "user-42" assert request.url == ( - f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}/items:keys" + 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( @@ -566,7 +566,7 @@ async def test_list_keys_defaults_to_desc_order() -> None: await store.list_keys() request = _sent_request(store) - assert request.url == f"{_BASE_URL}statestores/{_encode_segment('checkpoints')}/items:keys?api-version=v1&order=desc" + 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: 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 index db9aa9d32cbf..fc41756d4fb8 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_endpoint.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_endpoint.py @@ -48,10 +48,10 @@ def test_from_env_requires_variable(monkeypatch: pytest.MonkeyPatch) -> None: def test_build_url_appends_api_version() -> None: ep = FoundryStorageEndpoint(storage_base_url="https://x/storage/", api_version="v1") - assert ep.build_url("statestores") == "https://x/storage/statestores?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("statestores", after="it 1/2") - assert url == "https://x/storage/statestores?api-version=v1&after=it%201%2F2" + 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_policies.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py index 4117d021e80b..db466422ca93 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_storage_policies.py @@ -13,9 +13,9 @@ def test_mask_keeps_only_storage_path_and_api_version() -> None: - url = "https://proj.example.com/api/projects/secret/storage/statestores/abc/items:keys?api-version=v1&after=it_1" + 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/statestores/*/items:keys?api-version=v1" + 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