Python: Refactor ag-ui to clean up some patterns (#2363)

* Refactor ag-ui to clean up some patterns

* Mypy fixes

* Fix imports, typing, tests, logging.

* Fix test import error

* Fix imports again

* Fix thread handling

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

@@ -3,7 +3,7 @@
 """AgentFrameworkAgent wrapper for AG-UI protocol - Clean Architecture."""
 
 from collections.abc import AsyncGenerator
-from typing import Any
+from typing import Any, cast
 
 from ag_ui.core import BaseEvent
 from agent_framework import AgentProtocol
@@ -22,21 +22,48 @@ class AgentConfig:
 
     def __init__(
         self,
-        state_schema: dict[str, Any] | None = None,
+        state_schema: Any | None = None,
         predict_state_config: dict[str, dict[str, str]] | None = None,
         require_confirmation: bool = True,
     ):
         """Initialize agent configuration.
 
         Args:
-            state_schema: Optional state schema for state management
+            state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
             predict_state_config: Configuration for predictive state updates
             require_confirmation: Whether predictive updates require confirmation
         """
-        self.state_schema = state_schema or {}
+        self.state_schema = self._normalize_state_schema(state_schema)
         self.predict_state_config = predict_state_config or {}
         self.require_confirmation = require_confirmation
 
+    @staticmethod
+    def _normalize_state_schema(state_schema: Any | None) -> dict[str, Any]:
+        """Accept dict or Pydantic model/class and return a properties dict."""
+        if state_schema is None:
+            return {}
+
+        if isinstance(state_schema, dict):
+            return cast(dict[str, Any], state_schema)
+
+        base_model_type: type[Any] | None
+        try:
+            from pydantic import BaseModel as ImportedBaseModel
+
+            base_model_type = ImportedBaseModel
+        except Exception:  # pragma: no cover
+            base_model_type = None
+
+        if base_model_type is not None and isinstance(state_schema, base_model_type):
+            schema_dict = state_schema.__class__.model_json_schema()
+            return schema_dict.get("properties", {}) or {}
+
+        if base_model_type is not None and isinstance(state_schema, type) and issubclass(state_schema, base_model_type):
+            schema_dict = state_schema.model_json_schema()
+            return schema_dict.get("properties", {}) or {}
+
+        return {}
+
 
 class AgentFrameworkAgent:
     """Wraps Agent Framework agents for AG-UI protocol compatibility.
@@ -55,7 +82,7 @@ class AgentFrameworkAgent:
         agent: AgentProtocol,
         name: str | None = None,
         description: str | None = None,
-        state_schema: dict[str, Any] | None = None,
+        state_schema: Any | None = None,
         predict_state_config: dict[str, dict[str, str]] | None = None,
         require_confirmation: bool = True,
         orchestrators: list[Orchestrator] | None = None,
@@ -67,7 +94,7 @@ class AgentFrameworkAgent:
             agent: The Agent Framework agent to wrap
             name: Optional name for the agent
             description: Optional description
-            state_schema: Optional state schema for state management
+            state_schema: Optional state schema for state management; accepts dict or Pydantic model/class
             predict_state_config: Configuration for predictive state updates.
                 Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
             require_confirmation: Whether predictive updates require confirmation.

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

@@ -2,6 +2,7 @@
 
 """FastAPI endpoint creation for AG-UI agents."""
 
+import copy
 import logging
 from typing import Any
 
@@ -19,9 +20,10 @@ def add_agent_framework_fastapi_endpoint(
     app: FastAPI,
     agent: AgentProtocol | AgentFrameworkAgent,
     path: str = "/",
-    state_schema: dict[str, Any] | None = None,
+    state_schema: Any | None = None,
     predict_state_config: dict[str, dict[str, str]] | None = None,
     allow_origins: list[str] | None = None,
+    default_state: dict[str, Any] | None = None,
 ) -> None:
     """Add an AG-UI endpoint to a FastAPI app.
 
@@ -29,10 +31,11 @@ def add_agent_framework_fastapi_endpoint(
         app: The FastAPI application
         agent: The agent to expose (can be raw AgentProtocol or wrapped)
         path: The endpoint path
-        state_schema: Optional state schema for shared state management
+        state_schema: Optional state schema for shared state management; accepts dict or Pydantic model/class
         predict_state_config: Optional predictive state update configuration.
             Format: {"state_key": {"tool": "tool_name", "tool_argument": "arg_name"}}
         allow_origins: CORS origins (not yet implemented)
+        default_state: Optional initial state to seed when the client does not provide state keys
     """
     if isinstance(agent, AgentProtocol):
         wrapped_agent = AgentFrameworkAgent(
@@ -52,6 +55,11 @@ def add_agent_framework_fastapi_endpoint(
         """
         try:
             input_data = await request.json()
+            if default_state:
+                state = input_data.setdefault("state", {})
+                for key, value in default_state.items():
+                    if key not in state:
+                        state[key] = copy.deepcopy(value)
             logger.debug(
                 f"[{path}] Received request - Run ID: {input_data.get('run_id', 'no-run-id')}, "
                 f"Thread ID: {input_data.get('thread_id', 'no-thread-id')}, "

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_message_hygiene.py

@@ -0,0 +1,176 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Message hygiene utilities for orchestrators."""
+
+import json
+import logging
+from typing import Any
+
+from agent_framework import ChatMessage, FunctionCallContent, FunctionResultContent, TextContent
+
+logger = logging.getLogger(__name__)
+
+
+def sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
+    """Normalize tool ordering and inject synthetic results for AG-UI edge cases."""
+    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
+            }
+            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":
+            if pending_confirm_changes_id:
+                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:
+                        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
+                        continue
+                except (json.JSONDecodeError, KeyError) as exc:
+                    logger.debug("Could not parse user message as confirm_changes response: %s", type(exc).__name__)
+
+            if pending_tool_call_ids:
+                logger.info(
+                    f"User message arrived with {len(pending_tool_call_ids)} pending tool calls - injecting synthetic results"
+                )
+                for pending_call_id in pending_tool_call_ids:
+                    logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}")
+                    synthetic_result = ChatMessage(
+                        role="tool",
+                        contents=[
+                            FunctionResultContent(
+                                call_id=pending_call_id,
+                                result="Tool execution skipped - user provided follow-up message",
+                            )
+                        ],
+                    )
+                    sanitized.append(synthetic_result)
+                pending_tool_call_ids = None
+                pending_confirm_changes_id = None
+
+            sanitized.append(msg)
+            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
+                        if call_id == pending_confirm_changes_id:
+                            pending_confirm_changes_id = None
+                        break
+            if keep:
+                sanitized.append(msg)
+            continue
+
+        sanitized.append(msg)
+        pending_tool_call_ids = None
+        pending_confirm_changes_id = None
+
+    return sanitized
+
+
+def deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
+    """Remove duplicate messages while preserving order."""
+    seen_keys: dict[Any, int] = {}
+    unique_messages: list[ChatMessage] = []
+
+    for idx, msg in enumerate(messages):
+        role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+
+        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)
+
+            if key in seen_keys:
+                existing_idx = seen_keys[key]
+                existing_msg = unique_messages[existing_idx]
+
+                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
+
+                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)
+        ):
+            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:
+            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

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_state_manager.py

@@ -0,0 +1,102 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""State orchestration utilities."""
+
+import json
+from typing import Any
+
+from ag_ui.core import CustomEvent, EventType
+from agent_framework import ChatMessage, TextContent
+
+
+class StateManager:
+    """Coordinates state defaults, snapshots, and structured updates."""
+
+    def __init__(
+        self,
+        state_schema: dict[str, Any] | None,
+        predict_state_config: dict[str, dict[str, str]] | None,
+        require_confirmation: bool,
+    ) -> None:
+        self.state_schema = state_schema or {}
+        self.predict_state_config = predict_state_config or {}
+        self.require_confirmation = require_confirmation
+        self.current_state: dict[str, Any] = {}
+
+    def initialize(self, initial_state: dict[str, Any] | None) -> dict[str, Any]:
+        """Initialize state with schema defaults."""
+        self.current_state = (initial_state or {}).copy()
+        self._apply_schema_defaults()
+        return self.current_state
+
+    def predict_state_event(self) -> CustomEvent | None:
+        """Create predict-state custom event when configured."""
+        if not self.predict_state_config:
+            return None
+
+        predict_state_value = [
+            {
+                "state_key": state_key,
+                "tool": config["tool"],
+                "tool_argument": config["tool_argument"],
+            }
+            for state_key, config in self.predict_state_config.items()
+        ]
+
+        return CustomEvent(
+            type=EventType.CUSTOM,
+            name="PredictState",
+            value=predict_state_value,
+        )
+
+    def initial_snapshot_event(self, event_bridge: Any) -> Any:
+        """Emit initial snapshot when schema and state present."""
+        if not self.state_schema:
+            return None
+        self._apply_schema_defaults()
+        return event_bridge.create_state_snapshot_event(self.current_state)
+
+    def state_context_message(self, is_new_user_turn: bool, conversation_has_tool_calls: bool) -> ChatMessage | None:
+        """Inject state context only when starting a new user turn."""
+        if not self.current_state or not self.state_schema:
+            return None
+        if not is_new_user_turn or conversation_has_tool_calls:
+            return None
+
+        state_json = json.dumps(self.current_state, indent=2)
+        return ChatMessage(
+            role="system",
+            contents=[
+                TextContent(
+                    text=(
+                        "Current state of the application:\n"
+                        f"{state_json}\n\n"
+                        "When modifying state, you MUST include ALL existing data plus your changes.\n"
+                        "For example, if adding one new item to a list, include ALL existing items PLUS the one new item.\n"
+                        "Never replace existing data - always preserve and append or merge."
+                    )
+                )
+            ],
+        )
+
+    def extract_state_updates(self, response_dict: dict[str, Any]) -> dict[str, Any]:
+        """Extract state updates from structured response payloads."""
+        if self.state_schema:
+            return {key: response_dict[key] for key in self.state_schema.keys() if key in response_dict}
+        return {k: v for k, v in response_dict.items() if k != "message"}
+
+    def apply_state_updates(self, updates: dict[str, Any]) -> None:
+        """Merge state updates into current state."""
+        if not updates:
+            return
+        self.current_state.update(updates)
+
+    def _apply_schema_defaults(self) -> None:
+        """Fill missing state fields based on schema hints."""
+        for key, schema in self.state_schema.items():
+            if key in self.current_state:
+                continue
+            if isinstance(schema, dict) and schema.get("type") == "array":  # type: ignore
+                self.current_state[key] = []
+            else:
+                self.current_state[key] = {}

python/packages/ag-ui/agent_framework_ag_ui/_orchestration/_tooling.py

@@ -0,0 +1,80 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Tool handling helpers."""
+
+import logging
+from typing import Any
+
+from agent_framework import BaseChatClient, ChatAgent
+
+logger = logging.getLogger(__name__)
+
+
+def collect_server_tools(agent: Any) -> list[Any]:
+    """Collect server tools from ChatAgent or duck-typed agent."""
+    if isinstance(agent, ChatAgent):
+        tools_from_agent = 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}")
+        return server_tools
+
+    try:
+        chat_options_attr = getattr(agent, "chat_options", None)
+        if chat_options_attr is not None:
+            return getattr(chat_options_attr, "tools", None) or []
+    except AttributeError:
+        return []
+    return []
+
+
+def register_additional_client_tools(agent: Any, client_tools: list[Any] | None) -> None:
+    """Register client tools as additional declaration-only tools to avoid server execution."""
+    if not client_tools:
+        return
+
+    if isinstance(agent, ChatAgent):
+        chat_client = agent.chat_client
+        if isinstance(chat_client, BaseChatClient) and chat_client.function_invocation_configuration is not None:
+            chat_client.function_invocation_configuration.additional_tools = client_tools
+            logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
+        return
+
+    try:
+        chat_client_attr = getattr(agent, "chat_client", None)
+        if chat_client_attr is not None:
+            fic = getattr(chat_client_attr, "function_invocation_configuration", None)
+            if fic is not None:
+                fic.additional_tools = client_tools  # type: ignore[attr-defined]
+                logger.debug(
+                    f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
+                )
+    except AttributeError:
+        return
+
+
+def merge_tools(server_tools: list[Any], client_tools: list[Any] | None) -> list[Any] | None:
+    """Combine server and client tools without overriding server metadata."""
+    if not client_tools:
+        logger.info("[TOOLS] No client tools - not passing tools= parameter (using agent's configured tools)")
+        return None
+
+    server_tool_names = {getattr(tool, "name", None) for tool in server_tools}
+    unique_client_tools = [tool for tool in client_tools if getattr(tool, "name", None) not in server_tool_names]
+
+    if not unique_client_tools:
+        logger.info("[TOOLS] All client tools duplicate server tools - not passing tools= parameter")
+        return None
+
+    combined_tools: list[Any] = []
+    if server_tools:
+        combined_tools.extend(server_tools)
+    combined_tools.extend(unique_client_tools)
+    logger.info(
+        f"[TOOLS] Passing tools= parameter with {len(combined_tools)} tools "
+        f"({len(server_tools)} server + {len(unique_client_tools)} unique client)"
+    )
+    return combined_tools

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

@@ -21,7 +21,6 @@ from agent_framework import (
     AgentProtocol,
     AgentThread,
     ChatAgent,
-    ChatMessage,
     FunctionCallContent,
     FunctionResultContent,
     TextContent,
@@ -271,144 +270,29 @@ class DefaultOrchestrator(Orchestrator):
             AG-UI events
         """
         from ._events import AgentFrameworkEventBridge
+        from ._message_adapters import agui_messages_to_snapshot_format
+        from ._orchestration._message_hygiene import deduplicate_messages, sanitize_tool_history
+        from ._orchestration._state_manager import StateManager
+        from ._orchestration._tooling import (
+            collect_server_tools,
+            merge_tools,
+            register_additional_client_tools,
+        )
 
         logger.info(f"Starting default agent run for thread_id={context.thread_id}, run_id={context.run_id}")
 
-        # Initialize state tracking
-        initial_state = context.input_data.get("state", {})
-        current_state: dict[str, Any] = initial_state.copy() if initial_state else {}
-
-        # Check if agent uses structured outputs (response_format)
-        # Use isinstance to narrow type for proper attribute access
         response_format = None
         if isinstance(context.agent, ChatAgent):
             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
+        state_manager = StateManager(
+            state_schema=context.config.state_schema,
+            predict_state_config=context.config.predict_state_config,
+            require_confirmation=context.config.require_confirmation,
+        )
+        current_state = state_manager.initialize(context.input_data.get("state", {}))
 
-            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":
-                    # Check if this user message is a confirm_changes response (JSON with "accepted" field)
-                    # This must be checked BEFORE injecting synthetic results for pending tool calls
-                    if pending_confirm_changes_id:
-                        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}")
-
-                    # Before processing user message, check if there are pending tool calls without results
-                    # This happens when assistant made multiple tool calls but only some got results
-                    # This is checked AFTER confirm_changes special handling above
-                    if pending_tool_call_ids:
-                        logger.info(
-                            f"User message arrived with {len(pending_tool_call_ids)} pending tool calls - injecting synthetic results"
-                        )
-                        for pending_call_id in pending_tool_call_ids:
-                            logger.info(f"Injecting synthetic tool result for pending call_id={pending_call_id}")
-                            synthetic_result = ChatMessage(
-                                role="tool",
-                                contents=[
-                                    FunctionResultContent(
-                                        call_id=pending_call_id,
-                                        result="Tool execution skipped - user provided follow-up message",
-                                    )
-                                ],
-                            )
-                            sanitized.append(synthetic_result)
-                        pending_tool_call_ids = None
-                        pending_confirm_changes_id = None
-
-                    # Normal user message processing
-                    sanitized.append(msg)
-                    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
-                                # Note: We do NOT remove call_id from pending here.
-                                # This allows duplicate tool results to pass through sanitization
-                                # so the deduplicator can choose the best one (prefer non-empty results).
-                                # We only clear pending_tool_call_ids when a user message arrives.
-                                if call_id == pending_confirm_changes_id:
-                                    # For confirm_changes specifically, we do want to clear it
-                                    # since we only expect one response
-                                    pending_confirm_changes_id = None
-                                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,
             thread_id=context.thread_id,
@@ -421,42 +305,19 @@ class DefaultOrchestrator(Orchestrator):
 
         yield event_bridge.create_run_started_event()
 
-        # Emit PredictState custom event if we have predictive state config
-        if context.config.predict_state_config:
-            from ag_ui.core import CustomEvent, EventType
+        predict_event = state_manager.predict_state_event()
+        if predict_event:
+            yield predict_event
 
-            predict_state_value = [
-                {
-                    "state_key": state_key,
-                    "tool": config["tool"],
-                    "tool_argument": config["tool_argument"],
-                }
-                for state_key, config in context.config.predict_state_config.items()
-            ]
+        snapshot_event = state_manager.initial_snapshot_event(event_bridge)
+        if snapshot_event:
+            yield snapshot_event
 
-            yield CustomEvent(
-                type=EventType.CUSTOM,
-                name="PredictState",
-                value=predict_state_value,
-            )
-
-        # If we have a state schema, ensure we emit initial state snapshot
-        if context.config.state_schema:
-            # Initialize missing state fields with appropriate empty values based on schema type
-            for key, schema in context.config.state_schema.items():
-                if key not in current_state:
-                    # Default to empty object; use empty array if schema specifies "array" type
-                    current_state[key] = [] if isinstance(schema, dict) and schema.get("type") == "array" else {}  # type: ignore
-            yield event_bridge.create_state_snapshot_event(current_state)
-
-        # Create thread for context tracking
         thread = AgentThread()
         thread.metadata = {  # type: ignore[attr-defined]
             "ag_ui_thread_id": context.thread_id,
             "ag_ui_run_id": context.run_id,
         }
-
-        # Inject current state into thread metadata so agent can access it
         if current_state:
             thread.metadata["current_state"] = current_state  # type: ignore[attr-defined]
 
@@ -475,90 +336,24 @@ class DefaultOrchestrator(Orchestrator):
                 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}")
+                        logger.debug("    Content %s: %s - text_length=%s", j, content_type, len(content.text))
                     elif isinstance(content, FunctionCallContent):
-                        logger.debug(f"    Content {j}: {content_type} - {content.name}({content.arguments})")
-                    elif isinstance(content, FunctionResultContent):
+                        arg_length = len(str(content.arguments)) if content.arguments else 0
                         logger.debug(
-                            f"    Content {j}: {content_type} - call_id={content.call_id}, result={content.result}"
+                            "    Content %s: %s - %s args_length=%s", j, content_type, content.name, arg_length
+                        )
+                    elif isinstance(content, FunctionResultContent):
+                        result_preview = type(content.result).__name__ if content.result is not None else "None"
+                        logger.debug(
+                            "    Content %s: %s - call_id=%s, result_type=%s",
+                            j,
+                            content_type,
+                            content.call_id,
+                            result_preview,
                         )
                     else:
-                        logger.debug(f"    Content {j}: {content_type} - {content}")
+                        logger.debug(f"    Content {j}: {content_type}")
 
-        # 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)
 
@@ -575,66 +370,45 @@ class DefaultOrchestrator(Orchestrator):
                 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}")
