From 0cb3c63a137faad78c24403eb003b489a0042e0a Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Mon, 21 Sep 2026 15:07:58 +0200 Subject: [PATCH 1/5] Python: persist invocation sessions through configured storage --- ...2-python-invocation-session-persistence.md | 79 +++ python/packages/foundry_hosting/README.md | 51 +- .../_invocations.py | 151 ++++- .../packages/foundry_hosting/pyproject.toml | 1 + .../foundry_hosting/tests/test_invocations.py | 594 +++++++++++++++++- python/uv.lock | 2 + 6 files changed, 818 insertions(+), 60 deletions(-) create mode 100644 docs/decisions/0042-python-invocation-session-persistence.md diff --git a/docs/decisions/0042-python-invocation-session-persistence.md b/docs/decisions/0042-python-invocation-session-persistence.md new file mode 100644 index 00000000000..ba8f1dbb8df --- /dev/null +++ b/docs/decisions/0042-python-invocation-session-persistence.md @@ -0,0 +1,79 @@ +--- +status: proposed +contact: jpalvarezl +date: 2026-09-21 +deciders: jpalvarezl +--- + +# Persist Python invocation sessions through the existing session-store abstraction + +## Context and Problem Statement + +`InvocationsHostServer` keeps conversation sessions in a host-owned dictionary for the +process lifetime. Request-scoped agent factories do not change that ownership: even after +an agent is disposed, the host still owns the session. We need to release completed +request state without losing multi-turn conversation continuity. + +## Decision Drivers + +- Restore conversation state rather than silently evicting live sessions. +- Preserve hosted user/session partition boundaries and request-scoped agent ownership. +- Reuse existing storage and serialization contracts instead of adding another cache policy. +- Keep streaming, failure, and cancellation cleanup inside the session's coordination scope. +- Make persistence and retention behavior explicit to applications. + +## Considered Options + +- **Cap the in-memory dictionary and refuse new conversations.** Bounds retained entry + count, but needs a new capacity policy and eventually refuses new work until an + application deletes state. It does not bound the size of each session. +- **Evict idle sessions.** Releases memory but loses conversation state unless backed by + persistent storage. Adds a separate host retention policy and active-session tracking. +- **Use the existing session-store provider.** Releases host-owned session objects while + preserving restorable state. Adds storage I/O and requires serializable state, but + reuses the established hosting extension point and backend retention policy. + +## Decision Outcome + +Choose the existing session-store provider. `InvocationsHostServer` accepts +`agent_session_store_provider`, defaulting to `AgentSessionStoreProvider`, like Python +`ResponsesHostServer`. The default is Foundry storage when hosted and the SDK's file-backed +store locally. New default stores inherit the existing SDK retention policy: 30 days +since the last write. Existing stores and custom providers retain their own policies. +Absent or expired state starts a new conversation; there is no permanent tombstone registry. + +### Alignment with .NET + +This follows **.NET Foundry Responses hosting's persistence architecture**, not a purported +.NET Invocations cache limit: + +- [`AddFoundryResponses`](../../dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs) + defaults to `FoundryAgentSessionStore`. +- [`.NET FoundryAgentSessionStore`](../../dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs) + serializes sessions into the AgentServer state store, with a local fallback and the + platform's 30-day default item retention. +- The .NET in-memory session store remains a development/testing alternative; it does not + supply an eviction policy. The Invocations echo sample is stateless. + +Python reuses its own existing `StoreProvider[SessionStore]` contract. This is architectural +alignment, not shared storage-key formats, cross-language state compatibility, or identical +concurrency implementations. + +### Lifecycle and limitations + +Invocation storage keys hash an unambiguous partition encoding and use a versioned, +protocol-specific namespace; runtime `AgentSession.session_id` values remain unchanged. +Same-partition requests serialize within one host. Reference-counted lock entries exist +only for active/waiting requests and are removed after the final request leaves. + +The coordination scope includes stream closure, factory-agent cleanup, and session saving. +The host attempts to save state on success, execution failure, and interruption, without +retrying agent execution. Disconnect cleanup is shielded from the ASGI cancellation scope. +Persistence errors are reported; there is no success-shaped in-memory fallback. + +This removes cumulative host retention of completed session objects. It does not bound +individual session size, active request count, backend storage, or total process memory. +The SDK local backend reads its entire logical-store file per operation. The store API +does not provide distributed transactions, so concurrent updates across hosts still +require application-level coordination. Durable serialization also replaces Python +object-identity reuse with restored state, and custom state needs registered codecs. diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index a7a59803df4..7670f74c2c3 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -103,9 +103,9 @@ No additional partitioning configuration is required when using the default stor ### Agent Sessions -`ResponsesHostServer` persists the Agent Framework `AgentSession` durably. By default it -uses the `FoundryAgentSessionStore`, backed by Foundry storage when hosted and file-based -storage locally. Stored sessions are scoped under `agent_sessions`. +`ResponsesHostServer` and `InvocationsHostServer` persist the Agent Framework `AgentSession` +durably. By default they use `FoundryAgentSessionStore`, backed by Foundry storage when hosted +and file-based storage locally. Stored sessions are scoped under `agent_sessions`. See the [custom storage provider sample](../../samples/04-hosting/foundry-hosted-agents/responses/custom_storage/) for an example that uses an in-memory session store locally and Azure Cosmos DB when hosted. @@ -114,11 +114,50 @@ Native Responses refusal parts are stored as text carrying `additional_properties["model_output_kind"] == "refusal"` and emitted as `response.refusal.*` events when streamed back to clients. -`InvocationsHostServer` keeps sessions in memory. When hosted, `AgentSession.session_id` is an +#### Invocation sessions + +`InvocationsHostServer` loads session state for each request and saves it after agent and +stream cleanup, including when execution fails or is interrupted. It does not retain completed +sessions in a host-level cache. Storage errors are reported rather than silently starting a +new conversation or falling back to memory. After streaming headers have been sent, a +persistence failure terminates the stream and is logged; already delivered text cannot be +retracted. A failed request is not automatically retried. + +When hosted, `AgentSession.session_id` is an opaque composite identifier that preserves the boundaries between the platform session ID and user ID. Consumers must use it as a whole and must not parse it or depend on its internal -representation. Repeated requests for the same identifier pair reuse the session. Locally, -the platform session ID is used unchanged. +representation. Repeated requests for the same identifier pair restore the saved session. +Locally, the platform session ID is used unchanged. Storage keys are opaque, fixed-length, +and namespaced separately from Responses keys. + +New default stores expire saved sessions **30 days after the last write**. Each write renews +that period; a read does not. Existing stores keep their creation-time retention settings. +A missing, deleted, or expired session starts a fresh conversation on the next request, even +if the caller reuses its ID. A running request keeps its working session until cleanup and +does not evict other conversations. + +To select another store, pass `agent_session_store_provider`, implementing +`StoreProvider[SessionStore]`, to the constructor. Its `get_store` method receives the host +configuration and the request's platform context, as it does for `ResponsesHostServer`. +The provider owns storage retention and deletion. An explicitly supplied in-memory store is +still volatile and does not gain automatic eviction. + +**Migration from in-memory invocation sessions:** session state must now support +`AgentSession.to_dict()` / `AgentSession.from_dict()` serialization. Register codecs for +custom state types instead of storing arbitrary live Python objects. Requests restore +equivalent state, not the same Python object. Existing process-local sessions are not +migrated across deployment; the previous implementation also lost them on restart. +Local sessions now survive a host restart and are written under `AGENTSERVER_STATE_ROOT` +(or the default location above). Use separate storage roots or providers for independent +applications. + +Requests for the same session are serialized within one host through execution, streaming, +cleanup, and persistence. Different sessions can run concurrently. There is no cross-host +transaction or exactly-once execution guarantee; deployments with overlapping same-session +requests across workers must coordinate those requests. Removing the host cache does not +impose a byte-memory limit, request-admission limit, or storage quota. The SDK's local backend +reads a whole logical-store file per operation, so its transient memory and I/O can still +grow with stored data. ### Workflow checkpoints diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index 0cf7a6e97e0..abbd7f4f6bf 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,19 +1,54 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + +import asyncio +import hashlib import json +import logging from collections.abc import Awaitable, Callable from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager +from dataclasses import dataclass, field -from agent_framework import AgentSession, ResponseStream, SupportsAgentRun +from agent_framework import AgentSession, ResponseStream, SessionStore, SupportsAgentRun from agent_framework._telemetry import mark_feature_used -from azure.ai.agentserver.core import get_request_context +from anyio import CancelScope +from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context from azure.ai.agentserver.invocations import InvocationAgentServerHost from starlette.requests import Request from starlette.responses import Response, StreamingResponse +from starlette.types import Receive, Scope, Send from typing_extensions import Any, AsyncGenerator from ._agent_source import is_agent, resolve_agent, validate_agent_source from ._feature_usage import FeatureIndex +from ._state_store import AgentSessionStoreProvider, StoreProvider + +logger = logging.getLogger(__name__) + + +@dataclass +class _SessionLock: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + users: int = 0 + + +class _InvocationStreamingResponse(StreamingResponse): + def __init__(self, content: AsyncGenerator[str]) -> None: + super().__init__( + content, + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, + ) + self._content = content + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + try: + await super().__call__(scope, receive, send) + finally: + # Sending a chunk can fail while the generator is suspended at yield. + with CancelScope(shield=True): + await self._content.aclose() class InvocationsHostServer(InvocationAgentServerHost): @@ -24,6 +59,7 @@ def __init__( agent: SupportsAgentRun | Callable[[], SupportsAgentRun | Awaitable[SupportsAgentRun]], *, openapi_spec: dict[str, Any] | None = None, + agent_session_store_provider: StoreProvider[SessionStore] | None = None, **kwargs: Any, ) -> None: """Initialize an InvocationsHostServer. @@ -32,18 +68,23 @@ def __init__( agent: The agent to handle responses for, or a zero-argument sync or async callable that creates one for each request. Use a callable for agents that keep mutable state outside `AgentSession`. openapi_spec: The OpenAPI specification for the server. + agent_session_store_provider: Provider for conversation session storage. Defaults to Foundry storage + when hosted and the SDK's file-backed storage locally. New default stores expire sessions 30 days + after their last write. Custom providers control their own retention. **kwargs: Additional keyword arguments. This host will expect the request to be a JSON body with a "message" field. - The response from the host will be a JSON object with a "response" field containing - the agent's response and a "session_id" field containing the session ID. + The response contains the agent's text, or streamed text when "stream" is true. """ validate_agent_source(agent) super().__init__(openapi_spec=openapi_spec, **kwargs) self._agent = agent self._owns_request_agent = not is_agent(agent) - self._sessions: dict[str | tuple[str, str], AgentSession] = {} + self._session_store_provider = ( + AgentSessionStoreProvider() if agent_session_store_provider is None else agent_session_store_provider + ) + self._session_locks: dict[str | tuple[str, str], _SessionLock] = {} self.invoke_handler(self._handle_invoke) mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) @@ -81,10 +122,67 @@ def _partition_key(self) -> str | tuple[str, str]: @asynccontextmanager async def _request_agent(self) -> AsyncGenerator[SupportsAgentRun]: agent = await resolve_agent(self._agent) - async with AsyncExitStack() as resources: + resources = AsyncExitStack() + try: if self._owns_request_agent and isinstance(agent, AbstractAsyncContextManager): await resources.enter_async_context(agent) yield agent + except BaseException as exc: + with CancelScope(shield=True): + if not await resources.__aexit__(type(exc), exc, exc.__traceback__): + raise + else: + with CancelScope(shield=True): + await resources.aclose() + + @asynccontextmanager + async def _request_session( + self, partition_key: str | tuple[str, str], context: FoundryAgentRequestContext + ) -> AsyncGenerator[AgentSession]: + entry = self._session_locks.get(partition_key) + if entry is None: + entry = self._session_locks[partition_key] = _SessionLock() + entry.users += 1 + try: + async with entry.lock: + encoded_key = json.dumps(partition_key, separators=(",", ":")) + session_id = encoded_key if isinstance(partition_key, tuple) else partition_key + storage_key = "invocations:v1:" + hashlib.sha256(encoded_key.encode("utf-8")).hexdigest() + try: + store = self._session_store_provider.get_store(config=self.config, platform_context=context) + session = await store.get(storage_key) + if session is None: + session = AgentSession(session_id=session_id) + elif session.session_id != session_id: + raise ValueError("Stored invocation session does not match the requested session.") + except Exception: + logger.exception("Failed to load invocation session") + raise + + failure: BaseException | None = None + try: + yield session + except BaseException as exc: + failure = exc + raise + finally: + # Starlette cancels the streaming task's scope on disconnect. + with CancelScope(shield=True): + try: + await store.set(storage_key, session) + except Exception as exc: + logger.exception("Failed to persist invocation session") + if failure is None: + raise + if isinstance(failure, Exception): + raise RuntimeError("Invocation failed and session persistence also failed.") from exc + except asyncio.CancelledError: + logger.error("Invocation session persistence was cancelled") + raise + finally: + entry.users -= 1 + if entry.users == 0: + del self._session_locks[partition_key] async def _handle_invoke(self, request: Request) -> Response: """Invoke the agent with the given request.""" @@ -103,37 +201,34 @@ async def _handle_invoke(self, request: Request) -> Response: return StreamingResponse(content=error, status_code=400) return Response(content=error, status_code=400) - session = self._sessions.get(partition_key) - if session is None: - session_id = ( - json.dumps(partition_key, separators=(",", ":")) if isinstance(partition_key, tuple) else partition_key - ) - session = AgentSession(session_id=session_id) - self._sessions[partition_key] = session + context = get_request_context() if stream: async def stream_response() -> AsyncGenerator[str]: - async with self._request_agent() as agent: + async with ( + self._request_session(partition_key, context) as session, + self._request_agent() as agent, + ): stream = agent.run(user_message, session=session, stream=True) try: async for update in stream: if update.text: yield update.text finally: - if isinstance(stream, ResponseStream): - await stream.close() - else: - close = getattr(stream, "aclose", None) - if close is not None: - await close() - - return StreamingResponse( - stream_response(), - media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, - ) - - async with self._request_agent() as agent: + with CancelScope(shield=True): + if isinstance(stream, ResponseStream): + await stream.close() + else: + close = getattr(stream, "aclose", None) + if close is not None: + await close() + + return _InvocationStreamingResponse(stream_response()) + + async with ( + self._request_session(partition_key, context) as session, + self._request_agent() as agent, + ): response = await agent.run([user_message], session=session) return Response(content=response.text) diff --git a/python/packages/foundry_hosting/pyproject.toml b/python/packages/foundry_hosting/pyproject.toml index 3dd6acefa79..f63389f9bac 100644 --- a/python/packages/foundry_hosting/pyproject.toml +++ b/python/packages/foundry_hosting/pyproject.toml @@ -24,6 +24,7 @@ classifiers = [ ] dependencies = [ "agent-framework-core>=1.19.0,<2", + "anyio>=4,<5", "azure-ai-agentserver-core>=2.1.0,<3", "azure-ai-agentserver-responses>=2.2.0b1,<3", "azure-ai-agentserver-invocations>=1.1.0,<2", diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index 250799849a6..f88af3f6c6e 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -11,36 +11,69 @@ from __future__ import annotations +import asyncio import json -from collections.abc import AsyncIterator, Iterator +import weakref +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from itertools import product from typing import cast from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from agent_framework import ( + Agent, AgentResponse, AgentResponseUpdate, AgentSession, + BaseChatClient, + ChatMiddlewareLayer, + ChatResponse, + ChatResponseUpdate, Content, + FunctionInvocationLayer, + InMemoryHistoryProvider, Message, + ResponseStream, ServiceSessionId, + SessionStore, + tool, ) +from anyio import CancelScope from azure.ai.agentserver.core import ( + AgentConfig, FoundryAgentRequestContext, reset_request_context, set_request_context, ) -from starlette.requests import Request +from starlette.requests import ClientDisconnect, Request from starlette.responses import Response, StreamingResponse from typing_extensions import Any -from agent_framework_foundry_hosting import InvocationsHostServer +from agent_framework_foundry_hosting import InvocationsHostServer, StoreProvider + +pytestmark = pytest.mark.filterwarnings("ignore:.*SessionStore is experimental.*") # region Helpers +def _mock_session_store() -> MagicMock: + store = MagicMock(spec=SessionStore) + store.get.return_value = None + return store + + +class _SessionStoreProvider(StoreProvider[SessionStore]): + def __init__(self, store: SessionStore) -> None: + self.store = store + self.contexts: list[FoundryAgentRequestContext] = [] + + def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentRequestContext) -> SessionStore: + self.contexts.append(platform_context) + return self.store + + class _FakeAgent: """Minimal agent implementing the ``SupportsAgentRun`` protocol. @@ -54,6 +87,7 @@ def __init__( *, response: AgentResponse | None = None, stream_updates: list[AgentResponseUpdate] | None = None, + update_session: Callable[[AgentSession], None] | None = None, ) -> None: self.id = "fake-agent" self.name: str | None = "Fake Agent" @@ -61,6 +95,7 @@ def __init__( self._response = response self._stream_updates = stream_updates or [] self.calls: list[dict[str, Any]] = [] + self._update_session = update_session def run( self, @@ -71,6 +106,8 @@ def run( **kwargs: Any, ) -> Any: self.calls.append({"messages": messages, "stream": stream, "session": session}) + if session is not None and self._update_session is not None: + self._update_session(session) if stream: async def _gen() -> AsyncIterator[AgentResponseUpdate]: @@ -142,6 +179,7 @@ def _make_agent( *, response_text: str | None = None, stream_texts: list[str] | None = None, + update_session: Callable[[AgentSession], None] | None = None, ) -> _FakeAgent: """Build a ``_FakeAgent`` from plain text for non-streaming/streaming runs.""" response = None @@ -150,7 +188,7 @@ def _make_agent( stream_updates = None if stream_texts is not None: stream_updates = [AgentResponseUpdate(contents=[Content.from_text(t)]) for t in stream_texts] - return _FakeAgent(response=response, stream_updates=stream_updates) + return _FakeAgent(response=response, stream_updates=stream_updates, update_session=update_session) def _make_request(payload: dict[str, Any]) -> Request: @@ -186,6 +224,470 @@ async def _collect_stream(response: StreamingResponse) -> str: # endregion +class TestSessionLifecycle: + @pytest.mark.parametrize("stream", [False, True]) + async def test_http_requests_restore_sessions_without_changing_response_format(self, stream: bool) -> None: + turns: list[int] = [] + + def update_session(session: AgentSession) -> None: + session.state["turn"] = session.state.get("turn", 0) + 1 + turns.append(session.state["turn"]) + + for _ in range(2): + server = InvocationsHostServer( + _make_agent(response_text="ok", stream_texts=["o", "k"], update_session=update_session) + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: + response = await client.post( + "/invocations", + params={"agent_session_id": "session"}, + json={"message": "hello", "stream": stream}, + ) + assert response.status_code == 200 + assert response.text == "ok" + if stream: + assert response.headers["content-type"].startswith("text/event-stream") + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + assert turns == [1, 2] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_function_history_is_restored_without_reexecuting_completed_tool(self, stream: bool) -> None: + executed: list[str] = [] + + @tool + def remember() -> str: + """Record a single completed tool call.""" + executed.append("called") + return "remembered" + + class HistoryClient(FunctionInvocationLayer[Any], ChatMiddlewareLayer[Any], BaseChatClient[Any]): + def __init__(self) -> None: + super().__init__(middleware=[]) + self.calls: list[list[Message]] = [] + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.calls.append(list(messages)) + contents = ( + [Content.from_function_call(call_id="call-1", name="remember", arguments="{}")] + if len(self.calls) == 1 + else [Content.from_text("ok")] + ) + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=contents) + + async def complete() -> ChatResponse: + return ChatResponse(messages=[Message("assistant", contents=contents)]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else complete() + + client = HistoryClient() + for message in ["first", "second"]: + agent = Agent(client=client, tools=[remember], context_providers=[InMemoryHistoryProvider()]) + server = InvocationsHostServer(agent) + with _request_context(session_id="session"): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": message, "stream": stream}) + ) + if isinstance(response, StreamingResponse): + assert await _collect_stream(response) == "ok" + else: + assert bytes(response.body).decode() == "ok" + + assert executed == ["called"] + assert len(client.calls) == 3 + restored = [content for message in client.calls[-1] for content in message.contents] + assert sum(content.type == "function_call" for content in restored) == 1 + assert sum(content.type == "function_result" for content in restored) == 1 + assert client.calls[-1][0].text == "first" + assert client.calls[-1][-1].text == "second" + + async def test_mismatched_stored_session_is_not_replaced(self) -> None: + store = _mock_session_store() + store.get = AsyncMock(return_value=AgentSession(session_id="different")) + store.set = AsyncMock() + agent = _make_agent(response_text="ok") + server = InvocationsHostServer(agent, agent_session_store_provider=_SessionStoreProvider(store)) + with _request_context(session_id="session"), pytest.raises(ValueError, match="does not match"): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + assert agent.calls == [] + store.set.assert_not_awaited() + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_cancelled_persistence_is_reported_and_releases_coordination( + self, caplog: pytest.LogCaptureFixture + ) -> None: + store = _mock_session_store() + store.set = AsyncMock(side_effect=asyncio.CancelledError) + server = InvocationsHostServer( + _make_agent(response_text="ok"), agent_session_store_provider=_SessionStoreProvider(store) + ) + with _request_context(session_id="session"), pytest.raises(asyncio.CancelledError): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + assert "Invocation session persistence was cancelled" in caplog.text + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_factory_failure_does_not_retain_session(self, stream: bool) -> None: + def create_agent() -> _FakeAgent: + raise ValueError("factory failed") + + server = InvocationsHostServer(create_agent) + with _request_context(session_id="session"), pytest.raises(ValueError, match="factory failed"): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "hello", "stream": stream}) + ) + if isinstance(response, StreamingResponse): + await _collect_stream(response) + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_stream_close_failure_is_reported_after_persistence(self) -> None: + store = _mock_session_store() + store.set = AsyncMock() + agent = _make_agent(stream_texts=["ok"]) + iterator = MagicMock() + iterator.__aiter__ = MagicMock(return_value=iterator) + iterator.__anext__ = AsyncMock(side_effect=StopAsyncIteration) + iterator.aclose = AsyncMock(side_effect=ValueError("close failed")) + server = InvocationsHostServer(agent, agent_session_store_provider=_SessionStoreProvider(store)) + with patch.object(agent, "run", return_value=iterator), _request_context(session_id="session"): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "hello", "stream": True}) + ) + assert isinstance(response, StreamingResponse) + with pytest.raises(ValueError, match="close failed"): + await _collect_stream(response) + store.set.assert_awaited_once() + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_custom_store_receives_request_context_and_namespaced_bounded_keys(self) -> None: + store = _mock_session_store() + store.get = AsyncMock(return_value=None) + store.set = AsyncMock() + provider = _SessionStoreProvider(store) + server = InvocationsHostServer(_make_agent(response_text="ok"), agent_session_store_provider=provider) + + identities = ["x" * 1024, '["session","user"]'] + for identity in identities: + with _request_context(call_id="call", session_id=identity): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + server.config.is_hosted = True + with _request_context(call_id="hosted-call", session_id="session", user_id="user"): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + + keys = [call.args[0] for call in store.set.await_args_list] + assert len(set(keys)) == 3 + assert all(key.startswith("invocations:v1:") and len(key) <= 128 for key in keys) + assert [call.args[1].session_id for call in store.set.await_args_list] == [ + *identities, + '["session","user"]', + ] + assert provider.contexts[-1].user_id == "user" + assert provider.contexts[-1].call_id == "hosted-call" + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_unconsumed_stream_and_invalid_input_do_not_access_storage(self) -> None: + store = _mock_session_store() + provider = _SessionStoreProvider(store) + agent = _make_agent(stream_texts=["hello"]) + server = InvocationsHostServer(agent, agent_session_store_provider=provider) + with _request_context(session_id="session"): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "hello", "stream": True}) + ) + invalid = await server._handle_invoke(_make_request({})) # pyright: ignore[reportPrivateUsage] + assert invalid.status_code == 400 + assert isinstance(response, StreamingResponse) + await cast(Any, response.body_iterator).aclose() + assert provider.contexts == [] + assert agent.calls == [] + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + @pytest.mark.parametrize("cancel_waiter", [False, True]) + @pytest.mark.parametrize("stream", [False, True]) + async def test_same_session_serializes_and_waiter_cleanup_is_bounded( + self, cancel_waiter: bool, stream: bool + ) -> None: + entered = asyncio.Event() + release = asyncio.Event() + second_started = asyncio.Event() + turns: list[int] = [] + + class BlockingAgent(_FakeAgent): + def run( + self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any + ) -> Any: + async def complete() -> AgentResponse: + assert session is not None + turn = session.state.get("turn", 0) + 1 + if not turns: + entered.set() + await release.wait() + session.state["turn"] = turn + turns.append(turn) + return AgentResponse(messages=[Message("assistant", "ok")]) + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + await complete() + yield AgentResponseUpdate(contents=[Content.from_text("ok")]) + + return updates() if stream else complete() + + server = InvocationsHostServer(BlockingAgent()) + + async def invoke(*, second: bool = False) -> Response: + with _request_context(session_id="session"): + if second: + second_started.set() + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "hello", "stream": stream}) + ) + if isinstance(response, StreamingResponse): + await _collect_stream(response) + return response + + first = asyncio.create_task(invoke()) + await entered.wait() + second = asyncio.create_task(invoke(second=True)) + try: + await second_started.wait() + assert not second.done() + assert server._session_locks["session"].users == 2 # pyright: ignore[reportPrivateUsage] + if cancel_waiter: + second.cancel() + with pytest.raises(asyncio.CancelledError): + await second + assert server._session_locks["session"].users == 1 # pyright: ignore[reportPrivateUsage] + release.set() + await first + if not cancel_waiter: + await second + finally: + release.set() + await asyncio.gather(first, second, return_exceptions=True) + assert turns == ([1] if cancel_waiter else [1, 2]) + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_different_sessions_run_concurrently(self) -> None: + entered: set[str] = set() + both_entered = asyncio.Event() + + class ConcurrentAgent(_FakeAgent): + def run( + self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any + ) -> Any: + async def complete() -> AgentResponse: + assert session is not None + entered.add(session.session_id) + if len(entered) == 2: + both_entered.set() + await both_entered.wait() + return AgentResponse(messages=[Message("assistant", "ok")]) + + return complete() + + server = InvocationsHostServer(ConcurrentAgent()) + + async def invoke(session_id: str) -> None: + with _request_context(session_id=session_id): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + + await asyncio.wait_for(asyncio.gather(invoke("one"), invoke("two")), timeout=5) + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("failure", ["load", "save", "run", "run_and_save"]) + async def test_storage_and_agent_errors_are_not_successful_responses( + self, stream: bool, failure: str, caplog: pytest.LogCaptureFixture + ) -> None: + store = _mock_session_store() + store.get = ( + AsyncMock(side_effect=ValueError("load failed")) if failure == "load" else AsyncMock(return_value=None) + ) + store.set = AsyncMock(side_effect=ValueError("save failed")) if "save" in failure else AsyncMock() + + def update_session(session: AgentSession) -> None: + session.state["turn"] = 1 + if "run" in failure: + raise ValueError("run failed") + + agent = _make_agent(response_text="ok", stream_texts=["ok"], update_session=update_session) + server = InvocationsHostServer(agent, agent_session_store_provider=_SessionStoreProvider(store)) + expected = ( + "Invocation failed and session persistence also failed" + if failure == "run_and_save" + else (f"{failure} failed") + ) + with _request_context(session_id="session"), pytest.raises((ValueError, RuntimeError), match=expected): + result = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "hello", "stream": stream}) + ) + if isinstance(result, StreamingResponse): + await _collect_stream(result) + + if failure == "load": + assert agent.calls == [] + store.set.assert_not_awaited() + assert "Failed to load invocation session" in caplog.text + else: + assert store.set.await_args is not None + assert store.set.await_args.args[1].state == {"turn": 1} + if "save" in failure: + assert "Failed to persist invocation session" in caplog.text + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + @pytest.mark.parametrize("fail_save", [False, True]) + async def test_cancellation_preserves_state_and_cancellation_signal( + self, fail_save: bool, caplog: pytest.LogCaptureFixture + ) -> None: + entered = asyncio.Event() + store = _mock_session_store() + store.set = AsyncMock(side_effect=ValueError("save failed")) if fail_save else AsyncMock() + + class WaitingAgent(_FakeAgent): + def run( + self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any + ) -> Any: + async def complete() -> AgentResponse: + assert session is not None + session.state["started"] = True + entered.set() + await asyncio.Event().wait() + return AgentResponse() + + return complete() + + server = InvocationsHostServer(WaitingAgent(), agent_session_store_provider=_SessionStoreProvider(store)) + with _request_context(session_id="session"): + task = asyncio.create_task( + server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + ) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert store.set.await_args is not None + assert store.set.await_args.args[1].state == {"started": True} + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + if fail_save: + assert "Failed to persist invocation session" in caplog.text + + @pytest.mark.parametrize("spec_version", ["2.0", "2.4"]) + async def test_disconnect_closes_agent_and_saves_final_state(self, spec_version: str) -> None: + sent = asyncio.Event() + events: list[str] = [] + store = _mock_session_store() + + class StreamingAgent(_ContextAgent): + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + await asyncio.sleep(0) + events.append("exit") + + def run( + self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any + ) -> Any: + async def updates() -> AsyncIterator[AgentResponseUpdate]: + assert session is not None + try: + yield AgentResponseUpdate(contents=[Content.from_text("first")]) + await asyncio.Event().wait() + finally: + await asyncio.sleep(0) + session.state["closed"] = True + events.append("close") + + return updates() + + server = InvocationsHostServer( + lambda: StreamingAgent(events), agent_session_store_provider=_SessionStoreProvider(store) + ) + with _request_context(session_id="session"): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "hello", "stream": True}) + ) + + async def send(message: Any) -> None: + if message["type"] == "http.response.body": + sent.set() + if spec_version == "2.4": + raise OSError("disconnected") + await asyncio.Event().wait() + + async def receive() -> Any: + await sent.wait() + return {"type": "http.disconnect"} + + scope = {"type": "http", "asgi": {"spec_version": spec_version}} + if spec_version == "2.4": + with pytest.raises(ClientDisconnect): + await response(scope, receive, send) + else: + await response(scope, receive, send) + assert events == ["enter", "close", "exit"] + assert store.set.await_args is not None + assert store.set.await_args.args[1].state == {"closed": True} + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_cancelled_scope_still_awaits_persistence(self) -> None: + store = _mock_session_store() + saved = False + + async def save(session_id: str, session: AgentSession) -> None: + nonlocal saved + await asyncio.sleep(0) + saved = True + + store.set = AsyncMock(side_effect=save) + server = InvocationsHostServer( + _make_agent(response_text="ok"), agent_session_store_provider=_SessionStoreProvider(store) + ) + with _request_context(session_id="session"): + with CancelScope() as scope: + async with server._request_session("session", FoundryAgentRequestContext()): # pyright: ignore[reportPrivateUsage] + scope.cancel() + await asyncio.sleep(0) + assert saved + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_default_store_expiry_is_renewed_by_writes(self) -> None: + turns: list[int] = [] + + def update_session(session: AgentSession) -> None: + session.state["turn"] = session.state.get("turn", 0) + 1 + turns.append(session.state["turn"]) + + server = InvocationsHostServer(_make_agent(response_text="ok", update_session=update_session)) + day = 24 * 60 * 60 + with ( + _request_context(session_id="session"), + patch("azure.ai.agentserver.core.storage._local_state._now") as now, + ): + for timestamp in [0, 29 * day, 58 * day, 88 * day]: + now.return_value = timestamp + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + assert turns == [1, 2, 3, 1] + + async def test_unserializable_state_fails_explicitly(self) -> None: + def update_session(session: AgentSession) -> None: + session.state["unsupported"] = object() + + server = InvocationsHostServer(_make_agent(response_text="ok", update_session=update_session)) + with ( + _request_context(session_id="session"), + pytest.warns(RuntimeWarning, match="unsupported type 'object'"), + pytest.raises(TypeError), + ): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + # region Initialization @@ -193,7 +695,7 @@ class TestInit: def test_accepts_supports_agent_run(self) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) assert server._agent is not None # pyright: ignore[reportPrivateUsage] - assert server._sessions == {} # pyright: ignore[reportPrivateUsage] + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] @pytest.mark.parametrize("agent", [None, 42]) def test_rejects_invalid_agent_source(self, agent: Any) -> None: @@ -285,6 +787,38 @@ async def test_hosted_keys_and_session_ids_preserve_identifier_values(self) -> N class TestHandleInvoke: + async def test_completed_sessions_are_released_and_state_is_restored(self) -> None: + sessions: list[weakref.ReferenceType[AgentSession]] = [] + turns: list[int] = [] + + class CountingAgent(_FakeAgent): + def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: AgentSession | None = None, + **kwargs: Any, + ) -> Any: + assert session is not None + sessions.append(weakref.ref(session)) + turn = session.state.get("turn", 0) + 1 + session.state["turn"] = turn + turns.append(turn) + + async def complete() -> AgentResponse: + return AgentResponse(messages=[Message("assistant", "ok")]) + + return complete() + + server = InvocationsHostServer(CountingAgent()) + for session_id in ["first", "second", "third", "first"]: + with _request_context(session_id=session_id): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + + assert all(session() is None for session in sessions) + assert turns == [1, 1, 1, 2] + async def test_instance_context_lifetime_remains_caller_owned(self) -> None: events: list[str] = [] response = AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ok")])]) @@ -348,11 +882,12 @@ def create_agent() -> _FakeAgent: assert bytes(second.body).decode() == "agent-2" assert len(agents) == 2 assert agents[0] is not agents[1] - assert agents[0].calls[0]["session"] is agents[1].calls[0]["session"] + assert agents[0].calls[0]["session"] is not agents[1].calls[0]["session"] + assert agents[0].calls[0]["session"].session_id == agents[1].calls[0]["session"].session_id @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("hosted", [False, True]) - async def test_reusing_session_skips_serialization_and_construction(self, hosted: bool, stream: bool) -> None: + async def test_reusing_session_restores_identifier(self, hosted: bool, stream: bool) -> None: agent = _make_agent(response_text="ok", stream_texts=["ok"]) server = InvocationsHostServer(agent) server.config.is_hosted = hosted @@ -361,7 +896,6 @@ async def test_reusing_session_skips_serialization_and_construction(self, hosted with ( _request_context(call_id="call-1", session_id="sess-1", user_id="user-1"), - patch("agent_framework_foundry_hosting._invocations.json", wraps=json) as serializer, patch("agent_framework_foundry_hosting._invocations.AgentSession", wraps=AgentSession) as session_factory, ): for _ in range(2): @@ -371,10 +905,10 @@ async def test_reusing_session_skips_serialization_and_construction(self, hosted else: assert bytes(response.body).decode() == "ok" - assert agent.calls[0]["session"] is agent.calls[1]["session"] - assert agent.calls[0]["session"].session_id == expected_id - assert serializer.dumps.call_count == (1 if hosted else 0) + assert agent.calls[0]["session"] is not agent.calls[1]["session"] + assert agent.calls[0]["session"].session_id == agent.calls[1]["session"].session_id == expected_id session_factory.assert_called_once_with(session_id=expected_id) + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] async def test_missing_message_returns_400(self) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) @@ -411,10 +945,11 @@ async def test_non_streaming_returns_agent_text(self) -> None: assert isinstance(response, Response) assert response.status_code == 200 assert bytes(response.body).decode() == "Hello!" - # Agent is called with the message wrapped in a list and the cached session. + # Agent is called with the message wrapped in a list and the restored session. assert agent.calls[0]["messages"] == ["Hi"] assert agent.calls[0]["stream"] is False - assert agent.calls[0]["session"] is server._sessions["sess-1"] # pyright: ignore[reportPrivateUsage] + assert agent.calls[0]["session"].session_id == "sess-1" + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] async def test_streaming_yields_update_text(self) -> None: agent = _make_agent(stream_texts=["Hel", "lo", "!"]) @@ -429,19 +964,24 @@ async def test_streaming_yields_update_text(self) -> None: assert agent.calls[0]["messages"] == "Hi" assert agent.calls[0]["stream"] is True - async def test_session_is_reused_across_requests(self) -> None: - agent = _make_agent(response_text="ok") + async def test_session_state_is_restored_across_requests_and_hosts(self) -> None: + def update_session(session: AgentSession) -> None: + session.state["turn"] = session.state.get("turn", 0) + 1 + + agent = _make_agent(response_text="ok", update_session=update_session) server = InvocationsHostServer(agent) with _request_context(session_id="sess-1"): await server._handle_invoke(_make_request({"message": "one"})) # pyright: ignore[reportPrivateUsage] - first_session = server._sessions["sess-1"] # pyright: ignore[reportPrivateUsage] + first_session = agent.calls[-1]["session"] + server = InvocationsHostServer(agent) await server._handle_invoke(_make_request({"message": "two"})) # pyright: ignore[reportPrivateUsage] - second_session = server._sessions["sess-1"] # pyright: ignore[reportPrivateUsage] + second_session = agent.calls[-1]["session"] - assert first_session is second_session - assert list(server._sessions) == ["sess-1"] # pyright: ignore[reportPrivateUsage] - assert agent.calls[0]["session"] is agent.calls[1]["session"] + assert first_session is not second_session + assert first_session.state == {"turn": 1} + assert second_session.state == {"turn": 2} + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize( @@ -461,7 +1001,10 @@ async def test_hosted_sessions_preserve_identifier_boundaries( second_session_id: str, second_user_id: str, ) -> None: - agent = _make_agent(response_text="ok", stream_texts=["ok"]) + def update_session(session: AgentSession) -> None: + session.state.setdefault("turn", session.session_id) + + agent = _make_agent(response_text="ok", stream_texts=["ok"], update_session=update_session) server = InvocationsHostServer(agent) server.config.is_hosted = True identifiers = [(first_session_id, first_user_id), (second_session_id, second_user_id)] @@ -480,13 +1023,12 @@ async def test_hosted_sessions_preserve_identifier_boundaries( session = agent.calls[-1]["session"] assert isinstance(session, AgentSession) - assert session.state == {} - session.state["turn"] = (session_id, user_id) + assert session.state == {"turn": session.session_id} sessions.append(session) assert sessions[0] is not sessions[1] assert sessions[0].session_id != sessions[1].session_id - assert len(server._sessions) == 2 # pyright: ignore[reportPrivateUsage] + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] for (session_id, user_id), session in zip(identifiers, sessions): with _request_context(call_id="call-2", session_id=session_id, user_id=user_id): @@ -499,8 +1041,8 @@ async def test_hosted_sessions_preserve_identifier_boundaries( assert bytes(response.body).decode() == "ok" assert response.status_code == 200 - assert agent.calls[-1]["session"] is session - assert session.state == {"turn": (session_id, user_id)} + assert agent.calls[-1]["session"] is not session + assert agent.calls[-1]["session"].state == {"turn": session.session_id} # endregion diff --git a/python/uv.lock b/python/uv.lock index 65fe4e5f906..5666bf392b0 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -653,6 +653,7 @@ version = "1.0.0b260918" source = { editable = "packages/foundry_hosting" } dependencies = [ { name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "azure-ai-agentserver-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "azure-ai-agentserver-invocations", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "azure-ai-agentserver-responses", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -663,6 +664,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "agent-framework-core", editable = "packages/core" }, + { name = "anyio", specifier = ">=4,<5" }, { name = "azure-ai-agentserver-core", specifier = ">=2.1.0,<3" }, { name = "azure-ai-agentserver-invocations", specifier = ">=1.1.0,<2" }, { name = "azure-ai-agentserver-responses", specifier = ">=2.2.0b1,<3" }, From 5c9c367193adcf646dfcad3b130d49f26f304081 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Mon, 21 Sep 2026 15:19:16 +0200 Subject: [PATCH 2/5] docs: keep the hosting README focused on usage --- ...2-python-invocation-session-persistence.md | 25 ++++++++++ python/packages/foundry_hosting/README.md | 49 +++---------------- 2 files changed, 33 insertions(+), 41 deletions(-) diff --git a/docs/decisions/0042-python-invocation-session-persistence.md b/docs/decisions/0042-python-invocation-session-persistence.md index ba8f1dbb8df..9f59412427e 100644 --- a/docs/decisions/0042-python-invocation-session-persistence.md +++ b/docs/decisions/0042-python-invocation-session-persistence.md @@ -70,6 +70,14 @@ The coordination scope includes stream closure, factory-agent cleanup, and sessi The host attempts to save state on success, execution failure, and interruption, without retrying agent execution. Disconnect cleanup is shielded from the ASGI cancellation scope. Persistence errors are reported; there is no success-shaped in-memory fallback. +After streaming headers have been sent, a persistence failure terminates the stream and is +logged; already delivered text cannot be retracted. + +Starlette uses AnyIO cancellation scopes even when running on asyncio. `asyncio.shield()` +protects the child operation but still allows the awaiting request to unwind on cancellation. +Using `CancelScope(shield=True)` instead keeps cleanup awaited inside the request's ownership +scope before releasing session coordination. This does not replace the asyncio event loop. +AnyIO is already a transitive dependency and is declared directly because hosting imports it. This removes cumulative host retention of completed session objects. It does not bound individual session size, active request count, backend storage, or total process memory. @@ -77,3 +85,20 @@ The SDK local backend reads its entire logical-store file per operation. The sto does not provide distributed transactions, so concurrent updates across hosts still require application-level coordination. Durable serialization also replaces Python object-identity reuse with restored state, and custom state needs registered codecs. + +### Migration and provider configuration + +Existing process-local invocation sessions are not migrated across deployment; the previous +implementation also lost them on restart. Local sessions now survive host recreation and are +written under `AGENTSERVER_STATE_ROOT` or the SDK's default state directory. Independent +applications should use separate storage roots or providers. + +State must support `AgentSession.to_dict()` / `AgentSession.from_dict()`. Custom state types +need registered codecs rather than arbitrary live Python objects. Restored sessions preserve +state, not Python object identity. + +A custom provider's `get_store` receives host configuration and request platform context. +The provider owns retention and deletion; an explicitly selected in-memory store remains +volatile and does not gain automatic eviction. For new default stores, writes renew the +30-day expiry window and reads do not. Missing, deleted, or expired state starts fresh on +the next invocation, including when its caller reuses the same session ID. diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 7670f74c2c3..ceb9b8c5ff6 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -114,50 +114,17 @@ Native Responses refusal parts are stored as text carrying `additional_properties["model_output_kind"] == "refusal"` and emitted as `response.refusal.*` events when streamed back to clients. -#### Invocation sessions - -`InvocationsHostServer` loads session state for each request and saves it after agent and -stream cleanup, including when execution fails or is interrupted. It does not retain completed -sessions in a host-level cache. Storage errors are reported rather than silently starting a -new conversation or falling back to memory. After streaming headers have been sent, a -persistence failure terminates the stream and is logged; already delivered text cannot be -retracted. A failed request is not automatically retried. - -When hosted, `AgentSession.session_id` is an +`InvocationsHostServer` restores sessions from its configured store. When hosted, `AgentSession.session_id` is an opaque composite identifier that preserves the boundaries between the platform session ID and user ID. Consumers must use it as a whole and must not parse it or depend on its internal representation. Repeated requests for the same identifier pair restore the saved session. -Locally, the platform session ID is used unchanged. Storage keys are opaque, fixed-length, -and namespaced separately from Responses keys. - -New default stores expire saved sessions **30 days after the last write**. Each write renews -that period; a read does not. Existing stores keep their creation-time retention settings. -A missing, deleted, or expired session starts a fresh conversation on the next request, even -if the caller reuses its ID. A running request keeps its working session until cleanup and -does not evict other conversations. - -To select another store, pass `agent_session_store_provider`, implementing -`StoreProvider[SessionStore]`, to the constructor. Its `get_store` method receives the host -configuration and the request's platform context, as it does for `ResponsesHostServer`. -The provider owns storage retention and deletion. An explicitly supplied in-memory store is -still volatile and does not gain automatic eviction. - -**Migration from in-memory invocation sessions:** session state must now support -`AgentSession.to_dict()` / `AgentSession.from_dict()` serialization. Register codecs for -custom state types instead of storing arbitrary live Python objects. Requests restore -equivalent state, not the same Python object. Existing process-local sessions are not -migrated across deployment; the previous implementation also lost them on restart. -Local sessions now survive a host restart and are written under `AGENTSERVER_STATE_ROOT` -(or the default location above). Use separate storage roots or providers for independent -applications. - -Requests for the same session are serialized within one host through execution, streaming, -cleanup, and persistence. Different sessions can run concurrently. There is no cross-host -transaction or exactly-once execution guarantee; deployments with overlapping same-session -requests across workers must coordinate those requests. Removing the host cache does not -impose a byte-memory limit, request-admission limit, or storage quota. The SDK's local backend -reads a whole logical-store file per operation, so its transient memory and I/O can still -grow with stored data. +Locally, the platform session ID is used unchanged. + +Both hosts accept `agent_session_store_provider` to select a `StoreProvider[SessionStore]`. +Session state must support `AgentSession` serialization. New default stores expire saved +sessions 30 days after their last write; an invocation after expiry starts a fresh session. +Existing stores retain their creation-time settings, and custom providers own their retention +policies. ### Workflow checkpoints From ca4e92a7f310853dcfce0e304eed6dd38ae35008 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Mon, 21 Sep 2026 15:48:58 +0200 Subject: [PATCH 3/5] Python: address invocation persistence review feedback --- ...2-python-invocation-session-persistence.md | 27 ++- .../specs/004-python-function-calling-loop.md | 1 + python/packages/foundry_hosting/README.md | 7 +- .../_invocations.py | 101 +++++--- .../_state_store.py | 25 +- .../foundry_hosting/tests/test_invocations.py | 221 +++++++++++++++++- .../tests/test_invocations_int.py | 4 +- .../foundry_hosting/tests/test_state_store.py | 22 ++ .../invocations/basic/main.py | 2 +- 9 files changed, 360 insertions(+), 50 deletions(-) diff --git a/docs/decisions/0042-python-invocation-session-persistence.md b/docs/decisions/0042-python-invocation-session-persistence.md index 9f59412427e..8828d97c2c9 100644 --- a/docs/decisions/0042-python-invocation-session-persistence.md +++ b/docs/decisions/0042-python-invocation-session-persistence.md @@ -42,6 +42,11 @@ store locally. New default stores inherit the existing SDK retention policy: 30 since the last write. Existing stores and custom providers retain their own policies. Absent or expired state starts a new conversation; there is no permanent tombstone registry. +Invocations defaults to the separate `invocation_sessions` logical store. A prefix alone +cannot isolate keys from Responses, which accepts opaque conversation identifiers in its +`agent_sessions` store. The shared store adapter and provider accept `store_name` to select +the logical store without changing the Responses default. + ### Alignment with .NET This follows **.NET Foundry Responses hosting's persistence architecture**, not a purported @@ -61,8 +66,17 @@ concurrency implementations. ### Lifecycle and limitations -Invocation storage keys hash an unambiguous partition encoding and use a versioned, -protocol-specific namespace; runtime `AgentSession.session_id` values remain unchanged. +Invocation storage keys hash an unambiguous application-and-partition encoding and use a +versioned prefix within the separate logical store. Applications can set +`session_store_namespace`; otherwise the configured `FOUNDRY_AGENT_ID` takes precedence +over `FOUNDRY_AGENT_NAME`. Default storage rejects a missing identity rather than guessing +one or sharing an unnamed global application partition. Application identities must be +stable across restarts and distinct between independent applications. Agent versions and +per-process identifiers are not used as namespace defaults. + +New session IDs retain the existing hosted/local representation. Restored sessions keep +their stored ID: the generic `SessionStore` contract does not require that ID to match a +lookup key or the request's platform ID. Same-partition requests serialize within one host. Reference-counted lock entries exist only for active/waiting requests and are removed after the final request leaves. @@ -70,9 +84,14 @@ The coordination scope includes stream closure, factory-agent cleanup, and sessi The host attempts to save state on success, execution failure, and interruption, without retrying agent execution. Disconnect cleanup is shielded from the ASGI cancellation scope. Persistence errors are reported; there is no success-shaped in-memory fallback. +Combined execution/persistence failures retain both messages and the persistence cause. After streaming headers have been sent, a persistence failure terminates the stream and is logged; already delivered text cannot be retracted. +AgentServer can replace the response iterator with an SSE keep-alive wrapper whose pump +is still reading the agent stream. Disconnect cleanup closes the installed outer iterator +first, allowing that wrapper to cancel and join its pump, then closes the original stream. + Starlette uses AnyIO cancellation scopes even when running on asyncio. `asyncio.shield()` protects the child operation but still allows the awaiting request to unwind on cancellation. Using `CancelScope(shield=True)` instead keeps cleanup awaited inside the request's ownership @@ -91,13 +110,15 @@ object-identity reuse with restored state, and custom state needs registered cod Existing process-local invocation sessions are not migrated across deployment; the previous implementation also lost them on restart. Local sessions now survive host recreation and are written under `AGENTSERVER_STATE_ROOT` or the SDK's default state directory. Independent -applications should use separate storage roots or providers. +applications should configure distinct application namespaces. State must support `AgentSession.to_dict()` / `AgentSession.from_dict()`. Custom state types need registered codecs rather than arbitrary live Python objects. Restored sessions preserve state, not Python object identity. A custom provider's `get_store` receives host configuration and request platform context. +Providers explicitly supplied by the application are responsible for application and protocol +isolation, and may use their own scoping without `session_store_namespace`. The provider owns retention and deletion; an explicitly selected in-memory store remains volatile and does not gain automatic eviction. For new default stores, writes renew the 30-day expiry window and reads do not. Missing, deleted, or expired state starts fresh on diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index dc57a607cd6..a7110cb1db5 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -691,6 +691,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Pending hosted history replay | Stateless hosted approval requests remain replayable until a response is recorded, then both controls become inert. | `packages/openai/tests/openai/test_openai_chat_client.py::test_stateless_history_preserves_pending_hosted_approval_request_until_response` | | Non-history provider plus session | Local history is still auto-injected for approval resume. | `packages/core/tests/core/test_agents.py::test_non_history_context_provider_still_injects_inmemory` | | Hosted per-service-call persistence | A host-managed transcript remains available throughout a local function-call loop without being persisted into the framework session and replayed on the next hosted request. | `packages/foundry_hosting/tests/test_responses.py::TestAgentSessionPersistence::test_per_service_call_persistence_preserves_function_loop_history` | +| Invocations durable function history | A later invocation restores one completed call/result pair across host recreation without re-executing the tool; both streaming and non-streaming responses preserve the final text and model-bound conversation history. | `packages/foundry_hosting/tests/test_invocations.py::TestSessionLifecycle::test_function_history_is_restored_without_reexecuting_completed_tool` | | Invalidated per-service-call persistence | An invalidated service response, its partial calls, and any nonexistent results are not success-persisted, and invalid content does not advance session continuation state. | `packages/core/tests/core/test_function_invocation_logic.py::test_invalidated_response_is_not_persisted` | | Streaming message injection with per-service-call persistence | A streaming response rebuilt from updates mirrors the inner conversation id exactly, including clearing it, and keeps its internal marker, so the next iteration appends only the latest message rather than replaying the whole turn on top of provider-held history, and never persists a conversation id from an earlier injected service call. | `packages/core/tests/core/test_middleware_with_chat.py::TestChatMiddleware::test_message_injection_middleware_streaming_preserves_inner_continuation_state`, `test_message_injection_middleware_streaming_keeps_service_conversation_id_external`, `test_message_injection_middleware_streaming_clears_conversation_id_when_final_call_has_none`, `test_message_injection_middleware_conversation_id_matches_across_streaming_modes`, `packages/core/tests/core/test_harness_agent.py::test_streaming_harness_tool_call_does_not_duplicate_transcript` | | Service-side approval decision | Stored hosted request is skipped; the current approved or rejected hosted response is sent, while local approval controls are omitted from provider input. | `packages/openai/tests/openai/test_openai_chat_client.py::test_prepare_messages_strips_approval_request_but_keeps_response_under_storage`, `test_prepare_messages_drops_local_approval_controls` | diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index ceb9b8c5ff6..b3286174fa2 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -88,7 +88,7 @@ subdirectory instead. Each logical store is saved as one JSON file whose name is a URL-safe Base64 encoding of the store name. For example: -- Agent sessions: `YWdlbnRfc2Vzc2lvbnM.json` +- Responses agent sessions: `YWdlbnRfc2Vzc2lvbnM.json` - Function approvals: `ZnVuY3Rpb25fYXBwcm92YWxz.json` - Workflow checkpoints: one file per context, encoded from `checkpoints/` @@ -105,7 +105,8 @@ No additional partitioning configuration is required when using the default stor `ResponsesHostServer` and `InvocationsHostServer` persist the Agent Framework `AgentSession` durably. By default they use `FoundryAgentSessionStore`, backed by Foundry storage when hosted -and file-based storage locally. Stored sessions are scoped under `agent_sessions`. +and file-based storage locally. Responses sessions use the `agent_sessions` logical store; +Invocations sessions use the separate `invocation_sessions` store. See the [custom storage provider sample](../../samples/04-hosting/foundry-hosted-agents/responses/custom_storage/) for an example that uses an in-memory session store locally and Azure Cosmos DB when hosted. @@ -121,6 +122,8 @@ representation. Repeated requests for the same identifier pair restore the saved Locally, the platform session ID is used unchanged. Both hosts accept `agent_session_store_provider` to select a `StoreProvider[SessionStore]`. +For default Invocations storage, set `session_store_namespace` to a stable application-specific +name unless `FOUNDRY_AGENT_ID` or `FOUNDRY_AGENT_NAME` is configured. Session state must support `AgentSession` serialization. New default stores expire saved sessions 30 days after their last write; an invocation after expiry starts a fresh session. Existing stores retain their creation-time settings, and custom providers own their retention diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index abbd7f4f6bf..695651fbbd5 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -9,6 +9,7 @@ from collections.abc import Awaitable, Callable from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager from dataclasses import dataclass, field +from types import TracebackType from agent_framework import AgentSession, ResponseStream, SessionStore, SupportsAgentRun from agent_framework._telemetry import mark_feature_used @@ -33,6 +34,14 @@ class _SessionLock: users: int = 0 +class _ShieldedAsyncExitStack(AsyncExitStack): + async def __aexit__( + self, exc_type: type[BaseException] | None, exc: BaseException | None, traceback: TracebackType | None + ) -> bool | None: + with CancelScope(shield=True): + return await super().__aexit__(exc_type, exc, traceback) + + class _InvocationStreamingResponse(StreamingResponse): def __init__(self, content: AsyncGenerator[str]) -> None: super().__init__( @@ -48,7 +57,13 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: finally: # Sending a chunk can fail while the generator is suspended at yield. with CancelScope(shield=True): - await self._content.aclose() + try: + close = getattr(self.body_iterator, "aclose", None) + if close is not None: + await close() + finally: + if self.body_iterator is not self._content: + await self._content.aclose() class InvocationsHostServer(InvocationAgentServerHost): @@ -60,6 +75,7 @@ def __init__( *, openapi_spec: dict[str, Any] | None = None, agent_session_store_provider: StoreProvider[SessionStore] | None = None, + session_store_namespace: str | None = None, **kwargs: Any, ) -> None: """Initialize an InvocationsHostServer. @@ -71,18 +87,40 @@ def __init__( agent_session_store_provider: Provider for conversation session storage. Defaults to Foundry storage when hosted and the SDK's file-backed storage locally. New default stores expire sessions 30 days after their last write. Custom providers control their own retention. + session_store_namespace: Stable application identity for session storage. Defaults to the configured + Foundry agent ID or name. Required for the default store when neither is configured, including + unnamed local applications. Keep it stable across restarts and distinct for independent applications. + Custom providers without a namespace are responsible for application and protocol isolation. **kwargs: Additional keyword arguments. This host will expect the request to be a JSON body with a "message" field. The response contains the agent's text, or streamed text when "stream" is true. """ validate_agent_source(agent) + if session_store_namespace is not None and ( + not isinstance(session_store_namespace, str) or not session_store_namespace.strip() + ): + raise ValueError("session_store_namespace must be a non-empty string") super().__init__(openapi_spec=openapi_spec, **kwargs) self._agent = agent self._owns_request_agent = not is_agent(agent) + self._session_namespace = ( + session_store_namespace + if session_store_namespace is not None + else self.config.agent_guid or self.config.agent_name or None + ) + if agent_session_store_provider is None and ( + self._session_namespace is None or not self._session_namespace.strip() + ): + raise ValueError( + "The default invocation session store requires a stable application identity. " + "Set session_store_namespace or configure FOUNDRY_AGENT_ID or FOUNDRY_AGENT_NAME." + ) self._session_store_provider = ( - AgentSessionStoreProvider() if agent_session_store_provider is None else agent_session_store_provider + AgentSessionStoreProvider(store_name="invocation_sessions") + if agent_session_store_provider is None + else agent_session_store_provider ) self._session_locks: dict[str | tuple[str, str], _SessionLock] = {} self.invoke_handler(self._handle_invoke) @@ -122,18 +160,10 @@ def _partition_key(self) -> str | tuple[str, str]: @asynccontextmanager async def _request_agent(self) -> AsyncGenerator[SupportsAgentRun]: agent = await resolve_agent(self._agent) - resources = AsyncExitStack() - try: + async with _ShieldedAsyncExitStack() as resources: if self._owns_request_agent and isinstance(agent, AbstractAsyncContextManager): await resources.enter_async_context(agent) yield agent - except BaseException as exc: - with CancelScope(shield=True): - if not await resources.__aexit__(type(exc), exc, exc.__traceback__): - raise - else: - with CancelScope(shield=True): - await resources.aclose() @asynccontextmanager async def _request_session( @@ -145,40 +175,43 @@ async def _request_session( entry.users += 1 try: async with entry.lock: - encoded_key = json.dumps(partition_key, separators=(",", ":")) - session_id = encoded_key if isinstance(partition_key, tuple) else partition_key + session_id = ( + json.dumps(partition_key, separators=(",", ":")) + if isinstance(partition_key, tuple) + else partition_key + ) + encoded_key = json.dumps((self._session_namespace, partition_key), separators=(",", ":")) storage_key = "invocations:v1:" + hashlib.sha256(encoded_key.encode("utf-8")).hexdigest() try: store = self._session_store_provider.get_store(config=self.config, platform_context=context) session = await store.get(storage_key) if session is None: session = AgentSession(session_id=session_id) - elif session.session_id != session_id: - raise ValueError("Stored invocation session does not match the requested session.") except Exception: logger.exception("Failed to load invocation session") raise - failure: BaseException | None = None - try: - yield session - except BaseException as exc: - failure = exc - raise - finally: - # Starlette cancels the streaming task's scope on disconnect. - with CancelScope(shield=True): - try: - await store.set(storage_key, session) - except Exception as exc: - logger.exception("Failed to persist invocation session") - if failure is None: - raise - if isinstance(failure, Exception): - raise RuntimeError("Invocation failed and session persistence also failed.") from exc - except asyncio.CancelledError: - logger.error("Invocation session persistence was cancelled") + async def save_session( + exc_type: type[BaseException] | None, failure: BaseException | None, traceback: TracebackType | None + ) -> None: + try: + await store.set(storage_key, session) + except Exception as exc: + logger.exception("Failed to persist invocation session") + if failure is None: raise + if isinstance(failure, Exception): + raise RuntimeError( + f"Invocation failed: {str(failure) or type(failure).__name__}; " + f"session persistence also failed: {str(exc) or type(exc).__name__}" + ) from exc + except asyncio.CancelledError: + logger.error("Invocation session persistence was cancelled") + raise + + async with _ShieldedAsyncExitStack() as cleanup: + cleanup.push_async_exit(save_session) + yield session finally: entry.users -= 1 if entry.users == 0: diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index a88c984ead5..931342d768a 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py @@ -300,12 +300,21 @@ class FoundryAgentSessionStore(SessionStore): DEFAULT_ROOT_SCOPE = "agent_sessions" - def __init__(self, platform_context: FoundryAgentRequestContext) -> None: + def __init__(self, platform_context: FoundryAgentRequestContext, *, store_name: str | None = None) -> None: + """Initialize session storage. + + Args: + platform_context: The request-scoped platform context. + store_name: Logical state-store name. Defaults to `agent_sessions`. + """ + if store_name is not None and (not isinstance(store_name, str) or not store_name.strip()): + raise ValueError("store_name must be a non-empty string") self.platform_context = platform_context + self._store_name = self.DEFAULT_ROOT_SCOPE if store_name is None else store_name async def _get_store(self) -> FoundryStateStore: return await FoundryStateStore.get_or_create( - f"{self.DEFAULT_ROOT_SCOPE}", + self._store_name, user_isolation=True, ) @@ -334,9 +343,19 @@ class AgentSessionStoreProvider(StoreProvider[SessionStore]): This defaults to using the `FoundryAgentSessionStore` in all environments. """ + def __init__(self, *, store_name: str | None = None) -> None: + """Initialize the provider. + + Args: + store_name: Logical state-store name. Defaults to `agent_sessions`. + """ + if store_name is not None and (not isinstance(store_name, str) or not store_name.strip()): + raise ValueError("store_name must be a non-empty string") + self._store_name = store_name + def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentRequestContext) -> SessionStore: """Get agent session store for the requested hosting environment.""" - return FoundryAgentSessionStore(platform_context) + return FoundryAgentSessionStore(platform_context, store_name=self._store_name) # endregion Agent session persistence diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index f88af3f6c6e..0841a7d49fc 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -52,12 +52,20 @@ from typing_extensions import Any from agent_framework_foundry_hosting import InvocationsHostServer, StoreProvider +from agent_framework_foundry_hosting._state_store import FoundryAgentSessionStore pytestmark = pytest.mark.filterwarnings("ignore:.*SessionStore is experimental.*") # region Helpers +@pytest.fixture(autouse=True) +def configure_local_application(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FOUNDRY_AGENT_NAME", "invocation-test-application") + monkeypatch.delenv("FOUNDRY_AGENT_ID", raising=False) + monkeypatch.delenv("FOUNDRY_HOSTING_ENVIRONMENT", raising=False) + + def _mock_session_store() -> MagicMock: store = MagicMock(spec=SessionStore) store.get.return_value = None @@ -225,6 +233,134 @@ async def _collect_stream(response: StreamingResponse) -> str: class TestSessionLifecycle: + async def test_default_invocations_store_is_separate_from_responses(self) -> None: + storage_keys: list[str] = [] + save = FoundryAgentSessionStore.set + turns: list[int] = [] + + async def record_save(store: FoundryAgentSessionStore, key: str, session: AgentSession) -> None: + storage_keys.append(key) + await save(store, key, session) + + def update_session(session: AgentSession) -> None: + session.state["turn"] = session.state.get("turn", 0) + 1 + turns.append(session.state["turn"]) + + server = InvocationsHostServer(_make_agent(response_text="ok", update_session=update_session)) + with ( + _request_context(session_id="session"), + patch.object(FoundryAgentSessionStore, "set", record_save), + ): + await server._handle_invoke(_make_request({"message": "one"})) # pyright: ignore[reportPrivateUsage] + + responses_store = FoundryAgentSessionStore(FoundryAgentRequestContext()) + key = storage_keys[0] + assert await responses_store.get(key) is None + responses_session = AgentSession(session_id="responses-session") + responses_session.state["protocol"] = "responses" + await responses_store.set(key, responses_session) + with _request_context(session_id="session"): + await server._handle_invoke(_make_request({"message": "two"})) # pyright: ignore[reportPrivateUsage] + assert turns == [1, 2] + restored = await responses_store.get(key) + assert restored is not None + assert restored.state == {"protocol": "responses"} + + async def test_independent_application_names_do_not_share_sessions(self, monkeypatch: pytest.MonkeyPatch) -> None: + seen: list[dict[str, Any]] = [] + + for application in ["first-app", "second-app", "first-app"]: + monkeypatch.setenv("FOUNDRY_AGENT_NAME", application) + + def update_session(session: AgentSession, application: str = application) -> None: + seen.append(dict(session.state)) + session.state["application"] = application + + server = InvocationsHostServer(_make_agent(response_text="ok", update_session=update_session)) + with _request_context(session_id="shared-session"): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + + assert seen == [{}, {}, {"application": "first-app"}] + + @pytest.mark.parametrize("spec_version", ["2.0", "2.4"]) + async def test_keep_alive_disconnect_joins_pump_before_session_cleanup(self, spec_version: str) -> None: + entered = asyncio.Event() + disconnected = asyncio.Event() + saved: list[dict[str, Any]] = [] + pump: asyncio.Task[Any] | None = None + store = _mock_session_store() + + async def save(key: str, session: AgentSession) -> None: + saved.append(dict(session.state)) + + store.set.side_effect = save + + class IdleAgent(_FakeAgent): + def run( + self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any + ) -> Any: + async def updates() -> AsyncIterator[AgentResponseUpdate]: + nonlocal pump + assert session is not None + pump = asyncio.current_task() + entered.set() + try: + await asyncio.Event().wait() + yield AgentResponseUpdate(contents=[Content.from_text("unused")]) + finally: + await asyncio.sleep(0) + session.state["closed"] = True + + return updates() + + server = InvocationsHostServer(IdleAgent(), agent_session_store_provider=_SessionStoreProvider(store)) + server.config.sse_keepalive_interval = 1 + body = json.dumps({"message": "hello", "stream": True}).encode() + received = False + + async def receive() -> Any: + nonlocal received + if not received: + received = True + return {"type": "http.request", "body": body} + await disconnected.wait() + return {"type": "http.disconnect"} + + async def send(message: Any) -> None: + if message["type"] == "http.response.body" and message.get("body") == b": keep-alive\n\n": + assert entered.is_set() + if spec_version == "2.4": + raise OSError("disconnected while idle") + disconnected.set() + await asyncio.Event().wait() + + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": spec_version}, + "method": "POST", + "scheme": "http", + "path": "/invocations", + "root_path": "", + "query_string": b"agent_session_id=session", + "headers": [(b"content-type", b"application/json")], + "server": ("test", 80), + "client": ("test", 1234), + } + try: + if spec_version == "2.4": + with pytest.raises(ClientDisconnect): + await asyncio.wait_for(server(scope, receive, send), timeout=5) + else: + await asyncio.wait_for(server(scope, receive, send), timeout=5) + assert saved == [{"closed": True}] + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + assert pump is not None + assert pump.done() + finally: + if pump is not None and not pump.done(): + pump.cancel() + await asyncio.gather(pump, return_exceptions=True) + @pytest.mark.parametrize("stream", [False, True]) async def test_http_requests_restore_sessions_without_changing_response_format(self, stream: bool) -> None: turns: list[int] = [] @@ -309,16 +445,61 @@ async def complete() -> ChatResponse: assert client.calls[-1][0].text == "first" assert client.calls[-1][-1].text == "second" - async def test_mismatched_stored_session_is_not_replaced(self) -> None: + async def test_custom_store_can_restore_a_different_session_id(self) -> None: store = _mock_session_store() store.get = AsyncMock(return_value=AgentSession(session_id="different")) store.set = AsyncMock() agent = _make_agent(response_text="ok") server = InvocationsHostServer(agent, agent_session_store_provider=_SessionStoreProvider(store)) - with _request_context(session_id="session"), pytest.raises(ValueError, match="does not match"): + with _request_context(session_id="session"): await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - assert agent.calls == [] - store.set.assert_not_awaited() + assert agent.calls[0]["session"].session_id == "different" + store.set.assert_awaited_once() + assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + + async def test_combined_failures_include_types_for_empty_messages(self) -> None: + store = _mock_session_store() + store.set.side_effect = RuntimeError() + + def update_session(session: AgentSession) -> None: + raise ValueError + + server = InvocationsHostServer( + _make_agent(response_text="ok", update_session=update_session), + agent_session_store_provider=_SessionStoreProvider(store), + ) + with ( + _request_context(session_id="session"), + pytest.raises( + RuntimeError, match="Invocation failed: ValueError; session persistence also failed: RuntimeError" + ) as raised, + ): + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + assert raised.value.__cause__ is store.set.side_effect + + async def test_control_flow_exception_still_closes_request_agent(self) -> None: + class ControlFlowExit(BaseException): + pass + + events: list[str] = [] + error = ControlFlowExit() + store = _mock_session_store() + store.set.side_effect = ValueError("save failed") + + class ExitingAgent(_ContextAgent): + def run( + self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any + ) -> Any: + raise error + + server = InvocationsHostServer( + lambda: ExitingAgent(events), agent_session_store_provider=_SessionStoreProvider(store) + ) + with _request_context(session_id="session"), pytest.raises(ControlFlowExit) as raised: + await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] + assert raised.value is error + assert events == ["enter", "exit"] + store.set.assert_awaited_once() assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] async def test_cancelled_persistence_is_reported_and_releases_coordination( @@ -521,7 +702,7 @@ def update_session(session: AgentSession) -> None: agent = _make_agent(response_text="ok", stream_texts=["ok"], update_session=update_session) server = InvocationsHostServer(agent, agent_session_store_provider=_SessionStoreProvider(store)) expected = ( - "Invocation failed and session persistence also failed" + "Invocation failed: run failed; session persistence also failed: save failed" if failure == "run_and_save" else (f"{failure} failed") ) @@ -692,6 +873,36 @@ def update_session(session: AgentSession) -> None: class TestInit: + def test_default_store_requires_application_identity(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("FOUNDRY_AGENT_NAME") + with pytest.raises(ValueError, match="Set session_store_namespace"): + InvocationsHostServer(_make_agent(response_text="ok")) + + server = InvocationsHostServer(_make_agent(response_text="ok"), session_store_namespace="explicit-app") + assert server._session_namespace == "explicit-app" # pyright: ignore[reportPrivateUsage] + + def test_custom_store_can_own_application_isolation(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("FOUNDRY_AGENT_NAME") + server = InvocationsHostServer( + _make_agent(response_text="ok"), agent_session_store_provider=_SessionStoreProvider(_mock_session_store()) + ) + assert server._session_namespace is None # pyright: ignore[reportPrivateUsage] + + @pytest.mark.parametrize("namespace", ["", " \t", 123]) + def test_rejects_invalid_explicit_namespace(self, namespace: Any) -> None: + with pytest.raises(ValueError, match="session_store_namespace must be a non-empty string"): + InvocationsHostServer(_make_agent(response_text="ok"), session_store_namespace=namespace) + + def test_namespace_precedence_is_explicit_then_agent_id_then_name(self, monkeypatch: pytest.MonkeyPatch) -> None: + agent = _make_agent(response_text="ok") + named = InvocationsHostServer(agent) + assert named._session_namespace == "invocation-test-application" # pyright: ignore[reportPrivateUsage] + monkeypatch.setenv("FOUNDRY_AGENT_ID", "stable-agent-id") + identified = InvocationsHostServer(agent) + assert identified._session_namespace == "stable-agent-id" # pyright: ignore[reportPrivateUsage] + explicit = InvocationsHostServer(agent, session_store_namespace="explicit-app") + assert explicit._session_namespace == "explicit-app" # pyright: ignore[reportPrivateUsage] + def test_accepts_supports_agent_run(self) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) assert server._agent is not None # pyright: ignore[reportPrivateUsage] diff --git a/python/packages/foundry_hosting/tests/test_invocations_int.py b/python/packages/foundry_hosting/tests/test_invocations_int.py index 136760da74f..902ce319a88 100644 --- a/python/packages/foundry_hosting/tests/test_invocations_int.py +++ b/python/packages/foundry_hosting/tests/test_invocations_int.py @@ -57,7 +57,7 @@ def server() -> InvocationsHostServer: default_options={"store": False}, # pyrefly: ignore[bad-argument-type] ) - return InvocationsHostServer(agent) + return InvocationsHostServer(agent, session_store_namespace="invocations-integration") @tool @@ -78,7 +78,7 @@ def server_with_tools() -> InvocationsHostServer: default_options={"store": False}, # pyrefly: ignore[bad-argument-type] ) - return InvocationsHostServer(agent) + return InvocationsHostServer(agent, session_store_namespace="invocations-with-tools-integration") # --------------------------------------------------------------------------- diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 7ae5088a4d9..2c6d8389bc5 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -450,6 +450,28 @@ async def test_set_agent_session_uses_scoped_store() -> None: store.set_item.assert_awaited_once_with("storage-session-1", session.to_dict(), call_id="call-1") +async def test_agent_session_provider_selects_a_separate_logical_store() -> None: + store = _store() + with patch( + "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", + new=AsyncMock(return_value=store), + ) as get_or_create: + provider = AgentSessionStoreProvider(store_name="invocation_sessions") + storage = provider.get_store(config=_config(is_hosted=False), platform_context=_platform_context()) + await storage.set("storage-key", AgentSession(session_id="runtime-id")) + + get_or_create.assert_awaited_once_with("invocation_sessions", user_isolation=True) + assert store.set_item.call_args.args[1]["session_id"] == "runtime-id" + + +@pytest.mark.parametrize("store_name", ["", " \t", 123]) +def test_session_store_name_must_be_nonempty(store_name: Any) -> None: + with pytest.raises(ValueError, match="store_name must be a non-empty string"): + FoundryAgentSessionStore(_platform_context(), store_name=store_name) + with pytest.raises(ValueError, match="store_name must be a non-empty string"): + AgentSessionStoreProvider(store_name=store_name) + + async def test_get_agent_session_returns_deserialized_session() -> None: store = _store() session = AgentSession(session_id="agent-session-1") diff --git a/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py index 730600d4f98..7b859f39123 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py @@ -27,7 +27,7 @@ def main(): default_options={"store": False}, ) - server = InvocationsHostServer(agent) + server = InvocationsHostServer(agent, session_store_namespace="basic-invocations") server.run() From 841136fd29804fe4d9bb91a5697a2063d49b3397 Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Mon, 21 Sep 2026 15:54:13 +0200 Subject: [PATCH 4/5] docs: clarify session persistence guarantees --- python/packages/foundry_hosting/README.md | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index b3286174fa2..054728f1355 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -124,11 +124,17 @@ Locally, the platform session ID is used unchanged. Both hosts accept `agent_session_store_provider` to select a `StoreProvider[SessionStore]`. For default Invocations storage, set `session_store_namespace` to a stable application-specific name unless `FOUNDRY_AGENT_ID` or `FOUNDRY_AGENT_NAME` is configured. -Session state must support `AgentSession` serialization. New default stores expire saved -sessions 30 days after their last write; an invocation after expiry starts a fresh session. +Session state must support `AgentSession` serialization. Use `register_state_type()` codecs for +custom types; unsupported live objects fail during persistence. Restored sessions preserve +state, not Python object identity. New default stores expire saved sessions 30 days after +their last write; an invocation after expiry starts a fresh session. Existing stores retain their creation-time settings, and custom providers own their retention policies. +Same-session requests are coordinated only within one host instance. Multi-worker deployments +must coordinate overlapping requests for the same session; cross-host transactions and +exactly-once execution are not guaranteed. + ### Workflow checkpoints `ResponsesHostServer` persists workflow checkpoints durably. By default, it uses the From 790ef355821f2b19f9adbede573788b28441717b Mon Sep 17 00:00:00 2001 From: Jose Alvarez Date: Wed, 23 Sep 2026 15:35:41 +0200 Subject: [PATCH 5/5] Python: simplify invocation persistence around the existing store provider --- ...2-python-invocation-session-persistence.md | 125 ------ python/packages/foundry_hosting/README.md | 8 +- .../_invocations.py | 143 ++---- .../packages/foundry_hosting/pyproject.toml | 1 - .../foundry_hosting/tests/test_invocations.py | 424 +----------------- .../tests/test_invocations_int.py | 4 +- .../invocations/basic/main.py | 2 +- python/uv.lock | 2 - 8 files changed, 58 insertions(+), 651 deletions(-) delete mode 100644 docs/decisions/0042-python-invocation-session-persistence.md diff --git a/docs/decisions/0042-python-invocation-session-persistence.md b/docs/decisions/0042-python-invocation-session-persistence.md deleted file mode 100644 index 8828d97c2c9..00000000000 --- a/docs/decisions/0042-python-invocation-session-persistence.md +++ /dev/null @@ -1,125 +0,0 @@ ---- -status: proposed -contact: jpalvarezl -date: 2026-09-21 -deciders: jpalvarezl ---- - -# Persist Python invocation sessions through the existing session-store abstraction - -## Context and Problem Statement - -`InvocationsHostServer` keeps conversation sessions in a host-owned dictionary for the -process lifetime. Request-scoped agent factories do not change that ownership: even after -an agent is disposed, the host still owns the session. We need to release completed -request state without losing multi-turn conversation continuity. - -## Decision Drivers - -- Restore conversation state rather than silently evicting live sessions. -- Preserve hosted user/session partition boundaries and request-scoped agent ownership. -- Reuse existing storage and serialization contracts instead of adding another cache policy. -- Keep streaming, failure, and cancellation cleanup inside the session's coordination scope. -- Make persistence and retention behavior explicit to applications. - -## Considered Options - -- **Cap the in-memory dictionary and refuse new conversations.** Bounds retained entry - count, but needs a new capacity policy and eventually refuses new work until an - application deletes state. It does not bound the size of each session. -- **Evict idle sessions.** Releases memory but loses conversation state unless backed by - persistent storage. Adds a separate host retention policy and active-session tracking. -- **Use the existing session-store provider.** Releases host-owned session objects while - preserving restorable state. Adds storage I/O and requires serializable state, but - reuses the established hosting extension point and backend retention policy. - -## Decision Outcome - -Choose the existing session-store provider. `InvocationsHostServer` accepts -`agent_session_store_provider`, defaulting to `AgentSessionStoreProvider`, like Python -`ResponsesHostServer`. The default is Foundry storage when hosted and the SDK's file-backed -store locally. New default stores inherit the existing SDK retention policy: 30 days -since the last write. Existing stores and custom providers retain their own policies. -Absent or expired state starts a new conversation; there is no permanent tombstone registry. - -Invocations defaults to the separate `invocation_sessions` logical store. A prefix alone -cannot isolate keys from Responses, which accepts opaque conversation identifiers in its -`agent_sessions` store. The shared store adapter and provider accept `store_name` to select -the logical store without changing the Responses default. - -### Alignment with .NET - -This follows **.NET Foundry Responses hosting's persistence architecture**, not a purported -.NET Invocations cache limit: - -- [`AddFoundryResponses`](../../dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/ServiceCollectionExtensions.cs) - defaults to `FoundryAgentSessionStore`. -- [`.NET FoundryAgentSessionStore`](../../dotnet/src/Microsoft.Agents.AI.Foundry.Hosting/FoundryAgentSessionStore.cs) - serializes sessions into the AgentServer state store, with a local fallback and the - platform's 30-day default item retention. -- The .NET in-memory session store remains a development/testing alternative; it does not - supply an eviction policy. The Invocations echo sample is stateless. - -Python reuses its own existing `StoreProvider[SessionStore]` contract. This is architectural -alignment, not shared storage-key formats, cross-language state compatibility, or identical -concurrency implementations. - -### Lifecycle and limitations - -Invocation storage keys hash an unambiguous application-and-partition encoding and use a -versioned prefix within the separate logical store. Applications can set -`session_store_namespace`; otherwise the configured `FOUNDRY_AGENT_ID` takes precedence -over `FOUNDRY_AGENT_NAME`. Default storage rejects a missing identity rather than guessing -one or sharing an unnamed global application partition. Application identities must be -stable across restarts and distinct between independent applications. Agent versions and -per-process identifiers are not used as namespace defaults. - -New session IDs retain the existing hosted/local representation. Restored sessions keep -their stored ID: the generic `SessionStore` contract does not require that ID to match a -lookup key or the request's platform ID. -Same-partition requests serialize within one host. Reference-counted lock entries exist -only for active/waiting requests and are removed after the final request leaves. - -The coordination scope includes stream closure, factory-agent cleanup, and session saving. -The host attempts to save state on success, execution failure, and interruption, without -retrying agent execution. Disconnect cleanup is shielded from the ASGI cancellation scope. -Persistence errors are reported; there is no success-shaped in-memory fallback. -Combined execution/persistence failures retain both messages and the persistence cause. -After streaming headers have been sent, a persistence failure terminates the stream and is -logged; already delivered text cannot be retracted. - -AgentServer can replace the response iterator with an SSE keep-alive wrapper whose pump -is still reading the agent stream. Disconnect cleanup closes the installed outer iterator -first, allowing that wrapper to cancel and join its pump, then closes the original stream. - -Starlette uses AnyIO cancellation scopes even when running on asyncio. `asyncio.shield()` -protects the child operation but still allows the awaiting request to unwind on cancellation. -Using `CancelScope(shield=True)` instead keeps cleanup awaited inside the request's ownership -scope before releasing session coordination. This does not replace the asyncio event loop. -AnyIO is already a transitive dependency and is declared directly because hosting imports it. - -This removes cumulative host retention of completed session objects. It does not bound -individual session size, active request count, backend storage, or total process memory. -The SDK local backend reads its entire logical-store file per operation. The store API -does not provide distributed transactions, so concurrent updates across hosts still -require application-level coordination. Durable serialization also replaces Python -object-identity reuse with restored state, and custom state needs registered codecs. - -### Migration and provider configuration - -Existing process-local invocation sessions are not migrated across deployment; the previous -implementation also lost them on restart. Local sessions now survive host recreation and are -written under `AGENTSERVER_STATE_ROOT` or the SDK's default state directory. Independent -applications should configure distinct application namespaces. - -State must support `AgentSession.to_dict()` / `AgentSession.from_dict()`. Custom state types -need registered codecs rather than arbitrary live Python objects. Restored sessions preserve -state, not Python object identity. - -A custom provider's `get_store` receives host configuration and request platform context. -Providers explicitly supplied by the application are responsible for application and protocol -isolation, and may use their own scoping without `session_store_namespace`. -The provider owns retention and deletion; an explicitly selected in-memory store remains -volatile and does not gain automatic eviction. For new default stores, writes renew the -30-day expiry window and reads do not. Missing, deleted, or expired state starts fresh on -the next invocation, including when its caller reuses the same session ID. diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 054728f1355..8fc3a122fe4 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -122,8 +122,6 @@ representation. Repeated requests for the same identifier pair restore the saved Locally, the platform session ID is used unchanged. Both hosts accept `agent_session_store_provider` to select a `StoreProvider[SessionStore]`. -For default Invocations storage, set `session_store_namespace` to a stable application-specific -name unless `FOUNDRY_AGENT_ID` or `FOUNDRY_AGENT_NAME` is configured. Session state must support `AgentSession` serialization. Use `register_state_type()` codecs for custom types; unsupported live objects fail during persistence. Restored sessions preserve state, not Python object identity. New default stores expire saved sessions 30 days after @@ -131,9 +129,9 @@ their last write; an invocation after expiry starts a fresh session. Existing stores retain their creation-time settings, and custom providers own their retention policies. -Same-session requests are coordinated only within one host instance. Multi-worker deployments -must coordinate overlapping requests for the same session; cross-host transactions and -exactly-once execution are not guaranteed. +Applications must coordinate overlapping requests for the same session; the store does not +provide transactions or exactly-once execution. Independent local applications should use +separate state roots or store providers. ### Workflow checkpoints diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index 695651fbbd5..1f379d36a17 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -2,18 +2,14 @@ from __future__ import annotations -import asyncio -import hashlib import json import logging +import sys from collections.abc import Awaitable, Callable from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager -from dataclasses import dataclass, field -from types import TracebackType from agent_framework import AgentSession, ResponseStream, SessionStore, SupportsAgentRun from agent_framework._telemetry import mark_feature_used -from anyio import CancelScope from azure.ai.agentserver.core import FoundryAgentRequestContext, get_request_context from azure.ai.agentserver.invocations import InvocationAgentServerHost from starlette.requests import Request @@ -28,20 +24,6 @@ logger = logging.getLogger(__name__) -@dataclass -class _SessionLock: - lock: asyncio.Lock = field(default_factory=asyncio.Lock) - users: int = 0 - - -class _ShieldedAsyncExitStack(AsyncExitStack): - async def __aexit__( - self, exc_type: type[BaseException] | None, exc: BaseException | None, traceback: TracebackType | None - ) -> bool | None: - with CancelScope(shield=True): - return await super().__aexit__(exc_type, exc, traceback) - - class _InvocationStreamingResponse(StreamingResponse): def __init__(self, content: AsyncGenerator[str]) -> None: super().__init__( @@ -55,15 +37,14 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: try: await super().__call__(scope, receive, send) finally: - # Sending a chunk can fail while the generator is suspended at yield. - with CancelScope(shield=True): - try: - close = getattr(self.body_iterator, "aclose", None) - if close is not None: - await close() - finally: - if self.body_iterator is not self._content: - await self._content.aclose() + # The SDK may wrap the iterator in a keep-alive pump; join it before closing its source. + try: + close = getattr(self.body_iterator, "aclose", None) + if close is not None: + await close() + finally: + if self.body_iterator is not self._content: + await self._content.aclose() class InvocationsHostServer(InvocationAgentServerHost): @@ -75,7 +56,6 @@ def __init__( *, openapi_spec: dict[str, Any] | None = None, agent_session_store_provider: StoreProvider[SessionStore] | None = None, - session_store_namespace: str | None = None, **kwargs: Any, ) -> None: """Initialize an InvocationsHostServer. @@ -87,42 +67,21 @@ def __init__( agent_session_store_provider: Provider for conversation session storage. Defaults to Foundry storage when hosted and the SDK's file-backed storage locally. New default stores expire sessions 30 days after their last write. Custom providers control their own retention. - session_store_namespace: Stable application identity for session storage. Defaults to the configured - Foundry agent ID or name. Required for the default store when neither is configured, including - unnamed local applications. Keep it stable across restarts and distinct for independent applications. - Custom providers without a namespace are responsible for application and protocol isolation. **kwargs: Additional keyword arguments. This host will expect the request to be a JSON body with a "message" field. The response contains the agent's text, or streamed text when "stream" is true. """ validate_agent_source(agent) - if session_store_namespace is not None and ( - not isinstance(session_store_namespace, str) or not session_store_namespace.strip() - ): - raise ValueError("session_store_namespace must be a non-empty string") super().__init__(openapi_spec=openapi_spec, **kwargs) self._agent = agent self._owns_request_agent = not is_agent(agent) - self._session_namespace = ( - session_store_namespace - if session_store_namespace is not None - else self.config.agent_guid or self.config.agent_name or None - ) - if agent_session_store_provider is None and ( - self._session_namespace is None or not self._session_namespace.strip() - ): - raise ValueError( - "The default invocation session store requires a stable application identity. " - "Set session_store_namespace or configure FOUNDRY_AGENT_ID or FOUNDRY_AGENT_NAME." - ) self._session_store_provider = ( AgentSessionStoreProvider(store_name="invocation_sessions") if agent_session_store_provider is None else agent_session_store_provider ) - self._session_locks: dict[str | tuple[str, str], _SessionLock] = {} self.invoke_handler(self._handle_invoke) mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) @@ -160,7 +119,7 @@ def _partition_key(self) -> str | tuple[str, str]: @asynccontextmanager async def _request_agent(self) -> AsyncGenerator[SupportsAgentRun]: agent = await resolve_agent(self._agent) - async with _ShieldedAsyncExitStack() as resources: + async with AsyncExitStack() as resources: if self._owns_request_agent and isinstance(agent, AbstractAsyncContextManager): await resources.enter_async_context(agent) yield agent @@ -169,53 +128,32 @@ async def _request_agent(self) -> AsyncGenerator[SupportsAgentRun]: async def _request_session( self, partition_key: str | tuple[str, str], context: FoundryAgentRequestContext ) -> AsyncGenerator[AgentSession]: - entry = self._session_locks.get(partition_key) - if entry is None: - entry = self._session_locks[partition_key] = _SessionLock() - entry.users += 1 + session_id = ( + json.dumps(partition_key, separators=(",", ":")) if isinstance(partition_key, tuple) else partition_key + ) try: - async with entry.lock: - session_id = ( - json.dumps(partition_key, separators=(",", ":")) - if isinstance(partition_key, tuple) - else partition_key - ) - encoded_key = json.dumps((self._session_namespace, partition_key), separators=(",", ":")) - storage_key = "invocations:v1:" + hashlib.sha256(encoded_key.encode("utf-8")).hexdigest() - try: - store = self._session_store_provider.get_store(config=self.config, platform_context=context) - session = await store.get(storage_key) - if session is None: - session = AgentSession(session_id=session_id) - except Exception: - logger.exception("Failed to load invocation session") - raise - - async def save_session( - exc_type: type[BaseException] | None, failure: BaseException | None, traceback: TracebackType | None - ) -> None: - try: - await store.set(storage_key, session) - except Exception as exc: - logger.exception("Failed to persist invocation session") - if failure is None: - raise - if isinstance(failure, Exception): - raise RuntimeError( - f"Invocation failed: {str(failure) or type(failure).__name__}; " - f"session persistence also failed: {str(exc) or type(exc).__name__}" - ) from exc - except asyncio.CancelledError: - logger.error("Invocation session persistence was cancelled") - raise - - async with _ShieldedAsyncExitStack() as cleanup: - cleanup.push_async_exit(save_session) - yield session + store = self._session_store_provider.get_store(config=self.config, platform_context=context) + session = await store.get(session_id) + if session is None: + session = AgentSession(session_id=session_id) + except Exception: + logger.exception("Failed to load invocation session") + raise + try: + yield session finally: - entry.users -= 1 - if entry.users == 0: - del self._session_locks[partition_key] + failure = sys.exc_info()[1] + try: + await store.set(session_id, session) + except Exception as exc: + logger.exception("Failed to persist invocation session") + if failure is None: + raise + if isinstance(failure, Exception): + raise RuntimeError( + f"Invocation failed: {str(failure) or type(failure).__name__}; " + f"session persistence also failed: {str(exc) or type(exc).__name__}" + ) from exc async def _handle_invoke(self, request: Request) -> Response: """Invoke the agent with the given request.""" @@ -249,13 +187,12 @@ async def stream_response() -> AsyncGenerator[str]: if update.text: yield update.text finally: - with CancelScope(shield=True): - if isinstance(stream, ResponseStream): - await stream.close() - else: - close = getattr(stream, "aclose", None) - if close is not None: - await close() + if isinstance(stream, ResponseStream): + await stream.close() + else: + close = getattr(stream, "aclose", None) + if close is not None: + await close() return _InvocationStreamingResponse(stream_response()) diff --git a/python/packages/foundry_hosting/pyproject.toml b/python/packages/foundry_hosting/pyproject.toml index f63389f9bac..3dd6acefa79 100644 --- a/python/packages/foundry_hosting/pyproject.toml +++ b/python/packages/foundry_hosting/pyproject.toml @@ -24,7 +24,6 @@ classifiers = [ ] dependencies = [ "agent-framework-core>=1.19.0,<2", - "anyio>=4,<5", "azure-ai-agentserver-core>=2.1.0,<3", "azure-ai-agentserver-responses>=2.2.0b1,<3", "azure-ai-agentserver-invocations>=1.1.0,<2", diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index 0841a7d49fc..0849a96031f 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -20,7 +20,6 @@ from typing import cast from unittest.mock import AsyncMock, MagicMock, patch -import httpx import pytest from agent_framework import ( Agent, @@ -40,7 +39,6 @@ SessionStore, tool, ) -from anyio import CancelScope from azure.ai.agentserver.core import ( AgentConfig, FoundryAgentRequestContext, @@ -59,13 +57,6 @@ # region Helpers -@pytest.fixture(autouse=True) -def configure_local_application(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("FOUNDRY_AGENT_NAME", "invocation-test-application") - monkeypatch.delenv("FOUNDRY_AGENT_ID", raising=False) - monkeypatch.delenv("FOUNDRY_HOSTING_ENVIRONMENT", raising=False) - - def _mock_session_store() -> MagicMock: store = MagicMock(spec=SessionStore) store.get.return_value = None @@ -266,24 +257,9 @@ def update_session(session: AgentSession) -> None: assert restored is not None assert restored.state == {"protocol": "responses"} - async def test_independent_application_names_do_not_share_sessions(self, monkeypatch: pytest.MonkeyPatch) -> None: - seen: list[dict[str, Any]] = [] - - for application in ["first-app", "second-app", "first-app"]: - monkeypatch.setenv("FOUNDRY_AGENT_NAME", application) - - def update_session(session: AgentSession, application: str = application) -> None: - seen.append(dict(session.state)) - session.state["application"] = application - - server = InvocationsHostServer(_make_agent(response_text="ok", update_session=update_session)) - with _request_context(session_id="shared-session"): - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - - assert seen == [{}, {}, {"application": "first-app"}] - @pytest.mark.parametrize("spec_version", ["2.0", "2.4"]) - async def test_keep_alive_disconnect_joins_pump_before_session_cleanup(self, spec_version: str) -> None: + @pytest.mark.parametrize("keep_alive", [False, True]) + async def test_disconnect_closes_stream_before_session_save(self, spec_version: str, keep_alive: bool) -> None: entered = asyncio.Event() disconnected = asyncio.Event() saved: list[dict[str, Any]] = [] @@ -305,6 +281,8 @@ async def updates() -> AsyncIterator[AgentResponseUpdate]: pump = asyncio.current_task() entered.set() try: + if not keep_alive: + yield AgentResponseUpdate(contents=[Content.from_text("first")]) await asyncio.Event().wait() yield AgentResponseUpdate(contents=[Content.from_text("unused")]) finally: @@ -314,7 +292,7 @@ async def updates() -> AsyncIterator[AgentResponseUpdate]: return updates() server = InvocationsHostServer(IdleAgent(), agent_session_store_provider=_SessionStoreProvider(store)) - server.config.sse_keepalive_interval = 1 + server.config.sse_keepalive_interval = 1 if keep_alive else 0 body = json.dumps({"message": "hello", "stream": True}).encode() received = False @@ -327,7 +305,8 @@ async def receive() -> Any: return {"type": "http.disconnect"} async def send(message: Any) -> None: - if message["type"] == "http.response.body" and message.get("body") == b": keep-alive\n\n": + expected = b": keep-alive\n\n" if keep_alive else b"first" + if message["type"] == "http.response.body" and message.get("body") == expected: assert entered.is_set() if spec_version == "2.4": raise OSError("disconnected while idle") @@ -353,39 +332,14 @@ async def send(message: Any) -> None: else: await asyncio.wait_for(server(scope, receive, send), timeout=5) assert saved == [{"closed": True}] - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] assert pump is not None - assert pump.done() + if keep_alive: + assert pump.done() finally: - if pump is not None and not pump.done(): + if keep_alive and pump is not None and not pump.done(): pump.cancel() await asyncio.gather(pump, return_exceptions=True) - @pytest.mark.parametrize("stream", [False, True]) - async def test_http_requests_restore_sessions_without_changing_response_format(self, stream: bool) -> None: - turns: list[int] = [] - - def update_session(session: AgentSession) -> None: - session.state["turn"] = session.state.get("turn", 0) + 1 - turns.append(session.state["turn"]) - - for _ in range(2): - server = InvocationsHostServer( - _make_agent(response_text="ok", stream_texts=["o", "k"], update_session=update_session) - ) - async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: - response = await client.post( - "/invocations", - params={"agent_session_id": "session"}, - json={"message": "hello", "stream": stream}, - ) - assert response.status_code == 200 - assert response.text == "ok" - if stream: - assert response.headers["content-type"].startswith("text/event-stream") - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - assert turns == [1, 2] - @pytest.mark.parametrize("stream", [False, True]) async def test_function_history_is_restored_without_reexecuting_completed_tool(self, stream: bool) -> None: executed: list[str] = [] @@ -455,124 +409,8 @@ async def test_custom_store_can_restore_a_different_session_id(self) -> None: await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] assert agent.calls[0]["session"].session_id == "different" store.set.assert_awaited_once() - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - async def test_combined_failures_include_types_for_empty_messages(self) -> None: - store = _mock_session_store() - store.set.side_effect = RuntimeError() - - def update_session(session: AgentSession) -> None: - raise ValueError - - server = InvocationsHostServer( - _make_agent(response_text="ok", update_session=update_session), - agent_session_store_provider=_SessionStoreProvider(store), - ) - with ( - _request_context(session_id="session"), - pytest.raises( - RuntimeError, match="Invocation failed: ValueError; session persistence also failed: RuntimeError" - ) as raised, - ): - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - assert raised.value.__cause__ is store.set.side_effect - - async def test_control_flow_exception_still_closes_request_agent(self) -> None: - class ControlFlowExit(BaseException): - pass - - events: list[str] = [] - error = ControlFlowExit() - store = _mock_session_store() - store.set.side_effect = ValueError("save failed") - - class ExitingAgent(_ContextAgent): - def run( - self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any - ) -> Any: - raise error - - server = InvocationsHostServer( - lambda: ExitingAgent(events), agent_session_store_provider=_SessionStoreProvider(store) - ) - with _request_context(session_id="session"), pytest.raises(ControlFlowExit) as raised: - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - assert raised.value is error - assert events == ["enter", "exit"] - store.set.assert_awaited_once() - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - async def test_cancelled_persistence_is_reported_and_releases_coordination( - self, caplog: pytest.LogCaptureFixture - ) -> None: - store = _mock_session_store() - store.set = AsyncMock(side_effect=asyncio.CancelledError) - server = InvocationsHostServer( - _make_agent(response_text="ok"), agent_session_store_provider=_SessionStoreProvider(store) - ) - with _request_context(session_id="session"), pytest.raises(asyncio.CancelledError): - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - assert "Invocation session persistence was cancelled" in caplog.text - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - @pytest.mark.parametrize("stream", [False, True]) - async def test_factory_failure_does_not_retain_session(self, stream: bool) -> None: - def create_agent() -> _FakeAgent: - raise ValueError("factory failed") - - server = InvocationsHostServer(create_agent) - with _request_context(session_id="session"), pytest.raises(ValueError, match="factory failed"): - response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] - _make_request({"message": "hello", "stream": stream}) - ) - if isinstance(response, StreamingResponse): - await _collect_stream(response) - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - async def test_stream_close_failure_is_reported_after_persistence(self) -> None: - store = _mock_session_store() - store.set = AsyncMock() - agent = _make_agent(stream_texts=["ok"]) - iterator = MagicMock() - iterator.__aiter__ = MagicMock(return_value=iterator) - iterator.__anext__ = AsyncMock(side_effect=StopAsyncIteration) - iterator.aclose = AsyncMock(side_effect=ValueError("close failed")) - server = InvocationsHostServer(agent, agent_session_store_provider=_SessionStoreProvider(store)) - with patch.object(agent, "run", return_value=iterator), _request_context(session_id="session"): - response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] - _make_request({"message": "hello", "stream": True}) - ) - assert isinstance(response, StreamingResponse) - with pytest.raises(ValueError, match="close failed"): - await _collect_stream(response) - store.set.assert_awaited_once() - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - async def test_custom_store_receives_request_context_and_namespaced_bounded_keys(self) -> None: - store = _mock_session_store() - store.get = AsyncMock(return_value=None) - store.set = AsyncMock() - provider = _SessionStoreProvider(store) - server = InvocationsHostServer(_make_agent(response_text="ok"), agent_session_store_provider=provider) - - identities = ["x" * 1024, '["session","user"]'] - for identity in identities: - with _request_context(call_id="call", session_id=identity): - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - server.config.is_hosted = True - with _request_context(call_id="hosted-call", session_id="session", user_id="user"): - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - - keys = [call.args[0] for call in store.set.await_args_list] - assert len(set(keys)) == 3 - assert all(key.startswith("invocations:v1:") and len(key) <= 128 for key in keys) - assert [call.args[1].session_id for call in store.set.await_args_list] == [ - *identities, - '["session","user"]', - ] - assert provider.contexts[-1].user_id == "user" - assert provider.contexts[-1].call_id == "hosted-call" - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] + assert store.get.call_args.args == ("session",) + assert store.set.call_args.args[0] == "session" async def test_unconsumed_stream_and_invalid_input_do_not_access_storage(self) -> None: store = _mock_session_store() @@ -589,99 +427,6 @@ async def test_unconsumed_stream_and_invalid_input_do_not_access_storage(self) - await cast(Any, response.body_iterator).aclose() assert provider.contexts == [] assert agent.calls == [] - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - @pytest.mark.parametrize("cancel_waiter", [False, True]) - @pytest.mark.parametrize("stream", [False, True]) - async def test_same_session_serializes_and_waiter_cleanup_is_bounded( - self, cancel_waiter: bool, stream: bool - ) -> None: - entered = asyncio.Event() - release = asyncio.Event() - second_started = asyncio.Event() - turns: list[int] = [] - - class BlockingAgent(_FakeAgent): - def run( - self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any - ) -> Any: - async def complete() -> AgentResponse: - assert session is not None - turn = session.state.get("turn", 0) + 1 - if not turns: - entered.set() - await release.wait() - session.state["turn"] = turn - turns.append(turn) - return AgentResponse(messages=[Message("assistant", "ok")]) - - async def updates() -> AsyncIterator[AgentResponseUpdate]: - await complete() - yield AgentResponseUpdate(contents=[Content.from_text("ok")]) - - return updates() if stream else complete() - - server = InvocationsHostServer(BlockingAgent()) - - async def invoke(*, second: bool = False) -> Response: - with _request_context(session_id="session"): - if second: - second_started.set() - response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] - _make_request({"message": "hello", "stream": stream}) - ) - if isinstance(response, StreamingResponse): - await _collect_stream(response) - return response - - first = asyncio.create_task(invoke()) - await entered.wait() - second = asyncio.create_task(invoke(second=True)) - try: - await second_started.wait() - assert not second.done() - assert server._session_locks["session"].users == 2 # pyright: ignore[reportPrivateUsage] - if cancel_waiter: - second.cancel() - with pytest.raises(asyncio.CancelledError): - await second - assert server._session_locks["session"].users == 1 # pyright: ignore[reportPrivateUsage] - release.set() - await first - if not cancel_waiter: - await second - finally: - release.set() - await asyncio.gather(first, second, return_exceptions=True) - assert turns == ([1] if cancel_waiter else [1, 2]) - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - async def test_different_sessions_run_concurrently(self) -> None: - entered: set[str] = set() - both_entered = asyncio.Event() - - class ConcurrentAgent(_FakeAgent): - def run( - self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any - ) -> Any: - async def complete() -> AgentResponse: - assert session is not None - entered.add(session.session_id) - if len(entered) == 2: - both_entered.set() - await both_entered.wait() - return AgentResponse(messages=[Message("assistant", "ok")]) - - return complete() - - server = InvocationsHostServer(ConcurrentAgent()) - - async def invoke(session_id: str) -> None: - with _request_context(session_id=session_id): - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - - await asyncio.wait_for(asyncio.gather(invoke("one"), invoke("two")), timeout=5) - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("failure", ["load", "save", "run", "run_and_save"]) @@ -722,7 +467,6 @@ def update_session(session: AgentSession) -> None: assert store.set.await_args.args[1].state == {"turn": 1} if "save" in failure: assert "Failed to persist invocation session" in caplog.text - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] @pytest.mark.parametrize("fail_save", [False, True]) async def test_cancellation_preserves_state_and_cancellation_signal( @@ -756,157 +500,17 @@ async def complete() -> AgentResponse: await task assert store.set.await_args is not None assert store.set.await_args.args[1].state == {"started": True} - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] if fail_save: assert "Failed to persist invocation session" in caplog.text - @pytest.mark.parametrize("spec_version", ["2.0", "2.4"]) - async def test_disconnect_closes_agent_and_saves_final_state(self, spec_version: str) -> None: - sent = asyncio.Event() - events: list[str] = [] - store = _mock_session_store() - - class StreamingAgent(_ContextAgent): - async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: - await asyncio.sleep(0) - events.append("exit") - - def run( - self, messages: Any = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any - ) -> Any: - async def updates() -> AsyncIterator[AgentResponseUpdate]: - assert session is not None - try: - yield AgentResponseUpdate(contents=[Content.from_text("first")]) - await asyncio.Event().wait() - finally: - await asyncio.sleep(0) - session.state["closed"] = True - events.append("close") - - return updates() - - server = InvocationsHostServer( - lambda: StreamingAgent(events), agent_session_store_provider=_SessionStoreProvider(store) - ) - with _request_context(session_id="session"): - response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] - _make_request({"message": "hello", "stream": True}) - ) - - async def send(message: Any) -> None: - if message["type"] == "http.response.body": - sent.set() - if spec_version == "2.4": - raise OSError("disconnected") - await asyncio.Event().wait() - - async def receive() -> Any: - await sent.wait() - return {"type": "http.disconnect"} - - scope = {"type": "http", "asgi": {"spec_version": spec_version}} - if spec_version == "2.4": - with pytest.raises(ClientDisconnect): - await response(scope, receive, send) - else: - await response(scope, receive, send) - assert events == ["enter", "close", "exit"] - assert store.set.await_args is not None - assert store.set.await_args.args[1].state == {"closed": True} - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - async def test_cancelled_scope_still_awaits_persistence(self) -> None: - store = _mock_session_store() - saved = False - - async def save(session_id: str, session: AgentSession) -> None: - nonlocal saved - await asyncio.sleep(0) - saved = True - - store.set = AsyncMock(side_effect=save) - server = InvocationsHostServer( - _make_agent(response_text="ok"), agent_session_store_provider=_SessionStoreProvider(store) - ) - with _request_context(session_id="session"): - with CancelScope() as scope: - async with server._request_session("session", FoundryAgentRequestContext()): # pyright: ignore[reportPrivateUsage] - scope.cancel() - await asyncio.sleep(0) - assert saved - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - - async def test_default_store_expiry_is_renewed_by_writes(self) -> None: - turns: list[int] = [] - - def update_session(session: AgentSession) -> None: - session.state["turn"] = session.state.get("turn", 0) + 1 - turns.append(session.state["turn"]) - - server = InvocationsHostServer(_make_agent(response_text="ok", update_session=update_session)) - day = 24 * 60 * 60 - with ( - _request_context(session_id="session"), - patch("azure.ai.agentserver.core.storage._local_state._now") as now, - ): - for timestamp in [0, 29 * day, 58 * day, 88 * day]: - now.return_value = timestamp - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - assert turns == [1, 2, 3, 1] - - async def test_unserializable_state_fails_explicitly(self) -> None: - def update_session(session: AgentSession) -> None: - session.state["unsupported"] = object() - - server = InvocationsHostServer(_make_agent(response_text="ok", update_session=update_session)) - with ( - _request_context(session_id="session"), - pytest.warns(RuntimeWarning, match="unsupported type 'object'"), - pytest.raises(TypeError), - ): - await server._handle_invoke(_make_request({"message": "hello"})) # pyright: ignore[reportPrivateUsage] - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] - # region Initialization class TestInit: - def test_default_store_requires_application_identity(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("FOUNDRY_AGENT_NAME") - with pytest.raises(ValueError, match="Set session_store_namespace"): - InvocationsHostServer(_make_agent(response_text="ok")) - - server = InvocationsHostServer(_make_agent(response_text="ok"), session_store_namespace="explicit-app") - assert server._session_namespace == "explicit-app" # pyright: ignore[reportPrivateUsage] - - def test_custom_store_can_own_application_isolation(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("FOUNDRY_AGENT_NAME") - server = InvocationsHostServer( - _make_agent(response_text="ok"), agent_session_store_provider=_SessionStoreProvider(_mock_session_store()) - ) - assert server._session_namespace is None # pyright: ignore[reportPrivateUsage] - - @pytest.mark.parametrize("namespace", ["", " \t", 123]) - def test_rejects_invalid_explicit_namespace(self, namespace: Any) -> None: - with pytest.raises(ValueError, match="session_store_namespace must be a non-empty string"): - InvocationsHostServer(_make_agent(response_text="ok"), session_store_namespace=namespace) - - def test_namespace_precedence_is_explicit_then_agent_id_then_name(self, monkeypatch: pytest.MonkeyPatch) -> None: - agent = _make_agent(response_text="ok") - named = InvocationsHostServer(agent) - assert named._session_namespace == "invocation-test-application" # pyright: ignore[reportPrivateUsage] - monkeypatch.setenv("FOUNDRY_AGENT_ID", "stable-agent-id") - identified = InvocationsHostServer(agent) - assert identified._session_namespace == "stable-agent-id" # pyright: ignore[reportPrivateUsage] - explicit = InvocationsHostServer(agent, session_store_namespace="explicit-app") - assert explicit._session_namespace == "explicit-app" # pyright: ignore[reportPrivateUsage] - def test_accepts_supports_agent_run(self) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) assert server._agent is not None # pyright: ignore[reportPrivateUsage] - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] @pytest.mark.parametrize("agent", [None, 42]) def test_rejects_invalid_agent_source(self, agent: Any) -> None: @@ -1119,7 +723,6 @@ async def test_reusing_session_restores_identifier(self, hosted: bool, stream: b assert agent.calls[0]["session"] is not agent.calls[1]["session"] assert agent.calls[0]["session"].session_id == agent.calls[1]["session"].session_id == expected_id session_factory.assert_called_once_with(session_id=expected_id) - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] async def test_missing_message_returns_400(self) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) @@ -1160,7 +763,6 @@ async def test_non_streaming_returns_agent_text(self) -> None: assert agent.calls[0]["messages"] == ["Hi"] assert agent.calls[0]["stream"] is False assert agent.calls[0]["session"].session_id == "sess-1" - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] async def test_streaming_yields_update_text(self) -> None: agent = _make_agent(stream_texts=["Hel", "lo", "!"]) @@ -1192,7 +794,6 @@ def update_session(session: AgentSession) -> None: assert first_session is not second_session assert first_session.state == {"turn": 1} assert second_session.state == {"turn": 2} - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize( @@ -1239,7 +840,6 @@ def update_session(session: AgentSession) -> None: assert sessions[0] is not sessions[1] assert sessions[0].session_id != sessions[1].session_id - assert server._session_locks == {} # pyright: ignore[reportPrivateUsage] for (session_id, user_id), session in zip(identifiers, sessions): with _request_context(call_id="call-2", session_id=session_id, user_id=user_id): diff --git a/python/packages/foundry_hosting/tests/test_invocations_int.py b/python/packages/foundry_hosting/tests/test_invocations_int.py index 902ce319a88..136760da74f 100644 --- a/python/packages/foundry_hosting/tests/test_invocations_int.py +++ b/python/packages/foundry_hosting/tests/test_invocations_int.py @@ -57,7 +57,7 @@ def server() -> InvocationsHostServer: default_options={"store": False}, # pyrefly: ignore[bad-argument-type] ) - return InvocationsHostServer(agent, session_store_namespace="invocations-integration") + return InvocationsHostServer(agent) @tool @@ -78,7 +78,7 @@ def server_with_tools() -> InvocationsHostServer: default_options={"store": False}, # pyrefly: ignore[bad-argument-type] ) - return InvocationsHostServer(agent, session_store_namespace="invocations-with-tools-integration") + return InvocationsHostServer(agent) # --------------------------------------------------------------------------- diff --git a/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py index 7b859f39123..730600d4f98 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/main.py @@ -27,7 +27,7 @@ def main(): default_options={"store": False}, ) - server = InvocationsHostServer(agent, session_store_namespace="basic-invocations") + server = InvocationsHostServer(agent) server.run() diff --git a/python/uv.lock b/python/uv.lock index 5666bf392b0..65fe4e5f906 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -653,7 +653,6 @@ version = "1.0.0b260918" source = { editable = "packages/foundry_hosting" } dependencies = [ { name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, - { name = "anyio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "azure-ai-agentserver-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "azure-ai-agentserver-invocations", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "azure-ai-agentserver-responses", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -664,7 +663,6 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "agent-framework-core", editable = "packages/core" }, - { name = "anyio", specifier = ">=4,<5" }, { name = "azure-ai-agentserver-core", specifier = ">=2.1.0,<3" }, { name = "azure-ai-agentserver-invocations", specifier = ">=1.1.0,<2" }, { name = "azure-ai-agentserver-responses", specifier = ">=2.2.0b1,<3" },