From 7439f2c02a5f80e8088cfeff44ada709bc97922e Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:23:04 +0100 Subject: [PATCH 01/11] Python: Add request-scoped Foundry agent factories --- ...-python-foundry-request-agent-factories.md | 68 + .../workflow/test_functional_workflow.py | 112 ++ python/packages/foundry_hosting/README.md | 89 ++ .../_agent_factory.py | 110 ++ .../_invocations.py | 231 ++- .../_responses.py | 734 +++++++-- .../_state_store.py | 72 +- .../tests/test_invocations_factory.py | 754 +++++++++ .../foundry_hosting/tests/test_responses.py | 152 +- .../tests/test_responses_factory.py | 1401 +++++++++++++++++ .../foundry_hosting/tests/test_state_store.py | 26 + .../declarative_customer_support/README.md | 9 +- .../declarative_customer_support/main.py | 29 +- .../resilient_long_running_workflow/README.md | 7 +- .../resilient_long_running_workflow/main.py | 47 +- .../responses/workflows/README.md | 8 +- .../responses/workflows/main.py | 31 +- 17 files changed, 3618 insertions(+), 262 deletions(-) create mode 100644 docs/decisions/0040-python-foundry-request-agent-factories.md create mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py create mode 100644 python/packages/foundry_hosting/tests/test_invocations_factory.py create mode 100644 python/packages/foundry_hosting/tests/test_responses_factory.py diff --git a/docs/decisions/0040-python-foundry-request-agent-factories.md b/docs/decisions/0040-python-foundry-request-agent-factories.md new file mode 100644 index 00000000000..81f13d2ee57 --- /dev/null +++ b/docs/decisions/0040-python-foundry-request-agent-factories.md @@ -0,0 +1,68 @@ +--- +status: proposed +contact: RogerBarreto +date: 2026-09-10 +deciders: RogerBarreto +--- + +# Request-scoped agent factories for Python Foundry hosting + +## Context and Problem Statement + +Python Foundry hosts accept an agent instance and separate persisted state by user and conversation. +Some agent implementations also hold mutable execution state outside their `AgentSession`. +A workflow's executors, pending requests, and internal conversation therefore need their own lifetime, +independent of the server object's lifetime. + +## Decision Drivers + +- Preserve ordinary-agent instance callers and existing protocol formats. +- Create independent workflow execution objects for each request. +- Continue authorized conversations from stored state, including after process recreation. +- Avoid a core workflow redesign or a cache of live runtimes with expiration policies. + +## Considered Options + +| Option | Benefit | Cost | +| --- | --- | --- | +| Factory per request | Explicit ownership; same construction path for fresh calls and recovery | Applications rebuild mutable runtime objects and manage client ownership | +| Runtime cache per user and session | Avoids rebuilding objects on every call | Requires expiration, eviction, cleanup, concurrency controls, and a separate restart path | +| Move workflow runtime into a new core session type | Separates shared definitions from session state throughout the framework | Larger change to core execution and serialization contracts | +| Continue accepting shared workflow instances | No caller migration | Persisted-state isolation does not separate live workflow state | + +## Decision Outcome + +Add `agent_factory` to `ResponsesHostServer` and `InvocationsHostServer`. It is a zero-argument callable +returning an agent or an awaitable agent. Resolve it within the current platform request context and +keep the result local until execution, persistence, and streaming have finished. + +Retain `agent` for ordinary instances. Require factories for the two built-in workflow-agent types, +recognized by one private Foundry helper. Do not add a public workflow-recognition interface or attempt +to inspect arbitrary application wrappers. + +Responses continues using its conversation/response checkpoint scopes. Invocations adds workflow +checkpoint persistence scoped to user and invocation session without changing its text protocol. +Functional workflows need their own checkpoint adaptation because saved-input replay differs from +graph workflow restoration. + +Functional resilient Responses recovery is explicitly unsupported: checkpoints omit buffered output +from completed steps, so a hosting adapter cannot restore that output without rerunning application +work. Invocations retains its text-only exchange and rejects pending or interrupted functional +continuations; a new message after clean completion is supported. + +The host manages an agent's exposed async context manager. It does not recursively discover resources +inside executors or closures. Applications must construct fresh mutable runtime objects and give +shared clients a lifetime that outlasts every request using them. + +Serialize updates to the same workflow scope within a host process. Do not describe these locks as +distributed coordination or checkpoint replay as exactly-once execution. + +## Consequences + +Workflow callers migrate from `Host(workflow_agent)` to `Host(agent_factory=create_agent)`. +The factory must construct new mutable objects, not return the same instance or reuse stateful executors. +Stable workflow names, executor IDs, and serialization registrations are necessary for continuation. + +Ordinary-agent instance callers keep their existing lifecycle. Factories do not automatically save +arbitrary custom agent fields, and the Invocations text protocol still does not provide a complete +external approval exchange. A core session redesign remains a separate possible improvement. diff --git a/python/packages/core/tests/workflow/test_functional_workflow.py b/python/packages/core/tests/workflow/test_functional_workflow.py index 6f60aba0395..9e467988e66 100644 --- a/python/packages/core/tests/workflow/test_functional_workflow.py +++ b/python/packages/core/tests/workflow/test_functional_workflow.py @@ -717,6 +717,118 @@ async def wf(x: int, ctx: RunContext) -> str: class TestCheckpointing: + async def test_fresh_agent_checkpoint_replay_cannot_recover_buffered_step_output(self) -> None: + """A completed step's saved result does not preserve its buffered output events.""" + storage = InMemoryCheckpointStorage() + step_saved = asyncio.Event() + finish = asyncio.Event() + calls = 0 + + @step + async def emit_step(text: str) -> str: + nonlocal calls + calls += 1 + ctx = get_run_context() + assert ctx is not None + await ctx.add_event(WorkflowEvent("output", executor_id="emit_step", data=f"step:{text}")) + return text + + @workflow + async def buffered_workflow(text: str) -> str: + result = await emit_step(text) + step_saved.set() + await finish.wait() + return f"final:{result}" + + original = buffered_workflow.build().as_agent() + received: list[str] = [] + + async def consume_original() -> None: + async for update in original.run("original", stream=True, checkpoint_storage=storage): + if update.text: + received.append(update.text) + + consumer = asyncio.create_task(consume_original()) + try: + await asyncio.wait_for(step_saved.wait(), timeout=5) + checkpoint = await storage.get_latest(workflow_name=buffered_workflow.name) + assert checkpoint is not None + assert received == [] + finally: + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + + finish.set() + recovered = buffered_workflow.build().as_agent() + result = await recovered.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) + + assert calls == 1 + assert result.text == "final:original" + assert "step:original" not in result.text + + complete = await buffered_workflow.build().as_agent().run("original", checkpoint_storage=storage) + assert "step:original" in complete.text + assert "final:original" in complete.text + assert calls == 2 + + async def test_fresh_agent_completed_checkpoint_replays_old_input_instead_of_starting_new_turn(self) -> None: + calls: list[str] = [] + storage = InMemoryCheckpointStorage() + + @step + async def record(text: str) -> str: + calls.append(text) + return text + + @workflow + async def turns(text: str) -> str: + return await record(text) + + original = turns.build().as_agent() + assert (await original.run("first", checkpoint_storage=storage)).text == "first" + checkpoint = await storage.get_latest(workflow_name=turns.name) + assert checkpoint is not None + + restored = turns.build().as_agent() + replay = await restored.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) + assert replay.text == "first" + assert calls == ["first"] + with pytest.raises(ValueError, match="message.*checkpoint_id"): + await restored.run("second", checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) + + assert (await original.run("second", checkpoint_storage=storage)).text == "second" + assert calls == ["first", "second"] + + async def test_fresh_agent_resumes_pending_request_from_checkpoint_without_new_message(self) -> None: + storage = InMemoryCheckpointStorage() + + @workflow + async def review(text: str, ctx: RunContext) -> str: + answer = await ctx.request_info(text, response_type=str, request_id="review-request") + return f"{text}:{answer}" + + pending_agent = review.build().as_agent() + await pending_agent.run("original", checkpoint_storage=storage) + checkpoint = await storage.get_latest(workflow_name=review.name) + assert checkpoint is not None + assert "review-request" in checkpoint.pending_request_info_events + + restored = review.build().as_agent() + response = await restored.run( + checkpoint_id=checkpoint.checkpoint_id, + checkpoint_storage=storage, + responses={"review-request": "accepted"}, + ) + assert response.text == "original:accepted" + completed = next( + saved + for saved in await storage.list_checkpoints(workflow_name=review.name) + if saved.previous_checkpoint_id == checkpoint.checkpoint_id + ) + assert not restored.pending_requests + assert not completed.pending_request_info_events + async def test_checkpoint_save_and_restore(self): storage = InMemoryCheckpointStorage() diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 33264ec47c3..8956a2f481a 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -2,6 +2,64 @@ 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 + +Both hosts accept either an ordinary agent instance or an `agent_factory`, but not both. +Existing `ResponsesHostServer(agent)` and `InvocationsHostServer(agent)` calls remain supported for ordinary agents. +Pass `WorkflowAgent` and `FunctionalWorkflowAgent` through a factory instead of passing a live instance: + +```python +from agent_framework_foundry_hosting import InvocationsHostServer, ResponsesHostServer + + +def create_agent(): + # Build a new workflow, with new mutable agents and executors, here. + return build_workflow().as_agent() + + +server = ResponsesHostServer(agent_factory=create_agent) +# Alternatively, use the existing Invocations text protocol: +# server = InvocationsHostServer(agent_factory=create_agent) +``` + +A factory takes no arguments and can return an agent directly or await its construction: + +```python +async def create_agent(): + configuration = await load_configuration() + return build_workflow(configuration).as_agent() + + +server = ResponsesHostServer(agent_factory=create_agent) +``` + +The host calls the factory inside the current request's platform context, not during startup. The resulting agent +belongs to that request, including its entire stream. Responses recovery also creates a new agent through the factory. +Returning the same workflow instance repeatedly is not supported. Neither is building a new workflow around previously +used mutable executors. Keep workflow names, executor IDs, and serialized state type registrations stable so a new +instance can restore the previous instance's checkpoints. + +Factories are also useful for ordinary agents with request-specific configuration. They do not automatically persist +custom fields on an agent: state needed on the next request must use the supported session or checkpoint stores. +The host recognizes the two built-in workflow adapters and their subclasses; it cannot discover a workflow hidden +inside an arbitrary custom agent. + +### Resource ownership + +The host enters and exits a factory-created agent's async context manager when it has one. Resources remain available +until execution and streaming finish, and are released on completion or interruption. Factory code must clean up its +own partially constructed resources if it fails before returning an agent. + +Built-in workflow adapters do not automatically close every client or tool captured by their executors. Keep +application-owned, concurrency-safe clients open for the host lifetime, or return a workflow subclass with an async +context manager that owns its request-specific resources. Do not close a shared client at the end of one request. +An ordinary `Agent` context manager also manages its chat client and MCP tools, so their ownership must match its lifetime. + +See the [workflow sample](../../samples/04-hosting/foundry-hosted-agents/responses/workflows/) for fresh agents and +executors using an application-owned model client, and the +[recovery sample](../../samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/) +for request-local workflows restored after a process restart. + ## Conversation history `ResponsesHostServer` uses AgentServer response history as the model's conversation history by default: @@ -92,6 +150,37 @@ Native Responses refusal parts are stored as text carrying `FoundryCheckpointStore`, backed by Foundry storage when hosted and file-based storage locally. Stored checkpoints are scoped under `checkpoints`. +Each factory-created workflow starts with independent runtime state. Responses restores only the checkpoint belonging +to the current user and conversation or response chain. A `previous_response_id` without a saved workflow checkpoint +fails; a conversation with prior history but no workflow checkpoint also fails rather than silently starting over. +A new conversation with no history can start without a checkpoint. + +For resilient background Responses using a graph workflow, recovery uses the checkpoint associated with the persisted response output. +If no response checkpoint was recorded, it can use the latest workflow checkpoint. If execution stopped before any +workflow checkpoint was saved, recovery replays the original input using a fresh factory-created workflow. + +`FunctionalWorkflowAgent` does not support `resilient_background=True`. Its checkpoints retain completed step results +but not all output buffered inside those steps. Recovering from such a checkpoint could omit output; the host rejects +that configuration instead of silently losing output or rerunning application work. Functional workflows can use +the factory for normal Responses requests and supported pending-response continuation. + +### Invocations workflows + +`InvocationsHostServer(agent_factory=...)` persists workflow checkpoints for the platform user and invocation session. +A subsequent request with the same authorized session can restore its graph workflow in a new agent instance, including +after recreating the host. Ordinary-agent instance callers retain their existing in-memory session behavior. + +The wire format remains unchanged: requests use `message` and `stream`, and responses contain text or streamed text. +This helper does not add a structured exchange for external workflow approvals or pending requests. A plain text +message is not an approval response. Use Responses when callers need that structured exchange. + +Functional workflows accept a new message after a completed invocation, but pending or interrupted functional +continuation is rejected. They do not acquire graph-workflow recovery semantics by being passed through a factory. + +Requests updating the same workflow scope are serialized within one host; independent scopes can execute concurrently. +This is not a distributed lock across multiple host processes. Invocations does not automatically recover an +interrupted HTTP response, and checkpoints do not guarantee that external side effects execute exactly once. + ### Function approvals `ResponsesHostServer` persists function approvals durably. By default, it uses the diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py new file mode 100644 index 00000000000..587a2cc6c51 --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py @@ -0,0 +1,110 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Private request factory and execution-scope support.""" + +from __future__ import annotations + +import asyncio +import inspect +import weakref +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from typing import Any, TypeAlias, TypeGuard, cast + +from agent_framework import ( + AgentResponse, + AgentResponseUpdate, + FunctionalWorkflowAgent, + ResponseStream, + SupportsAgentRun, + WorkflowAgent, +) + +HostedAgent: TypeAlias = SupportsAgentRun | FunctionalWorkflowAgent +AgentFactory: TypeAlias = Callable[[], HostedAgent | Awaitable[HostedAgent]] +WorkflowAgentTypes: TypeAlias = WorkflowAgent | FunctionalWorkflowAgent + + +def is_workflow_agent(agent: object) -> TypeGuard[WorkflowAgentTypes]: + """Recognize the built-in workflow adapters, including subclasses.""" + return isinstance(agent, (WorkflowAgent, FunctionalWorkflowAgent)) + + +def validate_agent_source(agent: HostedAgent | None, agent_factory: AgentFactory | None) -> None: + if (agent is None) == (agent_factory is None): + raise ValueError("Provide exactly one of agent or agent_factory.") + if agent is not None and is_workflow_agent(agent): + raise TypeError( + "Workflow agents must be supplied through agent_factory. " + "The factory must create a new workflow and new mutable executors for each request." + ) + if agent_factory is not None and not callable(agent_factory): + raise TypeError("agent_factory must be a callable accepting no arguments, not a coroutine object.") + + +class AgentFactoryResolver: + """Resolve request agents without retaining completed workflow runtimes.""" + + def __init__(self, factory: AgentFactory) -> None: + self._factory = factory + self._seen: weakref.WeakValueDictionary[int, object] = weakref.WeakValueDictionary() + + async def resolve(self) -> HostedAgent: + result = self._factory() + agent = await result if inspect.isawaitable(result) else result + if not isinstance(agent, (SupportsAgentRun, FunctionalWorkflowAgent)): + raise TypeError("agent_factory must return an agent implementing SupportsAgentRun or a workflow agent.") + if is_workflow_agent(agent): + workflow = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow # pyright: ignore[reportPrivateUsage] + for value in (agent, workflow): + if self._seen.get(id(value)) is value: + raise RuntimeError( + "agent_factory reused a workflow agent or workflow. Create a new workflow and " + "new mutable executors for each request." + ) + for value in (agent, workflow): + self._seen[id(value)] = value + return agent + + +@dataclass +class _ScopeLock: + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + users: int = 0 + + +class ScopeLocks: + """Serialize a scope while retaining locks for both holders and waiters.""" + + def __init__(self) -> None: + self._entries: dict[tuple[str | None, ...], _ScopeLock] = {} + + @asynccontextmanager + async def hold(self, key: tuple[str | None, ...]) -> AsyncGenerator[None]: + entry = self._entries.setdefault(key, _ScopeLock()) + entry.users += 1 + try: + async with entry.lock: + yield + finally: + entry.users -= 1 + if not entry.users: + del self._entries[key] + + +async def close_run_iterator(iterator: AsyncIterator[Any]) -> None: + """Close the concrete run iterator, including known ResponseStream wrappers.""" + wrappers: list[ResponseStream[AgentResponseUpdate, AgentResponse]] = [] + current: Any = iterator + while isinstance(current, ResponseStream): + wrapper = cast(ResponseStream[AgentResponseUpdate, AgentResponse], current) + wrappers.append(wrapper) + current = wrapper._iterator # pyright: ignore[reportPrivateUsage] + try: + close = getattr(current, "aclose", None) + if close is not None: + await close() + finally: + for wrapper in reversed(wrappers): + await wrapper._run_cleanup_hooks() # pyright: ignore[reportPrivateUsage] 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 2b321c67046..49c084ed386 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,14 +1,67 @@ # Copyright (c) Microsoft. All rights reserved. -from agent_framework import AgentSession, SupportsAgentRun +from __future__ import annotations + +import json +import sys +from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, suppress +from typing import cast + +from agent_framework import ( + AgentSession, + CheckpointStorage, + FunctionalWorkflowAgent, + SessionStore, + WorkflowAgent, + WorkflowRunState, +) +from agent_framework._filesystem import _storage_key_segment # pyright: ignore[reportPrivateUsage] from agent_framework._telemetry import mark_feature_used +from anyio import CancelScope from azure.ai.agentserver.core import 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 Send from typing_extensions import Any, AsyncGenerator +from ._agent_factory import ( + AgentFactory, + AgentFactoryResolver, + HostedAgent, + ScopeLocks, + close_run_iterator, + is_workflow_agent, + validate_agent_source, +) from ._feature_usage import FeatureIndex +from ._state_store import ( + ContextScopedStoreProvider, + StoreProvider, + _CheckpointStorageWithErrors, # pyright: ignore[reportPrivateUsage] + _InvocationsAgentSessionStoreProvider, # pyright: ignore[reportPrivateUsage] + _InvocationsCheckpointStoreProvider, # pyright: ignore[reportPrivateUsage] +) + +_WORKFLOW_STATE_KEY = "_foundry_invocations_workflow" +_PENDING_REQUEST_ERROR = ( + "Invocations plain text cannot answer workflow pending requests or approvals. " + "Use a protocol that supports structured workflow responses." +) + + +class _InvocationStreamingResponse(StreamingResponse): + async def stream_response(self, send: Send) -> None: + # Starlette does not close a body iterator suspended at yield when send fails. + # Close it in the consumer task, where task-affine agent resources were entered. + try: + await super().stream_response(send) + except BaseException as exc: + with suppress(StopAsyncIteration): + await cast(AsyncGenerator[str], self.body_iterator).athrow(exc) + raise + finally: + await cast(AsyncGenerator[str], self.body_iterator).aclose() class InvocationsHostServer(InvocationAgentServerHost): @@ -16,25 +69,48 @@ class InvocationsHostServer(InvocationAgentServerHost): def __init__( self, - agent: SupportsAgentRun, + agent: HostedAgent | None = None, *, + agent_factory: AgentFactory | None = None, + checkpoint_store_provider: ContextScopedStoreProvider[CheckpointStorage] | None = None, + agent_session_store_provider: StoreProvider[SessionStore] | None = None, openapi_spec: dict[str, Any] | None = None, **kwargs: Any, ) -> None: """Initialize an InvocationsHostServer. Args: - agent: The agent to handle responses for. + agent: An ordinary agent instance. Supply workflows through agent_factory instead. + agent_factory: A sync or async callable creating an agent for each request. Recreate + workflows, mutable executors, providers, and tools. An async context manager + returned by the factory is entered and exited within that request. + checkpoint_store_provider: Optional provider for user/session-scoped workflow checkpoints. + agent_session_store_provider: Optional provider for persisted workflow provider state. openapi_spec: The OpenAPI specification for the server. **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. + Responses contain plain text; "stream": true streams text as text/event-stream. + Workflow factories must use stable workflow names and executor IDs across requests. + The text-only exchange cannot answer pending workflow requests or approvals. + Functional workflows support fresh messages after clean completion, but not + pending or interrupted continuation. Such continuation fails explicitly. + Factories own cleanup of nested resources not exposed by an async context manager. """ + validate_agent_source(agent, agent_factory) super().__init__(openapi_spec=openapi_spec, **kwargs) self._agent = agent + self._agent_resolver = AgentFactoryResolver(agent_factory) if agent_factory is not None else None + self._scope_locks = ScopeLocks() + self._checkpoint_storage_provider = ( + _InvocationsCheckpointStoreProvider() if checkpoint_store_provider is None else checkpoint_store_provider + ) + self._agent_session_storage_provider = ( + _InvocationsAgentSessionStoreProvider() + if agent_session_store_provider is None + else agent_session_store_provider + ) self._sessions: dict[str, AgentSession] = {} self.invoke_handler(self._handle_invoke) mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) @@ -70,6 +146,143 @@ def _partition_key(self) -> str: return context.session_id + @asynccontextmanager + async def _request_agent(self, scope: tuple[str | None, str]) -> AsyncGenerator[tuple[HostedAgent, CancelScope]]: + async with self._scope_locks.hold(scope): + # Enter before the agent's own task groups. Adding a new cancel scope + # around their exit would violate AnyIO's required scope nesting. + with CancelScope() as cleanup_scope: + agent = await cast(AgentFactoryResolver, self._agent_resolver).resolve() + resources = AsyncExitStack() + try: + if isinstance(agent, AbstractAsyncContextManager): + await resources.enter_async_context(agent) + yield agent, cleanup_scope + finally: + cleanup_scope.shield = True + exc_info = sys.exc_info() + if isinstance(exc_info[1], GeneratorExit): + await resources.aclose() + else: + await resources.__aexit__(*exc_info) + + @asynccontextmanager + async def _workflow_session( + self, agent: WorkflowAgent | FunctionalWorkflowAgent, storage_id: str + ) -> AsyncGenerator[tuple[AgentSession, CheckpointStorage]]: + context = get_request_context() + storage = _CheckpointStorageWithErrors( + self._checkpoint_storage_provider.get_store( + config=self.config, context_id=storage_id, platform_context=context + ) + ) + sessions = self._agent_session_storage_provider.get_store(config=self.config, platform_context=context) + workflow = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow # pyright: ignore[reportPrivateUsage] + kind = "graph" if isinstance(agent, WorkflowAgent) else "functional" + if isinstance(agent, WorkflowAgent): + storage.observe_workflow(agent) + session = await sessions.get(storage_id) + saved_marker = session.state.get(_WORKFLOW_STATE_KEY) if session is not None else None + marker = cast(dict[str, Any], saved_marker) if isinstance(saved_marker, dict) else None + if session is not None and ( + marker is None or marker.get("name") != workflow.name or marker.get("kind") != kind + ): + raise RuntimeError("The stored Invocations workflow name or kind does not match the factory result.") + checkpoint = await storage.get_latest(workflow_name=workflow.name) + if checkpoint is not None and checkpoint.workflow_name != workflow.name: + raise RuntimeError("The stored Invocations checkpoint does not match the workflow name.") + if marker is not None and checkpoint is None: + raise RuntimeError("The existing Invocations workflow session is missing its required checkpoint.") + if checkpoint is not None and session is None: + raise RuntimeError("The existing Invocations workflow checkpoint is missing its required agent session.") + if marker is not None and marker.get("checkpoint_failed"): + raise RuntimeError("The previous Invocations workflow run has incomplete checkpoint persistence.") + if isinstance(agent, FunctionalWorkflowAgent) and marker is not None and marker.get("completed") is not True: + raise RuntimeError("Invocations cannot continue a pending or interrupted functional workflow.") + if isinstance(agent, WorkflowAgent) and checkpoint is not None: + if checkpoint.pending_request_info_events: + raise RuntimeError(_PENDING_REQUEST_ERROR) + await agent.workflow.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) + if storage.save_error is not None: + raise storage.save_error + if agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS: + raise RuntimeError(_PENDING_REQUEST_ERROR) + if session is None: + session = AgentSession(session_id=storage_id) + marker = {"name": workflow.name, "kind": kind, "completed": False} + session.state[_WORKFLOW_STATE_KEY] = marker + # Record the attempt before execution. A failed first run must not look like a new session. + await sessions.set(storage_id, session) + try: + yield session, storage + if storage.save_error is not None: + raise storage.save_error + pending = ( + agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + if isinstance(agent, WorkflowAgent) + else bool(agent.pending_requests) + ) + if pending: + raise RuntimeError(_PENDING_REQUEST_ERROR) + if await storage.get_latest(workflow_name=workflow.name) is None: + raise RuntimeError("The Invocations workflow did not persist a required checkpoint.") + marker["completed"] = True + finally: + if storage.save_error is not None: + marker["checkpoint_failed"] = True + with CancelScope(shield=True): + await sessions.set(storage_id, session) + + @asynccontextmanager + async def _factory_session( + self, agent: HostedAgent, storage_id: str + ) -> AsyncGenerator[tuple[AgentSession, dict[str, Any]]]: + if is_workflow_agent(agent): + async with self._workflow_session(agent, storage_id) as (session, storage): + yield session, {"checkpoint_storage": storage} + else: + yield self._sessions.setdefault(storage_id, AgentSession(session_id=storage_id)), {} + + async def _handle_factory_invoke(self, user_message: Any, *, stream: bool) -> Response: + context = get_request_context() + scope = (context.user_id, cast(str, context.session_id)) + storage_id = _storage_key_segment(json.dumps(scope, ensure_ascii=False), encoded_prefix="~invocations-") + + if stream: + + async def stream_response() -> AsyncGenerator[str]: + async with ( + self._request_agent(scope) as (agent, cleanup_scope), + self._factory_session(agent, storage_id) as ( + session, + run_kwargs, + ), + ): + iterator = agent.run(user_message, session=session, stream=True, **run_kwargs).__aiter__() + try: + async for update in iterator: + if update.text: + yield update.text + finally: + cleanup_scope.shield = True + await close_run_iterator(iterator) + + return _InvocationStreamingResponse( + stream_response(), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, + ) + + async with ( + self._request_agent(scope) as (agent, _), + self._factory_session(agent, storage_id) as ( + session, + run_kwargs, + ), + ): + response = await agent.run([user_message], session=session, **run_kwargs) + return Response(content=response.text) + async def _handle_invoke(self, request: Request) -> Response: """Invoke the agent with the given request.""" try: @@ -87,12 +300,16 @@ async def _handle_invoke(self, request: Request) -> Response: return StreamingResponse(content=error, status_code=400) return Response(content=error, status_code=400) + if self._agent_resolver is not None: + return await self._handle_factory_invoke(user_message, stream=stream) + + agent = cast(HostedAgent, self._agent) session = self._sessions.setdefault(session_id, AgentSession(session_id=session_id)) if stream: async def stream_response() -> AsyncGenerator[str]: - async for update in self._agent.run(user_message, session=session, stream=True): + async for update in agent.run(user_message, session=session, stream=True): if update.text: yield update.text @@ -102,5 +319,5 @@ async def stream_response() -> AsyncGenerator[str]: headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) - response = await self._agent.run([user_message], session=session) + 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..f6092d53d46 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -9,30 +9,37 @@ import logging import os import re -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Generator, Mapping, Sequence +import sys +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, 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 from urllib.parse import urlparse from agent_framework import ( + AgentResponse, AgentResponseUpdate, + AgentSession, ChatOptions, CheckpointStorage, Content, ContextProvider, + FunctionalWorkflowAgent, HistoryProvider, InMemoryHistoryProvider, Message, RawAgent, + SessionContext, SessionStore, SupportsAgentRun, UsageDetails, WorkflowAgent, + WorkflowRunState, add_usage_details, ) from agent_framework._telemetry import mark_feature_used from agent_framework.exceptions import AgentFrameworkException +from anyio import CancelScope from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.responses import ( ResponseContext, @@ -73,6 +80,15 @@ from mcp import McpError from typing_extensions import Any +from ._agent_factory import ( + AgentFactory, + AgentFactoryResolver, + HostedAgent, + ScopeLocks, + close_run_iterator, + is_workflow_agent, + validate_agent_source, +) from ._feature_usage import FeatureIndex from ._state_store import ( AgentSessionStoreProvider, @@ -81,6 +97,7 @@ FunctionApprovalStore, FunctionApprovalStoreProvider, StoreProvider, + _CheckpointStorageWithErrors, # pyright: ignore[reportPrivateUsage] ) logger = logging.getLogger(__name__) @@ -130,6 +147,8 @@ def _create_response_event_stream(context: ResponseContext) -> ResponseEventStre _T = TypeVar("_T") +_WorkflowRunFactory = Callable[[], AsyncIterator[AgentResponseUpdate]] +_WorkflowRun = Callable[[_WorkflowRunFactory], AsyncIterator[AgentResponseUpdate]] # Sentinel put on the internal queue by _SignalledIterator's driver task to signal that the # wrapped iterator is exhausted (distinct from `None`, which is a valid item value). @@ -175,6 +194,7 @@ def __init__(self, iterator: AsyncIterator[_T], *events: asyncio.Event) -> None: self._queue: asyncio.Queue[Any] = asyncio.Queue(maxsize=1) # The background task that drives the wrapped iterator. self._driver: asyncio.Task[None] | None = None + self._closing = False @property def signalled(self) -> bool: @@ -190,16 +210,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: + # Consumer cancellation is forwarded explicitly by aclose(). Shield the + # driver from repeated AnyIO cancellation while its iterator unwinds. + with CancelScope(shield=True): + try: + while True: + try: + item: Any = await self._iterator.__anext__() + except StopAsyncIteration: + break + await self._queue.put(item) + finally: + await close_run_iterator(self._iterator) + except Exception as exc: + if self._closing: + raise + await self._queue.put(exc) + else: + await self._queue.put(_STOP_SENTINEL) async def __anext__(self) -> _T: if self._driver is None: @@ -213,15 +242,25 @@ async def __anext__(self) -> _T: await asyncio.wait([get_task, *waiters], return_when=asyncio.FIRST_COMPLETED) if any(waiter.done() for waiter in waiters): self._signalled = True + self._closing = True self._driver.cancel() - with suppress(BaseException): + with CancelScope(shield=True), suppress(asyncio.CancelledError): await self._driver get_task.cancel() with suppress(BaseException): await get_task raise StopAsyncIteration item = get_task.result() + except asyncio.CancelledError: + if not self._closing: + self._closing = True + self._driver.cancel() + raise finally: + if not get_task.done(): + get_task.cancel() + with suppress(asyncio.CancelledError): + await get_task for waiter in waiters: if not waiter.done(): waiter.cancel() @@ -242,8 +281,10 @@ async def aclose(self) -> None: """ if self._driver is None: return - self._driver.cancel() - with suppress(BaseException): + if not self._closing: + self._closing = True + self._driver.cancel() + with CancelScope(shield=True), suppress(asyncio.CancelledError): await self._driver @@ -373,13 +414,103 @@ def consent_url_from_error(exc: BaseException) -> list[ConsentError] | None: # region ResponsesHostServer +@dataclass(frozen=True) +class _AgentConfiguration: + workflow: bool + agent_server_history: bool + client_stores_by_default: bool + hosted_history: bool + + +def _validate_agent_configuration( + agent: HostedAgent, + history_source: Literal["agent_server", "agent"], + options: ResponsesServerOptions | None, +) -> _AgentConfiguration: + workflow = is_workflow_agent(agent) + if isinstance(agent, WorkflowAgent) 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." + ) + if isinstance(agent, FunctionalWorkflowAgent): + if agent._workflow._checkpoint_storage is not None: # pyright: ignore[reportPrivateUsage] + raise RuntimeError("The hosting infrastructure must manage the workflow checkpoint storage.") + if options and options.resilient_background: + raise RuntimeError( + "Functional workflow recovery is not supported: checkpoints do not preserve buffered step output. " + "Use a graph workflow for resilient_background=True." + ) + if options and options.resilient_background and not workflow: + 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 workflow: + raise RuntimeError( + "steerable_conversations=True is only supported for non-workflow agents. " + "Steering cannot be provided reliably for workflow agents." + ) + agent_server_history = history_source == "agent_server" + client_stores_by_default = False + if agent_server_history and not workflow: + 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, agent_server_history, client_stores_by_default, agent_server_history and not workflow + ) + + +def _initialize_agent_history(agent: HostedAgent, configuration: _AgentConfiguration) -> None: + if configuration.hosted_history and isinstance(agent, RawAgent): + 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)) + + class ResponsesHostServer(ResponsesAgentServerHost): """A responses server host for an agent.""" def __init__( self, - agent: SupportsAgentRun, + agent: HostedAgent | None = None, *, + agent_factory: AgentFactory | None = None, prefix: str = "", options: ResponsesServerOptions | None = None, store: ResponseProviderProtocol | None = None, @@ -392,7 +523,11 @@ def __init__( """Initialize a ResponsesHostServer. Args: - agent: The agent to handle responses for. + agent: An ordinary agent instance. Supply workflow agents through agent_factory. + agent_factory: A no-argument sync or async factory creating one agent per request. + Recreate mutable workflows and executors with stable names and IDs. Returned + async context managers remain open through execution, streaming, and persistence. + Functional workflows do not support resilient background recovery. prefix: The URL prefix for the server. options: Optional server options. store: Optional response store for input and history look up. @@ -442,92 +577,18 @@ def __init__( """ if history_source not in ("agent_server", "agent"): raise ValueError("history_source must be either 'agent_server' or '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 - - # No caller-owned agent state is mutated until all validation and base-host construction succeed. + validate_agent_source(agent, agent_factory) + configuration = _validate_agent_configuration(agent, history_source, options) if agent is not None else None 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._history_source: Literal["agent_server", "agent"] = history_source + self._host_options = options + self._configuration = configuration + self._resilient_background = bool(options and options.resilient_background) + self._agent = agent + self._agent_resolver = AgentFactoryResolver(agent_factory) if agent_factory is not None else None + self._scope_locks = ScopeLocks() + if agent is not None and configuration is not None: + _initialize_agent_history(agent, configuration) # Storage providers self._checkpoint_storage_provider = ( @@ -589,18 +650,83 @@ 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() + if self._agent_resolver is None: + async with AsyncExitStack() as resources: + inner = self._handle_prepared_response( + request, + context, + cancellation_signal, + response_event_stream, + cast(HostedAgent, self._agent), + cast(_AgentConfiguration, self._configuration), + resources, + ) + async with aclosing(inner): + async for event in inner: + yield event + return + terminal_event: ResponseStreamEvent | None = None + try: + async with AsyncExitStack() as locks: + platform_context = get_request_context() + if self.config.is_hosted and not platform_context.user_id: + raise RuntimeError("The hosted request context is missing user_id.") + scope_id = context.conversation_id or request.get("previous_response_id") or context.response_id + _validate_checkpoint_context_id(scope_id) + await locks.enter_async_context(self._scope_locks.hold((platform_context.user_id, scope_id))) + agent = await self._agent_resolver.resolve() + configuration = _validate_agent_configuration(agent, self._history_source, self._host_options) + _initialize_agent_history(agent, configuration) + # This scope must precede any task-affine MCP scopes entered by the agent. + with CancelScope() as cleanup_scope: + resources = AsyncExitStack() + try: + inner = self._handle_prepared_response( + request, + context, + cancellation_signal, + response_event_stream, + agent, + configuration, + resources, + ) + async with aclosing(inner): + 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: + cleanup_scope.shield = True + exc_info = sys.exc_info() + if isinstance(exc_info[1], GeneratorExit): + await resources.aclose() + else: + await resources.__aexit__(*exc_info) + if terminal_event is not None: + yield terminal_event + except Exception as ex: + logger.exception("Failed to prepare or release the request agent") + for event in self._emit_failure(response_event_stream, None, ex): + yield event + + async def _handle_prepared_response( + self, + request: CreateResponse, + context: ResponseContext, + cancellation_signal: asyncio.Event, + response_event_stream: ResponseEventStream, + agent: HostedAgent, + 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 @@ -608,7 +734,10 @@ async def _handle_response( # 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._agent_resolver is 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 +763,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 +778,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 = cast(SupportsAgentRun, agent).create_session() await session_storage.set(context.conversation_id or context.response_id, session) except Exception as save_error: logger.error( @@ -678,12 +807,20 @@ 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, 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, + cast(SupportsAgentRun, agent), + configuration, + ) try: async for event in inner: @@ -732,7 +869,7 @@ async def _load_input() -> list[Message]: 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 self._history_source != "agent_server": return [] history = await context.get_history() return await _output_items_to_messages(history, approval_storage=approval_storage) @@ -758,6 +895,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. @@ -797,7 +936,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 +955,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 +968,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 +977,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 +985,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), # pyright: ignore[reportUnknownMemberType] context.shutdown, cancellation_signal, ) @@ -867,12 +1006,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 " @@ -884,7 +1023,8 @@ async def _handle_inner_agent( request_failure = misconfigured try: if not stored_output_violation: - await session_storage.set(session_save_id, session) + with CancelScope(shield=True): + await session_storage.set(session_save_id, session) except Exception as save_error: save_failure = save_error if request_interrupted: @@ -912,11 +1052,112 @@ async def _handle_inner_workflow( response_event_stream: ResponseEventStream, tracker: _OutputItemTracker, cancellation_signal: asyncio.Event, + agent: HostedAgent, ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: - """Handle the creation of a response for a workflow agent.""" - if not isinstance(self._agent, WorkflowAgent): + if not is_workflow_agent(agent): raise RuntimeError("Agent is not a workflow agent.") + platform_context = get_request_context() + save_id = context.conversation_id or context.response_id + load_id = save_id if context.is_recovery else context.conversation_id or request.get("previous_response_id") + _validate_checkpoint_context_id(save_id) + if load_id is not None: + _validate_checkpoint_context_id(load_id) + sessions = self._session_storage_provider.get_store(config=self.config, platform_context=platform_context) + session = await sessions.get(load_id) if load_id is not None else None + workflow = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow # pyright: ignore[reportPrivateUsage] + kind = "graph" if isinstance(agent, WorkflowAgent) else "functional" + state_key = "_foundry_responses_workflow" + saved_marker = session.state.get(state_key) if session is not None else None + marker = cast(dict[str, Any], saved_marker) if isinstance(saved_marker, dict) else None + if session is not None and ( + marker is None or marker.get("name") != workflow.name or marker.get("kind") != kind + ): + raise RuntimeError("The stored Responses workflow name or kind does not match the factory result.") + if marker is not None and marker.get("checkpoint_failed"): + raise RuntimeError("The previous Responses workflow run has incomplete checkpoint persistence.") + had_session = session is not None + if session is None: + session = AgentSession() + prior_completed = marker is not None and marker.get("completed") is True + state_marker: dict[str, Any] = {"name": workflow.name, "kind": kind, "completed": False} + attempt_saved = False + attempt_started = False + + async def run_workflow(stream_factory: _WorkflowRunFactory) -> AsyncGenerator[AgentResponseUpdate]: + nonlocal attempt_saved, attempt_started + if cancellation_signal.is_set(): + return + if not attempt_started: + session.state[state_key] = state_marker + attempt_saved = True + await sessions.set(save_id, session) + if cancellation_signal.is_set(): + return + attempt_started = True + iterator = stream_factory() + try: + async for update in iterator: + yield update + finally: + await close_run_iterator(iterator) + + try: + if isinstance(agent, FunctionalWorkflowAgent): + inner = self._handle_functional_workflow( + request, context, tracker, cancellation_signal, agent, had_session, prior_completed, run_workflow + ) + else: + inner = self._handle_graph_workflow( + request, + context, + response_event_stream, + tracker, + cancellation_signal, + agent, + session if agent.context_providers else None, + had_session, + state_marker, + run_workflow, + ) + async with aclosing(inner): + async for event in inner: + if isinstance(event, ResponseCheckpointEvent): + await sessions.set(save_id, session) + yield event + if attempt_started and not cancellation_signal.is_set() and not context.shutdown.is_set(): + pending = ( + agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS + if isinstance(agent, WorkflowAgent) + else bool(agent.pending_requests) + ) + session.state[state_key]["completed"] = not pending + finally: + with CancelScope(shield=True): + if attempt_started: + await sessions.set(save_id, session) + elif attempt_saved: + # The initial write can finish just as cancellation arrives. + # No workflow work ran, so preserve the previous conversation. + if had_session and load_id == save_id: + session.state[state_key] = saved_marker + await sessions.set(save_id, session) + else: + await sessions.delete(save_id) + async def _handle_graph_workflow( + self, + request: CreateResponse, + context: ResponseContext, + response_event_stream: ResponseEventStream, + tracker: _OutputItemTracker, + cancellation_signal: asyncio.Event, + agent: WorkflowAgent, + session: AgentSession | None, + had_session: bool, + state_marker: dict[str, Any], + run_workflow: _WorkflowRun, + ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: + """Handle the creation of a response for a workflow agent.""" try: request_context = get_request_context() approval_storage = self._function_approval_storage_provider.get_store( @@ -936,15 +1177,25 @@ async def _handle_inner_workflow( # workflow from the last checkpoint. checkpoint_save_id = context.conversation_id or context.response_id _validate_checkpoint_context_id(checkpoint_save_id) - checkpoint_storage = self._checkpoint_storage_provider.get_store( - config=self.config, - context_id=checkpoint_save_id, - platform_context=request_context, + checkpoint_storage = _CheckpointStorageWithErrors( + self._checkpoint_storage_provider.get_store( + config=self.config, + context_id=checkpoint_save_id, + platform_context=request_context, + ) ) + checkpoint_storage.observe_workflow(agent) if context.is_recovery: if not self._resilient_background: raise RuntimeError("Recovery mode is only supported when resilient_background=True.") + persisted_messages = ( + await _output_items_to_messages( + context.persisted_response.get("output", []), approval_storage=approval_storage + ) + if context.persisted_response is not None and session is not None + else [] + ) # Resume from the workflow checkpoint durably paired with the last persisted response # snapshot (recorded in that snapshot's own metadata) -- NOT simply the latest workflow # checkpoint in storage, which may be ahead of what response.output actually reflects if @@ -952,19 +1203,36 @@ async def _handle_inner_workflow( checkpoint_id = response_event_stream.internal_metadata.get(_LATEST_CHECKPOINT_ID_KEY) 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 + run_stream = run_workflow( + lambda: self._resume_workflow_from_checkpoint( + checkpoint_id, + checkpoint_storage, + context.response_id, + agent, + session, + input_messages, + persisted_messages, + ) ) 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; " "resuming from the latest checkpoint", latest_checkpoint.checkpoint_id, ) - run_stream = self._resume_workflow_from_checkpoint( - latest_checkpoint.checkpoint_id, checkpoint_storage, context.response_id + recovery_checkpoint_id = latest_checkpoint.checkpoint_id + run_stream = run_workflow( + lambda: self._resume_workflow_from_checkpoint( + recovery_checkpoint_id, + checkpoint_storage, + context.response_id, + agent, + session, + input_messages, + persisted_messages, + ) ) else: # No checkpoint was ever paired with a persisted response snapshot (e.g. the crash @@ -974,10 +1242,13 @@ async def _handle_inner_workflow( logger.debug( "Serving recovery request with no prior workflow checkpoint; replaying original input" ) - run_stream = self._agent.run( - input_messages, - stream=True, - checkpoint_storage=checkpoint_storage, + run_stream = run_workflow( + lambda: agent.run( + input_messages, + stream=True, + session=session, + checkpoint_storage=checkpoint_storage, + ) ) else: # Determine the latest checkpoint (if any) so we can resume the @@ -991,18 +1262,26 @@ async def _handle_inner_workflow( if checkpoint_load_id is not None: _validate_checkpoint_context_id(checkpoint_load_id) if checkpoint_load_id != checkpoint_save_id: - restore_checkpoint_storage = self._checkpoint_storage_provider.get_store( - config=self.config, - context_id=checkpoint_load_id, - platform_context=request_context, + restore_checkpoint_storage = _CheckpointStorageWithErrors( + self._checkpoint_storage_provider.get_store( + config=self.config, + 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 raise RuntimeError( f"Cannot find an existing workflow checkpoint for previous_response_id={previous_response_id}." ) + if latest_checkpoint is None and ( + had_session or (context.conversation_id is not None and await context.get_history()) + ): + raise RuntimeError("The existing conversation is missing its required workflow checkpoint.") + if latest_checkpoint is not None and not had_session and agent.context_providers: + raise RuntimeError("The workflow checkpoint is missing its required outer agent session.") if latest_checkpoint is not None: # If we have a prior checkpoint, restore it first (drive the workflow @@ -1014,16 +1293,19 @@ 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 restores the workflow's 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_checkpoint_id = latest_checkpoint.checkpoint_id restore_iter = _SignalledIterator( - self._agent.run( - stream=True, - checkpoint_id=latest_checkpoint.checkpoint_id, - checkpoint_storage=restore_checkpoint_storage, + run_workflow( + lambda: agent.run( + stream=True, + checkpoint_id=restore_checkpoint_id, + checkpoint_storage=restore_checkpoint_storage, + ) ), context.shutdown, cancellation_signal, @@ -1031,6 +1313,12 @@ async def _handle_inner_workflow( async with aclosing(restore_iter): async for _ in restore_iter: pass + if restore_checkpoint_storage.save_error is not None: + state_marker["checkpoint_failed"] = True + raise restore_checkpoint_storage.save_error + if checkpoint_storage.save_error is not None: + state_marker["checkpoint_failed"] = True + raise checkpoint_storage.save_error if restore_iter.signalled: if context.shutdown.is_set(): await context.exit_for_recovery() @@ -1042,17 +1330,23 @@ async def _handle_inner_workflow( if cancellation_signal.is_set(): return - run_stream = self._agent.run( - input_messages, - stream=True, - checkpoint_storage=checkpoint_storage, + run_stream = run_workflow( + lambda: agent.run( + input_messages, + stream=True, + session=session, + checkpoint_storage=checkpoint_storage, + ) ) main_iter = _SignalledIterator(run_stream, context.shutdown, cancellation_signal) async with aclosing(main_iter): async for update in main_iter: + if checkpoint_storage.save_error is not None: + state_marker["checkpoint_failed"] = True + raise checkpoint_storage.save_error 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 @@ -1079,6 +1373,9 @@ async def _handle_inner_workflow( content, message_id=update.message_id, approval_storage=approval_storage ): yield event + if checkpoint_storage.save_error is not None: + state_marker["checkpoint_failed"] = True + raise checkpoint_storage.save_error # Cancellation needs no extra action here (the loop above already stopped); shutdown # does, but only if it's what actually stopped the loop, not a natural completion. if main_iter.signalled and context.shutdown.is_set(): @@ -1092,6 +1389,10 @@ async def _resume_workflow_from_checkpoint( checkpoint_id: str, checkpoint_storage: CheckpointStorage, response_id: str, + agent: WorkflowAgent, + session: AgentSession | None, + input_messages: list[Message], + persisted_messages: list[Message], ) -> AsyncGenerator[AgentResponseUpdate]: """Resume a crashed background workflow run, forwarding every event it produces. @@ -1106,9 +1407,24 @@ 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 + session_context: SessionContext | None = None + if session is not None: + session_context = SessionContext( + session_id=session.session_id, + service_session_id=session.service_session_id, + input_messages=input_messages, + options={}, + ) + for provider in agent.context_providers: + if isinstance(provider, HistoryProvider) and not provider.load_messages: + continue + await provider.before_run( + agent=agent, + session=session, + context=session_context, + state=session.state.setdefault(provider.source_id, {}), + ) + updates: list[AgentResponseUpdate] = [] async for event in agent.workflow.run( stream=True, checkpoint_id=checkpoint_id, @@ -1117,7 +1433,107 @@ async def _resume_workflow_from_checkpoint( for update in agent._convert_workflow_event_to_agent_response_updates( # pyright: ignore[reportPrivateUsage] response_id, event ): + if session_context is not None: + updates.append(update) yield update + if session_context is not None: + response = AgentResponse.from_updates(updates) + session_context._response = AgentResponse( # pyright: ignore[reportPrivateUsage] + messages=[*persisted_messages, *response.messages] + ) + await agent._run_after_providers(session=session, context=session_context) # pyright: ignore[reportPrivateUsage] + + async def _handle_functional_workflow( + self, + request: CreateResponse, + context: ResponseContext, + tracker: _OutputItemTracker, + cancellation_signal: asyncio.Event, + agent: FunctionalWorkflowAgent, + had_session: bool, + prior_completed: bool, + run_workflow: _WorkflowRun, + ) -> AsyncGenerator[ResponseStreamEvent]: + if context.is_recovery: + raise RuntimeError("Functional workflow recovery cannot preserve buffered step output.") + previous_id = request.get("previous_response_id") + if previous_id is not None and context.conversation_id is not None: + raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.") + platform_context = get_request_context() + save_id = context.conversation_id or context.response_id + load_id = context.conversation_id or previous_id + storage = self._checkpoint_storage_provider.get_store( + config=self.config, context_id=save_id, platform_context=platform_context + ) + restore_storage = storage + if load_id is not None and load_id != save_id: + restore_storage = self._checkpoint_storage_provider.get_store( + config=self.config, context_id=load_id, platform_context=platform_context + ) + checkpoint = await restore_storage.get_latest(workflow_name=agent._workflow.name) # pyright: ignore[reportPrivateUsage] + if checkpoint is None and ( + had_session + or previous_id is not None + or (context.conversation_id is not None and await context.get_history()) + ): + raise RuntimeError("The existing conversation is missing its required workflow checkpoint.") + if checkpoint is not None and not had_session: + raise RuntimeError("The functional workflow checkpoint is missing its required agent session.") + approval_storage = self._function_approval_storage_provider.get_store( + config=self.config, platform_context=platform_context + ) + messages = await _items_to_messages(await context.get_input_items(), approval_storage=approval_storage) + run_kwargs: dict[str, Any] = {"messages": messages, "checkpoint_storage": storage} + if checkpoint is not None and not prior_completed: + pending = checkpoint.pending_request_info_events + if not pending: + raise RuntimeError("Cannot continue an interrupted functional workflow without pending requests.") + responses: dict[str, Any] = {} + for message in messages: + for content in message.contents: + request_id = content.call_id if content.type == "function_result" else content.id + if request_id is None or request_id not in pending or request_id in responses: + raise ValueError("The input does not match an authorized pending functional workflow request.") + pending_request = pending[request_id] + if content.type == "function_result": + responses[request_id] = content if pending_request.response_type is Content else content.result + elif content.type == "function_approval_response" and pending_request.response_type is bool: + responses[request_id] = content.approved + elif content.type == "function_approval_response" and pending_request.response_type is Content: + responses[request_id] = content + else: + raise ValueError( + "This pending functional workflow request requires a matching function result." + ) + if not responses: + raise ValueError("Pending functional workflow requests require structured responses.") + run_kwargs = { + "checkpoint_storage": storage, + "checkpoint_id": checkpoint.checkpoint_id, + "responses": responses, + } + if cancellation_signal.is_set(): + return + + async def execute() -> AsyncGenerator[AgentResponseUpdate]: + if checkpoint is not None and not prior_completed: + # Copy the authorized checkpoint only once restoration actually starts. + await storage.save(checkpoint) + agent_stream = agent.run(stream=True, **run_kwargs) + try: + async for update in agent_stream: + yield update + finally: + await close_run_iterator(agent_stream) + + iterator = _SignalledIterator(run_workflow(execute), context.shutdown, cancellation_signal) + async with aclosing(iterator): + async for update in iterator: + for content in update.contents: + async for event in tracker.handle( + content, message_id=update.message_id, approval_storage=approval_storage + ): + yield event @staticmethod def _emit_failure( 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..aca8f4ca57f 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 @@ -1,9 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Callable from datetime import datetime -from typing import Generic, Protocol, TypeVar +from functools import wraps +from types import CoroutineType +from typing import Any, Generic, ParamSpec, Protocol, TypeVar from agent_framework import ( AgentSession, @@ -11,6 +15,7 @@ CheckpointStorage, Content, SessionStore, + WorkflowAgent, WorkflowCheckpoint, WorkflowCheckpointException, ) @@ -18,6 +23,49 @@ from azure.ai.agentserver.core.storage import FoundryStateStore, FoundryStorageConflictError StoreT = TypeVar("StoreT") +ResultT = TypeVar("ResultT") +ParametersT = ParamSpec("ParametersT") + + +class _CheckpointStorageWithErrors: # pyright: ignore[reportUnusedClass] + """Remember checkpoint write errors even when graph execution logs and ignores them.""" + + def __init__(self, storage: CheckpointStorage) -> None: + self._storage = storage + self.save_error: Exception | None = None + self.load = storage.load + self.list_checkpoints = storage.list_checkpoints + self.delete = storage.delete + self.get_latest = storage.get_latest + self.list_checkpoint_ids = storage.list_checkpoint_ids + + async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: + try: + return await self._storage.save(checkpoint) + except Exception as exc: + self.save_error = exc + raise + + def observe_workflow(self, agent: WorkflowAgent) -> None: + """Observe both stages the graph runner suppresses before and during checkpoint writes.""" + runner = agent.workflow._runner # pyright: ignore[reportPrivateUsage] + runner._prepare_checkpoint_state = self._observe( # pyright: ignore[reportPrivateUsage] + runner._prepare_checkpoint_state # pyright: ignore[reportPrivateUsage] + ) + runner.context.create_checkpoint = self._observe(runner.context.create_checkpoint) + + def _observe( + self, operation: Callable[ParametersT, CoroutineType[Any, Any, ResultT]] + ) -> Callable[ParametersT, CoroutineType[Any, Any, ResultT]]: + @wraps(operation) + async def observed(*args: ParametersT.args, **kwargs: ParametersT.kwargs) -> ResultT: + try: + return await operation(*args, **kwargs) + except Exception as exc: + self.save_error = exc + raise + + return observed class StoreProvider(ABC, Generic[StoreT]): @@ -339,4 +387,26 @@ def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentReques return FoundryAgentSessionStore(platform_context) +class _InvocationsCheckpointStore(FoundryCheckpointStore): + DEFAULT_ROOT_SCOPE = "invocations_checkpoints" + + +class _InvocationsAgentSessionStore(FoundryAgentSessionStore): + DEFAULT_ROOT_SCOPE = "invocations_agent_sessions" + + +class _InvocationsCheckpointStoreProvider(CheckpointStoreProvider): # pyright: ignore[reportUnusedClass] + def get_store( + self, *, config: AgentConfig, context_id: str, platform_context: FoundryAgentRequestContext + ) -> CheckpointStorage: + return _InvocationsCheckpointStore( + context_id, platform_context, allowed_checkpoint_types=self._allowed_checkpoint_types + ) + + +class _InvocationsAgentSessionStoreProvider(AgentSessionStoreProvider): # pyright: ignore[reportUnusedClass] + def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentRequestContext) -> SessionStore: + return _InvocationsAgentSessionStore(platform_context) + + # endregion Agent session persistence diff --git a/python/packages/foundry_hosting/tests/test_invocations_factory.py b/python/packages/foundry_hosting/tests/test_invocations_factory.py new file mode 100644 index 00000000000..05337b44f09 --- /dev/null +++ b/python/packages/foundry_hosting/tests/test_invocations_factory.py @@ -0,0 +1,754 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Request factory lifetime and workflow persistence for the text-only host.""" + +import asyncio +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock + +import pytest +from agent_framework import ( + Agent, + AgentExecutor, + AgentResponse, + AgentResponseUpdate, + AgentSession, + BaseChatClient, + ChatResponse, + ChatResponseUpdate, + Content, + Executor, + FunctionalWorkflowAgent, + InMemoryCheckpointStorage, + Message, + ResponseStream, + RunContext, + SessionStore, + WorkflowAgent, + WorkflowBuilder, + WorkflowCheckpointException, + WorkflowContext, + WorkflowEvent, + handler, + response_handler, + step, + workflow, +) +from anyio import CancelScope, create_task_group +from azure.ai.agentserver.core import ( + FoundryAgentRequestContext, + get_request_context, + reset_request_context, + set_request_context, +) +from starlette.requests import Request +from starlette.responses import StreamingResponse +from typing_extensions import AsyncGenerator, Never, Self + +from agent_framework_foundry_hosting import InvocationsHostServer + + +@contextmanager +def _context(user: str | None = "user", session: str = "session") -> Iterator[None]: + token = set_request_context(FoundryAgentRequestContext(user_id=user, session_id=session)) + try: + yield + finally: + reset_request_context(token) + + +def _request(message: str = "hello", *, stream: bool = False) -> Request: + request = MagicMock(spec=Request) + request.json = AsyncMock(return_value={"message": message, "stream": stream}) + return request + + +async def _invoke( + server: InvocationsHostServer, + message: str = "hello", + *, + stream: bool = False, + user: str | None = "user", + session: str = "session", +) -> str: + with _context(user, session): + response = await server._handle_invoke(_request(message, stream=stream)) + if isinstance(response, StreamingResponse): + chunks = [chunk async for chunk in response.body_iterator] + return "".join(chunk if isinstance(chunk, str) else bytes(chunk).decode() for chunk in chunks) + return bytes(response.body).decode() + + +class _OwnedAgent: + def __init__(self, events: list[str], *, fail: bool = False, wait: asyncio.Event | None = None) -> None: + self.id = "ordinary" + self.name: str | None = "ordinary" + self.description: str | None = None + self.events = events + self.fail = fail + self.wait = wait + self.owner: asyncio.Task[Any] | None = None + self.calls: list[Any] = [] + self.session: AgentSession | None = None + + async def __aenter__(self) -> "_OwnedAgent": + self.owner = asyncio.current_task() + self.events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + assert asyncio.current_task() is self.owner + await asyncio.sleep(0) + self.events.append("exit") + + def create_session(self, *, session_id: str | None = None) -> AgentSession: + return AgentSession(session_id=session_id) + + def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: + return AgentSession(session_id=session_id, service_session_id=service_session_id) + + def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any: + self.calls.append(messages) + self.session = session + + async def updates() -> AsyncGenerator[AgentResponseUpdate]: + try: + self.events.append("run") + yield AgentResponseUpdate(contents=[Content.from_text("a")]) + if self.wait is not None: + await self.wait.wait() + if self.fail: + raise RuntimeError("run failed") + yield AgentResponseUpdate(contents=[Content.from_text("")]) + yield AgentResponseUpdate(contents=[Content.from_text("b")]) + finally: + await asyncio.sleep(0) + self.events.append("iterator closed") + + async def response() -> AgentResponse: + self.events.append("run") + if self.wait is not None: + await self.wait.wait() + if self.fail: + raise RuntimeError("run failed") + return AgentResponse(messages=[Message("assistant", ["ab"])]) + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) if stream else response() + + +class _Stores: + def __init__(self) -> None: + self.checkpoints: dict[str, InMemoryCheckpointStorage] = {} + self.sessions = SessionStore() + self.checkpoint_provider = MagicMock() + self.checkpoint_provider.get_store.side_effect = self._checkpoint_store + self.session_provider = MagicMock() + self.session_provider.get_store.return_value = self.sessions + + def _checkpoint_store(self, *, context_id: str, **kwargs: Any) -> InMemoryCheckpointStorage: + assert get_request_context().session_id is not None + assert context_id.startswith("~invocations-") + assert "/" not in context_id and "\\" not in context_id + return self.checkpoints.setdefault(context_id, InMemoryCheckpointStorage()) + + def server(self, factory: Callable[..., Any]) -> InvocationsHostServer: + return InvocationsHostServer( + agent_factory=factory, + checkpoint_store_provider=self.checkpoint_provider, + agent_session_store_provider=self.session_provider, + ) + + +class _Counter(Executor): + def __init__(self) -> None: + super().__init__(id="counter") + self.count = 0 + + @handler + async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + self.count += 1 + # Include the outer provider history, not just executor checkpoint state. + await ctx.yield_output(f"{self.count}:{'|'.join(message.text for message in messages)}") + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"count": self.count} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.count = state["count"] + + +def _graph(name: str = "counter") -> WorkflowAgent: + return WorkflowBuilder(name=name, start_executor=_Counter()).build().as_agent() + + +@workflow(name="functional") +async def _functional(messages: Any) -> str: + return str(messages) + + +@workflow(name="functional-pending") +async def _functional_pending(messages: Any, ctx: RunContext) -> str: + await ctx.request_info("approve?", response_type=str) + return str(messages) + + +class _Pending(Executor): + @handler + async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + await ctx.request_info("approve?", response_type=str) + + @response_handler + async def answer(self, original_request: str, response: str, ctx: WorkflowContext[Never, str]) -> None: + await ctx.yield_output(response) + + +def _pending_graph() -> WorkflowAgent: + return WorkflowBuilder(name="pending", start_executor=_Pending(id="pending")).build().as_agent() + + +@pytest.mark.parametrize("factory", [_graph, lambda: _functional.build().as_agent()]) +def test_direct_workflow_instances_and_subclasses_require_factory(factory: Callable[..., Any]) -> None: + agent = factory() + with pytest.raises(TypeError, match="agent_factory"): + InvocationsHostServer(agent) + subclass: Any = type("CustomWorkflowAgent", (type(agent),), {}) + underlying = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow + with pytest.raises(TypeError, match="agent_factory"): + InvocationsHostServer(subclass(underlying)) + + +def test_factory_source_validation_and_no_startup_call() -> None: + factory = MagicMock() + with pytest.raises(ValueError, match="exactly one"): + InvocationsHostServer() + with pytest.raises(ValueError, match="exactly one"): + InvocationsHostServer(_OwnedAgent([]), agent_factory=factory) + with pytest.raises(TypeError, match="callable"): + InvocationsHostServer(agent_factory=cast(Any, 42)) + InvocationsHostServer(agent_factory=factory) + factory.assert_not_called() + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("awaitable", [False, True]) +async def test_factory_once_per_request_inside_context_preserves_text(stream: bool, awaitable: bool) -> None: + events: list[str] = [] + agents: list[_OwnedAgent] = [] + + def factory() -> Any: + assert get_request_context().user_id == "user" + agent = _OwnedAgent(events) + agents.append(agent) + + async def create() -> _OwnedAgent: + return agent + + return create() if awaitable else agent + + server = InvocationsHostServer(agent_factory=factory) + assert await _invoke(server, stream=stream) == "ab" + assert await _invoke(server, "second", stream=stream) == "ab" + assert len(agents) == 2 + assert agents[0].session is agents[1].session + assert agents[0].calls == ["hello" if stream else ["hello"]] + assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) * 2 + assert server._agent is None + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_async_factory_failure_and_invalid_result_release_lock(stream: bool) -> None: + calls = 0 + + async def factory() -> Any: + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("construction failed") + return object() + + server = InvocationsHostServer(agent_factory=factory) + with pytest.raises(RuntimeError, match="construction failed"): + await _invoke(server, stream=stream) + with pytest.raises(TypeError, match="agent_factory must return"): + await _invoke(server, stream=stream) + assert calls == 2 + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_run_failure_closes_owner_and_iterator(stream: bool) -> None: + events: list[str] = [] + server = InvocationsHostServer(agent_factory=lambda: _OwnedAgent(events, fail=True)) + with pytest.raises(RuntimeError, match="run failed"): + await _invoke(server, stream=stream) + assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) + assert not server._scope_locks._entries + + +async def test_stream_factory_and_resources_belong_to_asgi_consumer_task() -> None: + events: list[str] = [] + creation_task: asyncio.Task[Any] | None = None + + def factory() -> _OwnedAgent: + nonlocal creation_task + creation_task = asyncio.current_task() + return _OwnedAgent(events) + + server = InvocationsHostServer(agent_factory=factory) + with _context(): + response = await server._handle_invoke(_request(stream=True)) + assert creation_task is None + assert isinstance(response, StreamingResponse) + body: list[bytes] = [] + + async def send(message: Any) -> None: + if message["type"] == "http.response.body": + body.append(message["body"]) + + consumer = asyncio.create_task(response.stream_response(send)) + await consumer + assert creation_task is consumer + assert b"".join(body) == b"ab" + assert events == ["enter", "run", "iterator closed", "exit"] + + +@pytest.mark.parametrize("disconnect", [False, True]) +async def test_stream_disconnect_or_cancellation_closes_suspended_iterator(disconnect: bool) -> None: + events: list[str] = [] + sent = asyncio.Event() + server = InvocationsHostServer(agent_factory=lambda: _OwnedAgent(events)) + with _context(): + response = await server._handle_invoke(_request(stream=True)) + assert isinstance(response, StreamingResponse) + + async def send(message: Any) -> None: + if message["type"] == "http.response.body": + sent.set() + if disconnect: + raise OSError("disconnected") + await asyncio.Event().wait() + + consumer = asyncio.create_task(response.stream_response(send)) + await sent.wait() + if not disconnect: + consumer.cancel() + with pytest.raises(OSError if disconnect else asyncio.CancelledError): + await consumer + assert events == ["enter", "run", "iterator closed", "exit"] + assert not server._scope_locks._entries + + +async def test_cancelled_factory_construction_releases_scope() -> None: + started = asyncio.Event() + + async def factory() -> _OwnedAgent: + started.set() + await asyncio.Event().wait() + return _OwnedAgent([]) + + server = InvocationsHostServer(agent_factory=factory) + task = asyncio.create_task(_invoke(server)) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not server._scope_locks._entries + + +async def test_same_scope_serializes_stream_through_cleanup_and_cancelled_waiters() -> None: + events: list[str] = [] + release = asyncio.Event() + first_chunk = asyncio.Event() + created = 0 + + def factory() -> _OwnedAgent: + nonlocal created + created += 1 + return _OwnedAgent(events) + + server = InvocationsHostServer(agent_factory=factory) + with _context(): + response = await server._handle_invoke(_request(stream=True)) + assert isinstance(response, StreamingResponse) + + async def send(message: Any) -> None: + if message["type"] == "http.response.body" and message.get("body"): + first_chunk.set() + await release.wait() + + consumer = asyncio.create_task(response.stream_response(send)) + await first_chunk.wait() + cancelled = asyncio.create_task(_invoke(server)) + waiting = asyncio.create_task(_invoke(server)) + await asyncio.sleep(0) + assert created == 1 + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + assert len(server._scope_locks._entries) == 1 + assert await _invoke(server, session="independent") == "ab" + assert created == 2 + release.set() + await consumer + assert await waiting == "ab" + assert created == 3 + assert events == ["enter", "run", "enter", "run", "exit", "iterator closed", "exit", "enter", "run", "exit"] + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_graph_checkpoint_and_outer_history_survive_new_host(stream: bool) -> None: + stores = _Stores() + assert await _invoke(stores.server(_graph), "first", stream=stream) == "1:first" + second = await _invoke(stores.server(_graph), "second", stream=stream) + assert second == "2:first|1:first|second" + assert len(stores.checkpoints) == 1 + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_graph_scopes_isolate_users_sessions_and_unsafe_identifiers(stream: bool) -> None: + stores = _Stores() + server = stores.server(_graph) + scopes = [ + ("user", "../session"), + ("another-user", "../session"), + ("user", r"..\session"), + ("a:b", "c"), + ("b", "c:a"), + ("USER", "../session"), + ] + for user, session in scopes: + assert await _invoke(server, user, stream=stream, user=user, session=session) == f"1:{user}" + assert len(stores.checkpoints) == len(scopes) + for user, session in scopes: + assert (await _invoke(server, "next", stream=stream, user=user, session=session)).startswith("2:") + assert not server._sessions + + +@pytest.mark.parametrize("same_wrapper", [False, True]) +@pytest.mark.parametrize("functional", [False, True]) +async def test_reused_workflow_or_wrapper_is_rejected(same_wrapper: bool, functional: bool) -> None: + stores = _Stores() + agent = _functional.build().as_agent() if functional else _graph() + underlying = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow + server = stores.server(lambda: agent if same_wrapper else underlying.as_agent()) + await _invoke(server) + with pytest.raises(RuntimeError, match="reused a workflow"): + await _invoke(server, session="other") + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("damage", ["checkpoint", "session", "name", "kind"]) +async def test_missing_or_incompatible_continuation_never_restarts(damage: str) -> None: + stores = _Stores() + await _invoke(stores.server(_graph)) + storage_id, storage = next(iter(stores.checkpoints.items())) + if damage == "checkpoint": + for checkpoint in await storage.list_checkpoints(workflow_name="counter"): + await storage.delete(checkpoint.checkpoint_id) + elif damage == "session": + await stores.sessions.delete(storage_id) + + def factory() -> WorkflowAgent | FunctionalWorkflowAgent: + if damage == "kind": + return _functional.build().as_agent() + return _graph("different" if damage == "name" else "counter") + + with pytest.raises(RuntimeError, match="missing|does not match"): + await _invoke(stores.server(factory), "next") + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_graph_pending_request_cannot_be_answered_with_plain_text(stream: bool) -> None: + stores = _Stores() + server = stores.server(_pending_graph) + with pytest.raises(RuntimeError, match="plain text cannot answer"): + await _invoke(server, stream=stream) + with pytest.raises(RuntimeError, match="plain text cannot answer"): + await _invoke(stores.server(_pending_graph), "approved", stream=stream) + with pytest.raises(RuntimeError, match="plain text cannot answer"): + await _invoke(server, "approved", stream=stream, user="other") + assert len(stores.checkpoints) == 2 + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_functional_clean_completion_starts_new_message(stream: bool) -> None: + stores = _Stores() + + def factory() -> FunctionalWorkflowAgent: + return _functional.build().as_agent() + + assert await _invoke(stores.server(factory), "first", stream=stream) == ("first" if stream else "['first']") + storage = next(iter(stores.checkpoints.values())) + checkpoint = await storage.get_latest(workflow_name="functional") + assert checkpoint is not None + # A completed functional checkpoint can retain old pending events. Only the + # host's completion record distinguishes it from interrupted work. + checkpoint.pending_request_info_events["old"] = WorkflowEvent.request_info( + request_id="old", source_executor_id="functional", request_data="old", response_type=str + ) + await storage.save(checkpoint) + assert await _invoke(stores.server(factory), "next", stream=stream) == ("next" if stream else "['next']") + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_functional_pending_continuation_is_explicitly_unsupported(stream: bool) -> None: + stores = _Stores() + + def factory() -> FunctionalWorkflowAgent: + return _functional_pending.build().as_agent() + + with pytest.raises(RuntimeError, match="plain text cannot answer"): + await _invoke(stores.server(factory), stream=stream) + with pytest.raises(RuntimeError, match="pending or interrupted functional"): + await _invoke(stores.server(factory), "approved", stream=stream) + + +async def test_functional_interrupted_continuation_does_not_restart() -> None: + stores = _Stores() + calls = 0 + + @workflow(name="interrupted") + async def interrupted(messages: Any) -> str: + nonlocal calls + calls += 1 + raise RuntimeError("interrupted") + + def factory() -> FunctionalWorkflowAgent: + return interrupted.build().as_agent() + + with pytest.raises(RuntimeError, match="interrupted"): + await _invoke(stores.server(factory)) + with pytest.raises(RuntimeError, match="missing its required checkpoint"): + await _invoke(stores.server(factory)) + assert calls == 1 + + +async def test_session_save_failure_closes_workflow_owner(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + + class OwnedWorkflow(WorkflowAgent): + async def __aenter__(self) -> Self: + events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + events.append("exit") + + stores = _Stores() + monkeypatch.setattr(stores.sessions, "set", AsyncMock(side_effect=[None, RuntimeError("save failed")])) + server = stores.server(lambda: OwnedWorkflow(_graph().workflow)) + with pytest.raises(RuntimeError, match="save failed"): + await _invoke(server) + assert events == ["enter", "exit"] + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("cancel", [False, True]) +async def test_anyio_resource_scopes_exit_in_consumer_task_and_correct_order(cancel: bool) -> None: + events: list[str] = [] + sent = asyncio.Event() + + class TaskGroupAgent(_OwnedAgent): + async def __aenter__(self) -> Self: + await super().__aenter__() + self.group = create_task_group() + await self.group.__aenter__() + return self + + async def __aexit__(self, *args: Any) -> None: + await self.group.__aexit__(*args) + await super().__aexit__(*args) + + server = InvocationsHostServer(agent_factory=lambda: TaskGroupAgent(events)) + with _context(): + response = await server._handle_invoke(_request(stream=True)) + assert isinstance(response, StreamingResponse) + scope = CancelScope() + + async def send(message: Any) -> None: + if message["type"] == "http.response.body" and message.get("body"): + sent.set() + if cancel: + await asyncio.Event().wait() + + async def consume() -> None: + with scope: + await response.stream_response(send) + + consumer = asyncio.create_task(consume()) + await sent.wait() + if cancel: + scope.cancel() + await consumer + assert events == ["enter", "run", "iterator closed", "exit"] + assert not server._scope_locks._entries + + +class _TranscriptClient(BaseChatClient): + def __init__(self, transcripts: list[list[str]]) -> None: + super().__init__() + self.transcripts = transcripts + + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Any: + self.transcripts.append([message.text for message in messages]) + + async def updates() -> AsyncGenerator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("recorded")]) + + async def response() -> ChatResponse: + return ChatResponse(messages=[Message("assistant", ["recorded"])]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_real_agent_executor_transcript_is_isolated_and_restored(stream: bool) -> None: + stores = _Stores() + transcripts: list[list[str]] = [] + + def factory() -> WorkflowAgent: + agent = Agent(client=_TranscriptClient(transcripts), name="inner") + executor = AgentExecutor(agent, id="inner") + return WorkflowBuilder(name="transcript", start_executor=executor).build().as_agent() + + assert await _invoke(stores.server(factory), "private-first", stream=stream) == "recorded" + assert await _invoke(stores.server(factory), "unrelated", stream=stream, user="other") == "recorded" + assert transcripts[-1] == ["unrelated"] + assert await _invoke(stores.server(factory), "private-next", stream=stream) == "recorded" + assert "private-first" in transcripts[-1] + assert "private-next" in transcripts[-1] + assert "unrelated" not in transcripts[-1] + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_functional_interruption_with_checkpoint_rejects_new_message(stream: bool) -> None: + stores = _Stores() + calls = 0 + + @step + async def saved_step() -> str: + return "saved" + + @workflow(name="interrupted-after-step") + async def interrupted(messages: Any) -> str: + nonlocal calls + calls += 1 + await saved_step() + raise RuntimeError("failed after checkpoint") + + def factory() -> FunctionalWorkflowAgent: + return interrupted.build().as_agent() + + with pytest.raises(RuntimeError, match="failed after checkpoint"): + await _invoke(stores.server(factory), stream=stream) + with pytest.raises(RuntimeError, match="pending or interrupted functional"): + await _invoke(stores.server(factory), "next", stream=stream) + assert calls == 1 + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_cancellation_during_run_closes_resources(stream: bool) -> None: + events: list[str] = [] + server = InvocationsHostServer(agent_factory=lambda: _OwnedAgent(events, wait=asyncio.Event())) + task = asyncio.create_task(_invoke(server, stream=stream)) + await asyncio.sleep(0) + assert events == ["enter", "run"] + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) + assert not server._scope_locks._entries + + +async def test_invalid_request_does_not_construct_agent(monkeypatch: pytest.MonkeyPatch) -> None: + factory = MagicMock() + server = InvocationsHostServer(agent_factory=factory) + with _context(): + request = _request() + monkeypatch.setattr(request, "json", AsyncMock(return_value={})) + response = await server._handle_invoke(request) + assert response.status_code == 400 + with _context(): + server.config.is_hosted = True + with _context(user=None): + response = await server._handle_invoke(_request()) + assert response.status_code == 500 + factory.assert_not_called() + + +async def test_checkpoint_failure_closes_workflow_owner_and_preserves_attempt(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + + class OwnedWorkflow(WorkflowAgent): + async def __aenter__(self) -> Self: + events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + events.append("exit") + + stores = _Stores() + storage = InMemoryCheckpointStorage() + monkeypatch.setattr(storage, "save", AsyncMock(side_effect=RuntimeError("checkpoint failed"))) + stores.checkpoint_provider.get_store.side_effect = None + stores.checkpoint_provider.get_store.return_value = storage + server = stores.server(lambda: OwnedWorkflow(_graph().workflow)) + with pytest.raises(RuntimeError, match="checkpoint failed"): + await _invoke(server) + with pytest.raises(RuntimeError, match="missing its required checkpoint"): + await _invoke(server) + assert events == ["enter", "exit", "enter", "exit"] + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_checkpoint_preparation_failure_rejects_older_checkpoint_and_fresh_host_retry(stream: bool) -> None: + stores = _Stores() + calls: list[str] = [] + snapshots: list[int] = [] + agents: list[WorkflowAgent] = [] + + class FailingCheckpoint(_Counter): + @handler + async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + calls.extend(message.text for message in messages) + await super().count_message(messages, ctx) + + async def on_checkpoint_save(self) -> dict[str, Any]: + snapshots.append(self.count) + if self.count: + raise RuntimeError("executor checkpoint preparation failed") + return await super().on_checkpoint_save() + + def factory() -> WorkflowAgent: + agent = WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() + agents.append(agent) + return agent + + first = stores.server(factory) + with pytest.raises(WorkflowCheckpointException, match="Executor counter on_checkpoint_save failed"): + await _invoke(first, stream=stream) + storage_id, storage = next(iter(stores.checkpoints.items())) + checkpoint = await storage.get_latest(workflow_name="checkpoint-preparation") + assert checkpoint is not None + assert checkpoint.iteration_count == 0 + assert 0 in snapshots and 1 in snapshots + session = await stores.sessions.get(storage_id) + assert session is not None + assert session.state["_foundry_invocations_workflow"]["checkpoint_failed"] is True + + second = stores.server(factory) + with pytest.raises(RuntimeError, match="incomplete checkpoint persistence"): + await _invoke(second, "must-not-run", stream=stream) + assert calls == ["hello"] + assert len(agents) == 2 + assert agents[0] is not agents[1] + assert agents[0].workflow is not agents[1].workflow + assert not first._scope_locks._entries + assert not second._scope_locks._entries diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 97b076450bd..ca8aecf8eda 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -25,6 +25,7 @@ import pytest from agent_framework import ( Agent, + AgentExecutor, AgentExecutorRequest, AgentResponse, AgentResponseUpdate, @@ -36,6 +37,7 @@ Content, FunctionInvocationLayer, HistoryProvider, + InMemoryCheckpointStorage, InMemoryHistoryProvider, Message, RawAgent, @@ -360,11 +362,26 @@ async def set(self, session_id: str, session: AgentSession) -> None: _SESSION_STORE_UNSET = object() -def _make_server(agent: Any, **kwargs: Any) -> ResponsesHostServer: +def _make_server(agent: Any = None, **kwargs: Any) -> ResponsesHostServer: """Create a ResponsesHostServer, optionally replacing its private store for tests.""" session_store = kwargs.pop("session_store", _SESSION_STORE_UNSET) response_store = kwargs.pop("response_store", InMemoryResponseProvider()) - server = ResponsesHostServer(agent, store=response_store, **kwargs) + is_workflow = isinstance(agent, WorkflowAgent) or "agent_factory" in kwargs + if isinstance(agent, WorkflowAgent): + server = ResponsesHostServer(agent_factory=lambda: agent, store=response_store, **kwargs) + else: + server = ResponsesHostServer(agent, store=response_store, **kwargs) + if is_workflow: + if session_store is _SESSION_STORE_UNSET: + session_store = SessionStore() + checkpoints: dict[str, InMemoryCheckpointStorage] = {} + + def get_checkpoint_store(*, context_id: str, **kwargs: Any) -> InMemoryCheckpointStorage: + return checkpoints.setdefault(context_id, InMemoryCheckpointStorage()) + + checkpoint_provider = MagicMock(spec=CheckpointStoreProvider) + checkpoint_provider.get_store.side_effect = get_checkpoint_store + server._checkpoint_storage_provider = checkpoint_provider # pyright: ignore[reportPrivateUsage] if session_store is not _SESSION_STORE_UNSET: provider = MagicMock(spec=AgentSessionStoreProvider) provider.get_store.return_value = cast(SessionStore | None, session_store) @@ -970,9 +987,9 @@ def test_init_rejects_resilient_background_for_non_workflow_agent(self, tmp_path options=ResponsesServerOptions(resilient_background=True), ) - def test_init_rejects_steerable_conversations_for_workflow_agent(self) -> None: + def test_init_rejects_direct_workflow_agent(self) -> None: workflow_agent = _build_text_workflow_agent("hello from workflow") - with pytest.raises(RuntimeError, match="steerable_conversations"): + with pytest.raises(TypeError, match="agent_factory"): ResponsesHostServer( cast(SupportsAgentRun, workflow_agent), store=InMemoryResponseProvider(), @@ -4722,7 +4739,7 @@ async def test_workflow_rejects_invalid_checkpoint_scope( agent.workflow = MagicMock() agent.workflow.name = "workflow" agent.workflow._runner_context.has_checkpointing.return_value = False - server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) + server = ResponsesHostServer(agent_factory=lambda: agent, store=InMemoryResponseProvider()) context_kwargs: dict[str, Any] = {"response_id": "response-current", "mode_flags": MagicMock()} request = CreateResponse(model="m", input="hi") @@ -5570,7 +5587,8 @@ async def _aiter() -> AsyncIterator[AgentResponseUpdate]: 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() + inner_executor = AgentExecutor(inner, id="text-agent") + workflow = WorkflowBuilder(name="text-workflow", start_executor=start).add_edge(start, inner_executor).build() return WorkflowAgent(workflow=workflow, name="Text Workflow Agent") @@ -5653,7 +5671,10 @@ 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() + inner_executor = AgentExecutor(inner, id="multi-update-agent") + workflow = ( + WorkflowBuilder(name="multi-update-workflow", start_executor=start).add_edge(start, inner_executor).build() + ) return WorkflowAgent(workflow=workflow, name="Multi Update Workflow Agent"), inner @@ -5677,11 +5698,40 @@ def _build_approval_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, mock_agent).build() + inner_executor = AgentExecutor(mock_agent, id="approval-agent") + workflow = WorkflowBuilder(name="approval-workflow", start_executor=start).add_edge(start, inner_executor).build() workflow_agent = WorkflowAgent(workflow=workflow, name="Approval Workflow Agent") return workflow_agent, mock_agent +def _build_approval_workflow_factory( + *, approval_request_id: str, final_text: str +) -> tuple[Callable[[], WorkflowAgent], list[_ToolApprovalWorkflowAgentMock]]: + agents: list[_ToolApprovalWorkflowAgentMock] = [] + + def factory() -> WorkflowAgent: + workflow_agent, inner = _build_approval_workflow_agent( + approval_request_id=approval_request_id, final_text=final_text + ) + agents.append(inner) + return workflow_agent + + return factory, agents + + +def _build_multi_update_workflow_factory( + texts: Sequence[str], +) -> tuple[Callable[[], WorkflowAgent], list[_MultiUpdateWorkflowAgentMock]]: + agents: list[_MultiUpdateWorkflowAgentMock] = [] + + def factory() -> WorkflowAgent: + workflow_agent, inner = _build_multi_update_workflow_agent(texts) + agents.append(inner) + return workflow_agent + + return factory, agents + + class TestWorkflowAgentHosting: """End-to-end HTTP tests for ``ResponsesHostServer`` hosting a ``WorkflowAgent``. @@ -5770,23 +5820,18 @@ async def test_cancellation_signal_preempts_stuck_workflow_call(self) -> None: AsyncGenerator[Any, None], server._handle_response(request, context, cancellation_signal), # pyright: ignore[reportPrivateUsage] ) - await anext(handler) # response.created - await anext(handler) # response.in_progress - # Pull the first workflow event in the background so we can wait for the inner agent's - # run() to actually start (proving it's genuinely stuck on `gate`) before signalling -- - # otherwise cancellation could preempt the pull before the workflow even reaches it. - pending = asyncio.ensure_future(anext(handler)) + async def consume() -> list[Any]: + return [event async for event in handler] + + # Keep factory entry, iteration, and cleanup in the same task. + pending = asyncio.create_task(consume()) await asyncio.wait_for(inner.started.wait(), timeout=1.0) cancellation_signal.set() # Fires while the inner agent is stuck awaiting `gate`. - async def _drain() -> list[Any]: - first = await pending - return [first, *[event async for event in handler]] - # Bounded well below `gate` never being set: proves cancellation preempted the stuck # call instead of only being observed after it (eventually) produced an update. - events = await asyncio.wait_for(_drain(), timeout=1.0) + events = await asyncio.wait_for(pending, timeout=1.0) types = [event.get("type") for event in events if isinstance(event, Mapping)] assert "response.output_text.delta" not in types @@ -5817,12 +5862,12 @@ async def test_shutdown_signal_preempts_stuck_workflow_call(self, tmp_path: Path AsyncGenerator[Any, None], server._handle_response(request, context, cancellation_signal), # pyright: ignore[reportPrivateUsage] ) - await anext(handler) # response.created - await anext(handler) # response.in_progress - # Pull the first workflow event in the background so we can wait for the inner agent's - # run() to actually start (proving it's genuinely stuck on `gate`) before signalling. - pending = asyncio.ensure_future(anext(handler)) + async def consume() -> list[Any]: + return [event async for event in handler] + + # Keep factory entry, iteration, and cleanup in the same task. + pending = asyncio.create_task(consume()) await asyncio.wait_for(inner.started.wait(), timeout=1.0) context.shutdown.set() # Fires while the inner agent is stuck awaiting `gate`. @@ -5837,12 +5882,12 @@ async def test_cancellation_signal_set_before_turn_skips_new_input(self) -> None """Explicit-cancel: cancellation set before a continuation turn starts must skip that turn's new input entirely, whether caught by the restore-loop's own check or the standalone check guarding the start of a brand new workflow run.""" - workflow_agent, inner = _build_multi_update_workflow_agent(["hello"]) - server = _make_server(workflow_agent) + agent_factory, agents = _build_multi_update_workflow_factory(["hello"]) + server = _make_server(agent_factory=agent_factory) first = await _post(server, conversation_id="conv-1", stream=False) assert first.status_code == 200 - run_count_after_first_turn = inner.run_count + run_count_after_first_turn = sum(agent.run_count for agent in agents) assert run_count_after_first_turn == 1 request = CreateResponse(model="m", input="hi again", stream=True) @@ -5865,7 +5910,8 @@ async def test_cancellation_signal_set_before_turn_skips_new_input(self) -> None assert types[-1] == "response.completed" # At most the restore-only replay call happened; the new-turn call (which would deliver # "hi again") must never fire. - assert inner.run_count <= run_count_after_first_turn + 1 + assert len(agents) == 2 + assert sum(agent.run_count for agent in agents) <= run_count_after_first_turn + 1 async def test_shutdown_signal_set_before_restore_only_triggers_recovery(self, tmp_path: Path) -> None: """Shutdown observed while resuming a checkpoint (whether during the restore-only replay or @@ -5873,16 +5919,16 @@ async def test_shutdown_signal_set_before_restore_only_triggers_recovery(self, t ``exit_for_recovery()`` -- proving the post-loop ``signalled`` check (not a blind re-check of the flag) correctly gates this action so it doesn't also fire on a replay that merely finished naturally.""" - workflow_agent, inner = _build_multi_update_workflow_agent(["hello"]) + agent_factory, agents = _build_multi_update_workflow_factory(["hello"]) server = _make_server( - workflow_agent, + agent_factory=agent_factory, response_store=FileResponseStore(storage_dir=tmp_path), options=ResponsesServerOptions(resilient_background=True), ) first = await _post(server, conversation_id="conv-1", stream=False) assert first.status_code == 200 - run_count_after_first_turn = inner.run_count + run_count_after_first_turn = sum(agent.run_count for agent in agents) assert run_count_after_first_turn == 1 request = CreateResponse(model="m", input="hi again", stream=True) @@ -5903,7 +5949,8 @@ async def test_shutdown_signal_set_before_restore_only_triggers_recovery(self, t _ = [event async for event in handler] # Only the restore-only replay call may have happened; the new-turn call must never fire. - assert inner.run_count <= run_count_after_first_turn + 1 + assert len(agents) == 2 + assert sum(agent.run_count for agent in agents) <= run_count_after_first_turn + 1 async def test_previous_response_requires_existing_workflow_checkpoint(self) -> None: """A previous_response_id naming a scope with no checkpoint must fail loudly rather than @@ -5995,11 +6042,11 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None approval response back to the paused inner agent, and the inner agent emits the final assistant text. """ - workflow_agent, mock_agent = _build_approval_workflow_agent( + agent_factory, agents = _build_approval_workflow_factory( approval_request_id="apr_wf_rt", final_text="done with approval", ) - server = _make_server(workflow_agent) + server = _make_server(agent_factory=agent_factory) checkpoint_provider = server._checkpoint_storage_provider # pyright: ignore[reportPrivateUsage] with patch.object(checkpoint_provider, "get_store", wraps=checkpoint_provider.get_store) as get_store: @@ -6010,7 +6057,7 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None approval_items = [it for it in first_body["output"] if it["type"] == "mcp_approval_request"] assert len(approval_items) == 1 approval_request_id = approval_items[0]["id"] - assert mock_agent.run_count == 1 + assert sum(agent.run_count for agent in agents) == 1 second_payload: dict[str, Any] = { "model": "test-model", @@ -6037,7 +6084,8 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None # The inner agent must have been resumed (restore replay + new turn). # Restore call is a no-op for the mock (no input); the new-turn call # delivers the approval response, so run_count grows by at least 1. - assert mock_agent.run_count >= 2 + assert len(agents) == 2 + assert sum(agent.run_count for agent in agents) >= 2 # The final assistant text from the resumed inner agent surfaces in # the HTTP output. @@ -6055,7 +6103,7 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None # The new-turn invocation of the inner agent must have received the # approval response routed back through WorkflowAgent. approval_responses = [ - c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response" + c for m in agents[-1].last_run_messages for c in m.contents if c.type == "function_approval_response" ] assert len(approval_responses) == 1 assert approval_responses[0].approved is True @@ -6063,11 +6111,11 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None async def test_round_trip_approval_response_streaming(self) -> None: """Streaming variant of the round-trip: turn 2 is requested with ``stream=true`` and surfaces the resumed text as SSE events.""" - workflow_agent, mock_agent = _build_approval_workflow_agent( + agent_factory, agents = _build_approval_workflow_factory( approval_request_id="apr_wf_rt_st", final_text="streamed-done", ) - server = _make_server(workflow_agent) + server = _make_server(agent_factory=agent_factory) first = await _post(server, stream=False) first_body = first.json() @@ -6097,16 +6145,17 @@ async def test_round_trip_approval_response_streaming(self) -> None: text_done = [e for e in events if e["event"] == "response.output_text.done"] assert any("streamed-done" in e["data"]["text"] for e in text_done) - assert mock_agent.run_count >= 2 + assert len(agents) == 2 + assert sum(agent.run_count for agent in agents) >= 2 async def test_round_trip_approval_response_rejected(self) -> None: """Sending ``approve=False`` must surface as ``approved=False`` to the inner agent on resume.""" - workflow_agent, mock_agent = _build_approval_workflow_agent( + agent_factory, agents = _build_approval_workflow_factory( approval_request_id="apr_wf_reject", final_text="acknowledged", ) - server = _make_server(workflow_agent) + server = _make_server(agent_factory=agent_factory) first = await _post(server, stream=False) first_body = first.json() @@ -6131,7 +6180,7 @@ async def test_round_trip_approval_response_rejected(self) -> None: assert second.status_code == 200 approval_responses = [ - c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response" + c for m in agents[-1].last_run_messages for c in m.contents if c.type == "function_approval_response" ] assert len(approval_responses) == 1 assert approval_responses[0].approved is False @@ -6193,7 +6242,6 @@ async def _passthrough(items: Any, *, approval_storage: Any = None) -> list[Any] async def test_reads_overlap_and_preserve_order(self, monkeypatch: pytest.MonkeyPatch) -> None: self._identity_converters(monkeypatch) server = _make_server(_make_agent()) - server._uses_agent_server_history = True # pyright: ignore[reportPrivateUsage] history_msg = Message(role="assistant", contents=[Content.from_text("H")]) input_msg = Message(role="user", contents=[Content.from_text("I")]) @@ -6231,8 +6279,7 @@ async def get_history() -> list[Message]: async def test_history_read_skipped_without_agent_server_history(self, monkeypatch: pytest.MonkeyPatch) -> None: self._identity_converters(monkeypatch) - server = _make_server(_make_agent()) - server._uses_agent_server_history = False # pyright: ignore[reportPrivateUsage] + server = _make_server(_make_agent(), history_source="agent") input_msg = Message(role="user", contents=[Content.from_text("I")]) @@ -6251,7 +6298,6 @@ async def get_input_items() -> list[Message]: async def test_failed_read_cancels_and_drains_sibling(self, monkeypatch: pytest.MonkeyPatch) -> None: self._identity_converters(monkeypatch) server = _make_server(_make_agent()) - server._uses_agent_server_history = True # pyright: ignore[reportPrivateUsage] sibling_cancelled = asyncio.Event() @@ -6279,7 +6325,8 @@ async def get_history() -> list[Message]: # The still-blocked history read must have been cancelled, not left orphaned. await asyncio.wait_for(sibling_cancelled.wait(), timeout=1) - async def test_session_preparation_failure_cancels_pending_reads(self) -> None: + @pytest.mark.parametrize("use_factory", [False, True]) + async def test_session_preparation_failure_cancels_pending_reads(self, use_factory: bool) -> None: """If session preparation fails, the concurrently-launched read must be cancelled and drained by `_handle_inner_agent`, not left running as an orphan after the request fails.""" input_started = asyncio.Event() @@ -6293,7 +6340,14 @@ async def get(self, session_id: str) -> AgentSession | None: await input_started.wait() raise RuntimeError("session prep boom") - server = _make_server(_make_agent(), session_store=_GetFailsOnceReadStarted()) + server = ( + _make_server( + agent_factory=lambda: Agent(client=_RecordingHistoryClient()), + session_store=_GetFailsOnceReadStarted(), + ) + if use_factory + else _make_server(_make_agent(), session_store=_GetFailsOnceReadStarted()) + ) request = CreateResponse(model="m", input="hi", stream=True) # A previous_response_id makes session_load_id non-None so the failing get() is reached. request["previous_response_id"] = "resp-x" diff --git a/python/packages/foundry_hosting/tests/test_responses_factory.py b/python/packages/foundry_hosting/tests/test_responses_factory.py new file mode 100644 index 00000000000..03d57073542 --- /dev/null +++ b/python/packages/foundry_hosting/tests/test_responses_factory.py @@ -0,0 +1,1401 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Request-owned Responses agents, resources, and durable workflow continuation.""" + +import asyncio +import copy +import gc +import weakref +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Mapping, Sequence +from contextlib import aclosing, contextmanager +from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock + +import pytest +from agent_framework import ( + Agent, + AgentExecutor, + AgentResponse, + AgentResponseUpdate, + AgentSession, + BaseChatClient, + ChatResponse, + ChatResponseUpdate, + Content, + Executor, + FunctionalWorkflowAgent, + InMemoryCheckpointStorage, + InMemoryHistoryProvider, + Message, + ResponseStream, + RunContext, + SessionStore, + WorkflowAgent, + WorkflowBuilder, + WorkflowContext, + handler, + response_handler, + step, + workflow, +) +from anyio import CancelScope, create_task_group +from azure.ai.agentserver.core import ( + FoundryAgentRequestContext, + get_request_context, + reset_request_context, + set_request_context, +) +from azure.ai.agentserver.responses import ResponseContext, ResponsesServerOptions +from azure.ai.agentserver.responses.models import CreateResponse +from azure.ai.agentserver.responses.streaming._checkpoint import ResponseCheckpointEvent +from typing_extensions import Never, Self + +from agent_framework_foundry_hosting import ResponsesHostServer +from agent_framework_foundry_hosting._agent_factory import close_run_iterator + + +@contextmanager +def _platform(user: str = "alice") -> Iterator[None]: + token = set_request_context(FoundryAgentRequestContext(user_id=user, session_id="platform-session")) + try: + yield + finally: + reset_request_context(token) + + +def _context( + text: str = "hello", + *, + response: str = "response-1", + conversation: str | None = "conversation", + items: list[Any] | None = None, + history: list[Any] | None = None, +) -> ResponseContext: + context = ResponseContext(response_id=response, conversation_id=conversation, mode_flags=MagicMock()) + context.get_input_items = AsyncMock( # type: ignore[method-assign] + return_value=items + if items is not None + else [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}] + ) + context.get_history = AsyncMock(return_value=history or []) # type: ignore[method-assign] + return context + + +def _request(previous: str | None = None) -> CreateResponse: + request = CreateResponse(model="test-model", input="input resolved by context", stream=True) + if previous is not None: + request["previous_response_id"] = previous + return request + + +async def _collect( + server: ResponsesHostServer, + context: ResponseContext | None = None, + *, + user: str = "alice", + previous: str | None = None, +) -> list[Any]: + with _platform(user): + return [ + event async for event in server._handle_response(_request(previous), context or _context(), asyncio.Event()) + ] + + +def _types(events: list[Any]) -> list[str]: + return [event["type"] for event in events if isinstance(event, Mapping)] + + +def _text(events: list[Any]) -> str: + return "".join( + event["delta"] + for event in events + if isinstance(event, Mapping) and event.get("type") == "response.output_text.delta" + ) + + +def _failure(events: list[Any]) -> str: + assert _types(events)[-1] == "response.failed", events + assert _types(events).count("response.failed") == 1 + assert "response.completed" not in _types(events) + return events[-1]["response"]["error"]["message"] + + +class _ApprovalStore: + def __init__(self) -> None: + self.requests: dict[str, Content] = {} + + async def save_approval_request(self, request_id: str, content: Content) -> None: + self.requests[request_id] = content + + async def load_approval_request(self, request_id: str) -> Content | None: + return self.requests.get(request_id) + + +class _Stores: + def __init__(self) -> None: + self.sessions: dict[str | None, SessionStore] = {} + self.checkpoints: dict[tuple[str | None, str], InMemoryCheckpointStorage] = {} + self.approvals: dict[str | None, _ApprovalStore] = {} + self.session_provider = MagicMock() + self.session_provider.get_store.side_effect = self._sessions + self.checkpoint_provider = MagicMock() + self.checkpoint_provider.get_store.side_effect = self._checkpoints + self.approval_provider = MagicMock() + self.approval_provider.get_store.side_effect = self._approvals + + def _sessions(self, *, platform_context: FoundryAgentRequestContext, **kwargs: Any) -> SessionStore: + assert get_request_context().user_id == platform_context.user_id + return self.sessions.setdefault(platform_context.user_id, SessionStore()) + + def _checkpoints( + self, *, platform_context: FoundryAgentRequestContext, context_id: str, **kwargs: Any + ) -> InMemoryCheckpointStorage: + assert get_request_context().user_id == platform_context.user_id + return self.checkpoints.setdefault((platform_context.user_id, context_id), InMemoryCheckpointStorage()) + + def _approvals(self, *, platform_context: FoundryAgentRequestContext, **kwargs: Any) -> _ApprovalStore: + assert get_request_context().user_id == platform_context.user_id + return self.approvals.setdefault(platform_context.user_id, _ApprovalStore()) + + def server(self, factory: Callable[..., Any], **kwargs: Any) -> ResponsesHostServer: + return ResponsesHostServer( + agent_factory=factory, + history_source="agent", + checkpoint_store_provider=self.checkpoint_provider, + agent_session_store_provider=self.session_provider, + function_approval_store_provider=self.approval_provider, + **kwargs, + ) + + +class _OwnedAgent: + id = "ordinary" + name: str | None = "ordinary" + description: str | None = "Request resource test agent" + + def __init__(self, events: list[str], *, wait: asyncio.Event | None = None, fail: bool = False) -> None: + self.events = events + self.wait = wait + self.fail = fail + self.owner: asyncio.Task[Any] | None = None + self.session: AgentSession | None = None + + async def __aenter__(self) -> "_OwnedAgent": + self.owner = asyncio.current_task() + self.events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + assert asyncio.current_task() is self.owner + await asyncio.sleep(0) + self.events.append("exit") + + def create_session(self, *, session_id: str | None = None) -> AgentSession: + return AgentSession(session_id=session_id) + + def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: + return AgentSession(session_id=session_id, service_session_id=service_session_id) + + def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any: + assert stream + self.session = session + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + try: + self.events.append("run") + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text("first")]) + if self.wait is not None: + await self.wait.wait() + if self.fail: + raise RuntimeError("model failed") + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text("second")]) + finally: + await asyncio.sleep(0) + self.events.append("iterator closed") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + +class _Counter(Executor): + def __init__(self) -> None: + super().__init__(id="counter") + self.count = 0 + + @handler + async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + self.count += 1 + await ctx.yield_output(f"{self.count}:{'|'.join(message.text for message in messages)}") + + async def on_checkpoint_save(self) -> dict[str, Any]: + return {"count": self.count} + + async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: + self.count = state["count"] + + +def _graph(name: str = "counter", *, history: bool = False) -> WorkflowAgent: + providers = [InMemoryHistoryProvider(source_id="outer-history")] if history else [] + return WorkflowBuilder(name=name, start_executor=_Counter()).build().as_agent(context_providers=providers) + + +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("completed", [False, True]) +@pytest.mark.parametrize("stage", ["preset", "session_load", "input", "checkpoint_lookup", "attempt_save"]) +async def test_cancelled_workflow_preparation_preserves_session_and_next_request( + functional: bool, completed: bool, stage: str, monkeypatch: pytest.MonkeyPatch +) -> None: + await _cancel_workflow_preparation(functional, completed, stage, False, monkeypatch) + + +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("completed", [False, True]) +@pytest.mark.parametrize("stage", ["session_load", "input", "checkpoint_lookup", "attempt_save"]) +async def test_task_cancelled_workflow_preparation_preserves_session_and_next_request( + functional: bool, completed: bool, stage: str, monkeypatch: pytest.MonkeyPatch +) -> None: + await _cancel_workflow_preparation(functional, completed, stage, True, monkeypatch) + + +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("cancel_task", [False, True]) +async def test_cancelled_workflow_attempt_save_removes_unstarted_response_branch( + functional: bool, cancel_task: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + await _cancel_workflow_preparation(functional, True, "attempt_save", cancel_task, monkeypatch, chain=True) + + +async def _cancel_workflow_preparation( + functional: bool, + completed: bool, + stage: str, + cancel_task: bool, + monkeypatch: pytest.MonkeyPatch, + *, + chain: bool = False, +) -> None: + stores = _Stores() + factory = (lambda: _functional.build().as_agent()) if functional else _graph + stores.sessions["alice"] = SessionStore() + sessions = stores.sessions["alice"] + storage = InMemoryCheckpointStorage() + conversation = None if chain else "conversation" + saved_id = "initial" if chain else "conversation" + previous_id = "initial" if chain else None + stores.checkpoints[("alice", saved_id)] = storage + if completed: + initial = await _collect( + stores.server(factory), _context("before", response="initial", conversation=conversation) + ) + assert _types(initial)[-1] == "response.completed" + previous = await sessions.get(saved_id) + previous_snapshot = previous.to_dict() if previous is not None else None + checkpoint_ids = await storage.list_checkpoint_ids(workflow_name="functional" if functional else "counter") + context = _context("cancelled", response="cancelled", conversation=conversation) + cancellation_signal = asyncio.Event() + cancelled = False + consumer: asyncio.Task[list[Any]] | None = None + + async def consume() -> list[Any]: + with _platform(): + return [ + event + async for event in stores.server(factory)._handle_response( + _request(previous_id), context, cancellation_signal + ) + ] + + with monkeypatch.context() as patch: + if stage == "preset": + cancellation_signal.set() + else: + target, method = { + "session_load": (sessions, "get"), + "input": (context, "get_input_items"), + "checkpoint_lookup": (storage, "get_latest"), + "attempt_save": (sessions, "set"), + }[stage] + original = getattr(target, method) + + async def cancel_after_preparation(*args: Any, **kwargs: Any) -> Any: + nonlocal cancelled + result = await original(*args, **kwargs) + if not cancelled: + cancelled = True + if cancel_task: + assert consumer is not None + consumer.cancel() + else: + cancellation_signal.set() + await asyncio.sleep(0) + return result + + patch.setattr(target, method, cancel_after_preparation) + consumer = asyncio.create_task(consume()) + if cancel_task: + with pytest.raises(asyncio.CancelledError): + await consumer + else: + events = await consumer + assert not _text(events) + assert _types(events)[-1] == "response.completed" + current = await sessions.get(saved_id) + assert (current.to_dict() if current is not None else None) == previous_snapshot + assert await storage.list_checkpoint_ids(workflow_name="functional" if functional else "counter") == checkpoint_ids + if chain: + assert await sessions.get("cancelled") is None + + following = await _collect( + stores.server(factory), _context("next", response="following", conversation=conversation), previous=previous_id + ) + assert _types(following)[-1] == "response.completed", following + assert _text(following) == ("next" if functional else f"{2 if completed else 1}:next") + + +@pytest.mark.parametrize("functional", [False, True]) +async def test_cancelled_workflow_execution_without_checkpoint_preserves_incomplete_attempt( + functional: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + stores = _Stores() + storage = InMemoryCheckpointStorage() + stores.checkpoints[("alice", "conversation")] = storage + started = asyncio.Event() + release = asyncio.Event() + cancellation_signal = asyncio.Event() + calls: list[str] = [] + + @workflow(name="interrupted") + async def interrupted(messages: list[Message]) -> str: + calls.append(messages[0].text) + started.set() + await release.wait() + return messages[0].text + + factory = (lambda: interrupted.build().as_agent()) if functional else _graph + + async def consume() -> list[Any]: + with _platform(): + return [ + event + async for event in stores.server(factory)._handle_response( + _request(), _context("original"), cancellation_signal + ) + ] + + with monkeypatch.context() as patch: + if not functional: + original_save = storage.save + + async def blocked_checkpoint(checkpoint: Any) -> str: + started.set() + await release.wait() + return await original_save(checkpoint) + + patch.setattr(storage, "save", blocked_checkpoint) + consumer = asyncio.create_task(consume()) + await asyncio.wait_for(started.wait(), 2) + cancellation_signal.set() + await asyncio.wait_for(consumer, 2) + + session = await stores.sessions["alice"].get("conversation") + assert session is not None + assert session.state["_foundry_responses_workflow"]["completed"] is False + assert not await storage.list_checkpoint_ids(workflow_name="interrupted" if functional else "counter") + following = await _collect(stores.server(factory), _context("next", response="next")) + assert "missing its required workflow checkpoint" in _failure(following) + if functional: + assert calls == ["original"] + else: + context = _context("original") + context.is_recovery = True + recovered = await _collect( + stores.server(factory, options=ResponsesServerOptions(resilient_background=True)), context + ) + assert _types(recovered)[-1] == "response.completed" + assert _text(recovered) == "1:original" + + +@workflow(name="functional") +async def _functional(messages: list[Message]) -> str: + return "|".join(message.text for message in messages) + + +@workflow(name="functional-pending-string") +async def _pending_string(messages: list[Message], ctx: RunContext) -> str: + answer = await ctx.request_info("answer?", response_type=str) + return f"{messages[0].text}:{answer}" + + +@workflow(name="functional-pending-bool") +async def _pending_bool(messages: list[Message], ctx: RunContext) -> str: + answer = await ctx.request_info("approve?", response_type=bool) + return f"{messages[0].text}:{answer}" + + +class _Pending(Executor): + @handler + async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + await ctx.request_info(messages[0].text, response_type=str) + + @response_handler + async def answer(self, original_request: str, response: str, ctx: WorkflowContext[Never, str]) -> None: + await ctx.yield_output(f"{original_request}:{response}") + + +def _pending_graph() -> WorkflowAgent: + return WorkflowBuilder(name="pending", start_executor=_Pending(id="pending")).build().as_agent() + + +def test_constructor_validates_exactly_one_callable_without_constructing() -> None: + factory = MagicMock() + with pytest.raises(ValueError, match="exactly one"): + ResponsesHostServer() + with pytest.raises(ValueError, match="exactly one"): + ResponsesHostServer(_OwnedAgent([]), agent_factory=factory) + with pytest.raises(TypeError, match="callable"): + ResponsesHostServer(agent_factory=cast(Any, 42)) + server = _Stores().server(factory) + factory.assert_not_called() + assert server._agent is None + + +def test_constructor_rejects_coroutine_object_instead_of_factory() -> None: + async def factory() -> _OwnedAgent: + return _OwnedAgent([]) + + coroutine = factory() + try: + with pytest.raises(TypeError, match="callable"): + ResponsesHostServer(agent_factory=cast(Any, coroutine)) + finally: + coroutine.close() + + +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("subclass", [False, True]) +def test_direct_workflow_instances_require_factory(functional: bool, subclass: bool) -> None: + agent = _functional.build().as_agent() if functional else _graph() + if subclass: + underlying = agent._workflow if isinstance(agent, FunctionalWorkflowAgent) else agent.workflow + agent = type("CustomWorkflow", (type(agent),), {})(underlying) + with pytest.raises(TypeError, match="agent_factory"): + ResponsesHostServer(agent) + + +@pytest.mark.parametrize("kind", ["sync", "async", "awaitable-object"]) +async def test_factory_runs_once_per_request_not_startup_and_never_sets_instance_agent(kind: str) -> None: + events: list[str] = [] + agents: list[_OwnedAgent] = [] + + def create() -> _OwnedAgent: + assert get_request_context().user_id == "alice" + agent = _OwnedAgent(events) + agents.append(agent) + return agent + + async def create_async() -> _OwnedAgent: + return create() + + class AwaitableFactory: + def __call__(self) -> Any: + return create_async() + + factories: dict[str, Callable[..., Any]] = { + "sync": create, + "async": create_async, + "awaitable-object": AwaitableFactory(), + } + server = _Stores().server(factories[kind]) + assert not agents + for response in ("one", "two"): + result = await _collect(server, _context(response=response)) + assert _text(result) == "firstsecond" + assert _types(result)[-1] == "response.completed" + assert server._agent is None + assert server._agent_stack is None + await server._cleanup_agent() + assert len(agents) == 2 + assert events == ["enter", "run", "iterator closed", "exit"] * 2 + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("kind", ["invalid", "exception", "async-exception"]) +async def test_factory_failure_is_not_retried_and_releases_lock(kind: str) -> None: + calls = 0 + + def factory() -> Any: + nonlocal calls + calls += 1 + if kind == "invalid": + return object() + if kind == "exception": + raise RuntimeError("factory failed") + + async def failed() -> Any: + raise RuntimeError("factory failed") + + return failed() + + server = _Stores().server(factory) + message = _failure(await _collect(server)) + assert ("SupportsAgentRun" if kind == "invalid" else "factory failed") in message + assert calls == 1 + assert server._agent is None + assert not server._scope_locks._entries + + +async def test_cancelled_factory_is_not_retried_and_releases_lock() -> None: + entered = asyncio.Event() + calls = 0 + + async def factory() -> Any: + nonlocal calls + calls += 1 + entered.set() + await asyncio.Event().wait() + + server = _Stores().server(factory) + consumer = asyncio.create_task(_collect(server)) + await asyncio.wait_for(entered.wait(), 2) + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + assert calls == 1 + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("same_wrapper", [False, True]) +async def test_reusing_workflow_wrapper_or_underlying_workflow_fails(functional: bool, same_wrapper: bool) -> None: + agent = _functional.build().as_agent() if functional else _graph() + underlying = agent._workflow if isinstance(agent, FunctionalWorkflowAgent) else agent.workflow + server = _Stores().server(lambda: agent if same_wrapper else underlying.as_agent()) + assert _types(await _collect(server))[-1] == "response.completed" + assert "reused" in _failure(await _collect(server, _context(response="two"))) + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("functional", [False, True]) +async def test_completed_workflows_are_not_retained_by_factory_resolver(functional: bool) -> None: + refs: list[weakref.ReferenceType[Any]] = [] + + def factory() -> Any: + agent = _functional.build().as_agent() if functional else _graph() + underlying = agent._workflow if isinstance(agent, FunctionalWorkflowAgent) else agent.workflow + refs.extend([weakref.ref(agent), weakref.ref(underlying)]) + return agent + + server = _Stores().server(factory) + assert _types(await _collect(server))[-1] == "response.completed" + await asyncio.sleep(0) + gc.collect() + assert all(reference() is None for reference in refs) + assert server._agent_resolver is not None + assert not server._agent_resolver._seen + + +@pytest.mark.parametrize("finish", ["complete", "close", "cancel-signal", "model-error"]) +async def test_resources_stay_open_through_output_and_close_exactly_once(finish: str) -> None: + events: list[str] = [] + server = _Stores().server( + lambda: _OwnedAgent( + events, wait=asyncio.Event() if finish in ("close", "cancel-signal") else None, fail=finish == "model-error" + ) + ) + cancellation = asyncio.Event() + emitted: list[Any] = [] + with _platform(): + stream = cast(AsyncGenerator[Any, None], server._handle_response(_request(), _context(), cancellation)) + async with aclosing(stream): + async for event in stream: + emitted.append(event) + kind = event.get("type") if isinstance(event, Mapping) else None + if kind == "response.output_text.delta": + assert "enter" in events and "exit" not in events + if finish == "close": + break + if finish == "cancel-signal": + cancellation.set() + if kind in ("response.completed", "response.failed"): + assert events[-1] == "exit" + assert events == ["enter", "run", "iterator closed", "exit"] + assert not server._scope_locks._entries + if finish == "model-error": + assert "model failed" in _failure(emitted) + + +@pytest.mark.parametrize("model_failure", [False, True]) +async def test_session_persistence_failure_closes_resources_and_reports_both_errors( + model_failure: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + events: list[str] = [] + stores = _Stores() + sessions = SessionStore() + monkeypatch.setattr(sessions, "set", AsyncMock(side_effect=RuntimeError("session save failed"))) + stores.sessions["alice"] = sessions + server = stores.server(lambda: _OwnedAgent(events, fail=model_failure)) + message = _failure(await _collect(server)) + assert "session save failed" in message + if model_failure: + assert "model failed" in message + assert events == ["enter", "run", "iterator closed", "exit"] + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("cancel", [False, True]) +async def test_nested_task_affine_resources_close_in_consumer_task_under_anyio_cancellation( + cancel: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + events: list[str] = [] + sent = asyncio.Event() + scope = CancelScope() + stores = _Stores() + sessions = SessionStore() + save = sessions.set + + async def persist(session_id: str, session: AgentSession) -> None: + await asyncio.sleep(0) + await save(session_id, session) + events.append("saved") + + monkeypatch.setattr(sessions, "set", AsyncMock(side_effect=persist)) + stores.sessions["alice"] = sessions + + class NestedAgent(_OwnedAgent): + async def __aenter__(self) -> Self: + await super().__aenter__() + self.group = create_task_group() + await self.group.__aenter__() + return self + + async def __aexit__(self, *args: Any) -> None: + await self.group.__aexit__(*args) + await super().__aexit__(*args) + + server = stores.server(lambda: NestedAgent(events)) + + async def consume() -> None: + with _platform(), scope: + stream = cast(AsyncGenerator[Any, None], server._handle_response(_request(), _context(), asyncio.Event())) + async with aclosing(stream): + async for event in stream: + if isinstance(event, Mapping) and event.get("type") == "response.output_text.delta": + sent.set() + if cancel: + await asyncio.Event().wait() + + consumer = asyncio.create_task(consume()) + await asyncio.wait_for(sent.wait(), 2) + if cancel: + scope.cancel() + await asyncio.wait_for(consumer, 2) + assert events == ["enter", "run", "iterator closed", "saved", "exit"] + assert not server._scope_locks._entries + + +async def test_scope_lock_covers_last_output_cleanup_and_removes_cancelled_waiters() -> None: + events: list[str] = [] + exiting = asyncio.Event() + release = asyncio.Event() + constructed = 0 + + class SlowExit(_OwnedAgent): + async def __aexit__(self, *args: Any) -> None: + exiting.set() + await release.wait() + await super().__aexit__(*args) + + def factory() -> _OwnedAgent: + nonlocal constructed + constructed += 1 + return SlowExit(events) if constructed == 1 else _OwnedAgent(events) + + server = _Stores().server(factory) + first = asyncio.create_task(_collect(server)) + await asyncio.wait_for(exiting.wait(), 2) + second = asyncio.create_task(_collect(server, _context(response="two"))) + cancelled = asyncio.create_task(_collect(server, _context(response="three"))) + await asyncio.sleep(0) + assert constructed == 1 + assert not first.done() + assert next(iter(server._scope_locks._entries.values())).users == 3 + cancelled.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled + assert next(iter(server._scope_locks._entries.values())).users == 2 + release.set() + results = await asyncio.wait_for(asyncio.gather(first, second), 2) + assert all(_types(result)[-1] == "response.completed" for result in results) + assert constructed == 2 + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("other_user,other_scope", [("bob", "conversation"), ("alice", "other")]) +async def test_distinct_users_or_conversations_can_overlap(other_user: str, other_scope: str) -> None: + started = asyncio.Event() + release = asyncio.Event() + calls = 0 + + def factory() -> _OwnedAgent: + nonlocal calls + calls += 1 + if calls == 2: + started.set() + return _OwnedAgent([], wait=release) + + server = _Stores().server(factory) + first = asyncio.create_task(_collect(server)) + second = asyncio.create_task(_collect(server, _context(conversation=other_scope), user=other_user)) + try: + await asyncio.wait_for(started.wait(), 2) + finally: + release.set() + results = await asyncio.gather(first, second) + assert all(_types(result)[-1] == "response.completed" for result in results) + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("chain", [False, True]) +@pytest.mark.parametrize("history", [False, True]) +async def test_graph_state_and_explicit_outer_history_survive_new_hosts(chain: bool, history: bool) -> None: + stores = _Stores() + conversation = None if chain else "conversation" + first = await _collect(stores.server(lambda: _graph(history=history)), _context("first", conversation=conversation)) + assert _text(first) == "1:first" + second = await _collect( + stores.server(lambda: _graph(history=history)), + _context("second", response="response-2", conversation=conversation), + previous="response-1" if chain else None, + ) + assert _types(second)[-1] == "response.completed" + assert _text(second) == ("2:first|1:first|second" if history else "2:second") + assert ("alice", "response-2" if chain else "conversation") in stores.checkpoints + + +async def test_graph_stable_name_must_match_saved_marker_from_previous_host() -> None: + stores = _Stores() + assert _text(await _collect(stores.server(lambda: _graph("stable")))) == "1:hello" + result = await _collect(stores.server(lambda: _graph("new-random-name")), _context("next", response="two")) + assert "name or kind" in _failure(result) + assert not _text(result) + + +async def test_graph_with_outer_history_requires_saved_session_alongside_checkpoint() -> None: + stores = _Stores() + assert _text(await _collect(stores.server(lambda: _graph(history=True)))) == "1:hello" + stores.sessions.clear() + result = await _collect(stores.server(lambda: _graph(history=True)), _context("next", response="two")) + assert "missing its required outer agent session" in _failure(result) + assert not _text(result) + + +class _TranscriptClient(BaseChatClient): + def __init__(self, transcripts: list[list[str]]) -> None: + super().__init__() + self.transcripts = transcripts + + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Any: + assert stream + self.transcripts.append([message.text for message in messages]) + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("recorded")]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + +@pytest.mark.parametrize("chain", [False, True]) +async def test_agent_executor_transcript_isolates_users_scopes_and_continues_across_hosts(chain: bool) -> None: + stores = _Stores() + transcripts: list[list[str]] = [] + + def factory() -> WorkflowAgent: + agent = Agent(client=_TranscriptClient(transcripts), name="inner") + inner = AgentExecutor(agent, id="inner") + return WorkflowBuilder(name="transcript", start_executor=inner).build().as_agent() + + conversation = None if chain else "conversation" + assert _text(await _collect(stores.server(factory), _context("private", conversation=conversation))) == "recorded" + await _collect(stores.server(factory), _context("bob-only", conversation=conversation), user="bob") + assert transcripts[-1] == ["bob-only"] + await _collect(stores.server(factory), _context("other-scope", conversation="other")) + assert transcripts[-1] == ["other-scope"] + continued = await _collect( + stores.server(factory), + _context("next", response="response-2", conversation=conversation), + previous="response-1" if chain else None, + ) + assert _text(continued) == "recorded" + assert transcripts[-1] == ["private", "recorded", "next"] + + +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("damage", ["name", "checkpoint", "marker", "history"]) +async def test_missing_or_incompatible_workflow_state_does_not_restart(functional: bool, damage: str) -> None: + stores = _Stores() + factory = (lambda: _functional.build().as_agent()) if functional else _graph + assert _types(await _collect(stores.server(factory)))[-1] == "response.completed" + session = await stores.sessions["alice"].get("conversation") + assert session is not None + if damage == "name": + session.state["_foundry_responses_workflow"]["name"] = "different-name" + elif damage == "marker": + session.state.pop("_foundry_responses_workflow") + else: + stores.checkpoints.clear() + await stores.sessions["alice"].set("conversation", session) + context = _context("next", response="two", history=[{"type": "message"}] if damage == "history" else None) + message = _failure(await _collect(stores.server(factory), context)) + assert ("name or kind" if damage in ("name", "marker") else "missing its required workflow checkpoint") in message + + +async def test_functional_resilient_background_is_explicitly_rejected_without_running() -> None: + calls: list[str] = [] + + @workflow(name="not-recoverable") + async def functional(messages: Any) -> str: + calls.append("run") + return "unexpected" + + server = _Stores().server( + lambda: functional.build().as_agent(), options=ResponsesServerOptions(resilient_background=True) + ) + assert "buffered step output" in _failure(await _collect(server)) + assert not calls + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("chain", [False, True]) +async def test_completed_functional_workflow_starts_fresh_input_not_cached_previous_input(chain: bool) -> None: + stores = _Stores() + calls: list[str] = [] + + @step + async def record(text: str) -> str: + calls.append(text) + return text + + @workflow(name="fresh-functional") + async def functional(messages: list[Message]) -> str: + return await record(messages[0].text) + + conversation = None if chain else "conversation" + assert ( + _text( + await _collect( + stores.server(lambda: functional.build().as_agent()), _context("first", conversation=conversation) + ) + ) + == "first" + ) + assert ( + _text( + await _collect( + stores.server(lambda: functional.build().as_agent()), + _context("second", response="two", conversation=conversation), + previous="response-1" if chain else None, + ) + ) + == "second" + ) + assert calls == ["first", "second"] + + +@pytest.mark.parametrize("kind", ["graph", "functional-string", "functional-bool"]) +@pytest.mark.parametrize("chain", [False, True]) +async def test_authorized_pending_response_resumes_matching_checkpoint(kind: str, chain: bool) -> None: + stores = _Stores() + factory = { + "graph": _pending_graph, + "functional-string": lambda: _pending_string.build().as_agent(), + "functional-bool": lambda: _pending_bool.build().as_agent(), + }[kind] + conversation = None if chain else "conversation" + first = await _collect(stores.server(factory), _context("private", conversation=conversation)) + assert _types(first)[-1] == "response.completed" + checkpoint = await stores.checkpoints[("alice", conversation or "response-1")].get_latest( + workflow_name={ + "graph": "pending", + "functional-string": "functional-pending-string", + "functional-bool": "functional-pending-bool", + }[kind] + ) + assert checkpoint is not None + request_id = next(iter(checkpoint.pending_request_info_events)) + approval_id = None + if kind == "functional-bool": + approval_id = next(iter(stores.approvals["alice"].requests)) + assert stores.approvals["alice"].requests[approval_id].id == request_id + item = ( + {"type": "mcp_approval_response", "approval_request_id": approval_id, "approve": True} + if kind == "functional-bool" + else {"type": "function_call_output", "call_id": request_id, "output": "accepted"} + ) + resumed = await _collect( + stores.server(factory), + _context(response="two", conversation=conversation, items=[item]), + previous="response-1" if chain else None, + ) + assert _types(resumed)[-1] == "response.completed", resumed + assert _text(resumed) == ("private:True" if kind == "functional-bool" else "private:accepted") + + +@pytest.mark.parametrize("kind", ["text", "wrong-id", "approval-for-string", "cross-user"]) +async def test_functional_pending_response_requires_authorized_matching_type_and_user(kind: str) -> None: + stores = _Stores() + + def factory() -> FunctionalWorkflowAgent: + return _pending_string.build().as_agent() + + assert _types(await _collect(stores.server(factory)))[-1] == "response.completed" + checkpoint = await stores.checkpoints[("alice", "conversation")].get_latest( + workflow_name="functional-pending-string" + ) + assert checkpoint is not None + request_id = next(iter(checkpoint.pending_request_info_events)) + if kind == "approval-for-string": + approval_id = next(iter(stores.approvals["alice"].requests)) + items = [{"type": "mcp_approval_response", "approval_request_id": approval_id, "approve": True}] + elif kind == "text": + items = None + else: + items = [ + { + "type": "function_call_output", + "call_id": "wrong" if kind == "wrong-id" else request_id, + "output": "stolen", + } + ] + rejected = await _collect( + stores.server(factory), + _context("plain text", response="two", items=items), + user="bob" if kind == "cross-user" else "alice", + ) + if kind != "cross-user": + assert "pending functional workflow request" in _failure(rejected) + assert "hello:stolen" not in _text(rejected) + approved = await _collect( + stores.server(factory), + _context(response="three", items=[{"type": "function_call_output", "call_id": request_id, "output": "owner"}]), + ) + assert _text(approved) == "hello:owner" + + +async def test_graph_recovery_without_checkpoint_replays_original_input() -> None: + stores = _Stores() + context = _context("original") + context.is_recovery = True + server = stores.server(_graph, options=ResponsesServerOptions(resilient_background=True)) + events = await _collect(server, context) + assert _types(events)[-1] == "response.completed" + assert _text(events) == "1:original" + + +@pytest.mark.parametrize("functional", [False, True]) +async def test_workflow_final_session_save_failure_closes_request_resources( + functional: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + events: list[str] = [] + stores = _Stores() + sessions = SessionStore() + save = AsyncMock(side_effect=[None, RuntimeError("workflow session save failed")]) + monkeypatch.setattr(sessions, "set", save) + stores.sessions["alice"] = sessions + + class OwnedGraph(WorkflowAgent): + async def __aenter__(self) -> Self: + events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + await asyncio.sleep(0) + events.append("exit") + + class OwnedFunctional(FunctionalWorkflowAgent): + async def __aenter__(self) -> Self: + events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + await asyncio.sleep(0) + events.append("exit") + + server = stores.server( + lambda: OwnedFunctional(_functional.build()) if functional else OwnedGraph(_graph().workflow) + ) + assert "workflow session save failed" in _failure(await _collect(server)) + assert events == ["enter", "exit"] + assert save.await_count == 2 + assert not server._scope_locks._entries + + +async def test_resource_exit_failure_replaces_success_with_failure() -> None: + events: list[str] = [] + + class FailingExit(_OwnedAgent): + async def __aexit__(self, *args: Any) -> None: + await super().__aexit__(*args) + raise RuntimeError("resource exit failed") + + server = _Stores().server(lambda: FailingExit(events)) + result = await _collect(server) + assert _text(result) == "firstsecond" + assert "resource exit failed" in _failure(result) + assert events == ["enter", "run", "iterator closed", "exit"] + assert not server._scope_locks._entries + + +async def test_cancelling_streaming_task_persists_session_and_closes_iterator(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + emitted = asyncio.Event() + stores = _Stores() + sessions = SessionStore() + save = AsyncMock(wraps=sessions.set) + monkeypatch.setattr(sessions, "set", save) + stores.sessions["alice"] = sessions + server = stores.server(lambda: _OwnedAgent(events, wait=asyncio.Event())) + + async def consume() -> None: + with _platform(): + stream = cast(AsyncGenerator[Any, None], server._handle_response(_request(), _context(), asyncio.Event())) + async with aclosing(stream): + async for event in stream: + if isinstance(event, Mapping) and event.get("type") == "response.output_text.delta": + emitted.set() + + consumer = asyncio.create_task(consume()) + await asyncio.wait_for(emitted.wait(), 2) + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await consumer + assert events == ["enter", "run", "iterator closed", "exit"] + save.assert_awaited_once() + assert await sessions.get("conversation") is not None + assert not server._scope_locks._entries + + +async def test_anyio_cancellation_finishes_blocked_iterator_cleanup_before_owner_exit() -> None: + events: list[str] = [] + blocked = asyncio.Event() + scope = CancelScope() + + class TaskGroupAgent(_OwnedAgent): + async def __aenter__(self) -> Self: + await super().__aenter__() + self.group = create_task_group() + await self.group.__aenter__() + return self + + async def __aexit__(self, *args: Any) -> None: + await self.group.__aexit__(*args) + await super().__aexit__(*args) + + def run(self, messages: Any = None, *, stream: bool = False, **kwargs: Any) -> Any: + assert stream + + async def updates() -> AsyncIterator[AgentResponseUpdate]: + try: + events.append("run") + yield AgentResponseUpdate(role="assistant", contents=[Content.from_text("first")]) + blocked.set() + await asyncio.Event().wait() + finally: + events.append("cleanup started") + await asyncio.sleep(0) + events.append("cleanup finished") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + server = _Stores().server(lambda: TaskGroupAgent(events)) + + async def consume() -> None: + with scope: + await _collect(server) + + consumer = asyncio.create_task(consume()) + await asyncio.wait_for(blocked.wait(), 2) + scope.cancel() + await asyncio.wait_for(consumer, 2) + assert not server._scope_locks._entries + assert events == ["enter", "run", "cleanup started", "cleanup finished", "exit"] + + +async def test_checkpoint_preparation_failure_rejects_older_checkpoint_and_fresh_host_retry() -> None: + stores = _Stores() + calls: list[str] = [] + snapshots: list[int] = [] + agents: list[WorkflowAgent] = [] + + class FailingCheckpoint(_Counter): + @handler + async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + calls.extend(message.text for message in messages) + await super().count_message(messages, ctx) + + async def on_checkpoint_save(self) -> dict[str, Any]: + snapshots.append(self.count) + if self.count: + raise RuntimeError("executor checkpoint preparation failed") + return await super().on_checkpoint_save() + + def factory() -> WorkflowAgent: + agent = WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() + agents.append(agent) + return agent + + first = stores.server(factory) + assert "Executor counter on_checkpoint_save failed" in _failure(await _collect(first)) + storage = stores.checkpoints[("alice", "conversation")] + checkpoint = await storage.get_latest(workflow_name="checkpoint-preparation") + assert checkpoint is not None + assert checkpoint.iteration_count == 0 + assert 0 in snapshots and 1 in snapshots + session = await stores.sessions["alice"].get("conversation") + assert session is not None + assert session.state["_foundry_responses_workflow"]["checkpoint_failed"] is True + + second = stores.server(factory) + assert "incomplete checkpoint persistence" in _failure( + await _collect(second, _context("must-not-run", response="two")) + ) + assert calls == ["hello"] + assert len(agents) == 2 + assert agents[0] is not agents[1] + assert agents[0].workflow is not agents[1].workflow + assert not first._scope_locks._entries + assert not second._scope_locks._entries + + +@pytest.mark.parametrize("functional", [False, True]) +async def test_checkpoint_save_failure_does_not_report_success_and_closes_owner( + functional: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + events: list[str] = [] + storage = InMemoryCheckpointStorage() + save = AsyncMock(side_effect=RuntimeError("checkpoint save failed")) + monkeypatch.setattr(storage, "save", save) + stores = _Stores() + stores.checkpoints[("alice", "conversation")] = storage + + class OwnedGraph(WorkflowAgent): + async def __aenter__(self) -> Self: + events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + events.append("exit") + + class OwnedFunctional(FunctionalWorkflowAgent): + async def __aenter__(self) -> Self: + events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + events.append("exit") + + server = stores.server( + lambda: OwnedFunctional(_functional.build()) if functional else OwnedGraph(_graph().workflow) + ) + result = await _collect(server) + assert events == ["enter", "exit"] + assert save.await_count > 0 + assert not server._scope_locks._entries + assert "checkpoint save failed" in _failure(result) + + +async def test_interrupted_functional_step_does_not_restart_with_new_input() -> None: + stores = _Stores() + calls: list[str] = [] + + @step + async def saved() -> str: + return "saved" + + @workflow(name="interrupted-functional") + async def interrupted(messages: list[Message]) -> str: + calls.append(messages[0].text) + await saved() + raise RuntimeError("failed after step") + + def factory() -> FunctionalWorkflowAgent: + return interrupted.build().as_agent() + + assert "failed after step" in _failure(await _collect(stores.server(factory))) + assert "interrupted functional workflow" in _failure( + await _collect(stores.server(factory), _context("next", response="two")) + ) + assert calls == ["hello"] + + +@pytest.mark.parametrize("functional", [False, True]) +async def test_cross_user_previous_response_cannot_resume_another_users_pending_request(functional: bool) -> None: + stores = _Stores() + factory = (lambda: _pending_string.build().as_agent()) if functional else _pending_graph + first = await _collect(stores.server(factory), _context("alice-private", conversation=None)) + assert _types(first)[-1] == "response.completed" + checkpoint = await stores.checkpoints[("alice", "response-1")].get_latest( + workflow_name="functional-pending-string" if functional else "pending" + ) + assert checkpoint is not None + request_id = next(iter(checkpoint.pending_request_info_events)) + result = {"type": "function_call_output", "call_id": request_id, "output": "approved"} + rejected = await _collect( + stores.server(factory), + _context(response="two", conversation=None, items=[result]), + user="bob", + previous="response-1", + ) + assert "checkpoint" in _failure(rejected) + approved = await _collect( + stores.server(factory), _context(response="three", conversation=None, items=[result]), previous="response-1" + ) + assert _text(approved) == "alice-private:approved" + + +class _RecoveryStart(Executor): + @handler + async def start(self, messages: list[Message], ctx: WorkflowContext[str, str]) -> None: + await ctx.yield_output(f"first:{messages[0].text}") + await ctx.send_message(messages[0].text) + + +class _RecoveryEnd(Executor): + @handler + async def end(self, text: str, ctx: WorkflowContext[Never, str]) -> None: + await ctx.yield_output(f"last:{text}") + + +def _recovery_graph() -> WorkflowAgent: + start = _RecoveryStart(id="start") + end = _RecoveryEnd(id="end") + return WorkflowBuilder(name="recovery", start_executor=start).add_edge(start, end).build().as_agent() + + +@pytest.mark.parametrize("saved_input", [False, True]) +async def test_graph_recovery_preserves_explicit_outer_history_for_continuation(saved_input: bool) -> None: + stores = _Stores() + transcripts: list[list[str]] = [] + agents: list[WorkflowAgent] = [] + original = Message("user", ["original"]) + + class Start(Executor): + @handler + async def start(self, messages: list[Message], ctx: WorkflowContext[str]) -> None: + transcripts.append([message.text for message in messages]) + await ctx.send_message(messages[-1].text) + + def factory() -> WorkflowAgent: + start = Start(id="start") + end = _RecoveryEnd(id="end") + agent = ( + WorkflowBuilder(name="recovery-history", start_executor=start) + .add_edge(start, end) + .build() + .as_agent(context_providers=[InMemoryHistoryProvider(source_id="outer-history")]) + ) + agents.append(agent) + return agent + + interrupted = factory() + storage = InMemoryCheckpointStorage() + iterator = interrupted.workflow.run([original], stream=True, checkpoint_storage=storage).__aiter__() + try: + async for event in iterator: + if event.type == "superstep_completed": + break + finally: + await close_run_iterator(iterator) + checkpoint = await storage.get_latest(workflow_name="recovery-history") + assert checkpoint is not None + assert checkpoint.iteration_count == 1 + assert transcripts == [["original"]] + stores.checkpoints[("alice", "conversation")] = storage + session = AgentSession() + session.state["_foundry_responses_workflow"] = { + "name": "recovery-history", + "kind": "graph", + "completed": False, + } + if saved_input: + await InMemoryHistoryProvider(source_id="outer-history").save_messages( + session.session_id, [original], state=session.state.setdefault("outer-history", {}) + ) + stores.sessions["alice"] = SessionStore() + await stores.sessions["alice"].set("conversation", session) + + context = _context("original") + context.is_recovery = True + recovered = await _collect( + stores.server(factory, options=ResponsesServerOptions(resilient_background=True)), context + ) + assert _types(recovered)[-1] == "response.completed" + assert _text(recovered) == "last:original" + assert transcripts == [["original"]] + recovered_session = await stores.sessions["alice"].get("conversation") + assert recovered_session is not None + history = await InMemoryHistoryProvider(source_id="outer-history").get_messages( + recovered_session.session_id, state=recovered_session.state.get("outer-history") + ) + assert [(message.role, message.text) for message in history] == [ + ("user", "original"), + ("assistant", "last:original"), + ] + + continued = await _collect(stores.server(factory), _context("next", response="two")) + assert _types(continued)[-1] == "response.completed" + assert _text(continued) == "last:next" + assert transcripts == [["original"], ["original", "last:original", "next"]] + continued_session = await stores.sessions["alice"].get("conversation") + assert continued_session is not None + history = await InMemoryHistoryProvider(source_id="outer-history").get_messages( + continued_session.session_id, state=continued_session.state.get("outer-history") + ) + assert [message.text for message in history] == ["original", "last:original", "next", "last:next"] + assert len({id(agent.workflow) for agent in agents}) == 3 + + +@pytest.mark.parametrize("snapshot_kind", ["latest", "empty", "partial"]) +async def test_graph_recovery_selects_latest_or_response_paired_checkpoint( + snapshot_kind: str, monkeypatch: pytest.MonkeyPatch +) -> None: + stores = _Stores() + options = ResponsesServerOptions(resilient_background=True) + server = stores.server(_recovery_graph, options=options) + snapshots: list[Any] = [] + with _platform(): + async for event in server._handle_response(_request(), _context("original"), asyncio.Event()): + if isinstance(event, ResponseCheckpointEvent): + snapshots.append(copy.deepcopy(event.response)) + assert snapshots + storage = stores.checkpoints[("alice", "conversation")] + latest = await storage.get_latest(workflow_name="recovery") + assert latest is not None + load = AsyncMock(wraps=storage.load) + monkeypatch.setattr(storage, "load", load) + context = _context("must-not-replay") + context.is_recovery = True + if snapshot_kind != "latest": + context.persisted_response = ( + next(snapshot for snapshot in snapshots if snapshot.get("output")) + if snapshot_kind == "partial" + else snapshots[0] + ) + events = await _collect(stores.server(_recovery_graph, options=options), context) + assert _types(events)[-1] == "response.completed", events + assert "must-not-replay" not in _text(events) + load.assert_awaited_once() + if snapshot_kind != "latest": + assert load.await_args is not None + assert load.await_args.args[0] != latest.checkpoint_id + assert "original" in _text(events) + if snapshot_kind == "partial": + assert _text(events) == "last:original" + final_text = "".join( + content.get("text", "") + for item in events[-1]["response"]["output"] + for content in item.get("content", []) + ) + assert final_text == "first:originallast:original" + else: + load.assert_awaited_once_with(latest.checkpoint_id) diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index fc36d00373c..6b7f739afca 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -19,6 +19,8 @@ FoundryCheckpointStore, FoundryFunctionApprovalStore, FunctionApprovalStoreProvider, + _InvocationsAgentSessionStoreProvider, + _InvocationsCheckpointStoreProvider, ) @@ -79,6 +81,30 @@ def test_storage_providers_use_public_abstraction() -> None: assert issubclass(AgentSessionStoreProvider, StoreProvider) +@pytest.mark.parametrize("is_hosted", [False, True]) +async def test_invocations_namespaces_cannot_overlap_responses_records(is_hosted: bool) -> None: + store = _store() + config = _config(is_hosted=is_hosted) + context = _platform_context() + with patch( + "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", + new=AsyncMock(return_value=store), + ) as get_or_create: + for provider in (CheckpointStoreProvider(), _InvocationsCheckpointStoreProvider()): + await provider.get_store(config=config, context_id="same-id", platform_context=context).save( + _checkpoint("same-checkpoint") + ) + for provider in (AgentSessionStoreProvider(), _InvocationsAgentSessionStoreProvider()): + await provider.get_store(config=config, platform_context=context).set("same-id", AgentSession()) + assert [call.args[0] for call in get_or_create.await_args_list] == [ + "checkpoints/same-id", + "invocations_checkpoints/same-id", + "agent_sessions", + "invocations_agent_sessions", + ] + assert all(call.kwargs == {"user_isolation": True} for call in get_or_create.await_args_list) + + async def test_save_uses_context_scoped_store() -> None: store = _store() checkpoint = _checkpoint("checkpoint-1") 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..41f5233978c 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,14 @@ 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) supplies `ResponsesHostServer` with an `agent_factory`. Each request builds three new `Agent` +instances, registers them with a new `WorkflowFactory` so the YAML's `InvokeAzureAgent` actions can resolve them +by name, loads a new workflow, and wraps it with `.as_agent(...)`. The host restores the authorized conversation's +checkpoint into that instance when continuing a conversation. + +The `FoundryChatClient` is opened once in `main` and closed when the host exits. It is shared across requests, +but the agents and workflow executors are not. The YAML's stable identifiers allow later factory-created +workflows to restore earlier checkpoints. 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..438a2c2e944 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 @@ -1,10 +1,11 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import os 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 +87,10 @@ class TriageResponse(BaseModel): # --- Host setup ------------------------------------------------------------------ -def main() -> None: +def create_workflow_agent(client: FoundryChatClient) -> WorkflowAgent: + """Rebuild the YAML workflow and its agents for the current request.""" 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(), - ) - # The workflow's InvokeAzureAgent actions reference these agents by name. triage_agent = Agent( client=client, @@ -128,7 +124,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,8 +133,19 @@ def main() -> None: ), ) - ResponsesHostServer(workflow_agent).run() + +async def main() -> None: + """Share only the model client while each request gets a new workflow.""" + with DefaultAzureCredential() as credential: + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=credential, + ) + async with client.project_client, client.client: + server = ResponsesHostServer(agent_factory=lambda: create_workflow_agent(client)) + await server.run_async() if __name__ == "__main__": - main() + asyncio.run(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..45fcd706439 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 @@ -16,7 +16,7 @@ The workflow has three executors (see [main.py](main.py)): user's message. If no valid target is found, the workflow yields an error message instead of counting down. - **`CountdownExecutor`** decrements the target through a self-loop, sleeping for a second and yielding an output on each tick, to simulate a long-running operation. -- **`complete`** yields the workflow's final `"Countdown complete."` output once the countdown reaches zero. +- **`CompleteExecutor`** yields the workflow's final `"Countdown complete."` output once the countdown reaches zero. ### Agent Hosting @@ -26,6 +26,11 @@ Setting `resilient_background=True` in `ResponsesServerOptions` enables the fram workflow's progress and durably persist streamed output, so a background response can be recovered and resumed after a crash (see "Testing resiliency" below). +The host receives an `agent_factory`, which builds all three executors and the target-extraction agent anew for +each request, including recovery. The workflow name and executor IDs are stable across factory calls. The model +client is owned by `main`, shared across requests, and closed when the host exits. Conversation state comes from +the authorized checkpoint, not from a workflow instance left over from another request. + ## Running the Agent Host Follow the instructions in the [Running the Agent Host Locally](../../README.md#running-the-agent-host-locally) section of the README in the parent directory to run the agent host. 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..58775a9a4fe 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, handler from agent_framework.foundry import FoundryChatClient from agent_framework_foundry_hosting import ResponsesHostServer from azure.ai.agentserver.responses import ResponsesServerOptions @@ -70,19 +70,19 @@ async def countdown(self, target: int, ctx: WorkflowContext[int | str, str]) -> await ctx.send_message(target - 1, target_id=self.id) -@executor(id="complete") -async def complete(message: str, ctx: WorkflowContext[Never, str]) -> None: +class CompleteExecutor(Executor): """Yield the workflow's completion output.""" - await ctx.yield_output(message) + def __init__(self, id: str = "complete") -> None: + super().__init__(id=id) + + @handler + async def complete(self, 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", @@ -93,9 +93,10 @@ def build_workflow(): ) start = StartExecutor(target_agent) countdown = CountdownExecutor() + complete = CompleteExecutor() 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) @@ -103,17 +104,23 @@ def build_workflow(): ) -def main() -> None: +async 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") - server = ResponsesHostServer( - workflow_agent, - options=ResponsesServerOptions(resilient_background=True), - log_level="DEBUG", - ) - server.run() + with DefaultAzureCredential() as credential: + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=credential, + ) + async with client.project_client, client.client: + server = ResponsesHostServer( + agent_factory=lambda: build_workflow(client).as_agent(name="countdown-workflow"), + options=ResponsesServerOptions(resilient_background=True), + log_level="DEBUG", + ) + await server.run_async() if __name__ == "__main__": - main() + asyncio.run(main()) 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..b53fea7bb6c 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,13 @@ 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 via `.as_agent()` and supplied through `ResponsesHostServer(agent_factory=...)`. +Every request creates fresh agents, executors, and workflow state. The host restores the current conversation's +checkpoint into that new instance when continuing an existing conversation. + +The model client is opened once in `main` and closed when the host exits. Only that client is shared across requests; +the mutable agents and executors are created inside the factory. The workflow name and executor names remain stable +so newly created workflows can load earlier checkpoints. ## 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..204b6d7a7ea 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 @@ -1,8 +1,9 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio 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,13 +12,8 @@ load_dotenv() -def main(): - 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) -> WorkflowAgent: + """Create fresh agents and executors for one request, using the host-owned client.""" writer_agent = Agent( client=client, instructions=("You are an excellent slogan writer. You create new slogans based on the given topic."), @@ -48,8 +44,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,9 +58,19 @@ def main(): .as_agent() ) - server = ResponsesHostServer(workflow_agent) - server.run() + +async def main() -> None: + """Keep the model client open while request factories create independent workflows.""" + with DefaultAzureCredential() as credential: + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=credential, + ) + async with client.project_client, client.client: + server = ResponsesHostServer(agent_factory=lambda: create_workflow_agent(client)) + await server.run_async() if __name__ == "__main__": - main() + asyncio.run(main()) From bc653c342a7fda85a887481185ba22daf5243b54 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Mon, 14 Sep 2026 18:36:54 +0100 Subject: [PATCH 02/11] Python: Refine Foundry factory ownership and session persistence Validate functional response batches before state changes, persist ordinary Invocations factory sessions, and document fresh execution objects as the factory author responsibility. Stabilize recovery fixtures and correct CI test typing. --- ...-python-foundry-request-agent-factories.md | 2 + python/packages/foundry_hosting/README.md | 55 ++- .../_agent_factory.py | 17 +- .../_invocations.py | 20 +- .../_responses.py | 14 +- .../tests/test_agent_factory.py | 80 ++++ .../tests/test_invocations_factory.py | 382 ++++++++++++++++-- .../tests/test_responses_factory.py | 300 +++++++++++++- .../foundry_hosting/tests/test_state_store.py | 8 +- 9 files changed, 805 insertions(+), 73 deletions(-) create mode 100644 python/packages/foundry_hosting/tests/test_agent_factory.py diff --git a/docs/decisions/0040-python-foundry-request-agent-factories.md b/docs/decisions/0040-python-foundry-request-agent-factories.md index 81f13d2ee57..164e32788c3 100644 --- a/docs/decisions/0040-python-foundry-request-agent-factories.md +++ b/docs/decisions/0040-python-foundry-request-agent-factories.md @@ -61,6 +61,8 @@ distributed coordination or checkpoint replay as exactly-once execution. Workflow callers migrate from `Host(workflow_agent)` to `Host(agent_factory=create_agent)`. The factory must construct new mutable objects, not return the same instance or reuse stateful executors. +This is the factory implementer's responsibility. The host calls the factory per request but does not inspect +object identities, recursively check captured state, or require weak-reference support. Stable workflow names, executor IDs, and serialization registrations are necessary for continuation. Ordinary-agent instance callers keep their existing lifecycle. Factories do not automatically save diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 8956a2f481a..f070aaa46da 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -35,9 +35,33 @@ server = ResponsesHostServer(agent_factory=create_agent) The host calls the factory inside the current request's platform context, not during startup. The resulting agent belongs to that request, including its entire stream. Responses recovery also creates a new agent through the factory. -Returning the same workflow instance repeatedly is not supported. Neither is building a new workflow around previously -used mutable executors. Keep workflow names, executor IDs, and serialized state type registrations stable so a new -instance can restore the previous instance's checkpoints. +The factory implementer is responsible for creating fresh mutable workflows, executors, wrapped agents, and their +sessions inside the factory. Returning an existing workflow or rebuilding around shared mutable executors can carry +state across requests. The host does not inspect or track object identities, clone execution objects, or enforce +weak-reference support. + +For example, construct the agent and executor inside the factory, not once outside it: + +```python +from agent_framework import Agent, AgentExecutor, WorkflowBuilder +from agent_framework_foundry_hosting import ResponsesHostServer + + +def create_agent(): + agent = Agent(client=client, name="assistant") + executor = AgentExecutor(agent, id="assistant") + return WorkflowBuilder( + name="assistant-workflow", start_executor=executor + ).build().as_agent() + + +server = ResponsesHostServer(agent_factory=create_agent) +``` + +Here `client` is an application-owned model client whose concurrency and resource lifetime support sharing. +Any mutable objects captured by tools or functional workflow functions also remain the implementer's responsibility. +Keep workflow names, executor IDs, and serialized state type registrations stable so a new instance can restore the +previous instance's checkpoints. Factories are also useful for ordinary agents with request-specific configuration. They do not automatically persist custom fields on an agent: state needed on the next request must use the supported session or checkpoint stores. @@ -164,6 +188,29 @@ but not all output buffered inside those steps. Recovering from such a checkpoin that configuration instead of silently losing output or rerunning application work. Functional workflows can use the factory for normal Responses requests and supported pending-response continuation. +Functional continuation validates every supplied result against the authorized pending request and its declared +response type, using the same supported conversions as graph workflows. A string such as `"false"` is not a boolean +decision. Invalid types, unknown request IDs, and duplicate responses fail the whole submitted batch before the host +copies a checkpoint, records an execution attempt, or runs the workflow. The original pending state remains available +for a corrected retry. + +### Invocations ordinary factory sessions + +`InvocationsHostServer(agent_factory=...)` also persists ordinary agents' `AgentSession` through +`agent_session_store_provider`. The default provider uses Foundry storage when hosted and file-based storage locally. +Sessions are scoped by the platform user and invocation session, with ordinary factory records separated from workflow +metadata and Responses records. + +Each request loads the stored session or calls the new agent's `create_session` method, passes that session to the run, +and saves it when the run finishes or is interrupted. For streaming requests, saving occurs after the iterator closes. +The host holds its per-scope lock through execution, saving, and resource cleanup. A new host can continue the stored +session, but custom mutable agent fields are not persisted automatically. + +Factory sessions are not retained in the host's session dictionary. Retention, expiration, storage quotas, and any +in-memory caching are the configured storage provider's responsibility; the host adds no expiry or eviction policy. +Ordinary **instance** calls remain unchanged: they retain sessions in the host's unbounded in-memory dictionary and +do not use this persistence provider. Changing that legacy retention behavior is outside this factory feature. + ### Invocations workflows `InvocationsHostServer(agent_factory=...)` persists workflow checkpoints for the platform user and invocation session. @@ -177,7 +224,7 @@ message is not an approval response. Use Responses when callers need that struct Functional workflows accept a new message after a completed invocation, but pending or interrupted functional continuation is rejected. They do not acquire graph-workflow recovery semantics by being passed through a factory. -Requests updating the same workflow scope are serialized within one host; independent scopes can execute concurrently. +Requests updating the same factory scope are serialized within one host; independent scopes can execute concurrently. This is not a distributed lock across multiple host processes. Invocations does not automatically recover an interrupted HTTP response, and checkpoints do not guarantee that external side effects execute exactly once. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py index 587a2cc6c51..782dd4e276e 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py @@ -6,7 +6,6 @@ import asyncio import inspect -import weakref from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable from contextlib import asynccontextmanager from dataclasses import dataclass, field @@ -44,27 +43,19 @@ def validate_agent_source(agent: HostedAgent | None, agent_factory: AgentFactory class AgentFactoryResolver: - """Resolve request agents without retaining completed workflow runtimes.""" + """Resolve request agents without caching them or inspecting their execution objects. + + Factory authors are responsible for constructing independent mutable runtimes. + """ def __init__(self, factory: AgentFactory) -> None: self._factory = factory - self._seen: weakref.WeakValueDictionary[int, object] = weakref.WeakValueDictionary() async def resolve(self) -> HostedAgent: result = self._factory() agent = await result if inspect.isawaitable(result) else result if not isinstance(agent, (SupportsAgentRun, FunctionalWorkflowAgent)): raise TypeError("agent_factory must return an agent implementing SupportsAgentRun or a workflow agent.") - if is_workflow_agent(agent): - workflow = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow # pyright: ignore[reportPrivateUsage] - for value in (agent, workflow): - if self._seen.get(id(value)) is value: - raise RuntimeError( - "agent_factory reused a workflow agent or workflow. Create a new workflow and " - "new mutable executors for each request." - ) - for value in (agent, workflow): - self._seen[id(value)] = value 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 49c084ed386..c3cc7e1d284 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -12,6 +12,7 @@ CheckpointStorage, FunctionalWorkflowAgent, SessionStore, + SupportsAgentRun, WorkflowAgent, WorkflowRunState, ) @@ -85,7 +86,8 @@ def __init__( workflows, mutable executors, providers, and tools. An async context manager returned by the factory is entered and exited within that request. checkpoint_store_provider: Optional provider for user/session-scoped workflow checkpoints. - agent_session_store_provider: Optional provider for persisted workflow provider state. + agent_session_store_provider: Optional provider for persisted factory agent sessions, + including workflow provider state. Storage retention is controlled by the provider. openapi_spec: The OpenAPI specification for the server. **kwargs: Additional keyword arguments. @@ -95,6 +97,8 @@ def __init__( The text-only exchange cannot answer pending workflow requests or approvals. Functional workflows support fresh messages after clean completion, but not pending or interrupted continuation. Such continuation fails explicitly. + Ordinary factory sessions are persisted after each run, including interrupted runs, + rather than retained in this host. Ordinary agent instances retain sessions in memory. Factories own cleanup of nested resources not exposed by an async context manager. """ validate_agent_source(agent, agent_factory) @@ -241,7 +245,19 @@ async def _factory_session( async with self._workflow_session(agent, storage_id) as (session, storage): yield session, {"checkpoint_storage": storage} else: - yield self._sessions.setdefault(storage_id, AgentSession(session_id=storage_id)), {} + # Keep ordinary snapshots separate from workflow metadata in the same provider. + storage_id = f"ordinary-{storage_id}" + sessions = self._agent_session_storage_provider.get_store( + config=self.config, platform_context=get_request_context() + ) + session = await sessions.get(storage_id) + if session is None: + session = cast(SupportsAgentRun, agent).create_session(session_id=storage_id) + try: + yield session, {} + finally: + with CancelScope(shield=True): + await sessions.set(storage_id, session) async def _handle_factory_invoke(self, user_message: Any, *, stream: bool) -> Response: context = get_request_context() 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 f6092d53d46..e81592b5835 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -38,6 +38,7 @@ add_usage_details, ) from agent_framework._telemetry import mark_feature_used +from agent_framework._workflows._typing_utils import is_instance_of, try_coerce_to_type from agent_framework.exceptions import AgentFrameworkException from anyio import CancelScope from azure.ai.agentserver.core import get_request_context @@ -1496,15 +1497,22 @@ async def _handle_functional_workflow( raise ValueError("The input does not match an authorized pending functional workflow request.") pending_request = pending[request_id] if content.type == "function_result": - responses[request_id] = content if pending_request.response_type is Content else content.result + response = content if pending_request.response_type is Content else content.result elif content.type == "function_approval_response" and pending_request.response_type is bool: - responses[request_id] = content.approved + response = content.approved elif content.type == "function_approval_response" and pending_request.response_type is Content: - responses[request_id] = content + response = content else: raise ValueError( "This pending functional workflow request requires a matching function result." ) + response = try_coerce_to_type(response, pending_request.response_type) + if not is_instance_of(response, pending_request.response_type): + raise ValueError( + f"Response type mismatch for request ID {request_id}: " + f"expected {pending_request.response_type}, got {type(response)}" + ) + responses[request_id] = response if not responses: raise ValueError("Pending functional workflow requests require structured responses.") run_kwargs = { diff --git a/python/packages/foundry_hosting/tests/test_agent_factory.py b/python/packages/foundry_hosting/tests/test_agent_factory.py new file mode 100644 index 00000000000..96c1d44b402 --- /dev/null +++ b/python/packages/foundry_hosting/tests/test_agent_factory.py @@ -0,0 +1,80 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Factory resolution leaves execution-object ownership to the application.""" + +from typing import Any + +import pytest +from agent_framework import AgentExecutor, AgentResponse, AgentSession, Message, WorkflowAgent, WorkflowBuilder + +from agent_framework_foundry_hosting._agent_factory import AgentFactoryResolver + + +class _SlottedAgent: + __slots__ = ("description", "id", "name") + + def __init__(self) -> None: + self.id = "slotted" + self.name: str | None = "slotted" + self.description: str | None = None + + def create_session(self, *, session_id: str | None = None) -> AgentSession: + return AgentSession(session_id=session_id) + + def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: + return AgentSession(session_id=session_id, service_session_id=service_session_id) + + def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any: + async def response() -> AgentResponse: + return AgentResponse(messages=[Message("assistant", ["slotted"])]) + + return response() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_factory_does_not_require_weak_reference_support(asynchronous: bool) -> None: + calls = 0 + + def factory() -> WorkflowAgent: + nonlocal calls + calls += 1 + executor = AgentExecutor(_SlottedAgent(), id="inner") + return WorkflowBuilder(name="slotted-workflow", start_executor=executor).build().as_agent() + + async def async_factory() -> WorkflowAgent: + return factory() + + resolver = AgentFactoryResolver(async_factory if asynchronous else factory) + first = await resolver.resolve() + second = await resolver.resolve() + assert isinstance(first, WorkflowAgent) + assert isinstance(second, WorkflowAgent) + assert first is not second + assert first.workflow.executors["inner"] is not second.workflow.executors["inner"] + assert (await first.run("first")).text == "slotted" + assert (await second.run("second")).text == "slotted" + assert calls == 2 + + +@pytest.mark.parametrize("same_wrapper", [False, True]) +async def test_resolver_returns_factory_result_without_enforcing_object_ownership(same_wrapper: bool) -> None: + executor = AgentExecutor(_SlottedAgent(), id="inner") + agent = WorkflowBuilder(name="application-owned", start_executor=executor).build().as_agent() + calls = 0 + + def factory() -> WorkflowAgent: + nonlocal calls + calls += 1 + if same_wrapper: + return agent + return WorkflowBuilder(name="application-owned", start_executor=executor).build().as_agent() + + resolver = AgentFactoryResolver(factory) + first = await resolver.resolve() + second = await resolver.resolve() + assert isinstance(first, WorkflowAgent) + assert isinstance(second, WorkflowAgent) + assert first.workflow.executors["inner"] is executor + assert second.workflow.executors["inner"] is executor + assert (first is second) is same_wrapper + assert calls == 2 diff --git a/python/packages/foundry_hosting/tests/test_invocations_factory.py b/python/packages/foundry_hosting/tests/test_invocations_factory.py index 05337b44f09..a885c348415 100644 --- a/python/packages/foundry_hosting/tests/test_invocations_factory.py +++ b/python/packages/foundry_hosting/tests/test_invocations_factory.py @@ -3,6 +3,7 @@ """Request factory lifetime and workflow persistence for the text-only host.""" import asyncio +import json from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from typing import Any, cast @@ -36,6 +37,7 @@ step, workflow, ) +from agent_framework._filesystem import _storage_key_segment from anyio import CancelScope, create_task_group from azure.ai.agentserver.core import ( FoundryAgentRequestContext, @@ -139,9 +141,9 @@ async def response() -> AgentResponse: class _Stores: - def __init__(self) -> None: + def __init__(self, sessions: SessionStore | None = None) -> None: self.checkpoints: dict[str, InMemoryCheckpointStorage] = {} - self.sessions = SessionStore() + self.sessions = sessions if sessions is not None else SessionStore() self.checkpoint_provider = MagicMock() self.checkpoint_provider.get_store.side_effect = self._checkpoint_store self.session_provider = MagicMock() @@ -167,7 +169,11 @@ def __init__(self) -> None: self.count = 0 @handler - async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + async def count_message( + self, + messages: list[Message], + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: self.count += 1 # Include the outer provider history, not just executor checkpoint state. await ctx.yield_output(f"{self.count}:{'|'.join(message.text for message in messages)}") @@ -196,11 +202,16 @@ async def _functional_pending(messages: Any, ctx: RunContext) -> str: class _Pending(Executor): @handler - async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] await ctx.request_info("approve?", response_type=str) @response_handler - async def answer(self, original_request: str, response: str, ctx: WorkflowContext[Never, str]) -> None: + async def answer( + self, + original_request: str, + response: str, + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: await ctx.yield_output(response) @@ -247,11 +258,13 @@ async def create() -> _OwnedAgent: return create() if awaitable else agent - server = InvocationsHostServer(agent_factory=factory) + server = _Stores().server(factory) assert await _invoke(server, stream=stream) == "ab" assert await _invoke(server, "second", stream=stream) == "ab" assert len(agents) == 2 - assert agents[0].session is agents[1].session + assert agents[0].session is not None + assert agents[1].session is not None + assert agents[0].session.session_id == agents[1].session.session_id assert agents[0].calls == ["hello" if stream else ["hello"]] assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) * 2 assert server._agent is None @@ -281,7 +294,7 @@ async def factory() -> Any: @pytest.mark.parametrize("stream", [False, True]) async def test_run_failure_closes_owner_and_iterator(stream: bool) -> None: events: list[str] = [] - server = InvocationsHostServer(agent_factory=lambda: _OwnedAgent(events, fail=True)) + server = _Stores().server(lambda: _OwnedAgent(events, fail=True)) with pytest.raises(RuntimeError, match="run failed"): await _invoke(server, stream=stream) assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) @@ -297,7 +310,7 @@ def factory() -> _OwnedAgent: creation_task = asyncio.current_task() return _OwnedAgent(events) - server = InvocationsHostServer(agent_factory=factory) + server = _Stores().server(factory) with _context(): response = await server._handle_invoke(_request(stream=True)) assert creation_task is None @@ -319,7 +332,7 @@ async def send(message: Any) -> None: async def test_stream_disconnect_or_cancellation_closes_suspended_iterator(disconnect: bool) -> None: events: list[str] = [] sent = asyncio.Event() - server = InvocationsHostServer(agent_factory=lambda: _OwnedAgent(events)) + server = _Stores().server(lambda: _OwnedAgent(events)) with _context(): response = await server._handle_invoke(_request(stream=True)) assert isinstance(response, StreamingResponse) @@ -369,7 +382,7 @@ def factory() -> _OwnedAgent: created += 1 return _OwnedAgent(events) - server = InvocationsHostServer(agent_factory=factory) + server = _Stores().server(factory) with _context(): response = await server._handle_invoke(_request(stream=True)) assert isinstance(response, StreamingResponse) @@ -428,19 +441,6 @@ async def test_graph_scopes_isolate_users_sessions_and_unsafe_identifiers(stream assert not server._sessions -@pytest.mark.parametrize("same_wrapper", [False, True]) -@pytest.mark.parametrize("functional", [False, True]) -async def test_reused_workflow_or_wrapper_is_rejected(same_wrapper: bool, functional: bool) -> None: - stores = _Stores() - agent = _functional.build().as_agent() if functional else _graph() - underlying = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow - server = stores.server(lambda: agent if same_wrapper else underlying.as_agent()) - await _invoke(server) - with pytest.raises(RuntimeError, match="reused a workflow"): - await _invoke(server, session="other") - assert not server._scope_locks._entries - - @pytest.mark.parametrize("damage", ["checkpoint", "session", "name", "kind"]) async def test_missing_or_incompatible_continuation_never_restarts(damage: str) -> None: stores = _Stores() @@ -563,7 +563,7 @@ async def __aexit__(self, *args: Any) -> None: await self.group.__aexit__(*args) await super().__aexit__(*args) - server = InvocationsHostServer(agent_factory=lambda: TaskGroupAgent(events)) + server = _Stores().server(lambda: TaskGroupAgent(events)) with _context(): response = await server._handle_invoke(_request(stream=True)) assert isinstance(response, StreamingResponse) @@ -607,6 +607,330 @@ async def response() -> ChatResponse: return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() +class _SnapshotSessions(SessionStore): + def __init__(self) -> None: + self.snapshots: dict[str, str] = {} + + async def get(self, session_id: str) -> AgentSession | None: + snapshot = self.snapshots.get(session_id) + return AgentSession.from_dict(json.loads(snapshot)) if snapshot is not None else None + + async def set(self, session_id: str, session: AgentSession) -> None: + await asyncio.sleep(0) + self.snapshots[session_id] = json.dumps(session.to_dict()) + + +class _PersistedAgent(Agent): + def __init__(self, client: BaseChatClient, events: list[str]) -> None: + super().__init__(client=client, name="ordinary") + self.events = events + + async def __aenter__(self) -> Self: + self.events.append("enter") + return await super().__aenter__() + + async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: + await super().__aexit__(exc_type, exc_val, exc_tb) + self.events.append("exit") + + def create_session(self, *, session_id: str | None = None) -> AgentSession: + self.events.append("create") + session = super().create_session(session_id=session_id) + session.state["initialized"] = True + return session + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_ordinary_factory_sessions_persist_without_host_retention(stream: bool) -> None: + snapshots = _SnapshotSessions() + stores = _Stores(snapshots) + transcripts: list[list[str]] = [] + events: list[str] = [] + + def factory() -> Agent: + return _PersistedAgent(_TranscriptClient(transcripts), events) + + first = stores.server(factory) + for index in range(40): + assert await _invoke(first, f"first-{index}", session=str(index), stream=stream) == "recorded" + assert not first._sessions + assert not first._scope_locks._entries + assert len(snapshots.snapshots) == 40 + assert events.count("create") == 40 + second = stores.server(factory) + assert await _invoke(second, "next", session="0", stream=stream) == "recorded" + assert transcripts[-1] == ["first-0", "recorded", "next"] + assert events.count("create") == 40 + assert events.count("enter") == events.count("exit") == 41 + assert not second._sessions + assert not second._scope_locks._entries + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_ordinary_factory_persistence_isolates_authorized_scopes(stream: bool) -> None: + snapshots = _SnapshotSessions() + stores = _Stores(snapshots) + transcripts: list[list[str]] = [] + scopes = [ + ("user", "../session"), + ("other", "../session"), + ("user", r"..\session"), + ("a:b", "c"), + ("b", "c:a"), + ("USER", "../session"), + ] + for index, (user, session) in enumerate(scopes): + server = stores.server(lambda: Agent(client=_TranscriptClient(transcripts), name="ordinary")) + server.config.is_hosted = True + await _invoke(server, str(index), stream=stream, user=user, session=session) + assert transcripts[-1] == [str(index)] + for index, (user, session) in enumerate(scopes): + server = stores.server(lambda: Agent(client=_TranscriptClient(transcripts), name="ordinary")) + server.config.is_hosted = True + await _invoke(server, "next", stream=stream, user=user, session=session) + assert transcripts[-1] == [str(index), "recorded", "next"] + assert len(snapshots.snapshots) == len(scopes) + assert all("/" not in key and "\\" not in key for key in snapshots.snapshots) + assert stores.session_provider.get_store.call_args.kwargs["platform_context"].user_id == "USER" + stores.checkpoint_provider.get_store.assert_not_called() + + +async def test_ordinary_factory_namespace_does_not_replace_workflow_or_responses_sessions() -> None: + snapshots = _SnapshotSessions() + stores = _Stores(snapshots) + workflow_key = _storage_key_segment(json.dumps(("user", "session")), encoded_prefix="~invocations-") + # Responses uses conversation/response IDs, while Invocations encodes its user/session pair. + responses_session = AgentSession(session_id="session") + responses_session.state["responses"] = True + await snapshots.set("session", responses_session) + assert await _invoke(stores.server(_graph), "workflow-first") == "1:workflow-first" + original_records = dict(snapshots.snapshots) + transcripts: list[list[str]] = [] + await _invoke(stores.server(lambda: Agent(client=_TranscriptClient(transcripts), name="ordinary")), "ordinary") + assert transcripts == [["ordinary"]] + assert len(snapshots.snapshots) == 3 + assert {key: snapshots.snapshots[key] for key in original_records} == original_records + ordinary_key = next(key for key in snapshots.snapshots if key not in original_records) + assert ordinary_key != workflow_key + assert ordinary_key.startswith("ordinary-~invocations-") + assert await _invoke(stores.server(_graph), "workflow-next") == "2:workflow-first|1:workflow-first|workflow-next" + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("operation", ["get", "set"]) +async def test_ordinary_factory_store_failure_releases_resources_and_scope( + stream: bool, operation: str, monkeypatch: pytest.MonkeyPatch +) -> None: + snapshots = _SnapshotSessions() + stores = _Stores(snapshots) + transcripts: list[list[str]] = [] + events: list[str] = [] + original = getattr(snapshots, operation) + failing = AsyncMock(side_effect=RuntimeError(f"{operation} failed")) + monkeypatch.setattr(snapshots, operation, failing) + server = stores.server(lambda: _PersistedAgent(_TranscriptClient(transcripts), events)) + with pytest.raises(RuntimeError, match=f"{operation} failed"): + await _invoke(server, stream=stream) + assert events[-1] == "exit" + assert transcripts == ([] if operation == "get" else [["hello"]]) + assert not server._sessions + assert not server._scope_locks._entries + assert not snapshots.snapshots + monkeypatch.setattr(snapshots, operation, original) + assert await _invoke(server, "retry", stream=stream) == "recorded" + assert len(snapshots.snapshots) == 1 + + +class _InterruptedClient(BaseChatClient): + def __init__(self, started: asyncio.Event, events: list[str], *, fail: bool) -> None: + super().__init__() + self.started = started + self.events = events + self.fail = fail + + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Any: + async def interrupt() -> None: + self.started.set() + if self.fail: + raise RuntimeError("client failed") + await asyncio.Event().wait() + + async def updates() -> AsyncGenerator[ChatResponseUpdate]: + try: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("partial")]) + await interrupt() + finally: + self.events.append("client closed") + + async def response() -> ChatResponse: + await interrupt() + return ChatResponse(messages=[]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("failure", ["error", "asyncio", "anyio"]) +async def test_ordinary_factory_interruption_persists_session_before_resource_cleanup( + stream: bool, failure: str, monkeypatch: pytest.MonkeyPatch +) -> None: + snapshots = _SnapshotSessions() + stores = _Stores(snapshots) + events: list[str] = [] + started = asyncio.Event() + original_set = snapshots.set + + async def save(session_id: str, session: AgentSession) -> None: + await original_set(session_id, session) + events.append("saved") + + monkeypatch.setattr(snapshots, "set", save) + server = stores.server( + lambda: _PersistedAgent(_InterruptedClient(started, events, fail=failure == "error"), events) + ) + scope = CancelScope() + + async def consume() -> None: + with scope: + await _invoke(server, stream=stream) + + task = asyncio.create_task(consume()) + await started.wait() + if failure == "anyio": + scope.cancel() + await task + else: + if failure == "asyncio": + task.cancel() + with pytest.raises(RuntimeError if failure == "error" else asyncio.CancelledError): + await task + assert events[-2:] == ["saved", "exit"] + if stream: + assert events.index("client closed") < events.index("saved") + assert len(snapshots.snapshots) == 1 + saved = await snapshots.get(next(iter(snapshots.snapshots))) + assert saved is not None + assert saved.state["initialized"] is True + assert not server._sessions + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("interruption", ["close", "disconnect", "cancel"]) +async def test_ordinary_factory_suspended_stream_persists_on_close( + interruption: str, monkeypatch: pytest.MonkeyPatch +) -> None: + snapshots = _SnapshotSessions() + stores = _Stores(snapshots) + events: list[str] = [] + original_set = snapshots.set + + async def save(session_id: str, session: AgentSession) -> None: + await original_set(session_id, session) + events.append("saved") + + monkeypatch.setattr(snapshots, "set", save) + server = stores.server(lambda: _PersistedAgent(_TranscriptClient([]), events)) + with _context(): + response = await server._handle_invoke(_request(stream=True)) + assert isinstance(response, StreamingResponse) + assert not snapshots.snapshots + assert not events + if interruption == "close": + iterator = cast(AsyncGenerator[str], response.body_iterator) + assert await anext(iterator) == "recorded" + assert not snapshots.snapshots + await iterator.aclose() + else: + sent = asyncio.Event() + + async def send(message: Any) -> None: + if message["type"] == "http.response.body" and message.get("body"): + assert not snapshots.snapshots + sent.set() + if interruption == "disconnect": + raise OSError("disconnected") + await asyncio.Event().wait() + + consumer = asyncio.create_task(response.stream_response(send)) + await sent.wait() + if interruption == "cancel": + consumer.cancel() + with pytest.raises(OSError if interruption == "disconnect" else asyncio.CancelledError): + await consumer + assert events[-2:] == ["saved", "exit"] + assert len(snapshots.snapshots) == 1 + assert not server._sessions + assert not server._scope_locks._entries + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_ordinary_factory_lock_covers_save_and_allows_independent_scopes( + stream: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + snapshots = _SnapshotSessions() + stores = _Stores(snapshots) + transcripts: list[list[str]] = [] + events: list[str] = [] + saving = asyncio.Event() + release = asyncio.Event() + original_set = snapshots.set + + async def save(session_id: str, session: AgentSession) -> None: + if get_request_context().session_id == "session" and not saving.is_set(): + saving.set() + await release.wait() + await original_set(session_id, session) + + monkeypatch.setattr(snapshots, "set", save) + server = stores.server(lambda: _PersistedAgent(_TranscriptClient(transcripts), events)) + first = asyncio.create_task(_invoke(server, "first", stream=stream)) + waiting: asyncio.Task[str] | None = None + try: + await asyncio.wait_for(saving.wait(), timeout=2) + waiting = asyncio.create_task(_invoke(server, "second", stream=stream)) + await asyncio.sleep(0) + assert events.count("enter") == 1 + assert events.count("exit") == 0 + assert not first.done() + assert await asyncio.wait_for(_invoke(server, "independent", session="other"), timeout=2) == "recorded" + assert transcripts == [["first"], ["independent"]] + assert not waiting.done() + release.set() + assert await first == "recorded" + assert await waiting == "recorded" + assert transcripts[-1] == ["first", "recorded", "second"] + assert len(snapshots.snapshots) == 2 + assert not server._sessions + assert not server._scope_locks._entries + finally: + release.set() + await first + if waiting is not None: + await waiting + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_ordinary_instance_keeps_legacy_memory_and_ignores_session_provider(stream: bool) -> None: + provider = MagicMock() + provider.get_store.side_effect = AssertionError("instance must not access persisted storage") + transcripts: list[list[str]] = [] + agent = Agent(client=_TranscriptClient(transcripts), name="ordinary") + server = InvocationsHostServer(agent, agent_session_store_provider=provider) + await _invoke(server, "first", stream=stream) + original = server._sessions["session"] + await _invoke(server, "second", stream=stream) + assert server._sessions["session"] is original + assert transcripts[-1] == ["first", "recorded", "second"] + await _invoke(server, "independent", session="other", stream=stream) + assert len(server._sessions) == 2 + replacement = InvocationsHostServer(agent, agent_session_store_provider=provider) + await _invoke(replacement, "fresh", stream=stream) + assert transcripts[-1] == ["fresh"] + provider.get_store.assert_not_called() + + @pytest.mark.parametrize("stream", [False, True]) async def test_real_agent_executor_transcript_is_isolated_and_restored(stream: bool) -> None: stores = _Stores() @@ -655,7 +979,7 @@ def factory() -> FunctionalWorkflowAgent: @pytest.mark.parametrize("stream", [False, True]) async def test_cancellation_during_run_closes_resources(stream: bool) -> None: events: list[str] = [] - server = InvocationsHostServer(agent_factory=lambda: _OwnedAgent(events, wait=asyncio.Event())) + server = _Stores().server(lambda: _OwnedAgent(events, wait=asyncio.Event())) task = asyncio.create_task(_invoke(server, stream=stream)) await asyncio.sleep(0) assert events == ["enter", "run"] @@ -716,7 +1040,11 @@ async def test_checkpoint_preparation_failure_rejects_older_checkpoint_and_fresh class FailingCheckpoint(_Counter): @handler - async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + async def count_message( + self, + messages: list[Message], + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: calls.extend(message.text for message in messages) await super().count_message(messages, ctx) diff --git a/python/packages/foundry_hosting/tests/test_responses_factory.py b/python/packages/foundry_hosting/tests/test_responses_factory.py index 03d57073542..3462d2039b7 100644 --- a/python/packages/foundry_hosting/tests/test_responses_factory.py +++ b/python/packages/foundry_hosting/tests/test_responses_factory.py @@ -5,12 +5,15 @@ import asyncio import copy import gc +import json import weakref from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Mapping, Sequence from contextlib import aclosing, contextmanager +from dataclasses import dataclass from typing import Any, cast from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from agent_framework import ( Agent, @@ -45,7 +48,8 @@ reset_request_context, set_request_context, ) -from azure.ai.agentserver.responses import ResponseContext, ResponsesServerOptions +from azure.ai.agentserver.responses import InMemoryResponseProvider, ResponseContext, ResponsesServerOptions +from azure.ai.agentserver.responses.aio import ResponseEventStream from azure.ai.agentserver.responses.models import CreateResponse from azure.ai.agentserver.responses.streaming._checkpoint import ResponseCheckpointEvent from typing_extensions import Never, Self @@ -222,7 +226,11 @@ def __init__(self) -> None: self.count = 0 @handler - async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + async def count_message( + self, + messages: list[Message], + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: self.count += 1 await ctx.yield_output(f"{self.count}:{'|'.join(message.text for message in messages)}") @@ -433,11 +441,16 @@ async def _pending_bool(messages: list[Message], ctx: RunContext) -> str: class _Pending(Executor): @handler - async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] await ctx.request_info(messages[0].text, response_type=str) @response_handler - async def answer(self, original_request: str, response: str, ctx: WorkflowContext[Never, str]) -> None: + async def answer( + self, + original_request: str, + response: str, + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: await ctx.yield_output(f"{original_request}:{response}") @@ -563,17 +576,6 @@ async def factory() -> Any: assert not server._scope_locks._entries -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("same_wrapper", [False, True]) -async def test_reusing_workflow_wrapper_or_underlying_workflow_fails(functional: bool, same_wrapper: bool) -> None: - agent = _functional.build().as_agent() if functional else _graph() - underlying = agent._workflow if isinstance(agent, FunctionalWorkflowAgent) else agent.workflow - server = _Stores().server(lambda: agent if same_wrapper else underlying.as_agent()) - assert _types(await _collect(server))[-1] == "response.completed" - assert "reused" in _failure(await _collect(server, _context(response="two"))) - assert not server._scope_locks._entries - - @pytest.mark.parametrize("functional", [False, True]) async def test_completed_workflows_are_not_retained_by_factory_resolver(functional: bool) -> None: refs: list[weakref.ReferenceType[Any]] = [] @@ -589,8 +591,6 @@ def factory() -> Any: await asyncio.sleep(0) gc.collect() assert all(reference() is None for reference in refs) - assert server._agent_resolver is not None - assert not server._agent_resolver._seen @pytest.mark.parametrize("finish", ["complete", "close", "cancel-signal", "model-error"]) @@ -982,6 +982,246 @@ def factory() -> FunctionalWorkflowAgent: assert _text(approved) == "hello:owner" +@dataclass +class _FunctionalAnswer: + text: str + score: float + + +@pytest.mark.parametrize( + ("response_type", "value"), + [(bool, True), (bool, False), (str, "accepted"), (_FunctionalAnswer, {"text": "accepted", "score": 2})], + ids=["true", "false", "string", "structured"], +) +async def test_functional_response_coerces_normalized_results( + response_type: type, value: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + received: list[Any] = [] + + @workflow(name="typed-functional") + async def typed(messages: list[Message], ctx: RunContext) -> str: + answer = await ctx.request_info(messages[0].text, response_type=response_type, request_id="answer") + received.append(answer) + return repr(answer) + + stores = _Stores() + + def factory() -> FunctionalWorkflowAgent: + return typed.build().as_agent() + + assert _types(await _collect(stores.server(factory)))[-1] == "response.completed" + # The wire converter stringifies outputs. Exercise core structured coercion on + # normalized framework content without introducing JSON parsing into the host. + monkeypatch.setattr( + "agent_framework_foundry_hosting._responses._items_to_messages", + AsyncMock(return_value=[Message("tool", [Content("function_result", call_id="answer", result=value)])]), + ) + resumed = await _collect(stores.server(factory), _context(response="two")) + assert _types(resumed)[-1] == "response.completed", resumed + expected = _FunctionalAnswer("accepted", 2.0) if response_type is _FunctionalAnswer else value + assert received == [expected] + assert type(received[0]) is response_type + if isinstance(received[0], _FunctionalAnswer): + assert type(received[0].score) is float + assert _text(resumed) == repr(expected) + + +@pytest.mark.parametrize("chain", [False, True]) +@pytest.mark.parametrize( + "invalid", + ["bool-string", "bool-number", "bool-null", "string-bool", "approval-string", "unknown", "duplicate"], +) +async def test_functional_invalid_batch_preserves_checkpoint_and_allows_retry( + chain: bool, invalid: str, monkeypatch: pytest.MonkeyPatch +) -> None: + executions: list[str] = [] + received: list[Any] = [] + + @workflow(name="typed-functional-batch") + async def typed(messages: list[Message], ctx: RunContext) -> str: + executions.append(messages[0].text) + answers = await asyncio.gather( + ctx.request_info("text?", response_type=str, request_id="text"), + ctx.request_info("approve?", response_type=bool, request_id="decision"), + ) + received.extend(answers) + return f"{answers[0]}:{answers[1]}" + + stores = _Stores() + + def factory() -> FunctionalWorkflowAgent: + return typed.build().as_agent() + + conversation = None if chain else "conversation" + source_id = conversation or "response-1" + assert ( + _types(await _collect(stores.server(factory), _context(conversation=conversation)))[-1] == "response.completed" + ) + storage = stores.checkpoints[("alice", source_id)] + checkpoint = await storage.get_latest(workflow_name="typed-functional-batch") + assert checkpoint is not None + assert set(checkpoint.pending_request_info_events) == {"text", "decision"} + checkpoint_before = json.dumps(checkpoint.to_dict(), default=lambda value: value.to_dict(), sort_keys=True) + sessions = stores.sessions["alice"] + session = await sessions.get(source_id) + assert session is not None + state_before = copy.deepcopy(session.state) + source_save = AsyncMock(wraps=storage.save) + session_save = AsyncMock(wraps=sessions.set) + monkeypatch.setattr(storage, "save", source_save) + monkeypatch.setattr(sessions, "set", session_save) + contents = [ + Content.from_function_result("text", result="accepted"), + Content("function_result", call_id="decision", result=False), + ] + if invalid.startswith("bool-"): + contents[1] = Content( + "function_result", + call_id="decision", + result={"bool-string": "false", "bool-number": 1, "bool-null": None}[invalid], + ) + elif invalid == "string-bool": + contents[0] = Content("function_result", call_id="text", result=False) + elif invalid == "approval-string": + contents[1] = Content("function_approval_response", id="decision", approved=cast(Any, "false")) + elif invalid == "unknown": + contents.append(Content.from_function_result("not-authorized", result="extra")) + else: + contents.append(Content.from_function_result("text", result="duplicate")) + converted = AsyncMock(return_value=[Message("tool", contents)]) + monkeypatch.setattr("agent_framework_foundry_hosting._responses._items_to_messages", converted) + + rejected = await _collect( + stores.server(factory), + _context(response="rejected", conversation=conversation), + previous="response-1" if chain else None, + ) + error = _failure(rejected) + assert ("authorized pending" if invalid in ("unknown", "duplicate") else "Response type mismatch") in error + assert executions == ["hello"] + assert received == [] + source_save.assert_not_awaited() + session_save.assert_not_awaited() + unchanged = await storage.get_latest(workflow_name="typed-functional-batch") + assert unchanged is not None + assert json.dumps(unchanged.to_dict(), default=lambda value: value.to_dict(), sort_keys=True) == checkpoint_before + unchanged_session = await sessions.get(source_id) + assert unchanged_session is not None + assert unchanged_session.state == state_before + if chain: + assert await sessions.get("rejected") is None + assert not await stores.checkpoints[("alice", "rejected")].list_checkpoint_ids( + workflow_name="typed-functional-batch" + ) + + converted.return_value = [ + Message( + "tool", + [ + Content.from_function_result("text", result="accepted"), + Content("function_approval_response", id="decision", approved=False), + ], + ) + ] + resumed = await _collect( + stores.server(factory), + _context(response="retry", conversation=conversation), + previous="response-1" if chain else None, + ) + assert _types(resumed)[-1] == "response.completed", resumed + assert _text(resumed) == "accepted:False" + assert received == ["accepted", False] + assert executions == ["hello", "hello"] + + +@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) +@pytest.mark.parametrize("decision", [False, True]) +async def test_functional_response_http_rejects_string_for_bool_then_accepts_approval( + stream: bool, decision: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + stores = _Stores() + server = stores.server(lambda: _pending_bool.build().as_agent(), store=InMemoryResponseProvider()) + transport = httpx.ASGITransport(app=server) + + async def post(client: httpx.AsyncClient, items: Any) -> dict[str, Any]: + response = await client.post( + "/responses", + json={"model": "test-model", "input": items, "conversation": "conversation", "stream": stream}, + ) + assert response.status_code == 200, response.text + if not stream: + return response.json() + events = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] + terminal = [event for event in events if event["type"] in ("response.completed", "response.failed")] + assert len(terminal) == 1, events + return terminal[0]["response"] + + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + first = await post(client, "private") + assert first["status"] == "completed", first + user = next(iter(stores.sessions)) + storage = stores.checkpoints[(user, "conversation")] + checkpoint = await storage.get_latest(workflow_name="functional-pending-bool") + assert checkpoint is not None + request_id = next(iter(checkpoint.pending_request_info_events)) + approval_id = next(iter(stores.approvals[user].requests)) + session_save = AsyncMock(wraps=stores.sessions[user].set) + checkpoint_save = AsyncMock(wraps=storage.save) + monkeypatch.setattr(stores.sessions[user], "set", session_save) + monkeypatch.setattr(storage, "save", checkpoint_save) + rejected = await post(client, [{"type": "function_call_output", "call_id": request_id, "output": "false"}]) + assert rejected["status"] == "failed", rejected + assert "Response type mismatch" in rejected["error"]["message"] + session_save.assert_not_awaited() + checkpoint_save.assert_not_awaited() + resumed = await post( + client, [{"type": "mcp_approval_response", "approval_request_id": approval_id, "approve": decision}] + ) + assert resumed["status"] == "completed", resumed + assert [ + part["text"] + for item in resumed["output"] + if item["type"] == "message" + for part in item["content"] + if part["type"] == "output_text" + ] == [f"private:{decision}"] + + +@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) +async def test_functional_response_http_accepts_string_result(stream: bool) -> None: + stores = _Stores() + server = stores.server(lambda: _pending_string.build().as_agent(), store=InMemoryResponseProvider()) + transport = httpx.ASGITransport(app=server) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + first = await client.post("/responses", json={"input": "private", "conversation": "conversation"}) + assert first.status_code == 200 + assert first.json()["status"] == "completed" + user = next(iter(stores.sessions)) + checkpoint = await stores.checkpoints[(user, "conversation")].get_latest( + workflow_name="functional-pending-string" + ) + assert checkpoint is not None + request_id = next(iter(checkpoint.pending_request_info_events)) + response = await client.post( + "/responses", + json={ + "input": [{"type": "function_call_output", "call_id": request_id, "output": "accepted"}], + "conversation": "conversation", + "stream": stream, + }, + ) + assert response.status_code == 200, response.text + if stream: + events = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] + assert _types(events)[-1] == "response.completed", events + assert _text(events) == "private:accepted" + result = events[-1]["response"] + else: + result = response.json() + assert result["status"] == "completed", result + assert result["output"][0]["content"][0]["text"] == "private:accepted" + + async def test_graph_recovery_without_checkpoint_replays_original_input() -> None: stores = _Stores() context = _context("original") @@ -1129,7 +1369,11 @@ async def test_checkpoint_preparation_failure_rejects_older_checkpoint_and_fresh class FailingCheckpoint(_Counter): @handler - async def count_message(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: + async def count_message( + self, + messages: list[Message], + ctx: WorkflowContext[Never, str], # type: ignore[valid-type] + ) -> None: calls.extend(message.text for message in messages) await super().count_message(messages, ctx) @@ -1262,7 +1506,7 @@ async def start(self, messages: list[Message], ctx: WorkflowContext[str, str]) - class _RecoveryEnd(Executor): @handler - async def end(self, text: str, ctx: WorkflowContext[Never, str]) -> None: + async def end(self, text: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] await ctx.yield_output(f"last:{text}") @@ -1376,11 +1620,22 @@ async def test_graph_recovery_selects_latest_or_response_paired_checkpoint( context = _context("must-not-replay") context.is_recovery = True if snapshot_kind != "latest": - context.persisted_response = ( + snapshot = ( next(snapshot for snapshot in snapshots if snapshot.get("output")) if snapshot_kind == "partial" else snapshots[0] ) + # Pair the response with the matching workflow state explicitly, independent of + # how far the background iterator advanced while the host consumed its output. + checkpoints = await storage.list_checkpoints(workflow_name="recovery") + paired_checkpoint = next( + checkpoint + for checkpoint in checkpoints + if checkpoint.iteration_count == (1 if snapshot_kind == "partial" else 0) + ) + saved_response = ResponseEventStream(response=snapshot) + saved_response.internal_metadata["_last_checkpoint_id"] = paired_checkpoint.checkpoint_id + context.persisted_response = saved_response.checkpoint().response events = await _collect(stores.server(_recovery_graph, options=options), context) assert _types(events)[-1] == "response.completed", events assert "must-not-replay" not in _text(events) @@ -1388,6 +1643,11 @@ async def test_graph_recovery_selects_latest_or_response_paired_checkpoint( if snapshot_kind != "latest": assert load.await_args is not None assert load.await_args.args[0] != latest.checkpoint_id + assert context.persisted_response is not None + paired_checkpoint_id = ResponseEventStream(response=context.persisted_response).internal_metadata[ + "_last_checkpoint_id" + ] + load.assert_awaited_once_with(paired_checkpoint_id) assert "original" in _text(events) if snapshot_kind == "partial": assert _text(events) == "last:original" diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 6b7f739afca..6449c818ce5 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -90,12 +90,12 @@ async def test_invocations_namespaces_cannot_overlap_responses_records(is_hosted "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", new=AsyncMock(return_value=store), ) as get_or_create: - for provider in (CheckpointStoreProvider(), _InvocationsCheckpointStoreProvider()): - await provider.get_store(config=config, context_id="same-id", platform_context=context).save( + for checkpoint_provider in (CheckpointStoreProvider(), _InvocationsCheckpointStoreProvider()): + await checkpoint_provider.get_store(config=config, context_id="same-id", platform_context=context).save( _checkpoint("same-checkpoint") ) - for provider in (AgentSessionStoreProvider(), _InvocationsAgentSessionStoreProvider()): - await provider.get_store(config=config, platform_context=context).set("same-id", AgentSession()) + for session_provider in (AgentSessionStoreProvider(), _InvocationsAgentSessionStoreProvider()): + await session_provider.get_store(config=config, platform_context=context).set("same-id", AgentSession()) assert [call.args[0] for call in get_or_create.await_args_list] == [ "checkpoints/same-id", "invocations_checkpoints/same-id", From 2e1de12d475a7003a449d5f5989325f670b95747 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Tue, 15 Sep 2026 15:50:13 +0000 Subject: [PATCH 03/11] fix: align workflow continuation checks Require recorded completion for both Invocations workflow types. Preserve stored state when incomplete continuation is rejected. --- ...-python-foundry-request-agent-factories.md | 6 +- python/packages/foundry_hosting/README.md | 12 ++- .../_invocations.py | 12 +-- .../tests/test_invocations_factory.py | 89 +++++++++++++++++++ 4 files changed, 107 insertions(+), 12 deletions(-) diff --git a/docs/decisions/0040-python-foundry-request-agent-factories.md b/docs/decisions/0040-python-foundry-request-agent-factories.md index 164e32788c3..f1e5933859a 100644 --- a/docs/decisions/0040-python-foundry-request-agent-factories.md +++ b/docs/decisions/0040-python-foundry-request-agent-factories.md @@ -47,8 +47,10 @@ graph workflow restoration. Functional resilient Responses recovery is explicitly unsupported: checkpoints omit buffered output from completed steps, so a hosting adapter cannot restore that output without rerunning application -work. Invocations retains its text-only exchange and rejects pending or interrupted functional -continuations; a new message after clean completion is supported. +work. Invocations retains its text-only exchange and rejects pending or interrupted continuations +for both graph and functional workflows, before executing checkpointed work or a new message. +A new message after clean completion is supported. The persisted completion record tracks host +execution, not acknowledgment of HTTP delivery. The host manages an agent's exposed async context manager. It does not recursively discover resources inside executors or closures. Applications must construct fresh mutable runtime objects and give diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index f070aaa46da..dd4aef1c2e2 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -214,15 +214,19 @@ do not use this persistence provider. Changing that legacy retention behavior is ### Invocations workflows `InvocationsHostServer(agent_factory=...)` persists workflow checkpoints for the platform user and invocation session. -A subsequent request with the same authorized session can restore its graph workflow in a new agent instance, including -after recreating the host. Ordinary-agent instance callers retain their existing in-memory session behavior. +After a completed invocation, a subsequent request with the same authorized session can restore its graph workflow +in a new agent instance, including after recreating the host. Ordinary-agent instance callers retain their existing +in-memory session behavior. The wire format remains unchanged: requests use `message` and `stream`, and responses contain text or streamed text. This helper does not add a structured exchange for external workflow approvals or pending requests. A plain text message is not an approval response. Use Responses when callers need that structured exchange. -Functional workflows accept a new message after a completed invocation, but pending or interrupted functional -continuation is rejected. They do not acquire graph-workflow recovery semantics by being passed through a factory. +Both graph and functional workflows reject continuation when the previous invocation is pending or interrupted. +A saved checkpoint alone does not authorize resuming incomplete work: the host requires its persisted completion +record before restoring a graph checkpoint or accepting a new message. Rejection preserves the existing stored state. +After clean completion, graph workflows restore their state and functional workflows start the new message without +replaying the previous input. The completion record describes host execution, not acknowledgment of HTTP delivery. Requests updating the same factory scope are serialized within one host; independent scopes can execute concurrently. This is not a distributed lock across multiple host processes. Invocations does not automatically recover an 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 c3cc7e1d284..ae3a60e5af6 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -95,8 +95,8 @@ def __init__( Responses contain plain text; "stream": true streams text as text/event-stream. Workflow factories must use stable workflow names and executor IDs across requests. The text-only exchange cannot answer pending workflow requests or approvals. - Functional workflows support fresh messages after clean completion, but not - pending or interrupted continuation. Such continuation fails explicitly. + Both graph and functional workflows support fresh messages after clean completion, + but not pending or interrupted continuation. Such continuation fails explicitly. Ordinary factory sessions are persisted after each run, including interrupted runs, rather than retained in this host. Ordinary agent instances retain sessions in memory. Factories own cleanup of nested resources not exposed by an async context manager. @@ -201,11 +201,11 @@ async def _workflow_session( raise RuntimeError("The existing Invocations workflow checkpoint is missing its required agent session.") if marker is not None and marker.get("checkpoint_failed"): raise RuntimeError("The previous Invocations workflow run has incomplete checkpoint persistence.") - if isinstance(agent, FunctionalWorkflowAgent) and marker is not None and marker.get("completed") is not True: - raise RuntimeError("Invocations cannot continue a pending or interrupted functional workflow.") + if isinstance(agent, WorkflowAgent) and checkpoint is not None and checkpoint.pending_request_info_events: + raise RuntimeError(_PENDING_REQUEST_ERROR) + if marker is not None and marker.get("completed") is not True: + raise RuntimeError(f"Invocations cannot continue a pending or interrupted {kind} workflow.") if isinstance(agent, WorkflowAgent) and checkpoint is not None: - if checkpoint.pending_request_info_events: - raise RuntimeError(_PENDING_REQUEST_ERROR) await agent.workflow.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) if storage.save_error is not None: raise storage.save_error diff --git a/python/packages/foundry_hosting/tests/test_invocations_factory.py b/python/packages/foundry_hosting/tests/test_invocations_factory.py index a885c348415..230b688a53a 100644 --- a/python/packages/foundry_hosting/tests/test_invocations_factory.py +++ b/python/packages/foundry_hosting/tests/test_invocations_factory.py @@ -950,6 +950,95 @@ def factory() -> WorkflowAgent: assert "unrelated" not in transcripts[-1] +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("cancel", [False, True]) +@pytest.mark.parametrize("recreate_host", [False, True]) +@pytest.mark.parametrize("functional", [False, True]) +async def test_interruption_with_checkpoint_rejects_new_message( + stream: bool, cancel: bool, recreate_host: bool, functional: bool +) -> None: + stores = _Stores() + calls: list[str] = [] + started = asyncio.Event() + interrupt = True + workflow_name = "interrupted-workflow" + + async def finish_message(message: str) -> str: + calls.append(f"finish:{message}") + if interrupt: + started.set() + if cancel: + await asyncio.Event().wait() + raise RuntimeError("failed after checkpoint") + return message + + class Start(Executor): + @handler + async def start(self, messages: list[Message], ctx: WorkflowContext[str]) -> None: + calls.append(f"start:{messages[0].text}") + await ctx.send_message(messages[0].text) + + class Finish(Executor): + @handler + async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + await ctx.yield_output(await finish_message(message)) + + @step + async def saved_step(message: str) -> str: + calls.append(f"start:{message}") + return message + + @workflow(name=workflow_name) + async def interrupted(messages: str | list[str]) -> str: + message = messages if isinstance(messages, str) else messages[0] + return await finish_message(await saved_step(message)) + + def factory() -> WorkflowAgent | FunctionalWorkflowAgent: + if functional: + return interrupted.build().as_agent() + start = Start(id="start") + finish = Finish(id="finish") + return WorkflowBuilder(name=workflow_name, start_executor=start).add_edge(start, finish).build().as_agent() + + server = stores.server(factory) + if cancel: + task = asyncio.create_task(_invoke(server, "original", stream=stream)) + try: + await asyncio.wait_for(started.wait(), timeout=5) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + with pytest.raises(RuntimeError, match="failed after checkpoint"): + await _invoke(server, "original", stream=stream) + interrupt = False + + storage_id, storage = next(iter(stores.checkpoints.items())) + checkpoint = await storage.get_latest(workflow_name=workflow_name) + assert checkpoint is not None + if not functional: + assert checkpoint.iteration_count == 1 + assert checkpoint.messages + session = await stores.sessions.get(storage_id) + assert session is not None + assert session.state["_foundry_invocations_workflow"]["completed"] is False + saved_session = session.to_dict() + checkpoint_ids = await storage.list_checkpoint_ids(workflow_name=workflow_name) + + next_server = stores.server(factory) if recreate_host else server + kind = "functional" if functional else "graph" + with pytest.raises(RuntimeError, match=f"pending or interrupted {kind}"): + await _invoke(next_server, "must-not-run", stream=stream) + assert calls == ["start:original", "finish:original"] + assert await storage.list_checkpoint_ids(workflow_name=workflow_name) == checkpoint_ids + session = await stores.sessions.get(storage_id) + assert session is not None + assert session.to_dict() == saved_session + assert not server._scope_locks._entries + assert not next_server._scope_locks._entries + + @pytest.mark.parametrize("stream", [False, True]) async def test_functional_interruption_with_checkpoint_rejects_new_message(stream: bool) -> None: stores = _Stores() From 0db17c0e9d8e86662c1972ab7033c7605673318a Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Wed, 16 Sep 2026 11:34:19 +0000 Subject: [PATCH 04/11] refactor: simplify hosting checkpoint handling Keep checkpoint error policy with the workflow runtime. Add HTTP coverage for request-scoped factory behavior. --- ...-python-foundry-request-agent-factories.md | 7 +- python/packages/foundry_hosting/README.md | 9 +- .../_invocations.py | 19 +-- .../_responses.py | 38 +---- .../_state_store.py | 51 +----- .../tests/test_invocations_factory.py | 100 +++++++++--- .../tests/test_responses_factory.py | 150 +++++++++++++++--- 7 files changed, 230 insertions(+), 144 deletions(-) diff --git a/docs/decisions/0040-python-foundry-request-agent-factories.md b/docs/decisions/0040-python-foundry-request-agent-factories.md index f1e5933859a..c584c9a29a4 100644 --- a/docs/decisions/0040-python-foundry-request-agent-factories.md +++ b/docs/decisions/0040-python-foundry-request-agent-factories.md @@ -50,7 +50,12 @@ from completed steps, so a hosting adapter cannot restore that output without re work. Invocations retains its text-only exchange and rejects pending or interrupted continuations for both graph and functional workflows, before executing checkpointed work or a new message. A new message after clean completion is supported. The persisted completion record tracks host -execution, not acknowledgment of HTTP delivery. +execution, not checkpoint durability or acknowledgment of HTTP delivery. + +The hosts pass checkpoint storage to the workflow without changing its error policy. Graph workflows +can log checkpoint creation failures and continue; propagated runtime errors still fail the invocation. +The factory feature does not add a stricter persistence guarantee or intercept internal runner methods. +Continuation can therefore use an older saved checkpoint, or fail when the required checkpoint is missing. The host manages an agent's exposed async context manager. It does not recursively discover resources inside executors or closures. Applications must construct fresh mutable runtime objects and give diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index dd4aef1c2e2..2f999dee0b6 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -179,6 +179,12 @@ to the current user and conversation or response chain. A `previous_response_id` fails; a conversation with prior history but no workflow checkpoint also fails rather than silently starting over. A new conversation with no history can start without a checkpoint. +Both hosts use the configured checkpoint storage directly and retain the workflow runtime's checkpoint-error policy. +The graph runtime logs checkpoint creation failures and may complete without saving its latest state. Hosting does +not intercept runner methods or turn those logged failures into invocation failures. A completed invocation is not +proof of a successful checkpoint save: a later continuation may restore an older checkpoint, or fail if none exists. +Exceptions propagated by the runtime, including functional workflow checkpoint errors, still fail the invocation. + For resilient background Responses using a graph workflow, recovery uses the checkpoint associated with the persisted response output. If no response checkpoint was recorded, it can use the latest workflow checkpoint. If execution stopped before any workflow checkpoint was saved, recovery replays the original input using a fresh factory-created workflow. @@ -226,7 +232,8 @@ Both graph and functional workflows reject continuation when the previous invoca A saved checkpoint alone does not authorize resuming incomplete work: the host requires its persisted completion record before restoring a graph checkpoint or accepting a new message. Rejection preserves the existing stored state. After clean completion, graph workflows restore their state and functional workflows start the new message without -replaying the previous input. The completion record describes host execution, not acknowledgment of HTTP delivery. +replaying the previous input. The completion record describes host execution, not checkpoint durability or +acknowledgment of HTTP delivery. Requests updating the same factory scope are serialized within one host; independent scopes can execute concurrently. This is not a distributed lock across multiple host processes. Invocations does not automatically recover an 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 ae3a60e5af6..a89e6cb37b2 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -39,7 +39,6 @@ from ._state_store import ( ContextScopedStoreProvider, StoreProvider, - _CheckpointStorageWithErrors, # pyright: ignore[reportPrivateUsage] _InvocationsAgentSessionStoreProvider, # pyright: ignore[reportPrivateUsage] _InvocationsCheckpointStoreProvider, # pyright: ignore[reportPrivateUsage] ) @@ -175,16 +174,12 @@ async def _workflow_session( self, agent: WorkflowAgent | FunctionalWorkflowAgent, storage_id: str ) -> AsyncGenerator[tuple[AgentSession, CheckpointStorage]]: context = get_request_context() - storage = _CheckpointStorageWithErrors( - self._checkpoint_storage_provider.get_store( - config=self.config, context_id=storage_id, platform_context=context - ) + storage = self._checkpoint_storage_provider.get_store( + config=self.config, context_id=storage_id, platform_context=context ) sessions = self._agent_session_storage_provider.get_store(config=self.config, platform_context=context) workflow = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow # pyright: ignore[reportPrivateUsage] kind = "graph" if isinstance(agent, WorkflowAgent) else "functional" - if isinstance(agent, WorkflowAgent): - storage.observe_workflow(agent) session = await sessions.get(storage_id) saved_marker = session.state.get(_WORKFLOW_STATE_KEY) if session is not None else None marker = cast(dict[str, Any], saved_marker) if isinstance(saved_marker, dict) else None @@ -199,16 +194,12 @@ async def _workflow_session( raise RuntimeError("The existing Invocations workflow session is missing its required checkpoint.") if checkpoint is not None and session is None: raise RuntimeError("The existing Invocations workflow checkpoint is missing its required agent session.") - if marker is not None and marker.get("checkpoint_failed"): - raise RuntimeError("The previous Invocations workflow run has incomplete checkpoint persistence.") if isinstance(agent, WorkflowAgent) and checkpoint is not None and checkpoint.pending_request_info_events: raise RuntimeError(_PENDING_REQUEST_ERROR) if marker is not None and marker.get("completed") is not True: raise RuntimeError(f"Invocations cannot continue a pending or interrupted {kind} workflow.") if isinstance(agent, WorkflowAgent) and checkpoint is not None: await agent.workflow.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) - if storage.save_error is not None: - raise storage.save_error if agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS: raise RuntimeError(_PENDING_REQUEST_ERROR) if session is None: @@ -219,8 +210,6 @@ async def _workflow_session( await sessions.set(storage_id, session) try: yield session, storage - if storage.save_error is not None: - raise storage.save_error pending = ( agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS if isinstance(agent, WorkflowAgent) @@ -228,12 +217,8 @@ async def _workflow_session( ) if pending: raise RuntimeError(_PENDING_REQUEST_ERROR) - if await storage.get_latest(workflow_name=workflow.name) is None: - raise RuntimeError("The Invocations workflow did not persist a required checkpoint.") marker["completed"] = True finally: - if storage.save_error is not None: - marker["checkpoint_failed"] = True with CancelScope(shield=True): await sessions.set(storage_id, session) 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 e81592b5835..e43cf17d3f3 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -98,7 +98,6 @@ FunctionApprovalStore, FunctionApprovalStoreProvider, StoreProvider, - _CheckpointStorageWithErrors, # pyright: ignore[reportPrivateUsage] ) logger = logging.getLogger(__name__) @@ -1074,8 +1073,6 @@ async def _handle_inner_workflow( marker is None or marker.get("name") != workflow.name or marker.get("kind") != kind ): raise RuntimeError("The stored Responses workflow name or kind does not match the factory result.") - if marker is not None and marker.get("checkpoint_failed"): - raise RuntimeError("The previous Responses workflow run has incomplete checkpoint persistence.") had_session = session is not None if session is None: session = AgentSession() @@ -1117,7 +1114,6 @@ async def run_workflow(stream_factory: _WorkflowRunFactory) -> AsyncGenerator[Ag agent, session if agent.context_providers else None, had_session, - state_marker, run_workflow, ) async with aclosing(inner): @@ -1155,7 +1151,6 @@ async def _handle_graph_workflow( agent: WorkflowAgent, session: AgentSession | None, had_session: bool, - state_marker: dict[str, Any], run_workflow: _WorkflowRun, ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: """Handle the creation of a response for a workflow agent.""" @@ -1178,14 +1173,11 @@ async def _handle_graph_workflow( # workflow from the last checkpoint. checkpoint_save_id = context.conversation_id or context.response_id _validate_checkpoint_context_id(checkpoint_save_id) - checkpoint_storage = _CheckpointStorageWithErrors( - self._checkpoint_storage_provider.get_store( - config=self.config, - context_id=checkpoint_save_id, - platform_context=request_context, - ) + checkpoint_storage = self._checkpoint_storage_provider.get_store( + config=self.config, + context_id=checkpoint_save_id, + platform_context=request_context, ) - checkpoint_storage.observe_workflow(agent) if context.is_recovery: if not self._resilient_background: @@ -1263,12 +1255,10 @@ async def _handle_graph_workflow( if checkpoint_load_id is not None: _validate_checkpoint_context_id(checkpoint_load_id) if checkpoint_load_id != checkpoint_save_id: - restore_checkpoint_storage = _CheckpointStorageWithErrors( - self._checkpoint_storage_provider.get_store( - config=self.config, - context_id=checkpoint_load_id, - platform_context=request_context, - ) + restore_checkpoint_storage = self._checkpoint_storage_provider.get_store( + config=self.config, + context_id=checkpoint_load_id, + platform_context=request_context, ) latest_checkpoint = await restore_checkpoint_storage.get_latest(workflow_name=agent.workflow.name) @@ -1314,12 +1304,6 @@ async def _handle_graph_workflow( async with aclosing(restore_iter): async for _ in restore_iter: pass - if restore_checkpoint_storage.save_error is not None: - state_marker["checkpoint_failed"] = True - raise restore_checkpoint_storage.save_error - if checkpoint_storage.save_error is not None: - state_marker["checkpoint_failed"] = True - raise checkpoint_storage.save_error if restore_iter.signalled: if context.shutdown.is_set(): await context.exit_for_recovery() @@ -1343,9 +1327,6 @@ async def _handle_graph_workflow( main_iter = _SignalledIterator(run_stream, context.shutdown, cancellation_signal) async with aclosing(main_iter): async for update in main_iter: - if checkpoint_storage.save_error is not None: - state_marker["checkpoint_failed"] = True - raise checkpoint_storage.save_error if self._resilient_background: latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=agent.workflow.name) if ( @@ -1374,9 +1355,6 @@ async def _handle_graph_workflow( content, message_id=update.message_id, approval_storage=approval_storage ): yield event - if checkpoint_storage.save_error is not None: - state_marker["checkpoint_failed"] = True - raise checkpoint_storage.save_error # Cancellation needs no extra action here (the loop above already stopped); shutdown # does, but only if it's what actually stopped the loop, not a natural completion. if main_iter.signalled and context.shutdown.is_set(): 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 aca8f4ca57f..fd77bfa0169 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 @@ -1,13 +1,8 @@ # Copyright (c) Microsoft. All rights reserved. -from __future__ import annotations - from abc import ABC, abstractmethod -from collections.abc import Callable from datetime import datetime -from functools import wraps -from types import CoroutineType -from typing import Any, Generic, ParamSpec, Protocol, TypeVar +from typing import Generic, Protocol, TypeVar from agent_framework import ( AgentSession, @@ -15,7 +10,6 @@ CheckpointStorage, Content, SessionStore, - WorkflowAgent, WorkflowCheckpoint, WorkflowCheckpointException, ) @@ -23,49 +17,6 @@ from azure.ai.agentserver.core.storage import FoundryStateStore, FoundryStorageConflictError StoreT = TypeVar("StoreT") -ResultT = TypeVar("ResultT") -ParametersT = ParamSpec("ParametersT") - - -class _CheckpointStorageWithErrors: # pyright: ignore[reportUnusedClass] - """Remember checkpoint write errors even when graph execution logs and ignores them.""" - - def __init__(self, storage: CheckpointStorage) -> None: - self._storage = storage - self.save_error: Exception | None = None - self.load = storage.load - self.list_checkpoints = storage.list_checkpoints - self.delete = storage.delete - self.get_latest = storage.get_latest - self.list_checkpoint_ids = storage.list_checkpoint_ids - - async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: - try: - return await self._storage.save(checkpoint) - except Exception as exc: - self.save_error = exc - raise - - def observe_workflow(self, agent: WorkflowAgent) -> None: - """Observe both stages the graph runner suppresses before and during checkpoint writes.""" - runner = agent.workflow._runner # pyright: ignore[reportPrivateUsage] - runner._prepare_checkpoint_state = self._observe( # pyright: ignore[reportPrivateUsage] - runner._prepare_checkpoint_state # pyright: ignore[reportPrivateUsage] - ) - runner.context.create_checkpoint = self._observe(runner.context.create_checkpoint) - - def _observe( - self, operation: Callable[ParametersT, CoroutineType[Any, Any, ResultT]] - ) -> Callable[ParametersT, CoroutineType[Any, Any, ResultT]]: - @wraps(operation) - async def observed(*args: ParametersT.args, **kwargs: ParametersT.kwargs) -> ResultT: - try: - return await operation(*args, **kwargs) - except Exception as exc: - self.save_error = exc - raise - - return observed class StoreProvider(ABC, Generic[StoreT]): diff --git a/python/packages/foundry_hosting/tests/test_invocations_factory.py b/python/packages/foundry_hosting/tests/test_invocations_factory.py index 230b688a53a..0f43db18975 100644 --- a/python/packages/foundry_hosting/tests/test_invocations_factory.py +++ b/python/packages/foundry_hosting/tests/test_invocations_factory.py @@ -9,6 +9,7 @@ from typing import Any, cast from unittest.mock import AsyncMock, MagicMock +import httpx import pytest from agent_framework import ( Agent, @@ -29,7 +30,6 @@ SessionStore, WorkflowAgent, WorkflowBuilder, - WorkflowCheckpointException, WorkflowContext, WorkflowEvent, handler, @@ -607,6 +607,48 @@ async def response() -> ChatResponse: return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("use_agent_executor", [False, True]) +async def test_http_sessions_use_independent_workflow_state(stream: bool, use_agent_executor: bool) -> None: + stores = _Stores() + transcripts: list[list[str]] = [] + factory_scopes: list[tuple[str | None, str | None]] = [] + + class Remember(Executor): + @handler + async def remember(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + previous = ctx.get_state("previous") + ctx.set_state("previous", messages[0].text) + await ctx.yield_output(json.dumps(previous)) + + def factory() -> WorkflowAgent: + context = get_request_context() + factory_scopes.append((context.user_id, context.session_id)) + executor = ( + AgentExecutor(Agent(client=_TranscriptClient(transcripts), name="inner"), id="inner") + if use_agent_executor + else Remember(id="remember") + ) + return WorkflowBuilder(name="http-state", start_executor=executor).build().as_agent() + + server = stores.server(factory) + server.config.is_hosted = True + scopes = [("user-a", "session-a"), ("user-b", "session-b"), ("user-a", "session-c")] + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: + for (user, session), text in zip(scopes, ("first", "second", "independent"), strict=True): + response = await client.post( + "/invocations", + json={"message": text, "stream": stream}, + params={"agent_session_id": session}, + headers={"x-agent-user-id": user, "x-agent-foundry-call-id": f"call-{session}"}, + ) + assert response.status_code == 200, response.text + assert response.text == ("recorded" if use_agent_executor else "null") + if use_agent_executor: + assert transcripts[-1] == [text] + assert factory_scopes == scopes + + class _SnapshotSessions(SessionStore): def __init__(self) -> None: self.snapshots: dict[str, str] = {} @@ -1095,7 +1137,11 @@ async def test_invalid_request_does_not_construct_agent(monkeypatch: pytest.Monk factory.assert_not_called() -async def test_checkpoint_failure_closes_workflow_owner_and_preserves_attempt(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_checkpoint_save_failure_preserves_runtime_behavior_and_closes_owner( + functional: bool, stream: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: events: list[str] = [] class OwnedWorkflow(WorkflowAgent): @@ -1106,26 +1152,41 @@ async def __aenter__(self) -> Self: async def __aexit__(self, *args: Any) -> None: events.append("exit") + class OwnedFunctional(FunctionalWorkflowAgent): + async def __aenter__(self) -> Self: + events.append("enter") + return self + + async def __aexit__(self, *args: Any) -> None: + events.append("exit") + stores = _Stores() storage = InMemoryCheckpointStorage() monkeypatch.setattr(storage, "save", AsyncMock(side_effect=RuntimeError("checkpoint failed"))) stores.checkpoint_provider.get_store.side_effect = None stores.checkpoint_provider.get_store.return_value = storage - server = stores.server(lambda: OwnedWorkflow(_graph().workflow)) - with pytest.raises(RuntimeError, match="checkpoint failed"): - await _invoke(server) + server = stores.server( + lambda: OwnedFunctional(_functional.build()) if functional else OwnedWorkflow(_graph().workflow) + ) + if functional: + with pytest.raises(RuntimeError, match="checkpoint failed"): + await _invoke(server, stream=stream) + else: + assert await _invoke(server, stream=stream) == "1:hello" + assert "does not fail the workflow run" in caplog.text with pytest.raises(RuntimeError, match="missing its required checkpoint"): - await _invoke(server) + await _invoke(server, stream=stream) assert events == ["enter", "exit", "enter", "exit"] assert not server._scope_locks._entries @pytest.mark.parametrize("stream", [False, True]) -async def test_checkpoint_preparation_failure_rejects_older_checkpoint_and_fresh_host_retry(stream: bool) -> None: +async def test_checkpoint_preparation_failure_preserves_graph_runtime_behavior( + stream: bool, caplog: pytest.LogCaptureFixture +) -> None: stores = _Stores() calls: list[str] = [] snapshots: list[int] = [] - agents: list[WorkflowAgent] = [] class FailingCheckpoint(_Counter): @handler @@ -1144,13 +1205,11 @@ async def on_checkpoint_save(self) -> dict[str, Any]: return await super().on_checkpoint_save() def factory() -> WorkflowAgent: - agent = WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() - agents.append(agent) - return agent + return WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() - first = stores.server(factory) - with pytest.raises(WorkflowCheckpointException, match="Executor counter on_checkpoint_save failed"): - await _invoke(first, stream=stream) + server = stores.server(factory) + assert await _invoke(server, stream=stream) == "1:hello" + assert "does not fail the workflow run" in caplog.text storage_id, storage = next(iter(stores.checkpoints.items())) checkpoint = await storage.get_latest(workflow_name="checkpoint-preparation") assert checkpoint is not None @@ -1158,14 +1217,7 @@ def factory() -> WorkflowAgent: assert 0 in snapshots and 1 in snapshots session = await stores.sessions.get(storage_id) assert session is not None - assert session.state["_foundry_invocations_workflow"]["checkpoint_failed"] is True - - second = stores.server(factory) - with pytest.raises(RuntimeError, match="incomplete checkpoint persistence"): - await _invoke(second, "must-not-run", stream=stream) + assert session.state["_foundry_invocations_workflow"]["completed"] is True + assert "checkpoint_failed" not in session.state["_foundry_invocations_workflow"] assert calls == ["hello"] - assert len(agents) == 2 - assert agents[0] is not agents[1] - assert agents[0].workflow is not agents[1].workflow - assert not first._scope_locks._entries - assert not second._scope_locks._entries + assert not server._scope_locks._entries diff --git a/python/packages/foundry_hosting/tests/test_responses_factory.py b/python/packages/foundry_hosting/tests/test_responses_factory.py index 3462d2039b7..1981fff1567 100644 --- a/python/packages/foundry_hosting/tests/test_responses_factory.py +++ b/python/packages/foundry_hosting/tests/test_responses_factory.py @@ -105,6 +105,26 @@ async def _collect( ] +async def _post_http_response( + client: httpx.AsyncClient, items: Any, *, user: str, stream: bool, previous: str | None = None +) -> dict[str, Any]: + payload = {"model": "test-model", "input": items, "stream": stream} + if previous is not None: + payload["previous_response_id"] = previous + response = await client.post( + "/responses", + json=payload, + headers={"x-agent-user-id": user, "x-agent-foundry-call-id": f"call-{user}"}, + ) + assert response.status_code == 200, response.text + if not stream: + return response.json() + events = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] + terminal = [event for event in events if event["type"] in ("response.completed", "response.failed")] + assert len(terminal) == 1, events + return terminal[0]["response"] + + def _types(events: list[Any]) -> list[str]: return [event["type"] for event in events if isinstance(event, Mapping)] @@ -805,6 +825,47 @@ async def updates() -> AsyncIterator[ChatResponseUpdate]: return ResponseStream(updates(), finalizer=ChatResponse.from_updates) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("use_agent_executor", [False, True]) +async def test_fresh_http_requests_use_independent_workflow_state(stream: bool, use_agent_executor: bool) -> None: + stores = _Stores() + transcripts: list[list[str]] = [] + factory_users: list[str | None] = [] + + class Remember(Executor): + @handler + async def remember(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + previous = ctx.get_state("previous") + ctx.set_state("previous", messages[0].text) + await ctx.yield_output(json.dumps(previous)) + + def factory() -> WorkflowAgent: + factory_users.append(get_request_context().user_id) + executor = ( + AgentExecutor(Agent(client=_TranscriptClient(transcripts), name="inner"), id="inner") + if use_agent_executor + else Remember(id="remember") + ) + return WorkflowBuilder(name="http-state", start_executor=executor).build().as_agent() + + server = stores.server(factory, store=InMemoryResponseProvider()) + server.config.is_hosted = True + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: + for user, text in (("user-a", "first"), ("user-b", "second"), ("user-a", "independent")): + body = await _post_http_response(client, text, user=user, stream=stream) + assert body["status"] == "completed", body + assert [ + part["text"] + for item in body["output"] + if item["type"] == "message" + for part in item["content"] + if part["type"] == "output_text" + ] == ["recorded" if use_agent_executor else "null"] + if use_agent_executor: + assert transcripts[-1] == [text] + assert factory_users == ["user-a", "user-b", "user-a"] + + @pytest.mark.parametrize("chain", [False, True]) async def test_agent_executor_transcript_isolates_users_scopes_and_continues_across_hosts(chain: bool) -> None: stores = _Stores() @@ -1361,11 +1422,12 @@ async def consume() -> None: assert events == ["enter", "run", "cleanup started", "cleanup finished", "exit"] -async def test_checkpoint_preparation_failure_rejects_older_checkpoint_and_fresh_host_retry() -> None: +async def test_checkpoint_preparation_failure_preserves_graph_runtime_behavior( + caplog: pytest.LogCaptureFixture, +) -> None: stores = _Stores() calls: list[str] = [] snapshots: list[int] = [] - agents: list[WorkflowAgent] = [] class FailingCheckpoint(_Counter): @handler @@ -1384,12 +1446,13 @@ async def on_checkpoint_save(self) -> dict[str, Any]: return await super().on_checkpoint_save() def factory() -> WorkflowAgent: - agent = WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() - agents.append(agent) - return agent + return WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() - first = stores.server(factory) - assert "Executor counter on_checkpoint_save failed" in _failure(await _collect(first)) + server = stores.server(factory) + result = await _collect(server) + assert _types(result)[-1] == "response.completed" + assert _text(result) == "1:hello" + assert "does not fail the workflow run" in caplog.text storage = stores.checkpoints[("alice", "conversation")] checkpoint = await storage.get_latest(workflow_name="checkpoint-preparation") assert checkpoint is not None @@ -1397,23 +1460,15 @@ def factory() -> WorkflowAgent: assert 0 in snapshots and 1 in snapshots session = await stores.sessions["alice"].get("conversation") assert session is not None - assert session.state["_foundry_responses_workflow"]["checkpoint_failed"] is True - - second = stores.server(factory) - assert "incomplete checkpoint persistence" in _failure( - await _collect(second, _context("must-not-run", response="two")) - ) + assert session.state["_foundry_responses_workflow"]["completed"] is True + assert "checkpoint_failed" not in session.state["_foundry_responses_workflow"] assert calls == ["hello"] - assert len(agents) == 2 - assert agents[0] is not agents[1] - assert agents[0].workflow is not agents[1].workflow - assert not first._scope_locks._entries - assert not second._scope_locks._entries + assert not server._scope_locks._entries @pytest.mark.parametrize("functional", [False, True]) -async def test_checkpoint_save_failure_does_not_report_success_and_closes_owner( - functional: bool, monkeypatch: pytest.MonkeyPatch +async def test_checkpoint_save_failure_preserves_runtime_behavior_and_closes_owner( + functional: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: events: list[str] = [] storage = InMemoryCheckpointStorage() @@ -1445,7 +1500,13 @@ async def __aexit__(self, *args: Any) -> None: assert events == ["enter", "exit"] assert save.await_count > 0 assert not server._scope_locks._entries - assert "checkpoint save failed" in _failure(result) + if functional: + assert "checkpoint save failed" in _failure(result) + else: + assert _types(result)[-1] == "response.completed" + assert _text(result) == "1:hello" + assert "does not fail the workflow run" in caplog.text + assert await storage.get_latest(workflow_name="functional" if functional else "counter") is None async def test_interrupted_functional_step_does_not_restart_with_new_input() -> None: @@ -1472,6 +1533,53 @@ def factory() -> FunctionalWorkflowAgent: assert calls == ["hello"] +@pytest.mark.parametrize("functional", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_fresh_http_request_preserves_another_users_pending_request(functional: bool, stream: bool) -> None: + stores = _Stores() + completions: list[tuple[str | None, str]] = [] + + class Pending(Executor): + @handler + async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + await ctx.request_info(messages[0].text, response_type=str, request_id="answer") + + @response_handler + async def answer(self, original_request: str, response: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] + completions.append((get_request_context().user_id, original_request)) + await ctx.yield_output(f"{original_request}:{response}") + + @workflow(name="http-pending") + async def pending(messages: list[Message], ctx: RunContext) -> str: + answer = await ctx.request_info("answer?", response_type=str, request_id="answer") + completions.append((get_request_context().user_id, messages[0].text)) + return f"{messages[0].text}:{answer}" + + def factory() -> WorkflowAgent | FunctionalWorkflowAgent: + if functional: + return pending.build().as_agent() + return WorkflowBuilder(name="http-pending", start_executor=Pending(id="pending")).build().as_agent() + + server = stores.server(factory, store=InMemoryResponseProvider()) + server.config.is_hosted = True + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: + first = await _post_http_response(client, "first", user="user-a", stream=stream) + assert first["status"] == "completed", first + result = [{"type": "function_call_output", "call_id": "answer", "output": "accepted"}] + await _post_http_response(client, result, user="user-b", stream=stream) + assert completions == [] + resumed = await _post_http_response(client, result, user="user-a", stream=stream, previous=first["id"]) + assert resumed["status"] == "completed", resumed + assert [ + part["text"] + for item in resumed["output"] + if item["type"] == "message" + for part in item["content"] + if part["type"] == "output_text" + ] == ["first:accepted"] + assert completions == [("user-a", "first")] + + @pytest.mark.parametrize("functional", [False, True]) async def test_cross_user_previous_response_cannot_resume_another_users_pending_request(functional: bool) -> None: stores = _Stores() From a23359cf8715a326dfaf2df734a2323e473006df Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:50:32 +0100 Subject: [PATCH 05/11] Python: Simplify request-scoped Foundry agents --- ...-python-foundry-request-agent-factories.md | 77 - .../packages/core/agent_framework/_types.py | 14 + .../agent_framework/_workflows/_functional.py | 14 +- python/packages/core/tests/core/test_types.py | 38 + .../workflow/test_functional_workflow.py | 120 +- python/packages/foundry_hosting/README.md | 156 +- .../_agent_factory.py | 101 - .../_agent_source.py | 28 + .../_invocations.py | 249 +-- .../_responses.py | 611 ++---- .../_state_store.py | 23 +- .../tests/test_agent_factory.py | 80 - .../foundry_hosting/tests/test_invocations.py | 152 +- .../tests/test_invocations_factory.py | 1223 ------------ .../foundry_hosting/tests/test_responses.py | 217 +- .../tests/test_responses_factory.py | 1769 ----------------- .../foundry_hosting/tests/test_state_store.py | 35 +- .../declarative_customer_support/README.md | 11 +- .../declarative_customer_support/main.py | 29 +- .../resilient_long_running_workflow/README.md | 9 +- .../resilient_long_running_workflow/main.py | 42 +- .../responses/workflows/README.md | 11 +- .../responses/workflows/main.py | 25 +- 23 files changed, 619 insertions(+), 4415 deletions(-) delete mode 100644 docs/decisions/0040-python-foundry-request-agent-factories.md delete mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py create mode 100644 python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py delete mode 100644 python/packages/foundry_hosting/tests/test_agent_factory.py delete mode 100644 python/packages/foundry_hosting/tests/test_invocations_factory.py delete mode 100644 python/packages/foundry_hosting/tests/test_responses_factory.py diff --git a/docs/decisions/0040-python-foundry-request-agent-factories.md b/docs/decisions/0040-python-foundry-request-agent-factories.md deleted file mode 100644 index c584c9a29a4..00000000000 --- a/docs/decisions/0040-python-foundry-request-agent-factories.md +++ /dev/null @@ -1,77 +0,0 @@ ---- -status: proposed -contact: RogerBarreto -date: 2026-09-10 -deciders: RogerBarreto ---- - -# Request-scoped agent factories for Python Foundry hosting - -## Context and Problem Statement - -Python Foundry hosts accept an agent instance and separate persisted state by user and conversation. -Some agent implementations also hold mutable execution state outside their `AgentSession`. -A workflow's executors, pending requests, and internal conversation therefore need their own lifetime, -independent of the server object's lifetime. - -## Decision Drivers - -- Preserve ordinary-agent instance callers and existing protocol formats. -- Create independent workflow execution objects for each request. -- Continue authorized conversations from stored state, including after process recreation. -- Avoid a core workflow redesign or a cache of live runtimes with expiration policies. - -## Considered Options - -| Option | Benefit | Cost | -| --- | --- | --- | -| Factory per request | Explicit ownership; same construction path for fresh calls and recovery | Applications rebuild mutable runtime objects and manage client ownership | -| Runtime cache per user and session | Avoids rebuilding objects on every call | Requires expiration, eviction, cleanup, concurrency controls, and a separate restart path | -| Move workflow runtime into a new core session type | Separates shared definitions from session state throughout the framework | Larger change to core execution and serialization contracts | -| Continue accepting shared workflow instances | No caller migration | Persisted-state isolation does not separate live workflow state | - -## Decision Outcome - -Add `agent_factory` to `ResponsesHostServer` and `InvocationsHostServer`. It is a zero-argument callable -returning an agent or an awaitable agent. Resolve it within the current platform request context and -keep the result local until execution, persistence, and streaming have finished. - -Retain `agent` for ordinary instances. Require factories for the two built-in workflow-agent types, -recognized by one private Foundry helper. Do not add a public workflow-recognition interface or attempt -to inspect arbitrary application wrappers. - -Responses continues using its conversation/response checkpoint scopes. Invocations adds workflow -checkpoint persistence scoped to user and invocation session without changing its text protocol. -Functional workflows need their own checkpoint adaptation because saved-input replay differs from -graph workflow restoration. - -Functional resilient Responses recovery is explicitly unsupported: checkpoints omit buffered output -from completed steps, so a hosting adapter cannot restore that output without rerunning application -work. Invocations retains its text-only exchange and rejects pending or interrupted continuations -for both graph and functional workflows, before executing checkpointed work or a new message. -A new message after clean completion is supported. The persisted completion record tracks host -execution, not checkpoint durability or acknowledgment of HTTP delivery. - -The hosts pass checkpoint storage to the workflow without changing its error policy. Graph workflows -can log checkpoint creation failures and continue; propagated runtime errors still fail the invocation. -The factory feature does not add a stricter persistence guarantee or intercept internal runner methods. -Continuation can therefore use an older saved checkpoint, or fail when the required checkpoint is missing. - -The host manages an agent's exposed async context manager. It does not recursively discover resources -inside executors or closures. Applications must construct fresh mutable runtime objects and give -shared clients a lifetime that outlasts every request using them. - -Serialize updates to the same workflow scope within a host process. Do not describe these locks as -distributed coordination or checkpoint replay as exactly-once execution. - -## Consequences - -Workflow callers migrate from `Host(workflow_agent)` to `Host(agent_factory=create_agent)`. -The factory must construct new mutable objects, not return the same instance or reuse stateful executors. -This is the factory implementer's responsibility. The host calls the factory per request but does not inspect -object identities, recursively check captured state, or require weak-reference support. -Stable workflow names, executor IDs, and serialization registrations are necessary for continuation. - -Ordinary-agent instance callers keep their existing lifecycle. Factories do not automatically save -arbitrary custom agent fields, and the Invocations text protocol still does not provide a complete -external approval exchange. A core session redesign remains a separate possible improvement. diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 990c55c4ae9..db2600d5589 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -3520,6 +3520,20 @@ async def __anext__(self) -> UpdateT: update = await update return await self._record_update(update) + async def aclose(self) -> None: + """Close the active iterator and run cleanup hooks. + + This method is idempotent and also closes nested ``ResponseStream`` wrappers. + """ + try: + if self._iterator is not None: + close = getattr(self._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..eab19003e6f 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -4192,6 +4192,44 @@ async def async_hook(update: ChatResponseUpdate) -> ChatResponseUpdate: class TestResponseStreamCleanupHooks: """Tests for cleanup hooks (after stream consumption, before finalizer).""" + async def test_aclose_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(updates(), cleanup_hooks=[lambda: events.append("cleanup")]) + await anext(stream) + + await stream.aclose() + await stream.aclose() + + assert events == ["iterator", "cleanup"] + + async def test_aclose_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(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.aclose() + + 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 9e467988e66..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, @@ -717,118 +718,6 @@ async def wf(x: int, ctx: RunContext) -> str: class TestCheckpointing: - async def test_fresh_agent_checkpoint_replay_cannot_recover_buffered_step_output(self) -> None: - """A completed step's saved result does not preserve its buffered output events.""" - storage = InMemoryCheckpointStorage() - step_saved = asyncio.Event() - finish = asyncio.Event() - calls = 0 - - @step - async def emit_step(text: str) -> str: - nonlocal calls - calls += 1 - ctx = get_run_context() - assert ctx is not None - await ctx.add_event(WorkflowEvent("output", executor_id="emit_step", data=f"step:{text}")) - return text - - @workflow - async def buffered_workflow(text: str) -> str: - result = await emit_step(text) - step_saved.set() - await finish.wait() - return f"final:{result}" - - original = buffered_workflow.build().as_agent() - received: list[str] = [] - - async def consume_original() -> None: - async for update in original.run("original", stream=True, checkpoint_storage=storage): - if update.text: - received.append(update.text) - - consumer = asyncio.create_task(consume_original()) - try: - await asyncio.wait_for(step_saved.wait(), timeout=5) - checkpoint = await storage.get_latest(workflow_name=buffered_workflow.name) - assert checkpoint is not None - assert received == [] - finally: - consumer.cancel() - with pytest.raises(asyncio.CancelledError): - await consumer - - finish.set() - recovered = buffered_workflow.build().as_agent() - result = await recovered.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) - - assert calls == 1 - assert result.text == "final:original" - assert "step:original" not in result.text - - complete = await buffered_workflow.build().as_agent().run("original", checkpoint_storage=storage) - assert "step:original" in complete.text - assert "final:original" in complete.text - assert calls == 2 - - async def test_fresh_agent_completed_checkpoint_replays_old_input_instead_of_starting_new_turn(self) -> None: - calls: list[str] = [] - storage = InMemoryCheckpointStorage() - - @step - async def record(text: str) -> str: - calls.append(text) - return text - - @workflow - async def turns(text: str) -> str: - return await record(text) - - original = turns.build().as_agent() - assert (await original.run("first", checkpoint_storage=storage)).text == "first" - checkpoint = await storage.get_latest(workflow_name=turns.name) - assert checkpoint is not None - - restored = turns.build().as_agent() - replay = await restored.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) - assert replay.text == "first" - assert calls == ["first"] - with pytest.raises(ValueError, match="message.*checkpoint_id"): - await restored.run("second", checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) - - assert (await original.run("second", checkpoint_storage=storage)).text == "second" - assert calls == ["first", "second"] - - async def test_fresh_agent_resumes_pending_request_from_checkpoint_without_new_message(self) -> None: - storage = InMemoryCheckpointStorage() - - @workflow - async def review(text: str, ctx: RunContext) -> str: - answer = await ctx.request_info(text, response_type=str, request_id="review-request") - return f"{text}:{answer}" - - pending_agent = review.build().as_agent() - await pending_agent.run("original", checkpoint_storage=storage) - checkpoint = await storage.get_latest(workflow_name=review.name) - assert checkpoint is not None - assert "review-request" in checkpoint.pending_request_info_events - - restored = review.build().as_agent() - response = await restored.run( - checkpoint_id=checkpoint.checkpoint_id, - checkpoint_storage=storage, - responses={"review-request": "accepted"}, - ) - assert response.text == "original:accepted" - completed = next( - saved - for saved in await storage.list_checkpoints(workflow_name=review.name) - if saved.previous_checkpoint_id == checkpoint.checkpoint_id - ) - assert not restored.pending_requests - assert not completed.pending_request_info_events - async def test_checkpoint_save_and_restore(self): storage = InMemoryCheckpointStorage() @@ -1163,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 2f999dee0b6..8bb7c99eb7e 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -4,85 +4,29 @@ This package provides the integration of Agent Framework agents and workflows wi ## Agent instances and factories -Both hosts accept either an ordinary agent instance or an `agent_factory`, but not both. -Existing `ResponsesHostServer(agent)` and `InvocationsHostServer(agent)` calls remain supported for ordinary agents. -Pass `WorkflowAgent` and `FunctionalWorkflowAgent` through a factory instead of passing a live instance: +`ResponsesHostServer` and `InvocationsHostServer` accept an agent instance or a zero-argument callable through the +existing `agent` parameter. The callable may be synchronous or asynchronous and is invoked once for each request: ```python -from agent_framework_foundry_hosting import InvocationsHostServer, ResponsesHostServer - - -def create_agent(): - # Build a new workflow, with new mutable agents and executors, here. - return build_workflow().as_agent() - - -server = ResponsesHostServer(agent_factory=create_agent) -# Alternatively, use the existing Invocations text protocol: -# server = InvocationsHostServer(agent_factory=create_agent) +server = ResponsesHostServer(agent=create_agent) ``` -A factory takes no arguments and can return an agent directly or await its construction: +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 -async def create_agent(): - configuration = await load_configuration() - return build_workflow(configuration).as_agent() - - -server = ResponsesHostServer(agent_factory=create_agent) -``` - -The host calls the factory inside the current request's platform context, not during startup. The resulting agent -belongs to that request, including its entire stream. Responses recovery also creates a new agent through the factory. -The factory implementer is responsible for creating fresh mutable workflows, executors, wrapped agents, and their -sessions inside the factory. Returning an existing workflow or rebuilding around shared mutable executors can carry -state across requests. The host does not inspect or track object identities, clone execution objects, or enforce -weak-reference support. - -For example, construct the agent and executor inside the factory, not once outside it: - -```python -from agent_framework import Agent, AgentExecutor, WorkflowBuilder -from agent_framework_foundry_hosting import ResponsesHostServer - - def create_agent(): - agent = Agent(client=client, name="assistant") - executor = AgentExecutor(agent, id="assistant") - return WorkflowBuilder( - name="assistant-workflow", start_executor=executor - ).build().as_agent() + return build_workflow().as_agent() -server = ResponsesHostServer(agent_factory=create_agent) +server = ResponsesHostServer(agent=create_agent) ``` -Here `client` is an application-owned model client whose concurrency and resource lifetime support sharing. -Any mutable objects captured by tools or functional workflow functions also remain the implementer's responsibility. -Keep workflow names, executor IDs, and serialized state type registrations stable so a new instance can restore the -previous instance's checkpoints. - -Factories are also useful for ordinary agents with request-specific configuration. They do not automatically persist -custom fields on an agent: state needed on the next request must use the supported session or checkpoint stores. -The host recognizes the two built-in workflow adapters and their subclasses; it cannot discover a workflow hidden -inside an arbitrary custom agent. - -### Resource ownership - -The host enters and exits a factory-created agent's async context manager when it has one. Resources remain available -until execution and streaming finish, and are released on completion or interruption. Factory code must clean up its -own partially constructed resources if it fails before returning an agent. - -Built-in workflow adapters do not automatically close every client or tool captured by their executors. Keep -application-owned, concurrency-safe clients open for the host lifetime, or return a workflow subclass with an async -context manager that owns its request-specific resources. Do not close a shared client at the end of one request. -An ordinary `Agent` context manager also manages its chat client and MCP tools, so their ownership must match its lifetime. - -See the [workflow sample](../../samples/04-hosting/foundry-hosted-agents/responses/workflows/) for fresh agents and -executors using an application-owned model client, and the -[recovery sample](../../samples/04-hosting/foundry-hosted-agents/responses/resilient_long_running_workflow/) -for request-local workflows restored after a process restart. +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 @@ -103,8 +47,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: @@ -168,77 +113,18 @@ 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 +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. + ### Workflow checkpoints `ResponsesHostServer` persists workflow checkpoints durably. By default, it uses the `FoundryCheckpointStore`, backed by Foundry storage when hosted and file-based storage locally. Stored checkpoints are scoped under `checkpoints`. -Each factory-created workflow starts with independent runtime state. Responses restores only the checkpoint belonging -to the current user and conversation or response chain. A `previous_response_id` without a saved workflow checkpoint -fails; a conversation with prior history but no workflow checkpoint also fails rather than silently starting over. -A new conversation with no history can start without a checkpoint. - -Both hosts use the configured checkpoint storage directly and retain the workflow runtime's checkpoint-error policy. -The graph runtime logs checkpoint creation failures and may complete without saving its latest state. Hosting does -not intercept runner methods or turn those logged failures into invocation failures. A completed invocation is not -proof of a successful checkpoint save: a later continuation may restore an older checkpoint, or fail if none exists. -Exceptions propagated by the runtime, including functional workflow checkpoint errors, still fail the invocation. - -For resilient background Responses using a graph workflow, recovery uses the checkpoint associated with the persisted response output. -If no response checkpoint was recorded, it can use the latest workflow checkpoint. If execution stopped before any -workflow checkpoint was saved, recovery replays the original input using a fresh factory-created workflow. - -`FunctionalWorkflowAgent` does not support `resilient_background=True`. Its checkpoints retain completed step results -but not all output buffered inside those steps. Recovering from such a checkpoint could omit output; the host rejects -that configuration instead of silently losing output or rerunning application work. Functional workflows can use -the factory for normal Responses requests and supported pending-response continuation. - -Functional continuation validates every supplied result against the authorized pending request and its declared -response type, using the same supported conversions as graph workflows. A string such as `"false"` is not a boolean -decision. Invalid types, unknown request IDs, and duplicate responses fail the whole submitted batch before the host -copies a checkpoint, records an execution attempt, or runs the workflow. The original pending state remains available -for a corrected retry. - -### Invocations ordinary factory sessions - -`InvocationsHostServer(agent_factory=...)` also persists ordinary agents' `AgentSession` through -`agent_session_store_provider`. The default provider uses Foundry storage when hosted and file-based storage locally. -Sessions are scoped by the platform user and invocation session, with ordinary factory records separated from workflow -metadata and Responses records. - -Each request loads the stored session or calls the new agent's `create_session` method, passes that session to the run, -and saves it when the run finishes or is interrupted. For streaming requests, saving occurs after the iterator closes. -The host holds its per-scope lock through execution, saving, and resource cleanup. A new host can continue the stored -session, but custom mutable agent fields are not persisted automatically. - -Factory sessions are not retained in the host's session dictionary. Retention, expiration, storage quotas, and any -in-memory caching are the configured storage provider's responsibility; the host adds no expiry or eviction policy. -Ordinary **instance** calls remain unchanged: they retain sessions in the host's unbounded in-memory dictionary and -do not use this persistence provider. Changing that legacy retention behavior is outside this factory feature. - -### Invocations workflows - -`InvocationsHostServer(agent_factory=...)` persists workflow checkpoints for the platform user and invocation session. -After a completed invocation, a subsequent request with the same authorized session can restore its graph workflow -in a new agent instance, including after recreating the host. Ordinary-agent instance callers retain their existing -in-memory session behavior. - -The wire format remains unchanged: requests use `message` and `stream`, and responses contain text or streamed text. -This helper does not add a structured exchange for external workflow approvals or pending requests. A plain text -message is not an approval response. Use Responses when callers need that structured exchange. - -Both graph and functional workflows reject continuation when the previous invocation is pending or interrupted. -A saved checkpoint alone does not authorize resuming incomplete work: the host requires its persisted completion -record before restoring a graph checkpoint or accepting a new message. Rejection preserves the existing stored state. -After clean completion, graph workflows restore their state and functional workflows start the new message without -replaying the previous input. The completion record describes host execution, not checkpoint durability or -acknowledgment of HTTP delivery. - -Requests updating the same factory scope are serialized within one host; independent scopes can execute concurrently. -This is not a distributed lock across multiple host processes. Invocations does not automatically recover an -interrupted HTTP response, and checkpoints do not guarantee that external side effects execute exactly once. - ### Function approvals `ResponsesHostServer` persists function approvals durably. By default, it uses the diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py deleted file mode 100644 index 782dd4e276e..00000000000 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_factory.py +++ /dev/null @@ -1,101 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Private request factory and execution-scope support.""" - -from __future__ import annotations - -import asyncio -import inspect -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable -from contextlib import asynccontextmanager -from dataclasses import dataclass, field -from typing import Any, TypeAlias, TypeGuard, cast - -from agent_framework import ( - AgentResponse, - AgentResponseUpdate, - FunctionalWorkflowAgent, - ResponseStream, - SupportsAgentRun, - WorkflowAgent, -) - -HostedAgent: TypeAlias = SupportsAgentRun | FunctionalWorkflowAgent -AgentFactory: TypeAlias = Callable[[], HostedAgent | Awaitable[HostedAgent]] -WorkflowAgentTypes: TypeAlias = WorkflowAgent | FunctionalWorkflowAgent - - -def is_workflow_agent(agent: object) -> TypeGuard[WorkflowAgentTypes]: - """Recognize the built-in workflow adapters, including subclasses.""" - return isinstance(agent, (WorkflowAgent, FunctionalWorkflowAgent)) - - -def validate_agent_source(agent: HostedAgent | None, agent_factory: AgentFactory | None) -> None: - if (agent is None) == (agent_factory is None): - raise ValueError("Provide exactly one of agent or agent_factory.") - if agent is not None and is_workflow_agent(agent): - raise TypeError( - "Workflow agents must be supplied through agent_factory. " - "The factory must create a new workflow and new mutable executors for each request." - ) - if agent_factory is not None and not callable(agent_factory): - raise TypeError("agent_factory must be a callable accepting no arguments, not a coroutine object.") - - -class AgentFactoryResolver: - """Resolve request agents without caching them or inspecting their execution objects. - - Factory authors are responsible for constructing independent mutable runtimes. - """ - - def __init__(self, factory: AgentFactory) -> None: - self._factory = factory - - async def resolve(self) -> HostedAgent: - result = self._factory() - agent = await result if inspect.isawaitable(result) else result - if not isinstance(agent, (SupportsAgentRun, FunctionalWorkflowAgent)): - raise TypeError("agent_factory must return an agent implementing SupportsAgentRun or a workflow agent.") - return agent - - -@dataclass -class _ScopeLock: - lock: asyncio.Lock = field(default_factory=asyncio.Lock) - users: int = 0 - - -class ScopeLocks: - """Serialize a scope while retaining locks for both holders and waiters.""" - - def __init__(self) -> None: - self._entries: dict[tuple[str | None, ...], _ScopeLock] = {} - - @asynccontextmanager - async def hold(self, key: tuple[str | None, ...]) -> AsyncGenerator[None]: - entry = self._entries.setdefault(key, _ScopeLock()) - entry.users += 1 - try: - async with entry.lock: - yield - finally: - entry.users -= 1 - if not entry.users: - del self._entries[key] - - -async def close_run_iterator(iterator: AsyncIterator[Any]) -> None: - """Close the concrete run iterator, including known ResponseStream wrappers.""" - wrappers: list[ResponseStream[AgentResponseUpdate, AgentResponse]] = [] - current: Any = iterator - while isinstance(current, ResponseStream): - wrapper = cast(ResponseStream[AgentResponseUpdate, AgentResponse], current) - wrappers.append(wrapper) - current = wrapper._iterator # pyright: ignore[reportPrivateUsage] - try: - close = getattr(current, "aclose", None) - if close is not None: - await close() - finally: - for wrapper in reversed(wrappers): - await wrapper._run_cleanup_hooks() # pyright: ignore[reportPrivateUsage] 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..8722d9757f5 --- /dev/null +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py @@ -0,0 +1,28 @@ +# 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 hasattr(value, "run") and hasattr(value, "create_session") + + +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 a89e6cb37b2..f38ab174193 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,67 +1,18 @@ # Copyright (c) Microsoft. All rights reserved. -from __future__ import annotations - import json -import sys -from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, suppress -from typing import cast +from collections.abc import Awaitable, Callable -from agent_framework import ( - AgentSession, - CheckpointStorage, - FunctionalWorkflowAgent, - SessionStore, - SupportsAgentRun, - WorkflowAgent, - WorkflowRunState, -) -from agent_framework._filesystem import _storage_key_segment # pyright: ignore[reportPrivateUsage] +from agent_framework import AgentSession, SupportsAgentRun from agent_framework._telemetry import mark_feature_used -from anyio import CancelScope from azure.ai.agentserver.core import 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 Send from typing_extensions import Any, AsyncGenerator -from ._agent_factory import ( - AgentFactory, - AgentFactoryResolver, - HostedAgent, - ScopeLocks, - close_run_iterator, - is_workflow_agent, - validate_agent_source, -) +from ._agent_source import resolve_agent from ._feature_usage import FeatureIndex -from ._state_store import ( - ContextScopedStoreProvider, - StoreProvider, - _InvocationsAgentSessionStoreProvider, # pyright: ignore[reportPrivateUsage] - _InvocationsCheckpointStoreProvider, # pyright: ignore[reportPrivateUsage] -) - -_WORKFLOW_STATE_KEY = "_foundry_invocations_workflow" -_PENDING_REQUEST_ERROR = ( - "Invocations plain text cannot answer workflow pending requests or approvals. " - "Use a protocol that supports structured workflow responses." -) - - -class _InvocationStreamingResponse(StreamingResponse): - async def stream_response(self, send: Send) -> None: - # Starlette does not close a body iterator suspended at yield when send fails. - # Close it in the consumer task, where task-affine agent resources were entered. - try: - await super().stream_response(send) - except BaseException as exc: - with suppress(StopAsyncIteration): - await cast(AsyncGenerator[str], self.body_iterator).athrow(exc) - raise - finally: - await cast(AsyncGenerator[str], self.body_iterator).aclose() class InvocationsHostServer(InvocationAgentServerHost): @@ -69,60 +20,35 @@ class InvocationsHostServer(InvocationAgentServerHost): def __init__( self, - agent: HostedAgent | None = None, + agent: SupportsAgentRun | Callable[[], SupportsAgentRun | Awaitable[SupportsAgentRun]], *, - agent_factory: AgentFactory | None = None, - checkpoint_store_provider: ContextScopedStoreProvider[CheckpointStorage] | None = None, - agent_session_store_provider: StoreProvider[SessionStore] | None = None, openapi_spec: dict[str, Any] | None = None, **kwargs: Any, ) -> None: """Initialize an InvocationsHostServer. Args: - agent: An ordinary agent instance. Supply workflows through agent_factory instead. - agent_factory: A sync or async callable creating an agent for each request. Recreate - workflows, mutable executors, providers, and tools. An async context manager - returned by the factory is entered and exited within that request. - checkpoint_store_provider: Optional provider for user/session-scoped workflow checkpoints. - agent_session_store_provider: Optional provider for persisted factory agent sessions, - including workflow provider state. Storage retention is controlled by the provider. + 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. This host will expect the request to be a JSON body with a "message" field. - Responses contain plain text; "stream": true streams text as text/event-stream. - Workflow factories must use stable workflow names and executor IDs across requests. - The text-only exchange cannot answer pending workflow requests or approvals. - Both graph and functional workflows support fresh messages after clean completion, - but not pending or interrupted continuation. Such continuation fails explicitly. - Ordinary factory sessions are persisted after each run, including interrupted runs, - rather than retained in this host. Ordinary agent instances retain sessions in memory. - Factories own cleanup of nested resources not exposed by an async context manager. + 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, agent_factory) super().__init__(openapi_spec=openapi_spec, **kwargs) self._agent = agent - self._agent_resolver = AgentFactoryResolver(agent_factory) if agent_factory is not None else None - self._scope_locks = ScopeLocks() - self._checkpoint_storage_provider = ( - _InvocationsCheckpointStoreProvider() if checkpoint_store_provider is None else checkpoint_store_provider - ) - self._agent_session_storage_provider = ( - _InvocationsAgentSessionStoreProvider() - if agent_session_store_provider is None - else agent_session_store_provider - ) - self._sessions: dict[str, AgentSession] = {} + self._sessions: dict[str | tuple[str, str], AgentSession] = {} self.invoke_handler(self._handle_invoke) mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) - def _partition_key(self) -> str: + def _partition_key(self) -> str | tuple[str, str]: """Get the partition key for the current request. - A partition key is made up of the session ID and user ID. If the request is not - from a hosted environment, the partition key will be just the session ID. In the + A hosted partition key is a tuple containing the session ID and user ID, + preserving their boundaries. Locally, the key is just the session ID. In the Foundry hosted environment, the partition key is used to maintain isolation between different sessions and users, such that one user cannot access another user's sessions. @@ -140,7 +66,7 @@ def _partition_key(self) -> str: "The hosted environment is missing session_id or user_id in the request context. " "Please ensure that the request is coming from a valid Foundry platform service." ) - return f"{context.session_id}:{context.user_id}" + return context.session_id, context.user_id if not context.session_id: raise RuntimeError( @@ -149,145 +75,10 @@ def _partition_key(self) -> str: return context.session_id - @asynccontextmanager - async def _request_agent(self, scope: tuple[str | None, str]) -> AsyncGenerator[tuple[HostedAgent, CancelScope]]: - async with self._scope_locks.hold(scope): - # Enter before the agent's own task groups. Adding a new cancel scope - # around their exit would violate AnyIO's required scope nesting. - with CancelScope() as cleanup_scope: - agent = await cast(AgentFactoryResolver, self._agent_resolver).resolve() - resources = AsyncExitStack() - try: - if isinstance(agent, AbstractAsyncContextManager): - await resources.enter_async_context(agent) - yield agent, cleanup_scope - finally: - cleanup_scope.shield = True - exc_info = sys.exc_info() - if isinstance(exc_info[1], GeneratorExit): - await resources.aclose() - else: - await resources.__aexit__(*exc_info) - - @asynccontextmanager - async def _workflow_session( - self, agent: WorkflowAgent | FunctionalWorkflowAgent, storage_id: str - ) -> AsyncGenerator[tuple[AgentSession, CheckpointStorage]]: - context = get_request_context() - storage = self._checkpoint_storage_provider.get_store( - config=self.config, context_id=storage_id, platform_context=context - ) - sessions = self._agent_session_storage_provider.get_store(config=self.config, platform_context=context) - workflow = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow # pyright: ignore[reportPrivateUsage] - kind = "graph" if isinstance(agent, WorkflowAgent) else "functional" - session = await sessions.get(storage_id) - saved_marker = session.state.get(_WORKFLOW_STATE_KEY) if session is not None else None - marker = cast(dict[str, Any], saved_marker) if isinstance(saved_marker, dict) else None - if session is not None and ( - marker is None or marker.get("name") != workflow.name or marker.get("kind") != kind - ): - raise RuntimeError("The stored Invocations workflow name or kind does not match the factory result.") - checkpoint = await storage.get_latest(workflow_name=workflow.name) - if checkpoint is not None and checkpoint.workflow_name != workflow.name: - raise RuntimeError("The stored Invocations checkpoint does not match the workflow name.") - if marker is not None and checkpoint is None: - raise RuntimeError("The existing Invocations workflow session is missing its required checkpoint.") - if checkpoint is not None and session is None: - raise RuntimeError("The existing Invocations workflow checkpoint is missing its required agent session.") - if isinstance(agent, WorkflowAgent) and checkpoint is not None and checkpoint.pending_request_info_events: - raise RuntimeError(_PENDING_REQUEST_ERROR) - if marker is not None and marker.get("completed") is not True: - raise RuntimeError(f"Invocations cannot continue a pending or interrupted {kind} workflow.") - if isinstance(agent, WorkflowAgent) and checkpoint is not None: - await agent.workflow.run(checkpoint_id=checkpoint.checkpoint_id, checkpoint_storage=storage) - if agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS: - raise RuntimeError(_PENDING_REQUEST_ERROR) - if session is None: - session = AgentSession(session_id=storage_id) - marker = {"name": workflow.name, "kind": kind, "completed": False} - session.state[_WORKFLOW_STATE_KEY] = marker - # Record the attempt before execution. A failed first run must not look like a new session. - await sessions.set(storage_id, session) - try: - yield session, storage - pending = ( - agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS - if isinstance(agent, WorkflowAgent) - else bool(agent.pending_requests) - ) - if pending: - raise RuntimeError(_PENDING_REQUEST_ERROR) - marker["completed"] = True - finally: - with CancelScope(shield=True): - await sessions.set(storage_id, session) - - @asynccontextmanager - async def _factory_session( - self, agent: HostedAgent, storage_id: str - ) -> AsyncGenerator[tuple[AgentSession, dict[str, Any]]]: - if is_workflow_agent(agent): - async with self._workflow_session(agent, storage_id) as (session, storage): - yield session, {"checkpoint_storage": storage} - else: - # Keep ordinary snapshots separate from workflow metadata in the same provider. - storage_id = f"ordinary-{storage_id}" - sessions = self._agent_session_storage_provider.get_store( - config=self.config, platform_context=get_request_context() - ) - session = await sessions.get(storage_id) - if session is None: - session = cast(SupportsAgentRun, agent).create_session(session_id=storage_id) - try: - yield session, {} - finally: - with CancelScope(shield=True): - await sessions.set(storage_id, session) - - async def _handle_factory_invoke(self, user_message: Any, *, stream: bool) -> Response: - context = get_request_context() - scope = (context.user_id, cast(str, context.session_id)) - storage_id = _storage_key_segment(json.dumps(scope, ensure_ascii=False), encoded_prefix="~invocations-") - - if stream: - - async def stream_response() -> AsyncGenerator[str]: - async with ( - self._request_agent(scope) as (agent, cleanup_scope), - self._factory_session(agent, storage_id) as ( - session, - run_kwargs, - ), - ): - iterator = agent.run(user_message, session=session, stream=True, **run_kwargs).__aiter__() - try: - async for update in iterator: - if update.text: - yield update.text - finally: - cleanup_scope.shield = True - await close_run_iterator(iterator) - - return _InvocationStreamingResponse( - stream_response(), - media_type="text/event-stream", - headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, - ) - - async with ( - self._request_agent(scope) as (agent, _), - self._factory_session(agent, storage_id) as ( - session, - run_kwargs, - ), - ): - response = await agent.run([user_message], session=session, **run_kwargs) - return Response(content=response.text) - async def _handle_invoke(self, request: Request) -> Response: """Invoke the agent with the given request.""" try: - session_id = self._partition_key() + partition_key = self._partition_key() except Exception as e: return Response(content=str(e), status_code=500) @@ -301,11 +92,15 @@ async def _handle_invoke(self, request: Request) -> Response: return StreamingResponse(content=error, status_code=400) return Response(content=error, status_code=400) - if self._agent_resolver is not None: - return await self._handle_factory_invoke(user_message, stream=stream) + 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 - agent = cast(HostedAgent, self._agent) - session = self._sessions.setdefault(session_id, AgentSession(session_id=session_id)) + agent = await resolve_agent(self._agent) if stream: 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 e43cf17d3f3..e2535af5136 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -9,38 +9,39 @@ import logging import os import re -import sys -from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable, 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 from urllib.parse import urlparse from agent_framework import ( - AgentResponse, AgentResponseUpdate, - AgentSession, ChatOptions, CheckpointStorage, Content, ContextProvider, - FunctionalWorkflowAgent, HistoryProvider, InMemoryHistoryProvider, Message, RawAgent, - SessionContext, SessionStore, SupportsAgentRun, UsageDetails, WorkflowAgent, - WorkflowRunState, add_usage_details, ) from agent_framework._telemetry import mark_feature_used -from agent_framework._workflows._typing_utils import is_instance_of, try_coerce_to_type from agent_framework.exceptions import AgentFrameworkException -from anyio import CancelScope from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.responses import ( ResponseContext, @@ -81,15 +82,7 @@ from mcp import McpError from typing_extensions import Any -from ._agent_factory import ( - AgentFactory, - AgentFactoryResolver, - HostedAgent, - ScopeLocks, - close_run_iterator, - is_workflow_agent, - validate_agent_source, -) +from ._agent_source import is_agent, resolve_agent from ._feature_usage import FeatureIndex from ._state_store import ( AgentSessionStoreProvider, @@ -147,8 +140,6 @@ def _create_response_event_stream(context: ResponseContext) -> ResponseEventStre _T = TypeVar("_T") -_WorkflowRunFactory = Callable[[], AsyncIterator[AgentResponseUpdate]] -_WorkflowRun = Callable[[_WorkflowRunFactory], AsyncIterator[AgentResponseUpdate]] # Sentinel put on the internal queue by _SignalledIterator's driver task to signal that the # wrapped iterator is exhausted (distinct from `None`, which is a valid item value). @@ -170,9 +161,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 @@ -194,7 +184,6 @@ def __init__(self, iterator: AsyncIterator[_T], *events: asyncio.Event) -> None: self._queue: asyncio.Queue[Any] = asyncio.Queue(maxsize=1) # The background task that drives the wrapped iterator. self._driver: asyncio.Task[None] | None = None - self._closing = False @property def signalled(self) -> bool: @@ -211,24 +200,20 @@ def __aiter__(self) -> _SignalledIterator[_T]: async def _drive(self) -> None: """Pull items from the wrapped iterator into ``self._queue`` for the object's lifetime.""" try: - # Consumer cancellation is forwarded explicitly by aclose(). Shield the - # driver from repeated AnyIO cancellation while its iterator unwinds. - with CancelScope(shield=True): + while True: try: - while True: - try: - item: Any = await self._iterator.__anext__() - except StopAsyncIteration: - break - await self._queue.put(item) - finally: - await close_run_iterator(self._iterator) - except Exception as exc: - if self._closing: - raise - await self._queue.put(exc) - else: - await self._queue.put(_STOP_SENTINEL) + 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: + close = getattr(self._iterator, "aclose", None) + if close is not None: + await close() async def __anext__(self) -> _T: if self._driver is None: @@ -242,25 +227,15 @@ async def __anext__(self) -> _T: await asyncio.wait([get_task, *waiters], return_when=asyncio.FIRST_COMPLETED) if any(waiter.done() for waiter in waiters): self._signalled = True - self._closing = True self._driver.cancel() - with CancelScope(shield=True), suppress(asyncio.CancelledError): + with suppress(BaseException): await self._driver get_task.cancel() with suppress(BaseException): await get_task raise StopAsyncIteration item = get_task.result() - except asyncio.CancelledError: - if not self._closing: - self._closing = True - self._driver.cancel() - raise finally: - if not get_task.done(): - get_task.cancel() - with suppress(asyncio.CancelledError): - await get_task for waiter in waiters: if not waiter.done(): waiter.cancel() @@ -281,10 +256,8 @@ async def aclose(self) -> None: """ if self._driver is None: return - if not self._closing: - self._closing = True - self._driver.cancel() - with CancelScope(shield=True), suppress(asyncio.CancelledError): + self._driver.cancel() + with suppress(BaseException): await self._driver @@ -413,7 +386,6 @@ def consent_url_from_error(exc: BaseException) -> list[ConsentError] | None: # endregion Foundry Toolbox Auth integration -# region ResponsesHostServer @dataclass(frozen=True) class _AgentConfiguration: workflow: bool @@ -423,37 +395,32 @@ class _AgentConfiguration: def _validate_agent_configuration( - agent: HostedAgent, + agent: SupportsAgentRun, history_source: Literal["agent_server", "agent"], options: ResponsesServerOptions | None, ) -> _AgentConfiguration: - workflow = is_workflow_agent(agent) - if isinstance(agent, WorkflowAgent) and agent.workflow._runner_context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] + 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." ) - if isinstance(agent, FunctionalWorkflowAgent): - if agent._workflow._checkpoint_storage is not None: # pyright: ignore[reportPrivateUsage] - raise RuntimeError("The hosting infrastructure must manage the workflow checkpoint storage.") - if options and options.resilient_background: - raise RuntimeError( - "Functional workflow recovery is not supported: checkpoints do not preserve buffered step output. " - "Use a graph workflow for resilient_background=True." - ) - if options and options.resilient_background and not workflow: + + 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 workflow: + 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." ) - agent_server_history = history_source == "agent_server" + + uses_agent_server_history = history_source == "agent_server" client_stores_by_default = False - if agent_server_history and not workflow: + 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 " @@ -487,30 +454,35 @@ def _validate_agent_configuration( "STORES_BY_DEFAULT so hosting can enforce downstream storage behavior." ) client_stores_by_default = stores_by_default + return _AgentConfiguration( - workflow, agent_server_history, client_stores_by_default, agent_server_history and not workflow + 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: HostedAgent, configuration: _AgentConfiguration) -> None: - if configuration.hosted_history and isinstance(agent, RawAgent): - 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)) +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: HostedAgent | None = None, + agent: SupportsAgentRun | Callable[[], SupportsAgentRun | Awaitable[SupportsAgentRun]], *, - agent_factory: AgentFactory | None = None, prefix: str = "", options: ResponsesServerOptions | None = None, store: ResponseProviderProtocol | None = None, @@ -523,11 +495,8 @@ def __init__( """Initialize a ResponsesHostServer. Args: - agent: An ordinary agent instance. Supply workflow agents through agent_factory. - agent_factory: A no-argument sync or async factory creating one agent per request. - Recreate mutable workflows and executors with stable names and IDs. Returned - async context managers remain open through execution, streaming, and persistence. - Functional workflows do not support resilient background recovery. + 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. @@ -555,6 +524,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: @@ -577,18 +547,28 @@ 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, agent_factory) - configuration = _validate_agent_configuration(agent, history_source, options) if agent is not None else None + + 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._agent_source = agent + self._agent = resolved_agent + self._configuration = configuration self._history_source: Literal["agent_server", "agent"] = history_source self._host_options = options - self._configuration = configuration + 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) - self._agent = agent - self._agent_resolver = AgentFactoryResolver(agent_factory) if agent_factory is not None else None - self._scope_locks = ScopeLocks() - if agent is not None and configuration is not None: - _initialize_agent_history(agent, configuration) + if resolved_agent is not None and configuration is not None: + _initialize_agent_history(resolved_agent, configuration) # Storage providers self._checkpoint_storage_provider = ( @@ -627,10 +607,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 @@ -651,71 +634,41 @@ async def _handle_response( ) -> AsyncIterable[ResponseStreamEvent | ResponseCheckpointEvent]: """Handle the creation of a response.""" 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() - if self._agent_resolver is None: - async with AsyncExitStack() as resources: - inner = self._handle_prepared_response( - request, - context, - cancellation_signal, - response_event_stream, - cast(HostedAgent, self._agent), - cast(_AgentConfiguration, self._configuration), - resources, - ) - async with aclosing(inner): - async for event in inner: - yield event - return + terminal_event: ResponseStreamEvent | None = None - try: - async with AsyncExitStack() as locks: - platform_context = get_request_context() - if self.config.is_hosted and not platform_context.user_id: - raise RuntimeError("The hosted request context is missing user_id.") - scope_id = context.conversation_id or request.get("previous_response_id") or context.response_id - _validate_checkpoint_context_id(scope_id) - await locks.enter_async_context(self._scope_locks.hold((platform_context.user_id, scope_id))) - agent = await self._agent_resolver.resolve() - configuration = _validate_agent_configuration(agent, self._history_source, self._host_options) - _initialize_agent_history(agent, configuration) - # This scope must precede any task-affine MCP scopes entered by the agent. - with CancelScope() as cleanup_scope: - resources = AsyncExitStack() - try: - inner = self._handle_prepared_response( - request, - context, - cancellation_signal, - response_event_stream, - agent, - configuration, - resources, - ) - async with aclosing(inner): - 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: - cleanup_scope.shield = True - exc_info = sys.exc_info() - if isinstance(exc_info[1], GeneratorExit): - await resources.aclose() - else: - await resources.__aexit__(*exc_info) - if terminal_event is not None: - yield terminal_event - except Exception as ex: - logger.exception("Failed to prepare or release the request agent") - for event in self._emit_failure(response_event_stream, None, ex): - yield event + 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, + ) + async with aclosing(inner): + 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 + if terminal_event is not None: + yield terminal_event async def _handle_prepared_response( self, @@ -723,18 +676,17 @@ async def _handle_prepared_response( context: ResponseContext, cancellation_signal: asyncio.Event, response_event_stream: ResponseEventStream, - agent: HostedAgent, + 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: - if self._agent_resolver is None: + if self._configuration is not None: await self._ensure_agent_ready() elif isinstance(agent, AbstractAsyncContextManager): await resources.enter_async_context(agent) @@ -778,7 +730,7 @@ async def _handle_prepared_response( "Cannot find an existing agent session for " f"previous_response_id={previous_response_id}." ) - session = cast(SupportsAgentRun, 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( @@ -809,7 +761,12 @@ async def _handle_prepared_response( try: if configuration.workflow: inner = self._handle_inner_workflow( - request, context, response_event_stream, tracker, cancellation_signal, agent + request, + context, + response_event_stream, + tracker, + cancellation_signal, + cast(WorkflowAgent, agent), ) else: inner = self._handle_inner_agent( @@ -818,7 +775,7 @@ async def _handle_prepared_response( response_event_stream, tracker, cancellation_signal, - cast(SupportsAgentRun, agent), + agent, configuration, ) @@ -849,6 +806,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. @@ -863,13 +821,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 self._history_source != "agent_server": + if not uses_agent_server_history: return [] history = await context.get_history() return await _output_items_to_messages(history, approval_storage=approval_storage) @@ -925,7 +886,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") @@ -985,7 +950,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( - agent.run(stream=True, **run_kwargs), # pyright: ignore[reportUnknownMemberType] + agent.run(stream=True, **run_kwargs), # type: ignore[reportUnknownMemberType] context.shutdown, cancellation_signal, ) @@ -1023,8 +988,7 @@ async def _handle_inner_agent( request_failure = misconfigured try: if not stored_output_violation: - with CancelScope(shield=True): - await session_storage.set(session_save_id, session) + await session_storage.set(session_save_id, session) except Exception as save_error: save_failure = save_error if request_interrupted: @@ -1046,102 +1010,6 @@ async def _handle_inner_agent( raise save_failure async def _handle_inner_workflow( - self, - request: CreateResponse, - context: ResponseContext, - response_event_stream: ResponseEventStream, - tracker: _OutputItemTracker, - cancellation_signal: asyncio.Event, - agent: HostedAgent, - ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: - if not is_workflow_agent(agent): - raise RuntimeError("Agent is not a workflow agent.") - platform_context = get_request_context() - save_id = context.conversation_id or context.response_id - load_id = save_id if context.is_recovery else context.conversation_id or request.get("previous_response_id") - _validate_checkpoint_context_id(save_id) - if load_id is not None: - _validate_checkpoint_context_id(load_id) - sessions = self._session_storage_provider.get_store(config=self.config, platform_context=platform_context) - session = await sessions.get(load_id) if load_id is not None else None - workflow = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow # pyright: ignore[reportPrivateUsage] - kind = "graph" if isinstance(agent, WorkflowAgent) else "functional" - state_key = "_foundry_responses_workflow" - saved_marker = session.state.get(state_key) if session is not None else None - marker = cast(dict[str, Any], saved_marker) if isinstance(saved_marker, dict) else None - if session is not None and ( - marker is None or marker.get("name") != workflow.name or marker.get("kind") != kind - ): - raise RuntimeError("The stored Responses workflow name or kind does not match the factory result.") - had_session = session is not None - if session is None: - session = AgentSession() - prior_completed = marker is not None and marker.get("completed") is True - state_marker: dict[str, Any] = {"name": workflow.name, "kind": kind, "completed": False} - attempt_saved = False - attempt_started = False - - async def run_workflow(stream_factory: _WorkflowRunFactory) -> AsyncGenerator[AgentResponseUpdate]: - nonlocal attempt_saved, attempt_started - if cancellation_signal.is_set(): - return - if not attempt_started: - session.state[state_key] = state_marker - attempt_saved = True - await sessions.set(save_id, session) - if cancellation_signal.is_set(): - return - attempt_started = True - iterator = stream_factory() - try: - async for update in iterator: - yield update - finally: - await close_run_iterator(iterator) - - try: - if isinstance(agent, FunctionalWorkflowAgent): - inner = self._handle_functional_workflow( - request, context, tracker, cancellation_signal, agent, had_session, prior_completed, run_workflow - ) - else: - inner = self._handle_graph_workflow( - request, - context, - response_event_stream, - tracker, - cancellation_signal, - agent, - session if agent.context_providers else None, - had_session, - run_workflow, - ) - async with aclosing(inner): - async for event in inner: - if isinstance(event, ResponseCheckpointEvent): - await sessions.set(save_id, session) - yield event - if attempt_started and not cancellation_signal.is_set() and not context.shutdown.is_set(): - pending = ( - agent.workflow.status == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS - if isinstance(agent, WorkflowAgent) - else bool(agent.pending_requests) - ) - session.state[state_key]["completed"] = not pending - finally: - with CancelScope(shield=True): - if attempt_started: - await sessions.set(save_id, session) - elif attempt_saved: - # The initial write can finish just as cancellation arrives. - # No workflow work ran, so preserve the previous conversation. - if had_session and load_id == save_id: - session.state[state_key] = saved_marker - await sessions.set(save_id, session) - else: - await sessions.delete(save_id) - - async def _handle_graph_workflow( self, request: CreateResponse, context: ResponseContext, @@ -1149,9 +1017,6 @@ async def _handle_graph_workflow( tracker: _OutputItemTracker, cancellation_signal: asyncio.Event, agent: WorkflowAgent, - session: AgentSession | None, - had_session: bool, - run_workflow: _WorkflowRun, ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: """Handle the creation of a response for a workflow agent.""" try: @@ -1182,13 +1047,6 @@ async def _handle_graph_workflow( if context.is_recovery: if not self._resilient_background: raise RuntimeError("Recovery mode is only supported when resilient_background=True.") - persisted_messages = ( - await _output_items_to_messages( - context.persisted_response.get("output", []), approval_storage=approval_storage - ) - if context.persisted_response is not None and session is not None - else [] - ) # Resume from the workflow checkpoint durably paired with the last persisted response # snapshot (recorded in that snapshot's own metadata) -- NOT simply the latest workflow # checkpoint in storage, which may be ahead of what response.output actually reflects if @@ -1196,16 +1054,8 @@ async def _handle_graph_workflow( checkpoint_id = response_event_stream.internal_metadata.get(_LATEST_CHECKPOINT_ID_KEY) if checkpoint_id is not None: logger.debug("Serving recovery request from workflow checkpoint %s", checkpoint_id) - run_stream = run_workflow( - lambda: self._resume_workflow_from_checkpoint( - checkpoint_id, - checkpoint_storage, - context.response_id, - agent, - session, - input_messages, - persisted_messages, - ) + run_stream = self._resume_workflow_from_checkpoint( + checkpoint_id, checkpoint_storage, context.response_id, agent ) else: latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=agent.workflow.name) @@ -1215,17 +1065,11 @@ async def _handle_graph_workflow( "resuming from the latest checkpoint", latest_checkpoint.checkpoint_id, ) - recovery_checkpoint_id = latest_checkpoint.checkpoint_id - run_stream = run_workflow( - lambda: self._resume_workflow_from_checkpoint( - recovery_checkpoint_id, - checkpoint_storage, - context.response_id, - agent, - session, - input_messages, - persisted_messages, - ) + run_stream = self._resume_workflow_from_checkpoint( + 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 @@ -1235,13 +1079,10 @@ async def _handle_graph_workflow( logger.debug( "Serving recovery request with no prior workflow checkpoint; replaying original input" ) - run_stream = run_workflow( - lambda: agent.run( - input_messages, - stream=True, - session=session, - checkpoint_storage=checkpoint_storage, - ) + run_stream = agent.run( + input_messages, + stream=True, + checkpoint_storage=checkpoint_storage, ) else: # Determine the latest checkpoint (if any) so we can resume the @@ -1267,12 +1108,6 @@ async def _handle_graph_workflow( raise RuntimeError( f"Cannot find an existing workflow checkpoint for previous_response_id={previous_response_id}." ) - if latest_checkpoint is None and ( - had_session or (context.conversation_id is not None and await context.get_history()) - ): - raise RuntimeError("The existing conversation is missing its required workflow checkpoint.") - if latest_checkpoint is not None and not had_session and agent.context_providers: - raise RuntimeError("The workflow checkpoint is missing its required outer agent session.") if latest_checkpoint is not None: # If we have a prior checkpoint, restore it first (drive the workflow @@ -1284,19 +1119,16 @@ async def _handle_graph_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 restores the workflow's 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_checkpoint_id = latest_checkpoint.checkpoint_id restore_iter = _SignalledIterator( - run_workflow( - lambda: agent.run( - stream=True, - checkpoint_id=restore_checkpoint_id, - checkpoint_storage=restore_checkpoint_storage, - ) + agent.run( + stream=True, + checkpoint_id=latest_checkpoint.checkpoint_id, + checkpoint_storage=restore_checkpoint_storage, ), context.shutdown, cancellation_signal, @@ -1315,13 +1147,10 @@ async def _handle_graph_workflow( if cancellation_signal.is_set(): return - run_stream = run_workflow( - lambda: agent.run( - input_messages, - stream=True, - session=session, - checkpoint_storage=checkpoint_storage, - ) + run_stream = agent.run( + input_messages, + stream=True, + checkpoint_storage=checkpoint_storage, ) main_iter = _SignalledIterator(run_stream, context.shutdown, cancellation_signal) @@ -1369,9 +1198,6 @@ async def _resume_workflow_from_checkpoint( checkpoint_storage: CheckpointStorage, response_id: str, agent: WorkflowAgent, - session: AgentSession | None, - input_messages: list[Message], - persisted_messages: list[Message], ) -> AsyncGenerator[AgentResponseUpdate]: """Resume a crashed background workflow run, forwarding every event it produces. @@ -1386,24 +1212,6 @@ async def _resume_workflow_from_checkpoint( TODO(@taochen): #7677 """ - session_context: SessionContext | None = None - if session is not None: - session_context = SessionContext( - session_id=session.session_id, - service_session_id=session.service_session_id, - input_messages=input_messages, - options={}, - ) - for provider in agent.context_providers: - if isinstance(provider, HistoryProvider) and not provider.load_messages: - continue - await provider.before_run( - agent=agent, - session=session, - context=session_context, - state=session.state.setdefault(provider.source_id, {}), - ) - updates: list[AgentResponseUpdate] = [] async for event in agent.workflow.run( stream=True, checkpoint_id=checkpoint_id, @@ -1412,114 +1220,7 @@ async def _resume_workflow_from_checkpoint( for update in agent._convert_workflow_event_to_agent_response_updates( # pyright: ignore[reportPrivateUsage] response_id, event ): - if session_context is not None: - updates.append(update) yield update - if session_context is not None: - response = AgentResponse.from_updates(updates) - session_context._response = AgentResponse( # pyright: ignore[reportPrivateUsage] - messages=[*persisted_messages, *response.messages] - ) - await agent._run_after_providers(session=session, context=session_context) # pyright: ignore[reportPrivateUsage] - - async def _handle_functional_workflow( - self, - request: CreateResponse, - context: ResponseContext, - tracker: _OutputItemTracker, - cancellation_signal: asyncio.Event, - agent: FunctionalWorkflowAgent, - had_session: bool, - prior_completed: bool, - run_workflow: _WorkflowRun, - ) -> AsyncGenerator[ResponseStreamEvent]: - if context.is_recovery: - raise RuntimeError("Functional workflow recovery cannot preserve buffered step output.") - previous_id = request.get("previous_response_id") - if previous_id is not None and context.conversation_id is not None: - raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.") - platform_context = get_request_context() - save_id = context.conversation_id or context.response_id - load_id = context.conversation_id or previous_id - storage = self._checkpoint_storage_provider.get_store( - config=self.config, context_id=save_id, platform_context=platform_context - ) - restore_storage = storage - if load_id is not None and load_id != save_id: - restore_storage = self._checkpoint_storage_provider.get_store( - config=self.config, context_id=load_id, platform_context=platform_context - ) - checkpoint = await restore_storage.get_latest(workflow_name=agent._workflow.name) # pyright: ignore[reportPrivateUsage] - if checkpoint is None and ( - had_session - or previous_id is not None - or (context.conversation_id is not None and await context.get_history()) - ): - raise RuntimeError("The existing conversation is missing its required workflow checkpoint.") - if checkpoint is not None and not had_session: - raise RuntimeError("The functional workflow checkpoint is missing its required agent session.") - approval_storage = self._function_approval_storage_provider.get_store( - config=self.config, platform_context=platform_context - ) - messages = await _items_to_messages(await context.get_input_items(), approval_storage=approval_storage) - run_kwargs: dict[str, Any] = {"messages": messages, "checkpoint_storage": storage} - if checkpoint is not None and not prior_completed: - pending = checkpoint.pending_request_info_events - if not pending: - raise RuntimeError("Cannot continue an interrupted functional workflow without pending requests.") - responses: dict[str, Any] = {} - for message in messages: - for content in message.contents: - request_id = content.call_id if content.type == "function_result" else content.id - if request_id is None or request_id not in pending or request_id in responses: - raise ValueError("The input does not match an authorized pending functional workflow request.") - pending_request = pending[request_id] - if content.type == "function_result": - response = content if pending_request.response_type is Content else content.result - elif content.type == "function_approval_response" and pending_request.response_type is bool: - response = content.approved - elif content.type == "function_approval_response" and pending_request.response_type is Content: - response = content - else: - raise ValueError( - "This pending functional workflow request requires a matching function result." - ) - response = try_coerce_to_type(response, pending_request.response_type) - if not is_instance_of(response, pending_request.response_type): - raise ValueError( - f"Response type mismatch for request ID {request_id}: " - f"expected {pending_request.response_type}, got {type(response)}" - ) - responses[request_id] = response - if not responses: - raise ValueError("Pending functional workflow requests require structured responses.") - run_kwargs = { - "checkpoint_storage": storage, - "checkpoint_id": checkpoint.checkpoint_id, - "responses": responses, - } - if cancellation_signal.is_set(): - return - - async def execute() -> AsyncGenerator[AgentResponseUpdate]: - if checkpoint is not None and not prior_completed: - # Copy the authorized checkpoint only once restoration actually starts. - await storage.save(checkpoint) - agent_stream = agent.run(stream=True, **run_kwargs) - try: - async for update in agent_stream: - yield update - finally: - await close_run_iterator(agent_stream) - - iterator = _SignalledIterator(run_workflow(execute), context.shutdown, cancellation_signal) - async with aclosing(iterator): - async for update in iterator: - for content in update.contents: - async for event in tracker.handle( - content, message_id=update.message_id, approval_storage=approval_storage - ): - yield event @staticmethod def _emit_failure( 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 fd77bfa0169..a88c984ead5 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 @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. + from abc import ABC, abstractmethod from datetime import datetime from typing import Generic, Protocol, TypeVar @@ -338,26 +339,4 @@ def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentReques return FoundryAgentSessionStore(platform_context) -class _InvocationsCheckpointStore(FoundryCheckpointStore): - DEFAULT_ROOT_SCOPE = "invocations_checkpoints" - - -class _InvocationsAgentSessionStore(FoundryAgentSessionStore): - DEFAULT_ROOT_SCOPE = "invocations_agent_sessions" - - -class _InvocationsCheckpointStoreProvider(CheckpointStoreProvider): # pyright: ignore[reportUnusedClass] - def get_store( - self, *, config: AgentConfig, context_id: str, platform_context: FoundryAgentRequestContext - ) -> CheckpointStorage: - return _InvocationsCheckpointStore( - context_id, platform_context, allowed_checkpoint_types=self._allowed_checkpoint_types - ) - - -class _InvocationsAgentSessionStoreProvider(AgentSessionStoreProvider): # pyright: ignore[reportUnusedClass] - def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentRequestContext) -> SessionStore: - return _InvocationsAgentSessionStore(platform_context) - - # endregion Agent session persistence diff --git a/python/packages/foundry_hosting/tests/test_agent_factory.py b/python/packages/foundry_hosting/tests/test_agent_factory.py deleted file mode 100644 index 96c1d44b402..00000000000 --- a/python/packages/foundry_hosting/tests/test_agent_factory.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Factory resolution leaves execution-object ownership to the application.""" - -from typing import Any - -import pytest -from agent_framework import AgentExecutor, AgentResponse, AgentSession, Message, WorkflowAgent, WorkflowBuilder - -from agent_framework_foundry_hosting._agent_factory import AgentFactoryResolver - - -class _SlottedAgent: - __slots__ = ("description", "id", "name") - - def __init__(self) -> None: - self.id = "slotted" - self.name: str | None = "slotted" - self.description: str | None = None - - def create_session(self, *, session_id: str | None = None) -> AgentSession: - return AgentSession(session_id=session_id) - - def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: - return AgentSession(session_id=session_id, service_session_id=service_session_id) - - def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any: - async def response() -> AgentResponse: - return AgentResponse(messages=[Message("assistant", ["slotted"])]) - - return response() - - -@pytest.mark.parametrize("asynchronous", [False, True]) -async def test_factory_does_not_require_weak_reference_support(asynchronous: bool) -> None: - calls = 0 - - def factory() -> WorkflowAgent: - nonlocal calls - calls += 1 - executor = AgentExecutor(_SlottedAgent(), id="inner") - return WorkflowBuilder(name="slotted-workflow", start_executor=executor).build().as_agent() - - async def async_factory() -> WorkflowAgent: - return factory() - - resolver = AgentFactoryResolver(async_factory if asynchronous else factory) - first = await resolver.resolve() - second = await resolver.resolve() - assert isinstance(first, WorkflowAgent) - assert isinstance(second, WorkflowAgent) - assert first is not second - assert first.workflow.executors["inner"] is not second.workflow.executors["inner"] - assert (await first.run("first")).text == "slotted" - assert (await second.run("second")).text == "slotted" - assert calls == 2 - - -@pytest.mark.parametrize("same_wrapper", [False, True]) -async def test_resolver_returns_factory_result_without_enforcing_object_ownership(same_wrapper: bool) -> None: - executor = AgentExecutor(_SlottedAgent(), id="inner") - agent = WorkflowBuilder(name="application-owned", start_executor=executor).build().as_agent() - calls = 0 - - def factory() -> WorkflowAgent: - nonlocal calls - calls += 1 - if same_wrapper: - return agent - return WorkflowBuilder(name="application-owned", start_executor=executor).build().as_agent() - - resolver = AgentFactoryResolver(factory) - first = await resolver.resolve() - second = await resolver.resolve() - assert isinstance(first, WorkflowAgent) - assert isinstance(second, WorkflowAgent) - assert first.workflow.executors["inner"] is executor - assert second.workflow.executors["inner"] is executor - assert (first is second) is same_wrapper - assert calls == 2 diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index d5f8b180230..c322a55be3f 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -11,9 +11,11 @@ from __future__ import annotations +import json from collections.abc import AsyncIterator, Iterator from contextlib import contextmanager -from unittest.mock import AsyncMock, MagicMock +from itertools import product +from unittest.mock import AsyncMock, MagicMock, patch import pytest from agent_framework import ( @@ -169,11 +171,20 @@ def test_local_missing_session_id_raises(self) -> None: with _request_context(), pytest.raises(RuntimeError, match="missing session_id"): server._partition_key() # pyright: ignore[reportPrivateUsage] - def test_hosted_missing_user_id_raises(self) -> None: + def test_local_ignores_user_id(self) -> None: + server = InvocationsHostServer(_make_agent(response_text="hi")) + with _request_context(session_id="sess-1", user_id="user-1"): + assert server._partition_key() == "sess-1" # pyright: ignore[reportPrivateUsage] + + @pytest.mark.parametrize( + ("session_id", "user_id"), + [(None, "user-1"), ("", "user-1"), ("sess-1", None), ("sess-1", ""), (None, None)], + ) + def test_hosted_requires_both_identifiers(self, session_id: str | None, user_id: str | None) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) server.config.is_hosted = True with ( - _request_context(call_id="call-1", session_id="sess-1"), + _request_context(call_id="call-1", session_id=session_id, user_id=user_id), pytest.raises(RuntimeError, match="missing session_id or user_id"), ): server._partition_key() # pyright: ignore[reportPrivateUsage] @@ -182,7 +193,31 @@ def test_hosted_returns_composite_key(self) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) server.config.is_hosted = True with _request_context(call_id="call-1", session_id="sess-1", user_id="user-1"): - assert server._partition_key() == "sess-1:user-1" # pyright: ignore[reportPrivateUsage] + assert server._partition_key() == ("sess-1", "user-1") # pyright: ignore[reportPrivateUsage] + + async def test_hosted_keys_and_session_ids_preserve_identifier_values(self) -> None: + agent = _make_agent(response_text="hi") + server = InvocationsHostServer(agent) + server.config.is_hosted = True + identifiers = ["part", "part:part", "part,part", "[part]", 'part"\\', "part\n\t", "\u00e9", r"\u00e9", " part "] + keys: set[tuple[str, str]] = set() + request = _make_request({"message": "Hi"}) + + for session_id, user_id in product(identifiers, repeat=2): + with _request_context(call_id="call-1", session_id=session_id, user_id=user_id): + key = server._partition_key() # pyright: ignore[reportPrivateUsage] + response = await server._handle_invoke(request) # pyright: ignore[reportPrivateUsage] + + assert isinstance(key, tuple) + assert key == (session_id, user_id) + assert key not in keys + keys.add(key) + assert response.status_code == 200 + session = agent.calls[-1]["session"] + assert isinstance(session, AgentSession) + expected_id = json.dumps([session_id, user_id], separators=(",", ":")) + assert session.session_id == expected_id + assert session.to_dict()["session_id"] == expected_id # endregion @@ -192,6 +227,56 @@ def test_hosted_returns_composite_key(self) -> None: class TestHandleInvoke: + 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: + agent = _make_agent(response_text="ok", stream_texts=["ok"]) + server = InvocationsHostServer(agent) + server.config.is_hosted = hosted + request = _make_request({"message": "Hi", "stream": stream}) + expected_id = '["sess-1","user-1"]' if hosted else "sess-1" + + 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): + response = await server._handle_invoke(request) # pyright: ignore[reportPrivateUsage] + if isinstance(response, StreamingResponse): + assert await _collect_stream(response) == "ok" + 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) + session_factory.assert_called_once_with(session_id=expected_id) + async def test_missing_message_returns_400(self) -> None: server = InvocationsHostServer(_make_agent(response_text="hi")) request = _make_request({"stream": False}) @@ -259,5 +344,64 @@ async def test_session_is_reused_across_requests(self) -> None: assert list(server._sessions) == ["sess-1"] # pyright: ignore[reportPrivateUsage] assert agent.calls[0]["session"] is agent.calls[1]["session"] + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize( + ("first_session_id", "first_user_id", "second_session_id", "second_user_id"), + [ + ("session:segment", "user", "session", "segment:user"), + ("session,segment", "user", "session", "segment,user"), + ("session", "first-user", "session", "second-user"), + ("first-session", "user", "second-session", "user"), + ], + ) + async def test_hosted_sessions_preserve_identifier_boundaries( + self, + stream: bool, + first_session_id: str, + first_user_id: str, + second_session_id: str, + second_user_id: str, + ) -> None: + agent = _make_agent(response_text="ok", stream_texts=["ok"]) + server = InvocationsHostServer(agent) + server.config.is_hosted = True + identifiers = [(first_session_id, first_user_id), (second_session_id, second_user_id)] + sessions: list[AgentSession] = [] + + for session_id, user_id in identifiers: + with _request_context(call_id="call-1", session_id=session_id, user_id=user_id): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "Hi", "stream": stream}) + ) + if isinstance(response, StreamingResponse): + assert await _collect_stream(response) == "ok" + else: + assert bytes(response.body).decode() == "ok" + assert response.status_code == 200 + + session = agent.calls[-1]["session"] + assert isinstance(session, AgentSession) + assert session.state == {} + session.state["turn"] = (session_id, user_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): + response = await server._handle_invoke( # pyright: ignore[reportPrivateUsage] + _make_request({"message": "Continue", "stream": stream}) + ) + if isinstance(response, StreamingResponse): + assert await _collect_stream(response) == "ok" + else: + 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)} + # endregion diff --git a/python/packages/foundry_hosting/tests/test_invocations_factory.py b/python/packages/foundry_hosting/tests/test_invocations_factory.py deleted file mode 100644 index 0f43db18975..00000000000 --- a/python/packages/foundry_hosting/tests/test_invocations_factory.py +++ /dev/null @@ -1,1223 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Request factory lifetime and workflow persistence for the text-only host.""" - -import asyncio -import json -from collections.abc import Callable, Iterator, Mapping, Sequence -from contextlib import contextmanager -from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock - -import httpx -import pytest -from agent_framework import ( - Agent, - AgentExecutor, - AgentResponse, - AgentResponseUpdate, - AgentSession, - BaseChatClient, - ChatResponse, - ChatResponseUpdate, - Content, - Executor, - FunctionalWorkflowAgent, - InMemoryCheckpointStorage, - Message, - ResponseStream, - RunContext, - SessionStore, - WorkflowAgent, - WorkflowBuilder, - WorkflowContext, - WorkflowEvent, - handler, - response_handler, - step, - workflow, -) -from agent_framework._filesystem import _storage_key_segment -from anyio import CancelScope, create_task_group -from azure.ai.agentserver.core import ( - FoundryAgentRequestContext, - get_request_context, - reset_request_context, - set_request_context, -) -from starlette.requests import Request -from starlette.responses import StreamingResponse -from typing_extensions import AsyncGenerator, Never, Self - -from agent_framework_foundry_hosting import InvocationsHostServer - - -@contextmanager -def _context(user: str | None = "user", session: str = "session") -> Iterator[None]: - token = set_request_context(FoundryAgentRequestContext(user_id=user, session_id=session)) - try: - yield - finally: - reset_request_context(token) - - -def _request(message: str = "hello", *, stream: bool = False) -> Request: - request = MagicMock(spec=Request) - request.json = AsyncMock(return_value={"message": message, "stream": stream}) - return request - - -async def _invoke( - server: InvocationsHostServer, - message: str = "hello", - *, - stream: bool = False, - user: str | None = "user", - session: str = "session", -) -> str: - with _context(user, session): - response = await server._handle_invoke(_request(message, stream=stream)) - if isinstance(response, StreamingResponse): - chunks = [chunk async for chunk in response.body_iterator] - return "".join(chunk if isinstance(chunk, str) else bytes(chunk).decode() for chunk in chunks) - return bytes(response.body).decode() - - -class _OwnedAgent: - def __init__(self, events: list[str], *, fail: bool = False, wait: asyncio.Event | None = None) -> None: - self.id = "ordinary" - self.name: str | None = "ordinary" - self.description: str | None = None - self.events = events - self.fail = fail - self.wait = wait - self.owner: asyncio.Task[Any] | None = None - self.calls: list[Any] = [] - self.session: AgentSession | None = None - - async def __aenter__(self) -> "_OwnedAgent": - self.owner = asyncio.current_task() - self.events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - assert asyncio.current_task() is self.owner - await asyncio.sleep(0) - self.events.append("exit") - - def create_session(self, *, session_id: str | None = None) -> AgentSession: - return AgentSession(session_id=session_id) - - def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: - return AgentSession(session_id=session_id, service_session_id=service_session_id) - - def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any: - self.calls.append(messages) - self.session = session - - async def updates() -> AsyncGenerator[AgentResponseUpdate]: - try: - self.events.append("run") - yield AgentResponseUpdate(contents=[Content.from_text("a")]) - if self.wait is not None: - await self.wait.wait() - if self.fail: - raise RuntimeError("run failed") - yield AgentResponseUpdate(contents=[Content.from_text("")]) - yield AgentResponseUpdate(contents=[Content.from_text("b")]) - finally: - await asyncio.sleep(0) - self.events.append("iterator closed") - - async def response() -> AgentResponse: - self.events.append("run") - if self.wait is not None: - await self.wait.wait() - if self.fail: - raise RuntimeError("run failed") - return AgentResponse(messages=[Message("assistant", ["ab"])]) - - return ResponseStream(updates(), finalizer=AgentResponse.from_updates) if stream else response() - - -class _Stores: - def __init__(self, sessions: SessionStore | None = None) -> None: - self.checkpoints: dict[str, InMemoryCheckpointStorage] = {} - self.sessions = sessions if sessions is not None else SessionStore() - self.checkpoint_provider = MagicMock() - self.checkpoint_provider.get_store.side_effect = self._checkpoint_store - self.session_provider = MagicMock() - self.session_provider.get_store.return_value = self.sessions - - def _checkpoint_store(self, *, context_id: str, **kwargs: Any) -> InMemoryCheckpointStorage: - assert get_request_context().session_id is not None - assert context_id.startswith("~invocations-") - assert "/" not in context_id and "\\" not in context_id - return self.checkpoints.setdefault(context_id, InMemoryCheckpointStorage()) - - def server(self, factory: Callable[..., Any]) -> InvocationsHostServer: - return InvocationsHostServer( - agent_factory=factory, - checkpoint_store_provider=self.checkpoint_provider, - agent_session_store_provider=self.session_provider, - ) - - -class _Counter(Executor): - def __init__(self) -> None: - super().__init__(id="counter") - self.count = 0 - - @handler - async def count_message( - self, - messages: list[Message], - ctx: WorkflowContext[Never, str], # type: ignore[valid-type] - ) -> None: - self.count += 1 - # Include the outer provider history, not just executor checkpoint state. - await ctx.yield_output(f"{self.count}:{'|'.join(message.text for message in messages)}") - - async def on_checkpoint_save(self) -> dict[str, Any]: - return {"count": self.count} - - async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: - self.count = state["count"] - - -def _graph(name: str = "counter") -> WorkflowAgent: - return WorkflowBuilder(name=name, start_executor=_Counter()).build().as_agent() - - -@workflow(name="functional") -async def _functional(messages: Any) -> str: - return str(messages) - - -@workflow(name="functional-pending") -async def _functional_pending(messages: Any, ctx: RunContext) -> str: - await ctx.request_info("approve?", response_type=str) - return str(messages) - - -class _Pending(Executor): - @handler - async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - await ctx.request_info("approve?", response_type=str) - - @response_handler - async def answer( - self, - original_request: str, - response: str, - ctx: WorkflowContext[Never, str], # type: ignore[valid-type] - ) -> None: - await ctx.yield_output(response) - - -def _pending_graph() -> WorkflowAgent: - return WorkflowBuilder(name="pending", start_executor=_Pending(id="pending")).build().as_agent() - - -@pytest.mark.parametrize("factory", [_graph, lambda: _functional.build().as_agent()]) -def test_direct_workflow_instances_and_subclasses_require_factory(factory: Callable[..., Any]) -> None: - agent = factory() - with pytest.raises(TypeError, match="agent_factory"): - InvocationsHostServer(agent) - subclass: Any = type("CustomWorkflowAgent", (type(agent),), {}) - underlying = agent.workflow if isinstance(agent, WorkflowAgent) else agent._workflow - with pytest.raises(TypeError, match="agent_factory"): - InvocationsHostServer(subclass(underlying)) - - -def test_factory_source_validation_and_no_startup_call() -> None: - factory = MagicMock() - with pytest.raises(ValueError, match="exactly one"): - InvocationsHostServer() - with pytest.raises(ValueError, match="exactly one"): - InvocationsHostServer(_OwnedAgent([]), agent_factory=factory) - with pytest.raises(TypeError, match="callable"): - InvocationsHostServer(agent_factory=cast(Any, 42)) - InvocationsHostServer(agent_factory=factory) - factory.assert_not_called() - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("awaitable", [False, True]) -async def test_factory_once_per_request_inside_context_preserves_text(stream: bool, awaitable: bool) -> None: - events: list[str] = [] - agents: list[_OwnedAgent] = [] - - def factory() -> Any: - assert get_request_context().user_id == "user" - agent = _OwnedAgent(events) - agents.append(agent) - - async def create() -> _OwnedAgent: - return agent - - return create() if awaitable else agent - - server = _Stores().server(factory) - assert await _invoke(server, stream=stream) == "ab" - assert await _invoke(server, "second", stream=stream) == "ab" - assert len(agents) == 2 - assert agents[0].session is not None - assert agents[1].session is not None - assert agents[0].session.session_id == agents[1].session.session_id - assert agents[0].calls == ["hello" if stream else ["hello"]] - assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) * 2 - assert server._agent is None - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_async_factory_failure_and_invalid_result_release_lock(stream: bool) -> None: - calls = 0 - - async def factory() -> Any: - nonlocal calls - calls += 1 - if calls == 1: - raise RuntimeError("construction failed") - return object() - - server = InvocationsHostServer(agent_factory=factory) - with pytest.raises(RuntimeError, match="construction failed"): - await _invoke(server, stream=stream) - with pytest.raises(TypeError, match="agent_factory must return"): - await _invoke(server, stream=stream) - assert calls == 2 - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_run_failure_closes_owner_and_iterator(stream: bool) -> None: - events: list[str] = [] - server = _Stores().server(lambda: _OwnedAgent(events, fail=True)) - with pytest.raises(RuntimeError, match="run failed"): - await _invoke(server, stream=stream) - assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) - assert not server._scope_locks._entries - - -async def test_stream_factory_and_resources_belong_to_asgi_consumer_task() -> None: - events: list[str] = [] - creation_task: asyncio.Task[Any] | None = None - - def factory() -> _OwnedAgent: - nonlocal creation_task - creation_task = asyncio.current_task() - return _OwnedAgent(events) - - server = _Stores().server(factory) - with _context(): - response = await server._handle_invoke(_request(stream=True)) - assert creation_task is None - assert isinstance(response, StreamingResponse) - body: list[bytes] = [] - - async def send(message: Any) -> None: - if message["type"] == "http.response.body": - body.append(message["body"]) - - consumer = asyncio.create_task(response.stream_response(send)) - await consumer - assert creation_task is consumer - assert b"".join(body) == b"ab" - assert events == ["enter", "run", "iterator closed", "exit"] - - -@pytest.mark.parametrize("disconnect", [False, True]) -async def test_stream_disconnect_or_cancellation_closes_suspended_iterator(disconnect: bool) -> None: - events: list[str] = [] - sent = asyncio.Event() - server = _Stores().server(lambda: _OwnedAgent(events)) - with _context(): - response = await server._handle_invoke(_request(stream=True)) - assert isinstance(response, StreamingResponse) - - async def send(message: Any) -> None: - if message["type"] == "http.response.body": - sent.set() - if disconnect: - raise OSError("disconnected") - await asyncio.Event().wait() - - consumer = asyncio.create_task(response.stream_response(send)) - await sent.wait() - if not disconnect: - consumer.cancel() - with pytest.raises(OSError if disconnect else asyncio.CancelledError): - await consumer - assert events == ["enter", "run", "iterator closed", "exit"] - assert not server._scope_locks._entries - - -async def test_cancelled_factory_construction_releases_scope() -> None: - started = asyncio.Event() - - async def factory() -> _OwnedAgent: - started.set() - await asyncio.Event().wait() - return _OwnedAgent([]) - - server = InvocationsHostServer(agent_factory=factory) - task = asyncio.create_task(_invoke(server)) - await started.wait() - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert not server._scope_locks._entries - - -async def test_same_scope_serializes_stream_through_cleanup_and_cancelled_waiters() -> None: - events: list[str] = [] - release = asyncio.Event() - first_chunk = asyncio.Event() - created = 0 - - def factory() -> _OwnedAgent: - nonlocal created - created += 1 - return _OwnedAgent(events) - - server = _Stores().server(factory) - with _context(): - response = await server._handle_invoke(_request(stream=True)) - assert isinstance(response, StreamingResponse) - - async def send(message: Any) -> None: - if message["type"] == "http.response.body" and message.get("body"): - first_chunk.set() - await release.wait() - - consumer = asyncio.create_task(response.stream_response(send)) - await first_chunk.wait() - cancelled = asyncio.create_task(_invoke(server)) - waiting = asyncio.create_task(_invoke(server)) - await asyncio.sleep(0) - assert created == 1 - cancelled.cancel() - with pytest.raises(asyncio.CancelledError): - await cancelled - assert len(server._scope_locks._entries) == 1 - assert await _invoke(server, session="independent") == "ab" - assert created == 2 - release.set() - await consumer - assert await waiting == "ab" - assert created == 3 - assert events == ["enter", "run", "enter", "run", "exit", "iterator closed", "exit", "enter", "run", "exit"] - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_graph_checkpoint_and_outer_history_survive_new_host(stream: bool) -> None: - stores = _Stores() - assert await _invoke(stores.server(_graph), "first", stream=stream) == "1:first" - second = await _invoke(stores.server(_graph), "second", stream=stream) - assert second == "2:first|1:first|second" - assert len(stores.checkpoints) == 1 - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_graph_scopes_isolate_users_sessions_and_unsafe_identifiers(stream: bool) -> None: - stores = _Stores() - server = stores.server(_graph) - scopes = [ - ("user", "../session"), - ("another-user", "../session"), - ("user", r"..\session"), - ("a:b", "c"), - ("b", "c:a"), - ("USER", "../session"), - ] - for user, session in scopes: - assert await _invoke(server, user, stream=stream, user=user, session=session) == f"1:{user}" - assert len(stores.checkpoints) == len(scopes) - for user, session in scopes: - assert (await _invoke(server, "next", stream=stream, user=user, session=session)).startswith("2:") - assert not server._sessions - - -@pytest.mark.parametrize("damage", ["checkpoint", "session", "name", "kind"]) -async def test_missing_or_incompatible_continuation_never_restarts(damage: str) -> None: - stores = _Stores() - await _invoke(stores.server(_graph)) - storage_id, storage = next(iter(stores.checkpoints.items())) - if damage == "checkpoint": - for checkpoint in await storage.list_checkpoints(workflow_name="counter"): - await storage.delete(checkpoint.checkpoint_id) - elif damage == "session": - await stores.sessions.delete(storage_id) - - def factory() -> WorkflowAgent | FunctionalWorkflowAgent: - if damage == "kind": - return _functional.build().as_agent() - return _graph("different" if damage == "name" else "counter") - - with pytest.raises(RuntimeError, match="missing|does not match"): - await _invoke(stores.server(factory), "next") - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_graph_pending_request_cannot_be_answered_with_plain_text(stream: bool) -> None: - stores = _Stores() - server = stores.server(_pending_graph) - with pytest.raises(RuntimeError, match="plain text cannot answer"): - await _invoke(server, stream=stream) - with pytest.raises(RuntimeError, match="plain text cannot answer"): - await _invoke(stores.server(_pending_graph), "approved", stream=stream) - with pytest.raises(RuntimeError, match="plain text cannot answer"): - await _invoke(server, "approved", stream=stream, user="other") - assert len(stores.checkpoints) == 2 - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_functional_clean_completion_starts_new_message(stream: bool) -> None: - stores = _Stores() - - def factory() -> FunctionalWorkflowAgent: - return _functional.build().as_agent() - - assert await _invoke(stores.server(factory), "first", stream=stream) == ("first" if stream else "['first']") - storage = next(iter(stores.checkpoints.values())) - checkpoint = await storage.get_latest(workflow_name="functional") - assert checkpoint is not None - # A completed functional checkpoint can retain old pending events. Only the - # host's completion record distinguishes it from interrupted work. - checkpoint.pending_request_info_events["old"] = WorkflowEvent.request_info( - request_id="old", source_executor_id="functional", request_data="old", response_type=str - ) - await storage.save(checkpoint) - assert await _invoke(stores.server(factory), "next", stream=stream) == ("next" if stream else "['next']") - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_functional_pending_continuation_is_explicitly_unsupported(stream: bool) -> None: - stores = _Stores() - - def factory() -> FunctionalWorkflowAgent: - return _functional_pending.build().as_agent() - - with pytest.raises(RuntimeError, match="plain text cannot answer"): - await _invoke(stores.server(factory), stream=stream) - with pytest.raises(RuntimeError, match="pending or interrupted functional"): - await _invoke(stores.server(factory), "approved", stream=stream) - - -async def test_functional_interrupted_continuation_does_not_restart() -> None: - stores = _Stores() - calls = 0 - - @workflow(name="interrupted") - async def interrupted(messages: Any) -> str: - nonlocal calls - calls += 1 - raise RuntimeError("interrupted") - - def factory() -> FunctionalWorkflowAgent: - return interrupted.build().as_agent() - - with pytest.raises(RuntimeError, match="interrupted"): - await _invoke(stores.server(factory)) - with pytest.raises(RuntimeError, match="missing its required checkpoint"): - await _invoke(stores.server(factory)) - assert calls == 1 - - -async def test_session_save_failure_closes_workflow_owner(monkeypatch: pytest.MonkeyPatch) -> None: - events: list[str] = [] - - class OwnedWorkflow(WorkflowAgent): - async def __aenter__(self) -> Self: - events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - events.append("exit") - - stores = _Stores() - monkeypatch.setattr(stores.sessions, "set", AsyncMock(side_effect=[None, RuntimeError("save failed")])) - server = stores.server(lambda: OwnedWorkflow(_graph().workflow)) - with pytest.raises(RuntimeError, match="save failed"): - await _invoke(server) - assert events == ["enter", "exit"] - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("cancel", [False, True]) -async def test_anyio_resource_scopes_exit_in_consumer_task_and_correct_order(cancel: bool) -> None: - events: list[str] = [] - sent = asyncio.Event() - - class TaskGroupAgent(_OwnedAgent): - async def __aenter__(self) -> Self: - await super().__aenter__() - self.group = create_task_group() - await self.group.__aenter__() - return self - - async def __aexit__(self, *args: Any) -> None: - await self.group.__aexit__(*args) - await super().__aexit__(*args) - - server = _Stores().server(lambda: TaskGroupAgent(events)) - with _context(): - response = await server._handle_invoke(_request(stream=True)) - assert isinstance(response, StreamingResponse) - scope = CancelScope() - - async def send(message: Any) -> None: - if message["type"] == "http.response.body" and message.get("body"): - sent.set() - if cancel: - await asyncio.Event().wait() - - async def consume() -> None: - with scope: - await response.stream_response(send) - - consumer = asyncio.create_task(consume()) - await sent.wait() - if cancel: - scope.cancel() - await consumer - assert events == ["enter", "run", "iterator closed", "exit"] - assert not server._scope_locks._entries - - -class _TranscriptClient(BaseChatClient): - def __init__(self, transcripts: list[list[str]]) -> None: - super().__init__() - self.transcripts = transcripts - - def _inner_get_response( - self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any - ) -> Any: - self.transcripts.append([message.text for message in messages]) - - async def updates() -> AsyncGenerator[ChatResponseUpdate]: - yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("recorded")]) - - async def response() -> ChatResponse: - return ChatResponse(messages=[Message("assistant", ["recorded"])]) - - return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("use_agent_executor", [False, True]) -async def test_http_sessions_use_independent_workflow_state(stream: bool, use_agent_executor: bool) -> None: - stores = _Stores() - transcripts: list[list[str]] = [] - factory_scopes: list[tuple[str | None, str | None]] = [] - - class Remember(Executor): - @handler - async def remember(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - previous = ctx.get_state("previous") - ctx.set_state("previous", messages[0].text) - await ctx.yield_output(json.dumps(previous)) - - def factory() -> WorkflowAgent: - context = get_request_context() - factory_scopes.append((context.user_id, context.session_id)) - executor = ( - AgentExecutor(Agent(client=_TranscriptClient(transcripts), name="inner"), id="inner") - if use_agent_executor - else Remember(id="remember") - ) - return WorkflowBuilder(name="http-state", start_executor=executor).build().as_agent() - - server = stores.server(factory) - server.config.is_hosted = True - scopes = [("user-a", "session-a"), ("user-b", "session-b"), ("user-a", "session-c")] - async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: - for (user, session), text in zip(scopes, ("first", "second", "independent"), strict=True): - response = await client.post( - "/invocations", - json={"message": text, "stream": stream}, - params={"agent_session_id": session}, - headers={"x-agent-user-id": user, "x-agent-foundry-call-id": f"call-{session}"}, - ) - assert response.status_code == 200, response.text - assert response.text == ("recorded" if use_agent_executor else "null") - if use_agent_executor: - assert transcripts[-1] == [text] - assert factory_scopes == scopes - - -class _SnapshotSessions(SessionStore): - def __init__(self) -> None: - self.snapshots: dict[str, str] = {} - - async def get(self, session_id: str) -> AgentSession | None: - snapshot = self.snapshots.get(session_id) - return AgentSession.from_dict(json.loads(snapshot)) if snapshot is not None else None - - async def set(self, session_id: str, session: AgentSession) -> None: - await asyncio.sleep(0) - self.snapshots[session_id] = json.dumps(session.to_dict()) - - -class _PersistedAgent(Agent): - def __init__(self, client: BaseChatClient, events: list[str]) -> None: - super().__init__(client=client, name="ordinary") - self.events = events - - async def __aenter__(self) -> Self: - self.events.append("enter") - return await super().__aenter__() - - async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None: - await super().__aexit__(exc_type, exc_val, exc_tb) - self.events.append("exit") - - def create_session(self, *, session_id: str | None = None) -> AgentSession: - self.events.append("create") - session = super().create_session(session_id=session_id) - session.state["initialized"] = True - return session - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_ordinary_factory_sessions_persist_without_host_retention(stream: bool) -> None: - snapshots = _SnapshotSessions() - stores = _Stores(snapshots) - transcripts: list[list[str]] = [] - events: list[str] = [] - - def factory() -> Agent: - return _PersistedAgent(_TranscriptClient(transcripts), events) - - first = stores.server(factory) - for index in range(40): - assert await _invoke(first, f"first-{index}", session=str(index), stream=stream) == "recorded" - assert not first._sessions - assert not first._scope_locks._entries - assert len(snapshots.snapshots) == 40 - assert events.count("create") == 40 - second = stores.server(factory) - assert await _invoke(second, "next", session="0", stream=stream) == "recorded" - assert transcripts[-1] == ["first-0", "recorded", "next"] - assert events.count("create") == 40 - assert events.count("enter") == events.count("exit") == 41 - assert not second._sessions - assert not second._scope_locks._entries - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_ordinary_factory_persistence_isolates_authorized_scopes(stream: bool) -> None: - snapshots = _SnapshotSessions() - stores = _Stores(snapshots) - transcripts: list[list[str]] = [] - scopes = [ - ("user", "../session"), - ("other", "../session"), - ("user", r"..\session"), - ("a:b", "c"), - ("b", "c:a"), - ("USER", "../session"), - ] - for index, (user, session) in enumerate(scopes): - server = stores.server(lambda: Agent(client=_TranscriptClient(transcripts), name="ordinary")) - server.config.is_hosted = True - await _invoke(server, str(index), stream=stream, user=user, session=session) - assert transcripts[-1] == [str(index)] - for index, (user, session) in enumerate(scopes): - server = stores.server(lambda: Agent(client=_TranscriptClient(transcripts), name="ordinary")) - server.config.is_hosted = True - await _invoke(server, "next", stream=stream, user=user, session=session) - assert transcripts[-1] == [str(index), "recorded", "next"] - assert len(snapshots.snapshots) == len(scopes) - assert all("/" not in key and "\\" not in key for key in snapshots.snapshots) - assert stores.session_provider.get_store.call_args.kwargs["platform_context"].user_id == "USER" - stores.checkpoint_provider.get_store.assert_not_called() - - -async def test_ordinary_factory_namespace_does_not_replace_workflow_or_responses_sessions() -> None: - snapshots = _SnapshotSessions() - stores = _Stores(snapshots) - workflow_key = _storage_key_segment(json.dumps(("user", "session")), encoded_prefix="~invocations-") - # Responses uses conversation/response IDs, while Invocations encodes its user/session pair. - responses_session = AgentSession(session_id="session") - responses_session.state["responses"] = True - await snapshots.set("session", responses_session) - assert await _invoke(stores.server(_graph), "workflow-first") == "1:workflow-first" - original_records = dict(snapshots.snapshots) - transcripts: list[list[str]] = [] - await _invoke(stores.server(lambda: Agent(client=_TranscriptClient(transcripts), name="ordinary")), "ordinary") - assert transcripts == [["ordinary"]] - assert len(snapshots.snapshots) == 3 - assert {key: snapshots.snapshots[key] for key in original_records} == original_records - ordinary_key = next(key for key in snapshots.snapshots if key not in original_records) - assert ordinary_key != workflow_key - assert ordinary_key.startswith("ordinary-~invocations-") - assert await _invoke(stores.server(_graph), "workflow-next") == "2:workflow-first|1:workflow-first|workflow-next" - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("operation", ["get", "set"]) -async def test_ordinary_factory_store_failure_releases_resources_and_scope( - stream: bool, operation: str, monkeypatch: pytest.MonkeyPatch -) -> None: - snapshots = _SnapshotSessions() - stores = _Stores(snapshots) - transcripts: list[list[str]] = [] - events: list[str] = [] - original = getattr(snapshots, operation) - failing = AsyncMock(side_effect=RuntimeError(f"{operation} failed")) - monkeypatch.setattr(snapshots, operation, failing) - server = stores.server(lambda: _PersistedAgent(_TranscriptClient(transcripts), events)) - with pytest.raises(RuntimeError, match=f"{operation} failed"): - await _invoke(server, stream=stream) - assert events[-1] == "exit" - assert transcripts == ([] if operation == "get" else [["hello"]]) - assert not server._sessions - assert not server._scope_locks._entries - assert not snapshots.snapshots - monkeypatch.setattr(snapshots, operation, original) - assert await _invoke(server, "retry", stream=stream) == "recorded" - assert len(snapshots.snapshots) == 1 - - -class _InterruptedClient(BaseChatClient): - def __init__(self, started: asyncio.Event, events: list[str], *, fail: bool) -> None: - super().__init__() - self.started = started - self.events = events - self.fail = fail - - def _inner_get_response( - self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any - ) -> Any: - async def interrupt() -> None: - self.started.set() - if self.fail: - raise RuntimeError("client failed") - await asyncio.Event().wait() - - async def updates() -> AsyncGenerator[ChatResponseUpdate]: - try: - yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("partial")]) - await interrupt() - finally: - self.events.append("client closed") - - async def response() -> ChatResponse: - await interrupt() - return ChatResponse(messages=[]) - - return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else response() - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("failure", ["error", "asyncio", "anyio"]) -async def test_ordinary_factory_interruption_persists_session_before_resource_cleanup( - stream: bool, failure: str, monkeypatch: pytest.MonkeyPatch -) -> None: - snapshots = _SnapshotSessions() - stores = _Stores(snapshots) - events: list[str] = [] - started = asyncio.Event() - original_set = snapshots.set - - async def save(session_id: str, session: AgentSession) -> None: - await original_set(session_id, session) - events.append("saved") - - monkeypatch.setattr(snapshots, "set", save) - server = stores.server( - lambda: _PersistedAgent(_InterruptedClient(started, events, fail=failure == "error"), events) - ) - scope = CancelScope() - - async def consume() -> None: - with scope: - await _invoke(server, stream=stream) - - task = asyncio.create_task(consume()) - await started.wait() - if failure == "anyio": - scope.cancel() - await task - else: - if failure == "asyncio": - task.cancel() - with pytest.raises(RuntimeError if failure == "error" else asyncio.CancelledError): - await task - assert events[-2:] == ["saved", "exit"] - if stream: - assert events.index("client closed") < events.index("saved") - assert len(snapshots.snapshots) == 1 - saved = await snapshots.get(next(iter(snapshots.snapshots))) - assert saved is not None - assert saved.state["initialized"] is True - assert not server._sessions - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("interruption", ["close", "disconnect", "cancel"]) -async def test_ordinary_factory_suspended_stream_persists_on_close( - interruption: str, monkeypatch: pytest.MonkeyPatch -) -> None: - snapshots = _SnapshotSessions() - stores = _Stores(snapshots) - events: list[str] = [] - original_set = snapshots.set - - async def save(session_id: str, session: AgentSession) -> None: - await original_set(session_id, session) - events.append("saved") - - monkeypatch.setattr(snapshots, "set", save) - server = stores.server(lambda: _PersistedAgent(_TranscriptClient([]), events)) - with _context(): - response = await server._handle_invoke(_request(stream=True)) - assert isinstance(response, StreamingResponse) - assert not snapshots.snapshots - assert not events - if interruption == "close": - iterator = cast(AsyncGenerator[str], response.body_iterator) - assert await anext(iterator) == "recorded" - assert not snapshots.snapshots - await iterator.aclose() - else: - sent = asyncio.Event() - - async def send(message: Any) -> None: - if message["type"] == "http.response.body" and message.get("body"): - assert not snapshots.snapshots - sent.set() - if interruption == "disconnect": - raise OSError("disconnected") - await asyncio.Event().wait() - - consumer = asyncio.create_task(response.stream_response(send)) - await sent.wait() - if interruption == "cancel": - consumer.cancel() - with pytest.raises(OSError if interruption == "disconnect" else asyncio.CancelledError): - await consumer - assert events[-2:] == ["saved", "exit"] - assert len(snapshots.snapshots) == 1 - assert not server._sessions - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_ordinary_factory_lock_covers_save_and_allows_independent_scopes( - stream: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - snapshots = _SnapshotSessions() - stores = _Stores(snapshots) - transcripts: list[list[str]] = [] - events: list[str] = [] - saving = asyncio.Event() - release = asyncio.Event() - original_set = snapshots.set - - async def save(session_id: str, session: AgentSession) -> None: - if get_request_context().session_id == "session" and not saving.is_set(): - saving.set() - await release.wait() - await original_set(session_id, session) - - monkeypatch.setattr(snapshots, "set", save) - server = stores.server(lambda: _PersistedAgent(_TranscriptClient(transcripts), events)) - first = asyncio.create_task(_invoke(server, "first", stream=stream)) - waiting: asyncio.Task[str] | None = None - try: - await asyncio.wait_for(saving.wait(), timeout=2) - waiting = asyncio.create_task(_invoke(server, "second", stream=stream)) - await asyncio.sleep(0) - assert events.count("enter") == 1 - assert events.count("exit") == 0 - assert not first.done() - assert await asyncio.wait_for(_invoke(server, "independent", session="other"), timeout=2) == "recorded" - assert transcripts == [["first"], ["independent"]] - assert not waiting.done() - release.set() - assert await first == "recorded" - assert await waiting == "recorded" - assert transcripts[-1] == ["first", "recorded", "second"] - assert len(snapshots.snapshots) == 2 - assert not server._sessions - assert not server._scope_locks._entries - finally: - release.set() - await first - if waiting is not None: - await waiting - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_ordinary_instance_keeps_legacy_memory_and_ignores_session_provider(stream: bool) -> None: - provider = MagicMock() - provider.get_store.side_effect = AssertionError("instance must not access persisted storage") - transcripts: list[list[str]] = [] - agent = Agent(client=_TranscriptClient(transcripts), name="ordinary") - server = InvocationsHostServer(agent, agent_session_store_provider=provider) - await _invoke(server, "first", stream=stream) - original = server._sessions["session"] - await _invoke(server, "second", stream=stream) - assert server._sessions["session"] is original - assert transcripts[-1] == ["first", "recorded", "second"] - await _invoke(server, "independent", session="other", stream=stream) - assert len(server._sessions) == 2 - replacement = InvocationsHostServer(agent, agent_session_store_provider=provider) - await _invoke(replacement, "fresh", stream=stream) - assert transcripts[-1] == ["fresh"] - provider.get_store.assert_not_called() - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_real_agent_executor_transcript_is_isolated_and_restored(stream: bool) -> None: - stores = _Stores() - transcripts: list[list[str]] = [] - - def factory() -> WorkflowAgent: - agent = Agent(client=_TranscriptClient(transcripts), name="inner") - executor = AgentExecutor(agent, id="inner") - return WorkflowBuilder(name="transcript", start_executor=executor).build().as_agent() - - assert await _invoke(stores.server(factory), "private-first", stream=stream) == "recorded" - assert await _invoke(stores.server(factory), "unrelated", stream=stream, user="other") == "recorded" - assert transcripts[-1] == ["unrelated"] - assert await _invoke(stores.server(factory), "private-next", stream=stream) == "recorded" - assert "private-first" in transcripts[-1] - assert "private-next" in transcripts[-1] - assert "unrelated" not in transcripts[-1] - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("cancel", [False, True]) -@pytest.mark.parametrize("recreate_host", [False, True]) -@pytest.mark.parametrize("functional", [False, True]) -async def test_interruption_with_checkpoint_rejects_new_message( - stream: bool, cancel: bool, recreate_host: bool, functional: bool -) -> None: - stores = _Stores() - calls: list[str] = [] - started = asyncio.Event() - interrupt = True - workflow_name = "interrupted-workflow" - - async def finish_message(message: str) -> str: - calls.append(f"finish:{message}") - if interrupt: - started.set() - if cancel: - await asyncio.Event().wait() - raise RuntimeError("failed after checkpoint") - return message - - class Start(Executor): - @handler - async def start(self, messages: list[Message], ctx: WorkflowContext[str]) -> None: - calls.append(f"start:{messages[0].text}") - await ctx.send_message(messages[0].text) - - class Finish(Executor): - @handler - async def finish(self, message: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - await ctx.yield_output(await finish_message(message)) - - @step - async def saved_step(message: str) -> str: - calls.append(f"start:{message}") - return message - - @workflow(name=workflow_name) - async def interrupted(messages: str | list[str]) -> str: - message = messages if isinstance(messages, str) else messages[0] - return await finish_message(await saved_step(message)) - - def factory() -> WorkflowAgent | FunctionalWorkflowAgent: - if functional: - return interrupted.build().as_agent() - start = Start(id="start") - finish = Finish(id="finish") - return WorkflowBuilder(name=workflow_name, start_executor=start).add_edge(start, finish).build().as_agent() - - server = stores.server(factory) - if cancel: - task = asyncio.create_task(_invoke(server, "original", stream=stream)) - try: - await asyncio.wait_for(started.wait(), timeout=5) - finally: - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - else: - with pytest.raises(RuntimeError, match="failed after checkpoint"): - await _invoke(server, "original", stream=stream) - interrupt = False - - storage_id, storage = next(iter(stores.checkpoints.items())) - checkpoint = await storage.get_latest(workflow_name=workflow_name) - assert checkpoint is not None - if not functional: - assert checkpoint.iteration_count == 1 - assert checkpoint.messages - session = await stores.sessions.get(storage_id) - assert session is not None - assert session.state["_foundry_invocations_workflow"]["completed"] is False - saved_session = session.to_dict() - checkpoint_ids = await storage.list_checkpoint_ids(workflow_name=workflow_name) - - next_server = stores.server(factory) if recreate_host else server - kind = "functional" if functional else "graph" - with pytest.raises(RuntimeError, match=f"pending or interrupted {kind}"): - await _invoke(next_server, "must-not-run", stream=stream) - assert calls == ["start:original", "finish:original"] - assert await storage.list_checkpoint_ids(workflow_name=workflow_name) == checkpoint_ids - session = await stores.sessions.get(storage_id) - assert session is not None - assert session.to_dict() == saved_session - assert not server._scope_locks._entries - assert not next_server._scope_locks._entries - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_functional_interruption_with_checkpoint_rejects_new_message(stream: bool) -> None: - stores = _Stores() - calls = 0 - - @step - async def saved_step() -> str: - return "saved" - - @workflow(name="interrupted-after-step") - async def interrupted(messages: Any) -> str: - nonlocal calls - calls += 1 - await saved_step() - raise RuntimeError("failed after checkpoint") - - def factory() -> FunctionalWorkflowAgent: - return interrupted.build().as_agent() - - with pytest.raises(RuntimeError, match="failed after checkpoint"): - await _invoke(stores.server(factory), stream=stream) - with pytest.raises(RuntimeError, match="pending or interrupted functional"): - await _invoke(stores.server(factory), "next", stream=stream) - assert calls == 1 - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_cancellation_during_run_closes_resources(stream: bool) -> None: - events: list[str] = [] - server = _Stores().server(lambda: _OwnedAgent(events, wait=asyncio.Event())) - task = asyncio.create_task(_invoke(server, stream=stream)) - await asyncio.sleep(0) - assert events == ["enter", "run"] - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert events == (["enter", "run", "iterator closed", "exit"] if stream else ["enter", "run", "exit"]) - assert not server._scope_locks._entries - - -async def test_invalid_request_does_not_construct_agent(monkeypatch: pytest.MonkeyPatch) -> None: - factory = MagicMock() - server = InvocationsHostServer(agent_factory=factory) - with _context(): - request = _request() - monkeypatch.setattr(request, "json", AsyncMock(return_value={})) - response = await server._handle_invoke(request) - assert response.status_code == 400 - with _context(): - server.config.is_hosted = True - with _context(user=None): - response = await server._handle_invoke(_request()) - assert response.status_code == 500 - factory.assert_not_called() - - -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("stream", [False, True]) -async def test_checkpoint_save_failure_preserves_runtime_behavior_and_closes_owner( - functional: bool, stream: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - events: list[str] = [] - - class OwnedWorkflow(WorkflowAgent): - async def __aenter__(self) -> Self: - events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - events.append("exit") - - class OwnedFunctional(FunctionalWorkflowAgent): - async def __aenter__(self) -> Self: - events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - events.append("exit") - - stores = _Stores() - storage = InMemoryCheckpointStorage() - monkeypatch.setattr(storage, "save", AsyncMock(side_effect=RuntimeError("checkpoint failed"))) - stores.checkpoint_provider.get_store.side_effect = None - stores.checkpoint_provider.get_store.return_value = storage - server = stores.server( - lambda: OwnedFunctional(_functional.build()) if functional else OwnedWorkflow(_graph().workflow) - ) - if functional: - with pytest.raises(RuntimeError, match="checkpoint failed"): - await _invoke(server, stream=stream) - else: - assert await _invoke(server, stream=stream) == "1:hello" - assert "does not fail the workflow run" in caplog.text - with pytest.raises(RuntimeError, match="missing its required checkpoint"): - await _invoke(server, stream=stream) - assert events == ["enter", "exit", "enter", "exit"] - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("stream", [False, True]) -async def test_checkpoint_preparation_failure_preserves_graph_runtime_behavior( - stream: bool, caplog: pytest.LogCaptureFixture -) -> None: - stores = _Stores() - calls: list[str] = [] - snapshots: list[int] = [] - - class FailingCheckpoint(_Counter): - @handler - async def count_message( - self, - messages: list[Message], - ctx: WorkflowContext[Never, str], # type: ignore[valid-type] - ) -> None: - calls.extend(message.text for message in messages) - await super().count_message(messages, ctx) - - async def on_checkpoint_save(self) -> dict[str, Any]: - snapshots.append(self.count) - if self.count: - raise RuntimeError("executor checkpoint preparation failed") - return await super().on_checkpoint_save() - - def factory() -> WorkflowAgent: - return WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() - - server = stores.server(factory) - assert await _invoke(server, stream=stream) == "1:hello" - assert "does not fail the workflow run" in caplog.text - storage_id, storage = next(iter(stores.checkpoints.items())) - checkpoint = await storage.get_latest(workflow_name="checkpoint-preparation") - assert checkpoint is not None - assert checkpoint.iteration_count == 0 - assert 0 in snapshots and 1 in snapshots - session = await stores.sessions.get(storage_id) - assert session is not None - assert session.state["_foundry_invocations_workflow"]["completed"] is True - assert "checkpoint_failed" not in session.state["_foundry_invocations_workflow"] - assert calls == ["hello"] - assert not server._scope_locks._entries diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index ca8aecf8eda..e7873cd0ed3 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -25,7 +25,6 @@ import pytest from agent_framework import ( Agent, - AgentExecutor, AgentExecutorRequest, AgentResponse, AgentResponseUpdate, @@ -362,26 +361,11 @@ async def set(self, session_id: str, session: AgentSession) -> None: _SESSION_STORE_UNSET = object() -def _make_server(agent: Any = None, **kwargs: Any) -> ResponsesHostServer: +def _make_server(agent: Any, **kwargs: Any) -> ResponsesHostServer: """Create a ResponsesHostServer, optionally replacing its private store for tests.""" session_store = kwargs.pop("session_store", _SESSION_STORE_UNSET) response_store = kwargs.pop("response_store", InMemoryResponseProvider()) - is_workflow = isinstance(agent, WorkflowAgent) or "agent_factory" in kwargs - if isinstance(agent, WorkflowAgent): - server = ResponsesHostServer(agent_factory=lambda: agent, store=response_store, **kwargs) - else: - server = ResponsesHostServer(agent, store=response_store, **kwargs) - if is_workflow: - if session_store is _SESSION_STORE_UNSET: - session_store = SessionStore() - checkpoints: dict[str, InMemoryCheckpointStorage] = {} - - def get_checkpoint_store(*, context_id: str, **kwargs: Any) -> InMemoryCheckpointStorage: - return checkpoints.setdefault(context_id, InMemoryCheckpointStorage()) - - checkpoint_provider = MagicMock(spec=CheckpointStoreProvider) - checkpoint_provider.get_store.side_effect = get_checkpoint_store - server._checkpoint_storage_provider = checkpoint_provider # pyright: ignore[reportPrivateUsage] + server = ResponsesHostServer(agent, store=response_store, **kwargs) if session_store is not _SESSION_STORE_UNSET: provider = MagicMock(spec=AgentSessionStoreProvider) provider.get_store.return_value = cast(SessionStore | None, session_store) @@ -987,9 +971,9 @@ def test_init_rejects_resilient_background_for_non_workflow_agent(self, tmp_path options=ResponsesServerOptions(resilient_background=True), ) - def test_init_rejects_direct_workflow_agent(self) -> None: + def test_init_rejects_steerable_conversations_for_workflow_agent(self) -> None: workflow_agent = _build_text_workflow_agent("hello from workflow") - with pytest.raises(TypeError, match="agent_factory"): + with pytest.raises(RuntimeError, match="steerable_conversations"): ResponsesHostServer( cast(SupportsAgentRun, workflow_agent), store=InMemoryResponseProvider(), @@ -1489,6 +1473,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]: @@ -1497,7 +1482,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) @@ -1527,6 +1516,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 @@ -4739,7 +4729,7 @@ async def test_workflow_rejects_invalid_checkpoint_scope( agent.workflow = MagicMock() agent.workflow.name = "workflow" agent.workflow._runner_context.has_checkpointing.return_value = False - server = ResponsesHostServer(agent_factory=lambda: agent, store=InMemoryResponseProvider()) + server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) context_kwargs: dict[str, Any] = {"response_id": "response-current", "mode_flags": MagicMock()} request = CreateResponse(model="m", input="hi") @@ -4850,6 +4840,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")])]) @@ -5587,8 +5596,7 @@ async def _aiter() -> AsyncIterator[AgentResponseUpdate]: async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None: await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True)) - inner_executor = AgentExecutor(inner, id="text-agent") - workflow = WorkflowBuilder(name="text-workflow", start_executor=start).add_edge(start, inner_executor).build() + workflow = WorkflowBuilder(start_executor=start).add_edge(start, inner).build() return WorkflowAgent(workflow=workflow, name="Text Workflow Agent") @@ -5671,10 +5679,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)) - inner_executor = AgentExecutor(inner, id="multi-update-agent") - workflow = ( - WorkflowBuilder(name="multi-update-workflow", start_executor=start).add_edge(start, inner_executor).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 @@ -5698,48 +5703,60 @@ def _build_approval_workflow_agent( async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None: await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True)) - inner_executor = AgentExecutor(mock_agent, id="approval-agent") - workflow = WorkflowBuilder(name="approval-workflow", start_executor=start).add_edge(start, inner_executor).build() + workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build() workflow_agent = WorkflowAgent(workflow=workflow, name="Approval Workflow Agent") return workflow_agent, mock_agent -def _build_approval_workflow_factory( - *, approval_request_id: str, final_text: str -) -> tuple[Callable[[], WorkflowAgent], list[_ToolApprovalWorkflowAgentMock]]: - agents: list[_ToolApprovalWorkflowAgentMock] = [] +class TestWorkflowAgentHosting: + """End-to-end HTTP tests for ``ResponsesHostServer`` hosting a ``WorkflowAgent``. - def factory() -> WorkflowAgent: - workflow_agent, inner = _build_approval_workflow_agent( - approval_request_id=approval_request_id, final_text=final_text - ) - agents.append(inner) - return workflow_agent + These tests drive ``_handle_inner_workflow`` through the ASGI stack: + they exercise checkpoint write/restore (multi-turn) and the + tool-approval round-trip path, which is the primary differentiator + relative to the regular agent path. + """ - return factory, agents + 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 -def _build_multi_update_workflow_factory( - texts: Sequence[str], -) -> tuple[Callable[[], WorkflowAgent], list[_MultiUpdateWorkflowAgentMock]]: - agents: list[_MultiUpdateWorkflowAgentMock] = [] + server = _make_server(create_agent) - def factory() -> WorkflowAgent: - workflow_agent, inner = _build_multi_update_workflow_agent(texts) - agents.append(inner) - return workflow_agent + first = await _post(server, input_text="one") + second = await _post(server, input_text="two") - return factory, agents + 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] = [] -class TestWorkflowAgentHosting: - """End-to-end HTTP tests for ``ResponsesHostServer`` hosting a ``WorkflowAgent``. + 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 - These tests drive ``_handle_inner_workflow`` through the ASGI stack: - they exercise checkpoint write/restore (multi-turn) and the - tool-approval round-trip path, which is the primary differentiator - relative to the regular agent path. - """ + 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") @@ -5820,18 +5837,23 @@ async def test_cancellation_signal_preempts_stuck_workflow_call(self) -> None: AsyncGenerator[Any, None], server._handle_response(request, context, cancellation_signal), # pyright: ignore[reportPrivateUsage] ) + await anext(handler) # response.created + await anext(handler) # response.in_progress - async def consume() -> list[Any]: - return [event async for event in handler] - - # Keep factory entry, iteration, and cleanup in the same task. - pending = asyncio.create_task(consume()) + # Pull the first workflow event in the background so we can wait for the inner agent's + # run() to actually start (proving it's genuinely stuck on `gate`) before signalling -- + # otherwise cancellation could preempt the pull before the workflow even reaches it. + pending = asyncio.ensure_future(anext(handler)) await asyncio.wait_for(inner.started.wait(), timeout=1.0) cancellation_signal.set() # Fires while the inner agent is stuck awaiting `gate`. + async def _drain() -> list[Any]: + first = await pending + return [first, *[event async for event in handler]] + # Bounded well below `gate` never being set: proves cancellation preempted the stuck # call instead of only being observed after it (eventually) produced an update. - events = await asyncio.wait_for(pending, timeout=1.0) + events = await asyncio.wait_for(_drain(), timeout=1.0) types = [event.get("type") for event in events if isinstance(event, Mapping)] assert "response.output_text.delta" not in types @@ -5862,12 +5884,12 @@ async def test_shutdown_signal_preempts_stuck_workflow_call(self, tmp_path: Path AsyncGenerator[Any, None], server._handle_response(request, context, cancellation_signal), # pyright: ignore[reportPrivateUsage] ) + await anext(handler) # response.created + await anext(handler) # response.in_progress - async def consume() -> list[Any]: - return [event async for event in handler] - - # Keep factory entry, iteration, and cleanup in the same task. - pending = asyncio.create_task(consume()) + # Pull the first workflow event in the background so we can wait for the inner agent's + # run() to actually start (proving it's genuinely stuck on `gate`) before signalling. + pending = asyncio.ensure_future(anext(handler)) await asyncio.wait_for(inner.started.wait(), timeout=1.0) context.shutdown.set() # Fires while the inner agent is stuck awaiting `gate`. @@ -5882,12 +5904,12 @@ async def test_cancellation_signal_set_before_turn_skips_new_input(self) -> None """Explicit-cancel: cancellation set before a continuation turn starts must skip that turn's new input entirely, whether caught by the restore-loop's own check or the standalone check guarding the start of a brand new workflow run.""" - agent_factory, agents = _build_multi_update_workflow_factory(["hello"]) - server = _make_server(agent_factory=agent_factory) + workflow_agent, inner = _build_multi_update_workflow_agent(["hello"]) + server = _make_server(workflow_agent) first = await _post(server, conversation_id="conv-1", stream=False) assert first.status_code == 200 - run_count_after_first_turn = sum(agent.run_count for agent in agents) + run_count_after_first_turn = inner.run_count assert run_count_after_first_turn == 1 request = CreateResponse(model="m", input="hi again", stream=True) @@ -5910,8 +5932,7 @@ async def test_cancellation_signal_set_before_turn_skips_new_input(self) -> None assert types[-1] == "response.completed" # At most the restore-only replay call happened; the new-turn call (which would deliver # "hi again") must never fire. - assert len(agents) == 2 - assert sum(agent.run_count for agent in agents) <= run_count_after_first_turn + 1 + assert inner.run_count <= run_count_after_first_turn + 1 async def test_shutdown_signal_set_before_restore_only_triggers_recovery(self, tmp_path: Path) -> None: """Shutdown observed while resuming a checkpoint (whether during the restore-only replay or @@ -5919,16 +5940,16 @@ async def test_shutdown_signal_set_before_restore_only_triggers_recovery(self, t ``exit_for_recovery()`` -- proving the post-loop ``signalled`` check (not a blind re-check of the flag) correctly gates this action so it doesn't also fire on a replay that merely finished naturally.""" - agent_factory, agents = _build_multi_update_workflow_factory(["hello"]) + workflow_agent, inner = _build_multi_update_workflow_agent(["hello"]) server = _make_server( - agent_factory=agent_factory, + workflow_agent, response_store=FileResponseStore(storage_dir=tmp_path), options=ResponsesServerOptions(resilient_background=True), ) first = await _post(server, conversation_id="conv-1", stream=False) assert first.status_code == 200 - run_count_after_first_turn = sum(agent.run_count for agent in agents) + run_count_after_first_turn = inner.run_count assert run_count_after_first_turn == 1 request = CreateResponse(model="m", input="hi again", stream=True) @@ -5949,8 +5970,7 @@ async def test_shutdown_signal_set_before_restore_only_triggers_recovery(self, t _ = [event async for event in handler] # Only the restore-only replay call may have happened; the new-turn call must never fire. - assert len(agents) == 2 - assert sum(agent.run_count for agent in agents) <= run_count_after_first_turn + 1 + assert inner.run_count <= run_count_after_first_turn + 1 async def test_previous_response_requires_existing_workflow_checkpoint(self) -> None: """A previous_response_id naming a scope with no checkpoint must fail loudly rather than @@ -6042,11 +6062,11 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None approval response back to the paused inner agent, and the inner agent emits the final assistant text. """ - agent_factory, agents = _build_approval_workflow_factory( + workflow_agent, mock_agent = _build_approval_workflow_agent( approval_request_id="apr_wf_rt", final_text="done with approval", ) - server = _make_server(agent_factory=agent_factory) + server = _make_server(workflow_agent) checkpoint_provider = server._checkpoint_storage_provider # pyright: ignore[reportPrivateUsage] with patch.object(checkpoint_provider, "get_store", wraps=checkpoint_provider.get_store) as get_store: @@ -6057,7 +6077,7 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None approval_items = [it for it in first_body["output"] if it["type"] == "mcp_approval_request"] assert len(approval_items) == 1 approval_request_id = approval_items[0]["id"] - assert sum(agent.run_count for agent in agents) == 1 + assert mock_agent.run_count == 1 second_payload: dict[str, Any] = { "model": "test-model", @@ -6084,8 +6104,7 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None # The inner agent must have been resumed (restore replay + new turn). # Restore call is a no-op for the mock (no input); the new-turn call # delivers the approval response, so run_count grows by at least 1. - assert len(agents) == 2 - assert sum(agent.run_count for agent in agents) >= 2 + assert mock_agent.run_count >= 2 # The final assistant text from the resumed inner agent surfaces in # the HTTP output. @@ -6103,7 +6122,7 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None # The new-turn invocation of the inner agent must have received the # approval response routed back through WorkflowAgent. approval_responses = [ - c for m in agents[-1].last_run_messages for c in m.contents if c.type == "function_approval_response" + c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response" ] assert len(approval_responses) == 1 assert approval_responses[0].approved is True @@ -6111,11 +6130,11 @@ async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None async def test_round_trip_approval_response_streaming(self) -> None: """Streaming variant of the round-trip: turn 2 is requested with ``stream=true`` and surfaces the resumed text as SSE events.""" - agent_factory, agents = _build_approval_workflow_factory( + workflow_agent, mock_agent = _build_approval_workflow_agent( approval_request_id="apr_wf_rt_st", final_text="streamed-done", ) - server = _make_server(agent_factory=agent_factory) + server = _make_server(workflow_agent) first = await _post(server, stream=False) first_body = first.json() @@ -6145,17 +6164,16 @@ async def test_round_trip_approval_response_streaming(self) -> None: text_done = [e for e in events if e["event"] == "response.output_text.done"] assert any("streamed-done" in e["data"]["text"] for e in text_done) - assert len(agents) == 2 - assert sum(agent.run_count for agent in agents) >= 2 + assert mock_agent.run_count >= 2 async def test_round_trip_approval_response_rejected(self) -> None: """Sending ``approve=False`` must surface as ``approved=False`` to the inner agent on resume.""" - agent_factory, agents = _build_approval_workflow_factory( + workflow_agent, mock_agent = _build_approval_workflow_agent( approval_request_id="apr_wf_reject", final_text="acknowledged", ) - server = _make_server(agent_factory=agent_factory) + server = _make_server(workflow_agent) first = await _post(server, stream=False) first_body = first.json() @@ -6180,7 +6198,7 @@ async def test_round_trip_approval_response_rejected(self) -> None: assert second.status_code == 200 approval_responses = [ - c for m in agents[-1].last_run_messages for c in m.contents if c.type == "function_approval_response" + c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response" ] assert len(approval_responses) == 1 assert approval_responses[0].approved is False @@ -6242,6 +6260,7 @@ async def _passthrough(items: Any, *, approval_storage: Any = None) -> list[Any] async def test_reads_overlap_and_preserve_order(self, monkeypatch: pytest.MonkeyPatch) -> None: self._identity_converters(monkeypatch) server = _make_server(_make_agent()) + server._uses_agent_server_history = True # pyright: ignore[reportPrivateUsage] history_msg = Message(role="assistant", contents=[Content.from_text("H")]) input_msg = Message(role="user", contents=[Content.from_text("I")]) @@ -6279,7 +6298,8 @@ async def get_history() -> list[Message]: async def test_history_read_skipped_without_agent_server_history(self, monkeypatch: pytest.MonkeyPatch) -> None: self._identity_converters(monkeypatch) - server = _make_server(_make_agent(), history_source="agent") + server = _make_server(_make_agent()) + server._uses_agent_server_history = False # pyright: ignore[reportPrivateUsage] input_msg = Message(role="user", contents=[Content.from_text("I")]) @@ -6298,6 +6318,7 @@ async def get_input_items() -> list[Message]: async def test_failed_read_cancels_and_drains_sibling(self, monkeypatch: pytest.MonkeyPatch) -> None: self._identity_converters(monkeypatch) server = _make_server(_make_agent()) + server._uses_agent_server_history = True # pyright: ignore[reportPrivateUsage] sibling_cancelled = asyncio.Event() @@ -6325,8 +6346,7 @@ async def get_history() -> list[Message]: # The still-blocked history read must have been cancelled, not left orphaned. await asyncio.wait_for(sibling_cancelled.wait(), timeout=1) - @pytest.mark.parametrize("use_factory", [False, True]) - async def test_session_preparation_failure_cancels_pending_reads(self, use_factory: bool) -> None: + async def test_session_preparation_failure_cancels_pending_reads(self) -> None: """If session preparation fails, the concurrently-launched read must be cancelled and drained by `_handle_inner_agent`, not left running as an orphan after the request fails.""" input_started = asyncio.Event() @@ -6340,14 +6360,7 @@ async def get(self, session_id: str) -> AgentSession | None: await input_started.wait() raise RuntimeError("session prep boom") - server = ( - _make_server( - agent_factory=lambda: Agent(client=_RecordingHistoryClient()), - session_store=_GetFailsOnceReadStarted(), - ) - if use_factory - else _make_server(_make_agent(), session_store=_GetFailsOnceReadStarted()) - ) + server = _make_server(_make_agent(), session_store=_GetFailsOnceReadStarted()) request = CreateResponse(model="m", input="hi", stream=True) # A previous_response_id makes session_load_id non-None so the failing get() is reached. request["previous_response_id"] = "resp-x" diff --git a/python/packages/foundry_hosting/tests/test_responses_factory.py b/python/packages/foundry_hosting/tests/test_responses_factory.py deleted file mode 100644 index 1981fff1567..00000000000 --- a/python/packages/foundry_hosting/tests/test_responses_factory.py +++ /dev/null @@ -1,1769 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -"""Request-owned Responses agents, resources, and durable workflow continuation.""" - -import asyncio -import copy -import gc -import json -import weakref -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterator, Mapping, Sequence -from contextlib import aclosing, contextmanager -from dataclasses import dataclass -from typing import Any, cast -from unittest.mock import AsyncMock, MagicMock - -import httpx -import pytest -from agent_framework import ( - Agent, - AgentExecutor, - AgentResponse, - AgentResponseUpdate, - AgentSession, - BaseChatClient, - ChatResponse, - ChatResponseUpdate, - Content, - Executor, - FunctionalWorkflowAgent, - InMemoryCheckpointStorage, - InMemoryHistoryProvider, - Message, - ResponseStream, - RunContext, - SessionStore, - WorkflowAgent, - WorkflowBuilder, - WorkflowContext, - handler, - response_handler, - step, - workflow, -) -from anyio import CancelScope, create_task_group -from azure.ai.agentserver.core import ( - FoundryAgentRequestContext, - get_request_context, - reset_request_context, - set_request_context, -) -from azure.ai.agentserver.responses import InMemoryResponseProvider, ResponseContext, ResponsesServerOptions -from azure.ai.agentserver.responses.aio import ResponseEventStream -from azure.ai.agentserver.responses.models import CreateResponse -from azure.ai.agentserver.responses.streaming._checkpoint import ResponseCheckpointEvent -from typing_extensions import Never, Self - -from agent_framework_foundry_hosting import ResponsesHostServer -from agent_framework_foundry_hosting._agent_factory import close_run_iterator - - -@contextmanager -def _platform(user: str = "alice") -> Iterator[None]: - token = set_request_context(FoundryAgentRequestContext(user_id=user, session_id="platform-session")) - try: - yield - finally: - reset_request_context(token) - - -def _context( - text: str = "hello", - *, - response: str = "response-1", - conversation: str | None = "conversation", - items: list[Any] | None = None, - history: list[Any] | None = None, -) -> ResponseContext: - context = ResponseContext(response_id=response, conversation_id=conversation, mode_flags=MagicMock()) - context.get_input_items = AsyncMock( # type: ignore[method-assign] - return_value=items - if items is not None - else [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": text}]}] - ) - context.get_history = AsyncMock(return_value=history or []) # type: ignore[method-assign] - return context - - -def _request(previous: str | None = None) -> CreateResponse: - request = CreateResponse(model="test-model", input="input resolved by context", stream=True) - if previous is not None: - request["previous_response_id"] = previous - return request - - -async def _collect( - server: ResponsesHostServer, - context: ResponseContext | None = None, - *, - user: str = "alice", - previous: str | None = None, -) -> list[Any]: - with _platform(user): - return [ - event async for event in server._handle_response(_request(previous), context or _context(), asyncio.Event()) - ] - - -async def _post_http_response( - client: httpx.AsyncClient, items: Any, *, user: str, stream: bool, previous: str | None = None -) -> dict[str, Any]: - payload = {"model": "test-model", "input": items, "stream": stream} - if previous is not None: - payload["previous_response_id"] = previous - response = await client.post( - "/responses", - json=payload, - headers={"x-agent-user-id": user, "x-agent-foundry-call-id": f"call-{user}"}, - ) - assert response.status_code == 200, response.text - if not stream: - return response.json() - events = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] - terminal = [event for event in events if event["type"] in ("response.completed", "response.failed")] - assert len(terminal) == 1, events - return terminal[0]["response"] - - -def _types(events: list[Any]) -> list[str]: - return [event["type"] for event in events if isinstance(event, Mapping)] - - -def _text(events: list[Any]) -> str: - return "".join( - event["delta"] - for event in events - if isinstance(event, Mapping) and event.get("type") == "response.output_text.delta" - ) - - -def _failure(events: list[Any]) -> str: - assert _types(events)[-1] == "response.failed", events - assert _types(events).count("response.failed") == 1 - assert "response.completed" not in _types(events) - return events[-1]["response"]["error"]["message"] - - -class _ApprovalStore: - def __init__(self) -> None: - self.requests: dict[str, Content] = {} - - async def save_approval_request(self, request_id: str, content: Content) -> None: - self.requests[request_id] = content - - async def load_approval_request(self, request_id: str) -> Content | None: - return self.requests.get(request_id) - - -class _Stores: - def __init__(self) -> None: - self.sessions: dict[str | None, SessionStore] = {} - self.checkpoints: dict[tuple[str | None, str], InMemoryCheckpointStorage] = {} - self.approvals: dict[str | None, _ApprovalStore] = {} - self.session_provider = MagicMock() - self.session_provider.get_store.side_effect = self._sessions - self.checkpoint_provider = MagicMock() - self.checkpoint_provider.get_store.side_effect = self._checkpoints - self.approval_provider = MagicMock() - self.approval_provider.get_store.side_effect = self._approvals - - def _sessions(self, *, platform_context: FoundryAgentRequestContext, **kwargs: Any) -> SessionStore: - assert get_request_context().user_id == platform_context.user_id - return self.sessions.setdefault(platform_context.user_id, SessionStore()) - - def _checkpoints( - self, *, platform_context: FoundryAgentRequestContext, context_id: str, **kwargs: Any - ) -> InMemoryCheckpointStorage: - assert get_request_context().user_id == platform_context.user_id - return self.checkpoints.setdefault((platform_context.user_id, context_id), InMemoryCheckpointStorage()) - - def _approvals(self, *, platform_context: FoundryAgentRequestContext, **kwargs: Any) -> _ApprovalStore: - assert get_request_context().user_id == platform_context.user_id - return self.approvals.setdefault(platform_context.user_id, _ApprovalStore()) - - def server(self, factory: Callable[..., Any], **kwargs: Any) -> ResponsesHostServer: - return ResponsesHostServer( - agent_factory=factory, - history_source="agent", - checkpoint_store_provider=self.checkpoint_provider, - agent_session_store_provider=self.session_provider, - function_approval_store_provider=self.approval_provider, - **kwargs, - ) - - -class _OwnedAgent: - id = "ordinary" - name: str | None = "ordinary" - description: str | None = "Request resource test agent" - - def __init__(self, events: list[str], *, wait: asyncio.Event | None = None, fail: bool = False) -> None: - self.events = events - self.wait = wait - self.fail = fail - self.owner: asyncio.Task[Any] | None = None - self.session: AgentSession | None = None - - async def __aenter__(self) -> "_OwnedAgent": - self.owner = asyncio.current_task() - self.events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - assert asyncio.current_task() is self.owner - await asyncio.sleep(0) - self.events.append("exit") - - def create_session(self, *, session_id: str | None = None) -> AgentSession: - return AgentSession(session_id=session_id) - - def get_session(self, service_session_id: Any, *, session_id: str | None = None) -> AgentSession: - return AgentSession(session_id=session_id, service_session_id=service_session_id) - - def run(self, messages: Any = None, *, stream: bool = False, session: Any = None, **kwargs: Any) -> Any: - assert stream - self.session = session - - async def updates() -> AsyncIterator[AgentResponseUpdate]: - try: - self.events.append("run") - yield AgentResponseUpdate(role="assistant", contents=[Content.from_text("first")]) - if self.wait is not None: - await self.wait.wait() - if self.fail: - raise RuntimeError("model failed") - yield AgentResponseUpdate(role="assistant", contents=[Content.from_text("second")]) - finally: - await asyncio.sleep(0) - self.events.append("iterator closed") - - return ResponseStream(updates(), finalizer=AgentResponse.from_updates) - - -class _Counter(Executor): - def __init__(self) -> None: - super().__init__(id="counter") - self.count = 0 - - @handler - async def count_message( - self, - messages: list[Message], - ctx: WorkflowContext[Never, str], # type: ignore[valid-type] - ) -> None: - self.count += 1 - await ctx.yield_output(f"{self.count}:{'|'.join(message.text for message in messages)}") - - async def on_checkpoint_save(self) -> dict[str, Any]: - return {"count": self.count} - - async def on_checkpoint_restore(self, state: dict[str, Any]) -> None: - self.count = state["count"] - - -def _graph(name: str = "counter", *, history: bool = False) -> WorkflowAgent: - providers = [InMemoryHistoryProvider(source_id="outer-history")] if history else [] - return WorkflowBuilder(name=name, start_executor=_Counter()).build().as_agent(context_providers=providers) - - -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("completed", [False, True]) -@pytest.mark.parametrize("stage", ["preset", "session_load", "input", "checkpoint_lookup", "attempt_save"]) -async def test_cancelled_workflow_preparation_preserves_session_and_next_request( - functional: bool, completed: bool, stage: str, monkeypatch: pytest.MonkeyPatch -) -> None: - await _cancel_workflow_preparation(functional, completed, stage, False, monkeypatch) - - -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("completed", [False, True]) -@pytest.mark.parametrize("stage", ["session_load", "input", "checkpoint_lookup", "attempt_save"]) -async def test_task_cancelled_workflow_preparation_preserves_session_and_next_request( - functional: bool, completed: bool, stage: str, monkeypatch: pytest.MonkeyPatch -) -> None: - await _cancel_workflow_preparation(functional, completed, stage, True, monkeypatch) - - -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("cancel_task", [False, True]) -async def test_cancelled_workflow_attempt_save_removes_unstarted_response_branch( - functional: bool, cancel_task: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - await _cancel_workflow_preparation(functional, True, "attempt_save", cancel_task, monkeypatch, chain=True) - - -async def _cancel_workflow_preparation( - functional: bool, - completed: bool, - stage: str, - cancel_task: bool, - monkeypatch: pytest.MonkeyPatch, - *, - chain: bool = False, -) -> None: - stores = _Stores() - factory = (lambda: _functional.build().as_agent()) if functional else _graph - stores.sessions["alice"] = SessionStore() - sessions = stores.sessions["alice"] - storage = InMemoryCheckpointStorage() - conversation = None if chain else "conversation" - saved_id = "initial" if chain else "conversation" - previous_id = "initial" if chain else None - stores.checkpoints[("alice", saved_id)] = storage - if completed: - initial = await _collect( - stores.server(factory), _context("before", response="initial", conversation=conversation) - ) - assert _types(initial)[-1] == "response.completed" - previous = await sessions.get(saved_id) - previous_snapshot = previous.to_dict() if previous is not None else None - checkpoint_ids = await storage.list_checkpoint_ids(workflow_name="functional" if functional else "counter") - context = _context("cancelled", response="cancelled", conversation=conversation) - cancellation_signal = asyncio.Event() - cancelled = False - consumer: asyncio.Task[list[Any]] | None = None - - async def consume() -> list[Any]: - with _platform(): - return [ - event - async for event in stores.server(factory)._handle_response( - _request(previous_id), context, cancellation_signal - ) - ] - - with monkeypatch.context() as patch: - if stage == "preset": - cancellation_signal.set() - else: - target, method = { - "session_load": (sessions, "get"), - "input": (context, "get_input_items"), - "checkpoint_lookup": (storage, "get_latest"), - "attempt_save": (sessions, "set"), - }[stage] - original = getattr(target, method) - - async def cancel_after_preparation(*args: Any, **kwargs: Any) -> Any: - nonlocal cancelled - result = await original(*args, **kwargs) - if not cancelled: - cancelled = True - if cancel_task: - assert consumer is not None - consumer.cancel() - else: - cancellation_signal.set() - await asyncio.sleep(0) - return result - - patch.setattr(target, method, cancel_after_preparation) - consumer = asyncio.create_task(consume()) - if cancel_task: - with pytest.raises(asyncio.CancelledError): - await consumer - else: - events = await consumer - assert not _text(events) - assert _types(events)[-1] == "response.completed" - current = await sessions.get(saved_id) - assert (current.to_dict() if current is not None else None) == previous_snapshot - assert await storage.list_checkpoint_ids(workflow_name="functional" if functional else "counter") == checkpoint_ids - if chain: - assert await sessions.get("cancelled") is None - - following = await _collect( - stores.server(factory), _context("next", response="following", conversation=conversation), previous=previous_id - ) - assert _types(following)[-1] == "response.completed", following - assert _text(following) == ("next" if functional else f"{2 if completed else 1}:next") - - -@pytest.mark.parametrize("functional", [False, True]) -async def test_cancelled_workflow_execution_without_checkpoint_preserves_incomplete_attempt( - functional: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - stores = _Stores() - storage = InMemoryCheckpointStorage() - stores.checkpoints[("alice", "conversation")] = storage - started = asyncio.Event() - release = asyncio.Event() - cancellation_signal = asyncio.Event() - calls: list[str] = [] - - @workflow(name="interrupted") - async def interrupted(messages: list[Message]) -> str: - calls.append(messages[0].text) - started.set() - await release.wait() - return messages[0].text - - factory = (lambda: interrupted.build().as_agent()) if functional else _graph - - async def consume() -> list[Any]: - with _platform(): - return [ - event - async for event in stores.server(factory)._handle_response( - _request(), _context("original"), cancellation_signal - ) - ] - - with monkeypatch.context() as patch: - if not functional: - original_save = storage.save - - async def blocked_checkpoint(checkpoint: Any) -> str: - started.set() - await release.wait() - return await original_save(checkpoint) - - patch.setattr(storage, "save", blocked_checkpoint) - consumer = asyncio.create_task(consume()) - await asyncio.wait_for(started.wait(), 2) - cancellation_signal.set() - await asyncio.wait_for(consumer, 2) - - session = await stores.sessions["alice"].get("conversation") - assert session is not None - assert session.state["_foundry_responses_workflow"]["completed"] is False - assert not await storage.list_checkpoint_ids(workflow_name="interrupted" if functional else "counter") - following = await _collect(stores.server(factory), _context("next", response="next")) - assert "missing its required workflow checkpoint" in _failure(following) - if functional: - assert calls == ["original"] - else: - context = _context("original") - context.is_recovery = True - recovered = await _collect( - stores.server(factory, options=ResponsesServerOptions(resilient_background=True)), context - ) - assert _types(recovered)[-1] == "response.completed" - assert _text(recovered) == "1:original" - - -@workflow(name="functional") -async def _functional(messages: list[Message]) -> str: - return "|".join(message.text for message in messages) - - -@workflow(name="functional-pending-string") -async def _pending_string(messages: list[Message], ctx: RunContext) -> str: - answer = await ctx.request_info("answer?", response_type=str) - return f"{messages[0].text}:{answer}" - - -@workflow(name="functional-pending-bool") -async def _pending_bool(messages: list[Message], ctx: RunContext) -> str: - answer = await ctx.request_info("approve?", response_type=bool) - return f"{messages[0].text}:{answer}" - - -class _Pending(Executor): - @handler - async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - await ctx.request_info(messages[0].text, response_type=str) - - @response_handler - async def answer( - self, - original_request: str, - response: str, - ctx: WorkflowContext[Never, str], # type: ignore[valid-type] - ) -> None: - await ctx.yield_output(f"{original_request}:{response}") - - -def _pending_graph() -> WorkflowAgent: - return WorkflowBuilder(name="pending", start_executor=_Pending(id="pending")).build().as_agent() - - -def test_constructor_validates_exactly_one_callable_without_constructing() -> None: - factory = MagicMock() - with pytest.raises(ValueError, match="exactly one"): - ResponsesHostServer() - with pytest.raises(ValueError, match="exactly one"): - ResponsesHostServer(_OwnedAgent([]), agent_factory=factory) - with pytest.raises(TypeError, match="callable"): - ResponsesHostServer(agent_factory=cast(Any, 42)) - server = _Stores().server(factory) - factory.assert_not_called() - assert server._agent is None - - -def test_constructor_rejects_coroutine_object_instead_of_factory() -> None: - async def factory() -> _OwnedAgent: - return _OwnedAgent([]) - - coroutine = factory() - try: - with pytest.raises(TypeError, match="callable"): - ResponsesHostServer(agent_factory=cast(Any, coroutine)) - finally: - coroutine.close() - - -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("subclass", [False, True]) -def test_direct_workflow_instances_require_factory(functional: bool, subclass: bool) -> None: - agent = _functional.build().as_agent() if functional else _graph() - if subclass: - underlying = agent._workflow if isinstance(agent, FunctionalWorkflowAgent) else agent.workflow - agent = type("CustomWorkflow", (type(agent),), {})(underlying) - with pytest.raises(TypeError, match="agent_factory"): - ResponsesHostServer(agent) - - -@pytest.mark.parametrize("kind", ["sync", "async", "awaitable-object"]) -async def test_factory_runs_once_per_request_not_startup_and_never_sets_instance_agent(kind: str) -> None: - events: list[str] = [] - agents: list[_OwnedAgent] = [] - - def create() -> _OwnedAgent: - assert get_request_context().user_id == "alice" - agent = _OwnedAgent(events) - agents.append(agent) - return agent - - async def create_async() -> _OwnedAgent: - return create() - - class AwaitableFactory: - def __call__(self) -> Any: - return create_async() - - factories: dict[str, Callable[..., Any]] = { - "sync": create, - "async": create_async, - "awaitable-object": AwaitableFactory(), - } - server = _Stores().server(factories[kind]) - assert not agents - for response in ("one", "two"): - result = await _collect(server, _context(response=response)) - assert _text(result) == "firstsecond" - assert _types(result)[-1] == "response.completed" - assert server._agent is None - assert server._agent_stack is None - await server._cleanup_agent() - assert len(agents) == 2 - assert events == ["enter", "run", "iterator closed", "exit"] * 2 - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("kind", ["invalid", "exception", "async-exception"]) -async def test_factory_failure_is_not_retried_and_releases_lock(kind: str) -> None: - calls = 0 - - def factory() -> Any: - nonlocal calls - calls += 1 - if kind == "invalid": - return object() - if kind == "exception": - raise RuntimeError("factory failed") - - async def failed() -> Any: - raise RuntimeError("factory failed") - - return failed() - - server = _Stores().server(factory) - message = _failure(await _collect(server)) - assert ("SupportsAgentRun" if kind == "invalid" else "factory failed") in message - assert calls == 1 - assert server._agent is None - assert not server._scope_locks._entries - - -async def test_cancelled_factory_is_not_retried_and_releases_lock() -> None: - entered = asyncio.Event() - calls = 0 - - async def factory() -> Any: - nonlocal calls - calls += 1 - entered.set() - await asyncio.Event().wait() - - server = _Stores().server(factory) - consumer = asyncio.create_task(_collect(server)) - await asyncio.wait_for(entered.wait(), 2) - consumer.cancel() - with pytest.raises(asyncio.CancelledError): - await consumer - assert calls == 1 - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("functional", [False, True]) -async def test_completed_workflows_are_not_retained_by_factory_resolver(functional: bool) -> None: - refs: list[weakref.ReferenceType[Any]] = [] - - def factory() -> Any: - agent = _functional.build().as_agent() if functional else _graph() - underlying = agent._workflow if isinstance(agent, FunctionalWorkflowAgent) else agent.workflow - refs.extend([weakref.ref(agent), weakref.ref(underlying)]) - return agent - - server = _Stores().server(factory) - assert _types(await _collect(server))[-1] == "response.completed" - await asyncio.sleep(0) - gc.collect() - assert all(reference() is None for reference in refs) - - -@pytest.mark.parametrize("finish", ["complete", "close", "cancel-signal", "model-error"]) -async def test_resources_stay_open_through_output_and_close_exactly_once(finish: str) -> None: - events: list[str] = [] - server = _Stores().server( - lambda: _OwnedAgent( - events, wait=asyncio.Event() if finish in ("close", "cancel-signal") else None, fail=finish == "model-error" - ) - ) - cancellation = asyncio.Event() - emitted: list[Any] = [] - with _platform(): - stream = cast(AsyncGenerator[Any, None], server._handle_response(_request(), _context(), cancellation)) - async with aclosing(stream): - async for event in stream: - emitted.append(event) - kind = event.get("type") if isinstance(event, Mapping) else None - if kind == "response.output_text.delta": - assert "enter" in events and "exit" not in events - if finish == "close": - break - if finish == "cancel-signal": - cancellation.set() - if kind in ("response.completed", "response.failed"): - assert events[-1] == "exit" - assert events == ["enter", "run", "iterator closed", "exit"] - assert not server._scope_locks._entries - if finish == "model-error": - assert "model failed" in _failure(emitted) - - -@pytest.mark.parametrize("model_failure", [False, True]) -async def test_session_persistence_failure_closes_resources_and_reports_both_errors( - model_failure: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - events: list[str] = [] - stores = _Stores() - sessions = SessionStore() - monkeypatch.setattr(sessions, "set", AsyncMock(side_effect=RuntimeError("session save failed"))) - stores.sessions["alice"] = sessions - server = stores.server(lambda: _OwnedAgent(events, fail=model_failure)) - message = _failure(await _collect(server)) - assert "session save failed" in message - if model_failure: - assert "model failed" in message - assert events == ["enter", "run", "iterator closed", "exit"] - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("cancel", [False, True]) -async def test_nested_task_affine_resources_close_in_consumer_task_under_anyio_cancellation( - cancel: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - events: list[str] = [] - sent = asyncio.Event() - scope = CancelScope() - stores = _Stores() - sessions = SessionStore() - save = sessions.set - - async def persist(session_id: str, session: AgentSession) -> None: - await asyncio.sleep(0) - await save(session_id, session) - events.append("saved") - - monkeypatch.setattr(sessions, "set", AsyncMock(side_effect=persist)) - stores.sessions["alice"] = sessions - - class NestedAgent(_OwnedAgent): - async def __aenter__(self) -> Self: - await super().__aenter__() - self.group = create_task_group() - await self.group.__aenter__() - return self - - async def __aexit__(self, *args: Any) -> None: - await self.group.__aexit__(*args) - await super().__aexit__(*args) - - server = stores.server(lambda: NestedAgent(events)) - - async def consume() -> None: - with _platform(), scope: - stream = cast(AsyncGenerator[Any, None], server._handle_response(_request(), _context(), asyncio.Event())) - async with aclosing(stream): - async for event in stream: - if isinstance(event, Mapping) and event.get("type") == "response.output_text.delta": - sent.set() - if cancel: - await asyncio.Event().wait() - - consumer = asyncio.create_task(consume()) - await asyncio.wait_for(sent.wait(), 2) - if cancel: - scope.cancel() - await asyncio.wait_for(consumer, 2) - assert events == ["enter", "run", "iterator closed", "saved", "exit"] - assert not server._scope_locks._entries - - -async def test_scope_lock_covers_last_output_cleanup_and_removes_cancelled_waiters() -> None: - events: list[str] = [] - exiting = asyncio.Event() - release = asyncio.Event() - constructed = 0 - - class SlowExit(_OwnedAgent): - async def __aexit__(self, *args: Any) -> None: - exiting.set() - await release.wait() - await super().__aexit__(*args) - - def factory() -> _OwnedAgent: - nonlocal constructed - constructed += 1 - return SlowExit(events) if constructed == 1 else _OwnedAgent(events) - - server = _Stores().server(factory) - first = asyncio.create_task(_collect(server)) - await asyncio.wait_for(exiting.wait(), 2) - second = asyncio.create_task(_collect(server, _context(response="two"))) - cancelled = asyncio.create_task(_collect(server, _context(response="three"))) - await asyncio.sleep(0) - assert constructed == 1 - assert not first.done() - assert next(iter(server._scope_locks._entries.values())).users == 3 - cancelled.cancel() - with pytest.raises(asyncio.CancelledError): - await cancelled - assert next(iter(server._scope_locks._entries.values())).users == 2 - release.set() - results = await asyncio.wait_for(asyncio.gather(first, second), 2) - assert all(_types(result)[-1] == "response.completed" for result in results) - assert constructed == 2 - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("other_user,other_scope", [("bob", "conversation"), ("alice", "other")]) -async def test_distinct_users_or_conversations_can_overlap(other_user: str, other_scope: str) -> None: - started = asyncio.Event() - release = asyncio.Event() - calls = 0 - - def factory() -> _OwnedAgent: - nonlocal calls - calls += 1 - if calls == 2: - started.set() - return _OwnedAgent([], wait=release) - - server = _Stores().server(factory) - first = asyncio.create_task(_collect(server)) - second = asyncio.create_task(_collect(server, _context(conversation=other_scope), user=other_user)) - try: - await asyncio.wait_for(started.wait(), 2) - finally: - release.set() - results = await asyncio.gather(first, second) - assert all(_types(result)[-1] == "response.completed" for result in results) - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("chain", [False, True]) -@pytest.mark.parametrize("history", [False, True]) -async def test_graph_state_and_explicit_outer_history_survive_new_hosts(chain: bool, history: bool) -> None: - stores = _Stores() - conversation = None if chain else "conversation" - first = await _collect(stores.server(lambda: _graph(history=history)), _context("first", conversation=conversation)) - assert _text(first) == "1:first" - second = await _collect( - stores.server(lambda: _graph(history=history)), - _context("second", response="response-2", conversation=conversation), - previous="response-1" if chain else None, - ) - assert _types(second)[-1] == "response.completed" - assert _text(second) == ("2:first|1:first|second" if history else "2:second") - assert ("alice", "response-2" if chain else "conversation") in stores.checkpoints - - -async def test_graph_stable_name_must_match_saved_marker_from_previous_host() -> None: - stores = _Stores() - assert _text(await _collect(stores.server(lambda: _graph("stable")))) == "1:hello" - result = await _collect(stores.server(lambda: _graph("new-random-name")), _context("next", response="two")) - assert "name or kind" in _failure(result) - assert not _text(result) - - -async def test_graph_with_outer_history_requires_saved_session_alongside_checkpoint() -> None: - stores = _Stores() - assert _text(await _collect(stores.server(lambda: _graph(history=True)))) == "1:hello" - stores.sessions.clear() - result = await _collect(stores.server(lambda: _graph(history=True)), _context("next", response="two")) - assert "missing its required outer agent session" in _failure(result) - assert not _text(result) - - -class _TranscriptClient(BaseChatClient): - def __init__(self, transcripts: list[list[str]]) -> None: - super().__init__() - self.transcripts = transcripts - - def _inner_get_response( - self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any - ) -> Any: - assert stream - self.transcripts.append([message.text for message in messages]) - - async def updates() -> AsyncIterator[ChatResponseUpdate]: - yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("recorded")]) - - return ResponseStream(updates(), finalizer=ChatResponse.from_updates) - - -@pytest.mark.parametrize("stream", [False, True]) -@pytest.mark.parametrize("use_agent_executor", [False, True]) -async def test_fresh_http_requests_use_independent_workflow_state(stream: bool, use_agent_executor: bool) -> None: - stores = _Stores() - transcripts: list[list[str]] = [] - factory_users: list[str | None] = [] - - class Remember(Executor): - @handler - async def remember(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - previous = ctx.get_state("previous") - ctx.set_state("previous", messages[0].text) - await ctx.yield_output(json.dumps(previous)) - - def factory() -> WorkflowAgent: - factory_users.append(get_request_context().user_id) - executor = ( - AgentExecutor(Agent(client=_TranscriptClient(transcripts), name="inner"), id="inner") - if use_agent_executor - else Remember(id="remember") - ) - return WorkflowBuilder(name="http-state", start_executor=executor).build().as_agent() - - server = stores.server(factory, store=InMemoryResponseProvider()) - server.config.is_hosted = True - async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: - for user, text in (("user-a", "first"), ("user-b", "second"), ("user-a", "independent")): - body = await _post_http_response(client, text, user=user, stream=stream) - assert body["status"] == "completed", body - assert [ - part["text"] - for item in body["output"] - if item["type"] == "message" - for part in item["content"] - if part["type"] == "output_text" - ] == ["recorded" if use_agent_executor else "null"] - if use_agent_executor: - assert transcripts[-1] == [text] - assert factory_users == ["user-a", "user-b", "user-a"] - - -@pytest.mark.parametrize("chain", [False, True]) -async def test_agent_executor_transcript_isolates_users_scopes_and_continues_across_hosts(chain: bool) -> None: - stores = _Stores() - transcripts: list[list[str]] = [] - - def factory() -> WorkflowAgent: - agent = Agent(client=_TranscriptClient(transcripts), name="inner") - inner = AgentExecutor(agent, id="inner") - return WorkflowBuilder(name="transcript", start_executor=inner).build().as_agent() - - conversation = None if chain else "conversation" - assert _text(await _collect(stores.server(factory), _context("private", conversation=conversation))) == "recorded" - await _collect(stores.server(factory), _context("bob-only", conversation=conversation), user="bob") - assert transcripts[-1] == ["bob-only"] - await _collect(stores.server(factory), _context("other-scope", conversation="other")) - assert transcripts[-1] == ["other-scope"] - continued = await _collect( - stores.server(factory), - _context("next", response="response-2", conversation=conversation), - previous="response-1" if chain else None, - ) - assert _text(continued) == "recorded" - assert transcripts[-1] == ["private", "recorded", "next"] - - -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("damage", ["name", "checkpoint", "marker", "history"]) -async def test_missing_or_incompatible_workflow_state_does_not_restart(functional: bool, damage: str) -> None: - stores = _Stores() - factory = (lambda: _functional.build().as_agent()) if functional else _graph - assert _types(await _collect(stores.server(factory)))[-1] == "response.completed" - session = await stores.sessions["alice"].get("conversation") - assert session is not None - if damage == "name": - session.state["_foundry_responses_workflow"]["name"] = "different-name" - elif damage == "marker": - session.state.pop("_foundry_responses_workflow") - else: - stores.checkpoints.clear() - await stores.sessions["alice"].set("conversation", session) - context = _context("next", response="two", history=[{"type": "message"}] if damage == "history" else None) - message = _failure(await _collect(stores.server(factory), context)) - assert ("name or kind" if damage in ("name", "marker") else "missing its required workflow checkpoint") in message - - -async def test_functional_resilient_background_is_explicitly_rejected_without_running() -> None: - calls: list[str] = [] - - @workflow(name="not-recoverable") - async def functional(messages: Any) -> str: - calls.append("run") - return "unexpected" - - server = _Stores().server( - lambda: functional.build().as_agent(), options=ResponsesServerOptions(resilient_background=True) - ) - assert "buffered step output" in _failure(await _collect(server)) - assert not calls - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("chain", [False, True]) -async def test_completed_functional_workflow_starts_fresh_input_not_cached_previous_input(chain: bool) -> None: - stores = _Stores() - calls: list[str] = [] - - @step - async def record(text: str) -> str: - calls.append(text) - return text - - @workflow(name="fresh-functional") - async def functional(messages: list[Message]) -> str: - return await record(messages[0].text) - - conversation = None if chain else "conversation" - assert ( - _text( - await _collect( - stores.server(lambda: functional.build().as_agent()), _context("first", conversation=conversation) - ) - ) - == "first" - ) - assert ( - _text( - await _collect( - stores.server(lambda: functional.build().as_agent()), - _context("second", response="two", conversation=conversation), - previous="response-1" if chain else None, - ) - ) - == "second" - ) - assert calls == ["first", "second"] - - -@pytest.mark.parametrize("kind", ["graph", "functional-string", "functional-bool"]) -@pytest.mark.parametrize("chain", [False, True]) -async def test_authorized_pending_response_resumes_matching_checkpoint(kind: str, chain: bool) -> None: - stores = _Stores() - factory = { - "graph": _pending_graph, - "functional-string": lambda: _pending_string.build().as_agent(), - "functional-bool": lambda: _pending_bool.build().as_agent(), - }[kind] - conversation = None if chain else "conversation" - first = await _collect(stores.server(factory), _context("private", conversation=conversation)) - assert _types(first)[-1] == "response.completed" - checkpoint = await stores.checkpoints[("alice", conversation or "response-1")].get_latest( - workflow_name={ - "graph": "pending", - "functional-string": "functional-pending-string", - "functional-bool": "functional-pending-bool", - }[kind] - ) - assert checkpoint is not None - request_id = next(iter(checkpoint.pending_request_info_events)) - approval_id = None - if kind == "functional-bool": - approval_id = next(iter(stores.approvals["alice"].requests)) - assert stores.approvals["alice"].requests[approval_id].id == request_id - item = ( - {"type": "mcp_approval_response", "approval_request_id": approval_id, "approve": True} - if kind == "functional-bool" - else {"type": "function_call_output", "call_id": request_id, "output": "accepted"} - ) - resumed = await _collect( - stores.server(factory), - _context(response="two", conversation=conversation, items=[item]), - previous="response-1" if chain else None, - ) - assert _types(resumed)[-1] == "response.completed", resumed - assert _text(resumed) == ("private:True" if kind == "functional-bool" else "private:accepted") - - -@pytest.mark.parametrize("kind", ["text", "wrong-id", "approval-for-string", "cross-user"]) -async def test_functional_pending_response_requires_authorized_matching_type_and_user(kind: str) -> None: - stores = _Stores() - - def factory() -> FunctionalWorkflowAgent: - return _pending_string.build().as_agent() - - assert _types(await _collect(stores.server(factory)))[-1] == "response.completed" - checkpoint = await stores.checkpoints[("alice", "conversation")].get_latest( - workflow_name="functional-pending-string" - ) - assert checkpoint is not None - request_id = next(iter(checkpoint.pending_request_info_events)) - if kind == "approval-for-string": - approval_id = next(iter(stores.approvals["alice"].requests)) - items = [{"type": "mcp_approval_response", "approval_request_id": approval_id, "approve": True}] - elif kind == "text": - items = None - else: - items = [ - { - "type": "function_call_output", - "call_id": "wrong" if kind == "wrong-id" else request_id, - "output": "stolen", - } - ] - rejected = await _collect( - stores.server(factory), - _context("plain text", response="two", items=items), - user="bob" if kind == "cross-user" else "alice", - ) - if kind != "cross-user": - assert "pending functional workflow request" in _failure(rejected) - assert "hello:stolen" not in _text(rejected) - approved = await _collect( - stores.server(factory), - _context(response="three", items=[{"type": "function_call_output", "call_id": request_id, "output": "owner"}]), - ) - assert _text(approved) == "hello:owner" - - -@dataclass -class _FunctionalAnswer: - text: str - score: float - - -@pytest.mark.parametrize( - ("response_type", "value"), - [(bool, True), (bool, False), (str, "accepted"), (_FunctionalAnswer, {"text": "accepted", "score": 2})], - ids=["true", "false", "string", "structured"], -) -async def test_functional_response_coerces_normalized_results( - response_type: type, value: Any, monkeypatch: pytest.MonkeyPatch -) -> None: - received: list[Any] = [] - - @workflow(name="typed-functional") - async def typed(messages: list[Message], ctx: RunContext) -> str: - answer = await ctx.request_info(messages[0].text, response_type=response_type, request_id="answer") - received.append(answer) - return repr(answer) - - stores = _Stores() - - def factory() -> FunctionalWorkflowAgent: - return typed.build().as_agent() - - assert _types(await _collect(stores.server(factory)))[-1] == "response.completed" - # The wire converter stringifies outputs. Exercise core structured coercion on - # normalized framework content without introducing JSON parsing into the host. - monkeypatch.setattr( - "agent_framework_foundry_hosting._responses._items_to_messages", - AsyncMock(return_value=[Message("tool", [Content("function_result", call_id="answer", result=value)])]), - ) - resumed = await _collect(stores.server(factory), _context(response="two")) - assert _types(resumed)[-1] == "response.completed", resumed - expected = _FunctionalAnswer("accepted", 2.0) if response_type is _FunctionalAnswer else value - assert received == [expected] - assert type(received[0]) is response_type - if isinstance(received[0], _FunctionalAnswer): - assert type(received[0].score) is float - assert _text(resumed) == repr(expected) - - -@pytest.mark.parametrize("chain", [False, True]) -@pytest.mark.parametrize( - "invalid", - ["bool-string", "bool-number", "bool-null", "string-bool", "approval-string", "unknown", "duplicate"], -) -async def test_functional_invalid_batch_preserves_checkpoint_and_allows_retry( - chain: bool, invalid: str, monkeypatch: pytest.MonkeyPatch -) -> None: - executions: list[str] = [] - received: list[Any] = [] - - @workflow(name="typed-functional-batch") - async def typed(messages: list[Message], ctx: RunContext) -> str: - executions.append(messages[0].text) - answers = await asyncio.gather( - ctx.request_info("text?", response_type=str, request_id="text"), - ctx.request_info("approve?", response_type=bool, request_id="decision"), - ) - received.extend(answers) - return f"{answers[0]}:{answers[1]}" - - stores = _Stores() - - def factory() -> FunctionalWorkflowAgent: - return typed.build().as_agent() - - conversation = None if chain else "conversation" - source_id = conversation or "response-1" - assert ( - _types(await _collect(stores.server(factory), _context(conversation=conversation)))[-1] == "response.completed" - ) - storage = stores.checkpoints[("alice", source_id)] - checkpoint = await storage.get_latest(workflow_name="typed-functional-batch") - assert checkpoint is not None - assert set(checkpoint.pending_request_info_events) == {"text", "decision"} - checkpoint_before = json.dumps(checkpoint.to_dict(), default=lambda value: value.to_dict(), sort_keys=True) - sessions = stores.sessions["alice"] - session = await sessions.get(source_id) - assert session is not None - state_before = copy.deepcopy(session.state) - source_save = AsyncMock(wraps=storage.save) - session_save = AsyncMock(wraps=sessions.set) - monkeypatch.setattr(storage, "save", source_save) - monkeypatch.setattr(sessions, "set", session_save) - contents = [ - Content.from_function_result("text", result="accepted"), - Content("function_result", call_id="decision", result=False), - ] - if invalid.startswith("bool-"): - contents[1] = Content( - "function_result", - call_id="decision", - result={"bool-string": "false", "bool-number": 1, "bool-null": None}[invalid], - ) - elif invalid == "string-bool": - contents[0] = Content("function_result", call_id="text", result=False) - elif invalid == "approval-string": - contents[1] = Content("function_approval_response", id="decision", approved=cast(Any, "false")) - elif invalid == "unknown": - contents.append(Content.from_function_result("not-authorized", result="extra")) - else: - contents.append(Content.from_function_result("text", result="duplicate")) - converted = AsyncMock(return_value=[Message("tool", contents)]) - monkeypatch.setattr("agent_framework_foundry_hosting._responses._items_to_messages", converted) - - rejected = await _collect( - stores.server(factory), - _context(response="rejected", conversation=conversation), - previous="response-1" if chain else None, - ) - error = _failure(rejected) - assert ("authorized pending" if invalid in ("unknown", "duplicate") else "Response type mismatch") in error - assert executions == ["hello"] - assert received == [] - source_save.assert_not_awaited() - session_save.assert_not_awaited() - unchanged = await storage.get_latest(workflow_name="typed-functional-batch") - assert unchanged is not None - assert json.dumps(unchanged.to_dict(), default=lambda value: value.to_dict(), sort_keys=True) == checkpoint_before - unchanged_session = await sessions.get(source_id) - assert unchanged_session is not None - assert unchanged_session.state == state_before - if chain: - assert await sessions.get("rejected") is None - assert not await stores.checkpoints[("alice", "rejected")].list_checkpoint_ids( - workflow_name="typed-functional-batch" - ) - - converted.return_value = [ - Message( - "tool", - [ - Content.from_function_result("text", result="accepted"), - Content("function_approval_response", id="decision", approved=False), - ], - ) - ] - resumed = await _collect( - stores.server(factory), - _context(response="retry", conversation=conversation), - previous="response-1" if chain else None, - ) - assert _types(resumed)[-1] == "response.completed", resumed - assert _text(resumed) == "accepted:False" - assert received == ["accepted", False] - assert executions == ["hello", "hello"] - - -@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) -@pytest.mark.parametrize("decision", [False, True]) -async def test_functional_response_http_rejects_string_for_bool_then_accepts_approval( - stream: bool, decision: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - stores = _Stores() - server = stores.server(lambda: _pending_bool.build().as_agent(), store=InMemoryResponseProvider()) - transport = httpx.ASGITransport(app=server) - - async def post(client: httpx.AsyncClient, items: Any) -> dict[str, Any]: - response = await client.post( - "/responses", - json={"model": "test-model", "input": items, "conversation": "conversation", "stream": stream}, - ) - assert response.status_code == 200, response.text - if not stream: - return response.json() - events = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] - terminal = [event for event in events if event["type"] in ("response.completed", "response.failed")] - assert len(terminal) == 1, events - return terminal[0]["response"] - - async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: - first = await post(client, "private") - assert first["status"] == "completed", first - user = next(iter(stores.sessions)) - storage = stores.checkpoints[(user, "conversation")] - checkpoint = await storage.get_latest(workflow_name="functional-pending-bool") - assert checkpoint is not None - request_id = next(iter(checkpoint.pending_request_info_events)) - approval_id = next(iter(stores.approvals[user].requests)) - session_save = AsyncMock(wraps=stores.sessions[user].set) - checkpoint_save = AsyncMock(wraps=storage.save) - monkeypatch.setattr(stores.sessions[user], "set", session_save) - monkeypatch.setattr(storage, "save", checkpoint_save) - rejected = await post(client, [{"type": "function_call_output", "call_id": request_id, "output": "false"}]) - assert rejected["status"] == "failed", rejected - assert "Response type mismatch" in rejected["error"]["message"] - session_save.assert_not_awaited() - checkpoint_save.assert_not_awaited() - resumed = await post( - client, [{"type": "mcp_approval_response", "approval_request_id": approval_id, "approve": decision}] - ) - assert resumed["status"] == "completed", resumed - assert [ - part["text"] - for item in resumed["output"] - if item["type"] == "message" - for part in item["content"] - if part["type"] == "output_text" - ] == [f"private:{decision}"] - - -@pytest.mark.parametrize("stream", [False, True], ids=["nonstream", "stream"]) -async def test_functional_response_http_accepts_string_result(stream: bool) -> None: - stores = _Stores() - server = stores.server(lambda: _pending_string.build().as_agent(), store=InMemoryResponseProvider()) - transport = httpx.ASGITransport(app=server) - async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: - first = await client.post("/responses", json={"input": "private", "conversation": "conversation"}) - assert first.status_code == 200 - assert first.json()["status"] == "completed" - user = next(iter(stores.sessions)) - checkpoint = await stores.checkpoints[(user, "conversation")].get_latest( - workflow_name="functional-pending-string" - ) - assert checkpoint is not None - request_id = next(iter(checkpoint.pending_request_info_events)) - response = await client.post( - "/responses", - json={ - "input": [{"type": "function_call_output", "call_id": request_id, "output": "accepted"}], - "conversation": "conversation", - "stream": stream, - }, - ) - assert response.status_code == 200, response.text - if stream: - events = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] - assert _types(events)[-1] == "response.completed", events - assert _text(events) == "private:accepted" - result = events[-1]["response"] - else: - result = response.json() - assert result["status"] == "completed", result - assert result["output"][0]["content"][0]["text"] == "private:accepted" - - -async def test_graph_recovery_without_checkpoint_replays_original_input() -> None: - stores = _Stores() - context = _context("original") - context.is_recovery = True - server = stores.server(_graph, options=ResponsesServerOptions(resilient_background=True)) - events = await _collect(server, context) - assert _types(events)[-1] == "response.completed" - assert _text(events) == "1:original" - - -@pytest.mark.parametrize("functional", [False, True]) -async def test_workflow_final_session_save_failure_closes_request_resources( - functional: bool, monkeypatch: pytest.MonkeyPatch -) -> None: - events: list[str] = [] - stores = _Stores() - sessions = SessionStore() - save = AsyncMock(side_effect=[None, RuntimeError("workflow session save failed")]) - monkeypatch.setattr(sessions, "set", save) - stores.sessions["alice"] = sessions - - class OwnedGraph(WorkflowAgent): - async def __aenter__(self) -> Self: - events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - await asyncio.sleep(0) - events.append("exit") - - class OwnedFunctional(FunctionalWorkflowAgent): - async def __aenter__(self) -> Self: - events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - await asyncio.sleep(0) - events.append("exit") - - server = stores.server( - lambda: OwnedFunctional(_functional.build()) if functional else OwnedGraph(_graph().workflow) - ) - assert "workflow session save failed" in _failure(await _collect(server)) - assert events == ["enter", "exit"] - assert save.await_count == 2 - assert not server._scope_locks._entries - - -async def test_resource_exit_failure_replaces_success_with_failure() -> None: - events: list[str] = [] - - class FailingExit(_OwnedAgent): - async def __aexit__(self, *args: Any) -> None: - await super().__aexit__(*args) - raise RuntimeError("resource exit failed") - - server = _Stores().server(lambda: FailingExit(events)) - result = await _collect(server) - assert _text(result) == "firstsecond" - assert "resource exit failed" in _failure(result) - assert events == ["enter", "run", "iterator closed", "exit"] - assert not server._scope_locks._entries - - -async def test_cancelling_streaming_task_persists_session_and_closes_iterator(monkeypatch: pytest.MonkeyPatch) -> None: - events: list[str] = [] - emitted = asyncio.Event() - stores = _Stores() - sessions = SessionStore() - save = AsyncMock(wraps=sessions.set) - monkeypatch.setattr(sessions, "set", save) - stores.sessions["alice"] = sessions - server = stores.server(lambda: _OwnedAgent(events, wait=asyncio.Event())) - - async def consume() -> None: - with _platform(): - stream = cast(AsyncGenerator[Any, None], server._handle_response(_request(), _context(), asyncio.Event())) - async with aclosing(stream): - async for event in stream: - if isinstance(event, Mapping) and event.get("type") == "response.output_text.delta": - emitted.set() - - consumer = asyncio.create_task(consume()) - await asyncio.wait_for(emitted.wait(), 2) - consumer.cancel() - with pytest.raises(asyncio.CancelledError): - await consumer - assert events == ["enter", "run", "iterator closed", "exit"] - save.assert_awaited_once() - assert await sessions.get("conversation") is not None - assert not server._scope_locks._entries - - -async def test_anyio_cancellation_finishes_blocked_iterator_cleanup_before_owner_exit() -> None: - events: list[str] = [] - blocked = asyncio.Event() - scope = CancelScope() - - class TaskGroupAgent(_OwnedAgent): - async def __aenter__(self) -> Self: - await super().__aenter__() - self.group = create_task_group() - await self.group.__aenter__() - return self - - async def __aexit__(self, *args: Any) -> None: - await self.group.__aexit__(*args) - await super().__aexit__(*args) - - def run(self, messages: Any = None, *, stream: bool = False, **kwargs: Any) -> Any: - assert stream - - async def updates() -> AsyncIterator[AgentResponseUpdate]: - try: - events.append("run") - yield AgentResponseUpdate(role="assistant", contents=[Content.from_text("first")]) - blocked.set() - await asyncio.Event().wait() - finally: - events.append("cleanup started") - await asyncio.sleep(0) - events.append("cleanup finished") - - return ResponseStream(updates(), finalizer=AgentResponse.from_updates) - - server = _Stores().server(lambda: TaskGroupAgent(events)) - - async def consume() -> None: - with scope: - await _collect(server) - - consumer = asyncio.create_task(consume()) - await asyncio.wait_for(blocked.wait(), 2) - scope.cancel() - await asyncio.wait_for(consumer, 2) - assert not server._scope_locks._entries - assert events == ["enter", "run", "cleanup started", "cleanup finished", "exit"] - - -async def test_checkpoint_preparation_failure_preserves_graph_runtime_behavior( - caplog: pytest.LogCaptureFixture, -) -> None: - stores = _Stores() - calls: list[str] = [] - snapshots: list[int] = [] - - class FailingCheckpoint(_Counter): - @handler - async def count_message( - self, - messages: list[Message], - ctx: WorkflowContext[Never, str], # type: ignore[valid-type] - ) -> None: - calls.extend(message.text for message in messages) - await super().count_message(messages, ctx) - - async def on_checkpoint_save(self) -> dict[str, Any]: - snapshots.append(self.count) - if self.count: - raise RuntimeError("executor checkpoint preparation failed") - return await super().on_checkpoint_save() - - def factory() -> WorkflowAgent: - return WorkflowBuilder(name="checkpoint-preparation", start_executor=FailingCheckpoint()).build().as_agent() - - server = stores.server(factory) - result = await _collect(server) - assert _types(result)[-1] == "response.completed" - assert _text(result) == "1:hello" - assert "does not fail the workflow run" in caplog.text - storage = stores.checkpoints[("alice", "conversation")] - checkpoint = await storage.get_latest(workflow_name="checkpoint-preparation") - assert checkpoint is not None - assert checkpoint.iteration_count == 0 - assert 0 in snapshots and 1 in snapshots - session = await stores.sessions["alice"].get("conversation") - assert session is not None - assert session.state["_foundry_responses_workflow"]["completed"] is True - assert "checkpoint_failed" not in session.state["_foundry_responses_workflow"] - assert calls == ["hello"] - assert not server._scope_locks._entries - - -@pytest.mark.parametrize("functional", [False, True]) -async def test_checkpoint_save_failure_preserves_runtime_behavior_and_closes_owner( - functional: bool, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture -) -> None: - events: list[str] = [] - storage = InMemoryCheckpointStorage() - save = AsyncMock(side_effect=RuntimeError("checkpoint save failed")) - monkeypatch.setattr(storage, "save", save) - stores = _Stores() - stores.checkpoints[("alice", "conversation")] = storage - - class OwnedGraph(WorkflowAgent): - async def __aenter__(self) -> Self: - events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - events.append("exit") - - class OwnedFunctional(FunctionalWorkflowAgent): - async def __aenter__(self) -> Self: - events.append("enter") - return self - - async def __aexit__(self, *args: Any) -> None: - events.append("exit") - - server = stores.server( - lambda: OwnedFunctional(_functional.build()) if functional else OwnedGraph(_graph().workflow) - ) - result = await _collect(server) - assert events == ["enter", "exit"] - assert save.await_count > 0 - assert not server._scope_locks._entries - if functional: - assert "checkpoint save failed" in _failure(result) - else: - assert _types(result)[-1] == "response.completed" - assert _text(result) == "1:hello" - assert "does not fail the workflow run" in caplog.text - assert await storage.get_latest(workflow_name="functional" if functional else "counter") is None - - -async def test_interrupted_functional_step_does_not_restart_with_new_input() -> None: - stores = _Stores() - calls: list[str] = [] - - @step - async def saved() -> str: - return "saved" - - @workflow(name="interrupted-functional") - async def interrupted(messages: list[Message]) -> str: - calls.append(messages[0].text) - await saved() - raise RuntimeError("failed after step") - - def factory() -> FunctionalWorkflowAgent: - return interrupted.build().as_agent() - - assert "failed after step" in _failure(await _collect(stores.server(factory))) - assert "interrupted functional workflow" in _failure( - await _collect(stores.server(factory), _context("next", response="two")) - ) - assert calls == ["hello"] - - -@pytest.mark.parametrize("functional", [False, True]) -@pytest.mark.parametrize("stream", [False, True]) -async def test_fresh_http_request_preserves_another_users_pending_request(functional: bool, stream: bool) -> None: - stores = _Stores() - completions: list[tuple[str | None, str]] = [] - - class Pending(Executor): - @handler - async def ask(self, messages: list[Message], ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - await ctx.request_info(messages[0].text, response_type=str, request_id="answer") - - @response_handler - async def answer(self, original_request: str, response: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - completions.append((get_request_context().user_id, original_request)) - await ctx.yield_output(f"{original_request}:{response}") - - @workflow(name="http-pending") - async def pending(messages: list[Message], ctx: RunContext) -> str: - answer = await ctx.request_info("answer?", response_type=str, request_id="answer") - completions.append((get_request_context().user_id, messages[0].text)) - return f"{messages[0].text}:{answer}" - - def factory() -> WorkflowAgent | FunctionalWorkflowAgent: - if functional: - return pending.build().as_agent() - return WorkflowBuilder(name="http-pending", start_executor=Pending(id="pending")).build().as_agent() - - server = stores.server(factory, store=InMemoryResponseProvider()) - server.config.is_hosted = True - async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: - first = await _post_http_response(client, "first", user="user-a", stream=stream) - assert first["status"] == "completed", first - result = [{"type": "function_call_output", "call_id": "answer", "output": "accepted"}] - await _post_http_response(client, result, user="user-b", stream=stream) - assert completions == [] - resumed = await _post_http_response(client, result, user="user-a", stream=stream, previous=first["id"]) - assert resumed["status"] == "completed", resumed - assert [ - part["text"] - for item in resumed["output"] - if item["type"] == "message" - for part in item["content"] - if part["type"] == "output_text" - ] == ["first:accepted"] - assert completions == [("user-a", "first")] - - -@pytest.mark.parametrize("functional", [False, True]) -async def test_cross_user_previous_response_cannot_resume_another_users_pending_request(functional: bool) -> None: - stores = _Stores() - factory = (lambda: _pending_string.build().as_agent()) if functional else _pending_graph - first = await _collect(stores.server(factory), _context("alice-private", conversation=None)) - assert _types(first)[-1] == "response.completed" - checkpoint = await stores.checkpoints[("alice", "response-1")].get_latest( - workflow_name="functional-pending-string" if functional else "pending" - ) - assert checkpoint is not None - request_id = next(iter(checkpoint.pending_request_info_events)) - result = {"type": "function_call_output", "call_id": request_id, "output": "approved"} - rejected = await _collect( - stores.server(factory), - _context(response="two", conversation=None, items=[result]), - user="bob", - previous="response-1", - ) - assert "checkpoint" in _failure(rejected) - approved = await _collect( - stores.server(factory), _context(response="three", conversation=None, items=[result]), previous="response-1" - ) - assert _text(approved) == "alice-private:approved" - - -class _RecoveryStart(Executor): - @handler - async def start(self, messages: list[Message], ctx: WorkflowContext[str, str]) -> None: - await ctx.yield_output(f"first:{messages[0].text}") - await ctx.send_message(messages[0].text) - - -class _RecoveryEnd(Executor): - @handler - async def end(self, text: str, ctx: WorkflowContext[Never, str]) -> None: # type: ignore[valid-type] - await ctx.yield_output(f"last:{text}") - - -def _recovery_graph() -> WorkflowAgent: - start = _RecoveryStart(id="start") - end = _RecoveryEnd(id="end") - return WorkflowBuilder(name="recovery", start_executor=start).add_edge(start, end).build().as_agent() - - -@pytest.mark.parametrize("saved_input", [False, True]) -async def test_graph_recovery_preserves_explicit_outer_history_for_continuation(saved_input: bool) -> None: - stores = _Stores() - transcripts: list[list[str]] = [] - agents: list[WorkflowAgent] = [] - original = Message("user", ["original"]) - - class Start(Executor): - @handler - async def start(self, messages: list[Message], ctx: WorkflowContext[str]) -> None: - transcripts.append([message.text for message in messages]) - await ctx.send_message(messages[-1].text) - - def factory() -> WorkflowAgent: - start = Start(id="start") - end = _RecoveryEnd(id="end") - agent = ( - WorkflowBuilder(name="recovery-history", start_executor=start) - .add_edge(start, end) - .build() - .as_agent(context_providers=[InMemoryHistoryProvider(source_id="outer-history")]) - ) - agents.append(agent) - return agent - - interrupted = factory() - storage = InMemoryCheckpointStorage() - iterator = interrupted.workflow.run([original], stream=True, checkpoint_storage=storage).__aiter__() - try: - async for event in iterator: - if event.type == "superstep_completed": - break - finally: - await close_run_iterator(iterator) - checkpoint = await storage.get_latest(workflow_name="recovery-history") - assert checkpoint is not None - assert checkpoint.iteration_count == 1 - assert transcripts == [["original"]] - stores.checkpoints[("alice", "conversation")] = storage - session = AgentSession() - session.state["_foundry_responses_workflow"] = { - "name": "recovery-history", - "kind": "graph", - "completed": False, - } - if saved_input: - await InMemoryHistoryProvider(source_id="outer-history").save_messages( - session.session_id, [original], state=session.state.setdefault("outer-history", {}) - ) - stores.sessions["alice"] = SessionStore() - await stores.sessions["alice"].set("conversation", session) - - context = _context("original") - context.is_recovery = True - recovered = await _collect( - stores.server(factory, options=ResponsesServerOptions(resilient_background=True)), context - ) - assert _types(recovered)[-1] == "response.completed" - assert _text(recovered) == "last:original" - assert transcripts == [["original"]] - recovered_session = await stores.sessions["alice"].get("conversation") - assert recovered_session is not None - history = await InMemoryHistoryProvider(source_id="outer-history").get_messages( - recovered_session.session_id, state=recovered_session.state.get("outer-history") - ) - assert [(message.role, message.text) for message in history] == [ - ("user", "original"), - ("assistant", "last:original"), - ] - - continued = await _collect(stores.server(factory), _context("next", response="two")) - assert _types(continued)[-1] == "response.completed" - assert _text(continued) == "last:next" - assert transcripts == [["original"], ["original", "last:original", "next"]] - continued_session = await stores.sessions["alice"].get("conversation") - assert continued_session is not None - history = await InMemoryHistoryProvider(source_id="outer-history").get_messages( - continued_session.session_id, state=continued_session.state.get("outer-history") - ) - assert [message.text for message in history] == ["original", "last:original", "next", "last:next"] - assert len({id(agent.workflow) for agent in agents}) == 3 - - -@pytest.mark.parametrize("snapshot_kind", ["latest", "empty", "partial"]) -async def test_graph_recovery_selects_latest_or_response_paired_checkpoint( - snapshot_kind: str, monkeypatch: pytest.MonkeyPatch -) -> None: - stores = _Stores() - options = ResponsesServerOptions(resilient_background=True) - server = stores.server(_recovery_graph, options=options) - snapshots: list[Any] = [] - with _platform(): - async for event in server._handle_response(_request(), _context("original"), asyncio.Event()): - if isinstance(event, ResponseCheckpointEvent): - snapshots.append(copy.deepcopy(event.response)) - assert snapshots - storage = stores.checkpoints[("alice", "conversation")] - latest = await storage.get_latest(workflow_name="recovery") - assert latest is not None - load = AsyncMock(wraps=storage.load) - monkeypatch.setattr(storage, "load", load) - context = _context("must-not-replay") - context.is_recovery = True - if snapshot_kind != "latest": - snapshot = ( - next(snapshot for snapshot in snapshots if snapshot.get("output")) - if snapshot_kind == "partial" - else snapshots[0] - ) - # Pair the response with the matching workflow state explicitly, independent of - # how far the background iterator advanced while the host consumed its output. - checkpoints = await storage.list_checkpoints(workflow_name="recovery") - paired_checkpoint = next( - checkpoint - for checkpoint in checkpoints - if checkpoint.iteration_count == (1 if snapshot_kind == "partial" else 0) - ) - saved_response = ResponseEventStream(response=snapshot) - saved_response.internal_metadata["_last_checkpoint_id"] = paired_checkpoint.checkpoint_id - context.persisted_response = saved_response.checkpoint().response - events = await _collect(stores.server(_recovery_graph, options=options), context) - assert _types(events)[-1] == "response.completed", events - assert "must-not-replay" not in _text(events) - load.assert_awaited_once() - if snapshot_kind != "latest": - assert load.await_args is not None - assert load.await_args.args[0] != latest.checkpoint_id - assert context.persisted_response is not None - paired_checkpoint_id = ResponseEventStream(response=context.persisted_response).internal_metadata[ - "_last_checkpoint_id" - ] - load.assert_awaited_once_with(paired_checkpoint_id) - assert "original" in _text(events) - if snapshot_kind == "partial": - assert _text(events) == "last:original" - final_text = "".join( - content.get("text", "") - for item in events[-1]["response"]["output"] - for content in item.get("content", []) - ) - assert final_text == "first:originallast:original" - else: - load.assert_awaited_once_with(latest.checkpoint_id) diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 6449c818ce5..7ae5088a4d9 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -1,6 +1,8 @@ # Copyright (c) Microsoft. All rights reserved. +import os from collections.abc import Callable from dataclasses import dataclass +from pathlib import Path from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -19,8 +21,6 @@ FoundryCheckpointStore, FoundryFunctionApprovalStore, FunctionApprovalStoreProvider, - _InvocationsAgentSessionStoreProvider, - _InvocationsCheckpointStoreProvider, ) @@ -74,6 +74,13 @@ def _platform_context(call_id: str = "call-1", user_id: str = "user-1") -> Found return FoundryAgentRequestContext(call_id=call_id, user_id=user_id) +def test_local_agentserver_state_root_is_test_scoped(tmp_path: Path) -> None: + """Independent tests must not share the local fallback state directory.""" + state_root = Path(os.environ["AGENTSERVER_STATE_ROOT"]) + + assert state_root.is_relative_to(tmp_path) + + def test_storage_providers_use_public_abstraction() -> None: assert issubclass(CheckpointStoreProvider, ContextScopedStoreProvider) assert not issubclass(CheckpointStoreProvider, StoreProvider) @@ -81,30 +88,6 @@ def test_storage_providers_use_public_abstraction() -> None: assert issubclass(AgentSessionStoreProvider, StoreProvider) -@pytest.mark.parametrize("is_hosted", [False, True]) -async def test_invocations_namespaces_cannot_overlap_responses_records(is_hosted: bool) -> None: - store = _store() - config = _config(is_hosted=is_hosted) - context = _platform_context() - with patch( - "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", - new=AsyncMock(return_value=store), - ) as get_or_create: - for checkpoint_provider in (CheckpointStoreProvider(), _InvocationsCheckpointStoreProvider()): - await checkpoint_provider.get_store(config=config, context_id="same-id", platform_context=context).save( - _checkpoint("same-checkpoint") - ) - for session_provider in (AgentSessionStoreProvider(), _InvocationsAgentSessionStoreProvider()): - await session_provider.get_store(config=config, platform_context=context).set("same-id", AgentSession()) - assert [call.args[0] for call in get_or_create.await_args_list] == [ - "checkpoints/same-id", - "invocations_checkpoints/same-id", - "agent_sessions", - "invocations_agent_sessions", - ] - assert all(call.kwargs == {"user_isolation": True} for call in get_or_create.await_args_list) - - async def test_save_uses_context_scoped_store() -> None: store = _store() checkpoint = _checkpoint("checkpoint-1") 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 41f5233978c..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,14 +21,9 @@ Each user message re-runs the workflow from the trigger. Because `Workflow.as_ag ### Agent Hosting -[`main.py`](main.py) supplies `ResponsesHostServer` with an `agent_factory`. Each request builds three new `Agent` -instances, registers them with a new `WorkflowFactory` so the YAML's `InvokeAzureAgent` actions can resolve them -by name, loads a new workflow, and wraps it with `.as_agent(...)`. The host restores the authorized conversation's -checkpoint into that instance when continuing a conversation. - -The `FoundryChatClient` is opened once in `main` and closed when the host exits. It is shared across requests, -but the agents and workflow executors are not. The YAML's stable identifiers allow later factory-created -workflows to restore earlier checkpoints. +[`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 438a2c2e944..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 @@ -1,6 +1,5 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio import os from pathlib import Path from typing import Any, Literal @@ -87,10 +86,8 @@ class TriageResponse(BaseModel): # --- Host setup ------------------------------------------------------------------ -def create_workflow_agent(client: FoundryChatClient) -> WorkflowAgent: - """Rebuild the YAML workflow and its agents for the current request.""" - workflow_path = Path(__file__).parent / "workflow.yaml" - +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, @@ -134,18 +131,16 @@ def create_workflow_agent(client: FoundryChatClient) -> WorkflowAgent: ) -async def main() -> None: - """Share only the model client while each request gets a new workflow.""" - with DefaultAzureCredential() as credential: - client = FoundryChatClient( - project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], - model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], - credential=credential, - ) - async with client.project_client, client.client: - server = ResponsesHostServer(agent_factory=lambda: create_workflow_agent(client)) - await server.run_async() +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__": - asyncio.run(main()) + 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 45fcd706439..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 @@ -16,21 +16,18 @@ The workflow has three executors (see [main.py](main.py)): user's message. If no valid target is found, the workflow yields an error message instead of counting down. - **`CountdownExecutor`** decrements the target through a self-loop, sleeping for a second and yielding an output on each tick, to simulate a long-running operation. -- **`CompleteExecutor`** yields the workflow's final `"Countdown complete."` output once the countdown reaches zero. +- **`complete`** yields the workflow's final `"Countdown complete."` output once the countdown reaches zero. ### Agent Hosting 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). -The host receives an `agent_factory`, which builds all three executors and the target-extraction agent anew for -each request, including recovery. The workflow name and executor IDs are stable across factory calls. The model -client is owned by `main`, shared across requests, and closed when the host exits. Conversation state comes from -the authorized checkpoint, not from a workflow instance left over from another request. - ## Running the Agent Host Follow the instructions in the [Running the Agent Host Locally](../../README.md#running-the-agent-host-locally) section of the README in the parent directory to run the agent host. 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 58775a9a4fe..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, Workflow, WorkflowBuilder, WorkflowContext, 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 @@ -70,15 +70,10 @@ async def countdown(self, target: int, ctx: WorkflowContext[int | str, str]) -> await ctx.send_message(target - 1, target_id=self.id) -class CompleteExecutor(Executor): +@executor(id="complete") +async def complete(message: str, ctx: WorkflowContext[Never, str]) -> None: """Yield the workflow's completion output.""" - - def __init__(self, id: str = "complete") -> None: - super().__init__(id=id) - - @handler - async def complete(self, message: str, ctx: WorkflowContext[Never, str]) -> None: - await ctx.yield_output(message) + await ctx.yield_output(message) def build_workflow(client: FoundryChatClient) -> Workflow: @@ -93,7 +88,6 @@ def build_workflow(client: FoundryChatClient) -> Workflow: ) start = StartExecutor(target_agent) countdown = CountdownExecutor() - complete = CompleteExecutor() return ( WorkflowBuilder(name="countdown-workflow", start_executor=start, output_from="all") @@ -104,23 +98,21 @@ def build_workflow(client: FoundryChatClient) -> Workflow: ) -async def main() -> None: +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 - with DefaultAzureCredential() as credential: - client = FoundryChatClient( - project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], - model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], - credential=credential, - ) - async with client.project_client, client.client: - server = ResponsesHostServer( - agent_factory=lambda: build_workflow(client).as_agent(name="countdown-workflow"), - options=ResponsesServerOptions(resilient_background=True), - log_level="DEBUG", - ) - await server.run_async() + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=DefaultAzureCredential(), + ) + server = ResponsesHostServer( + agent=lambda: build_workflow(client).as_agent(name="countdown-workflow"), + options=ResponsesServerOptions(resilient_background=True), + log_level="DEBUG", + ) + server.run() if __name__ == "__main__": - asyncio.run(main()) + main() 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 b53fea7bb6c..325cb5ea255 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,13 +16,10 @@ See [main.py](main.py) for the full implementation. ### Agent Hosting -The workflow is exposed via `.as_agent()` and supplied through `ResponsesHostServer(agent_factory=...)`. -Every request creates fresh agents, executors, and workflow state. The host restores the current conversation's -checkpoint into that new instance when continuing an existing conversation. - -The model client is opened once in `main` and closed when the host exits. Only that client is shared across requests; -the mutable agents and executors are created inside the factory. The workflow name and executor names remain stable -so newly created workflows can load earlier checkpoints. +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 workflow, executors, and agents for each request while reusing the application-owned +`FoundryChatClient`. ## 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 204b6d7a7ea..54fb2fa283f 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 @@ -1,6 +1,5 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio import os from agent_framework import Agent, AgentExecutor, WorkflowAgent, WorkflowBuilder @@ -13,7 +12,7 @@ def create_workflow_agent(client: FoundryChatClient) -> WorkflowAgent: - """Create fresh agents and executors for one request, using the host-owned client.""" + """Create a fresh workflow agent for one hosted request.""" writer_agent = Agent( client=client, instructions=("You are an excellent slogan writer. You create new slogans based on the given topic."), @@ -59,18 +58,16 @@ def create_workflow_agent(client: FoundryChatClient) -> WorkflowAgent: ) -async def main() -> None: - """Keep the model client open while request factories create independent workflows.""" - with DefaultAzureCredential() as credential: - client = FoundryChatClient( - project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], - model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], - credential=credential, - ) - async with client.project_client, client.client: - server = ResponsesHostServer(agent_factory=lambda: create_workflow_agent(client)) - await server.run_async() +def main() -> None: + client = FoundryChatClient( + project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], + credential=DefaultAzureCredential(), + ) + + server = ResponsesHostServer(agent=lambda: create_workflow_agent(client)) + server.run() if __name__ == "__main__": - asyncio.run(main()) + main() From 187592c7f30e84167a18a48b45adb75a4445d1d8 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Thu, 17 Sep 2026 12:14:06 +0100 Subject: [PATCH 06/11] Python: Fix ResponseStream test typing --- python/packages/core/tests/core/test_types.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index eab19003e6f..4a7615b5912 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -4203,7 +4203,9 @@ async def updates() -> AsyncIterable[ChatResponseUpdate]: finally: events.append("iterator") - stream = ResponseStream(updates(), cleanup_hooks=[lambda: events.append("cleanup")]) + stream: ResponseStream[ChatResponseUpdate, Sequence[ChatResponseUpdate]] = ResponseStream( + updates(), cleanup_hooks=[lambda: events.append("cleanup")] + ) await anext(stream) await stream.aclose() @@ -4222,7 +4224,9 @@ async def updates() -> AsyncIterable[ChatResponseUpdate]: finally: events.append("iterator") - inner = ResponseStream(updates(), cleanup_hooks=[lambda: events.append("inner")]) + 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) From 83cfcdae983952d9bd9c9f2011b9885ac70b5f4c Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Thu, 17 Sep 2026 12:56:52 +0100 Subject: [PATCH 07/11] Python: Align request agent validation and stream cleanup --- python/packages/core/agent_framework/_types.py | 14 +++++++++----- python/packages/core/tests/core/test_types.py | 10 +++++----- .../_agent_source.py | 8 ++++++++ .../_invocations.py | 3 ++- .../_responses.py | 18 +++++++++++++----- .../foundry_hosting/tests/test_invocations.py | 5 +++++ .../foundry_hosting/tests/test_responses.py | 5 +++++ 7 files changed, 47 insertions(+), 16 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index db2600d5589..eef2e504259 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -3520,16 +3520,20 @@ async def __anext__(self) -> UpdateT: update = await update return await self._record_update(update) - async def aclose(self) -> None: + async def close(self) -> None: """Close the active iterator and run cleanup hooks. This method is idempotent and also closes nested ``ResponseStream`` wrappers. """ try: - if self._iterator is not None: - close = getattr(self._iterator, "aclose", None) - if close is not None: - await close() + 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() diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 4a7615b5912..80843c5e0f8 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -4192,7 +4192,7 @@ async def async_hook(update: ChatResponseUpdate) -> ChatResponseUpdate: class TestResponseStreamCleanupHooks: """Tests for cleanup hooks (after stream consumption, before finalizer).""" - async def test_aclose_closes_iterator_and_runs_cleanup_once(self) -> None: + 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] = [] @@ -4208,12 +4208,12 @@ async def updates() -> AsyncIterable[ChatResponseUpdate]: ) await anext(stream) - await stream.aclose() - await stream.aclose() + await stream.close() + await stream.close() assert events == ["iterator", "cleanup"] - async def test_aclose_closes_wrapped_stream(self) -> None: + async def test_close_closes_wrapped_stream(self) -> None: """Closing a wrapper releases the concrete inner stream.""" events: list[str] = [] @@ -4230,7 +4230,7 @@ async def updates() -> AsyncIterable[ChatResponseUpdate]: outer = inner.map(lambda update: update, _combine_updates).with_cleanup_hook(lambda: events.append("outer")) await anext(outer) - await outer.aclose() + await outer.close() assert events == ["iterator", "inner", "outer"] 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 index 8722d9757f5..50886d69cbc 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py @@ -15,6 +15,14 @@ def is_agent(value: object) -> TypeGuard[SupportsAgentRun]: return hasattr(value, "run") and hasattr(value, "create_session") +def validate_agent_source(source: object) -> None: + if is_agent(source): + return + if callable(source): + return + raise TypeError("agent must be an agent instance or a zero-argument callable that creates one.") + + async def resolve_agent(source: AgentSource) -> SupportsAgentRun: """Resolve an agent instance or request-scoped agent factory.""" if is_agent(source): 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 f38ab174193..f0cd102b636 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -11,7 +11,7 @@ from starlette.responses import Response, StreamingResponse from typing_extensions import Any, AsyncGenerator -from ._agent_source import resolve_agent +from ._agent_source import resolve_agent, validate_agent_source from ._feature_usage import FeatureIndex @@ -37,6 +37,7 @@ 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 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 e2535af5136..6f7b9ebcf29 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -34,6 +34,7 @@ InMemoryHistoryProvider, Message, RawAgent, + ResponseStream, SessionStore, SupportsAgentRun, UsageDetails, @@ -82,7 +83,7 @@ from mcp import McpError from typing_extensions import Any -from ._agent_source import is_agent, resolve_agent +from ._agent_source import is_agent, resolve_agent, validate_agent_source from ._feature_usage import FeatureIndex from ._state_store import ( AgentSessionStoreProvider, @@ -211,9 +212,13 @@ async def _drive(self) -> None: return await self._queue.put(item) finally: - close = getattr(self._iterator, "aclose", None) - if close is not None: - await close() + 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: @@ -547,6 +552,7 @@ 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) resolved_agent = agent if is_agent(agent) else None configuration = ( @@ -657,7 +663,7 @@ async def _handle_response( configuration, resources, ) - async with aclosing(inner): + try: async for event in inner: if isinstance(event, Mapping) and event.get("type") in ( "response.completed", @@ -667,6 +673,8 @@ async def _handle_response( terminal_event = event else: yield event + finally: + await inner.aclose() if terminal_event is not None: yield terminal_event diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index c322a55be3f..eab52f70b0b 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -153,6 +153,11 @@ 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) + # endregion diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index d0a9d8b9fa5..2276cbde45e 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -853,6 +853,11 @@ 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) + def test_init_basic(self) -> None: agent = _make_agent( response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])]) From dad3948bca7d72f3a275a4a384c9b6b1bcf9d216 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:19:18 +0100 Subject: [PATCH 08/11] Python: Manage request-scoped Invocations agents --- .../_invocations.py | 35 ++++++-- .../foundry_hosting/tests/test_invocations.py | 83 +++++++++++++++++++ 2 files changed, 110 insertions(+), 8 deletions(-) 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 f0cd102b636..0cf7a6e97e0 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -2,8 +2,9 @@ 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 @@ -11,7 +12,7 @@ from starlette.responses import Response, StreamingResponse from typing_extensions import Any, AsyncGenerator -from ._agent_source import resolve_agent, validate_agent_source +from ._agent_source import is_agent, resolve_agent, validate_agent_source from ._feature_usage import FeatureIndex @@ -41,6 +42,7 @@ def __init__( 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) @@ -76,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: @@ -101,14 +111,22 @@ async def _handle_invoke(self, request: Request) -> Response: session = AgentSession(session_id=session_id) self._sessions[partition_key] = session - agent = await resolve_agent(self._agent) - if stream: async def stream_response() -> AsyncGenerator[str]: - async for update in 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(), @@ -116,5 +134,6 @@ async def stream_response() -> AsyncGenerator[str]: headers={"Cache-Control": "no-cache", "Connection": "keep-alive"}, ) - response = await 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/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index eab52f70b0b..d08e0bf19a8 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, @@ -232,6 +274,47 @@ 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] = [] From 9714c4eb4b114262570e4d4ec47142f7d820b8ea Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:08:13 +0100 Subject: [PATCH 09/11] Python: Validate request agent callable shape --- .../agent_framework_foundry_hosting/_agent_source.py | 11 +++++++---- .../foundry_hosting/tests/test_invocations.py | 11 +++++++++++ .../packages/foundry_hosting/tests/test_responses.py | 7 +++++++ 3 files changed, 25 insertions(+), 4 deletions(-) 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 index 50886d69cbc..2335e141490 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py @@ -12,15 +12,18 @@ def is_agent(value: object) -> TypeGuard[SupportsAgentRun]: - return hasattr(value, "run") and hasattr(value, "create_session") + return not inspect.isclass(value) and hasattr(value, "run") and hasattr(value, "create_session") def validate_agent_source(source: object) -> None: if is_agent(source): return - if callable(source): - return - raise TypeError("agent must be an agent instance or a zero-argument callable that creates one.") + 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: diff --git a/python/packages/foundry_hosting/tests/test_invocations.py b/python/packages/foundry_hosting/tests/test_invocations.py index d08e0bf19a8..250799849a6 100644 --- a/python/packages/foundry_hosting/tests/test_invocations.py +++ b/python/packages/foundry_hosting/tests/test_invocations.py @@ -200,6 +200,17 @@ 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 diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 2276cbde45e..63f5df7e7f5 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -858,6 +858,13 @@ 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")])]) From dd72e86fd453b6e7ece49330621ee117a55ed738 Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:04:50 +0100 Subject: [PATCH 10/11] Python: Address request agent review feedback --- .../_agent_source.py | 2 +- .../foundry_hosting/tests/test_responses.py | 64 ++++++++++++++++--- .../responses/workflows/main.py | 16 ++--- 3 files changed, 64 insertions(+), 18 deletions(-) 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 index 2335e141490..5f9509b3fef 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_agent_source.py @@ -12,7 +12,7 @@ def is_agent(value: object) -> TypeGuard[SupportsAgentRun]: - return not inspect.isclass(value) and hasattr(value, "run") and hasattr(value, "create_session") + return not inspect.isclass(value) and isinstance(value, SupportsAgentRun) def validate_agent_source(source: object) -> None: diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 63f5df7e7f5..fb79cb7a619 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -141,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, @@ -153,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" @@ -165,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: @@ -209,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, @@ -3535,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" @@ -3547,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] @@ -4751,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" @@ -5253,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" @@ -5265,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 @@ -5304,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" @@ -5316,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 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 54fb2fa283f..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 @@ -11,8 +11,14 @@ load_dotenv() -def create_workflow_agent(client: FoundryChatClient) -> WorkflowAgent: +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"], + credential=DefaultAzureCredential(), + ) + writer_agent = Agent( client=client, instructions=("You are an excellent slogan writer. You create new slogans based on the given topic."), @@ -59,13 +65,7 @@ def create_workflow_agent(client: FoundryChatClient) -> WorkflowAgent: def main() -> None: - client = FoundryChatClient( - project_endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], - model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], - credential=DefaultAzureCredential(), - ) - - server = ResponsesHostServer(agent=lambda: create_workflow_agent(client)) + server = ResponsesHostServer(agent=create_workflow_agent) server.run() From ea10505b49e83c32400f1fabfef86764e4dab73d Mon Sep 17 00:00:00 2001 From: Roger Barreto <19890735+RogerBarreto@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:19:54 +0100 Subject: [PATCH 11/11] docs: Clarify request-scoped Foundry agent creation --- python/packages/foundry_hosting/README.md | 3 ++- .../foundry-hosted-agents/responses/workflows/README.md | 3 +-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index 8bb7c99eb7e..a7a59803df4 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -5,7 +5,8 @@ This package provides the integration of Agent Framework agents and workflows wi ## 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 and is invoked once for each request: +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) 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 325cb5ea255..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 @@ -18,8 +18,7 @@ See [main.py](main.py) for the full implementation. 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 workflow, executors, and agents for each request while reusing the application-owned -`FoundryChatClient`. +callable that builds a fresh `FoundryChatClient`, workflow, executors, and agents for each request. ## Running the Agent Host