+                        logger.info(f"    Content {j}: {content_type} - text_length={len(content.text)}")
                     elif isinstance(content, FunctionCallContent):
-                        logger.info(f"    Content {j}: {content_type} - {content.name}({content.arguments})")
+                        arg_length = len(str(content.arguments)) if content.arguments else 0
+                        logger.info("    Content %s: %s - %s args_length=%s", j, content_type, content.name, arg_length)
                     elif isinstance(content, FunctionResultContent):
+                        result_preview = type(content.result).__name__ if content.result is not None else "None"
                         logger.info(
-                            f"    Content {j}: {content_type} - call_id={content.call_id}, result={content.result}"
+                            "    Content %s: %s - call_id=%s, result_type=%s",
+                            j,
+                            content_type,
+                            content.call_id,
+                            result_preview,
                         )
                     else:
-                        logger.info(f"    Content {j}: {content_type} - {content}")
+                        logger.info(f"    Content {j}: {content_type}")
 
-        # 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 and this is a new user turn
         messages_to_run: list[Any] = []
-
-        # Check if the last message is from the user (new turn) vs assistant/tool (mid-execution)
         is_new_user_turn = False
         if provider_messages:
             last_msg = provider_messages[-1]
-            is_new_user_turn = last_msg.role.value == "user"
+            role_value = last_msg.role.value if hasattr(last_msg.role, "value") else str(last_msg.role)
+            is_new_user_turn = role_value == "user"
 
