From 0a4f457c02656b0c4217766a5949d28ef82536fb Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 3 Nov 2026 12:29:32 +0100 Subject: [PATCH 1/7] Python: Bind AG-UI workflow agent executors to one authorization scope An AgentExecutor owns one AgentSession for its lifetime, so a shared Workflow or WorkflowAgent that serves several authorization scopes shares that conversation and approval state between them. Bind every AgentExecutor of a workflow, including nested workflows, to the first authorization scope that runs it. A run in another scope now ends with a WORKFLOW_SCOPE_CONFLICT RUN_ERROR before it reaches the model. Workflow factories that build new executors for every call, single-user endpoints, and workflows without an AgentExecutor keep working unchanged. --- python/packages/ag-ui/AGENTS.md | 11 + python/packages/ag-ui/README.md | 24 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 11 + .../ag-ui/agent_framework_ag_ui/_endpoint.py | 13 +- .../_session_ownership.py | 73 +++++ .../ag-ui/agent_framework_ag_ui/_workflow.py | 20 +- .../ag_ui/test_workflow_scope_ownership.py | 291 ++++++++++++++++++ .../tests/ag_ui/workflow_scope_helpers.py | 187 +++++++++++ 8 files changed, 623 insertions(+), 7 deletions(-) create mode 100644 python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py create mode 100644 python/packages/ag-ui/tests/ag_ui/test_workflow_scope_ownership.py create mode 100644 python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 05486f1d285..c03a2f6979f 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -131,6 +131,17 @@ 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` owns one `AgentSession` (history, pending approvals, standing approval rules) that is not tied to a + caller. `_session_ownership.claim_agent_executors` binds every `AgentExecutor` of a workflow, including the ones in + nested `WorkflowExecutor` workflows and in the workflow behind a `WorkflowAgent` that an `AgentExecutor` wraps, to the + first non-`None` Snapshot Scope that runs it. The claim is one process-wide locked check-and-bind. + `AgentFrameworkWorkflow.run` and the `WorkflowAgent` branch of `run_agent_stream` call it before any workflow state is + read; 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 or + executors. It does not apply to workflows without an `AgentExecutor`, to a `None` scope, or to custom executors that + keep their own state, and it 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. - 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..b7a7e6e79c0 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,20 @@ 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` owns one `AgentSession` for its whole lifetime. 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` 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. Pass `AgentFrameworkWorkflow(workflow_factory=...)` with a factory that builds a new workflow, including new +`AgentExecutor` instances, for every call to serve several scopes. A factory that returns a workflow or executor 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. + ### 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..4a00638a85e 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,12 @@ _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, + claim_agent_executors, +) from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, _SNAPSHOT_SCOPE_INPUT_KEY, @@ -2983,6 +2989,11 @@ async def _run_agent_stream( thread_id = supplied_thread_id or str(uuid.uuid4()) run_id = supplied_run_id or str(uuid.uuid4()) snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY)) + if isinstance(agent, WorkflowAgent) and not claim_agent_executors(agent.workflow, 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 session_id = _session_id_for_thread( scope=snapshot_scope, thread_id=thread_id, 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..2ffb9a96f02 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -164,7 +164,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 +189,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`` 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) 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..a0f73d2c3d0 --- /dev/null +++ b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Authorization-scope ownership of the agent sessions that a workflow holds. + +An ``AgentExecutor`` owns one ``AgentSession`` for its whole lifetime. Conversation history, pending approval +requests, and standing approval rules live in that session, and the session is not tied to a caller. An executor +that serves two authorization scopes would therefore expose the conversation and approval state of one scope to the +other. AG-UI binds each executor to the first authorization scope that runs it and refuses every other scope. +""" + +from __future__ import annotations + +import threading +import weakref +from collections.abc import Iterator + +from agent_framework import AgentExecutor, 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 AgentExecutor instances that already belong to another authorization scope. Each " + "AgentExecutor owns one AgentSession, so serving a second scope would expose the first scope's conversation " + "and approval state. Pass AgentFrameworkWorkflow(workflow_factory=...) with a factory that builds a new " + "Workflow, including new AgentExecutor instances, for every call, or register the endpoint with " + "single_user=True when every caller is the same principal." +) + +_EXECUTOR_SCOPES: weakref.WeakKeyDictionary[AgentExecutor, str] = weakref.WeakKeyDictionary() +_CLAIM_LOCK = threading.Lock() + + +def _agent_executors(workflow: Workflow, 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. + """ + 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 + if isinstance(executor.agent, WorkflowAgent): + nested = executor.agent.workflow + elif isinstance(executor, WorkflowExecutor): + nested = executor.workflow + if nested is not None: + yield from _agent_executors(nested, visited) + + +def claim_agent_executors(workflow: Workflow, scope: str | None) -> bool: + """Bind every ``AgentExecutor`` of ``workflow`` 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`` means the run has no scope, so nothing is bound. + + Returns: + ``False`` when an executor already belongs to a different scope. Nothing is bound in that case. + """ + if scope is None: + return True + executors = list(_agent_executors(workflow)) + with _CLAIM_LOCK: + if any(_EXECUTOR_SCOPES.get(executor, scope) != scope for executor in executors): + return False + for executor in executors: + _EXECUTOR_SCOPES[executor] = scope + return True 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..43698cf5cc4 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,12 @@ _snapshot_messages_match, _snapshot_overlap_length, ) +from ._session_ownership import ( + SCOPE_CONFLICT_CODE, + SCOPE_CONFLICT_LOG_MESSAGE, + SCOPE_CONFLICT_MESSAGE, + claim_agent_executors, +) from ._snapshot_session import ThreadSnapshotSession, _event_messages_to_snapshot_dicts from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, @@ -671,8 +677,13 @@ 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`` instances, on every call. Each authorization scope and thread then keeps its own + executor sessions. name: Optional workflow name. description: Optional workflow description. snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless @@ -806,6 +817,11 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: ) return workflow = self._resolve_workflow(thread_id, snapshot_scope) + 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 checkpoint_storage = self.checkpoint_storage or _workflow_checkpoint_storage(workflow) owned_checkpoint_storage = ( _OwnedWorkflowCheckpointStorage( 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..ba073658d96 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_scope_ownership.py @@ -0,0 +1,291 @@ +# 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 gc +import weakref +from concurrent.futures import ThreadPoolExecutor +from typing import Any + +import pytest +from agent_framework import ( + AgentExecutor, + 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, + RecordingChatClient, + always_approve_resume, + assistant_text, + build_agent, + build_app, + 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 + +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) + + +# 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") + # Direct use without a scope has no principal to bind. + assert claim_agent_executors(workflow, None) + + +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_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..621cf83ece1 --- /dev/null +++ b/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py @@ -0,0 +1,187 @@ +# 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, + WorkflowBuilder, + 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 AGUIRequest, InMemoryAGUIThreadSnapshotStore, add_agent_framework_fastapi_endpoint + +ScopeMode = Literal["authorization", "snapshot"] + +REPLY_TEXT = "model-reply" +GATED_TOOL_NAME = "gated_action" +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]) -> Workflow: + """Build a new workflow around a new ``AgentExecutor``, which owns a new ``AgentSession``.""" + executor = AgentExecutor(agent, id="worker") + return WorkflowBuilder(start_executor=executor, output_from=[executor]).build() + + +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()}] From f51585ddf6f53993814e7ad80f2ace886c7a9bbf Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 3 Nov 2026 12:36:06 +0100 Subject: [PATCH 2/7] Python: Create executors per request in the hosting workflow guidance WorkflowBuilder.build() reuses the executors registered on the builder, so a builder target with cache_target=False returns a new Workflow for each request that still shares one AgentExecutor and its AgentSession. Document that in the hosting README and the WorkflowState docstring, create the workflow and its executors inside a per-request factory in the Responses workflow sample, add tests for both target shapes, and note the rule in the Python builder skill references. --- python/packages/hosting/README.md | 28 +++++++- .../hosting/agent_framework_hosting/_state.py | 14 +++- .../tests/hosting/_workflow_fixtures.py | 69 ++++++++++++++++++- .../hosting/tests/hosting/test_state.py | 37 ++++++++++ .../local_responses_workflow/README.md | 10 ++- .../local_responses_workflow/app.py | 40 +++++++---- .../references/hosting-and-deployment.md | 4 +- .../references/workflow-builder.md | 4 +- 8 files changed, 183 insertions(+), 23 deletions(-) diff --git a/python/packages/hosting/README.md b/python/packages/hosting/README.md index 1910c7c615c..42838e2d54a 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) for its whole lifetime. 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..e47c7936970 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`` for its whole lifetime, 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..a47b26ee61a 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` for its whole lifetime, 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..2c2704cfe01 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`` for its whole lifetime, 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. From 2ebcbd2cd0153444099ef0f9fd718bb369e97b15 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 3 Nov 2026 12:36:07 +0100 Subject: [PATCH 3/7] Python: Document that AgentExecutor and WorkflowAgent own one session Clarify in the AgentExecutor, WorkflowAgent, and ToolApprovalMiddleware docstrings that an instance keeps one session, with its history, pending approvals, and standing approval rules, for its whole lifetime and must not serve callers that need independent state. --- .../core/agent_framework/_harness/_tool_approval.py | 5 +++++ .../packages/core/agent_framework/_workflows/_agent.py | 9 ++++++++- .../core/agent_framework/_workflows/_agent_executor.py | 6 ++++++ 3 files changed, 19 insertions(+), 1 deletion(-) 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/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index e71a62a6eb3..58a98fbb2ba 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -111,7 +111,14 @@ 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 for its whole lifetime. 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. + """ # 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..eec0a64d6f9 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 owns one ``AgentSession`` for its whole lifetime, 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__( From e253d29142ebff424d360770fe66d6bbf4230f0f Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Tue, 3 Nov 2026 13:26:03 +0100 Subject: [PATCH 4/7] Python: Bind agent sessions and unscoped runs to one AG-UI owner Look through A2UI runners when claiming the workflow behind a WorkflowAgent. Bind the AgentSession that an AgentExecutor holds as well as the executor, so new executors that inject one shared session are refused for another scope. Treat a run without a scope as an owner of its own, so a workflow that ran without a scope is not handed to an authorization scope later, and the reverse. Claim only after a snapshot replay has returned, so a read-only replay does not take ownership. --- python/packages/ag-ui/AGENTS.md | 23 +-- python/packages/ag-ui/README.md | 14 +- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 14 +- .../ag-ui/agent_framework_ag_ui/_endpoint.py | 2 +- .../_session_ownership.py | 92 ++++++++--- .../ag-ui/agent_framework_ag_ui/_workflow.py | 16 +- .../ag_ui/test_workflow_scope_ownership.py | 148 +++++++++++++++++- .../tests/ag_ui/workflow_scope_helpers.py | 9 +- 8 files changed, 259 insertions(+), 59 deletions(-) diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index c03a2f6979f..9da9f134797 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -132,16 +132,19 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. 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` owns one `AgentSession` (history, pending approvals, standing approval rules) that is not tied to a - caller. `_session_ownership.claim_agent_executors` binds every `AgentExecutor` of a workflow, including the ones in - nested `WorkflowExecutor` workflows and in the workflow behind a `WorkflowAgent` that an `AgentExecutor` wraps, to the - first non-`None` Snapshot Scope that runs it. The claim is one process-wide locked check-and-bind. - `AgentFrameworkWorkflow.run` and the `WorkflowAgent` branch of `run_agent_stream` call it before any workflow state is - read; 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 or - executors. It does not apply to workflows without an `AgentExecutor`, to a `None` scope, or to custom executors that - keep their own state, and it 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. + caller. `_session_ownership.claim_agent_executors` binds every `AgentExecutor` of a workflow, and the `AgentSession` + each one holds, 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 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 does not apply to + workflows without an `AgentExecutor` or to custom executors that keep their own state, and it 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. - 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 b7a7e6e79c0..ed283ffe163 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -360,12 +360,14 @@ callers can reach the endpoint. An `AgentExecutor` owns one `AgentSession` for its whole lifetime. 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` 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. Pass `AgentFrameworkWorkflow(workflow_factory=...)` with a factory that builds a new workflow, including new -`AgentExecutor` instances, for every call to serve several scopes. A factory that returns a workflow or executor 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. +`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. 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 4a00638a85e..f0425311882 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 @@ -118,7 +118,7 @@ SCOPE_CONFLICT_CODE, SCOPE_CONFLICT_LOG_MESSAGE, SCOPE_CONFLICT_MESSAGE, - claim_agent_executors, + claim_workflow_agent, ) from ._snapshots import ( _DEFAULT_STATE_INPUT_KEY, @@ -2989,11 +2989,6 @@ async def _run_agent_stream( thread_id = supplied_thread_id or str(uuid.uuid4()) run_id = supplied_run_id or str(uuid.uuid4()) snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY)) - if isinstance(agent, WorkflowAgent) and not claim_agent_executors(agent.workflow, 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 session_id = _session_id_for_thread( scope=snapshot_scope, thread_id=thread_id, @@ -3040,6 +3035,13 @@ 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. + 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 2ffb9a96f02..9e7c1227f41 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -194,7 +194,7 @@ def add_agent_framework_fastapi_endpoint( 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`` instances, for every call to serve several scopes. + ``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 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 index a0f73d2c3d0..4608db219e8 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py @@ -3,9 +3,11 @@ """Authorization-scope ownership of the agent sessions that a workflow holds. An ``AgentExecutor`` owns one ``AgentSession`` for its whole lifetime. Conversation history, pending approval -requests, and standing approval rules live in that session, and the session is not tied to a caller. An executor -that serves two authorization scopes would therefore expose the conversation and approval state of one scope to the -other. AG-UI binds each executor to the first authorization scope that runs it and refuses every other scope. +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. """ from __future__ import annotations @@ -13,29 +15,52 @@ import threading import weakref from collections.abc import Iterator +from typing import Any, Final -from agent_framework import AgentExecutor, Workflow, WorkflowAgent, WorkflowExecutor +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 AgentExecutor instances that already belong to another authorization scope. Each " - "AgentExecutor owns one AgentSession, so serving a second scope would expose the first scope's conversation " - "and approval state. Pass AgentFrameworkWorkflow(workflow_factory=...) with a factory that builds a new " - "Workflow, including new AgentExecutor instances, for every call, or register the endpoint with " - "single_user=True when every caller is the same principal." + "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." ) -_EXECUTOR_SCOPES: weakref.WeakKeyDictionary[AgentExecutor, str] = weakref.WeakKeyDictionary() +# The owner of runs that carry no authorization scope. It never equals a scope string. +_UNSCOPED: Final = object() + +_OWNERS: weakref.WeakKeyDictionary[AgentExecutor | AgentSession, object] = weakref.WeakKeyDictionary() _CLAIM_LOCK = threading.Lock() -def _agent_executors(workflow: Workflow, visited: set[int] | None = None) -> Iterator[AgentExecutor]: +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. + ``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 @@ -44,30 +69,51 @@ def _agent_executors(workflow: Workflow, visited: set[int] | None = None) -> Ite nested: Workflow | None = None if isinstance(executor, AgentExecutor): yield executor - if isinstance(executor.agent, WorkflowAgent): - nested = executor.agent.workflow + 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`` to ``scope``. + """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`` means the run has no scope, so nothing is bound. + 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 already belongs to a different scope. Nothing is bound in that case. + ``False`` when an executor or session already belongs to a different owner. Nothing is bound in that case. """ - if scope is None: - return True - executors = list(_agent_executors(workflow)) + owner: object = _UNSCOPED if scope is None else scope + holders = list(_state_holders(workflow)) with _CLAIM_LOCK: - if any(_EXECUTOR_SCOPES.get(executor, scope) != scope for executor in executors): + if any(_OWNERS.get(holder, owner) != owner for holder in holders): return False - for executor in executors: - _EXECUTOR_SCOPES[executor] = scope + for holder in holders: + _OWNERS[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) 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 43698cf5cc4..8ec783ece43 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -682,8 +682,8 @@ def __init__( ``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`` instances, on every call. Each authorization scope and thread then keeps its own - executor sessions. + new ``AgentExecutor`` and ``AgentSession`` instances, on every call. Each authorization scope and + thread then keeps its own executor sessions. name: Optional workflow name. description: Optional workflow description. snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless @@ -817,11 +817,6 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: ) return workflow = self._resolve_workflow(thread_id, snapshot_scope) - 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 checkpoint_storage = self.checkpoint_storage or _workflow_checkpoint_storage(workflow) owned_checkpoint_storage = ( _OwnedWorkflowCheckpointStorage( @@ -896,6 +891,13 @@ 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. + 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_workflow_scope_ownership.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_scope_ownership.py index ba073658d96..213e69bac02 100644 --- 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 @@ -4,14 +4,16 @@ from __future__ import annotations +import asyncio import gc import weakref from concurrent.futures import ThreadPoolExecutor -from typing import Any +from typing import Any, cast import pytest from agent_framework import ( AgentExecutor, + SupportsAgentRun, Workflow, WorkflowBuilder, WorkflowContext, @@ -35,7 +37,7 @@ ) from agent_framework_ag_ui import AgentFrameworkAgent, AgentFrameworkWorkflow -from agent_framework_ag_ui._session_ownership import claim_agent_executors +from agent_framework_ag_ui._session_ownership import claim_agent_executors, claim_workflow_agent SCOPE_MODES = ["authorization", "snapshot"] CONFLICT_CODE = "WORKFLOW_SCOPE_CONFLICT" @@ -139,6 +141,89 @@ def test_workflow_agent_serves_only_the_first_authorization_scope(scope_mode: An 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) + + +# 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 @@ -247,8 +332,45 @@ def test_claim_binds_agent_executors_to_the_first_scope() -> None: assert claim_agent_executors(workflow, "scope-a") assert claim_agent_executors(workflow, "scope-a") assert not claim_agent_executors(workflow, "scope-b") - # Direct use without a scope has no principal to bind. + + +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_covers_agent_executors_inside_nested_workflows() -> None: @@ -268,6 +390,26 @@ def test_claim_covers_workflows_behind_a_wrapped_workflow_agent() -> None: 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())) 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 index 621cf83ece1..c30a6cd6e67 100644 --- a/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py +++ b/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py @@ -94,9 +94,12 @@ def gated_action(value: str) -> str: ) -def build_workflow(agent: Agent[Any]) -> Workflow: - """Build a new workflow around a new ``AgentExecutor``, which owns a new ``AgentSession``.""" - executor = AgentExecutor(agent, id="worker") +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() From fe6f14be99c7c892e20f4701d703a960326d1d12 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 4 Nov 2026 09:41:11 +0100 Subject: [PATCH 5/7] Python: Key AG-UI session ownership by object identity The ownership registry used a WeakKeyDictionary, which hashes its keys. An AgentSession or AgentExecutor subclass that defines __eq__ is unhashable, so every AG-UI run on it failed, and two distinct objects that compare equal could share one entry. Key the registry by id() with a weak reference, so user-defined hashing and equality never run and an entry disappears when its object is collected. Document what the ownership rule does not cover, such as state that a factory loads from storage into a new object, and say that an executor keeps its session between runs because a checkpoint restore replaces it. --- python/packages/ag-ui/AGENTS.md | 42 ++++++--- python/packages/ag-ui/README.md | 12 ++- .../_session_ownership.py | 59 ++++++++++-- .../ag-ui/agent_framework_ag_ui/_workflow.py | 4 +- .../ag-ui/tests/ag_ui/test_identity_owners.py | 94 +++++++++++++++++++ .../ag_ui/test_workflow_scope_ownership.py | 51 ++++++++++ .../tests/ag_ui/workflow_scope_helpers.py | 30 ++++++ .../core/agent_framework/_workflows/_agent.py | 2 +- .../_workflows/_agent_executor.py | 2 +- python/packages/hosting/README.md | 2 +- .../hosting/agent_framework_hosting/_state.py | 2 +- .../local_responses_workflow/README.md | 2 +- .../local_responses_workflow/app.py | 2 +- 13 files changed, 273 insertions(+), 31 deletions(-) create mode 100644 python/packages/ag-ui/tests/ag_ui/test_identity_owners.py diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 9da9f134797..27d5d789ca2 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -131,20 +131,34 @@ 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` owns one `AgentSession` (history, pending approvals, standing approval rules) that is not tied to a - caller. `_session_ownership.claim_agent_executors` binds every `AgentExecutor` of a workflow, and the `AgentSession` - each one holds, 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 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 does not apply to - workflows without an `AgentExecutor` or to custom executors that keep their own state, and it 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. +- 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`); + - workflows that another agent calls as a tool (`WorkflowAgent.as_tool()`); + - 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. - 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 ed283ffe163..7d1f790dd06 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -358,7 +358,7 @@ callers can reach the endpoint. ### Workflows That Hold Agent Sessions -An `AgentExecutor` owns one `AgentSession` for its whole lifetime. Conversation history, pending approval requests, and +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` @@ -372,6 +372,16 @@ scope; their pending interrupts stay fenced to the scope and thread that raised 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`, workflows that another agent +calls as a tool (`WorkflowAgent.as_tool()`), or state that sessions share outside the session object, for example a +`state` dictionary or a `service_session_id`. + ### 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/_session_ownership.py b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py index 4608db219e8..9d2737a1585 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py @@ -2,12 +2,14 @@ """Authorization-scope ownership of the agent sessions that a workflow holds. -An ``AgentExecutor`` owns one ``AgentSession`` for its whole lifetime. 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. +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. """ from __future__ import annotations @@ -32,7 +34,46 @@ # The owner of runs that carry no authorization scope. It never equals a scope string. _UNSCOPED: Final = object() -_OWNERS: weakref.WeakKeyDictionary[AgentExecutor | AgentSession, object] = weakref.WeakKeyDictionary() + +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() @@ -102,10 +143,10 @@ def claim_agent_executors(workflow: Workflow, scope: str | None) -> bool: owner: object = _UNSCOPED if scope is None else scope holders = list(_state_holders(workflow)) with _CLAIM_LOCK: - if any(_OWNERS.get(holder, owner) != owner for holder in holders): + if any(_OWNERS.owner_of(holder, owner) != owner for holder in holders): return False for holder in holders: - _OWNERS[holder] = owner + _OWNERS.bind(holder, owner) return True 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 8ec783ece43..ea1ea1ece7a 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -683,7 +683,9 @@ def __init__( 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. + 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 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_scope_ownership.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_scope_ownership.py index 213e69bac02..7554e867637 100644 --- 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 @@ -24,11 +24,14 @@ 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, @@ -177,6 +180,35 @@ def test_factory_that_shares_one_agent_session_between_scopes_is_refused(scope_m 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 @@ -373,6 +405,25 @@ def test_claim_allows_executors_of_one_workflow_to_share_a_session() -> None: 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() 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 index c30a6cd6e67..17fb13fe76d 100644 --- a/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py +++ b/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py @@ -103,6 +103,36 @@ def build_workflow(agent: Agent[Any], *, session: AgentSession | None = None) -> 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() + + async def _authenticated_scope(authorization: str = Header(default="")) -> str: principal = authorization.removeprefix("token-") if principal not in PRINCIPALS: diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 58a98fbb2ba..278a44ab65a 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -115,7 +115,7 @@ class WorkflowAgent(BaseAgent): 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 for its whole lifetime. One instance must therefore not serve callers that need independent + 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. """ diff --git a/python/packages/core/agent_framework/_workflows/_agent_executor.py b/python/packages/core/agent_framework/_workflows/_agent_executor.py index eec0a64d6f9..a434609299c 100644 --- a/python/packages/core/agent_framework/_workflows/_agent_executor.py +++ b/python/packages/core/agent_framework/_workflows/_agent_executor.py @@ -305,7 +305,7 @@ class AgentExecutor(Executor): The executor automatically detects the mode via WorkflowContext.is_streaming(). - Each executor owns one ``AgentSession`` for its whole lifetime, together with its cached messages and + 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 diff --git a/python/packages/hosting/README.md b/python/packages/hosting/README.md index 42838e2d54a..b0fd29f0bee 100644 --- a/python/packages/hosting/README.md +++ b/python/packages/hosting/README.md @@ -127,7 +127,7 @@ latest = await storage.get_latest(workflow_name=(await state.get_target()).name) 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) for its whole lifetime. Hosts that need +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`: diff --git a/python/packages/hosting/agent_framework_hosting/_state.py b/python/packages/hosting/agent_framework_hosting/_state.py index e47c7936970..72d5963feca 100644 --- a/python/packages/hosting/agent_framework_hosting/_state.py +++ b/python/packages/hosting/agent_framework_hosting/_state.py @@ -199,7 +199,7 @@ class WorkflowState(Generic[WorkflowT]): web frameworks and workflow/app code keep those concerns. A ``Workflow`` keeps its state between runs, and an ``AgentExecutor`` keeps - one ``AgentSession`` for its whole lifetime, so one target must not serve + 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 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 a47b26ee61a..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 @@ -23,7 +23,7 @@ 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` for its whole lifetime, so an executor created at module scope +`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. 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 2c2704cfe01..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 @@ -209,7 +209,7 @@ async def handle(self, response: AgentExecutorResponse, ctx: WorkflowContext[Any def create_workflow() -> Workflow: """Build the workflow and its executors for one request. - An ``AgentExecutor`` keeps one ``AgentSession`` for its whole lifetime, so every + 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. """ From 8b4e1ddda0a61dc838ecd5ae0ab345a9c9c76e16 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 4 Nov 2026 10:22:48 +0100 Subject: [PATCH 6/7] Python: Reject WorkflowAgent tools in AG-UI AG-UI does not support WorkflowAgent.as_tool(): every caller of an agent that exposes a workflow as a tool reaches the same workflow and its executor sessions, and the tool wrapper hides that workflow from the ownership claim. Record the agent that as_tool() wraps in a private attribute of the tool, keep it by reference when the tool is copied, and follow those links from AG-UI. The endpoint raises ValueError for an agent or workflow that exposes a WorkflowAgent as a tool, and a run ends with a WORKFLOW_AGENT_TOOL_UNSUPPORTED RUN_ERROR when a workflow factory builds one or a tool is added later. The rejection also applies with single_user=True. --- python/packages/ag-ui/AGENTS.md | 14 ++- python/packages/ag-ui/README.md | 13 ++- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 10 ++ .../ag-ui/agent_framework_ag_ui/_endpoint.py | 21 +++- .../_session_ownership.py | 73 ++++++++++++- .../ag-ui/agent_framework_ag_ui/_workflow.py | 10 ++ .../test_workflow_agent_tool_detection.py | 65 +++++++++++ .../test_workflow_agent_tool_rejection.py | 103 ++++++++++++++++++ .../tests/ag_ui/workflow_scope_helpers.py | 97 ++++++++++++++++- .../packages/core/agent_framework/_agents.py | 2 + .../packages/core/agent_framework/_tools.py | 2 + .../packages/core/tests/core/test_agents.py | 20 ++++ 12 files changed, 423 insertions(+), 7 deletions(-) create mode 100644 python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_detection.py create mode 100644 python/packages/ag-ui/tests/ag_ui/test_workflow_agent_tool_rejection.py diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 27d5d789ca2..5c7d7aa7bd6 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -147,7 +147,6 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. - 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`); - - workflows that another agent calls as a tool (`WorkflowAgent.as_tool()`); - executors that are not `AgentExecutor` but keep a session or cache (custom executors, the group chat orchestrator); - workflows without an `AgentExecutor`; @@ -159,6 +158,19 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. 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. 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 7d1f790dd06..4070744014a 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -378,9 +378,16 @@ same session id and state, so two scopes that load one stored session share its 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`, workflows that another agent -calls as a tool (`WorkflowAgent.as_tool()`), or state that sessions share outside the session object, for example a -`state` dictionary or a `service_session_id`. +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()`), because every +caller of that agent would reach the same workflow and its executor sessions. `add_agent_framework_fastapi_endpoint` +raises `ValueError` for an agent or workflow that exposes 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. Serve the +workflow with `AgentFrameworkWorkflow(workflow_factory=...)`, or build the workflow inside the tool function for every +call. AG-UI cannot detect a tool function of your own that calls a workflow, or a workflow tool that a context provider +adds while a run is prepared. ### Stable Single-User Scope 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 f0425311882..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 @@ -118,7 +118,11 @@ 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, @@ -3036,6 +3040,12 @@ async def _run_agent_stream( 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) 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 9e7c1227f41..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. @@ -224,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 index 9d2737a1585..c91be2927df 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py @@ -10,13 +10,16 @@ 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()``). The tool +reaches the executors of one shared workflow, and no scope owns them. ``workflow_agent_tools`` finds those tools. """ from __future__ import annotations import threading import weakref -from collections.abc import Iterator +from collections.abc import Iterator, Mapping, Sequence from typing import Any, Final from agent_framework import AgentExecutor, AgentSession, Workflow, WorkflowAgent, WorkflowExecutor @@ -30,6 +33,8 @@ "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() @@ -158,3 +163,69 @@ def claim_workflow_agent(agent: Any, scope: str | None) -> bool: """ 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 ea1ea1ece7a..5ddedb7ca96 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -53,7 +53,11 @@ 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 ( @@ -894,6 +898,12 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: 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) 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/workflow_scope_helpers.py b/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py index 17fb13fe76d..34c6e36b396 100644 --- a/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py +++ b/python/packages/ag-ui/tests/ag_ui/workflow_scope_helpers.py @@ -18,7 +18,9 @@ Message, ToolApprovalMiddleware, Workflow, + WorkflowAgent, WorkflowBuilder, + WorkflowExecutor, create_always_approve_tool_response, tool, ) @@ -27,12 +29,20 @@ 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 AGUIRequest, InMemoryAGUIThreadSnapshotStore, add_agent_framework_fastapi_endpoint +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") @@ -133,6 +143,91 @@ def build_unhashable_executor_workflow(agent: Agent[Any]) -> Workflow: 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: diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 38ee6a3cb38..3929a620183 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -885,6 +885,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/_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/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: From 68b7ad43f8c559b961ffb951cbf9af05b5e25302 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 4 Nov 2026 10:34:25 +0100 Subject: [PATCH 7/7] Python: Document that WorkflowAgent tools are unsupported in any form Say in the AG-UI README and AGENTS.md, the WorkflowAgent and as_tool() docstrings, and the ownership module that exposing a WorkflowAgent as a tool is not supported, including the forms that AG-UI cannot detect: a tool that a context provider adds while a run is prepared, and a tool function that calls a workflow from custom code. AG-UI gives the callers of those no isolation. --- python/packages/ag-ui/AGENTS.md | 9 +++++---- python/packages/ag-ui/README.md | 15 ++++++++------- .../agent_framework_ag_ui/_session_ownership.py | 5 +++-- python/packages/core/agent_framework/_agents.py | 3 +++ .../core/agent_framework/_workflows/_agent.py | 4 +++- 5 files changed, 22 insertions(+), 14 deletions(-) diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 5c7d7aa7bd6..ab01c7005a3 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -167,10 +167,11 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. 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. 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()`. + 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 4070744014a..0101059166a 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -381,13 +381,14 @@ id alone. To continue a conversation, use `checkpoint_storage`: AG-UI restores a 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()`), because every -caller of that agent would reach the same workflow and its executor sessions. `add_agent_framework_fastapi_endpoint` -raises `ValueError` for an agent or workflow that exposes 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. Serve the -workflow with `AgentFrameworkWorkflow(workflow_factory=...)`, or build the workflow inside the tool function for every -call. AG-UI cannot detect a tool function of your own that calls a workflow, or a workflow tool that a context provider -adds while a run is prepared. +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 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 index c91be2927df..9eb36b09af6 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_session_ownership.py @@ -11,8 +11,9 @@ 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()``). The tool -reaches the executors of one shared workflow, and no scope owns them. ``workflow_agent_tools`` finds those tools. +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 diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 3929a620183..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 diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index 278a44ab65a..5918dae8a59 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -117,7 +117,9 @@ class WorkflowAgent(BaseAgent): 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. + 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