Python: Fix ag-ui regressions (#2114)

* Bump ag-ui package version. Update CHANGELOG

* Fix ag-ui bugs

* Revert port test change

* Cleanup

* Intro factory funcs for samples

* Revert package ver change

python/.vscode/launch.json

@@ -16,7 +16,7 @@
             "name": "AG-UI Examples Server",
             "type": "debugpy",
             "request": "launch",
-            "module": "examples",
+            "module": "agent_framework_ag_ui_examples",
             "cwd": "${workspaceFolder}/packages/ag-ui",
             "console": "integratedTerminal",
             "justMyCode": false

python/packages/ag-ui/agent_framework_ag_ui/_events.py

@@ -85,6 +85,7 @@ class AgentFrameworkEventBridge:
         self.input_messages = input_messages or []
         self.pending_tool_calls: list[dict[str, Any]] = []  # Track tool calls for assistant message
         self.tool_results: list[dict[str, Any]] = []  # Track tool results
+        self.tool_calls_ended: set[str] = set()  # Track which tool calls have had ToolCallEndEvent emitted
 
     async def from_agent_run_update(self, update: AgentRunResponseUpdate) -> list[BaseEvent]:
         """
@@ -118,12 +119,14 @@ class AgentFrameworkEventBridge:
                         message_id=self.current_message_id,
                         role="assistant",
                     )
+                    logger.debug(f"Emitting TextMessageStartEvent with message_id={self.current_message_id}")
                     events.append(start_event)
 
                 event = TextMessageContentEvent(
                     message_id=self.current_message_id,
                     delta=content.text,
                 )
+                logger.debug(f"Emitting TextMessageContentEvent with delta: {content.text}")
                 events.append(event)
 
             elif isinstance(content, FunctionCallContent):
@@ -378,6 +381,7 @@ class AgentFrameworkEventBridge:
                     )
                     logger.info(f"Emitting ToolCallEndEvent for completed tool call '{content.call_id}'")
                     events.append(end_event)
+                    self.tool_calls_ended.add(content.call_id)  # Track that we emitted end event
 
                     # Log total StateDeltaEvent count for this tool call
                     if self.state_delta_count > 0:
@@ -617,6 +621,7 @@ class AgentFrameworkEventBridge:
                         f"Emitting ToolCallEndEvent for approval-required tool '{content.function_call.call_id}'"
                     )
                     events.append(end_event)
+                    self.tool_calls_ended.add(content.function_call.call_id)  # Track that we emitted end event
 
                 # Emit custom event for approval request
                 # Note: In AG-UI protocol, the frontend handles interrupts automatically

python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py

@@ -38,22 +38,69 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
     """
     result: list[ChatMessage] = []
     for msg in messages:
-        # Check for backend tool rendering results FIRST (may not have role field)
-        if "actionExecutionId" in msg or "actionName" in msg:
-            # Backend tool rendering - convert to FunctionResultContent
-            from agent_framework import FunctionResultContent
+        # Handle standard tool result messages early (role="tool") to preserve provider invariants
+        # This path maps AG‑UI tool messages to FunctionResultContent with the correct tool_call_id
+        role_str = msg.get("role", "user")
+        if role_str == "tool":
+            # Prefer explicit tool_call_id fields; fall back to backend fields only if necessary
+            tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
 
-            tool_call_id = msg.get("actionExecutionId", "")
+            # If no explicit tool_call_id, treat as backend tool rendering payloads where
+            # AG‑UI may send actionExecutionId/actionName. This must still map to the
+            # assistant's tool call id to satisfy provider requirements.
+            if not tool_call_id:
+                tool_call_id = msg.get("actionExecutionId") or ""
+
+            # Extract raw content text
+            result_content = msg.get("content")
+            if result_content is None:
+                result_content = msg.get("result", "")
+
+            # Distinguish approval payloads from actual tool results
+            is_approval = False
+            if isinstance(result_content, str) and result_content:
+                import json as _json
+
+                try:
+                    parsed = _json.loads(result_content)
+                    is_approval = isinstance(parsed, dict) and "accepted" in parsed
+                except Exception:
+                    is_approval = False
+
+            if is_approval:
+                # Approval responses should be treated as user messages to trigger human-in-the-loop flow
+                chat_msg = ChatMessage(
+                    role=Role.USER,
+                    contents=[TextContent(text=str(result_content))],
+                    additional_properties={"is_tool_result": True, "tool_call_id": str(tool_call_id or "")},
+                )
+                if "id" in msg:
+                    chat_msg.message_id = msg["id"]
+                result.append(chat_msg)
+                continue
+
+            chat_msg = ChatMessage(
+                role=Role.TOOL,
+                contents=[FunctionResultContent(call_id=str(tool_call_id), result=result_content)],
+            )
+            if "id" in msg:
+                chat_msg.message_id = msg["id"]
+            result.append(chat_msg)
+            continue
+
+        # Backend tool rendering payloads without an explicit role
+        # Prefer standard tool mapping above; this block only covers legacy/minimal payloads
+        if "actionExecutionId" in msg or "actionName" in msg:
+            # Prefer toolCallId if present; otherwise fall back to actionExecutionId
+            tool_call_id = msg.get("toolCallId") or msg.get("tool_call_id") or msg.get("actionExecutionId", "")
             result_content = msg.get("result", msg.get("content", ""))
 
             chat_msg = ChatMessage(
-                role=Role.TOOL,  # Tool results must be tool role
-                contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
+                role=Role.TOOL,
+                contents=[FunctionResultContent(call_id=str(tool_call_id), result=result_content)],
             )
-
             if "id" in msg:
                 chat_msg.message_id = msg["id"]
-
             result.append(chat_msg)
             continue
 
@@ -93,55 +140,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
             result.append(chat_msg)
             continue
 
-        role_str = msg.get("role", "user")
-
-        # Handle tool result messages (with role="tool")
-        if role_str == "tool":
-            # Check if this is a standard tool result (has tool_call_id or toolCallId)
-            tool_call_id = msg.get("tool_call_id") or msg.get("toolCallId")
-            result_content = msg.get("content", "")
-
-            # Distinguish between backend tool results and approval responses
-            # Approval responses have {"accepted": ...} structure
-            is_approval = False
-            if result_content:
-                import json
-
-                try:
-                    parsed_content = json.loads(result_content)
-                    is_approval = "accepted" in parsed_content
-                except (json.JSONDecodeError, TypeError):
-                    is_approval = False
-
-            # Backend tool results have non-empty content WITHOUT "accepted" field
-            if tool_call_id and result_content and not is_approval:
-                # Tool execution result - convert to FunctionResultContent with correct role
-                from agent_framework import FunctionResultContent
-
-                chat_msg = ChatMessage(
-                    role=Role.TOOL,
-                    contents=[FunctionResultContent(call_id=tool_call_id, result=result_content)],
-                )
-
-                if "id" in msg:
-                    chat_msg.message_id = msg["id"]
-
-                result.append(chat_msg)
-                continue
-            else:
-                # Human-in-the-loop approval response - mark for special handling
-                content = msg.get("content", "")
-                chat_msg = ChatMessage(
-                    role=Role.USER,  # Approval responses are user messages
-                    contents=[TextContent(text=content)],
-                    additional_properties={"is_tool_result": True, "tool_call_id": msg.get("toolCallId", "")},
-                )
-
-                if "id" in msg:
-                    chat_msg.message_id = msg["id"]
-
-                result.append(chat_msg)
-                continue
+        # No special handling required for assistant/plain messages here
 
         role = _AGUI_TO_FRAMEWORK_ROLE.get(role_str, Role.USER)
 

python/packages/ag-ui/agent_framework_ag_ui/_orchestrators.py

@@ -16,7 +16,15 @@ from ag_ui.core import (
     TextMessageEndEvent,
     TextMessageStartEvent,
 )
-from agent_framework import AgentProtocol, AgentThread, ChatAgent, TextContent
+from agent_framework import (
+    AgentProtocol,
+    AgentThread,
+    ChatAgent,
+    ChatMessage,
+    FunctionCallContent,
+    FunctionResultContent,
+    TextContent,
+)
 
 from ._utils import convert_agui_tools_to_agent_framework, generate_event_id
 
@@ -276,6 +284,98 @@ class DefaultOrchestrator(Orchestrator):
             response_format = context.agent.chat_options.response_format
         skip_text_content = response_format is not None
 
+        # Sanitizer: ensure tool results only follow assistant tool calls
+        # Also inject synthetic tool results for confirm_changes
+        def sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
+            sanitized: list[ChatMessage] = []
+            pending_tool_call_ids: set[str] | None = None
+            pending_confirm_changes_id: str | None = None
+
+            for msg in messages:
+                role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+
+                if role_value == "assistant":
+                    tool_ids = {
+                        str(content.call_id)
+                        for content in msg.contents or []
+                        if isinstance(content, FunctionCallContent) and content.call_id
+                    }
+                    # Check for confirm_changes tool call
+                    confirm_changes_call = None
+                    for content in msg.contents or []:
+                        if isinstance(content, FunctionCallContent) and content.name == "confirm_changes":
+                            confirm_changes_call = content
+                            break
+
+                    sanitized.append(msg)
+                    pending_tool_call_ids = tool_ids if tool_ids else None
+                    pending_confirm_changes_id = (
+                        str(confirm_changes_call.call_id)
+                        if confirm_changes_call and confirm_changes_call.call_id
+                        else None
+                    )
+                    continue
+
+                if role_value == "user" and pending_confirm_changes_id:
+                    # Check if this is a confirm_changes response (JSON with "accepted" field)
+                    user_text = ""
+                    for content in msg.contents or []:
+                        if isinstance(content, TextContent):
+                            user_text = content.text
+                            break
+
+                    try:
+                        parsed = json.loads(user_text)
+                        if "accepted" in parsed:
+                            # This is a confirm_changes response - inject synthetic tool result
+                            logger.info(
+                                f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}"
+                            )
+                            synthetic_result = ChatMessage(
+                                role="tool",
+                                contents=[
+                                    FunctionResultContent(
+                                        call_id=pending_confirm_changes_id,
+                                        result="Confirmed" if parsed.get("accepted") else "Rejected",
+                                    )
+                                ],
+                            )
+                            sanitized.append(synthetic_result)
+                            if pending_tool_call_ids:
+                                pending_tool_call_ids.discard(pending_confirm_changes_id)
+                            pending_confirm_changes_id = None
+                            # Don't add the user message to sanitized - it's been converted to tool result
+                            continue
+                    except (json.JSONDecodeError, KeyError) as e:
+                        # Failed to parse user message as confirm_changes response; continue normal processing
+                        logger.debug(f"Could not parse user message as confirm_changes response: {e}")
+
+                    # Not a confirm_changes response, continue normal processing
+                    sanitized.append(msg)
+                    pending_tool_call_ids = None
+                    pending_confirm_changes_id = None
+                    continue
+
+                if role_value == "tool":
+                    if not pending_tool_call_ids:
+                        continue
+                    keep = False
+                    for content in msg.contents or []:
+                        if isinstance(content, FunctionResultContent):
+                            call_id = str(content.call_id)
+                            if call_id in pending_tool_call_ids:
+                                keep = True
+                                break
+                    if keep:
+                        sanitized.append(msg)
+                    continue
+
+                sanitized.append(msg)
+                pending_tool_call_ids = None
+                pending_confirm_changes_id = None
+
+            return sanitized
+
         # Create event bridge
         event_bridge = AgentFrameworkEventBridge(
             run_id=context.run_id,
@@ -328,22 +428,151 @@ class DefaultOrchestrator(Orchestrator):
         if current_state:
             thread.metadata["current_state"] = current_state  # type: ignore[attr-defined]
 
-        # Add incoming AG-UI messages to the thread history
-        if context.messages:
-            await thread.on_new_messages(context.messages)
-
-        # Use the full incoming message batch to preserve tool-call adjacency
-        if not context.messages:
+        raw_messages = context.messages or []
+        if not raw_messages:
             logger.warning("No messages provided in AG-UI input")
             yield event_bridge.create_run_finished_event()
             return
 
+        logger.info(f"Received {len(raw_messages)} raw messages from client")
+        for i, msg in enumerate(raw_messages):
+            role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+            msg_id = getattr(msg, "message_id", None)
+            logger.info(f"  Raw message {i}: role={role}, id={msg_id}")
+            if hasattr(msg, "contents") and msg.contents:
+                for j, content in enumerate(msg.contents):
+                    content_type = type(content).__name__
+                    if isinstance(content, TextContent):
+                        logger.debug(f"    Content {j}: {content_type} - {content.text}")
+                    elif isinstance(content, FunctionCallContent):
+                        logger.debug(f"    Content {j}: {content_type} - {content.name}({content.arguments})")
+                    elif isinstance(content, FunctionResultContent):
+                        logger.debug(
+                            f"    Content {j}: {content_type} - call_id={content.call_id}, result={content.result}"
+                        )
+                    else:
+                        logger.debug(f"    Content {j}: {content_type} - {content}")
+
+        # After getting sanitized_messages, deduplicate them
+        def deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
+            """Remove duplicate messages while preserving order.
+
+            For tool results with the same call_id, prefer the one with actual data.
+            """
+            seen_keys: dict[Any, int] = {}  # key -> index in unique_messages (key can be various tuple types)
+            unique_messages: list[ChatMessage] = []
+
+            for idx, msg in enumerate(messages):
+                role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+
+                # For tool messages, use call_id as unique key
+                if role_value == "tool" and msg.contents and isinstance(msg.contents[0], FunctionResultContent):
+                    call_id = str(msg.contents[0].call_id)
+                    key: Any = (role_value, call_id)
+
+                    # Check if we already have this tool result
+                    if key in seen_keys:
+                        existing_idx = seen_keys[key]
+                        existing_msg = unique_messages[existing_idx]
+
+                        # Compare results - prefer non-empty over empty
+                        existing_result = None
+                        if existing_msg.contents and isinstance(existing_msg.contents[0], FunctionResultContent):
+                            existing_result = existing_msg.contents[0].result
+                        new_result = msg.contents[0].result
+
+                        # Replace if existing is empty/None and new has data
+                        if (not existing_result or existing_result == "") and new_result:
+                            logger.info(
+                                f"Replacing empty tool result at index {existing_idx} with data from index {idx}"
+                            )
+                            unique_messages[existing_idx] = msg
+                        else:
+                            logger.info(f"Skipping duplicate tool result at index {idx}: call_id={call_id}")
+                        continue
+
+                    seen_keys[key] = len(unique_messages)
+                    unique_messages.append(msg)
+
+                elif (
+                    role_value == "assistant"
+                    and msg.contents
+                    and any(isinstance(c, FunctionCallContent) for c in msg.contents)
+                ):
+                    # For assistant messages with tool_calls, use the tool call IDs
+                    tool_call_ids = tuple(
+                        sorted(str(c.call_id) for c in msg.contents if isinstance(c, FunctionCallContent) and c.call_id)
+                    )
+                    key = (role_value, tool_call_ids)
+
+                    if key in seen_keys:
+                        logger.info(f"Skipping duplicate assistant tool call at index {idx}")
+                        continue
+
+                    seen_keys[key] = len(unique_messages)
+                    unique_messages.append(msg)
+
+                else:
+                    # For other messages (system, user, assistant without tools), hash the content
+                    content_str = str([str(c) for c in msg.contents]) if msg.contents else ""
+                    key = (role_value, hash(content_str))
+
+                    if key in seen_keys:
+                        logger.info(f"Skipping duplicate message at index {idx}: role={role_value}")
+                        continue
+
+                    seen_keys[key] = len(unique_messages)
+                    unique_messages.append(msg)
+
+            return unique_messages
+
+        # Then use it:
+        sanitized_messages = sanitize_tool_history(raw_messages)
+        provider_messages = deduplicate_messages(sanitized_messages)
+
+        if not provider_messages:
+            logger.info("No provider-eligible messages after filtering; finishing run without invoking agent.")
+            yield event_bridge.create_run_finished_event()
+            return
+
+        logger.info(f"Processing {len(provider_messages)} provider messages after sanitization/deduplication")
+        for i, msg in enumerate(provider_messages):
+            role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+            logger.info(f"  Message {i}: role={role}")
+            if hasattr(msg, "contents") and msg.contents:
+                for j, content in enumerate(msg.contents):
+                    content_type = type(content).__name__
+                    if isinstance(content, TextContent):
+                        logger.info(f"    Content {j}: {content_type} - {content.text}")
+                    elif isinstance(content, FunctionCallContent):
+                        logger.info(f"    Content {j}: {content_type} - {content.name}({content.arguments})")
+                    elif isinstance(content, FunctionResultContent):
+                        logger.info(
+                            f"    Content {j}: {content_type} - call_id={content.call_id}, result={content.result}"
+                        )
+                    else:
+                        logger.info(f"    Content {j}: {content_type} - {content}")
+
+        # NOTE: For AG-UI, the client sends the full conversation history on each request.
+        # We should NOT add to thread.on_new_messages() as that would cause duplication.
+        # Instead, we pass messages directly to the agent via messages_to_run.
+
         # Inject current state as system message context if we have state
         messages_to_run: list[Any] = []
-        if current_state and context.config.state_schema:
-            state_json = json.dumps(current_state, indent=2)
-            from agent_framework import ChatMessage
 
+        conversation_has_tool_calls = False
+        logger.debug(f"Checking {len(provider_messages)} provider messages for tool calls")
+        for i, msg in enumerate(provider_messages):
+            logger.debug(
+                f"  Message {i}: role={msg.role.value}, contents={len(msg.contents) if hasattr(msg, 'contents') and msg.contents else 0}"
+            )
+        for msg in provider_messages:
+            if msg.role.value == "assistant" and hasattr(msg, "contents") and msg.contents:
+                if any(isinstance(content, FunctionCallContent) for content in msg.contents):
+                    conversation_has_tool_calls = True
+                    break
+        if current_state and context.config.state_schema and not conversation_has_tool_calls:
+            state_json = json.dumps(current_state, indent=2)
             state_context_msg = ChatMessage(
                 role="system",
                 contents=[
@@ -359,9 +588,9 @@ Never replace existing data - always append or merge."""
             )
             messages_to_run.append(state_context_msg)
 
-        # Preserve order from client to satisfy provider constraints (assistant tool_calls must
-        # immediately precede tool result messages). Using the full batch avoids reordering.
-        messages_to_run.extend(context.messages)
+        # Add all provider messages to messages_to_run
+        # AG-UI sends full conversation history on each request, so we pass it directly to the agent
+        messages_to_run.extend(provider_messages)
 
         # Handle client tools for hybrid execution
         # Client sends tool metadata, server merges with its own tools.
@@ -370,11 +599,23 @@ Never replace existing data - always append or merge."""
         from agent_framework import BaseChatClient
 
         client_tools = convert_agui_tools_to_agent_framework(context.input_data.get("tools"))
+        logger.info(f"[TOOLS] Client sent {len(client_tools) if client_tools else 0} tools")
+        if client_tools:
+            for tool in client_tools:
+                tool_name = getattr(tool, "name", "unknown")
+                declaration_only = getattr(tool, "declaration_only", None)
+                logger.info(f"[TOOLS]   - Client tool: {tool_name}, declaration_only={declaration_only}")
 
         # Extract server tools - use type narrowing when possible
         server_tools: list[Any] = []
         if isinstance(context.agent, ChatAgent):
-            server_tools = context.agent.chat_options.tools or []
+            tools_from_agent = context.agent.chat_options.tools
+            server_tools = list(tools_from_agent) if tools_from_agent else []
+            logger.info(f"[TOOLS] Agent has {len(server_tools)} configured tools")
+            for tool in server_tools:
+                tool_name = getattr(tool, "name", "unknown")
+                approval_mode = getattr(tool, "approval_mode", None)
+                logger.info(f"[TOOLS]   - {tool_name}: approval_mode={approval_mode}")
         else:
             # AgentProtocol allows duck-typed implementations - fallback to attribute access
             # This supports test mocks and custom agent implementations
@@ -412,15 +653,37 @@ Never replace existing data - always append or merge."""
                 except AttributeError:
                     pass
 
-        combined_tools: list[Any] = []
-        if server_tools:
-            combined_tools.extend(server_tools)
+        # For tools parameter: only pass if we have client tools to add
+        # If we pass tools=, it overrides the agent's configured tools and loses metadata like approval_mode
+        # So only pass tools when we need to add client tools on top of server tools
+        # IMPORTANT: Don't include client tools that duplicate server tools (same name)
+        tools_param = None
         if client_tools:
-            combined_tools.extend(client_tools)
+            # Get server tool names
+            server_tool_names = {getattr(tool, "name", None) for tool in server_tools}
+
+            # Filter out client tools that duplicate server tools
+            unique_client_tools = [
+                tool for tool in client_tools if getattr(tool, "name", None) not in server_tool_names
+            ]
+
+            if unique_client_tools:
+                combined_tools: list[Any] = []
+                if server_tools:
+                    combined_tools.extend(server_tools)
+                combined_tools.extend(unique_client_tools)
+                tools_param = combined_tools
+                logger.info(
+                    f"[TOOLS] Passing tools= parameter with {len(combined_tools)} tools ({len(server_tools)} server + {len(unique_client_tools)} unique client)"
+                )
+            else:
+                logger.info("[TOOLS] All client tools duplicate server tools - not passing tools= parameter")
+        else:
+            logger.info("[TOOLS] No client tools - not passing tools= parameter (using agent's configured tools)")
 
         # Collect all updates to get the final structured output
         all_updates: list[Any] = []
-        async for update in context.agent.run_stream(messages_to_run, thread=thread, tools=combined_tools or None):
+        async for update in context.agent.run_stream(messages_to_run, thread=thread, tools=tools_param):
             all_updates.append(update)
             events = await event_bridge.from_agent_run_update(update)
             for event in events:
@@ -432,6 +695,27 @@ Never replace existing data - always append or merge."""
             yield event_bridge.create_run_finished_event()
             return
 
+        # Check if there are pending tool calls (declaration-only tools that weren't executed)
+        # These need ToolCallEndEvent to signal the client to execute them
+        # Only emit for tool calls that haven't already had ToolCallEndEvent emitted
+        # (approval-required tools already had their end event emitted)
+        if event_bridge.pending_tool_calls:
+            pending_without_end = [
+                tc for tc in event_bridge.pending_tool_calls if tc.get("id") not in event_bridge.tool_calls_ended
+            ]
+            if pending_without_end:
+                logger.info(
+                    f"Found {len(pending_without_end)} pending tool calls without end event - emitting ToolCallEndEvent"
+                )
+                for tool_call in pending_without_end:
+                    tool_call_id = tool_call.get("id")
+                    if tool_call_id:
+                        from ag_ui.core import ToolCallEndEvent
+
+                        end_event = ToolCallEndEvent(tool_call_id=tool_call_id)
+                        logger.info(f"Emitting ToolCallEndEvent for declaration-only tool call '{tool_call_id}'")
+                        yield end_event
+
         # After streaming completes, check if agent has response_format and extract structured output
         if all_updates and response_format:
             from agent_framework import AgentRunResponse

python/packages/ag-ui/agent_framework_ag_ui_examples/README.md

@@ -10,11 +10,37 @@ pip install agent-framework-ag-ui
 
 ## Quick Start
 
+### Using Example Agents with Any Chat Client
+
+All example agents are factory functions that accept any `ChatClientProtocol`-compatible chat client:
+
+```python
+from fastapi import FastAPI
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.openai import OpenAIChatClient
+from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework_ag_ui_examples.agents import simple_agent, weather_agent
+
+app = FastAPI()
+
+# Option 1: Use Azure OpenAI
+azure_client = AzureOpenAIChatClient(model_id="gpt-4")
+add_agent_framework_fastapi_endpoint(app, simple_agent(azure_client), "/chat")
+
+# Option 2: Use OpenAI
+openai_client = OpenAIChatClient(model_id="gpt-4o")
+add_agent_framework_fastapi_endpoint(app, weather_agent(openai_client), "/weather")
+
+# Run with: uvicorn main:app --reload
+```
+
+### Creating Your Own Agent
+
 ```python
 from fastapi import FastAPI
 from agent_framework import ChatAgent
 from agent_framework.azure import AzureOpenAIChatClient
-from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
 
 # Create your agent
 agent = ChatAgent(
@@ -44,38 +70,97 @@ This integration supports all 7 AG-UI features:
 
 ## Examples
 
-Complete examples for all features are in the `examples/` directory:
+All example agents are implemented as **factory functions** that accept any chat client implementing `ChatClientProtocol`. This provides maximum flexibility to use Azure OpenAI, OpenAI, Anthropic, or any custom chat client implementation.
 
-- `examples/agents/simple_agent.py` - Basic agentic chat
-- `examples/agents/weather_agent.py` - Backend tool rendering
-- `examples/agents/task_planner_agent.py` - Human in the loop with approvals
-- `examples/agents/research_assistant_agent.py` - Agentic generative UI
-- `examples/agents/ui_generator_agent.py` - Tool-based generative UI
-- `examples/agents/recipe_agent.py` - Shared state management
-- `examples/agents/document_writer_agent.py` - Predictive state updates
-- `examples/server/main.py` - FastAPI server with all endpoints
+### Available Example Agents
 
-Run the example server:
+Complete examples for all AG-UI features are available:
 
-```bash
-cd examples/server
-uvicorn main:app --reload
+- `simple_agent(chat_client)` - Basic agentic chat (Feature 1)
+- `weather_agent(chat_client)` - Backend tool rendering (Feature 2)
+- `human_in_the_loop_agent(chat_client)` - Human-in-the-loop with step customization (Feature 3)
+- `task_steps_agent_wrapped(chat_client)` - Agentic generative UI with step execution (Feature 4)
+- `ui_generator_agent(chat_client)` - Tool-based generative UI (Feature 5)
+- `recipe_agent(chat_client)` - Shared state management (Feature 6)
+- `document_writer_agent(chat_client)` - Predictive state updates (Feature 7)
+- `research_assistant_agent(chat_client)` - Research with progress events
+- `task_planner_agent(chat_client)` - Task planning with approvals
+
+### Using Example Agents
+
+```python
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework.openai import OpenAIChatClient
+from agent_framework_ag_ui_examples.agents import (
+    simple_agent,
+    weather_agent,
+    recipe_agent,
+)
+
+# Create a chat client (use any ChatClientProtocol implementation)
+azure_client = AzureOpenAIChatClient(model_id="gpt-4")
+openai_client = OpenAIChatClient(model_id="gpt-4o")
+
+# Create agent instances by calling the factory functions
+agent1 = simple_agent(azure_client)
+agent2 = weather_agent(openai_client)
+agent3 = recipe_agent(azure_client)
 ```
 
-To enable debug logging:
+### Running the Example Server
+
+The example server demonstrates all 7 AG-UI features:
 
 ```bash
-ENABLE_DEBUG_LOGGING=1 uvicorn main:app --reload
+# Install the package
+pip install agent-framework-ag-ui
+
+# Run the example server
+python -m agent_framework_ag_ui_examples
+
+# Or with debug logging
+ENABLE_DEBUG_LOGGING=1 python -m agent_framework_ag_ui_examples
 ```
 
 The server exposes endpoints at:
-- `/agentic_chat`
-- `/backend_tool_rendering`
-- `/human_in_the_loop`
-- `/agentic_generative_ui`
-- `/tool_based_generative_ui`
-- `/shared_state`
-- `/predictive_state_updates`
+- `/agentic_chat` - Simple chat with `simple_agent`
+- `/backend_tool_rendering` - Weather tools with `weather_agent`
+- `/human_in_the_loop` - Step approval with `human_in_the_loop_agent`
+- `/agentic_generative_ui` - Task steps with `task_steps_agent_wrapped`
+- `/tool_based_generative_ui` - Custom UI components with `ui_generator_agent`
+- `/shared_state` - Recipe management with `recipe_agent`
+- `/predictive_state_updates` - Document writing with `document_writer_agent`
+
+### Complete FastAPI Example
+
+```python
+from fastapi import FastAPI
+from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+from agent_framework_ag_ui_examples.agents import (
+    simple_agent,
+    weather_agent,
+    human_in_the_loop_agent,
+    task_steps_agent_wrapped,
+    ui_generator_agent,
+    recipe_agent,
+    document_writer_agent,
+)
+
+app = FastAPI(title="AG-UI Examples")
+
+# Create a chat client (shared across all agents, or create individual ones)
+chat_client = AzureOpenAIChatClient(model_id="gpt-4")
+
+# Add all example endpoints
+add_agent_framework_fastapi_endpoint(app, simple_agent(chat_client), "/agentic_chat")
+add_agent_framework_fastapi_endpoint(app, weather_agent(chat_client), "/backend_tool_rendering")
+add_agent_framework_fastapi_endpoint(app, human_in_the_loop_agent(chat_client), "/human_in_the_loop")
+add_agent_framework_fastapi_endpoint(app, task_steps_agent_wrapped(chat_client), "/agentic_generative_ui")  # type: ignore[arg-type]
+add_agent_framework_fastapi_endpoint(app, ui_generator_agent(chat_client), "/tool_based_generative_ui")
+add_agent_framework_fastapi_endpoint(app, recipe_agent(chat_client), "/shared_state")
+add_agent_framework_fastapi_endpoint(app, document_writer_agent(chat_client), "/predictive_state_updates")
+```
 
 ## Architecture
 
@@ -97,6 +182,48 @@ The package uses a clean, orchestrator-based architecture:
 
 ## Advanced Usage
 
+### Creating Custom Agent Factories
+
+You can create your own agent factories following the same pattern as the examples:
+
+```python
+from agent_framework import ChatAgent, ai_function
+from agent_framework._clients import ChatClientProtocol
+from agent_framework_ag_ui import AgentFrameworkAgent
+
+@ai_function
+def my_tool(param: str) -> str:
+    """My custom tool."""
+    return f"Result: {param}"
+
+def my_custom_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+    """Create a custom agent with the specified chat client.
+    
+    Args:
+        chat_client: The chat client to use for the agent
+        
+    Returns:
+        A configured AgentFrameworkAgent instance
+    """
+    agent = ChatAgent(
+        name="my_custom_agent",
+        instructions="Custom instructions here",
+        chat_client=chat_client,
+        tools=[my_tool],
+    )
+    
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="MyCustomAgent",
+        description="My custom agent description",
+    )
+
+# Use it
+from agent_framework.azure import AzureOpenAIChatClient
+chat_client = AzureOpenAIChatClient()
+agent = my_custom_agent(chat_client)
+```
+
 ### Shared State
 
 State is injected as system messages and updated via predictive state updates:

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py

@@ -6,7 +6,7 @@ from .document_writer_agent import document_writer_agent
 from .human_in_the_loop_agent import human_in_the_loop_agent
 from .recipe_agent import recipe_agent
 from .research_assistant_agent import research_assistant_agent
-from .simple_agent import agent as simple_agent
+from .simple_agent import simple_agent
 from .task_planner_agent import task_planner_agent
 from .task_steps_agent import task_steps_agent_wrapped
 from .ui_generator_agent import ui_generator_agent

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/document_writer_agent.py

@@ -3,7 +3,7 @@
 """Example agent demonstrating predictive state updates with document writing."""
 
 from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 
 from agent_framework_ag_ui import AgentFrameworkAgent, DocumentWriterConfirmationStrategy
 
@@ -28,31 +28,43 @@ def write_document_local(document: str) -> str:
     return "Document written."
 
 
-agent = ChatAgent(
-    name="document_writer",
-    instructions=(
-        "You are a helpful assistant for writing documents. "
-        "To write the document, you MUST use the write_document_local tool. "
-        "You MUST write the full document, even when changing only a few words. "
-        "When you wrote the document, DO NOT repeat it as a message. "
-        "Just briefly summarize the changes you made. 2 sentences max. "
-        "\n\n"
-        "The current state of the document will be provided to you. "
-        "When editing, make minimal changes - do not change every word unless requested."
-    ),
-    chat_client=AzureOpenAIChatClient(),
-    tools=[write_document_local],
+_DOCUMENT_WRITER_INSTRUCTIONS = (
+    "You are a helpful assistant for writing documents. "
+    "To write the document, you MUST use the write_document_local tool. "
+    "You MUST write the full document, even when changing only a few words. "
+    "When you wrote the document, DO NOT repeat it as a message. "
+    "Just briefly summarize the changes you made. 2 sentences max. "
+    "\n\n"
+    "The current state of the document will be provided to you. "
+    "When editing, make minimal changes - do not change every word unless requested."
 )
 
-document_writer_agent = AgentFrameworkAgent(
-    agent=agent,
-    name="DocumentWriter",
-    description="Writes and edits documents with predictive state updates",
-    state_schema={
-        "document": {"type": "string", "description": "The current document content"},
-    },
-    predict_state_config={
-        "document": {"tool": "write_document_local", "tool_argument": "document"},
-    },
-    confirmation_strategy=DocumentWriterConfirmationStrategy(),
-)
+
+def document_writer_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+    """Create a document writer agent with predictive state updates.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with document writing capabilities
+    """
+    agent = ChatAgent(
+        name="document_writer",
+        instructions=_DOCUMENT_WRITER_INSTRUCTIONS,
+        chat_client=chat_client,
+        tools=[write_document_local],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="DocumentWriter",
+        description="Writes and edits documents with predictive state updates",
+        state_schema={
+            "document": {"type": "string", "description": "The current document content"},
+        },
+        predict_state_config={
+            "document": {"tool": "write_document_local", "tool_argument": "document"},
+        },
+        confirmation_strategy=DocumentWriterConfirmationStrategy(),
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/human_in_the_loop_agent.py

@@ -5,7 +5,7 @@
 from enum import Enum
 
 from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 from pydantic import BaseModel, Field
 
 
@@ -43,10 +43,18 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
     return f"Generated {len(steps)} execution steps for the task."
 
 
-# Create the human-in-the-loop agent using tool-based approach for predictive state
-human_in_the_loop_agent = ChatAgent(
-    name="human_in_the_loop_agent",
-    instructions="""You are a helpful assistant that can perform any task by breaking it down into steps.
+def human_in_the_loop_agent(chat_client: ChatClientProtocol) -> ChatAgent:
+    """Create a human-in-the-loop agent using tool-based approach for predictive state.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured ChatAgent instance with human-in-the-loop capabilities
+    """
+    return ChatAgent(
+        name="human_in_the_loop_agent",
+        instructions="""You are a helpful assistant that can perform any task by breaking it down into steps.
 
     When asked to perform a task, you MUST call the `generate_task_steps` function with the proper
     number of steps per the request.
@@ -71,6 +79,6 @@ human_in_the_loop_agent = ChatAgent(
     After calling the function, provide a brief acknowledgment like:
     "I've created a plan with 10 steps. You can customize which steps to enable before I proceed."
     """,
-    chat_client=AzureOpenAIChatClient(),
-    tools=[generate_task_steps],
-)
+        chat_client=chat_client,
+        tools=[generate_task_steps],
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/recipe_agent.py

@@ -5,7 +5,7 @@
 from enum import Enum
 
 from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 from pydantic import BaseModel, Field
 
 from agent_framework_ag_ui import AgentFrameworkAgent, RecipeConfirmationStrategy
@@ -67,10 +67,7 @@ def update_recipe(recipe: Recipe) -> str:
     return "Recipe updated."
 
 
-# Create the recipe agent using tool-based approach for streaming
-agent = ChatAgent(
-    name="recipe_agent",
-    instructions="""You are a helpful recipe assistant that creates and modifies recipes.
+_RECIPE_INSTRUCTIONS = """You are a helpful recipe assistant that creates and modifies recipes.
 
     CRITICAL RULES:
     1. You will receive the current recipe state in the system context
@@ -103,20 +100,34 @@ agent = ChatAgent(
     - Add aromatics: garlic, shallots
     - Add finishing touches: lemon zest, fresh parsley
     - Make instructions more detailed and professional
-    """,
-    chat_client=AzureOpenAIChatClient(),
-    tools=[update_recipe],
-)
+    """
 
-recipe_agent = AgentFrameworkAgent(
-    agent=agent,
-    name="RecipeAgent",
-    description="Creates and modifies recipes with streaming state updates",
-    state_schema={
-        "recipe": {"type": "object", "description": "The current recipe"},
-    },
-    predict_state_config={
-        "recipe": {"tool": "update_recipe", "tool_argument": "recipe"},
-    },
-    confirmation_strategy=RecipeConfirmationStrategy(),
-)
+
+def recipe_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+    """Create a recipe agent with streaming state updates.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with recipe management
+    """
+    agent = ChatAgent(
+        name="recipe_agent",
+        instructions=_RECIPE_INSTRUCTIONS,
+        chat_client=chat_client,
+        tools=[update_recipe],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="RecipeAgent",
+        description="Creates and modifies recipes with streaming state updates",
+        state_schema={
+            "recipe": {"type": "object", "description": "The current recipe"},
+        },
+        predict_state_config={
+            "recipe": {"tool": "update_recipe", "tool_argument": "recipe"},
+        },
+        confirmation_strategy=RecipeConfirmationStrategy(),
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/research_assistant_agent.py

@@ -5,7 +5,7 @@
 import asyncio
 
 from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 
 from agent_framework_ag_ui import AgentFrameworkAgent
 
@@ -82,19 +82,31 @@ async def analyze_data(dataset: str) -> str:
     return f"Analysis of '{dataset}':\n" + "\n".join(insights)
 
 
-agent = ChatAgent(
-    name="research_assistant",
-    instructions=(
-        "You are a research and analysis assistant. "
-        "You can research topics, create presentations, and analyze data. "
-        "Use the available tools to help users with their research needs."
-    ),
-    chat_client=AzureOpenAIChatClient(),
-    tools=[research_topic, create_presentation, analyze_data],
+_RESEARCH_ASSISTANT_INSTRUCTIONS = (
+    "You are a research and analysis assistant. "
+    "You can research topics, create presentations, and analyze data. "
+    "Use the available tools to help users with their research needs."
 )
 
-research_assistant_agent = AgentFrameworkAgent(
-    agent=agent,
-    name="ResearchAssistant",
-    description="Research assistant that emits progress events during task execution",
-)
+
+def research_assistant_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+    """Create a research assistant agent with progress events.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with research capabilities
+    """
+    agent = ChatAgent(
+        name="research_assistant",
+        instructions=_RESEARCH_ASSISTANT_INSTRUCTIONS,
+        chat_client=chat_client,
+        tools=[research_topic, create_presentation, analyze_data],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="ResearchAssistant",
+        description="Research assistant that emits progress events during task execution",
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/simple_agent.py

@@ -3,11 +3,20 @@
 """Simple agentic chat example (Feature 1: Agentic Chat)."""
 
 from agent_framework import ChatAgent
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 
-# Create a simple chat agent
-agent = ChatAgent(
-    name="simple_chat_agent",
-    instructions="You are a helpful assistant. Be concise and friendly.",
-    chat_client=AzureOpenAIChatClient(),
-)
+
+def simple_agent(chat_client: ChatClientProtocol) -> ChatAgent:
+    """Create a simple chat agent.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured ChatAgent instance
+    """
+    return ChatAgent(
+        name="simple_chat_agent",
+        instructions="You are a helpful assistant. Be concise and friendly.",
+        chat_client=chat_client,
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_planner_agent.py

@@ -3,7 +3,7 @@
 """Example agent demonstrating human-in-the-loop with function approvals."""
 
 from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 
 from agent_framework_ag_ui import AgentFrameworkAgent, TaskPlannerConfirmationStrategy
 
@@ -54,20 +54,32 @@ def book_meeting_room(room_name: str, date: str, start_time: str, end_time: str)
     return f"Meeting room '{room_name}' booked for {date} from {start_time} to {end_time}"
 
 
-agent = ChatAgent(
-    name="task_planner",
-    instructions=(
-        "You are a helpful assistant that plans and executes tasks. "
-        "You have access to calendar, email, and meeting room booking functions. "
-        "All of these actions require user approval before execution."
-    ),
-    chat_client=AzureOpenAIChatClient(),
-    tools=[create_calendar_event, send_email, book_meeting_room],
+_TASK_PLANNER_INSTRUCTIONS = (
+    "You are a helpful assistant that plans and executes tasks. "
+    "You have access to calendar, email, and meeting room booking functions. "
+    "All of these actions require user approval before execution."
 )
 
-task_planner_agent = AgentFrameworkAgent(
-    agent=agent,
-    name="TaskPlanner",
-    description="Plans and executes tasks with user approval",
-    confirmation_strategy=TaskPlannerConfirmationStrategy(),
-)
+
+def task_planner_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+    """Create a task planner agent with user approval for actions.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with task planning capabilities
+    """
+    agent = ChatAgent(
+        name="task_planner",
+        instructions=_TASK_PLANNER_INSTRUCTIONS,
+        chat_client=chat_client,
+        tools=[create_calendar_event, send_email, book_meeting_room],
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="TaskPlanner",
+        description="Plans and executes tasks with user approval",
+        confirmation_strategy=TaskPlannerConfirmationStrategy(),
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py

@@ -19,7 +19,7 @@ from ag_ui.core import (
     ToolCallStartEvent,
 )
 from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 from pydantic import BaseModel, Field
 
 from agent_framework_ag_ui import AgentFrameworkAgent
@@ -54,10 +54,18 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
     return "Steps generated."
 
 
-# Create the task steps agent using tool-based approach for streaming
-agent = ChatAgent(
-    name="task_steps_agent",
-    instructions="""You are a helpful assistant that breaks down tasks into actionable steps.
+def _create_task_steps_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+    """Create the task steps agent using tool-based approach for streaming.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance
+    """
+    agent = ChatAgent(
+        name="task_steps_agent",
+        instructions="""You are a helpful assistant that breaks down tasks into actionable steps.
 
     When asked to perform a task, you MUST:
     1. Use the generate_task_steps tool to create the steps
@@ -75,25 +83,25 @@ agent = ChatAgent(
     - "Installing platform"
     - "Adding finishing touches"
     """,
-    chat_client=AzureOpenAIChatClient(),
-    tools=[generate_task_steps],
-)
+        chat_client=chat_client,
+        tools=[generate_task_steps],
+    )
 
-task_steps_agent = AgentFrameworkAgent(
-    agent=agent,
-    name="TaskStepsAgent",
-    description="Generates task steps with streaming state updates",
-    state_schema={
-        "steps": {"type": "array", "description": "The list of task steps"},
-    },
-    predict_state_config={
-        "steps": {
-            "tool": "generate_task_steps",
-            "tool_argument": "steps",
-        }
-    },
-    require_confirmation=False,  # Agentic generative UI updates automatically without confirmation
-)
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="TaskStepsAgent",
+        description="Generates task steps with streaming state updates",
+        state_schema={
+            "steps": {"type": "array", "description": "The list of task steps"},
+        },
+        predict_state_config={
+            "steps": {
+                "tool": "generate_task_steps",
+                "tool_argument": "steps",
+            }
+        },
+        require_confirmation=False,  # Agentic generative UI updates automatically without confirmation
+    )
 
 
 # Wrap the agent's run method to add step execution simulation
@@ -131,7 +139,7 @@ class TaskStepsAgentWithExecution:
         logger.info("TaskStepsAgentWithExecution.run_agent() called - wrapper is active")
 
         # First, run the base agent to generate the plan - buffer text messages
-        final_state: dict[str, Any] | None = None
+        final_state: dict[str, Any] = {}
         run_finished_event: Any = None
         tool_call_id: str | None = None
         buffered_text_events: list[Any] = []  # Buffer text from first LLM call
@@ -142,9 +150,20 @@ class TaskStepsAgentWithExecution:
 
             match event:
                 case StateSnapshotEvent(snapshot=snapshot):
-                    final_state = snapshot
+                    final_state = snapshot.copy() if snapshot else {}
                     logger.info(f"Captured STATE_SNAPSHOT event with state: {final_state}")
                     yield event
+                case StateDeltaEvent(delta=delta):
+                    # Apply state delta to final_state
+                    if delta:
+                        for patch in delta:
+                            if patch.get("op") == "replace" and patch.get("path") == "/steps":
+                                final_state["steps"] = patch.get("value", [])
+                                logger.info(
+                                    f"Applied STATE_DELTA: updated steps to {len(final_state.get('steps', []))} items"
+                                )
+                    logger.info(f"Yielding event immediately: {event_type_str}")
+                    yield event
                 case RunFinishedEvent():
                     run_finished_event = event
                     logger.info("Captured RUN_FINISHED event - will send after step execution and summary")
@@ -314,5 +333,14 @@ class TaskStepsAgentWithExecution:
             yield run_finished_event
 
 
-# Export the wrapped agent
-task_steps_agent_wrapped = TaskStepsAgentWithExecution(task_steps_agent)
+def task_steps_agent_wrapped(chat_client: ChatClientProtocol) -> TaskStepsAgentWithExecution:
+    """Create a task steps agent with execution simulation.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A wrapped agent instance with step execution simulation
+    """
+    base_agent = _create_task_steps_agent(chat_client)
+    return TaskStepsAgentWithExecution(base_agent)

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py

@@ -4,23 +4,39 @@
 
 from typing import Any
 
-from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework import AIFunction, ChatAgent
+from agent_framework._clients import ChatClientProtocol
 
 from agent_framework_ag_ui import AgentFrameworkAgent
 
-
-@ai_function
-def generate_haiku(english: list[str], japanese: list[str], image_name: str | None, gradient: str) -> str:
-    """Generate a haiku with image and gradient background (FRONTEND_RENDER).
+# Declaration-only tools (func=None) - actual rendering happens on the client side
+generate_haiku = AIFunction[Any, str](
+    name="generate_haiku",
+    description="""Generate a haiku with image and gradient background (FRONTEND_RENDER).
 
     This tool generates UI for displaying a haiku with an image and gradient background.
-    The frontend should render this as a custom haiku component.
-
-    Args:
-        english: English haiku lines (exactly 3 lines)
-        japanese: Japanese haiku lines (exactly 3 lines)
-        image_name: Image filename for visual accompaniment. Must be one of:
+    The frontend should render this as a custom haiku component.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "english": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "English haiku lines (exactly 3 lines)",
+                "minItems": 3,
+                "maxItems": 3,
+            },
+            "japanese": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Japanese haiku lines (exactly 3 lines)",
+                "minItems": 3,
+                "maxItems": 3,
+            },
+            "image_name": {
+                "type": "string",
+                "description": """Image filename for visual accompaniment. Must be one of:
             - "Osaka_Castle_Turret_Stone_Wall_Pine_Trees_Daytime.jpg"
             - "Tokyo_Skyline_Night_Tokyo_Tower_Mount_Fuji_View.jpg"
             - "Itsukushima_Shrine_Miyajima_Floating_Torii_Gate_Sunset_Long_Exposure.jpg"
@@ -31,71 +47,100 @@ def generate_haiku(english: list[str], japanese: list[str], image_name: str | No
             - "Senso-ji_Temple_Asakusa_Cherry_Blossoms_Kimono_Umbrella.jpg"
             - "Cherry_Blossoms_Sakura_Night_View_City_Lights_Japan.jpg"
             - "Mount_Fuji_Lake_Reflection_Cherry_Blossoms_Sakura_Spring.jpg"
-        gradient: CSS gradient string for background (e.g., "linear-gradient(135deg, #667eea 0%, #764ba2 100%)")
+            """,
+            },
+            "gradient": {
+                "type": "string",
+                "description": 'CSS gradient string for background (e.g., "linear-gradient(135deg, #667eea 0%, #764ba2 100%)")',
+            },
+        },
+        "required": ["english", "japanese", "image_name", "gradient"],
+    },
+)
 