-        # Check if conversation has tool calls (indicates mid-execution)
         conversation_has_tool_calls = False
         for msg in provider_messages:
-            if msg.role.value == "assistant" and hasattr(msg, "contents") and msg.contents:
+            role_value = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+            if 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
 
-        # Only inject state context on new user turns AND when conversation doesn't have tool calls
-        # (tool calls indicate we're mid-execution, so state context was already injected)
-        if current_state and context.config.state_schema and is_new_user_turn and not conversation_has_tool_calls:
-            state_json = json.dumps(current_state, indent=2)
-            state_context_msg = ChatMessage(
-                role="system",
-                contents=[
-                    TextContent(
-                        text=f"""Current state of the application:
-        {state_json}
-
-        When modifying state, you MUST include ALL existing data plus your changes.
-        For example, if adding one new item to a list, include ALL existing items PLUS the one new item.
-        Never replace existing data - always preserve and append or merge."""
-                    )
-                ],
-            )
+        state_context_msg = state_manager.state_context_message(
+            is_new_user_turn=is_new_user_turn, conversation_has_tool_calls=conversation_has_tool_calls
+        )
+        if state_context_msg:
             messages_to_run.append(state_context_msg)
 
-        # 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.
-        # Client tools have func=None (declaration-only), so @use_function_invocation
-        # will return the function call without executing (passes back to client).
-        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:
@@ -643,85 +417,31 @@ class DefaultOrchestrator(Orchestrator):
                 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):
