diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 990c55c4ae9..eef2e504259 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -3520,6 +3520,24 @@ async def __anext__(self) -> UpdateT: update = await update return await self._record_update(update) + async def close(self) -> None: + """Close the active iterator and run cleanup hooks. + + This method is idempotent and also closes nested ``ResponseStream`` wrappers. + """ + try: + iterator: AsyncIterator[UpdateT] | None = self._iterator + if iterator is not None: + if isinstance(iterator, ResponseStream): + await cast(ResponseStream[UpdateT, Any], iterator).close() + else: + close = getattr(iterator, "aclose", None) + if close is not None: + await close() + finally: + self._consumed = True + await self._run_cleanup_hooks() + async def _resolve_stream_with_pull_contexts(self) -> AsyncIterable[UpdateT]: """Resolve the underlying stream while activating any registered pull context managers. diff --git a/python/packages/core/agent_framework/_workflows/_functional.py b/python/packages/core/agent_framework/_workflows/_functional.py index 14f7a132735..ae94a98692f 100644 --- a/python/packages/core/agent_framework/_workflows/_functional.py +++ b/python/packages/core/agent_framework/_workflows/_functional.py @@ -49,6 +49,7 @@ from copy import deepcopy from typing import Any, Generic, Literal, TypeVar, overload +from .._agents import BaseAgent from .._feature_stage import ExperimentalFeature, experimental from .._serialization import make_json_safe from .._types import AgentResponse, AgentResponseUpdate, ResponseStream @@ -1407,7 +1408,7 @@ def _decorator(fn: Callable[..., Awaitable[Any]]) -> FunctionalWorkflowDefinitio @experimental(feature_id=ExperimentalFeature.FUNCTIONAL_WORKFLOWS) -class FunctionalWorkflowAgent: +class FunctionalWorkflowAgent(BaseAgent): """Agent adapter for a :class:`FunctionalWorkflow`. Provides a ``run()`` method with the same overloaded signature as @@ -1452,10 +1453,13 @@ def __init__( # but not otherwise consumed. del kwargs self._workflow = workflow - self.name = name or workflow.name - self.id = f"FunctionalWorkflowAgent_{self.name}" - self.description: str | None = description if description is not None else workflow.description - self.context_providers: Sequence[Any] | None = context_providers + resolved_name = name or workflow.name + super().__init__( + id=f"FunctionalWorkflowAgent_{resolved_name}", + name=resolved_name, + description=description if description is not None else workflow.description, + context_providers=context_providers, + ) self._pending_requests: dict[str, WorkflowEvent[Any]] = {} @property diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 846c45e0b2a..80843c5e0f8 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -4192,6 +4192,48 @@ async def async_hook(update: ChatResponseUpdate) -> ChatResponseUpdate: class TestResponseStreamCleanupHooks: """Tests for cleanup hooks (after stream consumption, before finalizer).""" + async def test_close_closes_iterator_and_runs_cleanup_once(self) -> None: + """Closing a partially consumed stream releases its iterator and cleanup hooks.""" + events: list[str] = [] + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + try: + yield ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant") + yield ChatResponseUpdate(contents=[Content.from_text("second")], role="assistant") + finally: + events.append("iterator") + + stream: ResponseStream[ChatResponseUpdate, Sequence[ChatResponseUpdate]] = ResponseStream( + updates(), cleanup_hooks=[lambda: events.append("cleanup")] + ) + await anext(stream) + + await stream.close() + await stream.close() + + assert events == ["iterator", "cleanup"] + + async def test_close_closes_wrapped_stream(self) -> None: + """Closing a wrapper releases the concrete inner stream.""" + events: list[str] = [] + + async def updates() -> AsyncIterable[ChatResponseUpdate]: + try: + yield ChatResponseUpdate(contents=[Content.from_text("first")], role="assistant") + yield ChatResponseUpdate(contents=[Content.from_text("second")], role="assistant") + finally: + events.append("iterator") + + inner: ResponseStream[ChatResponseUpdate, Sequence[ChatResponseUpdate]] = ResponseStream( + updates(), cleanup_hooks=[lambda: events.append("inner")] + ) + outer = inner.map(lambda update: update, _combine_updates).with_cleanup_hook(lambda: events.append("outer")) + await anext(outer) + + await outer.close() + + assert events == ["iterator", "inner", "outer"] + async def test_cleanup_hook_called_after_iteration(self) -> None: """Cleanup hook is called after iteration completes.""" cleanup_called = {"value": False} diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 6f60aba0395..43c0ef84e37 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -26,6 +26,7 @@ InMemoryCheckpointStorage, RunContext, StepWrapper, + SupportsAgentRun, WorkflowEvent, WorkflowEventSource, WorkflowRunResult, @@ -1051,6 +1052,13 @@ async def wf(x: int) -> int: assert agent.id == "FunctionalWorkflowAgent_my_agent" assert agent.description == "A test workflow" + async def test_as_agent_implements_supports_agent_run(self): + @built_workflow + async def wf(x: int) -> int: + return x + + assert isinstance(wf.as_agent(), SupportsAgentRun) + # --------------------------------------------------------------------------- # Concurrent execution guard diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 9f85be244b5..a7a59803df4 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -2,6 +2,33 @@ This package provides the integration of Agent Framework agents and workflows with the Foundry Agent Server, which can be hosted on Foundry infrastructure. +## Agent instances and factories + +`ResponsesHostServer` and `InvocationsHostServer` accept an agent instance or a zero-argument callable through the +existing `agent` parameter. The callable may be synchronous or asynchronous, must return an object implementing +`SupportsAgentRun`, and is invoked once for each request: + +```python +server = ResponsesHostServer(agent=create_agent) +``` + +Passing an instance reuses that object for the lifetime of the host. Pass a callable when the agent keeps mutable state +outside `AgentSession`. In particular, a `WorkflowAgent` wraps a stateful workflow, so its callable should build a new +workflow, executors, and wrapped agents. Keep the workflow name and executor IDs stable so later requests can find and +restore its checkpoints: + +```python +def create_agent(): + return build_workflow().as_agent() + + +server = ResponsesHostServer(agent=create_agent) +``` + +The Responses host continues regular agents through its existing session store and workflow agents through its +existing checkpoint store. A callable does not make arbitrary instance fields persistent; state needed by later +requests must remain in the supported stores. + ## Conversation history `ResponsesHostServer` uses AgentServer response history as the model's conversation history by default: @@ -21,8 +48,9 @@ AgentServer history requires a framework `RawAgent` whose client declares the bo the agent's runtime options then let hosting enforce downstream storage behavior. Custom `SupportsAgentRun` implementations must use `history_source="agent"` because that protocol does not accept runtime chat options. -`ResponsesHostServer` owns the supplied agent instance and may add hosting-specific context providers. Do not reuse that -agent with another host or invoke it directly after constructing the server. +`ResponsesHostServer` owns a supplied agent instance and may add hosting-specific context providers. Do not reuse that +instance with another host or invoke it directly after constructing the server. An agent returned by a callable belongs +to that request. To preserve the agent's regular history and service-storage behavior, select the agent as the history source: diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py new file mode 100644 index 00000000000..5f9509b3fef --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import inspect +from collections.abc import Awaitable, Callable +from typing import TypeAlias, TypeGuard, cast + +from agent_framework import SupportsAgentRun + +AgentSource: TypeAlias = SupportsAgentRun | Callable[[], SupportsAgentRun | Awaitable[SupportsAgentRun]] + + +def is_agent(value: object) -> TypeGuard[SupportsAgentRun]: + return not inspect.isclass(value) and isinstance(value, SupportsAgentRun) + + +def validate_agent_source(source: object) -> None: + if is_agent(source): + return + if not callable(source): + raise TypeError("agent must be an agent instance or a zero-argument callable that creates one.") + try: + inspect.signature(source).bind() + except (TypeError, ValueError) as exc: + raise TypeError("agent callable must accept no arguments.") from exc + + +async def resolve_agent(source: AgentSource) -> SupportsAgentRun: + """Resolve an agent instance or request-scoped agent factory.""" + if is_agent(source): + return source + + factory = cast(Callable[[], SupportsAgentRun | Awaitable[SupportsAgentRun]], source) + result = factory() + agent = await result if inspect.isawaitable(result) else result + if not is_agent(agent): + raise TypeError("The agent factory must return an object implementing SupportsAgentRun.") + return agent 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 cd539da4fec..0cf7a6e97e0 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,8 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. import json +from collections.abc import Awaitable, Callable +from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager -from agent_framework import AgentSession, SupportsAgentRun +from agent_framework import AgentSession, ResponseStream, SupportsAgentRun from agent_framework._telemetry import mark_feature_used from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.invocations import InvocationAgentServerHost @@ -10,6 +12,7 @@ from starlette.responses import Response, StreamingResponse from typing_extensions import Any, AsyncGenerator +from ._agent_source import is_agent, resolve_agent, validate_agent_source from ._feature_usage import FeatureIndex @@ -18,7 +21,7 @@ class InvocationsHostServer(InvocationAgentServerHost): def __init__( self, - agent: SupportsAgentRun, + agent: SupportsAgentRun | Callable[[], SupportsAgentRun | Awaitable[SupportsAgentRun]], *, openapi_spec: dict[str, Any] | None = None, **kwargs: Any, @@ -26,7 +29,8 @@ def __init__( """Initialize an InvocationsHostServer. Args: - agent: The agent to handle responses for. + 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. **kwargs: Additional keyword arguments. @@ -34,9 +38,11 @@ def __init__( 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. """ + 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.invoke_handler(self._handle_invoke) mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) @@ -72,6 +78,14 @@ def _partition_key(self) -> str | tuple[str, str]: return context.session_id + @asynccontextmanager + async def _request_agent(self) -> AsyncGenerator[SupportsAgentRun]: + agent = await resolve_agent(self._agent) + async with AsyncExitStack() as resources: + if self._owns_request_agent and isinstance(agent, AbstractAsyncContextManager): + await resources.enter_async_context(agent) + yield agent + async def _handle_invoke(self, request: Request) -> Response: """Invoke the agent with the given request.""" try: @@ -100,9 +114,19 @@ async def _handle_invoke(self, request: Request) -> Response: if stream: async def stream_response() -> AsyncGenerator[str]: - async for update in self._agent.run(user_message, session=session, stream=True): - if update.text: - yield update.text + async with 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(), @@ -110,5 +134,6 @@ async def stream_response() -> AsyncGenerator[str]: headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) - response = await self._agent.run([user_message], session=session) + async with 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/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 77d56c5e631..6f7b9ebcf29 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -9,7 +9,16 @@ import logging import os import re -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Generator, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Generator, + Mapping, + Sequence, +) from contextlib import AbstractAsyncContextManager, AsyncExitStack, aclosing, suppress from dataclasses import asdict, dataclass, is_dataclass from typing import Generic, Literal, TypeGuard, TypeVar, cast @@ -25,6 +34,7 @@ InMemoryHistoryProvider, Message, RawAgent, + ResponseStream, SessionStore, SupportsAgentRun, UsageDetails, @@ -73,6 +83,7 @@ from mcp import McpError from typing_extensions import Any +from ._agent_source import is_agent, resolve_agent, validate_agent_source from ._feature_usage import FeatureIndex from ._state_store import ( AgentSessionStoreProvider, @@ -151,9 +162,8 @@ class _SignalledIterator(Generic[_T]): async context. If an event and a new item becomes ready at the same time, the event takes priority and the item - is discarded. Cancelling the background task while it's mid-call is also what actually interrupts - a suspended model/tool call, since ``ResponseStream`` (what ``SupportsAgentRun.run(stream=True)`` - returns) has no ``aclose()`` of its own. + is discarded. Cancelling the background task while it's mid-call interrupts a suspended model or + tool call, then the driver closes the underlying stream in ``finally``. Callers MUST drive this through ``contextlib.aclosing`` (or an equivalent try/finally calling ``aclose()``): ``__anext__`` only cancels the driver task on its own signalled/exhausted paths, so @@ -190,16 +200,25 @@ def __aiter__(self) -> _SignalledIterator[_T]: async def _drive(self) -> None: """Pull items from the wrapped iterator into ``self._queue`` for the object's lifetime.""" - while True: - try: - item: Any = await self._iterator.__anext__() - except StopAsyncIteration: - await self._queue.put(_STOP_SENTINEL) - return - except Exception as exc: - await self._queue.put(exc) - return - await self._queue.put(item) + try: + while True: + try: + item: Any = await self._iterator.__anext__() + except StopAsyncIteration: + await self._queue.put(_STOP_SENTINEL) + return + except Exception as exc: + await self._queue.put(exc) + return + await self._queue.put(item) + finally: + iterator: AsyncIterator[_T] = self._iterator + if isinstance(iterator, ResponseStream): + await cast(ResponseStream[_T, Any], iterator).close() + else: + close = getattr(iterator, "aclose", None) + if close is not None: + await close() async def __anext__(self) -> _T: if self._driver is None: @@ -372,13 +391,102 @@ def consent_url_from_error(exc: BaseException) -> list[ConsentError] | None: # endregion Foundry Toolbox Auth integration +@dataclass(frozen=True) +class _AgentConfiguration: + workflow: bool + agent_server_history: bool + client_stores_by_default: bool + hosted_history: bool + + +def _validate_agent_configuration( + agent: SupportsAgentRun, + history_source: Literal["agent_server", "agent"], + options: ResponsesServerOptions | None, +) -> _AgentConfiguration: + is_workflow_agent = isinstance(agent, WorkflowAgent) + if is_workflow_agent and agent.workflow._runner_context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] + raise RuntimeError( + "There should not be a checkpoint storage already present in the workflow agent. " + "The hosting infrastructure will manage checkpoints instead." + ) + + resilient_background = bool(options and options.resilient_background) + if resilient_background and not is_workflow_agent: + raise RuntimeError( + "resilient_background=True is only supported for workflow agents. " + "Crash recovery cannot be provided for non-workflow agents." + ) + if options and options.steerable_conversations and is_workflow_agent: + raise RuntimeError( + "steerable_conversations=True is only supported for non-workflow agents. " + "Steering cannot be provided reliably for workflow agents." + ) + + uses_agent_server_history = history_source == "agent_server" + client_stores_by_default = False + if uses_agent_server_history and not is_workflow_agent: + if not isinstance(agent, RawAgent): + raise RuntimeError( + "history_source='agent_server' requires a RawAgent so hosting can enforce downstream " + "storage options. Construct ResponsesHostServer with history_source='agent' for a custom " + "SupportsAgentRun implementation." + ) + for provider in agent.context_providers: + if isinstance(provider, HistoryProvider) and provider.load_messages: + if _is_hosted_responses_history_sentinel(provider): + continue + raise RuntimeError( + "AgentServer response history is enabled, but the agent has a HistoryProvider " + "with load_messages=True. Remove that provider or construct ResponsesHostServer " + "with history_source='agent' to use the agent's regular history setup." + ) + service_continuation_options = [ + name + for name in ("conversation_id", "previous_response_id", "conversation") + if agent.default_options.get(name) is not None + ] + if service_continuation_options: + raise RuntimeError( + "AgentServer response history is enabled, but the agent has downstream service continuation " + f"option(s): {', '.join(service_continuation_options)}. Remove them or construct " + "ResponsesHostServer with history_source='agent' to resume the downstream service conversation." + ) + stores_by_default = getattr(cast(Any, agent).client, "STORES_BY_DEFAULT", None) + if not isinstance(stores_by_default, bool): + raise RuntimeError( + "history_source='agent_server' requires the agent's chat client to declare " + "STORES_BY_DEFAULT so hosting can enforce downstream storage behavior." + ) + client_stores_by_default = stores_by_default + + return _AgentConfiguration( + workflow=is_workflow_agent, + agent_server_history=uses_agent_server_history, + client_stores_by_default=client_stores_by_default, + hosted_history=uses_agent_server_history and not is_workflow_agent, + ) + + +def _initialize_agent_history(agent: SupportsAgentRun, configuration: _AgentConfiguration) -> None: + if not configuration.hosted_history or not isinstance(agent, RawAgent): + return + if not configuration.client_stores_by_default: + agent.default_options.pop("store", None) + if not any( + _is_hosted_responses_history_sentinel(provider) + for provider in cast(Sequence[ContextProvider], agent.context_providers) + ): + agent.context_providers.append(InMemoryHistoryProvider(source_id=_HOSTED_RESPONSES_HISTORY_SOURCE_ID)) + + # region ResponsesHostServer class ResponsesHostServer(ResponsesAgentServerHost): """A responses server host for an agent.""" def __init__( self, - agent: SupportsAgentRun, + agent: SupportsAgentRun | Callable[[], SupportsAgentRun | Awaitable[SupportsAgentRun]], *, prefix: str = "", options: ResponsesServerOptions | None = None, @@ -392,7 +500,8 @@ def __init__( """Initialize a ResponsesHostServer. Args: - agent: The agent to handle responses for. + 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`. prefix: The URL prefix for the server. options: Optional server options. store: Optional response store for input and history look up. @@ -420,6 +529,7 @@ def __init__( `history_source="agent"` mode, is persisted by the configured session store. 3. The server owns the supplied agent instance and may add hosting-specific providers. Do not reuse the same agent with another host or invoke it directly after construction. + An agent returned by a callable belongs to that request. 4. Resiliency (resilient_background=True) is ONLY supported for workflows; constructing this server with a non-workflow agent and `resilient_background=True` raises `RuntimeError`. When resiliency is enabled, and the server crashes mid-response: @@ -442,92 +552,29 @@ def __init__( """ if history_source not in ("agent_server", "agent"): raise ValueError("history_source must be either 'agent_server' or 'agent'.") + validate_agent_source(agent) - is_workflow_agent = isinstance(agent, WorkflowAgent) - if is_workflow_agent and agent.workflow._runner_context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] - raise RuntimeError( - "There should not be a checkpoint storage already present in the workflow agent. " - "The hosting infrastructure will manage checkpoints instead." - ) - - resilient_background = bool(options and options.resilient_background) - if resilient_background and not is_workflow_agent: - raise RuntimeError( - "resilient_background=True is only supported for workflow agents. " - "Crash recovery cannot be provided for non-workflow agents." - ) - if options and options.steerable_conversations and is_workflow_agent: - raise RuntimeError( - "steerable_conversations=True is only supported for non-workflow agents. " - "Steering cannot be provided reliably for workflow agents." - ) - - uses_agent_server_history = history_source == "agent_server" - client_stores_by_default = False - if uses_agent_server_history and not is_workflow_agent: - if not isinstance(agent, RawAgent): - raise RuntimeError( - "history_source='agent_server' requires a RawAgent so hosting can enforce downstream " - "storage options. Construct ResponsesHostServer with history_source='agent' for a custom " - "SupportsAgentRun implementation." - ) - for provider in agent.context_providers: - if isinstance(provider, HistoryProvider) and provider.load_messages: - if _is_hosted_responses_history_sentinel(provider): - continue - raise RuntimeError( - "AgentServer response history is enabled, but the agent has a HistoryProvider " - "with load_messages=True. Remove that provider or construct ResponsesHostServer " - "with history_source='agent' to use the agent's regular history setup." - ) - service_continuation_options = [ - name - for name in ("conversation_id", "previous_response_id", "conversation") - if agent.default_options.get(name) is not None - ] - if service_continuation_options: - raise RuntimeError( - "AgentServer response history is enabled, but the agent has downstream service continuation " - f"option(s): {', '.join(service_continuation_options)}. Remove them or construct " - "ResponsesHostServer with history_source='agent' to resume the downstream service conversation." - ) - stores_by_default = getattr(cast(Any, agent).client, "STORES_BY_DEFAULT", None) - if not isinstance(stores_by_default, bool): - raise RuntimeError( - "history_source='agent_server' requires the agent's chat client to declare " - "STORES_BY_DEFAULT so hosting can enforce downstream storage behavior." - ) - client_stores_by_default = stores_by_default + resolved_agent = agent if is_agent(agent) else None + configuration = ( + _validate_agent_configuration(resolved_agent, history_source, options) + if resolved_agent is not None + else None + ) # No caller-owned agent state is mutated until all validation and base-host construction succeed. super().__init__(prefix=prefix, options=options, store=store, **kwargs) - self._uses_agent_server_history = uses_agent_server_history - self._client_stores_by_default = client_stores_by_default - self._is_workflow_agent = is_workflow_agent - self._resilient_background = resilient_background - - self._uses_hosted_responses_history = False - if self._uses_agent_server_history and not self._is_workflow_agent and isinstance(agent, RawAgent): - self._uses_hosted_responses_history = True - if not self._client_stores_by_default: - agent.default_options.pop("store", None) - if not any( - _is_hosted_responses_history_sentinel(provider) - for provider in cast(Sequence[ContextProvider], agent.context_providers) - ): - # The Responses provider already supplies the complete transcript on every - # call. Agent.run would otherwise mutate the same user-owned agent by - # auto-injecting its default InMemoryHistoryProvider. Install a transient - # buffer that carries history within a function-call loop, then discard its - # state before persisting the session so the transcript is not replayed twice. - agent.context_providers.append( - InMemoryHistoryProvider( - source_id=_HOSTED_RESPONSES_HISTORY_SOURCE_ID, - ) - ) - - self._agent: SupportsAgentRun = agent + self._agent_source = agent + self._agent = resolved_agent + self._configuration = configuration + self._history_source: Literal["agent_server", "agent"] = history_source + self._host_options = options + self._uses_agent_server_history = ( + configuration.agent_server_history if configuration is not None else history_source == "agent_server" + ) + self._resilient_background = bool(options and options.resilient_background) + if resolved_agent is not None and configuration is not None: + _initialize_agent_history(resolved_agent, configuration) # Storage providers self._checkpoint_storage_provider = ( @@ -566,10 +613,13 @@ async def _ensure_agent_ready(self) -> None: async with self._agent_init_lock: if self._agent_stack is not None: return + agent = self._agent + if agent is None: + raise RuntimeError("A request-scoped agent cannot use the server-lifetime initialization path.") stack = AsyncExitStack() try: - if isinstance(self._agent, AbstractAsyncContextManager): - await stack.enter_async_context(cast(AbstractAsyncContextManager[Any], self._agent)) + if isinstance(agent, AbstractAsyncContextManager): + await stack.enter_async_context(cast(AbstractAsyncContextManager[Any], agent)) except BaseException: await stack.aclose() raise @@ -589,26 +639,65 @@ async def _handle_response( cancellation_signal: asyncio.Event, ) -> AsyncIterable[ResponseStreamEvent | ResponseCheckpointEvent]: """Handle the creation of a response.""" - # Common per-request setup shared by the workflow and non-workflow paths: - # create the response stream and the streaming output-item tracker, emit - # the opening lifecycle events, and convert any exception raised while - # producing the response into a terminal ``response.failed`` event (which - # also drains the tracker so the SSE stream stays well-formed). response_event_stream = _create_response_event_stream(context) - if context.is_steered_turn: logger.debug("Serving steered turn (pending_input_count=%d)", context.pending_input_count) - yield response_event_stream.emit_created() yield response_event_stream.emit_in_progress() + terminal_event: ResponseStreamEvent | None = None + agent = await resolve_agent(self._agent_source) + configuration = self._configuration or _validate_agent_configuration( + agent, self._history_source, self._host_options + ) + if self._configuration is None: + _initialize_agent_history(agent, configuration) + + async with AsyncExitStack() as resources: + inner = self._handle_prepared_response( + request, + context, + cancellation_signal, + response_event_stream, + agent, + configuration, + resources, + ) + try: + async for event in inner: + if isinstance(event, Mapping) and event.get("type") in ( + "response.completed", + "response.incomplete", + "response.failed", + ): + terminal_event = event + else: + yield event + finally: + await inner.aclose() + if terminal_event is not None: + yield terminal_event + + async def _handle_prepared_response( + self, + request: CreateResponse, + context: ResponseContext, + cancellation_signal: asyncio.Event, + response_event_stream: ResponseEventStream, + agent: SupportsAgentRun, + configuration: _AgentConfiguration, + resources: AsyncExitStack, + ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: # Lazy-enter the agent (and any MCP tools it owns). The MCP client wraps gateway # consent failures (and other connection-time errors) in AgentFrameworkException; if # one of those is a consent error we surface the consent link to the client through # the already-opened response stream instead of failing the request. Other exception # types fall through to the outer handler below and become ``response.failed``. try: - await self._ensure_agent_ready() + if self._configuration is not None: + await self._ensure_agent_ready() + elif isinstance(agent, AbstractAsyncContextManager): + await resources.enter_async_context(agent) except AgentFrameworkException as ex: consent_errors_to_emit = consent_url_from_error(ex) if consent_errors_to_emit is None or len(consent_errors_to_emit) == 0: @@ -634,7 +723,7 @@ async def _handle_response( yield event return - if not self._is_workflow_agent: + if not configuration.workflow: try: request_context = get_request_context() session_storage = self._session_storage_provider.get_store( @@ -649,7 +738,7 @@ async def _handle_response( "Cannot find an existing agent session for " f"previous_response_id={previous_response_id}." ) - session = self._agent.create_session() + session = agent.create_session() await session_storage.set(context.conversation_id or context.response_id, session) except Exception as save_error: logger.error( @@ -678,12 +767,25 @@ async def _handle_response( tracker = _OutputItemTracker(response_event_stream) try: - if self._is_workflow_agent: + if configuration.workflow: inner = self._handle_inner_workflow( - request, context, response_event_stream, tracker, cancellation_signal + request, + context, + response_event_stream, + tracker, + cancellation_signal, + cast(WorkflowAgent, agent), ) else: - inner = self._handle_inner_agent(request, context, response_event_stream, tracker, cancellation_signal) + inner = self._handle_inner_agent( + request, + context, + response_event_stream, + tracker, + cancellation_signal, + agent, + configuration, + ) try: async for event in inner: @@ -712,6 +814,7 @@ async def _load_request_messages( context: ResponseContext, *, approval_storage: FunctionApprovalStore | None, + configuration: _AgentConfiguration | None = None, ) -> list[Message]: """Load the request's input and prior history concurrently, assembled for the run. @@ -726,13 +829,16 @@ async def _load_request_messages( need to know the storage-result ordering. If either read fails, the sibling task is cancelled and drained so no storage read is orphaned. """ + uses_agent_server_history = ( + configuration.agent_server_history if configuration is not None else self._uses_agent_server_history + ) async def _load_input() -> list[Message]: input_items = await context.get_input_items() return await _items_to_messages(input_items, approval_storage=approval_storage) async def _load_history() -> list[Message]: - if not self._uses_agent_server_history: + if not uses_agent_server_history: return [] history = await context.get_history() return await _output_items_to_messages(history, approval_storage=approval_storage) @@ -758,6 +864,8 @@ async def _handle_inner_agent( response_event_stream: ResponseEventStream, tracker: _OutputItemTracker, cancellation_signal: asyncio.Event, + agent: SupportsAgentRun, + configuration: _AgentConfiguration, ) -> AsyncGenerator[ResponseStreamEvent]: """Handle a regular (non-workflow) agent. @@ -786,7 +894,11 @@ async def _handle_inner_agent( # session load below. These are independent storage round-trips with no data dependency # between them, so overlapping them removes serial latency from the request critical path. request_messages_task = asyncio.ensure_future( - self._load_request_messages(context, approval_storage=approval_storage) + self._load_request_messages( + context, + approval_storage=approval_storage, + configuration=configuration, + ) ) previous_response_id = request.get("previous_response_id") @@ -797,7 +909,7 @@ async def _handle_inner_agent( raise RuntimeError( f"Cannot find an existing agent session for previous_response_id={previous_response_id}." ) - session = self._agent.create_session() + session = agent.create_session() session_save_id = context.conversation_id or context.response_id except BaseException as ex: # Session preparation failed (or the request was cancelled / the stream closed — @@ -816,7 +928,7 @@ async def _handle_inner_agent( request_interrupted = False try: - if self._uses_agent_server_history: + if configuration.agent_server_history: session.state.pop(_HOSTED_RESPONSES_HISTORY_SOURCE_ID, None) # A restored service ID belongs to the downstream model service. Replaying the # AgentServer transcript while resuming that service history would duplicate every @@ -829,8 +941,8 @@ async def _handle_inner_agent( "session": session, } chat_options, are_options_set = _to_chat_options(request) - if self._uses_agent_server_history: - if self._client_stores_by_default: + if configuration.agent_server_history: + if configuration.client_stores_by_default: # The response provider already owns the transcript used for this run. Keep a # storing downstream service stateless so it cannot become a second history source. chat_options["store"] = False @@ -838,7 +950,7 @@ async def _handle_inner_agent( # Do not pass a storage option to clients that do not advertise support for it. chat_options.pop("store", None) - if isinstance(self._agent, RawAgent): + if isinstance(agent, RawAgent): run_kwargs["options"] = chat_options elif are_options_set: logger.warning("Agent doesn't support runtime options. They will be ignored.") @@ -846,7 +958,7 @@ async def _handle_inner_agent( # Non-workflow agents can't be resilient, so there is no exit_for_recovery path here: # both shutdown and steering/cancel just wind the turn down once observed. agent_stream = _SignalledIterator( - self._agent.run(stream=True, **run_kwargs), # type: ignore[reportUnknownMemberType] + agent.run(stream=True, **run_kwargs), # type: ignore[reportUnknownMemberType] context.shutdown, cancellation_signal, ) @@ -867,12 +979,12 @@ async def _handle_inner_agent( exc_info=(type(ex), ex, ex.__traceback__), ) finally: - if self._uses_hosted_responses_history: + if configuration.hosted_history: session.state.pop(_HOSTED_RESPONSES_HISTORY_SOURCE_ID, None) # A service ID here means the client stored the turn despite the forced `store=False`. # Do not persist a session that could resume that unreconciled history on a later turn. - stored_output_violation = self._uses_agent_server_history and session.service_session_id is not None + stored_output_violation = configuration.agent_server_history and session.service_session_id is not None if stored_output_violation: misconfigured = RuntimeError( "The agent's chat client stored this turn server-side while AgentServer response history " @@ -912,11 +1024,9 @@ async def _handle_inner_workflow( response_event_stream: ResponseEventStream, tracker: _OutputItemTracker, cancellation_signal: asyncio.Event, + agent: WorkflowAgent, ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: """Handle the creation of a response for a workflow agent.""" - if not isinstance(self._agent, WorkflowAgent): - raise RuntimeError("Agent is not a workflow agent.") - try: request_context = get_request_context() approval_storage = self._function_approval_storage_provider.get_store( @@ -953,10 +1063,10 @@ async def _handle_inner_workflow( if checkpoint_id is not None: logger.debug("Serving recovery request from workflow checkpoint %s", checkpoint_id) run_stream = self._resume_workflow_from_checkpoint( - checkpoint_id, checkpoint_storage, context.response_id + checkpoint_id, checkpoint_storage, context.response_id, agent ) else: - latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=self._agent.workflow.name) + latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=agent.workflow.name) if latest_checkpoint is not None: logger.debug( "Found a workflow checkpoint %s but no prior response snapshot was durably persisted; " @@ -964,7 +1074,10 @@ async def _handle_inner_workflow( latest_checkpoint.checkpoint_id, ) run_stream = self._resume_workflow_from_checkpoint( - latest_checkpoint.checkpoint_id, checkpoint_storage, context.response_id + latest_checkpoint.checkpoint_id, + checkpoint_storage, + context.response_id, + agent, ) else: # No checkpoint was ever paired with a persisted response snapshot (e.g. the crash @@ -974,7 +1087,7 @@ async def _handle_inner_workflow( logger.debug( "Serving recovery request with no prior workflow checkpoint; replaying original input" ) - run_stream = self._agent.run( + run_stream = agent.run( input_messages, stream=True, checkpoint_storage=checkpoint_storage, @@ -996,7 +1109,7 @@ async def _handle_inner_workflow( context_id=checkpoint_load_id, platform_context=request_context, ) - latest_checkpoint = await restore_checkpoint_storage.get_latest(workflow_name=self._agent.workflow.name) + latest_checkpoint = await restore_checkpoint_storage.get_latest(workflow_name=agent.workflow.name) if latest_checkpoint is None and previous_response_id is not None: # A previous_response_id must have a prior workflow checkpoint to resume from @@ -1014,13 +1127,13 @@ async def _handle_inner_workflow( # If the restored checkpoint had pending request_info events, the # restore-only call replays them through # ``WorkflowAgent._convert_workflow_event_to_agent_response_updates`` - # and populates ``self._agent.pending_requests``. That is the correct + # and populates ``agent.pending_requests``. That is the correct # state: those requests are genuinely outstanding, and the next # ``run(input_messages, ...)`` call may contain ``function_call_output`` # items (carried as FunctionResult/FunctionApprovalResponse content) # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`. restore_iter = _SignalledIterator( - self._agent.run( + agent.run( stream=True, checkpoint_id=latest_checkpoint.checkpoint_id, checkpoint_storage=restore_checkpoint_storage, @@ -1042,7 +1155,7 @@ async def _handle_inner_workflow( if cancellation_signal.is_set(): return - run_stream = self._agent.run( + run_stream = agent.run( input_messages, stream=True, checkpoint_storage=checkpoint_storage, @@ -1052,7 +1165,7 @@ async def _handle_inner_workflow( async with aclosing(main_iter): async for update in main_iter: if self._resilient_background: - latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=self._agent.workflow.name) + latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=agent.workflow.name) if ( latest_checkpoint is not None and latest_checkpoint.checkpoint_id @@ -1092,6 +1205,7 @@ async def _resume_workflow_from_checkpoint( checkpoint_id: str, checkpoint_storage: CheckpointStorage, response_id: str, + agent: WorkflowAgent, ) -> AsyncGenerator[AgentResponseUpdate]: """Resume a crashed background workflow run, forwarding every event it produces. @@ -1106,9 +1220,6 @@ async def _resume_workflow_from_checkpoint( TODO(@taochen): #7677 """ - if not isinstance(self._agent, WorkflowAgent): - raise RuntimeError("Agent is not a workflow agent.") - agent = self._agent async for event in agent.workflow.run( stream=True, checkpoint_id=checkpoint_id, diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index 95857ab7d5b..250799849a6 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -15,6 +15,7 @@ from collections.abc import AsyncIterator, Iterator from contextlib import contextmanager from itertools import product +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -96,6 +97,47 @@ def get_session( return AgentSession(service_session_id=service_session_id, session_id=session_id) +class _ContextAgent(_FakeAgent): + def __init__( + self, + events: list[str], + *, + response: AgentResponse | None = None, + stream_updates: list[AgentResponseUpdate] | None = None, + ) -> None: + super().__init__(response=response, stream_updates=stream_updates) + self._events = events + + async def __aenter__(self) -> _ContextAgent: + self._events.append("enter") + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self._events.append("exit") + + def run( + self, + messages: Any = None, + *, + stream: bool = False, + session: AgentSession | None = None, + **kwargs: Any, + ) -> Any: + self._events.append("run") + result = super().run(messages, stream=stream, session=session, **kwargs) + if not stream: + return result + + async def _gen() -> AsyncIterator[AgentResponseUpdate]: + try: + async for update in result: + yield update + finally: + self._events.append("stream_close") + + return _gen() + + def _make_agent( *, response_text: str | None = None, @@ -153,6 +195,22 @@ def test_accepts_supports_agent_run(self) -> None: 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: + with pytest.raises(TypeError, match="agent must be an agent instance or a zero-argument callable"): + InvocationsHostServer(agent) + + def test_rejects_agent_class_requiring_constructor_arguments(self) -> None: + with pytest.raises(TypeError, match="agent callable must accept no arguments"): + InvocationsHostServer(cast(Any, _ContextAgent)) + + def test_rejects_factory_requiring_arguments(self) -> None: + def create_agent(name: str) -> _FakeAgent: + return _make_agent(response_text=name) + + with pytest.raises(TypeError, match="agent callable must accept no arguments"): + InvocationsHostServer(cast(Any, create_agent)) + # endregion @@ -227,6 +285,71 @@ async def test_hosted_keys_and_session_ids_preserve_identifier_values(self) -> N class TestHandleInvoke: + 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")])]) + server = InvocationsHostServer(_ContextAgent(events, response=response)) + + with _request_context(session_id="sess-1"): + await server._handle_invoke(_make_request({"message": "one"})) # pyright: ignore[reportPrivateUsage] + + assert events == ["run"] + + async def test_factory_agent_context_lifetime_non_streaming(self) -> None: + events: list[str] = [] + response = AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ok")])]) + server = InvocationsHostServer(lambda: _ContextAgent(events, response=response)) + + with _request_context(session_id="sess-1"): + result = await server._handle_invoke(_make_request({"message": "one"})) # pyright: ignore[reportPrivateUsage] + + assert bytes(result.body).decode() == "ok" + assert events == ["enter", "run", "exit"] + + async def test_factory_agent_context_lifetime_until_stream_closes(self) -> None: + events: list[str] = [] + updates = [ + AgentResponseUpdate(contents=[Content.from_text("one")]), + AgentResponseUpdate(contents=[Content.from_text("two")]), + ] + server = InvocationsHostServer(lambda: _ContextAgent(events, stream_updates=updates)) + + with _request_context(session_id="sess-1"): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "one", "stream": True}) + ) + + assert isinstance(response, StreamingResponse) + iterator = cast(Any, response.body_iterator) + assert await anext(iterator) == "one" + await iterator.aclose() + + assert events == ["enter", "run", "stream_close", "exit"] + + async def test_agent_callable_is_resolved_for_each_request(self) -> None: + agents: list[_FakeAgent] = [] + + def create_agent() -> _FakeAgent: + agent = _make_agent(response_text=f"agent-{len(agents) + 1}") + agents.append(agent) + return agent + + server = InvocationsHostServer(create_agent) + + with _request_context(session_id="sess-1"): + first = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "one"}) + ) + second = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "two"}) + ) + + assert bytes(first.body).decode() == "agent-1" + 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"] + @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: diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index dc712b1d223..fb79cb7a619 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -37,6 +37,7 @@ Content, FunctionInvocationLayer, HistoryProvider, + InMemoryCheckpointStorage, InMemoryHistoryProvider, Message, RawAgent, @@ -140,6 +141,42 @@ async def _raising_updates( raise RuntimeError(message) +class _AgentProtocolMock(MagicMock): + id = "test-agent" + name: str | None = "Test Agent" + description: str | None = "A mock agent for testing" + run: Any = None + create_session: Any = None + get_session: Any = None + + def __init__(self) -> None: + super().__init__() + self.run = MagicMock() + self.create_session = MagicMock(side_effect=lambda *, session_id=None: AgentSession(session_id=session_id)) + self.get_session = MagicMock( + side_effect=lambda service_session_id, *, session_id=None: AgentSession( + service_session_id=service_session_id, + session_id=session_id, + ) + ) + + +class _RawAgentMock(_AgentProtocolMock, RawAgent): + pass + + +class _WorkflowAgentMock(_AgentProtocolMock, WorkflowAgent): + _workflow_value: Any = None + + @property + def workflow(self) -> Any: + return self._workflow_value + + @workflow.setter + def workflow(self, value: Any) -> None: + self._workflow_value = value + + def _make_agent( *, response: AgentResponse | None = None, @@ -152,7 +189,7 @@ def _make_agent( tests that only care about complete output messages: the helper converts those messages into streamed updates. ``stream_updates`` is for tests that need explicit chunk boundaries to verify streaming event behavior. """ - agent = MagicMock(spec=RawAgent) if raw_agent else MagicMock() + agent = _RawAgentMock() if raw_agent else _AgentProtocolMock() agent.id = "test-agent" agent.name = "Test Agent" agent.description = "A mock agent for testing" @@ -164,7 +201,8 @@ def _make_agent( def create_session(*, session_id: str | None = None) -> AgentSession: return AgentSession(session_id=session_id) - agent.create_session.side_effect = create_session + agent.create_session = MagicMock(side_effect=create_session) + agent.run = MagicMock() if response is not None: @@ -208,6 +246,14 @@ def __init__(self) -> None: def create_session(self, *, session_id: str | None = None) -> AgentSession: return AgentSession(session_id=session_id) + def get_session( + self, + service_session_id: str | ServiceSessionId, + *, + session_id: str | None = None, + ) -> AgentSession: + return AgentSession(service_session_id=service_session_id, session_id=session_id) + def run( self, messages: Any = None, @@ -852,6 +898,18 @@ def test_stringify_mcp_output_extracts_only_text_content_mappings(self) -> None: class TestResponsesHostServerInit: + @pytest.mark.parametrize("agent", [None, 42]) + def test_init_rejects_invalid_agent_source(self, agent: Any) -> None: + with pytest.raises(TypeError, match="agent must be an agent instance or a zero-argument callable"): + ResponsesHostServer(agent) + + async def test_zero_argument_agent_class_is_resolved_as_factory(self) -> None: + server = _make_server(cast(Any, _StrictCustomAgent), history_source="agent") + + response = await _post(server) + + assert response.json()["status"] == "completed" + def test_init_basic(self) -> None: agent = _make_agent( response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) @@ -1487,6 +1545,7 @@ async def test_cancellation_signal_preempts_stuck_agent_call(self) -> None: response, not merely be checked between already-produced updates.""" store = SessionStore() gate = asyncio.Event() # Never set: simulates a model/tool call that never returns. + cleanup_called = asyncio.Event() agent = _make_agent() async def _stream_gen() -> AsyncIterator[AgentResponseUpdate]: @@ -1495,7 +1554,11 @@ async def _stream_gen() -> AsyncIterator[AgentResponseUpdate]: def run_streaming(*_args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: del _args, kwargs - return ResponseStream(_stream_gen(), finalizer=AgentResponse.from_updates) + return ResponseStream( + _stream_gen(), + finalizer=AgentResponse.from_updates, + cleanup_hooks=[cleanup_called.set], + ) agent.run = MagicMock(side_effect=run_streaming) server = _make_server(agent, session_store=store) @@ -1525,6 +1588,7 @@ async def _drain() -> list[Any]: types = [event.get("type") for event in events if isinstance(event, Mapping)] assert "response.output_text.delta" not in types assert types[-1] == "response.completed" + assert cleanup_called.is_set() async def test_consumer_failure_cancels_agent_stream_driver_task(self) -> None: """A crash in the consumer (`_OutputItemTracker.handle`) must not leave the background @@ -3516,7 +3580,7 @@ def _make_multi_response_agent( stream_updates_list: list[list[AgentResponseUpdate]] | None = None, ) -> MagicMock: """Create a mock agent that returns different responses on successive calls.""" - agent = MagicMock(spec=RawAgent) + agent = _RawAgentMock() agent.id = "test-agent" agent.name = "Test Agent" agent.description = "A mock agent for testing" @@ -3528,7 +3592,7 @@ def _make_multi_response_agent( def create_session(*, session_id: str | None = None) -> AgentSession: return AgentSession(session_id=session_id) - agent.create_session.side_effect = create_session + agent.create_session = MagicMock(side_effect=create_session) call_index = [0] @@ -4732,7 +4796,8 @@ async def test_workflow_rejects_invalid_checkpoint_scope( context_field: str, bad_id: str, ) -> None: - agent = MagicMock(spec=WorkflowAgent) + agent = _WorkflowAgentMock() + agent.run = MagicMock() agent.context_providers = [] agent.workflow = MagicMock() agent.workflow.name = "workflow" @@ -4848,6 +4913,25 @@ def test_returns_none_when_message_has_no_json(self) -> None: class TestAgentLifecycle: + async def test_factory_agent_is_entered_and_exited_for_each_request(self) -> None: + agents: list[MagicMock] = [] + + def create_agent() -> MagicMock: + agent = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) + ) + agents.append(agent) + return agent + + server = _make_server(create_agent) + + await _post(server, input_text="first", stream=False) + await _post(server, input_text="second", stream=False) + + assert len(agents) == 2 + assert [agent.__aenter__.await_count for agent in agents] == [1, 1] + assert [agent.__aexit__.await_count for agent in agents] == [1, 1] + async def test_agent_entered_lazily_on_first_request(self) -> None: agent = _make_agent( response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) @@ -5215,7 +5299,7 @@ async def _raise_stream() -> AsyncIterator[AgentResponseUpdate]: yield AgentResponseUpdate(contents=[Content.from_text("partial ")], role="assistant") raise RuntimeError("stream kaboom") - agent = MagicMock(spec=RawAgent) + agent = _RawAgentMock() agent.id = "test-agent" agent.name = "Test Agent" agent.description = "A mock agent for testing" @@ -5227,7 +5311,7 @@ async def _raise_stream() -> AsyncIterator[AgentResponseUpdate]: def create_session(*, session_id: str | None = None) -> AgentSession: return AgentSession(session_id=session_id) - agent.create_session.side_effect = create_session + agent.create_session = MagicMock(side_effect=create_session) def run_streaming(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: del args @@ -5266,7 +5350,7 @@ async def _raise_stream() -> AsyncIterator[AgentResponseUpdate]: yield AgentResponseUpdate(contents=[Content.from_text("hello ")], role="assistant") raise RuntimeError("mid-item kaboom") - agent = MagicMock(spec=RawAgent) + agent = _RawAgentMock() agent.id = "test-agent" agent.name = "Test Agent" agent.description = "A mock agent for testing" @@ -5278,7 +5362,7 @@ async def _raise_stream() -> AsyncIterator[AgentResponseUpdate]: def create_session(*, session_id: str | None = None) -> AgentSession: return AgentSession(session_id=session_id) - agent.create_session.side_effect = create_session + agent.create_session = MagicMock(side_effect=create_session) def run_streaming(*args: Any, **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse]: del args, kwargs @@ -5668,7 +5752,7 @@ def _build_multi_update_workflow_agent( async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None: await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True)) - workflow = WorkflowBuilder(start_executor=start).add_edge(start, inner).build() + workflow = WorkflowBuilder(name="multi-update-workflow", start_executor=start).add_edge(start, inner).build() return WorkflowAgent(workflow=workflow, name="Multi Update Workflow Agent"), inner @@ -5706,6 +5790,47 @@ class TestWorkflowAgentHosting: relative to the regular agent path. """ + async def test_async_factory_creates_workflow_agent_for_each_request(self) -> None: + created: list[tuple[WorkflowAgent, _MultiUpdateWorkflowAgentMock]] = [] + + async def create_agent() -> WorkflowAgent: + agent, inner = _build_multi_update_workflow_agent(["hello"]) + created.append((agent, inner)) + return agent + + server = _make_server(create_agent) + + first = await _post(server, input_text="one") + second = await _post(server, input_text="two") + + assert first.status_code == 200 + assert second.status_code == 200 + assert len(created) == 2 + assert created[0][0].workflow is not created[1][0].workflow + assert [inner.run_count for _, inner in created] == [1, 1] + + async def test_factory_workflow_restores_checkpoint_for_same_conversation(self) -> None: + runs: list[MagicMock] = [] + + def create_agent() -> WorkflowAgent: + agent, _ = _build_multi_update_workflow_agent(["hello"]) + run = MagicMock(wraps=agent.run) + cast(Any, agent).run = run + runs.append(run) + return agent + + checkpoint_storage = InMemoryCheckpointStorage() + checkpoint_provider = MagicMock(spec=CheckpointStoreProvider) + checkpoint_provider.get_store.return_value = checkpoint_storage + server = _make_server(create_agent, checkpoint_store_provider=checkpoint_provider) + + first = await _post(server, input_text="one", conversation_id="conversation-1") + second = await _post(server, input_text="two", conversation_id="conversation-1") + + assert first.status_code == 200 + assert second.status_code == 200 + assert [run.call_count for run in runs] == [1, 2] + async def test_basic_text_response(self) -> None: workflow_agent = _build_text_workflow_agent("hello from workflow") server = _make_server(workflow_agent) diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/README.md b/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/README.md index be3ff9a393b..f7a6a4ec380 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/README.md @@ -21,7 +21,9 @@ Each user message re-runs the workflow from the trigger. Because `Workflow.as_ag ### Agent Hosting -[`main.py`](main.py) builds three `Agent` instances on top of a shared `FoundryChatClient` (one per workflow role), registers them with the `WorkflowFactory` so the YAML's `InvokeAzureAgent` actions can resolve them by name, loads the workflow, wraps it with `.as_agent(...)`, and hands the agent to `ResponsesHostServer`, which provisions a REST API endpoint compatible with the OpenAI Responses protocol. +[`main.py`](main.py) gives `ResponsesHostServer` a callable that builds three `Agent` instances on top of a shared +`FoundryChatClient`, registers them with `WorkflowFactory`, loads the workflow, and wraps it with `.as_agent(...)`. +Each request receives a fresh workflow and agents, while the host restores the supported session and checkpoint state. The triage agent is configured with `response_format=TriageResponse` (a Pydantic model) so the workflow can read its structured fields via `Local.Triage.*`. The specialist agents are plain text and use `autoSend: true` to deliver their reply straight to the caller. diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/main.py b/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/main.py index 95b2d49a537..90b101ccf5f 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/declarative_customer_support/main.py @@ -4,7 +4,7 @@ from pathlib import Path from typing import Any, Literal -from agent_framework import Agent +from agent_framework import Agent, WorkflowAgent from agent_framework.foundry import FoundryChatClient, ResponsesHostServer from agent_framework_declarative import WorkflowFactory from agent_framework_openai import OpenAIChatOptions @@ -86,15 +86,8 @@ class TriageResponse(BaseModel): # --- Host setup ------------------------------------------------------------------ -def main() -> None: - workflow_path = Path(__file__).parent / "workflow.yaml" - - client = FoundryChatClient( - project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], - model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], - credential=DefaultAzureCredential(), - ) - +def create_workflow_agent(client: FoundryChatClient, workflow_path: Path) -> WorkflowAgent: + """Create a fresh declarative workflow agent for one hosted request.""" # The workflow's InvokeAzureAgent actions reference these agents by name. triage_agent = Agent( client=client, @@ -128,7 +121,7 @@ def main() -> None: # Wrap the declarative workflow as an AIAgent so it can be served behind # the Responses protocol. Each user turn re-runs the workflow with the # full conversation history available via Conversation.messages. - workflow_agent = workflow.as_agent( + return workflow.as_agent( name="declarative-customer-support", description=( "A multi-turn customer-support triage workflow that routes " @@ -137,7 +130,16 @@ def main() -> None: ), ) - ResponsesHostServer(workflow_agent).run() + +def main() -> None: + workflow_path = Path(__file__).parent / "workflow.yaml" + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=DefaultAzureCredential(), + ) + + ResponsesHostServer(agent=lambda: create_workflow_agent(client, workflow_path)).run() if __name__ == "__main__": diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/README.md b/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/README.md index cdf9df3eace..7c695fea3a1 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/README.md @@ -22,6 +22,8 @@ The workflow has three executors (see [main.py](main.py)): The workflow is hosted as an agent using the [Agent Framework](https://github.com/microsoft/agent-framework) `ResponsesHostServer`, which provisions a REST API endpoint compatible with the OpenAI Responses protocol. +The server receives a callable that creates a fresh workflow and executors for each request. Existing response +sessions and workflow checkpoints are restored by the host, while the application-owned `FoundryChatClient` is reused. Setting `resilient_background=True` in `ResponsesServerOptions` enables the framework to checkpoint the workflow's progress and durably persist streamed output, so a background response can be recovered and resumed after a crash (see "Testing resiliency" below). diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/main.py b/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/main.py index e941f23cd32..9a1b54aa26d 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/main.py @@ -15,7 +15,7 @@ import asyncio import os -from agent_framework import Agent, Executor, Message, WorkflowBuilder, WorkflowContext, executor, handler +from agent_framework import Agent, Executor, Message, Workflow, WorkflowBuilder, WorkflowContext, executor, handler from agent_framework.foundry import FoundryChatClient from agent_framework_foundry_hosting import ResponsesHostServer from azure.ai.agentserver.responses import ResponsesServerOptions @@ -76,13 +76,8 @@ async def complete(message: str, ctx: WorkflowContext[Never, str]) -> None: await ctx.yield_output(message) -def build_workflow(): +def build_workflow(client: FoundryChatClient) -> Workflow: """Build the target extraction, countdown, and completion workflow.""" - client = FoundryChatClient( - project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], - model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], - credential=DefaultAzureCredential(), - ) target_agent = Agent( client=client, name="counter_target_extractor", @@ -95,7 +90,7 @@ def build_workflow(): countdown = CountdownExecutor() return ( - WorkflowBuilder(start_executor=start, output_from="all") + WorkflowBuilder(name="countdown-workflow", start_executor=start, output_from="all") .add_edge(start, countdown) .add_edge(countdown, countdown) .add_edge(countdown, complete) @@ -106,9 +101,13 @@ def build_workflow(): def main() -> None: """Run the workflow as a durable Responses API host.""" print(f"PID: {os.getpid()}") # lets crash-recovery testing find and kill this process - workflow_agent = build_workflow().as_agent(name="countdown-workflow") + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=DefaultAzureCredential(), + ) server = ResponsesHostServer( - workflow_agent, + agent=lambda: build_workflow(client).as_agent(name="countdown-workflow"), options=ResponsesServerOptions(resilient_background=True), log_level="DEBUG", ) diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/README.md b/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/README.md index 608d1a564e8..173c5a2c811 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/README.md @@ -16,7 +16,9 @@ See [main.py](main.py) for the full implementation. ### Agent Hosting -The workflow is exposed as a single agent via `.as_agent()` and hosted using the [Agent Framework](https://github.com/microsoft/agent-framework) with the `ResponsesHostServer`, which provisions a REST API endpoint compatible with the OpenAI Responses protocol. +The workflow is exposed as an agent via `.as_agent()` and hosted using the +[Agent Framework](https://github.com/microsoft/agent-framework) with `ResponsesHostServer`. The host receives a +callable that builds a fresh `FoundryChatClient`, workflow, executors, and agents for each request. ## Running the Agent Host diff --git a/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/main.py b/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/main.py index e11a4655220..507b07722a0 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/main.py +++ b/python/samples/04-hosting/foundry-hosted-agents/responses/workflows/main.py @@ -2,7 +2,7 @@ import os -from agent_framework import Agent, AgentExecutor, WorkflowBuilder +from agent_framework import Agent, AgentExecutor, WorkflowAgent, WorkflowBuilder from agent_framework.foundry import FoundryChatClient, ResponsesHostServer from azure.identity import DefaultAzureCredential from dotenv import load_dotenv @@ -11,7 +11,8 @@ load_dotenv() -def main(): +def create_workflow_agent() -> WorkflowAgent: + """Create a fresh workflow agent for one hosted request.""" client = FoundryChatClient( project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], @@ -48,8 +49,9 @@ def main(): legal_executor = AgentExecutor(legal_agent, context_mode="last_agent") format_executor = AgentExecutor(format_agent, context_mode="last_agent") - workflow_agent = ( + return ( WorkflowBuilder( + name="slogan-workflow", start_executor=writer_executor, # Select only the formatted result as Workflow Output. # Unselected executor payloads are hidden unless selected as Intermediate Output. @@ -61,7 +63,9 @@ def main(): .as_agent() ) - server = ResponsesHostServer(workflow_agent) + +def main() -> None: + server = ResponsesHostServer(agent=create_workflow_agent) server.run()