-    Returns:
-        Haiku metadata for frontend rendering
-    """
-    return f"Haiku generated with image: {image_name}"
-
-
-@ai_function
-def create_chart(chart_type: str, data_points: list[dict[str, Any]], title: str) -> str:
-    """Create an interactive chart (FRONTEND_RENDER).
+create_chart = AIFunction[Any, str](
+    name="create_chart",
+    description="""Create an interactive chart (FRONTEND_RENDER).
 
     This tool creates chart specifications for frontend rendering.
-    The frontend should render this as an interactive chart component.
+    The frontend should render this as an interactive chart component.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "chart_type": {
+                "type": "string",
+                "description": "Type of chart (bar, line, pie, scatter)",
+            },
+            "data_points": {
+                "type": "array",
+                "items": {"type": "object"},
+                "description": "Data points for the chart",
+            },
+            "title": {
+                "type": "string",
+                "description": "Chart title",
+            },
+        },
+        "required": ["chart_type", "data_points", "title"],
+    },
+)
 
-    Args:
-        chart_type: Type of chart (bar, line, pie, scatter)
-        data_points: Data points for the chart
-        title: Chart title
-
-    Returns:
-        Chart specification for frontend rendering
-    """
-    return f"Chart '{title}' created with {len(data_points)} data points"
-
-
-@ai_function
-def display_timeline(events: list[dict[str, Any]], start_date: str, end_date: str) -> str:
-    """Display an interactive timeline (FRONTEND_RENDER).
+display_timeline = AIFunction[Any, str](
+    name="display_timeline",
+    description="""Display an interactive timeline (FRONTEND_RENDER).
 
     This tool creates timeline specifications for frontend rendering.
-    The frontend should render this as an interactive timeline component.
+    The frontend should render this as an interactive timeline component.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "events": {
+                "type": "array",
+                "items": {"type": "object"},
+                "description": "Events to display on the timeline",
+            },
+            "start_date": {
+                "type": "string",
+                "description": "Timeline start date",
+            },
+            "end_date": {
+                "type": "string",
+                "description": "Timeline end date",
+            },
+        },
+        "required": ["events", "start_date", "end_date"],
+    },
+)
 