-            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
-            try:
-                chat_options_attr = getattr(context.agent, "chat_options", None)
-                if chat_options_attr is not None:
-                    server_tools = getattr(chat_options_attr, "tools", None) or []
-            except AttributeError:
-                pass
+        server_tools = collect_server_tools(context.agent)
+        register_additional_client_tools(context.agent, client_tools)
+        tools_param = merge_tools(server_tools, client_tools)
 
-        # Register client tools as additional (declaration-only) so they are not executed on server
-        if client_tools:
-            if isinstance(context.agent, ChatAgent):
-                # Type-safe path for ChatAgent
-                chat_client = context.agent.chat_client
-                if (
-                    isinstance(chat_client, BaseChatClient)
-                    and chat_client.function_invocation_configuration is not None
-                ):
-                    chat_client.function_invocation_configuration.additional_tools = client_tools
-                    logger.debug(
-                        f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
-                    )
-            else:
-                # Fallback for AgentProtocol implementations (test mocks, custom agents)
-                try:
-                    chat_client_attr = getattr(context.agent, "chat_client", None)
-                    if chat_client_attr is not None:
-                        fic = getattr(chat_client_attr, "function_invocation_configuration", None)
-                        if fic is not None:
-                            fic.additional_tools = client_tools  # type: ignore[attr-defined]
-                            logger.debug(
-                                f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)"
-                            )
-                except AttributeError:
-                    pass
-
-        # 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:
-            # 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] = []
         update_count = 0
