diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 05486f1d285..ab01c7005a3 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -131,6 +131,47 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. `single_user=True` assigns one opaque authorization scope for the lifetime of that endpoint registration and must not be used for callers with distinct authorization. Use an explicit resolver when snapshots, checkpoints, or provider state must retain the same scope across process restarts. +- An `AgentExecutor` keeps one `AgentSession` (history, pending approvals, standing approval rules) between runs, and + that session is not tied to a caller. `_session_ownership.claim_agent_executors` binds every `AgentExecutor` of a + workflow, and the `AgentSession` each one holds when the run starts, to the first Snapshot Scope that runs it. It + looks into nested `WorkflowExecutor` workflows and into the workflow behind a `WorkflowAgent` that an + `AgentExecutor` wraps. A run without a scope is an owner of its own, so a workflow that ran without a scope refuses + a later scope and the reverse. The registry (`IdentityOwners`) keys objects by `id()` with a weak reference and + never calls their `__hash__` or `__eq__`, so unhashable subclasses work and equal objects stay apart. The claim is + one process-wide locked check-and-bind. `AgentFrameworkWorkflow.run` and the `WorkflowAgent` branch of + `run_agent_stream` (through `claim_workflow_agent`, which also looks through A2UI runners) call it after a snapshot + replay has returned and before any workflow runs, so a replay never takes ownership. Another scope receives + `RUN_STARTED` and `RUN_ERROR` with `WORKFLOW_SCOPE_CONFLICT` and never reaches the model. It covers a shared + `Workflow` or `WorkflowAgent`, and a `workflow_factory` that returns shared workflows, executors, or an injected + `AgentSession`. It follows live objects only and does not cover: + - state that a factory loads from storage into a new object (`AgentSession.from_dict()`, `SessionStore.get()`, a + history provider that keys data by `session_id`); + - sessions that alias state outside the session object (a shared `state` mapping or `service_session_id`); + - executors that are not `AgentExecutor` but keep a session or cache (custom executors, the group chat + orchestrator); + - workflows without an `AgentExecutor`; + - a session that an executor creates during a checkpoint restore (it stays reachable only through the bound + executor, and the next run of the same scope binds it). + + The rule does not partition the threads of one scope: `single_user=True` or a constant resolver shares one + conversation state across the threads of a shared workflow. Never reset sessions or clone executors to resolve a + conflict; the factory builds new executors and sessions. The factory receives only the thread id, so it must not + restore session state that is keyed by the thread id alone; use `checkpoint_storage`, which AG-UI loads only for the + scope and thread that saved it. +- AG-UI does not support a `WorkflowAgent` that an agent exposes as a tool (`WorkflowAgent.as_tool()`): every caller + of the agent would reach one workflow and its executor sessions, and the ownership claim cannot see that path. + `BaseAgent.as_tool()` keeps the agent it wraps in the private `_wrapped_agent` attribute of the tool, and + `_session_ownership.workflow_agent_tools` follows those links through the tools of an agent, the sub-agents they + delegate to, and the agents of every `AgentExecutor`, also through A2UI runners and nested workflows, with a cycle + guard. `add_agent_framework_fastapi_endpoint` raises `ValueError` (`_reject_workflow_agent_tools`) for what it can + see at registration. `AgentFrameworkWorkflow.run` and `run_agent_stream` end the run with `RUN_ERROR` + `WORKFLOW_AGENT_TOOL_UNSUPPORTED` before the ownership claim, which covers workflow factories and tools added after + registration. The rejection applies with `single_user=True` as well. A tool function that calls a workflow from + custom code, and a tool that a context provider adds while a run is prepared, are not detected. They are unsupported + all the same: exposing a `WorkflowAgent` as a tool is not supported in any form, and AG-UI gives no isolation to its + callers. The search needs a core whose `as_tool()` records `_wrapped_agent`; with an older core it finds nothing, so + raise this package's `agent-framework-core` floor to the first release that has it when you prepare a release. + Remove this check together with `WorkflowAgent.as_tool()`. - Request `state` is client-owned Shared State and is merged into `AgentSession.state`, minus a protected set: tool-approval state, history/context-provider namespaces, message-injection state, standalone FIDES security state, and provider-owned keys. diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 4a41670adba..0101059166a 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -65,9 +65,15 @@ app = FastAPI() add_agent_framework_fastapi_endpoint(app, workflow, "/", single_user=True) ``` +One workflow instance serves every request, so its executors keep their state across callers. A shared workflow that +holds an `AgentExecutor` (its `AgentSession` keeps history, pending approvals, and standing approval rules) serves only +the first authorization scope that runs it. Use `workflow_factory`, as shown next, to serve several scopes. + ### Server (Thread-Scoped WorkflowBuilder) -Use `workflow_factory` when your workflow keeps runtime state (for example pending `request_info` interrupts) and must be isolated per AG-UI thread: +Use `workflow_factory` when your workflow keeps runtime state (for example pending `request_info` interrupts or +`AgentExecutor` sessions) and must be isolated per authorization scope and AG-UI thread. Build the workflow and its +executors inside the factory: ```python from fastapi import FastAPI @@ -76,7 +82,7 @@ from agent_framework.ag_ui import AgentFrameworkWorkflow, add_agent_framework_fa def build_workflow_for_thread(thread_id: str) -> Workflow: - # Build a fresh workflow instance for each thread id. + # Build a fresh workflow instance, including new AgentExecutor instances, for each thread id. return WorkflowBuilder(start_executor=...).build() @@ -350,6 +356,40 @@ Authentication in the separate `dependencies` list can allow or reject a request server-side state. Do not combine authentication-only dependencies with `single_user=True` when different authenticated callers can reach the endpoint. +### Workflows That Hold Agent Sessions + +An `AgentExecutor` keeps one `AgentSession` between runs. Conversation history, pending approval requests, and +standing "always approve" rules live there, and the session is not tied to a caller. An endpoint therefore binds each +`AgentExecutor`, and the session it holds, to the first authorization scope that runs it. A shared `Workflow` or +`WorkflowAgent` that holds one serves only that scope. A run in another scope ends with a `WORKFLOW_SCOPE_CONFLICT` +`RUN_ERROR` before it reaches the model. A run without a scope owns what it runs in the same way, so a workflow that ran +locally is not handed to a scope later. A snapshot replay starts no workflow and does not take ownership. Pass +`AgentFrameworkWorkflow(workflow_factory=...)` with a factory that builds a new workflow, including new `AgentExecutor` +and `AgentSession` instances, for every call to serve several scopes. A factory that returns a workflow, an executor, or +a session it shares with another scope is refused the same way. Workflows without an `AgentExecutor` are not bound to a +scope; their pending interrupts stay fenced to the scope and thread that raised them. + +With `single_user=True`, or with a resolver that returns one constant scope, a shared workflow keeps one conversation +state for every thread. Use `workflow_factory` when threads must stay independent. + +The endpoint follows live `AgentExecutor` and `AgentSession` objects. It cannot see state that a factory loads from +storage into a new object. For example, `AgentSession.from_dict()` and `SessionStore.get()` return new objects with the +same session id and state, so two scopes that load one stored session share its history. The factory receives only the +thread id, and it runs once for each scope and thread. AG-UI does not pass it the authorization scope. A factory must +therefore start every session with fresh, independent state. Do not restore session state that is keyed by the thread +id alone. To continue a conversation, use `checkpoint_storage`: AG-UI restores a checkpoint only for the scope and +thread that saved it. The rule does not cover executors that are not `AgentExecutor`, or state that sessions share +outside the session object, for example a `state` dictionary or a `service_session_id`. + +AG-UI does not support an agent that exposes a `WorkflowAgent` as a tool (`WorkflowAgent.as_tool()`), however the tool +gets into the agent, because every caller of that agent would reach the same workflow and its executor sessions. AG-UI +rejects what it can see. `add_agent_framework_fastapi_endpoint` raises `ValueError` for an agent or workflow that lists +one, with `single_user=True` as well. When a workflow factory builds one, the run ends with a +`WORKFLOW_AGENT_TOOL_UNSUPPORTED` `RUN_ERROR` before it reaches the model. AG-UI cannot see a workflow tool that a +context provider adds while a run is prepared, or a tool function of your own that calls a workflow. Those stay +unsupported, and AG-UI does not isolate their callers. Serve the workflow with +`AgentFrameworkWorkflow(workflow_factory=...)`, or build the workflow inside the tool function for every call. + ### Stable Single-User Scope `single_user=True` creates an opaque scope for one endpoint registration, which is appropriate for local and diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 436607114a9..75a78aa49ab 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -114,6 +114,16 @@ _resolve_ui_payload, # type: ignore _track_tool_call_segment, # type: ignore ) +from ._session_ownership import ( + SCOPE_CONFLICT_CODE, + SCOPE_CONFLICT_LOG_MESSAGE, + SCOPE_CONFLICT_MESSAGE, + WORKFLOW_AGENT_TOOL_CODE, + WORKFLOW_AGENT_TOOL_MESSAGE, + claim_workflow_agent, + workflow_agent_tool_error, + workflow_agent_tools, +) from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, _SNAPSHOT_SCOPE_INPUT_KEY, @@ -3029,6 +3039,19 @@ async def _run_agent_stream( yield event return + # A replay above runs no agent, so only a run that reaches the workflow decides who owns its executors. + unsupported_tools = workflow_agent_tools(agent) + if unsupported_tools: + logger.error(workflow_agent_tool_error(agent.name or agent.id, unsupported_tools)) + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent(message=WORKFLOW_AGENT_TOOL_MESSAGE, code=WORKFLOW_AGENT_TOOL_CODE) + return + if not claim_workflow_agent(agent, snapshot_scope): + logger.error(SCOPE_CONFLICT_LOG_MESSAGE, agent.name or agent.id) + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent(message=SCOPE_CONFLICT_MESSAGE, code=SCOPE_CONFLICT_CODE) + return + snapshot_seed_messages: list[dict[str, Any]] | None = None if stored_snapshot is not None: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py index c8248e07fd1..c107bbbd534 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -23,6 +23,7 @@ from ._agent import AgentFrameworkAgent from ._approval_state import _APPROVAL_SCOPE_INPUT_KEY from ._run_common import _is_snapshot_hydration_request +from ._session_ownership import workflow_agent_tool_error, workflow_agent_tools from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, _SNAPSHOT_SCOPE_INPUT_KEY, @@ -120,6 +121,20 @@ def _single_user_scope() -> str: return f"{_SINGLE_USER_SCOPE_PREFIX}{uuid.uuid4().hex}" +def _reject_workflow_agent_tools(protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow) -> None: + """Refuse a runner that exposes a ``WorkflowAgent`` as a tool, so that the mistake shows at startup.""" + # A runner subclass that skips the base constructor has nothing to inspect, and a workflow factory builds its + # workflow for each request, so a run checks what the factory built. + if isinstance(protocol_runner, AgentFrameworkAgent): + subject = getattr(protocol_runner, "agent", None) + else: + subject = getattr(protocol_runner, "workflow", None) + tool_names = workflow_agent_tools(subject) + if tool_names: + runner_name = getattr(protocol_runner, "name", None) or type(protocol_runner).__name__ + raise ValueError(workflow_agent_tool_error(runner_name, tool_names)) + + def add_agent_framework_fastapi_endpoint( app: FastAPI, agent: SupportsAgentRun | AgentFrameworkAgent | Workflow | AgentFrameworkWorkflow, @@ -145,7 +160,9 @@ def add_agent_framework_fastapi_endpoint( Args: app: The FastAPI application - agent: The agent to expose (can be raw SupportsAgentRun or wrapped) + agent: The agent to expose (can be raw SupportsAgentRun or wrapped). The endpoint raises ``ValueError`` for an + agent, or a workflow, that exposes a ``WorkflowAgent`` as a tool (``WorkflowAgent.as_tool()``), and a run + ends with a ``WORKFLOW_AGENT_TOOL_UNSUPPORTED`` ``RUN_ERROR`` when a workflow factory builds one. path: The endpoint path state_schema: Optional state schema for shared state management; accepts dict or Pydantic model/class predict_state_config: Optional predictive state update configuration. @@ -164,7 +181,8 @@ def add_agent_framework_fastapi_endpoint( authorization_scope_resolver for new FastAPI applications because it can consume authenticated dependency results. This resolver also scopes internal Agent Sessions, in-memory workflow_factory instances, and approval state. It must return a non-empty string derived from authorized request context; invalid results - fail the request before accessing state or invoking the runner. + fail the request before accessing state or invoking the runner. The shared-workflow rule described for + authorization_scope_resolver applies to this resolver as well. checkpoint_storage: Optional workflow checkpoint storage, applied when the endpoint exposes a workflow. When provided, each run creates a checkpoint at the end of every superstep, and a run may resume from a persisted checkpoint by supplying its id in the AG-UI forwarded props @@ -188,13 +206,19 @@ def add_agent_framework_fastapi_endpoint( authorization_scope_resolver: Optional FastAPI dependency callable that returns a non-empty authorization scope derived from authenticated request context. The callable may declare normal FastAPI dependencies, headers, security schemes, or the current Request. Its result scopes snapshots, internal Agent Sessions, - workflow instances and checkpoints, detached-run admission, and server-side approval state. + workflow instances and checkpoints, detached-run admission, and server-side approval state. An + ``AgentExecutor`` owns one ``AgentSession`` (history, pending approvals, standing approval rules), so a + shared ``Workflow`` or ``WorkflowAgent`` that holds one serves only the first scope that runs it. A run in + another scope ends with a ``WORKFLOW_SCOPE_CONFLICT`` ``RUN_ERROR`` before it reaches the model. Pass + ``AgentFrameworkWorkflow(workflow_factory=...)`` with a factory that builds a new workflow, including new + ``AgentExecutor`` and ``AgentSession`` instances, for every call to serve several scopes. single_user: Explicitly declare that every caller of this endpoint belongs to one authorization principal. The endpoint receives one opaque scope that remains stable for that registration and separates equal route paths registered by other applications. Use an explicit resolver when scope continuity must survive process restarts or when storage is shared beyond one registration. A single-user application can return its own stable application scope from authorization_scope_resolver. Never enable this mode for callers with distinct - authorization. + authorization. A shared workflow keeps one conversation state for every thread of that principal; use + ``workflow_factory`` when threads must stay independent. """ _validate_keepalive_seconds(keepalive_seconds) _validate_detached_run_options(max_detached_runs, detached_run_timeout_seconds) @@ -217,6 +241,8 @@ def add_agent_framework_fastapi_endpoint( else: raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.") + _reject_workflow_agent_tools(protocol_runner) + _validate_authorization_scope_configuration( authorization_scope_resolver=authorization_scope_resolver, snapshot_scope_resolver=snapshot_scope_resolver, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py new file mode 100644 index 00000000000..9eb36b09af6 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py @@ -0,0 +1,232 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Authorization-scope ownership of the agent sessions that a workflow holds. + +An ``AgentExecutor`` keeps one ``AgentSession`` between runs. Conversation history, pending approval requests, and +standing approval rules live in that session, and the session is not tied to a caller. An executor, or a session, +that serves two authorization scopes would therefore expose the conversation and approval state of one scope to the +other. AG-UI binds each executor, and the session it holds, to the first authorization scope that runs it and +refuses every other scope. A run without a scope is an owner of its own, so a workflow that ran locally is not +handed to an authorization scope later, and the reverse. + +The binding follows live objects. It does not see state that an application loads from storage into a new object. + +AG-UI also refuses an agent that exposes a ``WorkflowAgent`` as a tool (``WorkflowAgent.as_tool()``), which is not +supported in any form. The tool reaches the executors of one shared workflow, and no scope owns them. +``workflow_agent_tools`` finds the tools that AG-UI can see. The others stay unsupported. +""" + +from __future__ import annotations + +import threading +import weakref +from collections.abc import Iterator, Mapping, Sequence +from typing import Any, Final + +from agent_framework import AgentExecutor, AgentSession, Workflow, WorkflowAgent, WorkflowExecutor + +SCOPE_CONFLICT_CODE = "WORKFLOW_SCOPE_CONFLICT" +SCOPE_CONFLICT_MESSAGE = "This workflow is bound to a different authorization scope." +SCOPE_CONFLICT_LOG_MESSAGE = ( + "Workflow %r holds an AgentExecutor or AgentSession that already serves a different authorization scope, " + "or runs without a scope. An AgentSession keeps the conversation and approval state of one scope, so serving " + "another scope would expose that state. Pass AgentFrameworkWorkflow(workflow_factory=...) with a factory that " + "builds a new Workflow, including new AgentExecutor and AgentSession instances, for every call, or register " + "the endpoint with single_user=True when every caller is the same principal." +) +WORKFLOW_AGENT_TOOL_CODE = "WORKFLOW_AGENT_TOOL_UNSUPPORTED" +WORKFLOW_AGENT_TOOL_MESSAGE = "This agent exposes a workflow as a tool, which AG-UI does not support." + +# The owner of runs that carry no authorization scope. It never equals a scope string. +_UNSCOPED: Final = object() + + +class IdentityOwners: + """Owners of objects, keyed by object identity. + + The registry never calls ``__hash__`` or ``__eq__`` of the objects it tracks. An unhashable object works, and two + distinct objects that compare equal stay apart. It holds weak references only, and an entry disappears when its + object is collected. Callers serialize ``bind`` and ``owner_of`` themselves. The cleanup callback takes no lock, + because it can run at any time, including while a caller holds one. + """ + + def __init__(self) -> None: + self._entries: dict[int, tuple[weakref.ref[Any], object]] = {} + + def __len__(self) -> int: + return len(self._entries) + + def owner_of(self, subject: object, default: object) -> object: + """Return the owner of ``subject``, or ``default`` when it has none.""" + entry = self._entries.get(id(subject)) + # A dead or different referent means the identity was reused after the original object was collected. + return entry[1] if entry is not None and entry[0]() is subject else default + + def bind(self, subject: object, owner: object) -> None: + """Bind ``subject`` to ``owner`` and replace any previous owner.""" + key = id(subject) + entries = self._entries + entry = entries.get(key) + if entry is not None and entry[0]() is subject: + entries[key] = (entry[0], owner) + return + + def forget(reference: weakref.ref[Any]) -> None: + current = entries.get(key) + if current is not None and current[0] is reference: + entries.pop(key, None) + + entries[key] = (weakref.ref(subject, forget), owner) + + +_OWNERS = IdentityOwners() +_CLAIM_LOCK = threading.Lock() + + +def _unwrap_a2ui_runners(agent: Any) -> Any: + """Return the agent inside the A2UI runners that wrap ``agent``.""" + try: + from ._a2ui import A2UIAgent + except ImportError: + return agent + while isinstance(agent, A2UIAgent): + agent = agent.inner_agent + return agent + + +def _workflow_behind(agent: Any) -> Workflow | None: + """Return the workflow that ``agent`` runs, looking through A2UI runners, or ``None``.""" + unwrapped = _unwrap_a2ui_runners(agent) + return unwrapped.workflow if isinstance(unwrapped, WorkflowAgent) else None + + +def _agent_executors(workflow: object, visited: set[int] | None = None) -> Iterator[AgentExecutor]: + """Yield the ``AgentExecutor`` instances of ``workflow`` and of every workflow nested in it. + + Nested workflows are sub-workflows of ``WorkflowExecutor`` and the workflow behind a ``WorkflowAgent`` that an + ``AgentExecutor`` wraps. An object that is not a ``Workflow`` has no executors to inspect. + """ + if not isinstance(workflow, Workflow): + return + visited = set() if visited is None else visited + if id(workflow) in visited: + return + visited.add(id(workflow)) + for executor in workflow.get_executors_list(): + nested: Workflow | None = None + if isinstance(executor, AgentExecutor): + yield executor + nested = _workflow_behind(executor.agent) + elif isinstance(executor, WorkflowExecutor): + nested = executor.workflow + if nested is not None: + yield from _agent_executors(nested, visited) + + +def _state_holders(workflow: Workflow) -> Iterator[AgentExecutor | AgentSession]: + """Yield what keeps conversation and approval state: each ``AgentExecutor`` and the session it holds. + + Executors can share a session, for example to share one service-side conversation inside one workflow. + """ + for executor in _agent_executors(workflow): + yield executor + session = getattr(executor, "_session", None) + if isinstance(session, AgentSession): + yield session + + +def claim_agent_executors(workflow: Workflow, scope: str | None) -> bool: + """Bind every ``AgentExecutor`` of ``workflow``, and the session each one holds, to ``scope``. + + Args: + workflow: The workflow that is about to run, including the workflows nested in it. + scope: The authorization scope of the run. ``None`` stands for runs without a scope, which own what they + run just as a scope does. + + Returns: + ``False`` when an executor or session already belongs to a different owner. Nothing is bound in that case. + """ + owner: object = _UNSCOPED if scope is None else scope + holders = list(_state_holders(workflow)) + with _CLAIM_LOCK: + if any(_OWNERS.owner_of(holder, owner) != owner for holder in holders): + return False + for holder in holders: + _OWNERS.bind(holder, owner) + return True + + +def claim_workflow_agent(agent: Any, scope: str | None) -> bool: + """Bind the executors behind ``agent`` to ``scope`` when ``agent`` runs a workflow. + + Looks through A2UI runners. An agent that does not run a workflow keeps its state in sessions that AG-UI + already scopes itself, so it needs no claim. + """ + workflow = _workflow_behind(agent) + return workflow is None or claim_agent_executors(workflow, scope) + + +def _tool_delegates(agent: Any) -> Iterator[tuple[str, Any]]: + """Yield the name of each tool of ``agent`` that ``as_tool()`` made, with the agent that tool delegates to.""" + options = getattr(agent, "default_options", None) + tools = options.get("tools") if isinstance(options, Mapping) else None + if not isinstance(tools, (list, tuple)): + return + for tool in tools: + delegate = getattr(tool, "_wrapped_agent", None) + if delegate is not None: + yield str(getattr(tool, "name", "")), delegate + + +def workflow_agent_tools(subject: object) -> list[str]: + """Return the names of the tools that wrap a ``WorkflowAgent`` and that ``subject`` can reach. + + The search follows the tools that ``as_tool()`` made into the agents they delegate to, and the agents of every + ``AgentExecutor`` of a workflow, and it looks through A2UI runners. It stops at a workflow tool, because AG-UI + refuses that tool whatever the workflow holds. It reads the tools that an agent lists in its default options. It + does not see a tool that a context provider adds while a run is prepared, or a tool function that calls a workflow + from custom code. + + Args: + subject: An agent or a workflow. Any other object has no tools to search. + + Returns: + The tool names, without duplicates, in the order that the search found them. + """ + found: list[str] = [] + seen: set[int] = set() + + def visit_workflow(workflow: object) -> None: + for executor in _agent_executors(workflow): + visit_agent(executor.agent) + + def visit_agent(agent: Any) -> None: + agent = _unwrap_a2ui_runners(agent) + if id(agent) in seen: + return + seen.add(id(agent)) + workflow = _workflow_behind(agent) + if workflow is not None: + visit_workflow(workflow) + for tool_name, delegate in _tool_delegates(agent): + if _workflow_behind(delegate) is not None: + found.append(tool_name) + else: + visit_agent(delegate) + + if isinstance(subject, Workflow): + visit_workflow(subject) + else: + visit_agent(subject) + return list(dict.fromkeys(found)) + + +def workflow_agent_tool_error(runner_name: object, tool_names: Sequence[str]) -> str: + """Explain why AG-UI refuses ``runner_name``, which exposes a ``WorkflowAgent`` through the tools ``tool_names``.""" + tools = ", ".join(repr(name) for name in tool_names) + return ( + "AG-UI does not support an agent that exposes a WorkflowAgent as a tool (WorkflowAgent.as_tool()), because " + f"every caller of {runner_name!r} would share the state of that workflow. Remove the tool: {tools}. " + "Serve the workflow with AgentFrameworkWorkflow(workflow_factory=...), or build the workflow inside the tool " + "function for every call." + ) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index 4b56f672098..5ddedb7ca96 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -49,6 +49,16 @@ _snapshot_messages_match, _snapshot_overlap_length, ) +from ._session_ownership import ( + SCOPE_CONFLICT_CODE, + SCOPE_CONFLICT_LOG_MESSAGE, + SCOPE_CONFLICT_MESSAGE, + WORKFLOW_AGENT_TOOL_CODE, + WORKFLOW_AGENT_TOOL_MESSAGE, + claim_agent_executors, + workflow_agent_tool_error, + workflow_agent_tools, +) from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, @@ -671,8 +681,15 @@ def __init__( """Initialize the AG-UI workflow wrapper. Args: - workflow: Optional workflow instance to expose. - workflow_factory: Optional factory for thread-scoped workflow instances. + workflow: Optional workflow instance to expose. One instance serves every request, so its executors keep + their state across callers. When requests carry an authorization scope, a workflow that holds + ``AgentExecutor`` instances serves only the first scope that runs it; a run in another scope ends + with a ``WORKFLOW_SCOPE_CONFLICT`` ``RUN_ERROR``. + workflow_factory: Optional factory for thread-scoped workflow instances. Build a new workflow, including + new ``AgentExecutor`` and ``AgentSession`` instances, on every call. Each authorization scope and + thread then keeps its own executor sessions. The factory receives only the thread id and runs once + for each scope and thread. Start every session with fresh state, and do not restore session state + that is keyed by the thread id alone. Use ``checkpoint_storage`` to continue a conversation. name: Optional workflow name. description: Optional workflow description. snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless @@ -880,6 +897,19 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: yield event return + # A replay above starts no workflow, so only a run that reaches the executors decides who owns them. + unsupported_tools = workflow_agent_tools(workflow) + if unsupported_tools: + logger.error(workflow_agent_tool_error(self.name, unsupported_tools)) + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent(message=WORKFLOW_AGENT_TOOL_MESSAGE, code=WORKFLOW_AGENT_TOOL_CODE) + return + if not claim_agent_executors(workflow, snapshot_scope): + logger.error(SCOPE_CONFLICT_LOG_MESSAGE, self.name) + yield RunStartedEvent(run_id=run_id, thread_id=thread_id) + yield RunErrorEvent(message=SCOPE_CONFLICT_MESSAGE, code=SCOPE_CONFLICT_CODE) + return + # Seed follow-up turns so the workflow runs with the full persisted thread # history instead of just the latest request messages. if stored_snapshot is not None and resume_payload is None: diff --git a/python/packages/ag-ui/tests/ag_ui/test_identity_owners.py b/python/packages/ag-ui/tests/ag_ui/test_identity_owners.py new file mode 100644 index 00000000000..cfab4046710 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_identity_owners.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests of the identity-keyed owner registry that AG-UI uses to bind executors and sessions to a scope.""" + +from __future__ import annotations + +import gc +import weakref + +from agent_framework import AgentSession + +from agent_framework_ag_ui._session_ownership import IdentityOwners + + +def test_identity_owners_find_an_object_by_identity() -> None: + owners = IdentityOwners() + first, second = AgentSession(), AgentSession() + + owners.bind(first, "scope-a") + + assert owners.owner_of(first, None) == "scope-a" + assert owners.owner_of(second, None) is None + + +def test_identity_owners_forget_an_object_when_it_is_collected() -> None: + owners = IdentityOwners() + subject = AgentSession() + owners.bind(subject, "scope-a") + assert len(owners) == 1 + + del subject + gc.collect() + + assert len(owners) == 0 + + +def test_identity_owners_do_not_keep_an_object_alive() -> None: + owners = IdentityOwners() + subject = AgentSession() + reference = weakref.ref(subject) + owners.bind(subject, "scope-a") + + del subject + gc.collect() + + assert reference() is None + + +def test_identity_owners_rebinding_keeps_one_entry_with_the_latest_owner() -> None: + owners = IdentityOwners() + subject = AgentSession() + + owners.bind(subject, "scope-a") + owners.bind(subject, "scope-b") + + assert len(owners) == 1 + assert owners.owner_of(subject, None) == "scope-b" + + +def test_identity_owners_ignore_an_entry_that_belongs_to_another_object() -> None: + owners = IdentityOwners() + first, second = AgentSession(), AgentSession() + owners.bind(first, "scope-a") + # Move the entry of `first` under the key of `second`, as if `second` had reused the identity of a dead object. + owners._entries[id(second)] = owners._entries.pop(id(first)) # pyright: ignore[reportPrivateUsage] + + assert owners.owner_of(second, None) is None + + +def test_identity_owners_replace_an_entry_that_belongs_to_another_object() -> None: + owners = IdentityOwners() + first, second = AgentSession(), AgentSession() + owners.bind(first, "scope-a") + owners._entries[id(second)] = owners._entries.pop(id(first)) # pyright: ignore[reportPrivateUsage] + + owners.bind(second, "scope-b") + + assert len(owners) == 1 + assert owners.owner_of(second, None) == "scope-b" + + +def test_identity_owners_obsolete_cleanup_keeps_the_entry_of_a_newer_object() -> None: + owners = IdentityOwners() + old, new = AgentSession(), AgentSession() + owners.bind(old, "scope-a") + # Keep the cleanup of `old` alive so that it runs when `old` dies, then hand its key to `new`. + old_reference = owners._entries[id(old)][0] # pyright: ignore[reportPrivateUsage] + owners._entries[id(old)] = (weakref.ref(new), "scope-b") # pyright: ignore[reportPrivateUsage] + + del old + gc.collect() + + assert old_reference() is None + assert len(owners) == 1 diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_detection.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_detection.py new file mode 100644 index 00000000000..aa496128b5b --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_detection.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests of the search for tools that wrap a WorkflowAgent, which AG-UI refuses to serve.""" + +from __future__ import annotations + +import pytest +from workflow_scope_helpers import ( # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + WORKFLOW_TOOL_NAME, + RecordingChatClient, + build_agent, + build_agent_with_tools, + build_runner_with_a_workflow_tool, + build_workflow, + build_workflow_tool, +) + +from agent_framework_ag_ui._session_ownership import workflow_agent_tools + +# The forms that are an agent or a workflow themselves. The two wrapper forms are unwrapped by the endpoint. +SEARCHABLE_FORMS = ["agent", "sub_agent", "executor_agent", "workflow_agent_runner", "nested_workflow"] + + +@pytest.mark.parametrize("form", SEARCHABLE_FORMS) +def test_workflow_agent_tools_names_the_tool_in_every_form(form: str) -> None: + assert workflow_agent_tools(build_runner_with_a_workflow_tool(form)) == [WORKFLOW_TOOL_NAME] + + +def test_workflow_agent_tools_ignores_the_tools_of_ordinary_agents() -> None: + outer = build_agent_with_tools(build_agent(RecordingChatClient()).as_tool(name="sub_agent_tool")) + + assert workflow_agent_tools(outer) == [] + + +def test_workflow_agent_tools_finds_nothing_in_a_workflow_without_tools() -> None: + assert workflow_agent_tools(build_workflow(build_agent(RecordingChatClient()))) == [] + + +def test_workflow_agent_tools_finds_nothing_in_an_object_that_is_neither_an_agent_nor_a_workflow() -> None: + assert workflow_agent_tools(object()) == [] + + +def test_workflow_agent_tools_terminates_when_agents_delegate_to_each_other() -> None: + first = build_agent_with_tools(name="first") + second = build_agent_with_tools(first.as_tool(name="first_tool"), name="second") + first.default_options["tools"] = [second.as_tool(name="second_tool")] + + assert workflow_agent_tools(first) == [] + + +def test_workflow_agent_tools_reports_a_workflow_tool_that_sits_behind_a_cycle() -> None: + first = build_agent_with_tools(name="first") + second = build_agent_with_tools(first.as_tool(name="first_tool"), build_workflow_tool(), name="second") + first.default_options["tools"] = [second.as_tool(name="second_tool")] + + assert workflow_agent_tools(first) == [WORKFLOW_TOOL_NAME] + + +def test_workflow_agent_tools_looks_through_a2ui_runners() -> None: + pytest.importorskip("ag_ui_a2ui_toolkit") + from agent_framework_ag_ui._a2ui import enable_a2ui + + outer = build_agent_with_tools(build_workflow_tool()) + + assert workflow_agent_tools(enable_a2ui(outer, object())) == [WORKFLOW_TOOL_NAME] diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_rejection.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_rejection.py new file mode 100644 index 00000000000..b0c34083983 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_rejection.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""AG-UI endpoint tests: an agent that exposes a WorkflowAgent as a tool is rejected, not served.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from fastapi.testclient import TestClient +from workflow_scope_helpers import ( # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + OUTER_REPLY_TEXT, + RUNNER_FORMS_WITH_A_WORKFLOW_TOOL, + WORKFLOW_TOOL_NAME, + RecordingChatClient, + assistant_text, + build_agent, + build_agent_with_tools, + build_app, + build_runner_with_a_workflow_tool, + build_workflow_tool, + build_workflow_whose_agent_has_tools, + event_types, + post_run, + run_error_codes, +) + +from agent_framework_ag_ui import AgentFrameworkWorkflow + +REJECTED_CODE = "WORKFLOW_AGENT_TOOL_UNSUPPORTED" +SCOPE_MODES = ["authorization", "snapshot", "single_user"] + + +def _assert_rejected(events: list[dict[str, Any]]) -> None: + assert event_types(events) == ["RUN_STARTED", "RUN_ERROR"] + assert run_error_codes(events) == [REJECTED_CODE] + + +@pytest.mark.parametrize("form", RUNNER_FORMS_WITH_A_WORKFLOW_TOOL) +def test_endpoint_rejects_a_runner_that_exposes_a_workflow_agent_as_a_tool(form: str) -> None: + with pytest.raises(ValueError, match=r"WorkflowAgent\.as_tool\(\)") as rejection: + build_app(build_runner_with_a_workflow_tool(form), "single_user") + + assert WORKFLOW_TOOL_NAME in str(rejection.value) + + +def test_endpoint_registers_a_workflow_factory_without_calling_it() -> None: + calls: list[str] = [] + + def factory(thread_id: str) -> Any: + calls.append(thread_id) + return build_workflow_whose_agent_has_tools(build_workflow_tool()) + + build_app(AgentFrameworkWorkflow(workflow_factory=factory), "single_user") + + # A factory builds its workflow per request, so only a run can find out what it exposes. + assert calls == [] + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_workflow_factory_that_builds_a_workflow_agent_tool_ends_the_run(scope_mode: Any) -> None: + inner_client = RecordingChatClient() + runner = AgentFrameworkWorkflow( + workflow_factory=lambda _thread_id: build_workflow_whose_agent_has_tools(build_workflow_tool(inner_client)) + ) + + with TestClient(build_app(runner, scope_mode)) as http: + events = post_run(http, "alice", "thread-a", "ALICE_ONLY") + + _assert_rejected(events) + assert inner_client.calls == [] + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_agent_that_gains_a_workflow_agent_tool_after_registration_ends_the_run(scope_mode: Any) -> None: + inner_client = RecordingChatClient() + agent = build_agent_with_tools() + app = build_app(agent, scope_mode) + agent.default_options["tools"] = [build_workflow_tool(inner_client)] + + with TestClient(app) as http: + events = post_run(http, "alice", "thread-a", "ALICE_ONLY") + + _assert_rejected(events) + assert inner_client.calls == [] + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_endpoint_still_serves_an_agent_that_exposes_an_ordinary_agent_as_a_tool(scope_mode: Any) -> None: + inner_client = RecordingChatClient() + sub_agent_tool = build_agent(inner_client).as_tool(name="sub_agent_tool") + outer = build_agent_with_tools(sub_agent_tool, tool_name="sub_agent_tool") + + with TestClient(build_app(outer, scope_mode)) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + for events in (alice, bob): + assert run_error_codes(events) == [] + assert OUTER_REPLY_TEXT in assistant_text(events) + # Each call reached the sub-agent, which starts from a new session every time. + assert len(inner_client.calls) == 2 + assert not inner_client.saw("ALICE_ONLY", call=1) diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_scope_ownership.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_scope_ownership.py new file mode 100644 index 00000000000..7554e867637 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_scope_ownership.py @@ -0,0 +1,484 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""AG-UI endpoint tests: workflow-owned agent sessions stay bound to one authorization scope.""" + +from __future__ import annotations + +import asyncio +import gc +import weakref +from concurrent.futures import ThreadPoolExecutor +from typing import Any, cast + +import pytest +from agent_framework import ( + AgentExecutor, + SupportsAgentRun, + Workflow, + WorkflowBuilder, + WorkflowContext, + WorkflowExecutor, + executor, +) +from fastapi.testclient import TestClient +from workflow_scope_helpers import ( # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + INVOKE_PREFIX, + REPLY_TEXT, + EqualSession, + RecordingChatClient, + UnhashableSession, + always_approve_resume, + assistant_text, + build_agent, + build_app, + build_unhashable_executor_workflow, + build_workflow, + event_types, + interrupts, + post_run, + run_error_codes, +) + +from agent_framework_ag_ui import AgentFrameworkAgent, AgentFrameworkWorkflow +from agent_framework_ag_ui._session_ownership import claim_agent_executors, claim_workflow_agent + +SCOPE_MODES = ["authorization", "snapshot"] +CONFLICT_CODE = "WORKFLOW_SCOPE_CONFLICT" + + +def _assert_completed(events: list[dict[str, Any]]) -> None: + """A run reached the model, produced its reply, and finished without an error.""" + assert run_error_codes(events) == [] + assert REPLY_TEXT in assistant_text(events) + assert "RUN_FINISHED" in event_types(events) + + +def _assert_scope_conflict(events: list[dict[str, Any]]) -> None: + assert event_types(events) == ["RUN_STARTED", "RUN_ERROR"] + assert run_error_codes(events) == [CONFLICT_CODE] + + +# region shared instances serve one scope + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +@pytest.mark.parametrize("persist_snapshots", [False, True]) +@pytest.mark.parametrize("same_thread_id", [False, True]) +def test_shared_workflow_serves_only_the_first_authorization_scope( + scope_mode: Any, persist_snapshots: bool, same_thread_id: bool +) -> None: + client = RecordingChatClient() + app = build_app(build_workflow(build_agent(client)), scope_mode, persist_snapshots=persist_snapshots) + + with TestClient(app) as http: + alice_first = post_run(http, "alice", "thread-a", "ALICE_FIRST") + bob = post_run(http, "bob", "thread-a" if same_thread_id else "thread-b", "BOB_ONLY") + alice_second = post_run(http, "alice", "thread-a", "ALICE_SECOND") + + _assert_completed(alice_first) + _assert_completed(alice_second) + _assert_scope_conflict(bob) + # Only Alice's two runs reached the model, so Bob's request and Alice's history never met. + assert len(client.calls) == 2 + assert not any("BOB_ONLY" in text for call in client.calls for _, text in call) + assert client.saw("ALICE_FIRST", call=1) + assert client.sessions[0] is client.sessions[1] + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +@pytest.mark.parametrize("share", ["one_workflow", "one_builder"]) +def test_workflow_factory_result_shared_between_scopes_is_refused(scope_mode: Any, share: str) -> None: + client = RecordingChatClient() + agent = build_agent(client) + shared_workflow = build_workflow(agent) + shared_builder = WorkflowBuilder(start_executor=AgentExecutor(agent, id="worker")) + runner = AgentFrameworkWorkflow( + workflow_factory=lambda _thread_id: shared_workflow if share == "one_workflow" else shared_builder.build() + ) + + with TestClient(build_app(runner, scope_mode)) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + _assert_completed(alice) + _assert_scope_conflict(bob) + assert len(client.calls) == 1 + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_factory_that_wraps_a_shared_workflow_agent_is_refused(scope_mode: Any) -> None: + client = RecordingChatClient() + shared_workflow_agent = build_workflow(build_agent(client)).as_agent() + + def factory(_thread_id: str) -> Workflow: + outer = AgentExecutor(shared_workflow_agent, id="outer") + return WorkflowBuilder(start_executor=outer, output_from=[outer]).build() + + runner = AgentFrameworkWorkflow(workflow_factory=factory) + + with TestClient(build_app(runner, scope_mode)) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + _assert_completed(alice) + _assert_scope_conflict(bob) + assert len(client.calls) == 1 + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +@pytest.mark.parametrize("wrap", ["workflow_agent", "wrapped_workflow_agent"]) +def test_workflow_agent_serves_only_the_first_authorization_scope(scope_mode: Any, wrap: str) -> None: + client = RecordingChatClient() + workflow_agent = build_workflow(build_agent(client)).as_agent() + runner = workflow_agent if wrap == "workflow_agent" else AgentFrameworkAgent(agent=workflow_agent) + + with TestClient(build_app(runner, scope_mode)) as http: + alice_first = post_run(http, "alice", "thread-a", "ALICE_FIRST") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + alice_second = post_run(http, "alice", "thread-a", "ALICE_SECOND") + + _assert_completed(alice_first) + _assert_completed(alice_second) + _assert_scope_conflict(bob) + assert len(client.calls) == 2 + assert not any("BOB_ONLY" in text for call in client.calls for _, text in call) + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_a2ui_wrapped_workflow_agent_serves_only_the_first_authorization_scope(scope_mode: Any) -> None: + pytest.importorskip("ag_ui_a2ui_toolkit") + from agent_framework_ag_ui._a2ui import enable_a2ui + + client = RecordingChatClient() + workflow_agent = build_workflow(build_agent(client)).as_agent() + runner = AgentFrameworkAgent(agent=cast(SupportsAgentRun, enable_a2ui(workflow_agent, object()))) + + with TestClient(build_app(runner, scope_mode)) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + _assert_completed(alice) + _assert_scope_conflict(bob) + assert len(client.calls) == 1 + assert not any("BOB_ONLY" in text for call in client.calls for _, text in call) + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_factory_that_shares_one_agent_session_between_scopes_is_refused(scope_mode: Any) -> None: + client = RecordingChatClient() + agent = build_agent(client) + shared_session = agent.create_session() + runner = AgentFrameworkWorkflow(workflow_factory=lambda _thread_id: build_workflow(agent, session=shared_session)) + + with TestClient(build_app(runner, scope_mode)) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + _assert_completed(alice) + _assert_scope_conflict(bob) + assert len(client.calls) == 1 + assert not any("ALICE_ONLY" in text for call in client.calls[1:] for _, text in call) + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_shared_workflow_with_an_unhashable_session_serves_only_the_first_authorization_scope(scope_mode: Any) -> None: + client = RecordingChatClient() + app = build_app(build_workflow(build_agent(client), session=UnhashableSession()), scope_mode) + + with TestClient(app) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + # Ownership never hashes a session, so a session class that defines __eq__ alone still runs. + _assert_completed(alice) + _assert_scope_conflict(bob) + assert len(client.calls) == 1 + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_shared_workflow_with_an_unhashable_executor_serves_only_the_first_authorization_scope(scope_mode: Any) -> None: + client = RecordingChatClient() + app = build_app(build_unhashable_executor_workflow(build_agent(client)), scope_mode) + + with TestClient(app) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + _assert_completed(alice) + _assert_scope_conflict(bob) + assert len(client.calls) == 1 + + +# endregion + +# region requests that start no workflow, and runs without a scope + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +@pytest.mark.parametrize("runner_kind", ["workflow", "workflow_agent"]) +def test_snapshot_replay_does_not_bind_the_executors(scope_mode: Any, runner_kind: str) -> None: + client = RecordingChatClient() + workflow = build_workflow(build_agent(client)) + runner = workflow if runner_kind == "workflow" else workflow.as_agent() + + with TestClient(build_app(runner, scope_mode, persist_snapshots=True)) as http: + # A read-only replay starts no workflow and reaches no model, so it does not decide who owns the executors. + replay = post_run(http, "alice", "thread-a") + bob = post_run(http, "bob", "thread-b", "BOB_FIRST") + alice = post_run(http, "alice", "thread-a", "ALICE_LATER") + + assert event_types(replay) == ["RUN_STARTED", "RUN_FINISHED"] + _assert_completed(bob) + _assert_scope_conflict(alice) + assert len(client.calls) == 1 + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_workflow_run_without_a_scope_blocks_a_later_scoped_endpoint(scope_mode: Any) -> None: + client = RecordingChatClient() + runner = AgentFrameworkWorkflow(workflow=build_workflow(build_agent(client))) + + async def run_without_scope() -> list[Any]: + request = { + "thread_id": "thread-local", + "run_id": "run-local", + "messages": [{"id": "message-local", "role": "user", "content": "LOCAL_ONLY"}], + } + return [event async for event in runner.run(request)] + + local_events = asyncio.run(run_without_scope()) + + assert any(type(event).__name__ == "RunFinishedEvent" for event in local_events) + with TestClient(build_app(runner, scope_mode)) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + + _assert_scope_conflict(alice) + assert len(client.calls) == 1 + + +# endregion + +# region intentional persistence and per-scope construction keep working + + +@pytest.mark.parametrize("scope_mode", ["single_user", "authorization"]) +def test_shared_workflow_keeps_one_state_for_every_thread_of_its_scope(scope_mode: Any) -> None: + client = RecordingChatClient() + app = build_app(build_workflow(build_agent(client)), scope_mode) + + with TestClient(app) as http: + first = post_run(http, "alice", "thread-a", "FIRST_THREAD") + second = post_run(http, "alice", "thread-b", "SECOND_THREAD") + + _assert_completed(first) + _assert_completed(second) + assert client.saw("FIRST_THREAD", call=1) + assert client.sessions[0] is client.sessions[1] + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +@pytest.mark.parametrize("persist_snapshots", [False, True]) +@pytest.mark.parametrize("same_thread_id", [False, True]) +def test_workflow_factory_isolates_scopes_and_keeps_each_scope_history( + scope_mode: Any, persist_snapshots: bool, same_thread_id: bool +) -> None: + client = RecordingChatClient() + agent = build_agent(client) + runner = AgentFrameworkWorkflow(workflow_factory=lambda _thread_id: build_workflow(agent)) + app = build_app(runner, scope_mode, persist_snapshots=persist_snapshots) + + with TestClient(app) as http: + alice_first = post_run(http, "alice", "thread-a", "ALICE_FIRST") + bob_first = post_run(http, "bob", "thread-a" if same_thread_id else "thread-b", "BOB_FIRST") + alice_second = post_run(http, "alice", "thread-a", "ALICE_SECOND") + + for events in (alice_first, bob_first, alice_second): + _assert_completed(events) + assert len(client.calls) == 3 + # Bob starts from his own state; Alice continues from hers and never sees Bob's. + assert not client.saw("ALICE_FIRST", call=1) + assert client.saw("ALICE_FIRST", call=2) + assert not client.saw("BOB_FIRST", call=2) + assert client.sessions[0] is not client.sessions[1] + assert client.sessions[0] is client.sessions[2] + + +@pytest.mark.parametrize("scope_mode", SCOPE_MODES) +def test_shared_workflow_without_agent_executor_is_not_bound_to_a_scope(scope_mode: Any) -> None: + @executor(id="responder") + async def responder(message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + await ctx.yield_output("Workflow response") + + app = build_app(WorkflowBuilder(start_executor=responder).build(), scope_mode) + + with TestClient(app) as http: + alice = post_run(http, "alice", "thread-a", "ALICE_ONLY") + bob = post_run(http, "bob", "thread-b", "BOB_ONLY") + + for events in (alice, bob): + assert run_error_codes(events) == [] + assert "Workflow response" in assistant_text(events) + + +# endregion + +# region approvals follow the same ownership + + +@pytest.mark.parametrize("isolation", ["factory", "shared"]) +def test_standing_approval_does_not_reach_another_scope(isolation: str) -> None: + client = RecordingChatClient() + agent = build_agent(client, standing_approvals=True) + runner: Any = ( + AgentFrameworkWorkflow(workflow_factory=lambda _thread_id: build_workflow(agent)) + if isolation == "factory" + else build_workflow(agent) + ) + + with TestClient(build_app(runner, "authorization")) as http: + first = post_run(http, "alice", "thread-a", f"{INVOKE_PREFIX}ALICE_FIRST") + (interrupt,) = interrupts(first) + approved = post_run(http, "alice", "thread-a", resume=always_approve_resume(interrupt)) + # Alice's standing rule skips her next approval prompt. + repeated = post_run(http, "alice", "thread-a", f"{INVOKE_PREFIX}ALICE_SECOND") + bob = post_run(http, "bob", "thread-b", f"{INVOKE_PREFIX}BOB_ONLY") + + assert run_error_codes(approved) == [] + assert interrupts(repeated) == [] + assert client.executed == ["ALICE_FIRST", "ALICE_SECOND"] + if isolation == "factory": + assert len(interrupts(bob)) == 1 + else: + _assert_scope_conflict(bob) + + +# endregion + +# region claim_agent_executors + + +def test_claim_binds_agent_executors_to_the_first_scope() -> None: + workflow = build_workflow(build_agent(RecordingChatClient())) + + assert claim_agent_executors(workflow, "scope-a") + assert claim_agent_executors(workflow, "scope-a") + assert not claim_agent_executors(workflow, "scope-b") + + +def test_claim_binds_unscoped_runs_to_their_own_owner() -> None: + workflow = build_workflow(build_agent(RecordingChatClient())) + + assert claim_agent_executors(workflow, None) + assert claim_agent_executors(workflow, None) + assert not claim_agent_executors(workflow, "scope-a") + + +def test_claim_refuses_an_unscoped_run_on_scoped_executors() -> None: + workflow = build_workflow(build_agent(RecordingChatClient())) + + assert claim_agent_executors(workflow, "scope-a") + assert not claim_agent_executors(workflow, None) + + +def test_claim_binds_the_session_of_every_agent_executor() -> None: + agent = build_agent(RecordingChatClient()) + shared_session = agent.create_session() + first = build_workflow(agent, session=shared_session) + second = build_workflow(agent, session=shared_session) + + assert claim_agent_executors(first, "scope-a") + # New executors that inject the session of a bound executor are refused for another scope. + assert not claim_agent_executors(second, "scope-b") + assert claim_agent_executors(second, "scope-a") + + +def test_claim_allows_executors_of_one_workflow_to_share_a_session() -> None: + agent = build_agent(RecordingChatClient()) + shared_session = agent.create_session() + writer = AgentExecutor(agent, id="writer", session=shared_session) + reviewer = AgentExecutor(agent, id="reviewer", session=shared_session) + workflow = WorkflowBuilder(start_executor=writer, output_from=[reviewer]).add_edge(writer, reviewer).build() + + assert claim_agent_executors(workflow, "scope-a") + assert claim_agent_executors(workflow, "scope-a") + assert not claim_agent_executors(workflow, "scope-b") + + +def test_claim_binds_an_agent_executor_subclass_that_is_unhashable() -> None: + workflow = build_unhashable_executor_workflow(build_agent(RecordingChatClient())) + + assert claim_agent_executors(workflow, "scope-a") + assert not claim_agent_executors(workflow, "scope-b") + + +def test_claim_keeps_distinct_sessions_that_compare_equal_apart() -> None: + agent = build_agent(RecordingChatClient()) + first = build_workflow(agent, session=EqualSession()) + second = build_workflow(agent, session=EqualSession()) + + assert claim_agent_executors(first, "scope-a") + # Equality and hash of a session never decide who owns it, so the second session is not the first one. + assert claim_agent_executors(second, "scope-b") + assert not claim_agent_executors(first, "scope-b") + assert not claim_agent_executors(second, "scope-a") + + +def test_claim_covers_agent_executors_inside_nested_workflows() -> None: + inner = build_workflow(build_agent(RecordingChatClient())) + outer = WorkflowBuilder(start_executor=WorkflowExecutor(inner, id="inner")).build() + + assert claim_agent_executors(outer, "scope-a") + assert not claim_agent_executors(inner, "scope-b") + assert not claim_agent_executors(outer, "scope-b") + + +def test_claim_covers_workflows_behind_a_wrapped_workflow_agent() -> None: + inner = build_workflow(build_agent(RecordingChatClient())) + outer = WorkflowBuilder(start_executor=AgentExecutor(inner.as_agent(), id="outer")).build() + + assert claim_agent_executors(outer, "scope-a") + assert not claim_agent_executors(inner, "scope-b") + + +def test_claim_workflow_agent_ignores_agents_that_run_no_workflow() -> None: + agent = build_agent(RecordingChatClient()) + + assert claim_workflow_agent(agent, "scope-a") + assert claim_workflow_agent(agent, "scope-b") + assert claim_workflow_agent(agent, None) + + +def test_claim_workflow_agent_looks_through_a2ui_runners() -> None: + pytest.importorskip("ag_ui_a2ui_toolkit") + from agent_framework_ag_ui._a2ui import enable_a2ui + + workflow = build_workflow(build_agent(RecordingChatClient())) + wrapped = enable_a2ui(enable_a2ui(workflow.as_agent(), object()), object()) + + assert claim_workflow_agent(wrapped, "scope-a") + assert not claim_agent_executors(workflow, "scope-b") + assert not claim_workflow_agent(wrapped, "scope-b") + + +def test_claim_is_atomic_across_threads() -> None: + workflow = build_workflow(build_agent(RecordingChatClient())) + + with ThreadPoolExecutor(max_workers=16) as pool: + results = list(pool.map(lambda index: claim_agent_executors(workflow, f"scope-{index}"), range(64))) + + assert results.count(True) == 1 + + +def test_claim_does_not_keep_executors_alive() -> None: + workflow = build_workflow(build_agent(RecordingChatClient())) + executor_ref = weakref.ref(workflow.get_start_executor()) + assert claim_agent_executors(workflow, "scope-a") + + del workflow + gc.collect() + + assert executor_ref() is None + + +# endregion diff --git a/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py b/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py new file mode 100644 index 00000000000..34c6e36b396 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py @@ -0,0 +1,315 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Helpers for AG-UI tests of workflow-owned agent sessions and authorization scopes.""" + +from __future__ import annotations + +import uuid +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any, Literal + +from agent_framework import ( + Agent, + AgentExecutor, + AgentSession, + ChatResponseUpdate, + Content, + InMemoryHistoryProvider, + Message, + ToolApprovalMiddleware, + Workflow, + WorkflowAgent, + WorkflowBuilder, + WorkflowExecutor, + create_always_approve_tool_response, + tool, +) +from conftest import StreamingChatClientStub # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] +from fastapi import FastAPI, Header, HTTPException +from fastapi.testclient import TestClient +from sse_helpers import parse_sse_response # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] + +from agent_framework_ag_ui import ( + AgentFrameworkAgent, + AgentFrameworkWorkflow, + AGUIRequest, + InMemoryAGUIThreadSnapshotStore, + add_agent_framework_fastapi_endpoint, +) + +ScopeMode = Literal["authorization", "snapshot"] + +REPLY_TEXT = "model-reply" +OUTER_REPLY_TEXT = "outer-reply" +GATED_TOOL_NAME = "gated_action" +WORKFLOW_TOOL_NAME = "worker_tool" +INVOKE_PREFIX = "invoke:" +PRINCIPALS = ("alice", "bob") + + +class RecordingChatClient(StreamingChatClientStub): + """Stub chat client that records what every model call received and which session was behind it. + + A user message that starts with ``invoke:`` makes the model request the approval-gated tool. + """ + + def __init__(self) -> None: + super().__init__(self._respond) + self.calls: list[list[tuple[str, str]]] = [] + self.sessions: list[AgentSession | None] = [] + self.executed: list[str] = [] + + async def _respond( + self, + messages: Sequence[Message], + options: Mapping[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + self.calls.append([(message.role, message.text) for message in messages]) + self.sessions.append(self.last_session) + last = messages[-1] if messages else None + if last is not None and last.role == "user" and last.text.startswith(INVOKE_PREFIX): + contents = [ + Content.from_function_call( + call_id=str(uuid.uuid4()), + id=str(uuid.uuid4()), + name=GATED_TOOL_NAME, + arguments={"value": last.text.removeprefix(INVOKE_PREFIX)}, + ) + ] + else: + contents = [Content.from_text(REPLY_TEXT)] + yield ChatResponseUpdate(role="assistant", contents=contents, message_id=str(uuid.uuid4())) + + def saw(self, marker: str, *, call: int = -1) -> bool: + """Whether the model-bound messages of one call (the last by default) contain ``marker``.""" + return any(marker in text for _, text in self.calls[call]) + + +def build_agent(client: RecordingChatClient, *, standing_approvals: bool = False) -> Agent[Any]: + """Build an agent whose history lives in its session and whose gated tool records its executions.""" + + @tool(name=GATED_TOOL_NAME, approval_mode="always_require") + def gated_action(value: str) -> str: + client.executed.append(value) + return f"completed:{value}" + + return Agent( + client=client, + name="worker", + tools=[gated_action], + context_providers=[InMemoryHistoryProvider(source_id="history")], + middleware=[ToolApprovalMiddleware()] if standing_approvals else [], + ) + + +def build_workflow(agent: Agent[Any], *, session: AgentSession | None = None) -> Workflow: + """Build a new workflow around a new ``AgentExecutor``. + + The executor owns a new ``AgentSession`` unless ``session`` injects an existing one. + """ + executor = AgentExecutor(agent, id="worker", session=session) + return WorkflowBuilder(start_executor=executor, output_from=[executor]).build() + + +class UnhashableSession(AgentSession): + """Session that compares by value. A class that defines ``__eq__`` alone is unhashable.""" + + def __eq__(self, other: object) -> bool: + return isinstance(other, UnhashableSession) and self.session_id == other.session_id + + +class EqualSession(AgentSession): + """Session whose instances all compare equal and share one hash, whatever their state.""" + + def __eq__(self, other: object) -> bool: + return isinstance(other, EqualSession) + + def __hash__(self) -> int: + return 0 + + +class UnhashableExecutor(AgentExecutor): + """Executor that compares by id. A class that defines ``__eq__`` alone is unhashable.""" + + def __eq__(self, other: object) -> bool: + return isinstance(other, UnhashableExecutor) and self.id == other.id + + +def build_unhashable_executor_workflow(agent: Agent[Any]) -> Workflow: + """Build a new workflow around an ``AgentExecutor`` subclass that is unhashable.""" + executor = UnhashableExecutor(agent, id="worker") + return WorkflowBuilder(start_executor=executor, output_from=[executor]).build() + + +class ToolCallingClient(StreamingChatClientStub): + """Stub chat client that calls one tool with the user's text, then answers once the tool result is in.""" + + def __init__(self, tool_name: str = WORKFLOW_TOOL_NAME) -> None: + super().__init__(self._respond) + self.tool_name = tool_name + self.calls = 0 + + async def _respond( + self, + messages: Sequence[Message], + options: Mapping[str, Any], + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del options, kwargs + self.calls += 1 + last = messages[-1] + if last.role == "tool": + contents = [Content.from_text(OUTER_REPLY_TEXT)] + else: + contents = [ + Content.from_function_call( + call_id=str(uuid.uuid4()), + id=str(uuid.uuid4()), + name=self.tool_name, + arguments={"task": last.text}, + ) + ] + yield ChatResponseUpdate(role="assistant", contents=contents, message_id=str(uuid.uuid4())) + + +def build_workflow_agent(client: RecordingChatClient) -> WorkflowAgent: + """Build a ``WorkflowAgent`` around a new ``AgentExecutor``, which keeps its own session.""" + return build_workflow(build_agent(client)).as_agent(name="worker_agent") + + +def build_workflow_tool(client: RecordingChatClient | None = None) -> Any: + """Build the tool that ``WorkflowAgent.as_tool()`` returns for a new workflow.""" + return build_workflow_agent(client or RecordingChatClient()).as_tool(name=WORKFLOW_TOOL_NAME) + + +def build_agent_with_tools(*tools: Any, name: str = "outer", tool_name: str = WORKFLOW_TOOL_NAME) -> Agent[Any]: + """Build an ordinary agent that exposes ``tools``. Its model calls the tool called ``tool_name``, then answers.""" + return Agent(client=ToolCallingClient(tool_name), name=name, tools=list(tools)) + + +def build_workflow_whose_agent_has_tools(*tools: Any) -> Workflow: + """Build a workflow with one ``AgentExecutor`` whose agent exposes ``tools``.""" + executor = AgentExecutor(build_agent_with_tools(*tools, name="holder"), id="holder") + return WorkflowBuilder(start_executor=executor, output_from=[executor]).build() + + +RUNNER_FORMS_WITH_A_WORKFLOW_TOOL = [ + "agent", + "wrapped_agent", + "sub_agent", + "executor_agent", + "workflow_runner", + "workflow_agent_runner", + "nested_workflow", +] + + +def build_runner_with_a_workflow_tool(form: str) -> Any: + """Build something that an endpoint can serve and that exposes a ``WorkflowAgent`` as a tool, in a given form.""" + tool_wrapping_a_workflow = build_workflow_tool() + outer = build_agent_with_tools(tool_wrapping_a_workflow) + if form == "agent": + return outer + if form == "wrapped_agent": + return AgentFrameworkAgent(agent=outer) + if form == "sub_agent": + return build_agent_with_tools(outer.as_tool(name="outer_tool"), name="coordinator") + holder = build_workflow_whose_agent_has_tools(tool_wrapping_a_workflow) + if form == "executor_agent": + return holder + if form == "workflow_runner": + return AgentFrameworkWorkflow(workflow=holder) + if form == "workflow_agent_runner": + return holder.as_agent() + if form == "nested_workflow": + return WorkflowBuilder(start_executor=WorkflowExecutor(holder, id="inner")).build() + raise ValueError(f"Unknown form: {form}") + + +async def _authenticated_scope(authorization: str = Header(default="")) -> str: + principal = authorization.removeprefix("token-") + if principal not in PRINCIPALS: + raise HTTPException(status_code=401) + return f"principal:{principal}" + + +def _forwarded_scope(request: AGUIRequest) -> str: + assert request.forwarded_props is not None + return f"principal:{request.forwarded_props['principal']}" + + +def build_app( + runner: Any, + scope_mode: ScopeMode | Literal["single_user"], + *, + persist_snapshots: bool = False, +) -> FastAPI: + """Register ``runner`` with the authorization mode under test.""" + app = FastAPI() + options: dict[str, Any] = {"keepalive_seconds": None} + if persist_snapshots: + options["snapshot_store"] = InMemoryAGUIThreadSnapshotStore() + if scope_mode == "authorization": + options["authorization_scope_resolver"] = _authenticated_scope + elif scope_mode == "snapshot": + options["snapshot_scope_resolver"] = _forwarded_scope + else: + options["single_user"] = True + add_agent_framework_fastapi_endpoint(app, runner, **options) + return app + + +def post_run( + http: TestClient, + principal: str, + thread_id: str, + text: str | None = None, + *, + resume: list[dict[str, Any]] | None = None, +) -> list[dict[str, Any]]: + """POST one AG-UI run as ``principal`` with a fresh message id and consume the whole SSE stream.""" + body: dict[str, Any] = { + "threadId": thread_id, + "runId": f"run-{uuid.uuid4()}", + "messages": [] if text is None else [{"id": f"message-{uuid.uuid4()}", "role": "user", "content": text}], + "state": {}, + "tools": [], + "context": [], + "forwardedProps": {"principal": principal}, + } + if resume is not None: + body["resume"] = resume + with http.stream("POST", "/", json=body, headers={"Authorization": f"token-{principal}"}) as response: + content = b"".join(response.iter_bytes()) + assert response.status_code == 200 + return parse_sse_response(content) + + +def event_types(events: Sequence[Mapping[str, Any]]) -> list[str]: + return [event["type"] for event in events] + + +def run_error_codes(events: Sequence[Mapping[str, Any]]) -> list[str | None]: + return [event.get("code") for event in events if event["type"] == "RUN_ERROR"] + + +def assistant_text(events: Sequence[Mapping[str, Any]]) -> str: + return "".join(event["delta"] for event in events if event["type"] == "TEXT_MESSAGE_CONTENT") + + +def interrupts(events: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]: + return [ + interrupt + for event in events + if event["type"] == "RUN_FINISHED" + for interrupt in event.get("outcome", {}).get("interrupts", []) + ] + + +def always_approve_resume(interrupt: Mapping[str, Any]) -> list[dict[str, Any]]: + """Answer an approval interrupt with a standing "always approve" decision for the gated tool.""" + request = Content.from_dict(interrupt["metadata"]["agent_framework"]["value"]) + response = create_always_approve_tool_response(request) + return [{"interruptId": interrupt["id"], "status": "resolved", "payload": response.to_dict()}] diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 38ee6a3cb38..fb1e94f3841 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -704,6 +704,9 @@ def as_tool( :attr:`FunctionInvocationContext.parent_service_session_state_keys` when the parent session is non-empty. Automatic function calling supplies it. + ``WorkflowAgent.as_tool()`` is not supported, because every caller of + the tool would share the state of one workflow. + Examples: .. code-block:: python @@ -885,6 +888,8 @@ async def _agent_wrapper(ctx: FunctionInvocationContext, **kwargs: Any) -> str: func=_agent_wrapper, input_model=input_schema, approval_mode=approval_mode, + # Private link for hosts that must know which agent a tool delegates to. Serialization skips it. + _wrapped_agent=self, ) diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index 8a6d22606d8..6bd976fe054 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -394,6 +394,11 @@ class ToolApprovalMiddleware(AgentMiddleware): upgraded agent may use an improved version of a tool for which the old approval is still appropriate. The developer must decide whether each standing approval should carry over to the changed tool. + + The rules are not tied to a caller either. Callers that share a session + share its standing approvals, for example when one ``AgentExecutor`` + serves several callers. Give each caller its own session, or its own + workflow and executors. """ def __init__( diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 8aa177ec117..68affb64e05 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -664,6 +664,8 @@ class WeatherArgs(BaseModel): "_schema_supplied", "_invoke_sync_on_event_loop", } + # A tool that ``BaseAgent.as_tool()`` made keeps pointing at its agent when copied, as its function does. + _SHALLOW_COPY_FIELDS: ClassVar[set[str]] = {"raw_representation", "_wrapped_agent"} def __init__( self, diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index e71a62a6eb3..5918dae8a59 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -111,7 +111,16 @@ def should_forward(self, role: str | None, message_id: str | None = None) -> boo class WorkflowAgent(BaseAgent): - """An `Agent` subclass that wraps a workflow and exposes it as an agent.""" + """An `Agent` subclass that wraps a workflow and exposes it as an agent. + + A ``WorkflowAgent`` wraps one stateful workflow. The ``session`` passed to :meth:`run` applies to the + wrapper only. It is not passed to the executors inside the workflow, and an ``AgentExecutor`` keeps its + own session between runs. One instance must therefore not serve callers that need independent + state, and while one caller's request is pending, other runs of the same instance fail until it is + resolved. Create a new workflow and ``WorkflowAgent`` for each caller or conversation. Exposing a + ``WorkflowAgent`` as a tool with ``as_tool()`` is not supported, because every caller of the tool would share + this workflow. + """ # Class variable for the request info function name REQUEST_INFO_FUNCTION_NAME: ClassVar[str] = "request_info" diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index 15a62177b74..a434609299c 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -304,6 +304,12 @@ class AgentExecutor(Executor): response to be sent downstream. The executor automatically detects the mode via WorkflowContext.is_streaming(). + + Each executor keeps one ``AgentSession`` between runs, together with its cached messages and + pending user-input requests. Conversation history, pending approvals, and standing approval rules live in + that session, and the session is not tied to a caller. Do not let one executor, or a workflow that holds + it, serve callers that need independent state. Create the workflow and its executors for each caller or + conversation instead, and persist continuation through checkpoints. """ def __init__( diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 28924882830..8a515a5a55f 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -4,6 +4,7 @@ import asyncio import contextlib +import copy import inspect import json import logging @@ -1816,6 +1817,25 @@ async def test_chat_agent_as_tool_basic(client: SupportsChatGetResponse) -> None assert tool.input_model is None +async def test_chat_agent_as_tool_records_the_agent_it_delegates_to(client: SupportsChatGetResponse) -> None: + """Hosts that refuse some delegates read a private link, and serialization leaves it out.""" + agent = Agent(client=client, name="TestAgent", description="Test agent for as_tool") + + tool = agent.as_tool() + + assert getattr(tool, "_wrapped_agent", None) is agent + assert "_wrapped_agent" not in tool.to_dict() + + +async def test_chat_agent_as_tool_copies_keep_pointing_at_the_wrapped_agent(client: SupportsChatGetResponse) -> None: + """A copy of the tool must not copy the agent it delegates to, just as the tool function does not.""" + agent = Agent(client=client, name="TestAgent", description="Test agent for as_tool") + tool = agent.as_tool() + + assert getattr(copy.copy(tool), "_wrapped_agent", None) is agent + assert getattr(copy.deepcopy(tool), "_wrapped_agent", None) is agent + + async def test_chat_agent_as_tool_custom_parameters( client: SupportsChatGetResponse, ) -> None: diff --git a/python/packages/hosting/README.md b/python/packages/hosting/README.md index 1910c7c615c..b0fd29f0bee 100644 --- a/python/packages/hosting/README.md +++ b/python/packages/hosting/README.md @@ -125,9 +125,24 @@ result = await (await state.get_target()).run("Hello", checkpoint_storage=storag latest = await storage.get_latest(workflow_name=(await state.get_target()).name) ``` -A `Workflow` instance allows only one active run. Hosts that need simultaneous -workflow runs should pass a factory or builder and set `cache_target=False` so -each request receives a fresh workflow instance. +A `Workflow` instance allows only one active run, and it keeps its state between +runs. An `AgentExecutor` keeps one `AgentSession` (history, pending approvals, +and standing approval rules) between runs. Hosts that need +simultaneous or independent runs should pass a factory that creates the +workflow and its executors on every call, and set `cache_target=False`: + +```python +from agent_framework import Agent, AgentExecutor, Workflow, WorkflowBuilder +from agent_framework_hosting import WorkflowState + + +def create_workflow() -> Workflow: + writer = AgentExecutor(Agent(client=client, name="writer")) + return WorkflowBuilder(start_executor=writer).build() + + +state = WorkflowState(create_workflow, cache_target=False) +``` `WorkflowState` also accepts an unbuilt workflow builder directly: @@ -144,6 +159,13 @@ This is structural: orchestration builders from `agent_framework_orchestrations` and `MagenticBuilder`) also work because they expose the same zero-argument `build() -> Workflow` method. +With `cache_target=False`, a builder returns a new `Workflow` for each call, but +the new workflow reuses the executors registered on the builder. An +`AgentExecutor` created before the builder therefore carries its session from +one request to the next. Use a builder for independent requests only when it +creates its executors in `build()`. The orchestration builders do this for the +agents you give them, but they share any executor instance you pass in. + Cross-channel identity linking, multicast delivery, background runs, continuation tokens, and durable delivery runners are follow-up enhancements, not part of this v1 state surface. diff --git a/python/packages/hosting/agent_framework_hosting/_state.py b/python/packages/hosting/agent_framework_hosting/_state.py index 34f02cb3afd..72d5963feca 100644 --- a/python/packages/hosting/agent_framework_hosting/_state.py +++ b/python/packages/hosting/agent_framework_hosting/_state.py @@ -197,6 +197,16 @@ class WorkflowState(Generic[WorkflowT]): Holds the Agent Framework workflow target. Does not own routes, middleware, protocol dispatch, native SDK calls, or checkpoint storage -- web frameworks and workflow/app code keep those concerns. + + A ``Workflow`` keeps its state between runs, and an ``AgentExecutor`` keeps + one ``AgentSession`` between runs, so one target must not serve + callers that need independent state. For independent requests, pass a + factory that creates the workflow and its executors on every call and set + ``cache_target=False``. A builder target returns a new ``Workflow`` per + call but reuses the executors registered on it; only a builder that + creates its executors in ``build()`` gives each call new executors. The + orchestration builders do this for the agents they receive, not for + executor instances you pass in. """ def __init__( @@ -217,7 +227,9 @@ def __init__( Keyword Args: cache_target: Whether to cache a resolved callable/awaitable/built target. Defaults to ``True`` so expensive target setup - happens once. + happens once. ``False`` resolves a new target for every call; + the new target still shares whatever the factory or builder + reuses, such as executors created outside it. Raises: ValueError: If ``cache_target=False`` is used with a one-shot diff --git a/python/packages/hosting/tests/hosting/_workflow_fixtures.py b/python/packages/hosting/tests/hosting/_workflow_fixtures.py index 72addb7dff8..48fe349c6a0 100644 --- a/python/packages/hosting/tests/hosting/_workflow_fixtures.py +++ b/python/packages/hosting/tests/hosting/_workflow_fixtures.py @@ -7,9 +7,26 @@ rather than stringified forms. """ +from collections.abc import AsyncIterator, Awaitable, Mapping, Sequence from typing import Any -from agent_framework import Executor, Workflow, WorkflowBuilder, WorkflowContext, handler +from agent_framework import ( + Agent, + AgentExecutor, + BaseChatClient, + ChatOptions, + ChatResponse, + ChatResponseUpdate, + Content, + Executor, + InMemoryHistoryProvider, + Message, + ResponseStream, + Workflow, + WorkflowBuilder, + WorkflowContext, + handler, +) class _UpperExecutor(Executor): @@ -48,3 +65,53 @@ async def handle(self, text: str, ctx: WorkflowContext[Any, str]) -> None: def build_multi_chunk_workflow() -> Workflow: return WorkflowBuilder(start_executor=_MultiChunkExecutor(id="multi")).build() + + +class RecordingClient(BaseChatClient[ChatOptions[None]]): + """Chat client that records the text of every model-bound message and always replies ``response``.""" + + def __init__(self) -> None: + super().__init__() + self.calls: list[list[str]] = [] + + def _inner_get_response( + self, + *, + messages: Sequence[Message], + stream: bool = False, + options: Mapping[str, Any], + **kwargs: Any, + ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]: + self.calls.append([message.text for message in messages]) + if stream: + + async def updates() -> AsyncIterator[ChatResponseUpdate]: + yield ChatResponseUpdate(role="assistant", contents=[Content.from_text("response")]) + + return ResponseStream(updates(), finalizer=ChatResponse.from_updates) + + async def get_response() -> ChatResponse: + return ChatResponse(messages=Message("assistant", ["response"])) + + return get_response() + + +def agent_executor_workflow_builder(client: RecordingClient) -> WorkflowBuilder: + """Return an *unbuilt* builder whose ``AgentExecutor`` is created once and shared by every build.""" + agent = Agent(client=client, name="worker", context_providers=[InMemoryHistoryProvider(source_id="history")]) + executor = AgentExecutor(agent, id="worker") + return WorkflowBuilder(start_executor=executor, output_from=[executor]) + + +def build_agent_executor_workflow(client: RecordingClient) -> Workflow: + """Build a workflow around a new ``AgentExecutor``, which owns a new ``AgentSession``.""" + return agent_executor_workflow_builder(client).build() + + +async def run_workflow(workflow: Workflow, text: str, *, stream: bool) -> None: + """Run ``workflow`` once, streaming or not, and drain the result.""" + if stream: + async for _ in workflow.run(text, stream=True): + pass + else: + await workflow.run(text) diff --git a/python/packages/hosting/tests/hosting/test_state.py b/python/packages/hosting/tests/hosting/test_state.py index ab6223dde17..ec1cbc80a65 100644 --- a/python/packages/hosting/tests/hosting/test_state.py +++ b/python/packages/hosting/tests/hosting/test_state.py @@ -351,3 +351,40 @@ def build(self) -> Workflow: state: WorkflowState[Workflow] = WorkflowState(_FakeOrchestrationBuilder()) assert await state.get_target() is workflow + + @pytest.mark.parametrize("stream", [False, True]) + async def test_factory_that_creates_executors_gives_each_target_its_own_agent_session(self, stream: bool) -> None: + client = _workflow_fixture("RecordingClient")() + build_workflow = _workflow_fixture("build_agent_executor_workflow") + run_workflow = _workflow_fixture("run_workflow") + state: WorkflowState[Workflow] = WorkflowState(lambda: build_workflow(client), cache_target=False) + + first = await state.get_target() + second = await state.get_target() + await run_workflow(first, "FIRST_CALLER", stream=stream) + await run_workflow(second, "SECOND_CALLER", stream=stream) + + assert first.get_start_executor() is not second.get_start_executor() + assert client.calls == [["FIRST_CALLER"], ["SECOND_CALLER"]] + + @pytest.mark.parametrize("stream", [False, True]) + async def test_builder_target_reuses_the_executors_registered_on_it(self, stream: bool) -> None: + """``cache_target=False`` returns a new ``Workflow`` per call, not new executors. + + An ``AgentExecutor`` registered on the builder keeps its ``AgentSession`` across every + target, so one request's history reaches the next. The README tells hosts to create + executors inside a factory when requests must stay independent. + """ + client = _workflow_fixture("RecordingClient")() + run_workflow = _workflow_fixture("run_workflow") + builder = _workflow_fixture("agent_executor_workflow_builder")(client) + state: WorkflowState[Workflow] = WorkflowState(builder, cache_target=False) + + first = await state.get_target() + second = await state.get_target() + await run_workflow(first, "FIRST_CALLER", stream=stream) + await run_workflow(second, "SECOND_CALLER", stream=stream) + + assert first is not second + assert first.get_start_executor() is second.get_start_executor() + assert client.calls[-1] == ["FIRST_CALLER", "response", "SECOND_CALLER"] diff --git a/python/samples/04-hosting/af-hosting/local_responses_workflow/README.md b/python/samples/04-hosting/af-hosting/local_responses_workflow/README.md index b0680a5e600..a760c301a79 100644 --- a/python/samples/04-hosting/af-hosting/local_responses_workflow/README.md +++ b/python/samples/04-hosting/af-hosting/local_responses_workflow/README.md @@ -19,9 +19,13 @@ stored under the sample's local `storage/` root. Checkpoints are scoped into per-continuation buckets so a "latest checkpoint" lookup cannot cross conversations. -The sample passes a builder with `cache_target=False`, producing a fresh -workflow instance for every request. This is required for simultaneous runs -because one `Workflow` instance intentionally permits only one active run. +The sample passes a factory with `cache_target=False`. The factory creates the +workflow and its `AgentExecutor` for every request. This is required for +simultaneous runs, because one `Workflow` instance intentionally permits only +one active run. It also keeps requests independent: an `AgentExecutor` keeps one +`AgentSession` between runs, so an executor created at module scope +would carry one request's history into the next. The workflow name and executor +ids stay stable, so a checkpoint restores into the next request's new workflow. ## Production readiness diff --git a/python/samples/04-hosting/af-hosting/local_responses_workflow/app.py b/python/samples/04-hosting/af-hosting/local_responses_workflow/app.py index 026b56e7e12..d674ca53fed 100644 --- a/python/samples/04-hosting/af-hosting/local_responses_workflow/app.py +++ b/python/samples/04-hosting/af-hosting/local_responses_workflow/app.py @@ -7,6 +7,7 @@ 1. ``agent-framework-hosting-responses`` converts the Responses request body to Agent Framework run values and renders the final response payload. 2. ``agent-framework-hosting`` resolves the workflow target via ``WorkflowState``. + A factory creates the workflow and its executors for every request. 3. FastAPI owns the route, request parsing, policy decisions, response object, and file-backed checkpoint cursor. @@ -62,6 +63,7 @@ Executor, FileCheckpointStorage, Message, + Workflow, WorkflowBuilder, WorkflowContext, handler, @@ -202,23 +204,35 @@ async def handle(self, response: AgentExecutorResponse, ctx: WorkflowContext[Any client = FoundryChatClient(credential=DefaultAzureCredential()) -writer = Agent( - client=client, - name="writer", - instructions="You are an excellent slogan writer. Create one short slogan from the given brief.", -) -writer_ex = AgentExecutor(writer, context_mode="last_agent") -formatter_ex = TerminalFormatter(id="terminal_formatter") -workflow_builder = WorkflowBuilder( - name="local_responses_slogan_workflow", - start_executor=writer_ex, - output_from=[formatter_ex], -).add_edge(writer_ex, formatter_ex) + +def create_workflow() -> Workflow: + """Build the workflow and its executors for one request. + + An ``AgentExecutor`` keeps one ``AgentSession`` between runs, so every + request needs its own executors. The workflow name and the executor ids stay + stable, which lets a checkpoint restore into the next request's new workflow. + """ + writer = Agent( + client=client, + name="writer", + instructions="You are an excellent slogan writer. Create one short slogan from the given brief.", + ) + writer_ex = AgentExecutor(writer, context_mode="last_agent") + formatter_ex = TerminalFormatter(id="terminal_formatter") + return ( + WorkflowBuilder( + name="local_responses_slogan_workflow", + start_executor=writer_ex, + output_from=[formatter_ex], + ) + .add_edge(writer_ex, formatter_ex) + .build() + ) app = FastAPI() -state = WorkflowState(workflow_builder, cache_target=False) +state = WorkflowState(create_workflow, cache_target=False) @app.post("/responses", response_model=None) diff --git a/skills/agent-framework-builder-python/references/hosting-and-deployment.md b/skills/agent-framework-builder-python/references/hosting-and-deployment.md index 480c09fc374..732b600240d 100644 --- a/skills/agent-framework-builder-python/references/hosting-and-deployment.md +++ b/skills/agent-framework-builder-python/references/hosting-and-deployment.md @@ -23,7 +23,9 @@ into application code. - Reuse an agent instance only when all mutable state is session-scoped. Otherwise use a request-owned factory. - Prefer native workflow hosting. A shared `workflow.as_agent()` instance can leak workflow state across requests. - Build fresh workflows and request-owned resources while keeping workflow and executor identities stable for - checkpoint restoration. + checkpoint restoration. A builder reused across requests returns new workflows that share its registered + executors, including each `AgentExecutor` and its `AgentSession`; create those executors inside the per-request + factory. - Make session, checkpoint, response, identity, approval, storage, and cleanup ownership explicit. Read the selected host's public docstrings and maintained sample for constructor options, parsers, history modes, diff --git a/skills/agent-framework-builder-python/references/workflow-builder.md b/skills/agent-framework-builder-python/references/workflow-builder.md index c121e97c66e..cb3039a86ed 100644 --- a/skills/agent-framework-builder-python/references/workflow-builder.md +++ b/skills/agent-framework-builder-python/references/workflow-builder.md @@ -19,7 +19,9 @@ docstrings for signatures and execution options. ## Stable design rules - A functional workflow definition must be built before running; each built workflow owns mutable execution state. -- Build fresh workflows when executors or wrapped agents hold run-specific state. +- Build fresh workflows when executors or wrapped agents hold run-specific state. A reused builder returns new + workflows that share its registered executors, so create stateful executors, such as `AgentExecutor` with its + `AgentSession`, inside the per-request factory. - Builders that accept `SupportsAgentRun` can use custom or remote agent implementations. `HandoffBuilder` accepts concrete `Agent` participants only. - Use standard sequential, concurrent, handoff, group-chat, or Magentic builders before recreating those patterns.