-    Args:
-        events: Events to display on the timeline
-        start_date: Timeline start date
-        end_date: Timeline end date
-
-    Returns:
-        Timeline specification for frontend rendering
-    """
-    return f"Timeline created with {len(events)} events from {start_date} to {end_date}"
-
-
-@ai_function
-def show_comparison_table(items: list[dict[str, Any]], columns: list[str]) -> str:
-    """Show a comparison table (FRONTEND_RENDER).
+show_comparison_table = AIFunction[Any, str](
+    name="show_comparison_table",
+    description="""Show a comparison table (FRONTEND_RENDER).
 
     This tool creates table specifications for frontend rendering.
-    The frontend should render this as an interactive comparison table.
-
-    Args:
-        items: Items to compare
-        columns: Column names
-
-    Returns:
-        Table specification for frontend rendering
-    """
-    return f"Comparison table created with {len(items)} items and {len(columns)} columns"
+    The frontend should render this as an interactive comparison table.""",
+    func=None,  # Makes declaration_only=True so client renders the UI
+    input_model={
+        "type": "object",
+        "properties": {
+            "items": {
+                "type": "array",
+                "items": {"type": "object"},
+                "description": "Items to compare",
+            },
+            "columns": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Column names",
+            },
+        },
+        "required": ["items", "columns"],
+    },
+)
 
 
-# Create the UI generator agent using tool-based approach with forced tool usage
-agent = ChatAgent(
-    name="ui_generator",
-    instructions="""You MUST use the provided tools to generate content. Never respond with plain text descriptions.
+_UI_GENERATOR_INSTRUCTIONS = """You MUST use the provided tools to generate content. Never respond with plain text descriptions.
 
     For haiku requests:
     - Call generate_haiku tool with all 4 required parameters
@@ -105,15 +150,29 @@ agent = ChatAgent(
     - gradient: CSS gradient string
 
     For other requests, use the appropriate tool (create_chart, display_timeline, show_comparison_table).
-    """,
-    chat_client=AzureOpenAIChatClient(),
-    tools=[generate_haiku, create_chart, display_timeline, show_comparison_table],
-    # Force tool usage - the LLM MUST call a tool, cannot respond with plain text
-    chat_options={"tool_choice": "required"},
-)
+    """
 
-ui_generator_agent = AgentFrameworkAgent(
-    agent=agent,
-    name="UIGenerator",
-    description="Generates custom UI components through tool calls",
-)
+
+def ui_generator_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+    """Create a UI generator agent with frontend rendering tools.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured AgentFrameworkAgent instance with UI generation tools
+    """
+    agent = ChatAgent(
+        name="ui_generator",
+        instructions=_UI_GENERATOR_INSTRUCTIONS,
+        chat_client=chat_client,
+        tools=[generate_haiku, create_chart, display_timeline, show_comparison_table],
+        # Force tool usage - the LLM MUST call a tool, cannot respond with plain text
+        chat_options={"tool_choice": "required"},
+    )
+
+    return AgentFrameworkAgent(
+        agent=agent,
+        name="UIGenerator",
+        description="Generates custom UI components through tool calls",
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_agent.py

@@ -5,7 +5,7 @@
 from typing import Any
 
 from agent_framework import ChatAgent, ai_function
-from agent_framework.azure import AzureOpenAIChatClient
+from agent_framework._clients import ChatClientProtocol
 
 
 @ai_function
@@ -58,14 +58,22 @@ def get_forecast(location: str, days: int = 3) -> str:
     return f"{days}-day forecast for {location}:\n" + "\n".join(forecast)
 
 
-# Create the weather agent
-weather_agent = ChatAgent(
-    name="weather_agent",
-    instructions=(
-        "You are a helpful weather assistant. "
-        "Use the get_weather and get_forecast functions to help users with weather information. "
-        "Always provide friendly and informative responses."
-    ),
-    chat_client=AzureOpenAIChatClient(),
-    tools=[get_weather, get_forecast],
-)
+def weather_agent(chat_client: ChatClientProtocol) -> ChatAgent:
+    """Create a weather agent with get_weather and get_forecast tools.
+
+    Args:
+        chat_client: The chat client to use for the agent
+
+    Returns:
+        A configured ChatAgent instance with weather tools
+    """
+    return ChatAgent(
+        name="weather_agent",
+        instructions=(
+            "You are a helpful weather assistant. "
+            "Use the get_weather and get_forecast functions to help users with weather information. "
+            "Always provide friendly and informative responses."
+        ),
+        chat_client=chat_client,
+        tools=[get_weather, get_forecast],
+    )

python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/__init__.py

@@ -1,3 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-"""API endpoints for AG-UI examples."""

python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/backend_tool_rendering.py

@@ -2,6 +2,7 @@
 
 """Backend tool rendering endpoint."""
 
+from agent_framework.azure import AzureOpenAIChatClient
 from fastapi import FastAPI
 
 from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
@@ -15,8 +16,11 @@ def register_backend_tool_rendering(app: FastAPI) -> None:
     Args:
         app: The FastAPI application.
     """