-        async for update in context.agent.run_stream(messages_to_run, thread=thread, tools=tools_param):
+        # Prepare metadata for chat client (Azure requires string values)
+        safe_metadata: dict[str, Any] = {}
+        thread_metadata = getattr(thread, "metadata", None)
+        if thread_metadata:
+            for key, value in thread_metadata.items():
+                value_str = value if isinstance(value, str) else json.dumps(value)
+                if len(value_str) > 512:
+                    value_str = value_str[:512]
+                safe_metadata[key] = value_str
+
+        run_kwargs: dict[str, Any] = {
+            "thread": thread,
+            "tools": tools_param,
+            "metadata": safe_metadata,
+        }
+        if safe_metadata:
+            run_kwargs["store"] = True
+
+        async for update in context.agent.run_stream(messages_to_run, **run_kwargs):
             update_count += 1
             logger.info(f"[STREAM] Received update #{update_count} from agent")
             all_updates.append(update)
@@ -733,23 +453,19 @@ class DefaultOrchestrator(Orchestrator):
 
         logger.info(f"[STREAM] Agent stream completed. Total updates: {update_count}")
 
-        # After agent completes, check if we should stop (waiting for user to confirm changes)
         if event_bridge.should_stop_after_confirm:
             logger.info("Stopping run after confirm_changes - waiting for user response")
             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"
