diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 692ebcfd06..120f0d90b0 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -335,7 +335,10 @@ A frontend can then hydrate the latest stored snapshot for the scoped thread: Endpoint configuration requires `snapshot_scope_resolver` whenever a snapshot store is configured, including when the store is already set on a pre-wrapped `AgentFrameworkAgent` or `AgentFrameworkWorkflow`. The resolver returns -the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key. +the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key. When using +`AgentFrameworkWorkflow(workflow_factory=...)`, the same resolver also scopes the in-memory workflow cache even +without a snapshot store; provide it in multi-user deployments so two users who submit the same `threadId` do not +share a live `Workflow` instance. For hosted agents, request Shared State is also available through `AgentSession.state` during that run, whether or not snapshot persistence is configured. Request values are untrusted per-run context: they overlay ordinary restored 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 d1b184c964..5486802d64 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_endpoint.py @@ -111,7 +111,8 @@ def add_agent_framework_fastapi_endpoint( snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence is opt-in and requires an explicit Snapshot Scope resolver. snapshot_scope_resolver: Optional resolver for the application-defined Snapshot Scope. Required whenever - a snapshot store is configured because an AG-UI Thread id is not an authorization boundary. + a snapshot store is configured because an AG-UI Thread id is not an authorization boundary. Also scopes + in-memory workflow_factory instances when provided without a snapshot store. keepalive_seconds: Endpoint SSE keepalive interval in seconds. Defaults to 15. Positive values emit fixed SSE comments while the stream is open. None disables keepalive and preserves the non-keepalive response path. Keepalive comments are transport traffic and do not change AG-UI events. @@ -150,15 +151,13 @@ async def agent_endpoint(request_body: AGUIRequest) -> Response: """ try: input_data = request_body.model_dump(exclude_none=True) - snapshot_persistence_active = False + snapshot_persistence_active = _get_snapshot_store(protocol_runner) is not None if snapshot_scope_resolver is not None: snapshot_scope = snapshot_scope_resolver(request_body) if isawaitable(snapshot_scope): snapshot_scope = await snapshot_scope input_data[_APPROVAL_SCOPE_INPUT_KEY] = snapshot_scope - if _get_snapshot_store(protocol_runner) is not None: - input_data[_SNAPSHOT_SCOPE_INPUT_KEY] = snapshot_scope - snapshot_persistence_active = True + input_data[_SNAPSHOT_SCOPE_INPUT_KEY] = snapshot_scope if default_state: if snapshot_persistence_active: # Defer default application to the runner so defaults only fill keys 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 d558a65a30..27b3f4fa62 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -248,8 +248,9 @@ def __init__( self.workflow = workflow self._workflow_factory = workflow_factory # Cache keyed by (snapshot_scope, thread_id): the Snapshot Scope is the - # authorization boundary, so the same thread id under different scopes - # must never share an in-memory workflow instance. + # authorization boundary for both snapshots and in-memory workflow_factory + # instances, so the same thread id under different scopes must never share + # mutable workflow state. self._workflow_by_thread: dict[tuple[str | None, str], Workflow] = {} self.name = name if name is not None else getattr(workflow, "name", "workflow") self.description = description if description is not None else getattr(workflow, "description", "") diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index eb060c9315..303a307d91 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -40,6 +40,7 @@ from fastapi.testclient import TestClient from agent_framework_ag_ui import ( + AGUIRequest, AGUIThreadSnapshot, InMemoryAGUIThreadSnapshotStore, add_agent_framework_fastapi_endpoint, @@ -5160,3 +5161,75 @@ def factory(thread_id: str) -> Any: runner.clear_thread_workflow("thread-1") assert runner._resolve_workflow("thread-1", "tenant-b") is not workflow_b + + +async def test_workflow_factory_cache_is_scoped_by_resolver_without_snapshot_store(): + """Snapshot Scope resolver scopes live workflow_factory instances even without snapshot persistence.""" + + @executor(id="responder") + async def responder(message: Any, ctx: WorkflowContext[Any, Any]) -> None: + del message + await ctx.yield_output("Workflow response") + + created_workflows: list[Any] = [] + + def factory(thread_id: str) -> Any: + del thread_id + workflow = WorkflowBuilder(start_executor=responder).build() + created_workflows.append(workflow) + return workflow + + def resolve_scope(request: AGUIRequest) -> str: + forwarded_props = request.forwarded_props + assert forwarded_props is not None + tenant = forwarded_props["tenant"] + assert isinstance(tenant, str) + return tenant + + app = FastAPI() + runner = AgentFrameworkWorkflow(workflow_factory=factory) + add_agent_framework_fastapi_endpoint( + app, + runner, + path="/workflow", + snapshot_scope_resolver=resolve_scope, + ) + client = TestClient(app) + + response_a = client.post( + "/workflow", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "Hello tenant A"}], + "forwardedProps": {"tenant": "tenant-a"}, + }, + ) + response_b = client.post( + "/workflow", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "Hello tenant B"}], + "forwardedProps": {"tenant": "tenant-b"}, + }, + ) + response_a_again = client.post( + "/workflow", + json={ + "thread_id": "thread-1", + "messages": [{"role": "user", "content": "Hello tenant A again"}], + "forwardedProps": {"tenant": "tenant-a"}, + }, + ) + + assert response_a.status_code == 200 + assert response_b.status_code == 200 + assert response_a_again.status_code == 200 + assert len(created_workflows) == 2 + assert ( + runner._resolve_workflow("thread-1", "tenant-a") # pyright: ignore[reportPrivateUsage] + is created_workflows[0] + ) + assert ( + runner._resolve_workflow("thread-1", "tenant-b") # pyright: ignore[reportPrivateUsage] + is created_workflows[1] + )