+    # Create a chat client and call the factory function
+    chat_client = AzureOpenAIChatClient()
+
     add_agent_framework_fastapi_endpoint(
         app,
-        weather_agent,
+        weather_agent(chat_client),
         "/backend_tool_rendering",
     )

python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py

@@ -6,6 +6,7 @@ import logging
 import os
 
 import uvicorn
+from agent_framework.azure import AzureOpenAIChatClient
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
@@ -14,8 +15,8 @@ from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
 from ..agents.document_writer_agent import document_writer_agent
 from ..agents.human_in_the_loop_agent import human_in_the_loop_agent
 from ..agents.recipe_agent import recipe_agent
-from ..agents.simple_agent import agent as simple_agent
-from ..agents.task_steps_agent import task_steps_agent_wrapped as task_steps_agent  # Custom wrapper
+from ..agents.simple_agent import simple_agent
+from ..agents.task_steps_agent import task_steps_agent_wrapped
 from ..agents.ui_generator_agent import ui_generator_agent
 from ..agents.weather_agent import weather_agent
 
@@ -58,38 +59,42 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Create a shared chat client for all agents
+# You can use different chat clients for different agents if needed
+chat_client = AzureOpenAIChatClient()
+
 # Agentic Chat - basic chat agent
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=simple_agent,
+    agent=simple_agent(chat_client),
     path="/agentic_chat",
 )
 
 # Backend Tool Rendering - agent with tools
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=weather_agent,
+    agent=weather_agent(chat_client),
     path="/backend_tool_rendering",
 )
 
 # Shared State - recipe agent with structured output
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=recipe_agent,
+    agent=recipe_agent(chat_client),
     path="/shared_state",
 )
 
 # Predictive State Updates - document writer with predictive state
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=document_writer_agent,
+    agent=document_writer_agent(chat_client),
     path="/predictive_state_updates",
 )
 
 # Human in the Loop - human-in-the-loop agent with step customization
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=human_in_the_loop_agent,
+    agent=human_in_the_loop_agent(chat_client),
     path="/human_in_the_loop",
     state_schema={"steps": {"type": "array"}},
     predict_state_config={"steps": {"tool": "generate_task_steps", "tool_argument": "steps"}},