+                    "Found %s pending tool calls without end event - emitting ToolCallEndEvent",
+                    len(pending_without_end),
                 )
                 for tool_call in pending_without_end:
                     tool_call_id = tool_call.get("id")
@@ -760,76 +476,47 @@ class DefaultOrchestrator(Orchestrator):
                         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
             from pydantic import BaseModel
 
             logger.info(f"Processing structured output, update count: {len(all_updates)}")
-
-            # Convert streaming updates to final response to get the structured output
             final_response = AgentRunResponse.from_agent_run_response_updates(
                 all_updates, output_format_type=response_format
             )
 
             if final_response.value and isinstance(final_response.value, BaseModel):
-                # Convert Pydantic model to dict
                 response_dict = final_response.value.model_dump(mode="json", exclude_none=True)
-                logger.info(f"Received structured output: {list(response_dict.keys())}")
+                logger.info(f"Received structured output keys: {list(response_dict.keys())}")
 
-                # Extract state fields based on state_schema
-                state_updates: dict[str, Any] = {}
-
-                if context.config.state_schema:
-                    # Use state_schema to determine which fields are state
-                    for state_key in context.config.state_schema.keys():
-                        if state_key in response_dict:
-                            state_updates[state_key] = response_dict[state_key]
-                else:
-                    # No schema: treat all non-message fields as state
-                    state_updates = {k: v for k, v in response_dict.items() if k != "message"}
-
-                # Apply state updates if any found
+                state_updates = state_manager.extract_state_updates(response_dict)
                 if state_updates:
