diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index f6f61f6de83..913328e255e 100644 --- a/python/packages/foundry_hosting/README.md +++ b/python/packages/foundry_hosting/README.md @@ -422,3 +422,100 @@ locally. Stored checkpoints are scoped under `checkpoints`. `ResponsesHostServer` persists function approvals durably. By default, it uses the `FoundryFunctionApprovalStore`, backed by Foundry storage when hosted and file-based storage locally. Stored approvals are scoped under `function_approvals`. + +## Native Invocations workflows + +Host a native workflow with an explicit application parser rather than wrapping it +with `workflow.as_agent()`: + +```python +server = InvocationsHostServer( + workflow=build_workflow, + parse_request=parse_request, + checkpoint_store_provider=CheckpointStoreProvider( + allowed_checkpoint_types=["my_app:Ticket", "my_app:TicketState"] + ), +) +``` + +The sync/async parser receives the raw Starlette `Request` and must return either +`WorkflowTurn(input=typed_input, stream=...)` or +`WorkflowTurn(responses={request_id: typed_reply, ...}, stream=...)`. There is no +default workflow JSON schema. Validate your application input and reply types, +including duplicate fields or decisions, in the parser; the host additionally +validates the entire reply batch against its exact pending checkpoint. + +A sync/async `build_workflow(request)` factory returns a **built**, fresh `Workflow`, +with stable name, executor IDs and graph identity. Build fresh executors, agents, +clients, providers and tools for each request; allocation must not execute tools or +workflow side effects. Hosted workflows require a factory. Local built instances +are single-use and remain application-owned. Factory-owned agents are entered and +closed per turn; applications must explicitly own external credentials/providers +through their resource-managing agent subclasses. Do not configure another +checkpoint store on `WorkflowBuilder` or share mutable graphs between requests. + +Native `AgentExecutor` graphs support `Agent` and constrained bare `RawAgent`. +For `Agent`, hosting applies options through a request-owned agent middleware +policy, disables inner service storage and preserves the developer's defaults. +A bare `RawAgent` must come from a validated fresh factory with a client declaring +its storage capability. A storing raw client requires explicit factory default +`store=False`; requested model overrides must already be materialized and matched +in that fresh agent's defaults. Unsafe/custom agents, mismatched overrides and +unexpected downstream continuation are rejected before successful output or +commit. The host never forwards duplicate `options=` through client kwargs, and +caller JSON cannot assert factory ownership. + +The workflow uses one fixed Invocations lineage inside the trusted user and +platform sandbox scope. Later turns restore the **exact** committed checkpoint, +not a workflow name's unrelated latest checkpoint. Graph changes, stale writers, +forged/cross-scope/replayed decisions and duplicate authority fail before +executor/response-handler dispatch. Reply with the complete pending batch: +partial, unallowlisted, or invalid batches do not claim or consume any usable +authority, so the complete valid batch can be retried. Native factories +and explicit `client_kwargs`/`function_invocation_kwargs` keep their existing +purposes; they do not introduce arbitrary entry-state or per-executor run-option +APIs. Host-controlled identity, storage and execution controls cannot be supplied +as caller generation settings. + +Non-streaming responses are JSON `{"output": [...]}` containing typed `output` and +`request_info` event objects. JSON primitives, dataclasses, Pydantic models and +framework `Message`, `Content` and response values retain their supported encoding. +All nested mappings, including mappings returned by model/dataclass/framework +serializers, require string keys before JSON conversion so key coercion cannot +silently collide. Private provider continuation tokens are omitted; +unsupported/nonfinite/circular values fail explicitly. Streaming emits framed +`output`/`request_info` SSE, +then `done` with the **sandbox** `session_id`. Live output is provisional until +`done`; pending authority and following frames are buffered until the exact cursor +is conditionally committed. Snapshot retention is bounded by both event count and +encoded bytes; exceeding either limit fails the turn without `done`. A conflict, +encoding failure, execution failure or +interruption cannot emit successful completion. Failed claimed turns are blocked +to avoid replaying uncertain effects; start a new sandbox instead of blindly +retrying non-idempotent work. Checkpoints do not provide exactly-once tool effects. +If a connection drops after commit, retrying the same trusted user/sandbox/call ID +replays that response's complete stored snapshot without executing the workflow or +replacing its head. A different call ID starts normal turn validation and cannot +retrieve that snapshot. +Local callers may pass one built workflow for a one-shot non-pausing run. Any +workflow that can pause for external input must use a request-aware factory so a +fresh graph can restore and consume the durable reply on the next request. +After a newer Invocations head is durably committed, the host reclaims completed +ancestor response records and checkpoints oldest-first, retaining only current, +pending, active or blocked authority. Cleanup failures are logged and retried by +a later successful turn without changing the already committed result; default +store expiry remains a backstop. Custom checkpoint providers own their conditional +delete and retention behavior. + +The existing Invocations sandbox-routing and trusted user/call requirements apply, +including the verified `agent_session_id` query alternative when the platform +does not configure `FOUNDRY_AGENT_SESSION_ID`. Body/header IDs cannot establish +sandbox routing or select private checkpoints. The agent path still uses its +separate `invocation_sessions` store, existing hooks/defaults and JSON/`delta` SSE +contract. Agent-only `legacy_wire_format=True`, `prepare_options` and +`agent_session_store_provider` are rejected for native workflows rather than +silently reinterpreted. + +See the [typed Ticket/review example](../../samples/04-hosting/foundry-hosted-agents/invocations/basic/README.md#native-workflow-with-typed-tickets) +for the parser, stable graph, checkpoint type allowlist, pending replies, +JSON/SSE and local/hosted migration guidance. 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 a6f81fee53b..0c038cc2f5b 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -6,17 +6,35 @@ import inspect import json import logging +import math import sys +import uuid import warnings import weakref -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager from contextvars import Token from copy import deepcopy +from dataclasses import asdict, is_dataclass from typing import cast -from agent_framework import AgentSession, ResponseStream, SessionStore, SupportsAgentRun +from agent_framework import ( + AgentResponse, + AgentResponseUpdate, + AgentSession, + ChatResponse, + ChatResponseUpdate, + CheckpointStorage, + ResponseStream, + SessionStore, + SupportsAgentRun, + Workflow, + WorkflowCheckpointException, + WorkflowEvent, + WorkflowExecutor, +) from agent_framework._telemetry import mark_feature_used +from anyio import CancelScope from azure.ai.agentserver.core import ( FoundryAgentRequestContext, get_request_context, @@ -25,6 +43,7 @@ ) from azure.ai.agentserver.core.storage import FoundryStorageConflictError, FoundryStoragePreconditionError from azure.ai.agentserver.invocations import InvocationAgentServerHost +from pydantic import BaseModel from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.types import Receive, Scope, Send @@ -32,13 +51,21 @@ from ._agent_source import AgentSource, is_agent, resolve_agent, validate_agent_source from ._feature_usage import FeatureIndex -from ._request import InvocationRun, UnsupportedOptions, validate_request_options, validate_unsupported_options +from ._request import ( + InvocationRun, + UnsupportedOptions, + WorkflowTurn, + validate_request_options, + validate_unsupported_options, +) from ._scope import FoundryRequestScope -from ._state_store import AgentSessionStoreProvider, StoreProvider +from ._state_store import AgentSessionStoreProvider, CheckpointStoreProvider, ContextScopedStoreProvider, StoreProvider +from ._workflow_source import WorkflowResolver, WorkflowSource, workflow_agents +from ._workflow_state import HostedWorkflowRun, WorkflowBlockedError, WorkflowConflictError logger = logging.getLogger(__name__) -InvocationParser = Callable[[Request], InvocationRun | Awaitable[InvocationRun]] +InvocationParser = Callable[[Request], InvocationRun | WorkflowTurn[Any] | Awaitable[InvocationRun | WorkflowTurn[Any]]] InvocationOptionsHook = Callable[[Request, dict[str, Any]], Mapping[str, Any] | Awaitable[Mapping[str, Any]]] _AGENT_CONTROLLED_FIELDS = frozenset({ @@ -52,12 +79,170 @@ "tokenizer", "tools", }) +_MAX_WORKFLOW_SNAPSHOT_EVENTS = 4096 +_MAX_WORKFLOW_SNAPSHOT_BYTES = 4 * 1024 * 1024 class _UnsupportedAgentOptions(TypeError): """The agent cannot accept the caller's run options under the selected policy.""" +class _WorkflowOutputError(TypeError): + """Workflow application data cannot be represented as JSON.""" + + +class _WorkflowRequestError(ValueError): + """A validated application request cannot be dispatched to the workflow.""" + + +def _validate_workflow_mapping_keys(value: Any, path: str, active: set[int]) -> None: + if isinstance(value, Mapping): + mapping = cast(Mapping[object, Any], value) + identifier = id(mapping) + if identifier in active: + raise _WorkflowOutputError(f"Workflow output at {path} is not JSON-serializable: circular mapping.") + active.add(identifier) + try: + for index, (key, item) in enumerate(mapping.items()): + if not isinstance(key, str): + raise _WorkflowOutputError( + f"Workflow output at {path} is not JSON-serializable: " + f"mapping keys must be strings, found {type(key).__name__}." + ) + _validate_workflow_mapping_keys(item, f"{path}.value[{index}]", active) + finally: + active.remove(identifier) + elif isinstance(value, (list, tuple)): + sequence = cast(Sequence[Any], value) + identifier = id(sequence) + if identifier in active: + raise _WorkflowOutputError(f"Workflow output at {path} is not JSON-serializable: circular sequence.") + active.add(identifier) + try: + for index, item in enumerate(sequence): + _validate_workflow_mapping_keys(item, f"{path}[{index}]", active) + finally: + active.remove(identifier) + + +def _workflow_serialized_value(value: Any, path: str) -> Any: + if isinstance(value, (AgentResponse, AgentResponseUpdate, ChatResponse, ChatResponseUpdate)): + return value.to_dict(exclude={"continuation_token"}) + if isinstance(value, BaseModel): + _validate_workflow_mapping_keys(value.model_dump(mode="python"), f"{path}.model", set()) + return value.model_dump(mode="json") + if is_dataclass(value) and not isinstance(value, type): + return asdict(value) + to_dict = getattr(value, "to_dict", None) + if callable(to_dict): + return to_dict() + raise _WorkflowOutputError( + f"Workflow output at {path} is not JSON-serializable: unsupported type {type(value).__name__}." + ) + + +def _workflow_json_value(value: Any, path: str, active: set[int]) -> Any: + if value is None or isinstance(value, (str, int, bool)): + return value + if isinstance(value, float): + if not math.isfinite(value): + raise _WorkflowOutputError(f"Workflow output at {path} is not JSON-serializable: number must be finite.") + return value + if isinstance(value, Mapping): + mapping = cast(Mapping[object, Any], value) + identifier = id(mapping) + if identifier in active: + raise _WorkflowOutputError(f"Workflow output at {path} is not JSON-serializable: circular mapping.") + active.add(identifier) + try: + result: dict[str, Any] = {} + for index, (key, item) in enumerate(mapping.items()): + if not isinstance(key, str): + raise _WorkflowOutputError( + f"Workflow output at {path} is not JSON-serializable: " + f"mapping keys must be strings, found {type(key).__name__}." + ) + result[key] = _workflow_json_value(item, f"{path}.value[{index}]", active) + return result + finally: + active.remove(identifier) + if isinstance(value, (list, tuple)): + sequence = cast(Sequence[Any], value) + identifier = id(sequence) + if identifier in active: + raise _WorkflowOutputError(f"Workflow output at {path} is not JSON-serializable: circular sequence.") + active.add(identifier) + try: + return [_workflow_json_value(item, f"{path}[{index}]", active) for index, item in enumerate(sequence)] + finally: + active.remove(identifier) + identifier = id(value) + if identifier in active: + raise _WorkflowOutputError(f"Workflow output at {path} is not JSON-serializable: circular serialized value.") + active.add(identifier) + try: + serialized = _workflow_serialized_value(value, path) + return _workflow_json_value(serialized, f"{path}.serialized", active) + finally: + active.remove(identifier) + + +def _workflow_output(value: Any) -> Any: + """Normalize typed application data without text coercion or skipped values.""" + try: + normalized = _workflow_json_value(value, "$", set()) + return json.loads(json.dumps(normalized, allow_nan=False, ensure_ascii=False).encode("utf-8")) + except _WorkflowOutputError: + raise + except (TypeError, ValueError, RecursionError) as exc: + raise _WorkflowOutputError(f"Workflow output type {type(value).__name__} is not JSON-serializable.") from exc + + +def _workflow_frame(event: WorkflowEvent[Any]) -> dict[str, Any] | None: + if event.type not in ("output", "request_info"): + return None + frame: dict[str, Any] = {"type": event.type, "data": event.data} + if event.type == "request_info": + frame["request_id"] = event.request_id + frame["source_executor_id"] = event.source_executor_id + else: + frame["executor_id"] = event.executor_id + return cast(dict[str, Any], _workflow_output(frame)) + + +def _workflow_snapshot_frames(snapshot: Mapping[str, Any] | None) -> list[dict[str, Any]]: + if snapshot is None or set(snapshot) != {"output"} or not isinstance(snapshot["output"], list): + raise RuntimeError("The stored Invocations workflow output snapshot is invalid.") + frames: list[dict[str, Any]] = [] + for item in cast(list[Any], snapshot["output"]): + if not isinstance(item, Mapping): + raise RuntimeError("The stored Invocations workflow output snapshot is invalid.") + frame = cast(dict[str, Any], _workflow_output(item)) + if frame.get("type") == "output": + if set(frame) != {"type", "data", "executor_id"}: + raise RuntimeError("The stored Invocations workflow output snapshot is invalid.") + elif frame.get("type") == "request_info": + if set(frame) != {"type", "data", "request_id", "source_executor_id"}: + raise RuntimeError("The stored Invocations workflow output snapshot is invalid.") + else: + raise RuntimeError("The stored Invocations workflow output snapshot is invalid.") + frames.append(frame) + return frames + + +def _workflow_frame_size(frame: Mapping[str, Any]) -> int: + return len(json.dumps(frame, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8")) + + +def _workflow_can_pause(workflow: Workflow) -> bool: + for executor in workflow.get_executors_list(): + if getattr(executor, "is_request_response_capable", False) is True: + return True + if isinstance(executor, WorkflowExecutor) and _workflow_can_pause(executor.workflow): + return True + return False + + def _sse(event: str, data: Mapping[str, Any]) -> str: return f"event: {event}\ndata: {json.dumps(data)}\n\n" @@ -71,6 +256,29 @@ def _invocation_failure(exc: Exception) -> tuple[str, int, str]: return "Agent invocation failed.", 500, "invocation_failed" +def _workflow_failure(exc: Exception) -> tuple[str, int, str]: + if isinstance(exc, _WorkflowOutputError): + return str(exc), 500, "invalid_workflow_output" + cause: BaseException | None = exc + while cause is not None: + if isinstance(cause, WorkflowBlockedError): + return ( + "This workflow turn was interrupted or failed; start a new Foundry sandbox.", + 409, + "workflow_blocked", + ) + if isinstance(cause, WorkflowConflictError): + return str(cause), 409, "workflow_conflict" + cause = cause.__cause__ + if _invocation_failure(exc)[1] == 409: + return "Another request advanced this workflow; reload before retrying.", 409, "workflow_conflict" + if isinstance(exc, _WorkflowRequestError): + return str(exc), 400, "invalid_workflow_turn" + if isinstance(exc, PermissionError): + return "Workflow state is not available in this trusted scope.", 403, "workflow_scope" + return "Workflow invocation failed.", 500, "workflow_invocation_failed" + + class _InvocationStreamingResponse(StreamingResponse): def __init__(self, content: AsyncGenerator[str]) -> None: super().__init__( @@ -94,15 +302,40 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: await self._content.aclose() +class _WorkflowStreamingResponse(_InvocationStreamingResponse): + def __init__(self, content: AsyncGenerator[str], *, keep_alive: bool) -> None: + super().__init__(content) + self._keep_alive = keep_alive + + async def stream_response(self, send: Send) -> None: + try: + await super().stream_response(send) + finally: + with CancelScope(shield=True): + if self._keep_alive and self.body_iterator is not self._content: + close = getattr(self.body_iterator, "aclose", None) + if close is None: + raise TypeError("The Invocations SDK stream iterator must support aclose.") + await close() + try: + await self._content.aclose() + finally: + # Without the SDK pump, finish its wrapper in the task that created its context tokens. + async for _ in self.body_iterator: + pass + + class InvocationsHostServer(InvocationAgentServerHost): - """Host an agent with durable sessions and application-defined Invocations input.""" + """Host an agent or native workflow with durable, application-defined Invocations turns.""" def __init__( self, - agent: AgentSource, + agent: AgentSource | None = None, *, + workflow: WorkflowSource[Request] | None = None, openapi_spec: dict[str, Any] | None = None, agent_session_store_provider: StoreProvider[SessionStore] | None = None, + checkpoint_store_provider: ContextScopedStoreProvider[CheckpointStorage] | None = None, parse_request: InvocationParser | None = None, prepare_options: InvocationOptionsHook | None = None, unsupported_options: UnsupportedOptions = "warn", @@ -114,30 +347,61 @@ def __init__( Args: 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`. + workflow: A built native workflow or a sync/async request-aware factory returning a freshly built + workflow with stable graph/executor IDs. Hosted workflows require a factory. Supply exactly + one of agent or workflow; local workflow instances are single-use and remain application-owned. openapi_spec: The OpenAPI specification for the server. agent_session_store_provider: Provider for conversation session storage. Defaults to Foundry storage when hosted and the SDK's file-backed storage locally. New default stores expire sessions 30 days after their last write. Custom providers control their own retention. - parse_request: Optional sync or async parser returning an `InvocationRun` from application JSON. - Without one, accepts a JSON object with `message`, optional `options`, and optional `stream`. + checkpoint_store_provider: Native workflow checkpoint storage. The default accepts the host's + verified sandbox/user scope. Custom providers must preserve scoped conditional-write semantics. + parse_request: Sync or async parser returning `InvocationRun` for agents or `WorkflowTurn` for + workflows. Native workflows require an explicit parser for typed input or complete pending + reply batches. Without an agent parser, accepts `message`, optional `options` and `stream`. prepare_options: Optional sync or async hook to filter or replace a copy of caller generation options. Tool context and agent execution controls must remain in developer-owned agent configuration. unsupported_options: `"warn"` (default), `"ignore"`, or `"error"` for agents without runtime options. legacy_wire_format: Opt into the deprecated plain-text response and raw streaming chunks instead of - the default JSON response and framed `delta`/`done`/`error` server-sent events. + the default agent JSON response and framed `delta`/`done`/`error` server-sent events. + This agent-only compatibility mode is incompatible with native workflows. **kwargs: Additional keyword arguments. """ - validate_agent_source(agent) + if agent is not None and workflow is not None: + raise TypeError("Provide exactly one of agent or workflow.") + if agent is None and workflow is None: + raise TypeError( + "Provide exactly one of agent or workflow; " + "agent must be an agent instance or a zero-argument callable that creates one." + ) + if workflow is None: + validate_agent_source(agent) + elif parse_request is None: + raise TypeError("parse_request is required when hosting a native Workflow.") if parse_request is not None and not callable(parse_request): raise TypeError("parse_request must be a callable.") if prepare_options is not None and not callable(prepare_options): raise TypeError("prepare_options must be a callable.") if not isinstance(legacy_wire_format, bool): raise TypeError("legacy_wire_format must be a boolean.") + if workflow is not None: + if legacy_wire_format: + raise ValueError("legacy_wire_format=True is incompatible with native workflows.") + if prepare_options is not None: + raise ValueError("prepare_options is agent-only; use the workflow parser's explicit client_kwargs.") + if agent_session_store_provider is not None: + raise ValueError( + "agent_session_store_provider is agent-only; use checkpoint_store_provider for workflows." + ) + self._workflow_resolver: WorkflowResolver[Request] | None = ( + WorkflowResolver(workflow) if workflow is not None else None + ) super().__init__(openapi_spec=openapi_spec, **kwargs) + if self._workflow_resolver is not None and self.config.is_hosted and not self._workflow_resolver.is_factory: + raise ValueError("Hosted native workflows require a request-aware factory that creates fresh executors.") self._agent = agent - self._owns_request_agent = not is_agent(agent) + self._owns_request_agent = agent is not None and not is_agent(agent) self._parse_request = parse_request self._prepare_options = prepare_options self._unsupported_options = validate_unsupported_options(unsupported_options) @@ -147,6 +411,9 @@ def __init__( if agent_session_store_provider is None else agent_session_store_provider ) + self._checkpoint_storage_provider = ( + CheckpointStoreProvider() if checkpoint_store_provider is None else checkpoint_store_provider + ) self._session_locks: weakref.WeakValueDictionary[str | tuple[str, str], asyncio.Lock] = ( weakref.WeakValueDictionary() ) @@ -228,6 +495,8 @@ def _hosted_scope(self, request: Request, context: FoundryAgentRequestContext) - @asynccontextmanager async def _request_agent(self) -> AsyncGenerator[SupportsAgentRun]: + if self._agent is None: + raise RuntimeError("No agent is configured for Invocations.") agent = await resolve_agent(self._agent) async with AsyncExitStack() as resources: if self._owns_request_agent and isinstance(agent, AbstractAsyncContextManager): @@ -273,11 +542,14 @@ async def _request_session( f"session persistence also failed: {str(exc) or type(exc).__name__}" ) from exc - async def _parse(self, request: Request) -> InvocationRun: + async def _parse(self, request: Request) -> InvocationRun | WorkflowTurn[Any]: if self._parse_request is not None: result = self._parse_request(request) parsed = await result if inspect.isawaitable(result) else result - if not isinstance(parsed, InvocationRun): + if self._workflow_resolver is not None: + if not isinstance(parsed, WorkflowTurn): + raise TypeError("parse_request must return WorkflowTurn.") + elif not isinstance(parsed, InvocationRun): raise TypeError("parse_request must return InvocationRun.") return parsed @@ -339,7 +611,7 @@ async def _close_interrupted_stream(stream: object) -> None: await close() async def _handle_invoke(self, request: Request) -> Response: - """Invoke the agent with the given request.""" + """Invoke the configured agent or native workflow with a trusted request.""" context = get_request_context() try: hosted_scope = self._hosted_scope(request, context) if self.config.is_hosted else None @@ -350,13 +622,17 @@ async def _handle_invoke(self, request: Request) -> Response: try: parsed = await self._parse(request) - options = await self._options(request, parsed) + options = await self._options(request, parsed) if isinstance(parsed, InvocationRun) else {} except (TypeError, ValueError) as exc: return JSONResponse({"error": str(exc)}, status_code=400) except Exception: logger.exception("Failed to prepare Invocations request") return JSONResponse({"error": "Failed to prepare invocation request."}, status_code=500) + if isinstance(parsed, WorkflowTurn): + scope = hosted_scope or FoundryRequestScope.from_context(self.config, context) + return await self._handle_workflow(request, parsed, scope, context, partition_key) + if parsed.stream: if options and self._unsupported_options == "error" and is_agent(self._agent): try: @@ -440,3 +716,178 @@ async def stream_response() -> AsyncGenerator[str]: logger.exception("Invocations agent request failed") message, status, _ = _invocation_failure(exc) return JSONResponse({"error": message}, status_code=status) + + async def _handle_workflow( + self, + request: Request, + turn: WorkflowTurn[Any], + scope: FoundryRequestScope, + context: FoundryAgentRequestContext, + partition_key: str | tuple[str, str], + ) -> Response: + if turn.stream: + + async def stream_response() -> AsyncGenerator[str]: + token: Token[FoundryAgentRequestContext] | None = set_request_context(context) + events = self._workflow_events(request, turn, scope, context, partition_key) + try: + async for frame in events: + wire_frame = _sse(frame["type"], frame) + reset_request_context(token) + token = None + yield wire_frame + token = set_request_context(context) + reset_request_context(token) + token = None + yield _sse("done", {"session_id": scope.session_id}) + except Exception as exc: + logger.exception("Invocations workflow stream failed") + message, status, code = _workflow_failure(exc) + if token is not None: + reset_request_context(token) + token = None + yield _sse("error", {"message": message, "code": code, "status": status}) + finally: + if token is None: + token = set_request_context(context) + try: + with CancelScope(shield=True): + await events.aclose() + finally: + reset_request_context(token) + + return _WorkflowStreamingResponse(stream_response(), keep_alive=self.config.sse_keepalive_interval > 0) + + try: + frames = [frame async for frame in self._workflow_events(request, turn, scope, context, partition_key)] + return JSONResponse({"output": frames}) + except Exception as exc: + logger.exception("Invocations workflow request failed") + message, status, _ = _workflow_failure(exc) + return JSONResponse({"error": message}, status_code=status) + + async def _workflow_events( + self, + request: Request, + turn: WorkflowTurn[Any], + scope: FoundryRequestScope, + context: FoundryAgentRequestContext, + partition_key: str | tuple[str, str], + ) -> AsyncGenerator[dict[str, Any]]: + resolver = self._workflow_resolver + if resolver is None: + raise RuntimeError("No native workflow is configured for Invocations.") + frames: list[dict[str, Any]] = [] + snapshot_bytes = 0 + deferred: list[dict[str, Any]] = [] + defer = False + lock = self._session_locks.setdefault(partition_key, asyncio.Lock()) + async with lock: + workflow = await resolver.resolve(request) + if not resolver.is_factory and _workflow_can_pause(workflow): + raise _WorkflowRequestError("A pausing native workflow requires a fresh request-aware factory.") + run = await HostedWorkflowRun.prepare( + workflow, + scope=scope, + response_id=context.call_id or uuid.uuid4().hex, + config=self.config, + platform_context=context, + checkpoint_store_provider=self._checkpoint_storage_provider, + lineage_id="invocations", + stored=resolver.is_factory, + replay_completed=True, + fresh_factory=resolver.is_factory, + reclaim_superseded=True, + ) + if run.has_completed_output: + for frame in _workflow_snapshot_frames(run.snapshot): + yield frame + return + try: + turn = run.validate_turn(turn) + except (TypeError, ValueError) as exc: + raise _WorkflowRequestError(str(exc)) from exc + if ( + run.stored + and turn.input is not None + and isinstance(self._checkpoint_storage_provider, CheckpointStoreProvider) + ): + try: + self._checkpoint_storage_provider.validate_checkpoint_value(turn.input) + except WorkflowCheckpointException as exc: + raise _WorkflowRequestError( + "Stored native workflow input uses an application type that is not checkpoint-allowlisted. " + 'Configure CheckpointStoreProvider(allowed_checkpoint_types=["module:qualname"]).' + ) from exc + if ( + run.stored + and turn.responses is not None + and isinstance(self._checkpoint_storage_provider, CheckpointStoreProvider) + ): + try: + for reply in turn.responses.values(): + self._checkpoint_storage_provider.validate_checkpoint_value(reply) + except WorkflowCheckpointException as exc: + raise _WorkflowRequestError( + "Stored native workflow reply uses an application type that is not checkpoint-allowlisted. " + 'Configure CheckpointStoreProvider(allowed_checkpoint_types=["module:qualname"]).' + ) from exc + committed = False + resources = AsyncExitStack() + try: + await run.claim() + try: + if resolver.is_factory: + entered: set[int] = set() + for agent in workflow_agents(workflow): + if id(agent) not in entered and isinstance(agent, AbstractAsyncContextManager): + entered.add(id(agent)) + await resources.enter_async_context(agent) + events = run.events(turn) + resources.push_async_callback(events.aclose) + async for event in events: + if event.type == "request_info" and not resolver.is_factory: + raise _WorkflowRequestError( + "Native workflow pauses require a request-aware factory that rebuilds the workflow." + ) + frame = _workflow_frame(event) + if frame is None: + continue + frame_size = _workflow_frame_size(frame) + if ( + len(frames) >= _MAX_WORKFLOW_SNAPSHOT_EVENTS + or snapshot_bytes + frame_size > _MAX_WORKFLOW_SNAPSHOT_BYTES + ): + raise _WorkflowOutputError( + "Workflow output exceeds the bounded Invocations snapshot limit." + ) + snapshot_bytes += frame_size + frames.append(frame) + defer = defer or event.type == "request_info" + if defer: + deferred.append(frame) + else: + yield frame + finally: + # Starlette cancellation must not interrupt core/resource cleanup. + with CancelScope(shield=True): + await resources.__aexit__(*sys.exc_info()) + await run.commit({"output": frames}) + committed = True + finally: + if not committed: + failure = sys.exc_info()[1] + try: + with CancelScope(shield=True): + await run.abort() + except Exception as exc: + logger.exception("Failed to persist interrupted Invocations workflow state") + if failure is None: + raise + if isinstance(failure, Exception): + raise RuntimeError( + "Workflow invocation failed and its interrupted state could not be persisted." + ) from exc + # Pending reply authority is durable, and a reply need not wait on this stream's lock. + for frame in deferred: + yield frame 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 3ffd6f89f5e..487d6882cd0 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 @@ -2,9 +2,10 @@ import hashlib +import json from abc import ABC, abstractmethod from datetime import datetime -from typing import Generic, Protocol, TypeVar +from typing import Any, Generic, Protocol, TypeVar from agent_framework import ( AgentSession, @@ -27,6 +28,12 @@ StoreT = TypeVar("StoreT") +def _encoded_checkpoint_hash(value: Any) -> str: + """Hash the exact JSON-compatible representation acknowledged by storage.""" + payload = json.dumps(value, sort_keys=True, allow_nan=False, separators=(",", ":")).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + def _store_scope(config: AgentConfig, platform_context: FoundryAgentRequestContext) -> FoundryRequestScope | None: if not config.is_hosted: return None @@ -148,6 +155,14 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: await store.set_item(checkpoint.checkpoint_id, encoded_checkpoint, call_id=self.platform_context.call_id) return checkpoint.checkpoint_id + async def _load_encoded(self, checkpoint_id: CheckpointID) -> Any: + store = await self._get_store() + async with store: + item = await store.get_item(checkpoint_id, call_id=self.platform_context.call_id) + if item is None: + raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") + return item.value + async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: """Load a workflow checkpoint from the store. @@ -162,12 +177,21 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint: """ from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value - store = await self._get_store() - async with store: - item = await store.get_item(checkpoint_id, call_id=self.platform_context.call_id) - if item is None: - raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}") - return WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value, allowed_types=self._allowed_types)) + encoded = await self._load_encoded(checkpoint_id) + return WorkflowCheckpoint.from_dict(decode_checkpoint_value(encoded, allowed_types=self._allowed_types)) + + async def load_with_hash(self, checkpoint_id: CheckpointID) -> tuple[WorkflowCheckpoint, str]: + """Load a checkpoint and hash its exact persisted encoding. + + Pickle is not a canonical serialization across Python versions or repeated + encodes. Hashing the acknowledged storage value preserves exact tamper + detection without rejecting an unchanged checkpoint after deserialization. + """ + from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value + + encoded = await self._load_encoded(checkpoint_id) + checkpoint = WorkflowCheckpoint.from_dict(decode_checkpoint_value(encoded, allowed_types=self._allowed_types)) + return checkpoint, _encoded_checkpoint_hash(encoded) async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: """List all workflow checkpoints for a given workflow name.""" diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_state.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_state.py index 529e760805e..442eadb64e0 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_state.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_state.py @@ -6,8 +6,9 @@ import hashlib import json +import logging import secrets -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping from dataclasses import asdict, dataclass, replace from typing import Any, Literal, cast @@ -51,6 +52,7 @@ _CONFLICT = "Another request advanced this workflow. Reload its current response or start a fresh lineage." _BLOCKED = "This workflow turn was interrupted or failed; start a fresh workflow lineage to avoid replaying effects." +logger = logging.getLogger(__name__) class WorkflowConflictError(RuntimeError): @@ -254,6 +256,20 @@ async def save_head(self, head: WorkflowHead, *, expected_etag: str | None) -> s ) return await self._write(key, asdict(head), expected_etag) + async def delete_response(self, response_id: str, *, expected_etag: str) -> bool: + """Conditionally delete a superseded response record.""" + store = await self._get_store() + async with store: + try: + deleted = await store.delete_item( + _key("response", response_id), + if_match=expected_etag, + call_id=self.scope.call_id, + ) + except (FoundryStorageConflictError, FoundryStoragePreconditionError) as exc: + raise WorkflowConflictError(_CONFLICT) from exc + return deleted.id is not None + async def _write(self, key: str, value: dict[str, Any], expected_etag: str | None) -> str: store = await self._get_store() async with store: @@ -298,6 +314,21 @@ async def load(self, checkpoint_id: str) -> WorkflowCheckpoint: state[key] = fresh.get(key, {}) return replace(checkpoint, state=state) + async def load_persisted(self, checkpoint_id: str) -> tuple[WorkflowCheckpoint, str]: + """Load a checkpoint with a process-stable integrity hash when supported.""" + load_with_hash = getattr(cast(Any, self.storage), "load_with_hash", None) + if callable(load_with_hash): + loader = cast( + Callable[[str], Awaitable[tuple[WorkflowCheckpoint, str]]], + load_with_hash, + ) + checkpoint, checkpoint_hash = await loader(checkpoint_id) + if not isinstance(checkpoint, WorkflowCheckpoint) or not isinstance(checkpoint_hash, str): + raise TypeError("Checkpoint storage returned an invalid checkpoint integrity pair.") + return checkpoint, checkpoint_hash + checkpoint = await self.storage.load(checkpoint_id) + return checkpoint, _checkpoint_hash(checkpoint) + async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: return await self.storage.list_checkpoints(workflow_name=workflow_name) @@ -327,12 +358,14 @@ def __init__( stored: bool, recovery: bool, fresh_factory: bool, + reclaim_superseded: bool, ) -> None: self.workflow = workflow self.scope = scope self.stored = stored self.recovery = recovery self.fresh_factory = fresh_factory + self.reclaim_superseded = reclaim_superseded self.store = FoundryWorkflowBindingStore(scope) self.binding: WorkflowBinding self._head: WorkflowHead | None = None @@ -364,7 +397,9 @@ async def prepare( lineage_id: str | None = None, stored: bool = True, recovery: bool = False, + replay_completed: bool = False, fresh_factory: bool = False, + reclaim_superseded: bool = False, ) -> HostedWorkflowRun: """Read exact continuation state without claiming, consuming replies, or executing.""" if not workflow.name or not workflow.graph_signature_hash: @@ -373,7 +408,14 @@ async def prepare( raise ValueError("A workflow turn cannot combine previous_response_id and conversation.") if not stored and (recovery or previous_response_id is not None or conversation_id is not None): raise ValueError("store=false workflow requests cannot continue stored state; start a one-shot turn.") - run = cls(workflow, scope, stored=stored, recovery=recovery, fresh_factory=fresh_factory) + run = cls( + workflow, + scope, + stored=stored, + recovery=recovery, + fresh_factory=fresh_factory, + reclaim_superseded=reclaim_superseded, + ) run.binding = WorkflowBinding( response_id=response_id, checkpoint_id=None, @@ -388,6 +430,9 @@ async def prepare( return run own, own_etag = await run.store.get_response(response_id) + if not recovery and replay_completed and own is not None and own.status == "completed": + recovery = True + run.recovery = True if recovery and own is not None: run._record, run._record_etag = own, own_etag run.binding = own.binding @@ -511,12 +556,12 @@ async def prepare( ) if previous.checkpoint_id is None: raise ValueError("The workflow response has no acknowledged checkpoint.") - checkpoint = await storage.load(previous.checkpoint_id) + checkpoint, checkpoint_hash = await run._storage.load_persisted(previous.checkpoint_id) if ( checkpoint.checkpoint_id != previous.checkpoint_id or checkpoint.workflow_name != workflow.name or checkpoint.graph_signature_hash != workflow.graph_signature_hash - or _checkpoint_hash(checkpoint) != previous.checkpoint_hash + or checkpoint_hash != previous.checkpoint_hash ): raise ValueError("The exact workflow checkpoint does not match its response binding.") run._checkpoint = checkpoint @@ -686,6 +731,7 @@ async def claim(self) -> None: ) self._record_etag = await self.store.save_response(record, expected_etag=self._record_etag) self._record = record + await self._reclaim_completed_ancestors() async def assert_claim(self) -> None: if not self._claimed: @@ -755,7 +801,7 @@ async def stage( if not self.stored: return snapshot_value = _json_object(snapshot) - checkpoint = await self.get_checkpoint(checkpoint_id or self.checkpoint_id) + checkpoint, checkpoint_hash = await self._get_checkpoint_with_hash(checkpoint_id or self.checkpoint_id) if checkpoint is None: raise RuntimeError("Output cannot be paired before a core checkpoint is acknowledged.") if approvals is not None and any( @@ -765,9 +811,7 @@ async def stage( await self.assert_claim() if self._record is None: raise RuntimeError("The workflow response pair has not been claimed.") - binding = replace( - self.binding, checkpoint_id=checkpoint.checkpoint_id, checkpoint_hash=_checkpoint_hash(checkpoint) - ) + binding = replace(self.binding, checkpoint_id=checkpoint.checkpoint_id, checkpoint_hash=checkpoint_hash) pending_approvals = { wire_id: request_id for wire_id, request_id in (approvals if approvals is not None else self._record.approvals or {}).items() @@ -785,14 +829,20 @@ async def stage( async def get_checkpoint(self, checkpoint_id: str | None) -> WorkflowCheckpoint | None: """Read a checkpoint acknowledged in this run, including a producer-stamped stream boundary.""" + checkpoint, _ = await self._get_checkpoint_with_hash(checkpoint_id) + return checkpoint + + async def _get_checkpoint_with_hash( + self, checkpoint_id: str | None + ) -> tuple[WorkflowCheckpoint | None, str | None]: if checkpoint_id is None: - return None + return None, None if self._storage is None or ( checkpoint_id not in self._storage.acknowledged_ids and (self._checkpoint is None or checkpoint_id != self._checkpoint.checkpoint_id) ): raise ValueError("Output pairing requires a checkpoint acknowledged by this exact workflow turn.") - checkpoint = await self._storage.storage.load(checkpoint_id) + checkpoint, checkpoint_hash = await self._storage.load_persisted(checkpoint_id) if ( checkpoint.checkpoint_id != checkpoint_id or checkpoint.workflow_name != self.workflow.name @@ -800,7 +850,7 @@ async def get_checkpoint(self, checkpoint_id: str | None) -> WorkflowCheckpoint ): raise ValueError("The acknowledged workflow checkpoint identity changed.") validate_workflow_provider_state(self.workflow, checkpoint) - return checkpoint + return checkpoint, checkpoint_hash async def commit(self, snapshot: Mapping[str, Any] | None = None) -> None: """Commit only fully finalized, encoded output, before protocol success or done.""" @@ -823,6 +873,44 @@ async def commit(self, snapshot: Mapping[str, Any] | None = None) -> None: self._head = head self._claimed = False + async def _reclaim_completed_ancestors(self) -> None: + """Best-effort reclaim after the new head is durably committed. + + Cleanup is oldest-first. If a deletion fails, newer records retain the + chain needed for a later successful commit to retry cleanup. + """ + if not self.reclaim_superseded or self._storage is None: + return + response_id = self.binding.previous_response_id + records: list[tuple[WorkflowRecord, str]] = [] + visited: set[str] = set() + try: + current_checkpoint_id = self.binding.checkpoint_id + checkpoint_ids = await self._storage.storage.list_checkpoint_ids(workflow_name=self.binding.workflow_name) + for checkpoint_id in checkpoint_ids: + if current_checkpoint_id is None or checkpoint_id != current_checkpoint_id: + await self._storage.storage.delete(checkpoint_id) + while response_id is not None: + if response_id in visited: + raise ValueError("Invalid cycle in superseded workflow response lineage.") + visited.add(response_id) + record, etag = await self.store.get_response(response_id) + if record is None or etag is None: + break + if ( + record.status != "completed" + or record.binding.scope_key != self.scope.storage_key + or record.binding.lineage_id != self.binding.lineage_id + or record.binding.conversation_id != self.binding.conversation_id + ): + break + records.append((record, etag)) + response_id = record.binding.previous_response_id + for record, etag in reversed(records): + await self.store.delete_response(record.binding.response_id, expected_etag=etag) + except Exception: + logger.exception("Failed to reclaim superseded native workflow state; a later commit can retry cleanup") + async def abort(self) -> None: """Block interrupted/failed authority instead of permitting unsafe reply or effect replay.""" if not self.stored or not self._claimed or self.already_completed: diff --git a/python/packages/foundry_hosting/tests/foundry_hosting/test_invocations_workflow.py b/python/packages/foundry_hosting/tests/foundry_hosting/test_invocations_workflow.py new file mode 100644 index 00000000000..2ccd0d59ef4 --- /dev/null +++ b/python/packages/foundry_hosting/tests/foundry_hosting/test_invocations_workflow.py @@ -0,0 +1,1547 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Protocol-specific regressions for native Invocations workflows.""" + +from __future__ import annotations + +import asyncio +import gc +import json +import uuid +from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import date +from typing import Any, cast +from unittest.mock import patch + +import httpx +import pytest +from agent_framework import ( + Agent, + AgentExecutor, + AgentResponse, + AgentResponseUpdate, + BaseChatClient, + ChatMiddlewareLayer, + ChatOptions, + ChatResponse, + ChatResponseUpdate, + Content, + Executor, + FunctionInvocationContext, + FunctionInvocationLayer, + InMemoryHistoryProvider, + Message, + RawAgent, + ResponseStream, + Workflow, + WorkflowBuilder, + WorkflowContext, + handler, + response_handler, + tool, +) +from azure.ai.agentserver.core import ( + FoundryAgentRequestContext, + get_request_context, + reset_request_context, + set_request_context, +) +from pydantic import BaseModel +from starlette.requests import ClientDisconnect, Request +from starlette.responses import Response, StreamingResponse + +from agent_framework_foundry_hosting import ( + CheckpointStoreProvider, + FoundryRequestScope, + InvocationsHostServer, + WorkflowSource, + WorkflowTurn, +) +from agent_framework_foundry_hosting._invocations import _workflow_output, _WorkflowOutputError +from agent_framework_foundry_hosting._state_store import FoundryCheckpointStore +from agent_framework_foundry_hosting._workflow_state import FoundryWorkflowBindingStore, WorkflowHead + + +@dataclass(frozen=True) +class Ticket: + ticket_id: str + question: str + + +class TicketModel(BaseModel): + ticket_id: str + due: date + + +@dataclass(frozen=True) +class MappingEnvelope: + payload: Mapping[Any, str] + + +class MappingModel(BaseModel): + payload: dict[Any, str] + + +class MappingSerializer: + def to_dict(self) -> dict[Any, str]: + return {1: "private-token", "1": "safe"} + + +@dataclass(frozen=True) +class TicketState: + turns: int = 0 + last_ticket: str | None = None + + +@dataclass(frozen=True) +class TicketReview: + ticket_id: str + reviewer: str + + +@dataclass(frozen=True) +class TicketDecision: + approved: bool + + +@dataclass(frozen=True) +class UnallowlistedReply: + value: str + + +class _TicketExecutor(Executor): + def __init__(self, calls: list[tuple[Ticket, str | None]], *, id: str = "ticket-start") -> None: + super().__init__(id=id) + self.calls = calls + + @handler + async def handle(self, ticket: Ticket, ctx: WorkflowContext[None, dict[str, Any]]) -> None: + assert isinstance(ticket, Ticket) + self.calls.append((ticket, get_request_context().call_id)) + state = ctx.get_state("tickets", TicketState()) + assert isinstance(state, TicketState) + ctx.set_state("tickets", TicketState(state.turns + 1, ticket.ticket_id)) + await ctx.yield_output({"ticket": ticket, "turn": state.turns + 1, "previous_ticket": state.last_ticket}) + + +class _ReviewExecutor(Executor): + def __init__(self, replies: list[tuple[str, bool]], starts: list[str] | None = None) -> None: + super().__init__(id="ticket-review") + self.replies = replies + self.starts = starts + + @handler + async def handle(self, ticket: Ticket, ctx: WorkflowContext[None, dict[str, Any]]) -> None: + if self.starts is not None: + self.starts.append(ticket.ticket_id) + for reviewer in ("support", "security"): + await ctx.request_info(TicketReview(ticket.ticket_id, reviewer), response_type=TicketDecision) + + @response_handler + async def review( + self, request: TicketReview, response: TicketDecision, ctx: WorkflowContext[None, dict[str, Any]] + ) -> None: + self.replies.append((request.reviewer, response.approved)) + await ctx.yield_output({ + "ticket_id": request.ticket_id, + "reviewer": request.reviewer, + "approved": response.approved, + }) + + +class _AnyReplyExecutor(Executor): + def __init__(self, replies: list[object]) -> None: + super().__init__(id="any-reply") + self.replies = replies + + @handler + async def handle(self, message: str, ctx: WorkflowContext[None, object]) -> None: + await ctx.request_info(f"{message}:first", response_type=object, request_id="first") + await ctx.request_info(f"{message}:second", response_type=object, request_id="second") + + @response_handler + async def decide(self, request: str, response: object, ctx: WorkflowContext[None, object]) -> None: + self.replies.append(response) + await ctx.yield_output({"request": request, "accepted": True}) + + +class _ApprovalExecutor(Executor): + def __init__(self, decisions: list[bool]) -> None: + super().__init__(id="approval") + self.decisions = decisions + + @handler + async def handle(self, ticket: Ticket, ctx: WorkflowContext[None, dict[str, Any]]) -> None: + for reviewer in ("support", "security"): + function = Content.from_function_call( + call_id=uuid.uuid4().hex, + name="approve_ticket", + arguments={"ticket_id": ticket.ticket_id, "reviewer": reviewer}, + ) + await ctx.request_info( + Content.from_function_approval_request(uuid.uuid4().hex, function), + response_type=Content, + ) + + @response_handler + async def decide(self, request: Content, response: Content, ctx: WorkflowContext[None, dict[str, Any]]) -> None: + assert isinstance(response.approved, bool) + self.decisions.append(response.approved) + await ctx.yield_output({"approved": response.approved}) + + +class _RegistryExecutor(Executor): + def __init__(self, agent: Agent[Any]) -> None: + super().__init__("registry") + self._agents = {"registered": agent} + + @handler + async def handle(self, ticket: Ticket, ctx: WorkflowContext[None, dict[str, Any]]) -> None: + await ctx.yield_output({"ticket_id": ticket.ticket_id}) + + +class _ContextClient( + FunctionInvocationLayer[ChatOptions], ChatMiddlewareLayer[ChatOptions], BaseChatClient[ChatOptions] +): + STORES_BY_DEFAULT = True + + def __init__( + self, + calls: list[dict[str, Any]], + resources: list[tuple[str, str | None]], + ) -> None: + super().__init__(middleware=[]) + self.calls = calls + self.resources = resources + self.round = 0 + + async def __aenter__(self) -> _ContextClient: + self.resources.append(("enter", get_request_context().call_id)) + return self + + async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: + self.resources.append(("exit", get_request_context().call_id)) + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.round += 1 + self.calls.append({ + "options": dict(options), + "kwargs": dict(kwargs), + "call_id": get_request_context().call_id, + }) + contents = ( + [Content.from_function_call(uuid.uuid4().hex, "inspect_runtime", arguments={})] + if self.round == 1 + else [Content.from_text("checked")] + ) + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=contents) + + async def complete() -> ChatResponse: + return ChatResponse(messages=[Message("assistant", contents=contents)]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else complete() + + +def _request(payload: Any, *, session_id: str = "session") -> Request: + body = json.dumps(payload).encode() + + async def receive() -> Any: + return {"type": "http.request", "body": body} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/invocations", + "query_string": f"agent_session_id={session_id}".encode(), + "headers": [(b"content-type", b"application/json")], + }, + receive, + ) + + +@contextmanager +def _context(*, session_id: str = "session", call_id: str | None = None, user_id: str | None = None) -> Iterator[None]: + token = set_request_context(FoundryAgentRequestContext(session_id=session_id, call_id=call_id, user_id=user_id)) + try: + yield + finally: + reset_request_context(token) + + +async def _parse_ticket(request: Request) -> WorkflowTurn[Ticket]: + payload = await request.json() + if "responses" in payload: + return WorkflowTurn( + responses={ + request_id: TicketDecision(approved=value["approved"]) + for request_id, value in payload["responses"].items() + }, + stream=payload.get("stream", False), + ) + return WorkflowTurn( + input=Ticket(payload["ticket_id"], payload.get("question", "Help")), + stream=payload.get("stream", False), + ) + + +def _host( + factory: WorkflowSource[Request], + *, + parser: Any = _parse_ticket, + checkpoints: CheckpointStoreProvider | None = None, +) -> InvocationsHostServer: + return InvocationsHostServer( + workflow=factory, + parse_request=parser, + checkpoint_store_provider=checkpoints + or CheckpointStoreProvider( + allowed_checkpoint_types=[ + f"{__name__}:{value.__qualname__}" + for value in (Ticket, TicketModel, TicketState, TicketReview, TicketDecision) + ] + ), + configure_observability=None, + ) + + +def _sse_events(body: str) -> list[dict[str, Any]]: + frames: list[dict[str, Any]] = [] + for frame in body.split("\n\n"): + if frame.startswith("event: "): + event, data = frame.split("\n", 1) + frames.append({"event": event.removeprefix("event: "), **json.loads(data.removeprefix("data: "))}) + return frames + + +async def _events(response: Response) -> list[dict[str, Any]]: + if isinstance(response, StreamingResponse): + body = "".join([ + chunk if isinstance(chunk, str) else bytes(chunk).decode() async for chunk in response.body_iterator + ]) + return _sse_events(body) + payload = json.loads(bytes(response.body)) + if response.status_code != 200: + return [{"event": "error", **payload, "status": response.status_code}] + assert response.media_type == "application/json" + return [{"event": item["type"], **item} for item in payload["output"]] + + +async def _invoke( + server: InvocationsHostServer, + payload: Any, + *, + session_id: str = "session", + call_id: str | None = None, + user_id: str | None = None, +) -> tuple[Response, list[dict[str, Any]]]: + with _context(session_id=session_id, call_id=call_id, user_id=user_id): + response = await server._handle_invoke(_request(payload, session_id=session_id)) + # Deliberately collect after the original request context has been removed. + return response, await _events(response) + + +class TestWorkflowOutput: + @pytest.mark.parametrize( + ("value", "expected"), + [ + (None, None), + (True, True), + (3, 3), + (0.25, 0.25), + ("ready", "ready"), + (("one", 2), ["one", 2]), + (Ticket("T-1", "Help"), {"ticket_id": "T-1", "question": "Help"}), + ( + TicketModel(ticket_id="T-1", due=date(2026, 10, 1)), + {"ticket_id": "T-1", "due": "2026-10-01"}, + ), + ( + {"tickets": [Ticket("T-1", "Help"), Ticket("T-2", "Thanks")]}, + {"tickets": [{"ticket_id": "T-1", "question": "Help"}, {"ticket_id": "T-2", "question": "Thanks"}]}, + ), + ], + ) + def test_encodes_application_values(self, value: Any, expected: Any) -> None: + assert _workflow_output(value) == expected + + @pytest.mark.parametrize( + "value", + [ + Content.from_text("typed content"), + Message("assistant", ["typed message"]), + AgentResponse(messages=[Message("assistant", ["typed response"])]), + ], + ) + def test_preserves_framework_output_instead_of_flattening_to_text(self, value: Any) -> None: + assert _workflow_output(value) == value.to_dict() + assert isinstance(_workflow_output(value), dict) + + @pytest.mark.parametrize( + "value", [object(), {"nested": object()}, [object()], float("nan"), float("inf"), Ticket, "\ud800"] + ) + def test_rejects_unsupported_values_without_stringifying_or_dropping_them(self, value: Any) -> None: + with pytest.raises(_WorkflowOutputError, match="not JSON-serializable"): + _workflow_output(value) + + def test_rejects_circular_output(self) -> None: + circular: list[Any] = [] + circular.append(circular) + with pytest.raises(_WorkflowOutputError, match="not JSON-serializable"): + _workflow_output(circular) + + def test_rejects_invalid_nested_framework_data(self) -> None: + content = Content.from_text("text", additional_properties={"invalid": object()}) + with pytest.raises(_WorkflowOutputError, match="not JSON-serializable"): + _workflow_output(content) + + @pytest.mark.parametrize( + "value", + [ + {"nested": {1: "private-token", "1": "safe"}}, + MappingEnvelope({1: "private-token", "1": "safe"}), + MappingModel(payload={1: "private-token", "1": "safe"}), + Content.from_text( + "typed content", + additional_properties=cast(Any, {1: "private-token", "1": "safe"}), + ), + MappingSerializer(), + ], + ) + def test_rejects_non_string_mapping_keys_before_json_coercion_without_leaking_values(self, value: Any) -> None: + with pytest.raises(_WorkflowOutputError, match=r"at \$.*mapping keys must be strings") as raised: + _workflow_output(value) + assert "private-token" not in str(raised.value) + assert "'1'" not in str(raised.value) + + @pytest.mark.parametrize("response_type", [AgentResponse, AgentResponseUpdate, ChatResponse, ChatResponseUpdate]) + def test_does_not_expose_private_provider_continuation(self, response_type: Any) -> None: + response = response_type(continuation_token={"private": "provider-token"}) + assert "continuation_token" not in _workflow_output(response) + assert response.continuation_token == {"private": "provider-token"} + + +class TestNativeWorkflowConfiguration: + def test_requires_one_source_and_an_explicit_parser(self) -> None: + def create(_request: Request) -> Workflow: + return WorkflowBuilder(name="tickets", start_executor=_TicketExecutor([])).build() + + invalid_agent: Any = object() + with pytest.raises(TypeError, match="exactly one"): + InvocationsHostServer(agent=invalid_agent, workflow=create, parse_request=_parse_ticket) + with pytest.raises(TypeError, match="parse_request is required"): + InvocationsHostServer(workflow=create) + + def test_workflow_rejects_agent_legacy_wire_format(self) -> None: + with pytest.raises(ValueError, match="legacy_wire_format.*native"): + InvocationsHostServer( + workflow=lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor([])).build(), + parse_request=_parse_ticket, + legacy_wire_format=True, + ) + + @pytest.mark.parametrize("stream", [False, True]) + async def test_parser_must_return_workflow_turn_before_building_or_loading_state(self, stream: bool) -> None: + created: list[Request] = [] + + def create(request: Request) -> Workflow: + created.append(request) + return WorkflowBuilder(name="tickets", start_executor=_TicketExecutor([])).build() + + server = _host(create, parser=lambda _: {"input": "untyped", "stream": stream}) + response, events = await _invoke(server, {}) + assert response.status_code == 400 + assert "WorkflowTurn" in events[0]["error"] + assert created == [] + + @pytest.mark.parametrize( + "field", ["additional_function_arguments", "instructions", "tools", "middleware", "session"] + ) + @pytest.mark.parametrize("shape", ["flat", "options", "executor-options"]) + async def test_client_kwargs_cannot_override_execution_controls(self, field: str, shape: str) -> None: + created: list[Request] = [] + calls: list[tuple[Ticket, str | None]] = [] + + def create(request: Request) -> Workflow: + created.append(request) + return WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build() + + unsafe: dict[str, Any] = {field: {"caller": "forged"}} + if shape == "options": + unsafe = {"options": unsafe} + elif shape == "executor-options": + unsafe = {"ticket-start": {"options": unsafe}} + server = _host(create, parser=lambda _: WorkflowTurn(input=Ticket("T-1", "Help"), client_kwargs=unsafe)) + response, _ = await _invoke(server, {}) + assert response.status_code == 400 + assert len(created) == 1 + assert calls == [] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_built_workflow_allows_one_shot_output_but_rejects_pauses_without_state(self, stream: bool) -> None: + calls: list[tuple[Ticket, str | None]] = [] + one_shot = WorkflowBuilder(name="one-shot", start_executor=_TicketExecutor(calls)).build() + response, events = await _invoke(_host(one_shot), {"ticket_id": "T-1", "stream": stream}) + assert response.status_code == 200 + assert [event["event"] for event in events if event["event"] == "output"] == ["output"] + assert calls == [(Ticket("T-1", "Help"), None)] + store = FoundryWorkflowBindingStore(FoundryRequestScope("session", None, None, False)) + assert await store.get_head("invocations", None) == (None, None) + + pausing = WorkflowBuilder(name="pausing", start_executor=_ReviewExecutor([])).build() + response, events = await _invoke(_host(pausing), {"ticket_id": "T-2", "stream": stream}) + assert response.status_code == (200 if stream else 400) + assert [event["event"] for event in events] == ["error"] + message = events[0].get("message") or events[0].get("error") + assert isinstance(message, str) and "request-aware factory" in message + assert await store.get_head("invocations", None) == (None, None) + + +class TestNativeWorkflowTurns: + async def test_checkpoint_allowlist_failure_happens_before_claim_and_dispatch(self) -> None: + calls: list[tuple[Ticket, str | None]] = [] + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build() + + response, events = await _invoke(_host(factory, checkpoints=CheckpointStoreProvider()), {"ticket_id": "T-1"}) + assert response.status_code == 400 + assert "checkpoint-allowlisted" in events[0]["error"] + assert calls == [] + allowed = CheckpointStoreProvider( + allowed_checkpoint_types=[f"{__name__}:{value.__qualname__}" for value in (Ticket, TicketState)] + ) + response, events = await _invoke(_host(factory, checkpoints=allowed), {"ticket_id": "T-1"}) + assert response.status_code == 200 + assert events[0]["data"]["turn"] == 1 + assert len(calls) == 1 + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("async_factory", [False, True]) + async def test_typed_input_state_and_output_survive_fresh_factories_and_host_replacement( + self, stream: bool, async_factory: bool + ) -> None: + calls: list[tuple[Ticket, str | None]] = [] + created: list[tuple[Request, _TicketExecutor, str | None]] = [] + + def create(request: Request) -> Workflow: + executor = _TicketExecutor(calls) + created.append((request, executor, get_request_context().call_id)) + return WorkflowBuilder(name="tickets", start_executor=executor).build() + + async def create_async(request: Request) -> Workflow: + await asyncio.sleep(0) + return create(request) + + factory: Any = create_async if async_factory else create + for number in (1, 2, 3): + server = _host(factory) + response, events = await _invoke( + server, {"ticket_id": f"T-{number}", "stream": stream}, call_id=f"call-{number}" + ) + assert response.status_code == 200 + outputs = [item["data"] for item in events if item["event"] == "output"] + assert outputs == [ + { + "ticket": {"ticket_id": f"T-{number}", "question": "Help"}, + "turn": number, + "previous_ticket": f"T-{number - 1}" if number > 1 else None, + } + ] + if stream: + assert response.media_type == "text/event-stream" + assert events[-1] == {"event": "done", "session_id": "session"} + assert len(server._session_locks) == 0 + + assert [call_id for _, call_id in calls] == ["call-1", "call-2", "call-3"] + assert [call_id for _, _, call_id in created] == ["call-1", "call-2", "call-3"] + assert len({id(executor) for _, executor, _ in created}) == 3 + assert all(isinstance(request, Request) for request, _, _ in created) + + async def test_local_turn_ids_are_host_owned_when_no_call_id_exists(self) -> None: + calls: list[tuple[Ticket, str | None]] = [] + server = _host(lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build()) + for number in (1, 2): + response, events = await _invoke(server, {"ticket_id": f"T-{number}", "response_id": "caller-collision"}) + assert response.status_code == 200 + assert events[0]["data"]["turn"] == number + assert [call_id for _, call_id in calls] == [None, None] + + async def test_trusted_call_ids_not_body_ids_control_records_and_duplicate_calls_do_not_dispatch(self) -> None: + calls: list[tuple[Ticket, str | None]] = [] + server = _host(lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build()) + server.config.is_hosted = True + server.config.session_id = "" + payload = { + "ticket_id": "T-1", + "response_id": "caller-collision", + "previous_response_id": "forged-parent", + "agent_session_id": "forged-sandbox", + "session_id": "forged-session", + "user_id": "forged-user", + "call_id": "forged-call", + } + for number in (1, 2): + response, events = await _invoke( + server, payload, session_id="sandbox", user_id="user", call_id=f"trusted-{number}" + ) + assert response.status_code == 200 + assert events[0]["data"]["turn"] == number + response, events = await _invoke(server, payload, session_id="sandbox", user_id="user", call_id="trusted-2") + assert response.status_code == 200 + assert events[0]["data"]["turn"] == 2 + assert [call_id for _, call_id in calls] == ["trusted-1", "trusted-2"] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_native_output_keeps_pydantic_application_json(self, stream: bool) -> None: + class ModelOutput(Executor): + @handler + async def handle(self, ticket: Ticket, ctx: WorkflowContext[None, TicketModel]) -> None: + await ctx.yield_output(TicketModel(ticket_id=ticket.ticket_id, due=date(2026, 10, 1))) + + server = _host(lambda _: WorkflowBuilder(name="model-output", start_executor=ModelOutput("start")).build()) + response, events = await _invoke(server, {"ticket_id": "T-1", "stream": stream}) + assert response.status_code == 200 + assert [item["data"] for item in events if item["event"] == "output"] == [ + { + "ticket_id": "T-1", + "due": "2026-10-01", + } + ] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_hosted_workflow_isolates_same_user_sandboxes_and_other_users(self, stream: bool) -> None: + calls: list[tuple[Ticket, str | None]] = [] + identities = [("sandbox-a", "user-a"), ("sandbox-b", "user-a"), ("sandbox-a", "user-b")] + for round_number in (1, 2): + for session_id, user_id in identities: + server = _host(lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build()) + server.config.is_hosted = True + server.config.session_id = "" + response, events = await _invoke( + server, + {"ticket_id": f"T-{round_number}", "stream": stream}, + session_id=session_id, + user_id=user_id, + call_id=f"call-{round_number}-{session_id}-{user_id}", + ) + assert response.status_code == 200 + outputs = [item["data"] for item in events if item["event"] == "output"] + assert outputs[0]["turn"] == round_number + + async def test_restore_uses_exact_cursor_instead_of_an_unrelated_later_checkpoint(self) -> None: + calls: list[tuple[Ticket, str | None]] = [] + stores: list[Any] = [] + + class RecordingCheckpoints(CheckpointStoreProvider): + def get_store_for_scope( + self, *, scope: Any, context_id: str, platform_context: FoundryAgentRequestContext + ) -> Any: + store = super().get_store_for_scope( + scope=scope, context_id=context_id, platform_context=platform_context + ) + stores.append(store) + return store + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build() + + provider = RecordingCheckpoints( + allowed_checkpoint_types=[f"{__name__}:{value.__qualname__}" for value in (Ticket, TicketState)] + ) + first, _ = await _invoke(_host(factory, checkpoints=provider), {"ticket_id": "T-1"}) + assert first.status_code == 200 + unrelated = WorkflowBuilder(name="tickets", start_executor=_TicketExecutor([])).build() + with _context(): + await unrelated.run(Ticket("unrelated", "Different run"), checkpoint_storage=stores[0]) + latest = await stores[0].get_latest(workflow_name="tickets") + assert latest.checkpoint_id == unrelated.get_last_checkpoint_id() + with patch.object( + FoundryCheckpointStore, "get_latest", side_effect=AssertionError("must restore exact cursor") + ): + response, events = await _invoke(_host(factory, checkpoints=provider), {"ticket_id": "T-2"}) + assert response.status_code == 200 + assert events[0]["data"]["turn"] == 2 + assert events[0]["data"]["previous_ticket"] == "T-1" + + async def test_multiple_host_restarts_retain_only_current_response_and_checkpoint(self) -> None: + calls: list[tuple[Ticket, str | None]] = [] + stores: list[FoundryCheckpointStore] = [] + + class RecordingCheckpoints(CheckpointStoreProvider): + def get_store_for_scope( + self, *, scope: Any, context_id: str, platform_context: FoundryAgentRequestContext + ) -> FoundryCheckpointStore: + store = cast( + FoundryCheckpointStore, + super().get_store_for_scope( + scope=scope, + context_id=context_id, + platform_context=platform_context, + ), + ) + stores.append(store) + return store + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build() + + provider = RecordingCheckpoints( + allowed_checkpoint_types=[f"{__name__}:{value.__qualname__}" for value in (Ticket, TicketState)] + ) + for number in range(1, 6): + response, events = await _invoke( + _host(factory, checkpoints=provider), + {"ticket_id": f"T-{number}"}, + session_id="retained", + call_id=f"call-{number}", + ) + assert response.status_code == 200 + assert events[0]["data"]["turn"] == number + + binding_store = FoundryWorkflowBindingStore(FoundryRequestScope("retained", None, None, False)) + for number in range(1, 5): + assert await binding_store.get_response(f"call-{number}") == (None, None) + current, _ = await binding_store.get_response("call-5") + head, _ = await binding_store.get_head("invocations", None) + assert current is not None and current.status == "completed" + assert head is not None and head.binding == current.binding + raw_bindings = await binding_store._get_store() # pyright: ignore[reportPrivateUsage] + async with raw_bindings: + keys = await raw_bindings.list_keys(call_id=None) + assert len(keys.keys) == 2 # Current response record plus the fixed lineage head. + # The current turn retains its entry and final checkpoints; prior turns are reclaimed. + assert len(await stores[-1].list_checkpoint_ids(workflow_name="tickets")) == 2 + + before = list(calls) + replay, events = await _invoke( + _host(factory, checkpoints=provider), + {"ticket_id": "T-5"}, + session_id="retained", + call_id="call-5", + ) + assert replay.status_code == 200 + assert events[0]["data"]["turn"] == 5 + assert calls == before + + @pytest.mark.parametrize("change", ["name", "executor"]) + async def test_graph_changes_fail_before_executor_dispatch(self, change: str) -> None: + calls: list[tuple[Ticket, str | None]] = [] + first = _host(lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build()) + assert (await _invoke(first, {"ticket_id": "first"}))[0].status_code == 200 + changed = _host( + lambda _: WorkflowBuilder( + name="renamed" if change == "name" else "tickets", + start_executor=_TicketExecutor(calls, id="renamed" if change == "executor" else "ticket-start"), + ).build() + ) + response, _ = await _invoke(changed, {"ticket_id": "second"}) + assert response.status_code == 409 + assert len(calls) == 1 + + @pytest.mark.parametrize("stream", [False, True]) + async def test_nonserializable_output_is_a_visible_failure_without_done(self, stream: bool) -> None: + class UnsupportedOutput(Executor): + @handler + async def handle(self, _: Ticket, ctx: WorkflowContext[None, Any]) -> None: + await ctx.yield_output({"unsupported": object()}) + + server = _host( + lambda _: WorkflowBuilder(name="invalid-output", start_executor=UnsupportedOutput("start")).build() + ) + response, events = await _invoke(server, {"ticket_id": "T-1", "stream": stream}) + assert response.status_code == (200 if stream else 500) + assert [item["event"] for item in events] == ["error"] + message = events[0].get("message") or events[0].get("error") + assert isinstance(message, str) + assert "JSON-serializable" in message + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("error_type", [RuntimeError, ValueError]) + async def test_workflow_failure_is_sanitized_and_never_reports_done( + self, stream: bool, error_type: type[Exception] + ) -> None: + class Broken(Executor): + @handler + async def handle(self, _: Ticket, ctx: WorkflowContext[None, Any]) -> None: + raise error_type("private provider-token and checkpoint-id") + + server = _host(lambda _: WorkflowBuilder(name="broken", start_executor=Broken("start")).build()) + response, events = await _invoke(server, {"ticket_id": "T-1", "stream": stream}) + assert response.status_code == (200 if stream else 500) + assert [item["event"] for item in events] == ["error"] + assert "private" not in json.dumps(events) + + +class TestNativeWorkflowPendingReplies: + @pytest.mark.parametrize("stream", [False, True]) + async def test_json_approval_metadata_is_bound_to_exact_pending_calls_before_dispatch(self, stream: bool) -> None: + decisions: list[bool] = [] + + async def parser(request: Request) -> WorkflowTurn[Ticket]: + payload = await request.json() + if "responses" in payload: + return WorkflowTurn( + responses={key: Content.from_dict(value) for key, value in payload["responses"].items()}, + stream=payload.get("stream", False), + ) + return WorkflowTurn(input=Ticket(payload["ticket_id"], "Help"), stream=payload.get("stream", False)) + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="approvals", start_executor=_ApprovalExecutor(decisions)).build() + + _, events = await _invoke(_host(factory, parser=parser), {"ticket_id": "T-1", "stream": stream}) + pending = [event for event in events if event["event"] == "request_info"] + assert len(pending) == 2 + valid = { + item["request_id"]: Content.from_dict(item["data"]).to_function_approval_response(True).to_dict() + for item in pending + } + request_ids = list(valid) + duplicated = {request_ids[0]: valid[request_ids[0]], request_ids[1]: valid[request_ids[0]]} + response, rejected = await _invoke(_host(factory, parser=parser), {"responses": duplicated, "stream": stream}) + assert response.status_code == (200 if stream else 400) + assert [event["event"] for event in rejected] == ["error"] + assert decisions == [] + forged = json.loads(json.dumps(valid)) + forged[request_ids[0]]["function_call"]["arguments"] = {"ticket_id": "forged"} + response, rejected = await _invoke(_host(factory, parser=parser), {"responses": forged, "stream": stream}) + assert response.status_code == (200 if stream else 400) + assert [event["event"] for event in rejected] == ["error"] + assert decisions == [] + response, _ = await _invoke(_host(factory, parser=parser), {"responses": valid, "stream": stream}) + assert response.status_code == 200 + assert decisions == [True, True] + response, replayed = await _invoke(_host(factory, parser=parser), {"responses": valid, "stream": stream}) + assert response.status_code == (200 if stream else 400) + assert [event["event"] for event in replayed] == ["error"] + assert decisions == [True, True] + + async def test_unallowlisted_reply_batch_is_rejected_before_claim_and_valid_retry_succeeds(self) -> None: + replies: list[object] = [] + provider = CheckpointStoreProvider() + + async def parser(request: Request) -> WorkflowTurn[str]: + payload = await request.json() + if "responses" not in payload: + return WorkflowTurn(input=payload["input"]) + values = { + request_id: ( + UnallowlistedReply(value["value"]) + if isinstance(value, dict) and value.get("kind") == "custom" + else value + ) + for request_id, value in payload["responses"].items() + } + return WorkflowTurn(responses=values) + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="any-reply", start_executor=_AnyReplyExecutor(replies)).build() + + response, events = await _invoke(_host(factory, parser=parser, checkpoints=provider), {"input": "review"}) + assert response.status_code == 200 + request_ids = [event["request_id"] for event in events if event["event"] == "request_info"] + assert request_ids == ["first", "second"] + store = FoundryWorkflowBindingStore(FoundryRequestScope("session", None, None, False)) + before, before_etag = await store.get_head("invocations", None) + + invalid = { + "responses": { + "first": {"accepted": True}, + "second": {"kind": "custom", "value": "private-token"}, + } + } + response, rejected = await _invoke(_host(factory, parser=parser, checkpoints=provider), invalid) + assert response.status_code == 400 + assert [event["event"] for event in rejected] == ["error"] + assert "checkpoint-allowlisted" in rejected[0]["error"] + assert "private-token" not in rejected[0]["error"] + after, after_etag = await store.get_head("invocations", None) + assert (after, after_etag) == (before, before_etag) + assert replies == [] + + valid = {"responses": {"first": {"accepted": True}, "second": {"accepted": False}}} + response, _ = await _invoke(_host(factory, parser=parser, checkpoints=provider), valid) + assert response.status_code == 200 + assert replies == [{"accepted": True}, {"accepted": False}] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_partial_replies_do_not_consume_authority_and_complete_replies_survive_restart( + self, stream: bool + ) -> None: + replies: list[tuple[str, bool]] = [] + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="reviews", start_executor=_ReviewExecutor(replies)).build() + + response, events = await _invoke(_host(factory), {"ticket_id": "T-1", "stream": stream}) + assert response.status_code == 200 + pending = [item for item in events if item["event"] == "request_info"] + assert len(pending) == 2 + ids = {item["data"]["reviewer"]: item["request_id"] for item in pending} + first_reply = {"responses": {ids["support"]: {"approved": True}}, "stream": stream} + response, events = await _invoke(_host(factory), first_reply) + assert response.status_code == (200 if stream else 400) + assert [item["event"] for item in events] == ["error"] + assert replies == [] + + complete = { + "responses": {ids["support"]: {"approved": True}, ids["security"]: {"approved": False}}, + "stream": stream, + } + response, events = await _invoke(_host(factory), complete) + assert response.status_code == 200 + assert sorted(replies) == [("security", False), ("support", True)] + assert not [item for item in events if item["event"] == "request_info"] + + response, events = await _invoke(_host(factory), complete) + assert response.status_code == (200 if stream else 400) + assert [item["event"] for item in events] == ["error"] + assert len(replies) == 2 + + @pytest.mark.parametrize("identity", ["other-sandbox", "other-user", "forged-request"]) + async def test_pending_reply_cannot_cross_scope_or_forge_a_request(self, identity: str) -> None: + replies: list[tuple[str, bool]] = [] + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="reviews", start_executor=_ReviewExecutor(replies)).build() + + server = _host(factory) + server.config.is_hosted = True + server.config.session_id = "" + response, events = await _invoke( + server, {"ticket_id": "T-1"}, session_id="sandbox-a", user_id="user-a", call_id="first" + ) + assert response.status_code == 200 + request_id = next(item["request_id"] for item in events if item["event"] == "request_info") + other = _host(factory) + other.config.is_hosted = True + other.config.session_id = "" + response, _ = await _invoke( + other, + {"responses": {"forged" if identity == "forged-request" else request_id: {"approved": True}}}, + session_id="sandbox-b" if identity == "other-sandbox" else "sandbox-a", + user_id="user-b" if identity == "other-user" else "user-a", + call_id="second", + ) + assert response.status_code == 400 + assert replies == [] + + async def test_new_input_does_not_discard_pending_replies(self) -> None: + replies: list[tuple[str, bool]] = [] + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="reviews", start_executor=_ReviewExecutor(replies)).build() + + first, events = await _invoke(_host(factory), {"ticket_id": "T-1"}) + assert first.status_code == 200 + request_ids = [item["request_id"] for item in events if item["event"] == "request_info"] + response, _ = await _invoke(_host(factory), {"ticket_id": "T-2"}) + assert response.status_code == 400 + response, _ = await _invoke( + _host(factory), {"responses": {request_id: {"approved": True} for request_id in request_ids}} + ) + assert response.status_code == 200 + assert len(replies) == 2 + + async def test_invalid_mixed_reply_batch_does_not_consume_valid_authority(self) -> None: + replies: list[tuple[str, bool]] = [] + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="reviews", start_executor=_ReviewExecutor(replies)).build() + + _, events = await _invoke(_host(factory), {"ticket_id": "T-1"}) + ids = [item["request_id"] for item in events if item["event"] == "request_info"] + response, _ = await _invoke( + _host(factory), {"responses": {ids[0]: {"approved": True}, "forged": {"approved": True}}} + ) + assert response.status_code == 400 + assert replies == [] + response, _ = await _invoke(_host(factory), {"responses": {key: {"approved": True} for key in ids}}) + assert response.status_code == 200 + assert len(replies) == 2 + + +class TestNativeWorkflowConcurrency: + async def test_partial_claim_is_repaired_and_reported_as_blocked(self) -> None: + scope = FoundryRequestScope("session", None, None, False) + store = FoundryWorkflowBindingStore(scope) + owner = "interrupted-owner" + await store.save_head( + WorkflowHead("invocations", None, response_id="interrupted", owner=owner), + expected_etag=None, + ) + server = _host(lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor([])).build()) + response, events = await _invoke(server, {"ticket_id": "T-1"}) + assert response.status_code == 409 + assert events[0]["event"] == "error" + assert events[0]["status"] == 409 + record, _ = await store.get_response("interrupted") + head, _ = await store.get_head("invocations", None) + assert record is not None and record.status == "blocked" and record.owner == owner + assert head is not None and head.blocked + + @pytest.mark.parametrize("same_host", [False, True]) + async def test_same_session_serializes_locally_or_fails_claim_cas_before_dispatch(self, same_host: bool) -> None: + entered = asyncio.Event() + release = asyncio.Event() + runs: list[str] = [] + + class Waiting(Executor): + @handler + async def handle(self, ticket: Ticket, ctx: WorkflowContext[None, str]) -> None: + runs.append(ticket.ticket_id) + if ticket.ticket_id == "first": + entered.set() + await release.wait() + await ctx.yield_output(ticket.ticket_id) + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="concurrent", start_executor=Waiting("start")).build() + + server = _host(factory) + first = asyncio.create_task(_invoke(server, {"ticket_id": "first"}, call_id="call-first")) + second: asyncio.Task[Any] | None = None + try: + await asyncio.wait_for(entered.wait(), timeout=5) + second = asyncio.create_task( + _invoke(server if same_host else _host(factory), {"ticket_id": "second"}, call_id="call-second") + ) + if same_host: + await asyncio.sleep(0.05) + assert runs == ["first"] + assert not second.done() + else: + second_response, events = await asyncio.wait_for(second, timeout=5) + assert second_response.status_code == 409 + assert [item["event"] for item in events] == ["error"] + assert runs == ["first"] + release.set() + assert (await asyncio.wait_for(first, timeout=5))[0].status_code == 200 + if same_host: + assert (await asyncio.wait_for(second, timeout=5))[0].status_code == 200 + assert runs == ["first", "second"] + assert len(server._session_locks) == 0 + finally: + release.set() + for task in (first, second): + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(*(task for task in (first, second) if task is not None), return_exceptions=True) + + +class TestNativeWorkflowAgentExecutors: + async def test_registry_backed_agent_resources_are_owned_and_reuse_is_rejected(self) -> None: + resources: list[tuple[str, str | None]] = [] + + def fresh(_request: Request) -> Workflow: + options: ChatOptions = {"store": False} + agent = Agent(client=_ContextClient([], resources), name="registered", default_options=options) + return WorkflowBuilder(name="registry", start_executor=_RegistryExecutor(agent)).build() + + for number in (1, 2): + response, _ = await _invoke(_host(fresh), {"ticket_id": f"T-{number}"}, call_id=f"call-{number}") + assert response.status_code == 200 + assert resources == [ + ("enter", "call-1"), + ("exit", "call-1"), + ("enter", "call-2"), + ("exit", "call-2"), + ] + + options: ChatOptions = {"store": False} + shared = Agent(client=_ContextClient([], resources), name="shared", default_options=options) + + def reused(_request: Request) -> Workflow: + return WorkflowBuilder(name="registry", start_executor=_RegistryExecutor(shared)).build() + + server = _host(reused) + assert (await _invoke(server, {"ticket_id": "T-3"}, session_id="reused"))[0].status_code == 200 + response, _ = await _invoke(server, {"ticket_id": "T-4"}, session_id="reused") + assert response.status_code == 500 + + @pytest.mark.parametrize("failure", ["store", "overrides"]) + async def test_raw_factory_rejects_unsafe_or_unmaterialized_defaults_before_resources_and_dispatch( + self, failure: str + ) -> None: + calls: list[dict[str, Any]] = [] + resources: list[tuple[str, str | None]] = [] + + def factory(_request: Request) -> Workflow: + options: ChatOptions = { + "store": failure == "store", + "temperature": 0.2 if failure == "overrides" else 0.7, + } + agent = RawAgent( + client=_ContextClient(calls, resources), + name="raw", + default_options=options, + ) + return WorkflowBuilder(name="raw-agent", start_executor=AgentExecutor(agent, id="responder")).build() + + server = _host( + factory, + parser=lambda _: WorkflowTurn(input="hello", client_kwargs={"options": {"temperature": 0.7}}), + ) + response, _ = await _invoke(server, {}) + assert response.status_code == 400 + assert calls == resources == [] + + async def test_borrowed_raw_workflow_cannot_forge_fresh_factory_authority(self) -> None: + calls: list[dict[str, Any]] = [] + resources: list[tuple[str, str | None]] = [] + options: ChatOptions = {"store": False, "temperature": 0.7} + agent = RawAgent(client=_ContextClient(calls, resources), name="raw", default_options=options) + workflow = WorkflowBuilder(name="raw-agent", start_executor=AgentExecutor(agent, id="responder")).build() + server = _host( + workflow, + parser=lambda _: WorkflowTurn(input="hello", client_kwargs={"options": {"temperature": 0.7}}), + ) + response, _ = await _invoke(server, {"fresh_factory": True}) + assert response.status_code == 400 + assert calls == resources == [] + + @pytest.mark.parametrize("bare_raw", [False, True]) + @pytest.mark.parametrize("stream", [False, True]) + async def test_unfinished_provider_tokens_are_not_published_or_committed( + self, bare_raw: bool, stream: bool + ) -> None: + class BackgroundClient(_ContextClient): + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Any: + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate( + role="assistant", contents=[Content.from_text("unfinished")], continuation_token={} + ) + + async def complete() -> ChatResponse: + return ChatResponse(messages=[Message("assistant", ["unfinished"])], continuation_token={}) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) if stream else complete() + + def factory(_request: Request) -> Workflow: + agent_type = RawAgent if bare_raw else Agent + options: ChatOptions = {"store": False} + agent = agent_type(client=BackgroundClient([], []), default_options=options) + return WorkflowBuilder(name="background", start_executor=AgentExecutor(agent, id="responder")).build() + + response, events = await _invoke( + _host(factory, parser=lambda _: WorkflowTurn(input="hello", stream=stream)), {} + ) + assert response.status_code == (200 if stream else 500) + assert [event["event"] for event in events] == ["error"] + assert "continuation_token" not in json.dumps(events) + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("bare_raw", [False, True]) + async def test_real_agent_client_options_tool_context_and_owned_resources_refresh_each_turn( + self, stream: bool, bare_raw: bool + ) -> None: + calls: list[dict[str, Any]] = [] + resources: list[tuple[str, str | None]] = [] + tools: list[tuple[str | None, str | None]] = [] + agents: list[RawAgent[Any]] = [] + defaults: list[dict[str, Any]] = [] + + async def factory(request: Request) -> Workflow: + @tool + def inspect_runtime(ctx: FunctionInvocationContext) -> str: + """Inspect only the explicit trusted function context.""" + tools.append((get_request_context().call_id, ctx.kwargs.get("tenant_id"))) + return "checked" + + payload = await request.json() + agent_type = RawAgent if bare_raw else Agent + options: ChatOptions = ( + {"store": False, "temperature": payload["temperature"]} + if bare_raw + else {"store": True, "temperature": 0.2} + ) + agent = agent_type( + client=_ContextClient(calls, resources), + name="responder", + tools=[inspect_runtime], + context_providers=[InMemoryHistoryProvider()], + default_options=options, + ) + agents.append(agent) + defaults.append(dict(agent.default_options)) + return WorkflowBuilder(name="real-agent", start_executor=AgentExecutor(agent, id="responder")).build() + + async def parser(request: Request) -> WorkflowTurn[str]: + payload = await request.json() + return WorkflowTurn( + input=payload["message"], + stream=payload.get("stream", False), + client_kwargs={"options": {"temperature": payload["temperature"]}, "application_marker": "retained"}, + function_invocation_kwargs={"tenant_id": "trusted-tenant"}, + ) + + for number in (1, 2): + server = _host(factory, parser=parser) + server.config.is_hosted = True + server.config.session_id = "" + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: + response = await client.post( + "/invocations", + params={"agent_session_id": "session"}, + headers={"x-agent-user-id": "trusted-user", "x-agent-foundry-call-id": f"trusted-{number}"}, + json={"message": "inspect", "temperature": 0.7, "stream": stream}, + ) + assert response.status_code == 200 + if stream: + assert _sse_events(response.text)[-1] == {"event": "done", "session_id": "session"} + else: + assert any(event["type"] == "output" for event in response.json()["output"]) + assert tools == [("trusted-1", "trusted-tenant"), ("trusted-2", "trusted-tenant")] + assert resources == [ + ("enter", "trusted-1"), + ("exit", "trusted-1"), + ("enter", "trusted-2"), + ("exit", "trusted-2"), + ] + assert [call["call_id"] for call in calls] == ["trusted-1", "trusted-1", "trusted-2", "trusted-2"] + assert all(call["options"]["temperature"] == 0.7 and call["options"]["store"] is False for call in calls) + assert all( + call["kwargs"]["application_marker"] == "retained" and "options" not in call["kwargs"] for call in calls + ) + expected_defaults = {"store": False, "temperature": 0.7} if bare_raw else {"store": True, "temperature": 0.2} + assert all(agent.default_options == original for agent, original in zip(agents, defaults)) + assert all( + all(agent.default_options[key] == value for key, value in expected_defaults.items()) for agent in agents + ) + + +class TestNativeWorkflowStreamingLifecycle: + @pytest.mark.parametrize("stream", [False, True]) + async def test_final_head_cas_failure_never_publishes_pending_authority_or_success(self, stream: bool) -> None: + replies: list[tuple[str, bool]] = [] + save_head = FoundryWorkflowBindingStore.save_head + + async def lose_final_cas( + store: FoundryWorkflowBindingStore, head: WorkflowHead, *, expected_etag: str | None + ) -> str: + return await save_head( + store, + head, + expected_etag="deliberately-stale-etag" if head.response_id is None else expected_etag, + ) + + server = _host(lambda _: WorkflowBuilder(name="reviews", start_executor=_ReviewExecutor(replies)).build()) + with patch.object(FoundryWorkflowBindingStore, "save_head", lose_final_cas): + response, events = await _invoke(server, {"ticket_id": "T-1", "stream": stream}) + assert response.status_code == (200 if stream else 409) + assert [item["event"] for item in events] == ["error"] + assert events[0]["status"] == 409 + assert replies == [] + + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize("limit_kind", ["events", "bytes"]) + async def test_snapshot_retention_limits_abort_without_done_or_cross_session_impact( + self, stream: bool, limit_kind: str + ) -> None: + class Fanout(Executor): + @handler + async def handle(self, _: Ticket, ctx: WorkflowContext[None, str]) -> None: + values = ["one", "two", "three"] if limit_kind == "events" else ["x" * 512] + for value in values: + await ctx.yield_output(value) + + server = _host(lambda _: WorkflowBuilder(name="fanout", start_executor=Fanout("start")).build()) + event_limit = 2 if limit_kind == "events" else 100 + byte_limit = 4096 if limit_kind == "events" else 128 + with ( + patch( + "agent_framework_foundry_hosting._invocations._MAX_WORKFLOW_SNAPSHOT_EVENTS", + event_limit, + ), + patch( + "agent_framework_foundry_hosting._invocations._MAX_WORKFLOW_SNAPSHOT_BYTES", + byte_limit, + ), + ): + response, events = await _invoke(server, {"ticket_id": "T-1", "stream": stream}, session_id="limited") + assert response.status_code == (200 if stream else 500) + assert events[-1]["event"] == "error" + message = events[-1].get("message") or events[-1].get("error") + assert isinstance(message, str) + assert "bounded Invocations snapshot limit" in message + assert not any(event["event"] == "done" for event in events) + store = FoundryWorkflowBindingStore(FoundryRequestScope("limited", None, None, False)) + head, _ = await store.get_head("invocations", None) + assert head is not None and head.blocked + + other, other_events = await _invoke( + _host(lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor([])).build()), + {"ticket_id": "T-2"}, + session_id="unrelated", + ) + assert other.status_code == 200 + assert other_events[0]["event"] == "output" + + async def test_same_call_replays_complete_committed_pending_batch_after_first_frame_disconnect(self) -> None: + replies: list[tuple[str, bool]] = [] + starts: list[str] = [] + delivered: list[dict[str, Any]] = [] + + def factory(_request: Request) -> Workflow: + return WorkflowBuilder(name="reviews", start_executor=_ReviewExecutor(replies, starts)).build() + + server = _host(factory) + server.config.is_hosted = True + server.config.session_id = "" + server.config.sse_keepalive_interval = 0 + received = False + body = json.dumps({"ticket_id": "T-1", "stream": True}).encode() + + async def receive() -> Any: + nonlocal received + if not received: + received = True + return {"type": "http.request", "body": body} + await asyncio.Event().wait() + return {"type": "http.disconnect"} + + async def send(message: Any) -> None: + if message["type"] == "http.response.body": + delivered.extend(_sse_events(message.get("body", b"").decode())) + if len(delivered) == 1: + raise OSError("disconnected after the first durable pending frame") + + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.4"}, + "method": "POST", + "scheme": "http", + "path": "/invocations", + "root_path": "", + "query_string": b"agent_session_id=session", + "headers": [ + (b"content-type", b"application/json"), + (b"x-agent-foundry-call-id", b"pending-call"), + (b"x-agent-user-id", b"user-a"), + ], + "server": ("test", 80), + "client": ("test", 1234), + } + with pytest.raises(ClientDisconnect): + await asyncio.wait_for(server(scope, receive, send), timeout=5) + assert [item["event"] for item in delivered] == ["request_info"] + assert starts == ["T-1"] + trusted_scope = FoundryRequestScope("session", "user-a", "inspect", True) + store = FoundryWorkflowBindingStore(trusted_scope) + before, before_etag = await store.get_head("invocations", None) + + def hosted() -> InvocationsHostServer: + replacement = _host(factory) + replacement.config.is_hosted = True + replacement.config.session_id = "" + return replacement + + headers = {"x-agent-foundry-call-id": "pending-call", "x-agent-user-id": "user-a"} + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=hosted()), base_url="http://test") as client: + replay_json = await client.post( + "/invocations", + params={"agent_session_id": "session"}, + headers=headers, + json={"ticket_id": "T-1", "stream": False}, + ) + assert replay_json.status_code == 200 + complete = replay_json.json()["output"] + assert [item["type"] for item in complete] == ["request_info", "request_info"] + assert complete[0]["request_id"] == delivered[0]["request_id"] + assert [item["data"]["reviewer"] for item in complete] == ["support", "security"] + assert starts == ["T-1"] + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=hosted()), base_url="http://test") as client: + replay_sse = await client.post( + "/invocations", + params={"agent_session_id": "session"}, + headers=headers, + json={"ticket_id": "T-1", "stream": True}, + ) + replay_events = _sse_events(replay_sse.text) + assert [item["event"] for item in replay_events] == ["request_info", "request_info", "done"] + assert [item["request_id"] for item in replay_events[:-1]] == [item["request_id"] for item in complete] + assert starts == ["T-1"] + after, after_etag = await store.get_head("invocations", None) + assert (after, after_etag) == (before, before_etag) + + different_headers = {"x-agent-foundry-call-id": "different-call", "x-agent-user-id": "user-a"} + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=hosted()), base_url="http://test") as client: + different = await client.post( + "/invocations", + params={"agent_session_id": "session"}, + headers=different_headers, + json={"ticket_id": "T-1"}, + ) + assert different.status_code == 400 + assert starts == ["T-1"] + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=hosted()), base_url="http://test") as client: + cross_scope = await client.post( + "/invocations", + params={"agent_session_id": "other-session"}, + headers=headers, + json={"ticket_id": "T-1"}, + ) + assert cross_scope.status_code == 200 + assert [item["request_id"] for item in cross_scope.json()["output"]] != [ + item["request_id"] for item in complete + ] + assert starts == ["T-1", "T-1"] + + reply_headers = {"x-agent-foundry-call-id": "reply-call", "x-agent-user-id": "user-a"} + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=hosted()), base_url="http://test") as client: + response = await client.post( + "/invocations", + params={"agent_session_id": "session"}, + headers=reply_headers, + json={"responses": {item["request_id"]: {"approved": True} for item in complete}}, + ) + assert response.status_code == 200 + assert sorted(replies) == [("security", True), ("support", True)] + + @pytest.mark.parametrize("spec_version", ["2.0", "2.4"]) + @pytest.mark.parametrize("keep_alive", [False, True]) + @pytest.mark.parametrize("executor_kind", ["typed", "agent", "raw"]) + async def test_sdk_disconnect_closes_core_stream_without_success( + self, spec_version: str, keep_alive: bool, executor_kind: str, caplog: pytest.LogCaptureFixture + ) -> None: + entered = asyncio.Event() + closed = asyncio.Event() + disconnected = asyncio.Event() + messages: list[Any] = [] + cleanup_calls: list[str | None] = [] + factory_calls: list[str | None] = [] + client_resources: list[tuple[str, str | None]] = [] + + class Idle(Executor): + @handler + async def handle(self, _: Ticket, ctx: WorkflowContext[None, str]) -> None: + entered.set() + try: + if not keep_alive: + await ctx.yield_output("first") + await asyncio.Event().wait() + finally: + cleanup_calls.append(get_request_context().call_id) + closed.set() + + class SlowClient(_ContextClient): + def _inner_get_response( + self, *, messages: Sequence[Message], stream: bool, options: Mapping[str, Any], **kwargs: Any + ) -> Any: + async def updates() -> AsyncIterator[ChatResponseUpdate]: + entered.set() + try: + if not keep_alive: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("first")]) + await asyncio.Event().wait() + finally: + cleanup_calls.append(get_request_context().call_id) + closed.set() + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + def factory(_request: Request) -> Workflow: + factory_calls.append(get_request_context().call_id) + if executor_kind == "typed": + return WorkflowBuilder(name="idle", start_executor=Idle("start")).build() + agent_type = RawAgent if executor_kind == "raw" else Agent + options: ChatOptions = {"store": False} + agent = agent_type(client=SlowClient([], client_resources), default_options=options) + return WorkflowBuilder(name="idle", start_executor=AgentExecutor(agent, id="start")).build() + + server = _host( + factory, + parser=_parse_ticket if executor_kind == "typed" else lambda _: WorkflowTurn(input="start", stream=True), + ) + server.config.sse_keepalive_interval = 1 if keep_alive else 0 + body = json.dumps({"ticket_id": "T-1", "stream": True}).encode() + received = False + + async def receive() -> Any: + nonlocal received + if not received: + received = True + return {"type": "http.request", "body": body} + await disconnected.wait() + return {"type": "http.disconnect"} + + async def send(message: Any) -> None: + messages.append(message) + expected = b": keep-alive\n\n" if keep_alive else b"event: output" + if message["type"] == "http.response.body" and expected in message.get("body", b""): + assert entered.is_set() + if spec_version == "2.4": + raise OSError("disconnected") + disconnected.set() + await asyncio.Event().wait() + + scope = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": spec_version}, + "method": "POST", + "scheme": "http", + "path": "/invocations", + "root_path": "", + "query_string": b"agent_session_id=session", + "headers": [ + (b"content-type", b"application/json"), + (b"x-agent-foundry-call-id", b"trusted-call"), + (b"x-agent-user-id", b"trusted-user"), + ], + "server": ("test", 80), + "client": ("test", 1234), + } + with _context(session_id="ambient", call_id="ambient-call"): + ambient = get_request_context() + if spec_version == "2.4": + with pytest.raises(ClientDisconnect): + await asyncio.wait_for(server(scope, receive, send), timeout=5) + else: + await asyncio.wait_for(server(scope, receive, send), timeout=5) + assert get_request_context() is ambient + assert closed.is_set() + assert factory_calls == cleanup_calls == ["trusted-call"] + if executor_kind != "typed": + assert client_resources == [("enter", "trusted-call"), ("exit", "trusted-call")] + assert not any(b"event: done" in message.get("body", b"") for message in messages) + assert all(not lock.locked() for lock in server._session_locks.values()) + # Cancelled-task tracebacks can retain released locks until their reference cycles are collected. + gc.collect() + assert len(server._session_locks) == 0 + assert "different Context" not in caplog.text + + async def test_json_and_sse_use_the_real_invocations_asgi_route(self) -> None: + calls: list[tuple[Ticket, str | None]] = [] + server = _host(lambda _: WorkflowBuilder(name="tickets", start_executor=_TicketExecutor(calls)).build()) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as client: + first = await client.post("/invocations", json={"ticket_id": "T-1"}) + assert first.status_code == 200 + assert first.headers["content-type"] == "application/json" + session_id = first.headers["x-agent-session-id"] + assert first.json()["output"][0]["data"]["turn"] == 1 + second = await client.post( + "/invocations", + params={"agent_session_id": session_id}, + json={"ticket_id": "T-2", "stream": True}, + ) + assert second.status_code == 200 + assert second.headers["content-type"].startswith("text/event-stream") + frames = _sse_events(second.text) + assert frames[0]["data"]["turn"] == 2 + assert frames[-1] == {"event": "done", "session_id": session_id} diff --git a/python/packages/foundry_hosting/tests/foundry_hosting/test_workflow_state.py b/python/packages/foundry_hosting/tests/foundry_hosting/test_workflow_state.py index 7dba626a47d..bfa49c9fd44 100644 --- a/python/packages/foundry_hosting/tests/foundry_hosting/test_workflow_state.py +++ b/python/packages/foundry_hosting/tests/foundry_hosting/test_workflow_state.py @@ -402,6 +402,59 @@ async def test_exact_pairs_restore_state_and_reject_old_lineage() -> None: assert head is not None and head.binding == second.binding +async def test_completed_response_replay_is_explicit_and_does_not_advance_the_head() -> None: + first = await _prepare("one", workflow=_workflow(_Review())) + await _finish(first, WorkflowTurn(input="review")) + store = FoundryWorkflowBindingStore(_scope()) + before, before_etag = await store.get_head("one", None) + with pytest.raises(WorkflowConflictError, match="already exists"): + await _prepare("one", workflow=_workflow(_Review())) + + replay = await _prepare("one", workflow=_workflow(_Review()), replay_completed=True) + assert replay.has_completed_output + assert replay.snapshot is not None + assert replay.snapshot["output"] == [] + after, after_etag = await store.get_head("one", None) + assert (after, after_etag) == (before, before_etag) + + +async def test_default_store_integrity_hashes_the_persisted_encoding_without_repickling() -> None: + first = await _prepare("one", workflow=_workflow(_Approval())) + with patch( + "agent_framework_foundry_hosting._workflow_state._checkpoint_hash", + side_effect=AssertionError("default storage must not re-pickle checkpoints for integrity"), + ): + await _finish(first, WorkflowTurn(input="/safe")) + second = await _prepare("two", workflow=_workflow(_Approval()), previous_response_id="one") + assert set(second.pending_requests) == {"approval-1"} + + +async def test_superseded_cleanup_failure_is_retried_after_later_commit( + caplog: pytest.LogCaptureFixture, +) -> None: + first = await _prepare("one", reclaim_superseded=True) + assert await _finish(first, WorkflowTurn(input=1)) == [1] + second = await _prepare("two", previous_response_id="one", reclaim_superseded=True) + with patch.object( + FoundryWorkflowBindingStore, + "delete_response", + AsyncMock(side_effect=RuntimeError("private cleanup detail")), + ): + assert await _finish(second, WorkflowTurn(input=1)) == [2] + store = FoundryWorkflowBindingStore(_scope()) + old, _ = await store.get_response("one") + current, _ = await store.get_response("two") + assert old is not None and current is not None and current.status == "completed" + assert "Failed to reclaim superseded native workflow state" in caplog.text + + third = await _prepare("three", previous_response_id="two", reclaim_superseded=True) + assert await _finish(third, WorkflowTurn(input=1)) == [3] + assert await store.get_response("one") == (None, None) + assert await store.get_response("two") == (None, None) + final, _ = await store.get_response("three") + assert final is not None and final.status == "completed" + + @pytest.mark.parametrize("scope", [_scope(user="other"), _scope(sandbox="other")]) async def test_checkpoint_and_reply_scope_isolation(scope: FoundryRequestScope) -> None: first = await _prepare("one", workflow=_workflow(_Review())) diff --git a/python/samples/04-hosting/foundry-hosted-agents/README.md b/python/samples/04-hosting/foundry-hosted-agents/README.md index c907a4058d1..50ce84c621c 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/README.md @@ -79,6 +79,7 @@ An opaque call ID correlates one request; it is not a storage namespace. | 1 | [Basic](invocations/basic/) | An Invocations agent with a custom JSON parser, durable MAF history, and JSON/SSE responses. | | 2 | [Break Glass](invocations/break_glass/) | An agent demonstrating a "break glass" scenario where customizations of the API behaviors are needed, allowing for more direct control over how requests and responses are handled by the hosting layer. | | 3 | [Telegram](invocations/telegram/) | A Telegram bot routed through API Management to a direct-code hosted agent, with streaming responses and durable Cosmos DB history. | +| 4 | [Native Workflow](invocations/basic/README.md#native-workflow-with-typed-tickets) | An explicit JSON parser maps typed tickets and pending review replies to a freshly built native workflow, with scoped exact-checkpoint continuation and JSON/SSE output. | ## Running the Agent Host Locally diff --git a/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/README.md b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/README.md index baed33d8836..37283c6edce 100644 --- a/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/README.md +++ b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/README.md @@ -80,3 +80,113 @@ breaking change. Follow the [parent deployment guide](../../README.md#deploying-the-agent-to-foundry) when deploying this example to Foundry. + +## Native workflow with typed tickets + +[`workflow.py`](workflow.py) is an alternative entrypoint, not a wrapper around +the agent above. It hosts a native `Workflow` with +`InvocationsHostServer(workflow=build_workflow, parse_request=parse_request)`; +there is no `workflow.as_agent()`. Its deterministic ticket/review executors +need no model deployment or Azure credentials locally. + +The application parser validates its own JSON and returns a +`WorkflowTurn(input=Ticket(...), stream=...)` or +`WorkflowTurn(responses={request_id: TicketDecision(...)}, stream=...)`. +The parser checks reply types; the host separately checks that each reply +belongs to the exact pending checkpoint in the trusted sandbox and user. +It rejects unknown fields rather than forwarding caller-controlled +instructions, tools, middleware, session IDs or tool context as run options. +When adding agents or tools, keep their defaults and trusted context in the +request-aware factory. `client_kwargs` and `function_invocation_kwargs` keep +their existing workflow meanings; they are not arbitrary entry-state updates. + +Use the current workspace until the coordinated hosting beta containing the +native workflow API is published. From `python/`, run: + +```bash +uv run --no-sync python samples/04-hosting/foundry-hosted-agents/invocations/basic/workflow.py +``` + +After that release, the script's PEP 723 metadata also supports +`uv run --script --prerelease=allow workflow.py`. Do not change a +`pyproject.toml` to install sample-only imports. This entrypoint does not read +a `.env` file. + +### Application JSON and SSE + +For an immediate typed result, opt out of review: + +```bash +curl -i -X POST http://localhost:8088/invocations \ + -H "Content-Type: application/json" \ + -d '{"ticket_id":"T-1","question":"Where is my order?","require_review":false}' +``` + +Non-streaming workflow responses are application JSON, with typed payloads in +the `output` event list. A ticket produces a result containing its +`ticket_id`, persisted `turn` number and `status`. Dataclasses, Pydantic models +and framework `Message`, `Content` and `AgentResponse` values retain their +supported structured encoding; they are not coerced to text. Unsupported +objects fail visibly instead of being stringified or omitted. + +By default, tickets pause for a review. Send a new ticket with +`"stream": true` to receive framed `request_info` SSE containing a `request_id` +and the typed ticket review. Reuse that ID in the next request: + +```bash +curl -i -N -X POST \ + "http://localhost:8088/invocations?agent_session_id=" \ + -H "Content-Type: application/json" \ + -d '{"responses":{"":{"approved":true}},"stream":true}' +``` + +The reply emits an `output` event with status `approved` (or `rejected` for +`false`). A final `done` event contains the platform `session_id`, and is +emitted only after the exact workflow cursor has been conditionally saved. +Live output is provisional until `done`. Authority-bearing `request_info` +frames and any following frames are buffered until the cursor is committed, +so emitted request IDs refer to durable pending state. +If the connection drops after commit, retry the request with the same trusted +user, sandbox and platform call ID to replay the complete stored response +without executing the ticket workflow again. +Failures produce `error`, not `done`. Non-streaming pauses expose the same +`request_info` data in the JSON event list. Native workflows do not support +`legacy_wire_format=True`; the old plain-text/raw-chunk opt-in applies only +to the agent entrypoint. + +### Checkpoints, fresh factories and pending replies + +`build_workflow(request)` builds fresh executors on every request. The graph +name `ticket-workflow` and executor ID `ticket_start` remain stable. Typed +`TicketState` is saved in workflow state through `ctx.set_state`, and restored +from the exact prior checkpoint even after replacing the host process. The +checkpoint provider explicitly allowlists this example's application types +for decoding; it does not load arbitrary caller-specified Python classes. +An incompatible graph is rejected before dispatch. Do not configure a +second checkpoint store on `WorkflowBuilder` or share a mutable workflow +instance between requests. + +Finish pending reviews before submitting another ticket. Replies cannot +invent a request ID, cross users or sandboxes, replay an already-consumed +request, or reuse a stale checkpoint's authority. If a graph has multiple +pending requests, answer the **complete batch** in one turn. Partial or +invalid batches are rejected before any reply is consumed or handler runs; +resubmit the complete valid batch. A failed or interrupted turn does not +report successful completion. Its claimed lineage is blocked instead of +replaying possibly completed effects; start a new Foundry sandbox. +Conditional storage +conflicts are visible; checkpoints do not provide exactly-once external tool +side effects, so use idempotent application actions. + +Continue using the **platform** sandbox ID from `x-agent-session-id` in the +`agent_session_id` query. It routes to the Foundry sandbox, whereas MAF +checkpoints and pending request IDs identify state *inside* that sandbox. +Body/header IDs cannot change platform routing. Hosted requests also require +trusted user and call IDs; without `FOUNDRY_AGENT_SESSION_ID`, use an explicitly +created hosted session and a matching routed query, as described above. +Local requests retain the SDK's single-user session fallback. + +For hosted deployment, use the parent guide with `workflow.py` as the Python +entrypoint (or the container command). The existing basic manifest and +Dockerfile still launch `main.py`; change the entrypoint deliberately rather +than expecting this alternative to replace the agent sample. diff --git a/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/workflow.py b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/workflow.py new file mode 100644 index 00000000000..a810f3ec933 --- /dev/null +++ b/python/samples/04-hosting/foundry-hosted-agents/invocations/basic/workflow.py @@ -0,0 +1,166 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "agent-framework-core", +# "agent-framework-foundry-hosting", +# ] +# /// +# Copyright (c) Microsoft. All rights reserved. + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Literal + +from agent_framework import Executor, Workflow, WorkflowBuilder, WorkflowContext, handler, response_handler +from agent_framework_foundry_hosting import CheckpointStoreProvider, InvocationsHostServer, WorkflowTurn +from starlette.requests import Request + +"""Host typed tickets and durable review replies as a native Invocations workflow. + +This deterministic example needs no model or Azure credentials to run locally. +The host restores workflow state, not a shared Workflow or workflow.as_agent(). +""" + + +@dataclass(frozen=True) +class Ticket: + ticket_id: str + question: str + require_review: bool = True + + +@dataclass(frozen=True) +class TicketState: + turn: int = 0 + last_ticket_id: str | None = None + + +@dataclass(frozen=True) +class TicketReview: + ticket_id: str + question: str + turn: int + + +@dataclass(frozen=True) +class TicketDecision: + approved: bool + + +@dataclass(frozen=True) +class TicketResult: + ticket_id: str + turn: int + status: Literal["received", "approved", "rejected"] + + +# 1. Use a typed start executor and checkpoint-supported application state. +class TicketStart(Executor): + def __init__(self) -> None: + super().__init__(id="ticket_start") + + @handler + async def run(self, ticket: Ticket, ctx: WorkflowContext[None, TicketResult]) -> None: + state = ctx.get_state("tickets", TicketState()) + if not isinstance(state, TicketState): + raise TypeError("The restored tickets state must be TicketState.") + turn = state.turn + 1 + ctx.set_state("tickets", TicketState(turn=turn, last_ticket_id=ticket.ticket_id)) + if ticket.require_review: + await ctx.request_info( + TicketReview(ticket.ticket_id, ticket.question, turn), + response_type=TicketDecision, + ) + else: + await ctx.yield_output(TicketResult(ticket.ticket_id, turn, "received")) + + @response_handler + async def review( + self, + request: TicketReview, + decision: TicketDecision, + ctx: WorkflowContext[None, TicketResult], + ) -> None: + await ctx.yield_output( + TicketResult(request.ticket_id, request.turn, "approved" if decision.approved else "rejected") + ) + + +# 2. Build fresh executors for each request, with stable graph and executor IDs. +def build_workflow(_request: Request) -> Workflow: + return WorkflowBuilder(name="ticket-workflow", start_executor=TicketStart()).build() + + +# 3. The application owns JSON validation and typed input/reply conversion. +def unique_json_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + values: dict[str, Any] = {} + for key, value in pairs: + if key in values: + raise ValueError("Duplicate JSON fields are not allowed.") + values[key] = value + return values + + +async def parse_request(request: Request) -> WorkflowTurn[Ticket]: + payload: Any = json.loads(await request.body(), object_pairs_hook=unique_json_object) + if not isinstance(payload, dict): + raise ValueError("The request must be a JSON object.") + stream = payload.get("stream", False) + if not isinstance(stream, bool): + raise ValueError("stream must be a boolean.") + + if "responses" in payload: + if set(payload) - {"responses", "stream"}: + raise ValueError("A review reply accepts only responses and stream.") + responses = payload["responses"] + if not isinstance(responses, dict) or not responses: + raise ValueError("responses must be a nonempty object keyed by pending request IDs.") + decisions: dict[str, TicketDecision] = {} + for request_id, response in responses.items(): + if not isinstance(request_id, str) or not request_id: + raise ValueError("Each response needs a nonempty request ID.") + if ( + not isinstance(response, dict) + or set(response) != {"approved"} + or not isinstance(response["approved"], bool) + ): + raise ValueError("Each response must contain exactly one boolean approved field.") + decisions[request_id] = TicketDecision(approved=response["approved"]) + return WorkflowTurn(responses=decisions, stream=stream) + + if set(payload) - {"ticket_id", "question", "require_review", "stream"}: + raise ValueError("A ticket accepts only ticket_id, question, require_review and stream.") + ticket_id = payload.get("ticket_id") + question = payload.get("question") + require_review = payload.get("require_review", True) + if not isinstance(ticket_id, str) or not ticket_id.strip() or not isinstance(question, str) or not question.strip(): + raise ValueError("ticket_id and question must be nonempty strings.") + if not isinstance(require_review, bool): + raise ValueError("require_review must be a boolean.") + return WorkflowTurn(input=Ticket(ticket_id, question, require_review), stream=stream) + + +# 4. Allow only these application types when decoding durable checkpoints. +def main() -> None: + checkpoints = CheckpointStoreProvider( + allowed_checkpoint_types=[ + f"{__name__}:{value.__qualname__}" + for value in (Ticket, TicketState, TicketReview, TicketDecision, TicketResult) + ] + ) + InvocationsHostServer( + workflow=build_workflow, + parse_request=parse_request, + checkpoint_store_provider=checkpoints, + ).run() + + +if __name__ == "__main__": + main() + +# A new ticket normally emits request_info with its typed TicketReview payload. +# Reply with {"responses": {"": {"approved": true}}} in the same sandbox. +# require_review=false produces a typed TicketResult immediately. +# stream=true emits framed request_info/output, then done only after the cursor is saved.