@@ -98,23 +103,26 @@ add_agent_framework_fastapi_endpoint(
 # Agentic Generative UI - task steps agent with streaming state updates
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=task_steps_agent,  # type: ignore[arg-type]
+    agent=task_steps_agent_wrapped(chat_client),  # type: ignore[arg-type]
     path="/agentic_generative_ui",
 )
 
 # Tool-based Generative UI - UI generator with frontend-rendered tools
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=ui_generator_agent,
+    agent=ui_generator_agent(chat_client),
     path="/tool_based_generative_ui",
 )
 
 
 def main():
     """Run the server."""
-    port = int(os.getenv("PORT", "8888"))
+    port = int(os.getenv("PORT", "8887"))
     host = os.getenv("HOST", "127.0.0.1")
 
+    print(f"\nAG-UI Examples Server starting on http://{host}:{port}")
+    print("Set ENABLE_DEBUG_LOGGING=1 for detailed request logging\n")
+
     # Use log_config=None to prevent uvicorn from reconfiguring logging
     # This preserves our file + console logging setup
     uvicorn.run(

python/packages/ag-ui/tests/test_agent_wrapper_comprehensive.py

@@ -505,17 +505,20 @@ async def test_error_handling_with_exception():
 
 
 async def test_json_decode_error_in_tool_result():
-    """Test handling of JSONDecodeError when parsing tool result."""
+    """Test handling of orphaned tool result - should be sanitized out."""
     from agent_framework_ag_ui import AgentFrameworkAgent
 
     class MockChatClient:
         async def get_streaming_response(self, messages, chat_options, **kwargs):
-            yield ChatResponseUpdate(contents=[TextContent(text="Fallback response")])
+            # Should not be called since orphaned tool result is dropped
+            if False:
+                yield
+            raise AssertionError("ChatClient should not be called with orphaned tool result")
 
     agent = ChatAgent(name="test_agent", instructions="Test", chat_client=MockChatClient())
     wrapper = AgentFrameworkAgent(agent=agent)
 
-    # Send invalid JSON as tool result
+    # Send invalid JSON as tool result without preceding tool call
     input_data = {
         "messages": [
             {
@@ -530,10 +533,12 @@ async def test_json_decode_error_in_tool_result():
     async for event in wrapper.run_agent(input_data):
         events.append(event)
 
-    # Should fall through to normal agent processing
+    # Orphaned tool result should be sanitized out
+    # Only run lifecycle events should be emitted, no text/tool events
     text_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
-    assert len(text_events) > 0
-    assert text_events[0].delta == "Fallback response"
+    tool_events = [e for e in events if e.type.startswith("TOOL_CALL")]
+    assert len(text_events) == 0
+    assert len(tool_events) == 0
 
 
 async def test_suppressed_summary_with_document_state():

python/packages/ag-ui/tests/test_orchestrators_coverage.py

@@ -0,0 +1,811 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Comprehensive tests for orchestrator coverage."""
+
+from collections.abc import AsyncGenerator
+from types import SimpleNamespace
+from typing import Any
+
+from agent_framework import (
+    AgentRunResponseUpdate,
+    ChatMessage,
+    TextContent,
+    ai_function,
+)
+from pydantic import BaseModel
+
+from agent_framework_ag_ui._agent import AgentConfig
+from agent_framework_ag_ui._orchestrators import (
+    DefaultOrchestrator,
+    ExecutionContext,
+    HumanInTheLoopOrchestrator,
+)
+
+
+@ai_function(approval_mode="always_require")
+def approval_tool(param: str) -> str:
+    """Tool requiring approval."""
+    return f"executed: {param}"
+
+
+class MockAgent:
+    """Mock agent for testing."""
+
+    def __init__(self, updates: list[AgentRunResponseUpdate] | None = None) -> None:
+        self.updates = updates or [AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")]
+        self.chat_options = SimpleNamespace(tools=[approval_tool], response_format=None)
+        self.chat_client = SimpleNamespace(function_invocation_configuration=None)
+        self.messages_received: list[Any] = []
+        self.tools_received: list[Any] | None = None
+
+    async def run_stream(
+        self,
+        messages: list[Any],
+        *,
+        thread: Any = None,
+        tools: list[Any] | None = None,
+    ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
+        self.messages_received = messages
+        self.tools_received = tools
+        for update in self.updates:
+            yield update
+
+
+async def test_human_in_the_loop_json_decode_error() -> None:
+    """Test HumanInTheLoopOrchestrator handles invalid JSON in tool result."""
+    orchestrator = HumanInTheLoopOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": [{"type": "text", "text": "not valid json {"}],
+            }
+        ],
+    }
+
+    messages = [
+        ChatMessage(
+            role="tool",
+            contents=[TextContent(text="not valid json {")],
+            additional_properties={"is_tool_result": True},
+        )
+    ]
+
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=MockAgent(),
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    assert orchestrator.can_handle(context)
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should emit RunErrorEvent for invalid JSON
+    error_events = [e for e in events if e.type == "RUN_ERROR"]
+    assert len(error_events) == 1
+    assert "Invalid tool result format" in error_events[0].message
+
+
+async def test_sanitize_tool_history_confirm_changes() -> None:
+    """Test sanitize_tool_history logic for confirm_changes synthetic result."""
+    from agent_framework import ChatMessage, FunctionCallContent, TextContent
+
+    # Create messages that will trigger confirm_changes synthetic result injection
+    messages = [
+        ChatMessage(
+            role="assistant",
+            contents=[
+                FunctionCallContent(
+                    name="confirm_changes",
+                    call_id="call_confirm_123",
+                    arguments='{"changes": "test"}',
+                )
+            ],
+        ),
+        ChatMessage(
+            role="user",
+            contents=[TextContent(text='{"accepted": true}')],
+        ),
+    ]
+
+    # The sanitize_tool_history function is internal to DefaultOrchestrator.run
+    # We'll test it indirectly by checking the orchestrator processes it correctly
+    orchestrator = DefaultOrchestrator()
+
+    # Use pre-constructed ChatMessage objects to bypass message adapter
+    input_data = {"messages": []}
+
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    # Override the messages property to use our pre-constructed messages
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Agent should receive synthetic tool result
+    assert len(agent.messages_received) > 0
+    tool_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
+    ]
+    assert len(tool_messages) == 1
+    assert str(tool_messages[0].contents[0].call_id) == "call_confirm_123"
+    assert tool_messages[0].contents[0].result == "Confirmed"
+
+
+async def test_sanitize_tool_history_orphaned_tool_result() -> None:
+    """Test sanitize_tool_history removes orphaned tool results."""
+    from agent_framework import ChatMessage, FunctionResultContent, TextContent
+
+    # Tool result without preceding assistant tool call
+    messages = [
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="orphan_123", result="orphaned data")],
+        ),
+        ChatMessage(
+            role="user",
+            contents=[TextContent(text="Hello")],
+        ),
+    ]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Orphaned tool result should be filtered out
+    tool_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
+    ]
+    assert len(tool_messages) == 0
+
+
+async def test_orphaned_tool_result_sanitization() -> None:
+    """Test that orphaned tool results are filtered out."""
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": [{"type": "tool_result", "tool_call_id": "orphan_123", "content": "result"}],
+            },
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Hello"}],
+            },
+        ],
+    }
+
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Orphaned tool result should be filtered, only user message remains
+    tool_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
+    ]
+    assert len(tool_messages) == 0
+
+
+async def test_deduplicate_messages_empty_tool_results() -> None:
+    """Test deduplicate_messages prefers non-empty tool results."""
+    from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent
+
+    messages = [
+        ChatMessage(
+            role="assistant",
+            contents=[FunctionCallContent(name="test_tool", call_id="call_789", arguments="{}")],
+        ),
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="call_789", result="")],
+        ),
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="call_789", result="real data")],
+        ),
+    ]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should have only one tool result with actual data
+    tool_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
+    ]
+    assert len(tool_messages) == 1
+    assert tool_messages[0].contents[0].result == "real data"
+
+
+async def test_deduplicate_messages_duplicate_assistant_tool_calls() -> None:
+    """Test deduplicate_messages removes duplicate assistant tool call messages."""
+    from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent
+
+    messages = [
+        ChatMessage(
+            role="assistant",
+            contents=[FunctionCallContent(name="test_tool", call_id="call_abc", arguments="{}")],
+        ),
+        ChatMessage(
+            role="assistant",
+            contents=[FunctionCallContent(name="test_tool", call_id="call_abc", arguments="{}")],
+        ),
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="call_abc", result="result")],
+        ),
+    ]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should have only one assistant message
+    assistant_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "assistant"
+    ]
+    assert len(assistant_messages) == 1
+
+
+async def test_deduplicate_messages_duplicate_system_messages() -> None:
+    """Test that deduplication logic is invoked for system messages."""
+    from agent_framework import ChatMessage, TextContent
+
+    messages = [
+        ChatMessage(
+            role="system",
+            contents=[TextContent(text="You are a helpful assistant.")],
+        ),
+        ChatMessage(
+            role="system",
+            contents=[TextContent(text="You are a helpful assistant.")],
+        ),
+        ChatMessage(
+            role="user",
+            contents=[TextContent(text="Hello")],
+        ),
+    ]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Deduplication uses hash() which may not deduplicate identical content
+    # This test verifies deduplication logic runs without errors
+    system_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
+    ]
+    # At least one system message should be present
+    assert len(system_messages) >= 1
+
+
+async def test_state_context_injection() -> None:
+    """Test state context message injection for first request."""
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Hello"}],
+            }
+        ],
+        "state": {"items": ["apple", "banana"]},
+    }
+
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(state_schema={"items": {"type": "array"}}),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should inject system message with current state
+    system_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
+    ]
+    assert len(system_messages) == 1
+    assert "apple" in system_messages[0].contents[0].text
+    assert "banana" in system_messages[0].contents[0].text
+
+
+async def test_no_state_context_injection_with_tool_calls() -> None:
+    """Test state context is NOT injected if conversation has tool calls."""
+    from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent
+
+    messages = [
+        ChatMessage(
+            role="assistant",
+            contents=[FunctionCallContent(name="get_weather", call_id="call_xyz", arguments="{}")],
+        ),
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="call_xyz", result="sunny")],
+        ),
+        ChatMessage(
+            role="user",
+            contents=[TextContent(text="Thanks")],
+        ),
+    ]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": [], "state": {"weather": "sunny"}}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(state_schema={"weather": {"type": "string"}}),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should NOT inject state context system message since conversation has tool calls
+    system_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "system"
+    ]
+    assert len(system_messages) == 0
+
+
+async def test_structured_output_processing() -> None:
+    """Test structured output extraction and state update."""
+
+    class RecipeState(BaseModel):
+        ingredients: list[str]
+        message: str
+
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Add tomato"}],
+            }
+        ],
+    }
+
+    # Agent with structured output
+    agent = MockAgent(
+        updates=[
+            AgentRunResponseUpdate(
+                contents=[TextContent(text='{"ingredients": ["tomato"], "message": "Added tomato"}')],
+                role="assistant",
+            )
+        ]
+    )
+    agent.chat_options.response_format = RecipeState
+
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(state_schema={"ingredients": {"type": "array"}}),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should emit StateSnapshotEvent with ingredients
+    state_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(state_events) >= 1
+
+    # Should emit TextMessage with message field
+    text_content_events = [e for e in events if e.type == "TEXT_MESSAGE_CONTENT"]
+    assert len(text_content_events) >= 1
+    assert any("Added tomato" in e.delta for e in text_content_events)
+
+
+async def test_duplicate_client_tools_filtered() -> None:
+    """Test that client tools duplicating server tools are filtered out."""
+
+    @ai_function
+    def get_weather(location: str) -> str:
+        """Get weather for location."""
+        return f"Weather in {location}"
+
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Hello"}],
+            }
+        ],
+        "tools": [
+            {
+                "name": "get_weather",
+                "description": "Client weather tool.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"location": {"type": "string"}},
+                    "required": ["location"],
+                },
+            }
+        ],
+    }
+
+    agent = MockAgent()
+    agent.chat_options.tools = [get_weather]
+
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # tools parameter should not be passed since client tool duplicates server tool
+    assert agent.tools_received is None
+
+
+async def test_unique_client_tools_merged() -> None:
+    """Test that unique client tools are merged with server tools."""
+
+    @ai_function
+    def server_tool() -> str:
+        """Server tool."""
+        return "server"
+
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Hello"}],
+            }
+        ],
+        "tools": [
+            {
+                "name": "client_tool",
+                "description": "Unique client tool.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"param": {"type": "string"}},
+                    "required": ["param"],
+                },
+            }
+        ],
+    }
+
+    agent = MockAgent()
+    agent.chat_options.tools = [server_tool]
+
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # tools parameter should be passed with both server and client tools
+    assert agent.tools_received is not None
+    tool_names = [getattr(tool, "name", None) for tool in agent.tools_received]
+    assert "server_tool" in tool_names
+    assert "client_tool" in tool_names
+
+
+async def test_empty_messages_handling() -> None:
+    """Test orchestrator handles empty message list gracefully."""
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {"messages": []}
+
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should emit run lifecycle events but not call agent
+    assert len(agent.messages_received) == 0
+    run_started = [e for e in events if e.type == "RUN_STARTED"]
+    run_finished = [e for e in events if e.type == "RUN_FINISHED"]
+    assert len(run_started) == 1
+    assert len(run_finished) == 1
+
+
+async def test_all_messages_filtered_handling() -> None:
+    """Test orchestrator handles case where all messages are filtered out."""
+    orchestrator = DefaultOrchestrator()
+
+    input_data = {
+        "messages": [
+            {
+                "role": "tool",
+                "content": [{"type": "tool_result", "tool_call_id": "orphan", "content": "data"}],
+            }
+        ]
+    }
+
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should finish without calling agent
+    assert len(agent.messages_received) == 0
+    run_finished = [e for e in events if e.type == "RUN_FINISHED"]
+    assert len(run_finished) == 1
+
+
+async def test_confirm_changes_with_invalid_json_fallback() -> None:
+    """Test confirm_changes with invalid JSON falls back to normal processing."""
+    from agent_framework import ChatMessage, FunctionCallContent, TextContent
+
+    messages = [
+        ChatMessage(
+            role="assistant",
+            contents=[
+                FunctionCallContent(
+                    name="confirm_changes",
+                    call_id="call_confirm_invalid",
+                    arguments='{"changes": "test"}',
+                )
+            ],
+        ),
+        ChatMessage(
+            role="user",
+            contents=[TextContent(text="invalid json {")],
+        ),
+    ]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Invalid JSON should fall back - user message should be included
+    user_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "user"
+    ]
+    assert len(user_messages) == 1
+
+
+async def test_tool_result_kept_when_call_id_matches() -> None:
+    """Test tool result is kept when call_id matches pending tool calls."""
+    from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent
+
+    messages = [
+        ChatMessage(
+            role="assistant",
+            contents=[FunctionCallContent(name="get_data", call_id="call_match", arguments="{}")],
+        ),
+        ChatMessage(
+            role="tool",
+            contents=[FunctionResultContent(call_id="call_match", result="data")],
+        ),
+    ]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Tool result should be kept
+    tool_messages = [
+        msg
+        for msg in agent.messages_received
+        if (msg.role.value if hasattr(msg.role, "value") else str(msg.role)) == "tool"
+    ]
+    assert len(tool_messages) == 1
+    assert tool_messages[0].contents[0].result == "data"
+
+
+async def test_agent_protocol_fallback_paths() -> None:
+    """Test fallback paths for non-ChatAgent implementations."""
+
+    class CustomAgent:
+        """Custom agent without ChatAgent type."""
+
+        def __init__(self) -> None:
+            self.chat_options = SimpleNamespace(tools=[], response_format=None)
+            self.chat_client = SimpleNamespace(function_invocation_configuration=SimpleNamespace())
+            self.messages_received: list[Any] = []
+
+        async def run_stream(
+            self,
+            messages: list[Any],
+            *,
+            thread: Any = None,
+            tools: list[Any] | None = None,
+        ) -> AsyncGenerator[AgentRunResponseUpdate, None]:
+            self.messages_received = messages
+            yield AgentRunResponseUpdate(contents=[TextContent(text="response")], role="assistant")
+
+    from agent_framework import ChatMessage, TextContent
+
+    messages = [ChatMessage(role="user", contents=[TextContent(text="Hello")])]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+    agent = CustomAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,  # type: ignore
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should work with custom agent implementation
+    assert len(agent.messages_received) > 0
+
+
+async def test_initial_state_snapshot_with_array_schema() -> None:
+    """Test state initialization with array type schema."""
+    from agent_framework import ChatMessage, TextContent
+
+    messages = [ChatMessage(role="user", contents=[TextContent(text="Hello")])]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": [], "state": {}}
+    agent = MockAgent()
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(state_schema={"items": {"type": "array"}}),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Should emit state snapshot with empty array for items
+    state_events = [e for e in events if e.type == "STATE_SNAPSHOT"]
+    assert len(state_events) >= 1
+
+
+async def test_response_format_skip_text_content() -> None:
+    """Test that response_format causes skip_text_content to be set."""
+
+    class OutputModel(BaseModel):
+        result: str
+
+    from agent_framework import ChatMessage, TextContent
+
+    messages = [ChatMessage(role="user", contents=[TextContent(text="Hello")])]
+
+    orchestrator = DefaultOrchestrator()
+    input_data = {"messages": []}
+
+    agent = MockAgent()
+    agent.chat_options.response_format = OutputModel
+
+    context = ExecutionContext(
+        input_data=input_data,
+        agent=agent,
+        config=AgentConfig(),
+    )
+    context._messages = messages
+
+    events = []
+    async for event in orchestrator.run(context):
+        events.append(event)
+
+    # Test passes if no errors occur - verifies response_format code path
+    assert len(events) > 0