-                    current_state.update(state_updates)
-
-                    # Emit StateSnapshotEvent with the updated state
+                    state_manager.apply_state_updates(state_updates)
                     state_snapshot = event_bridge.create_state_snapshot_event(current_state)
                     yield state_snapshot
                     logger.info(f"Emitted StateSnapshotEvent with updates: {list(state_updates.keys())}")
 
-                # If there's a message field, emit it as chat text
                 if "message" in response_dict and response_dict["message"]:
                     message_id = generate_event_id()
                     yield TextMessageStartEvent(message_id=message_id, role="assistant")
                     yield TextMessageContentEvent(message_id=message_id, delta=response_dict["message"])
                     yield TextMessageEndEvent(message_id=message_id)
-                    logger.info(f"Emitted conversational message: {response_dict['message'][:100]}...")
+                    logger.info(f"Emitted conversational message with length={len(response_dict['message'])}")
 
         logger.info(f"[FINALIZE] Checking for unclosed message. current_message_id={event_bridge.current_message_id}")
         if event_bridge.current_message_id:
             logger.info(f"[FINALIZE] Emitting TextMessageEndEvent for message_id={event_bridge.current_message_id}")
             yield event_bridge.create_message_end_event(event_bridge.current_message_id)
 
-            # Emit MessagesSnapshotEvent to persist the final assistant text message
-            from ._message_adapters import agui_messages_to_snapshot_format
-
-            # Build the final assistant message with accumulated text content
             assistant_text_message = {
                 "id": event_bridge.current_message_id,
                 "role": "assistant",
                 "content": event_bridge.accumulated_text_content,
             }
 
-            # Convert input messages to snapshot format (normalize content structure)
-            # event_bridge.input_messages are already in AG-UI format, just need normalization
             converted_input_messages = agui_messages_to_snapshot_format(event_bridge.input_messages)
-
-            # Build complete messages array
-            # Include: input messages + any pending tool calls/results + final text message
             all_messages = converted_input_messages.copy()
 
-            # Add assistant message with tool calls if any
             if event_bridge.pending_tool_calls:
                 tool_call_message = {
                     "id": generate_event_id(),
@@ -838,18 +525,16 @@ class DefaultOrchestrator(Orchestrator):
                 }
                 all_messages.append(tool_call_message)
 
-            # Add tool results if any
             all_messages.extend(event_bridge.tool_results.copy())
-
-            # Add final text message
             all_messages.append(assistant_text_message)
 
             messages_snapshot = MessagesSnapshotEvent(
                 messages=all_messages,  # type: ignore[arg-type]
             )
             logger.info(
-                f"[FINALIZE] Emitting MessagesSnapshotEvent with {len(all_messages)} messages "
-                f"(text content length: {len(event_bridge.accumulated_text_content)})"
+                "[FINALIZE] Emitting MessagesSnapshotEvent with %s messages (text content length: %s)",
+                len(all_messages),
+                len(event_bridge.accumulated_text_content),
             )
             yield messages_snapshot
         else: