Python: (ag-ui): Add Workflow Support, Harden Streaming Semantics, and add Dynamic Handoff Demo (#3911)

* fix Workflow.as_agent() streaming regression in ag-ui

* Address PR feedback

* workflows wip

* wip

* wip

* Workflow AG-UI demo

* Fixes for handoff workflow demo

* Fixes to workflows support in AG-UI

* Fixes

* Add headers to some demo files

* Fix comment

* Fixes for store

* Make _input_schema lazy-loaded

* fix mypy

* revert session change to handoff only for now

---------

Co-authored-by: Eduard van Valkenburg <eavanvalkenburg@users.noreply.github.com>

python/packages/core/agent_framework/_tools.py

@@ -14,6 +14,7 @@ from collections.abc import (
     Mapping,
     Sequence,
 )
+from contextlib import suppress
 from functools import partial, wraps
 from time import perf_counter, time_ns
 from typing import (
@@ -288,15 +289,18 @@ class FunctionTool(SerializationMixin):
         self.func = func
         self._instance = None  # Store the instance for bound methods
 
+        # Initialize schema cache (will be lazily populated)
+        self._input_schema_cached: dict[str, Any] | None = None
+
         # Track if schema was supplied as JSON dict (for optimization)
         if isinstance(input_model, Mapping):
             self._schema_supplied = True
-            self._input_schema: dict[str, Any] = dict(input_model)
+            self._input_schema_cached = dict(input_model)
             self.input_model: type[BaseModel] | None = None
         else:
             self._schema_supplied = False
             self.input_model = self._resolve_input_model(input_model)
-            self._input_schema = self.input_model.model_json_schema()
+            # Defer schema generation to avoid issues with forward references
         self._cached_parameters: dict[str, Any] | None = None
         self.approval_mode = approval_mode or "never_require"
         if max_invocations is not None and max_invocations < 1:
@@ -546,6 +550,19 @@ class FunctionTool(SerializationMixin):
                 self._invocation_duration_histogram.record(duration, attributes=attributes)
                 logger.info("Function duration: %fs", duration)
 
+    @property
+    def _input_schema(self) -> dict[str, Any]:
+        """Get the input schema, generating it lazily if needed."""
+        if self._input_schema_cached is None:
+            if self.input_model is not None:
+                # Try to rebuild the model in case it has forward references
+                with suppress(Exception):
+                    self.input_model.model_rebuild(force=True, raise_errors=False)
+                self._input_schema_cached = self.input_model.model_json_schema()
+            else:
+                self._input_schema_cached = {}
+        return self._input_schema_cached
+
     def parameters(self) -> dict[str, Any]:
         """Create the JSON schema of the parameters.
 

python/packages/core/agent_framework/ag_ui/__init__.py

@@ -20,10 +20,9 @@ IMPORT_PATH = "agent_framework_ag_ui"
 PACKAGE_NAME = "agent-framework-ag-ui"
 _IMPORTS = [
     "AgentFrameworkAgent",
+    "AgentFrameworkWorkflow",
     "add_agent_framework_fastapi_endpoint",
     "AGUIChatClient",
-    "AGUIEventConverter",
-    "AGUIHttpService",
 ]
 
 

python/packages/core/agent_framework/ag_ui/__init__.pyi

@@ -2,9 +2,11 @@
 
 from agent_framework_ag_ui import (
     AgentFrameworkAgent,
+    AgentFrameworkWorkflow,
     AGUIChatClient,
     AGUIEventConverter,
     AGUIHttpService,
+    __version__,
     add_agent_framework_fastapi_endpoint,
 )
 
@@ -13,5 +15,7 @@ __all__ = [
     "AGUIEventConverter",
     "AGUIHttpService",
     "AgentFrameworkAgent",
+    "AgentFrameworkWorkflow",
+    "__version__",
     "add_agent_framework_fastapi_endpoint",
 ]

python/packages/core/tests/workflow/test_full_conversation.py

@@ -362,9 +362,7 @@ async def test_run_request_with_full_history_clears_service_session_id() -> None
     """Replaying a full conversation (including function calls) via AgentExecutorRequest must
     clear service_session_id so the API does not receive both previous_response_id and the
     same function-call items in input — which would cause a 'Duplicate item' API error."""
-    tool_agent = _ToolHistoryAgent(
-        id="tool_agent", name="ToolAgent", summary_text="Done."
-    )
+    tool_agent = _ToolHistoryAgent(id="tool_agent", name="ToolAgent", summary_text="Done.")
     tool_exec = AgentExecutor(tool_agent, id="tool_agent")
 
     spy_agent = _SessionIdCapturingAgent(id="spy_agent", name="SpyAgent")
@@ -393,9 +391,7 @@ async def test_from_response_preserves_service_session_id() -> None:
     """from_response hands off a prior agent's full conversation to the next executor.
     The receiving executor's service_session_id is preserved so the API can continue
     the conversation using previous_response_id."""
-    tool_agent = _ToolHistoryAgent(
-        id="tool_agent2", name="ToolAgent", summary_text="Done."
-    )
+    tool_agent = _ToolHistoryAgent(id="tool_agent2", name="ToolAgent", summary_text="Done.")
     tool_exec = AgentExecutor(tool_agent, id="tool_agent2")
 
     spy_agent = _SessionIdCapturingAgent(id="spy_agent2", name="SpyAgent")
@@ -403,11 +399,7 @@ async def test_from_response_preserves_service_session_id() -> None:
     # Simulate a prior run on the spy executor.
     spy_exec._session.service_session_id = "resp_PREVIOUS_RUN"  # pyright: ignore[reportPrivateUsage]
 
-    wf = (
-        WorkflowBuilder(start_executor=tool_exec, output_executors=[spy_exec])
-        .add_edge(tool_exec, spy_exec)
-        .build()
-    )
+    wf = WorkflowBuilder(start_executor=tool_exec, output_executors=[spy_exec]).add_edge(tool_exec, spy_exec).build()
 
     result = await wf.run("start")
     assert result.get_outputs() is not None