diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index 96b41d6d38f..583a93fb16c 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -699,6 +699,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 a7a59803df4..8fc3a122fe4 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/` @@ -103,9 +103,10 @@ 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. 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. @@ -114,11 +115,23 @@ 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 +`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 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. + +Both hosts accept `agent_session_store_provider` to select a `StoreProvider[SessionStore]`. +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. + +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 0cf7a6e97e0..1f379d36a17 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,50 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import json +import logging +import sys from collections.abc import Awaitable, Callable from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager -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 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__) + + +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: + # 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): @@ -24,6 +55,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 +64,24 @@ 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(store_name="invocation_sessions") + if agent_session_store_provider is None + else agent_session_store_provider + ) self.invoke_handler(self._handle_invoke) mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) @@ -86,6 +124,37 @@ async def _request_agent(self) -> AsyncGenerator[SupportsAgentRun]: await resources.enter_async_context(agent) yield agent + @asynccontextmanager + async def _request_session( + self, partition_key: str | tuple[str, str], context: FoundryAgentRequestContext + ) -> AsyncGenerator[AgentSession]: + session_id = ( + json.dumps(partition_key, separators=(",", ":")) if isinstance(partition_key, tuple) else partition_key + ) + try: + 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: + 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.""" try: @@ -103,18 +172,15 @@ 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: @@ -128,12 +194,11 @@ async def stream_response() -> AsyncGenerator[str]: if close is not None: await close() - return StreamingResponse( - stream_response(), - media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, - ) + return _InvocationStreamingResponse(stream_response()) - async with self._request_agent() as agent: + 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/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 250799849a6..0849a96031f 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -11,8 +11,10 @@ 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 @@ -20,27 +22,57 @@ import pytest from agent_framework import ( + Agent, AgentResponse, AgentResponseUpdate, AgentSession, + BaseChatClient, + ChatMiddlewareLayer, + ChatResponse, + ChatResponseUpdate, Content, + FunctionInvocationLayer, + InMemoryHistoryProvider, Message, + ResponseStream, ServiceSessionId, + SessionStore, + tool, ) 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 +from agent_framework_foundry_hosting._state_store import FoundryAgentSessionStore + +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 +86,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 +94,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 +105,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 +178,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 +187,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 +223,287 @@ async def _collect_stream(response: StreamingResponse) -> str: # endregion +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"} + + @pytest.mark.parametrize("spec_version", ["2.0", "2.4"]) + @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]] = [] + 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: + if not keep_alive: + yield AgentResponseUpdate(contents=[Content.from_text("first")]) + 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 if keep_alive else 0 + 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: + 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") + 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 pump is not None + if keep_alive: + assert pump.done() + finally: + 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_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_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"): + 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 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() + 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 == [] + + @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: run failed; session persistence also failed: save 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 + + @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} + if fail_save: + assert "Failed to persist invocation session" in caplog.text + + # region Initialization @@ -193,7 +511,6 @@ 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] @pytest.mark.parametrize("agent", [None, 42]) def test_rejects_invalid_agent_source(self, agent: Any) -> None: @@ -285,6 +602,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 +697,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 +711,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,9 +720,8 @@ 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) async def test_missing_message_returns_400(self) -> None: @@ -411,10 +759,10 @@ 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" async def test_streaming_yields_update_text(self) -> None: agent = _make_agent(stream_texts=["Hel", "lo", "!"]) @@ -429,19 +777,23 @@ 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} @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize( @@ -461,7 +813,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 +835,11 @@ 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] 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 +852,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/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")