Python: [BREAKING] Simplify API: ChatAgent -> Agent, ChatMessage -> Message (#3747)

* [BREAKING] Rename ChatAgent -> Agent, ChatMessage -> Message, ChatClientProtocol -> SupportsChatGetResponse

Simplify the public API by removing redundant 'Chat' prefix from core types:
- ChatAgent -> Agent
- RawChatAgent -> RawAgent
- ChatMessage -> Message
- ChatClientProtocol -> SupportsChatGetResponse

Also renamed internal WorkflowMessage (was Message in _runner_context) to avoid collision.

No backward compatibility aliases - this is a clean breaking change.

* [BREAKING] Rename Agent chat_client parameter to client

* Fix rebase issues: WorkflowMessage references and broken markdown links

* Fix formatting and lint issues from code quality checks

* Fix import ordering in workflow sample files

* fixed rebase

* Fix test failures: use WorkflowMessage and A2AMessage after ChatMessage→Message rename

- Replace Message(data=..., source_id=...) with WorkflowMessage(...) in workflow tests
- Fix isinstance check in A2A agent to use A2AMessage instead of Message
- Fix import in test_workflow_observability.py (Message→WorkflowMessage)

* Fix lint, fmt, and sample errors after ChatMessage→Message rename

- Auto-fix 70+ ruff lint issues across samples (ChatMessage→Message refs)
- Fix HostedVectorStoreContent→Content.from_hosted_vector_store in file search sample
- Fix _normalize_messages→normalize_messages in custom agent sample
- Fix context.terminate→raise MiddlewareTermination in middleware samples
- Fix with_update_hook→with_transform_hook in override middleware sample
- Add TOptions_co import back to custom_chat_client sample
- Add noqa for FastAPI File() default in chatkit sample
- Fix B023 loop variable capture in weather agent sample

* fix: update Agent constructor calls from chat_client to client in declaration-only tool tests

* fix: add register_cleanup to devui lazy-loading proxy and type stub

* fixed tests and updated new pieces

* fix agui typevar

* fix merge errors

* fix merge conflicts

* fiux merge

* Remove unused links

---------

Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>

python/packages/ag-ui/README.md

@@ -14,15 +14,15 @@ pip install agent-framework-ag-ui
 
 ```python
 from fastapi import FastAPI
-from agent_framework import ChatAgent
+from agent_framework import Agent
 from agent_framework.azure import AzureOpenAIChatClient
 from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 
 # Create your agent
-agent = ChatAgent(
+agent = Agent(
     name="my_agent",
     instructions="You are a helpful assistant.",
-    chat_client=AzureOpenAIChatClient(
+    client=AzureOpenAIChatClient(
         endpoint="https://your-resource.openai.azure.com/",
         deployment_name="gpt-4o-mini",
         api_key="your-api-key",
@@ -58,7 +58,7 @@ The `AGUIChatClient` supports:
 - Streaming and non-streaming responses
 - Hybrid tool execution (client-side + server-side tools)
 - Automatic thread management for conversation continuity
-- Integration with `ChatAgent` for client-side history management
+- Integration with `Agent` for client-side history management
 
 ## Documentation
 
@@ -91,7 +91,7 @@ The AG-UI endpoint does not enforce authentication by default. **For production
 import os
 from fastapi import Depends, FastAPI, HTTPException, Security
 from fastapi.security import APIKeyHeader
-from agent_framework import ChatAgent
+from agent_framework import Agent
 from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 
 # Configure API key authentication
@@ -104,7 +104,7 @@ async def verify_api_key(api_key: str | None = Security(API_KEY_HEADER)) -> None
         raise HTTPException(status_code=401, detail="Invalid or missing API key")
 
 # Create agent and app
-agent = ChatAgent(name="my_agent", instructions="...", chat_client=...)
+agent = Agent(name="my_agent", instructions="...", client=...)
 app = FastAPI()
 
 # Register endpoint WITH authentication

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

@@ -15,11 +15,11 @@ from typing import TYPE_CHECKING, Any, Generic, TypedDict, cast
 import httpx
 from agent_framework import (
     BaseChatClient,
-    ChatMessage,
     ChatResponse,
     ChatResponseUpdate,
     Content,
     FunctionTool,
+    Message,
     ResponseStream,
 )
 from agent_framework._middleware import ChatMiddlewareLayer
@@ -69,10 +69,10 @@ AGUIChatOptionsT = TypeVar(
 )
 
 
-def _apply_server_function_call_unwrap(chat_client: BaseChatClientT) -> BaseChatClientT:
+def _apply_server_function_call_unwrap(client: BaseChatClientT) -> BaseChatClientT:
     """Class decorator that unwraps server-side function calls after tool handling."""
 
-    original_get_response = chat_client.get_response
+    original_get_response = client.get_response
 
     @wraps(original_get_response)
     def response_wrapper(
@@ -105,8 +105,8 @@ def _apply_server_function_call_unwrap(chat_client: BaseChatClientT) -> BaseChat
         _unwrap_server_function_call_contents(cast(MutableSequence[Content | dict[str, Any]], update.contents))
         return update
 
-    chat_client.get_response = response_wrapper  # type: ignore[assignment]
-    return chat_client
+    client.get_response = response_wrapper  # type: ignore[assignment]
+    return client
 
 
 @_apply_server_function_call_unwrap
@@ -130,8 +130,8 @@ class AGUIChatClient(
         This client sends exactly the messages it receives to the server. It does NOT
         automatically maintain conversation history. The server must handle history via thread_id.
 
-        For stateless servers: Use ChatAgent wrapper which will send full message history on each
-        request. However, even with ChatAgent, the server must echo back all context for the
+        For stateless servers: Use Agent wrapper which will send full message history on each
+        request. However, even with Agent, the server must echo back all context for the
         agent to maintain history across turns.
 
     Important: Tool Handling (Hybrid Execution - matches .NET)
@@ -140,7 +140,7 @@ class AGUIChatClient(
         3. When LLM calls a client tool, function invocation executes it locally
         4. Both client and server tools work together (hybrid pattern)
 
-        The wrapping ChatAgent's function invocation handles client tool execution
+        The wrapping Agent's function invocation handles client tool execution
         automatically when the server's LLM decides to call them.
 
     Examples:
@@ -162,18 +162,18 @@ class AGUIChatClient(
                 metadata={"thread_id": thread_id}
             )
 
-        Recommended usage with ChatAgent (client manages history):
+        Recommended usage with Agent (client manages history):
 
         .. code-block:: python
 
-            from agent_framework import ChatAgent
+            from agent_framework import Agent
             from agent_framework.ag_ui import AGUIChatClient
 
             client = AGUIChatClient(endpoint="http://localhost:8888/")
-            agent = ChatAgent(name="assistant", client=client)
+            agent = Agent(name="assistant", client=client)
             thread = await agent.get_new_thread()
 
-            # ChatAgent automatically maintains history and sends full context
+            # Agent automatically maintains history and sends full context
             response = await agent.run("Hello!", thread=thread)
             response2 = await agent.run("How are you?", thread=thread)
 
@@ -282,9 +282,7 @@ class AGUIChatClient(
         logger = get_logger()
         logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")
 
-    def _extract_state_from_messages(
-        self, messages: Sequence[ChatMessage]
-    ) -> tuple[list[ChatMessage], dict[str, Any] | None]:
+    def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[list[Message], dict[str, Any] | None]:
         """Extract state from last message if present.
 
         Args:
@@ -319,11 +317,11 @@ class AGUIChatClient(
 
         return list(messages), None
 
-    def _convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[dict[str, Any]]:
+    def _convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
         """Convert Agent Framework messages to AG-UI format.
 
         Args:
-            messages: List of ChatMessage objects
+            messages: List of Message objects
 
         Returns:
             List of AG-UI formatted message dictionaries
@@ -353,7 +351,7 @@ class AGUIChatClient(
     def _inner_get_response(
         self,
         *,
-        messages: Sequence[ChatMessage],
+        messages: Sequence[Message],
         stream: bool = False,
         options: Mapping[str, Any],
         **kwargs: Any,
@@ -393,7 +391,7 @@ class AGUIChatClient(
     async def _streaming_impl(
         self,
         *,
-        messages: Sequence[ChatMessage],
+        messages: Sequence[Message],
         options: Mapping[str, Any],
         **kwargs: Any,
     ) -> AsyncIterable[ChatResponseUpdate]:
@@ -415,7 +413,7 @@ class AGUIChatClient(
         agui_messages = self._convert_messages_to_agui_format(messages_to_send)
 
         # Send client tools to server so LLM knows about them
-        # Client tools execute via ChatAgent's function invocation wrapper
+        # Client tools execute via Agent's function invocation wrapper
         agui_tools = convert_tools_to_agui_format(options.get("tools"))
 
         # Build set of client tool names (matches .NET clientToolSet)

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

@@ -9,8 +9,8 @@ import logging
 from typing import Any, cast
 
 from agent_framework import (
-    ChatMessage,
     Content,
+    Message,
     prepare_function_call_results,
 )
 
@@ -25,9 +25,9 @@ from ._utils import (
 logger = logging.getLogger(__name__)
 
 
-def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
+def _sanitize_tool_history(messages: list[Message]) -> list[Message]:
     """Normalize tool ordering and inject synthetic results for AG-UI edge cases."""
-    sanitized: list[ChatMessage] = []
+    sanitized: list[Message] = []
     pending_tool_call_ids: set[str] | None = None
     pending_confirm_changes_id: str | None = None
 
@@ -60,7 +60,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
                 ]
                 if filtered_contents:
                     # Create a new message without confirm_changes to avoid mutating the input
-                    filtered_msg = ChatMessage(role=msg.role, contents=filtered_contents)
+                    filtered_msg = Message(role=msg.role, contents=filtered_contents)
                     sanitized.append(filtered_msg)
                 # If no contents left after filtering, don't append anything
 
@@ -99,7 +99,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
 
             if pending_confirm_changes_id and approval_accepted is not None:
                 logger.info(f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}")
-                synthetic_result = ChatMessage(
+                synthetic_result = Message(
                     role="tool",
                     contents=[
                         Content.from_function_result(
@@ -128,7 +128,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
                         logger.info(
                             f"Injecting synthetic tool result for confirm_changes call_id={pending_confirm_changes_id}"
                         )
-                        synthetic_result = ChatMessage(
+                        synthetic_result = Message(
                             role="tool",
                             contents=[
                                 Content.from_function_result(
@@ -152,7 +152,7 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
                 )
                 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(
+                    synthetic_result = Message(
                         role="tool",
                         contents=[
                             Content.from_function_result(
@@ -196,10 +196,10 @@ def _sanitize_tool_history(messages: list[ChatMessage]) -> list[ChatMessage]:
     return sanitized
 
 
-def _deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
+def _deduplicate_messages(messages: list[Message]) -> list[Message]:
     """Remove duplicate messages while preserving order."""
     seen_keys: dict[Any, int] = {}
-    unique_messages: list[ChatMessage] = []
+    unique_messages: list[Message] = []
 
     for idx, msg in enumerate(messages):
         role_value = get_role_value(msg)
@@ -256,7 +256,7 @@ def _deduplicate_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
 
 def normalize_agui_input_messages(
     messages: list[dict[str, Any]],
-) -> tuple[list[ChatMessage], list[dict[str, Any]]]:
+) -> tuple[list[Message], list[dict[str, Any]]]:
     """Normalize raw AG-UI messages into provider and snapshot formats."""
     provider_messages = agui_messages_to_agent_framework(messages)
     provider_messages = _sanitize_tool_history(provider_messages)
@@ -265,14 +265,14 @@ def normalize_agui_input_messages(
     return provider_messages, snapshot_messages
 
 
-def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[ChatMessage]:
+def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Message]:
     """Convert AG-UI messages to Agent Framework format.
 
     Args:
         messages: List of AG-UI messages
 
     Returns:
-        List of Agent Framework ChatMessage objects
+        List of Agent Framework Message objects
     """
 
     def _update_tool_call_arguments(
@@ -367,7 +367,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
         allowed_keys = set(original_args.keys())
         return {key: value for key, value in modified_args.items() if key in allowed_keys}
 
-    result: list[ChatMessage] = []
+    result: list[Message] = []
     for msg in messages:
         # Handle standard tool result messages early (role="tool") to preserve provider invariants
         # This path maps AG‑UI tool messages to function_result content with the correct tool_call_id
@@ -480,7 +480,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
                                 merged_args["steps"] = merged_steps
                         state_args = merged_args
 
-                        # Update the ChatMessage tool call with only enabled steps (for LLM context).
+                        # Update the Message tool call with only enabled steps (for LLM context).
                         # The LLM should only see the steps that were actually approved/executed.
                         updated_args_for_llm = (
                             json.dumps(filtered_args)
@@ -510,14 +510,14 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
                         function_call=func_call_for_approval,
                         additional_properties={"ag_ui_state_args": state_args} if state_args else None,
                     )
-                    chat_msg = ChatMessage(
+                    chat_msg = Message(
                         role="user",
                         contents=[approval_response],
                     )
                 else:
                     # No matching function call found - this is likely a confirm_changes approval
                     # Keep the old behavior for backwards compatibility
-                    chat_msg = ChatMessage(
+                    chat_msg = Message(
                         role="user",
                         contents=[Content.from_text(text=approval_payload_text)],
                         additional_properties={"is_tool_result": True, "tool_call_id": str(tool_call_id or "")},
@@ -537,7 +537,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
                 func_result = result_content
             else:
                 func_result = str(result_content)
-            chat_msg = ChatMessage(
+            chat_msg = Message(
                 role="tool",
                 contents=[Content.from_function_result(call_id=str(tool_call_id), result=func_result)],
             )
@@ -553,7 +553,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
             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(
+            chat_msg = Message(
                 role="tool",
                 contents=[Content.from_function_result(call_id=str(tool_call_id), result=result_content)],
             )
@@ -592,7 +592,7 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
                             arguments=arguments,
                         )
                     )
-            chat_msg = ChatMessage(role="assistant", contents=contents)
+            chat_msg = Message(role="assistant", contents=contents)
             if "id" in msg:
                 chat_msg.message_id = msg["id"]
             result.append(chat_msg)
@@ -622,14 +622,14 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
                 )
                 approval_contents.append(approval_response)
 
-            chat_msg = ChatMessage(role=role, contents=approval_contents)  # type: ignore[call-overload]
+            chat_msg = Message(role=role, contents=approval_contents)  # type: ignore[call-overload]
         else:
             # Regular text message
             content = msg.get("content", "")
             if isinstance(content, str):
-                chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=content)])  # type: ignore[call-overload]
+                chat_msg = Message(role=role, contents=[Content.from_text(text=content)])  # type: ignore[call-overload]
             else:
-                chat_msg = ChatMessage(role=role, contents=[Content.from_text(text=str(content))])  # type: ignore[call-overload]
+                chat_msg = Message(role=role, contents=[Content.from_text(text=str(content))])  # type: ignore[call-overload]
 
         if "id" in msg:
             chat_msg.message_id = msg["id"]
@@ -639,11 +639,11 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Cha
     return result
 
 
-def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str, Any]]) -> list[dict[str, Any]]:
+def agent_framework_messages_to_agui(messages: list[Message] | list[dict[str, Any]]) -> list[dict[str, Any]]:
     """Convert Agent Framework messages to AG-UI format.
 
     Args:
-        messages: List of Agent Framework ChatMessage objects or AG-UI dicts (already converted)
+        messages: List of Agent Framework Message objects or AG-UI dicts (already converted)
 
     Returns:
         List of AG-UI message dictionaries
@@ -672,7 +672,7 @@ def agent_framework_messages_to_agui(messages: list[ChatMessage] | list[dict[str
             result.append(normalized_msg)
             continue
 
-        # Convert ChatMessage to AG-UI format
+        # Convert Message to AG-UI format
         role_value: str = msg.role if hasattr(msg.role, "value") else msg.role  # type: ignore[assignment]
         role = FRAMEWORK_TO_AGUI_ROLE.get(role_value, "user")
 

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

@@ -13,8 +13,8 @@ import logging
 from typing import Any
 
 from agent_framework import (
-    ChatMessage,
     Content,
+    Message,
 )
 
 from .._utils import get_role_value
@@ -22,7 +22,7 @@ from .._utils import get_role_value
 logger = logging.getLogger(__name__)
 
 
-def pending_tool_call_ids(messages: list[ChatMessage]) -> set[str]:
+def pending_tool_call_ids(messages: list[Message]) -> set[str]:
     """Get IDs of tool calls without corresponding results.
 
     Args:
@@ -42,7 +42,7 @@ def pending_tool_call_ids(messages: list[ChatMessage]) -> set[str]:
     return pending_ids - resolved_ids
 
 
-def is_state_context_message(message: ChatMessage) -> bool:
+def is_state_context_message(message: Message) -> bool:
     """Check if a message is a state context system message.
 
     Args:
@@ -178,7 +178,7 @@ def build_safe_metadata(thread_metadata: dict[str, Any] | None) -> dict[str, Any
     return safe_metadata
 
 
-def latest_approval_response(messages: list[ChatMessage]) -> Content | None:
+def latest_approval_response(messages: list[Message]) -> Content | None:
     """Get the latest approval response from messages.
 
     Args:

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

@@ -39,7 +39,7 @@ def collect_server_tools(agent: SupportsAgentRun) -> list[Any]:
     functions need to be included for tool execution during approval flows.
 
     Args:
-        agent: Agent instance to collect tools from. Works with ChatAgent
+        agent: Agent instance to collect tools from. Works with Agent
             or any agent with default_options and optional mcp_tools attributes.
 
     Returns:
@@ -53,7 +53,7 @@ def collect_server_tools(agent: SupportsAgentRun) -> list[Any]:
     tools_from_agent = default_options.get("tools") if isinstance(default_options, dict) else None
     server_tools = list(tools_from_agent) if tools_from_agent else []
 
-    # Include functions from connected MCP tools (only available on ChatAgent)
+    # Include functions from connected MCP tools (only available on Agent)
     mcp_tools = getattr(agent, "mcp_tools", None)
     if mcp_tools:
         server_tools.extend(_collect_mcp_tool_functions(mcp_tools))
@@ -70,19 +70,19 @@ def register_additional_client_tools(agent: SupportsAgentRun, client_tools: list
     """Register client tools as additional declaration-only tools to avoid server execution.
 
     Args:
-        agent: Agent instance to register tools on. Works with ChatAgent
-            or any agent with a chat_client attribute.
+        agent: Agent instance to register tools on. Works with Agent
+            or any agent with a client attribute.
         client_tools: List of client tools to register.
     """
     if not client_tools:
         return
 
-    chat_client = getattr(agent, "chat_client", None)
-    if chat_client is None:
+    client = getattr(agent, "client", None)
+    if client is None:
         return
 
-    if isinstance(chat_client, BaseChatClient) and chat_client.function_invocation_configuration is not None:  # type: ignore[attr-defined]
-        chat_client.function_invocation_configuration["additional_tools"] = client_tools  # type: ignore[attr-defined]
+    if isinstance(client, BaseChatClient) and client.function_invocation_configuration is not None:  # type: ignore[attr-defined]
+        client.function_invocation_configuration["additional_tools"] = client_tools  # type: ignore[attr-defined]
         logger.debug(f"[TOOLS] Registered {len(client_tools)} client tools as additional_tools (declaration-only)")
 
 

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

@@ -28,8 +28,8 @@ from ag_ui.core import (
 )
 from agent_framework import (
     AgentThread,
-    ChatMessage,
     Content,
+    Message,
     SupportsAgentRun,
     prepare_function_call_results,
 )
@@ -195,7 +195,7 @@ class FlowState:
 def _create_state_context_message(
     current_state: dict[str, Any],
     state_schema: dict[str, Any],
-) -> ChatMessage | None:
+) -> Message | None:
     """Create a system message with current state context.
 
     This injects the current state into the conversation so the model
@@ -206,13 +206,13 @@ def _create_state_context_message(
         state_schema: The state schema (used to determine if injection is needed)
 
     Returns:
-        ChatMessage with state context, or None if not needed
+        Message with state context, or None if not needed
     """
     if not current_state or not state_schema:
         return None
 
     state_json = json.dumps(current_state, indent=2)
-    return ChatMessage(
+    return Message(
         role="system",
         contents=[
             Content.from_text(
@@ -229,10 +229,10 @@ def _create_state_context_message(
 
 
 def _inject_state_context(
-    messages: list[ChatMessage],
+    messages: list[Message],
     current_state: dict[str, Any],
     state_schema: dict[str, Any],
-) -> list[ChatMessage]:
+) -> list[Message]:
     """Inject state context message into messages if appropriate.
 
     The state context is injected before the last user message to give
@@ -592,7 +592,7 @@ async def _resolve_approval_responses(
     Args:
         messages: List of messages (will be modified in place)
         tools: List of available tools
-        agent: The agent instance (to get chat_client and config)
+        agent: The agent instance (to get client and config)
         run_kwargs: Kwargs for tool execution
     """
     fcc_todo = _collect_approval_responses(messages)
@@ -605,12 +605,10 @@ async def _resolve_approval_responses(
 
     # Execute approved tool calls
     if approved_responses and tools:
-        chat_client = getattr(agent, "chat_client", None)
-        config = normalize_function_invocation_configuration(
-            getattr(chat_client, "function_invocation_configuration", None)
-        )
+        client = getattr(agent, "client", None)
+        config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None))
         middleware_pipeline = FunctionMiddlewarePipeline(
-            *getattr(chat_client, "function_middleware", ()),
+            *getattr(client, "function_middleware", ()),
             *run_kwargs.get("middleware", ()),
         )
         # Filter out AG-UI-specific kwargs that should not be passed to tool execution
@@ -672,7 +670,7 @@ def _convert_approval_results_to_tool_messages(messages: list[Any]) -> None:
     This modifies the messages list in place.
 
     Args:
-        messages: List of ChatMessage objects to process
+        messages: List of Message objects to process
     """
     result: list[Any] = []
 
@@ -694,11 +692,11 @@ def _convert_approval_results_to_tool_messages(messages: list[Any]) -> None:
 
         # Tool messages first (right after the preceding assistant message per OpenAI requirements)
         for func_result in function_results:
-            result.append(ChatMessage(role="tool", contents=[func_result]))
+            result.append(Message(role="tool", contents=[func_result]))
 
         # Then user message with remaining content (if any)
         if other_contents:
-            result.append(ChatMessage(role=msg.role, contents=other_contents))
+            result.append(Message(role=msg.role, contents=other_contents))
 
     messages[:] = result
 
@@ -793,9 +791,9 @@ async def run_agent_stream(
     # Check for structured output mode (skip text content)
     skip_text = False
     response_format = None
-    from agent_framework import ChatAgent
+    from agent_framework import Agent
 
-    if isinstance(agent, ChatAgent):
+    if isinstance(agent, Agent):
         response_format = agent.default_options.get("response_format")
         skip_text = response_format is not None
 

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

@@ -12,7 +12,7 @@ pip install agent-framework-ag-ui
 
 ### Using Example Agents with Any Chat Client
 
-All example agents are factory functions that accept any `ChatClientProtocol`-compatible chat client:
+All example agents are factory functions that accept any `SupportsChatGetResponse`-compatible chat client:
 
 ```python
 from fastapi import FastAPI
@@ -38,15 +38,15 @@ add_agent_framework_fastapi_endpoint(app, weather_agent(openai_client), "/weathe
 
 ```python
 from fastapi import FastAPI
-from agent_framework import ChatAgent
+from agent_framework import Agent
 from agent_framework.azure import AzureOpenAIChatClient
 from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 
 # Create your agent
-agent = ChatAgent(
+agent = Agent(
     name="my_agent",
     instructions="You are a helpful assistant.",
-    chat_client=AzureOpenAIChatClient(model_id="gpt-4o"),
+    client=AzureOpenAIChatClient(model_id="gpt-4o"),
 )
 
 # Create FastAPI app and add AG-UI endpoint
@@ -70,21 +70,21 @@ This integration supports all 7 AG-UI features:
 
 ## Examples
 
-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.
+All example agents are implemented as **factory functions** that accept any chat client implementing `SupportsChatGetResponse`. This provides maximum flexibility to use Azure OpenAI, OpenAI, Anthropic, or any custom chat client implementation.
 
 ### Available Example Agents
 
 Complete examples for all AG-UI features are available:
 
-- `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
+- `simple_agent(client)` - Basic agentic chat (Feature 1)
+- `weather_agent(client)` - Backend tool rendering (Feature 2)
+- `human_in_the_loop_agent(client)` - Human-in-the-loop with step customization (Feature 3)
+- `task_steps_agent_wrapped(client)` - Agentic generative UI with step execution (Feature 4)
+- `ui_generator_agent(client)` - Tool-based generative UI (Feature 5)
+- `recipe_agent(client)` - Shared state management (Feature 6)
+- `document_writer_agent(client)` - Predictive state updates (Feature 7)
+- `research_assistant_agent(client)` - Research with progress events
+- `task_planner_agent(client)` - Task planning with approvals
 
 ### Using Example Agents
 
@@ -97,7 +97,7 @@ from agent_framework_ag_ui_examples.agents import (
     recipe_agent,
 )
 
-# Create a chat client (use any ChatClientProtocol implementation)
+# Create a chat client (use any SupportsChatGetResponse implementation)
 azure_client = AzureOpenAIChatClient(model_id="gpt-4")
 openai_client = OpenAIChatClient(model_id="gpt-4o")
 
@@ -150,16 +150,16 @@ from agent_framework_ag_ui_examples.agents import (
 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")
+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")
+add_agent_framework_fastapi_endpoint(app, simple_agent(client), "/agentic_chat")
+add_agent_framework_fastapi_endpoint(app, weather_agent(client), "/backend_tool_rendering")
+add_agent_framework_fastapi_endpoint(app, human_in_the_loop_agent(client), "/human_in_the_loop")
+add_agent_framework_fastapi_endpoint(app, task_steps_agent_wrapped(client), "/agentic_generative_ui")  # type: ignore[arg-type]
+add_agent_framework_fastapi_endpoint(app, ui_generator_agent(client), "/tool_based_generative_ui")
+add_agent_framework_fastapi_endpoint(app, recipe_agent(client), "/shared_state")
+add_agent_framework_fastapi_endpoint(app, document_writer_agent(client), "/predictive_state_updates")
 ```
 
 ## Architecture
@@ -187,8 +187,8 @@ The package uses a clean, orchestrator-based architecture:
 You can create your own agent factories following the same pattern as the examples:
 
 ```python
-from agent_framework import ChatAgent, tool
-from agent_framework import ChatClientProtocol
+from agent_framework import Agent, tool
+from agent_framework import SupportsChatGetResponse
 from agent_framework.ag_ui import AgentFrameworkAgent
 
 @tool
@@ -196,19 +196,19 @@ def my_tool(param: str) -> str:
     """My custom tool."""
     return f"Result: {param}"
 
-def my_custom_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+def my_custom_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
     """Create a custom agent with the specified chat client.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
         A configured AgentFrameworkAgent instance
     """
-    agent = ChatAgent(
+    agent = Agent(
         name="my_custom_agent",
         instructions="Custom instructions here",
-        chat_client=chat_client,
+        client=client,
         tools=[my_tool],
     )
 
@@ -220,8 +220,8 @@ def my_custom_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
 
 # Use it
 from agent_framework.azure import AzureOpenAIChatClient
-chat_client = AzureOpenAIChatClient()
-agent = my_custom_agent(chat_client)
+client = AzureOpenAIChatClient()
+agent = my_custom_agent(client)
 ```
 
 ### Shared State
@@ -229,14 +229,14 @@ agent = my_custom_agent(chat_client)
 State is injected as system messages and updated via predictive state updates:
 
 ```python
-from agent_framework import ChatAgent
+from agent_framework import Agent
 from agent_framework.azure import AzureOpenAIChatClient
 from agent_framework.ag_ui import AgentFrameworkAgent
 
 # Create your agent
-agent = ChatAgent(
+agent = Agent(
     name="recipe_agent",
-    chat_client=AzureOpenAIChatClient(model_id="gpt-4o"),
+    client=AzureOpenAIChatClient(model_id="gpt-4o"),
 )
 
 state_schema = {
@@ -266,14 +266,14 @@ wrapped_agent = AgentFrameworkAgent(
 Predictive state updates automatically stream tool arguments as optimistic state updates:
 
 ```python
-from agent_framework import ChatAgent
+from agent_framework import Agent
 from agent_framework.azure import AzureOpenAIChatClient
 from agent_framework.ag_ui import AgentFrameworkAgent
 
 # Create your agent
-agent = ChatAgent(
+agent = Agent(
     name="document_writer",
-    chat_client=AzureOpenAIChatClient(model_id="gpt-4o"),
+    client=AzureOpenAIChatClient(model_id="gpt-4o"),
 )
 
 predict_state_config = {

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

@@ -4,7 +4,7 @@
 
 from __future__ import annotations
 
-from agent_framework import ChatAgent, ChatClientProtocol, tool
+from agent_framework import Agent, SupportsChatGetResponse, tool
 from agent_framework.ag_ui import AgentFrameworkAgent
 
 
@@ -40,19 +40,19 @@ _DOCUMENT_WRITER_INSTRUCTIONS = (
 )
 
 
-def document_writer_agent(chat_client: ChatClientProtocol) -> AgentFrameworkAgent:
+def document_writer_agent(client: SupportsChatGetResponse) -> AgentFrameworkAgent:
     """Create a document writer agent with predictive state updates.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
         A configured AgentFrameworkAgent instance with document writing capabilities
     """
-    agent = ChatAgent(
+    agent = Agent(
         name="document_writer",
         instructions=_DOCUMENT_WRITER_INSTRUCTIONS,
-        chat_client=chat_client,
+        client=client,
         tools=[write_document],
     )
 

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 typing import Any
 
-from agent_framework import ChatAgent, ChatClientProtocol, tool
+from agent_framework import Agent, SupportsChatGetResponse, tool
 from pydantic import BaseModel, Field
 
 
@@ -43,16 +43,16 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
     return f"Generated {len(steps)} execution steps for the task."
 
 
-def human_in_the_loop_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
+def human_in_the_loop_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
     """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
+        client: The chat client to use for the agent
 
     Returns:
-        A configured ChatAgent instance with human-in-the-loop capabilities
+        A configured Agent instance with human-in-the-loop capabilities
     """
-    return ChatAgent(
+    return Agent(
         name="human_in_the_loop_agent",
         instructions="""You are a helpful assistant that can perform any task by breaking it down into steps.
 
@@ -81,6 +81,6 @@ def human_in_the_loop_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[A
     After the user approves and the function executes, THEN provide a brief acknowledgment like:
     "The plan has been created with X steps selected."
     """,
-        chat_client=chat_client,
+        client=client,
         tools=[generate_task_steps],
     )

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

@@ -7,7 +7,7 @@ from __future__ import annotations
 from enum import Enum
 from typing import Any
 
-from agent_framework import ChatAgent, ChatClientProtocol, tool
+from agent_framework import Agent, SupportsChatGetResponse, tool
 from agent_framework.ag_ui import AgentFrameworkAgent
 from pydantic import BaseModel, Field
 
@@ -104,19 +104,19 @@ _RECIPE_INSTRUCTIONS = """You are a helpful recipe assistant that creates and mo
     """
 
 
-def recipe_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
+def recipe_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
     """Create a recipe agent with streaming state updates.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
         A configured AgentFrameworkAgent instance with recipe management
     """
-    agent = ChatAgent(
+    agent = Agent(
         name="recipe_agent",
         instructions=_RECIPE_INSTRUCTIONS,
-        chat_client=chat_client,
+        client=client,
         tools=[update_recipe],
     )
 

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

@@ -5,7 +5,7 @@
 import asyncio
 from typing import Any
 
-from agent_framework import ChatAgent, ChatClientProtocol, tool
+from agent_framework import Agent, SupportsChatGetResponse, tool
 from agent_framework.ag_ui import AgentFrameworkAgent
 
 
@@ -88,19 +88,19 @@ _RESEARCH_ASSISTANT_INSTRUCTIONS = (
 )
 
 
-def research_assistant_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
+def research_assistant_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
     """Create a research assistant agent.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
         A configured AgentFrameworkAgent instance with research capabilities
     """
-    agent = ChatAgent(
+    agent = Agent(
         name="research_assistant",
         instructions=_RESEARCH_ASSISTANT_INSTRUCTIONS,
-        chat_client=chat_client,
+        client=client,
         tools=[research_topic, create_presentation, analyze_data],
     )
 

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

@@ -4,20 +4,20 @@
 
 from typing import Any
 
-from agent_framework import ChatAgent, ChatClientProtocol
+from agent_framework import Agent, SupportsChatGetResponse
 
 
-def simple_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
+def simple_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
     """Create a simple chat agent.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
-        A configured ChatAgent instance
+        A configured Agent instance
     """
-    return ChatAgent[Any](
+    return Agent[Any](
         name="simple_chat_agent",
         instructions="You are a helpful assistant. Be concise and friendly.",
-        chat_client=chat_client,
+        client=client,
     )

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

@@ -4,7 +4,7 @@
 
 from typing import Any
 
-from agent_framework import ChatAgent, ChatClientProtocol, tool
+from agent_framework import Agent, SupportsChatGetResponse, tool
 from agent_framework.ag_ui import AgentFrameworkAgent
 
 
@@ -61,19 +61,19 @@ _TASK_PLANNER_INSTRUCTIONS = (
 )
 
 
-def task_planner_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
+def task_planner_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
     """Create a task planner agent with user approval for actions.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
         A configured AgentFrameworkAgent instance with task planning capabilities
     """
-    agent = ChatAgent(
+    agent = Agent(
         name="task_planner",
         instructions=_TASK_PLANNER_INSTRUCTIONS,
-        chat_client=chat_client,
+        client=client,
         tools=[create_calendar_event, send_email, book_meeting_room],
     )
 

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

@@ -20,7 +20,7 @@ from ag_ui.core import (
     TextMessageStartEvent,
     ToolCallStartEvent,
 )
-from agent_framework import ChatAgent, ChatClientProtocol, ChatMessage, Content, tool
+from agent_framework import Agent, Content, Message, SupportsChatGetResponse, tool
 from agent_framework.ag_ui import AgentFrameworkAgent
 from pydantic import BaseModel, Field
 
@@ -54,16 +54,16 @@ def generate_task_steps(steps: list[TaskStep]) -> str:
     return "Steps generated."
 
 
-def _create_task_steps_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrameworkAgent:
+def _create_task_steps_agent(client: SupportsChatGetResponse[Any]) -> AgentFrameworkAgent:
     """Create the task steps agent using tool-based approach for streaming.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
         A configured AgentFrameworkAgent instance
     """
-    agent = ChatAgent[Any](
+    agent = Agent[Any](
         name="task_steps_agent",
         instructions="""You are a helpful assistant that breaks down tasks into actionable steps.
 
@@ -83,7 +83,7 @@ def _create_task_steps_agent(chat_client: ChatClientProtocol[Any]) -> AgentFrame
     - "Installing platform"
     - "Adding finishing touches"
     """,
-        chat_client=chat_client,
+        client=client,
         tools=[generate_task_steps],
     )
 
@@ -220,30 +220,30 @@ class TaskStepsAgentWithExecution:
 
             # Get the underlying chat agent and client
             chat_agent = self._base_agent.agent  # type: ignore
-            chat_client = chat_agent.chat_client  # type: ignore
+            client = chat_agent.client  # type: ignore
 
             # Build messages for summary call
 
             original_messages = input_data.get("messages", [])
 
-            # Convert to ChatMessage objects if needed
-            messages: list[ChatMessage] = []
+            # Convert to Message objects if needed
+            messages: list[Message] = []
             for msg in original_messages:
                 if isinstance(msg, dict):
                     content_str = msg.get("content", "")
                     if isinstance(content_str, str):
                         messages.append(
-                            ChatMessage(
+                            Message(
                                 role=msg.get("role", "user"),
                                 contents=[Content.from_text(text=content_str)],
                             )
                         )
-                elif isinstance(msg, ChatMessage):
+                elif isinstance(msg, Message):
                     messages.append(msg)
 
             # Add completion message
             messages.append(
-                ChatMessage(
+                Message(
                     role="user",
                     contents=[
                         Content.from_text(
@@ -270,7 +270,7 @@ class TaskStepsAgentWithExecution:
 
                 # Stream completion
                 accumulated_text = ""
-                async for chunk in chat_client.get_response(messages=messages, stream=True):
+                async for chunk in client.get_response(messages=messages, stream=True):
                     # chunk is ChatResponseUpdate
                     if hasattr(chunk, "text") and chunk.text:
                         accumulated_text += chunk.text
@@ -332,14 +332,14 @@ class TaskStepsAgentWithExecution:
             yield run_finished_event
 
 
-def task_steps_agent_wrapped(chat_client: ChatClientProtocol[Any]) -> TaskStepsAgentWithExecution:
+def task_steps_agent_wrapped(client: SupportsChatGetResponse[Any]) -> TaskStepsAgentWithExecution:
     """Create a task steps agent with execution simulation.
 
     Args:
-        chat_client: The chat client to use for the agent
+        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)
+    base_agent = _create_task_steps_agent(client)
     return TaskStepsAgentWithExecution(base_agent)

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

@@ -7,7 +7,7 @@ from __future__ import annotations
 import sys
 from typing import TYPE_CHECKING, Any, TypedDict
 
-from agent_framework import ChatAgent, ChatClientProtocol, FunctionTool
+from agent_framework import Agent, FunctionTool, SupportsChatGetResponse
 from agent_framework.ag_ui import AgentFrameworkAgent
 
 if sys.version_info >= (3, 13):
@@ -168,19 +168,19 @@ _UI_GENERATOR_INSTRUCTIONS = """You MUST use the provided tools to generate cont
 OptionsT = TypeVar("OptionsT", bound=TypedDict, default="ChatOptions")  # type: ignore[valid-type]
 
 
-def ui_generator_agent(chat_client: ChatClientProtocol[OptionsT]) -> AgentFrameworkAgent:
+def ui_generator_agent(client: SupportsChatGetResponse[OptionsT]) -> AgentFrameworkAgent:
     """Create a UI generator agent with custom React component rendering.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
         A configured AgentFrameworkAgent instance with UI generation capabilities
     """
-    agent = ChatAgent(
+    agent = Agent(
         name="ui_generator",
         instructions=_UI_GENERATOR_INSTRUCTIONS,
-        chat_client=chat_client,
+        client=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
         default_options={"tool_choice": "required"},  # type: ignore

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

@@ -6,7 +6,7 @@ from __future__ import annotations
 
 from typing import Any
 
-from agent_framework import ChatAgent, ChatClientProtocol, tool
+from agent_framework import Agent, SupportsChatGetResponse, tool
 
 
 @tool
@@ -59,16 +59,16 @@ def get_forecast(location: str, days: int = 3) -> str:
     return f"{days}-day forecast for {location}:\n" + "\n".join(forecast)
 
 
-def weather_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
+def weather_agent(client: SupportsChatGetResponse[Any]) -> Agent[Any]:
     """Create a weather agent with get_weather and get_forecast tools.
 
     Args:
-        chat_client: The chat client to use for the agent
+        client: The chat client to use for the agent
 
     Returns:
-        A configured ChatAgent instance with weather tools
+        A configured Agent instance with weather tools
     """
-    return ChatAgent[Any](
+    return Agent[Any](
         name="weather_agent",
         instructions=(
             "You are a helpful weather assistant. "
@@ -76,6 +76,6 @@ def weather_agent(chat_client: ChatClientProtocol[Any]) -> ChatAgent[Any]:
             "Always provide friendly and informative responses. "
             "First return the weather result, and then return details about the forecast."
         ),
-        chat_client=chat_client,
+        client=client,
         tools=[get_weather, get_forecast],
     )

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

@@ -4,7 +4,7 @@
 
 from typing import Any, cast
 
-from agent_framework._clients import ChatClientProtocol
+from agent_framework._clients import SupportsChatGetResponse
 from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from agent_framework.azure import AzureOpenAIChatClient
 from fastapi import FastAPI
@@ -19,10 +19,10 @@ def register_backend_tool_rendering(app: FastAPI) -> None:
         app: The FastAPI application.
     """
     # Create a chat client and call the factory function
-    chat_client = cast(ChatClientProtocol[Any], AzureOpenAIChatClient())
+    client = cast(SupportsChatGetResponse[Any], AzureOpenAIChatClient())
 
     add_agent_framework_fastapi_endpoint(
         app,
-        weather_agent(chat_client),
+        weather_agent(client),
         "/backend_tool_rendering",
     )

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

@@ -10,7 +10,7 @@ from typing import cast
 
 import uvicorn
 from agent_framework import ChatOptions
-from agent_framework._clients import ChatClientProtocol
+from agent_framework._clients import SupportsChatGetResponse
 from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from agent_framework.anthropic import AnthropicClient
 from agent_framework.azure import AzureOpenAIChatClient
@@ -67,43 +67,43 @@ app.add_middleware(
 # Create a shared chat client for all agents
 # You can use different chat clients for different agents if needed
 # Set CHAT_CLIENT=anthropic to use Anthropic, defaults to Azure OpenAI
-chat_client: ChatClientProtocol[ChatOptions] = cast(
-    ChatClientProtocol[ChatOptions],
+client: SupportsChatGetResponse[ChatOptions] = cast(
+    SupportsChatGetResponse[ChatOptions],
     AnthropicClient() if os.getenv("CHAT_CLIENT", "").lower() == "anthropic" else AzureOpenAIChatClient(),
 )
 
 # Agentic Chat - basic chat agent
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=simple_agent(chat_client),
+    agent=simple_agent(client),
     path="/agentic_chat",
 )
 
 # Backend Tool Rendering - agent with tools
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=weather_agent(chat_client),
+    agent=weather_agent(client),
     path="/backend_tool_rendering",
 )
 
 # Shared State - recipe agent with structured output
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=recipe_agent(chat_client),
+    agent=recipe_agent(client),
     path="/shared_state",
 )
 
 # Predictive State Updates - document writer with predictive state
 add_agent_framework_fastapi_endpoint(
     app=app,
-    agent=document_writer_agent(chat_client),
+    agent=document_writer_agent(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(chat_client),
+    agent=human_in_the_loop_agent(client),
     path="/human_in_the_loop",
     state_schema={"steps": {"type": "array"}},
     predict_state_config={"steps": {"tool": "generate_task_steps", "tool_argument": "steps"}},
@@ -112,14 +112,14 @@ 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_wrapped(chat_client),  # type: ignore[arg-type]
+    agent=task_steps_agent_wrapped(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(chat_client),
+    agent=ui_generator_agent(client),
     path="/tool_based_generative_ui",
 )
 

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

@@ -35,9 +35,9 @@ python client_advanced.py
 
 **Note:** This example shows direct `AGUIChatClient` usage. Tool execution and conversation continuity depend on server-side configuration and capabilities.
 
-### ChatAgent Integration (`client_with_agent.py`)
+### Agent Integration (`client_with_agent.py`)
 
-Best practice example using `ChatAgent` wrapper with **AgentThread**
+Best practice example using `Agent` wrapper with **AgentThread**
 - **AgentThread** maintains conversation state
 - Client-side conversation history management via `thread.message_store`
 - **Hybrid tool execution**: client-side + server-side tools simultaneously
@@ -77,7 +77,7 @@ The AG-UI protocol supports two approaches to conversation history:
    - Full message history sent with each request
    - Works with any AG-UI server (stateful or stateless)
 
-The `ChatAgent` wrapper (used in client_with_agent.py) collects messages from local storage and sends the full history to `AGUIChatClient`, which then forwards everything to the server.
+The `Agent` wrapper (used in client_with_agent.py) collects messages from local storage and sends the full history to `AGUIChatClient`, which then forwards everything to the server.
 
 ### Tool/Function Calling
 
@@ -91,14 +91,14 @@ Client defines:           Server defines:
 
 User: "What's the weather in SF and what time is it?"
     ↓
-ChatAgent sends: full history + tool definitions for get_weather, read_sensors
+Agent sends: full history + tool definitions for get_weather, read_sensors
     ↓
 Server LLM decides: "I need get_weather('SF') and get_current_time()"
     ↓
 Server executes get_current_time() → "2025-11-11 14:30:00 UTC"
 Server sends function call request → get_weather('SF')
     ↓
-ChatAgent intercepts get_weather call → executes locally
+Agent intercepts get_weather call → executes locally
     ↓
 Client sends result → "Sunny, 72°F"
     ↓
@@ -110,7 +110,7 @@ Client receives final response
 **How it works:**
 
 1. **Client-Side Tools** (`client_with_agent.py`):
-   - Tools defined in ChatAgent's `tools` parameter execute locally
+   - Tools defined in Agent's `tools` parameter execute locally
    - Tool metadata (name, description, schema) sent to server for planning
    - When server requests client tool → client intercepts → executes locally → sends result
 
@@ -126,7 +126,7 @@ Client receives final response
    - Client tools execute client-side
 
 **Direct AGUIChatClient Usage** (client_advanced.py):
-Even without ChatAgent wrapper, client-side tools work:
+Even without Agent wrapper, client-side tools work:
 - Tools passed in ChatOptions execute locally
 - Server can also have its own tools
 - Hybrid execution works automatically
@@ -184,7 +184,7 @@ Create a file named `server.py`:
 
 import os
 
-from agent_framework import ChatAgent
+from agent_framework import Agent
 from agent_framework.azure import AzureOpenAIChatClient
 from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from fastapi import FastAPI
@@ -202,10 +202,10 @@ if not api_key:
     raise ValueError("AZURE_OPENAI_API_KEY environment variable is required")
 
 # Create the AI agent
-agent = ChatAgent(
+agent = Agent(
     name="AGUIAssistant",
     instructions="You are a helpful assistant.",
-    chat_client=AzureOpenAIChatClient(
+    client=AzureOpenAIChatClient(
         endpoint=endpoint,
         deployment_name=deployment_name,
         api_key=api_key,
@@ -227,7 +227,7 @@ if __name__ == "__main__":
 ### Key Concepts
 
 - **`add_agent_framework_fastapi_endpoint`**: Registers the AG-UI endpoint with automatic request/response handling and SSE streaming
-- **`ChatAgent`**: The agent that will handle incoming requests
+- **`Agent`**: The agent that will handle incoming requests
 - **FastAPI Integration**: Uses FastAPI's native async support for streaming responses
 - **Instructions**: The agent is created with default instructions, which can be overridden by client messages
 - **Configuration**: `AzureOpenAIChatClient` can read from environment variables (`AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`, `AZURE_OPENAI_API_KEY`) or accept parameters directly
@@ -236,10 +236,10 @@ if __name__ == "__main__":
 
 ```python
 # No need to read environment variables manually
-agent = ChatAgent(
+agent = Agent(
     name="AGUIAssistant",
     instructions="You are a helpful assistant.",
-    chat_client=AzureOpenAIChatClient(),  # Reads from environment automatically
+    client=AzureOpenAIChatClient(),  # Reads from environment automatically
 )
 ```
 
@@ -354,7 +354,7 @@ if __name__ == "__main__":
 - **Thread Management**: Pass `thread_id` in metadata to maintain conversation context across requests
 - **Streaming Responses**: Use `get_response(..., stream=True)` for real-time streaming or `get_response(..., stream=False)` for non-streaming
 - **Context Manager**: Use `async with` for automatic cleanup of HTTP connections
-- **Standard Interface**: Works with all Agent Framework patterns (ChatAgent, tools, etc.)
+- **Standard Interface**: Works with all Agent Framework patterns (Agent, tools, etc.)
 - **Hybrid Tool Execution**: Supports both client-side and server-side tools executing together in the same conversation
 
 ### Configure and Run the Client

python/packages/ag-ui/getting_started/client_advanced.py

@@ -114,15 +114,15 @@ async def non_streaming_example(client: AGUIChatClient, thread_id: str | None =
 async def tool_example(client: AGUIChatClient, thread_id: str | None = None):
     """Demonstrate sending tool definitions to the server.
 
-    IMPORTANT: When using AGUIChatClient directly (without ChatAgent wrapper):
+    IMPORTANT: When using AGUIChatClient directly (without Agent wrapper):
     - Tools are sent as DEFINITIONS only
     - No automatic client-side execution (no function invocation middleware)
     - Server must have matching tool implementations to execute them
 
     For CLIENT-SIDE tool execution (like .NET AGUIClient sample):
-    - Use ChatAgent wrapper with tools
+    - Use Agent wrapper with tools
     - See client_with_agent.py for the hybrid pattern
-    - ChatAgent middleware intercepts and executes client tools locally
+    - Agent middleware intercepts and executes client tools locally
     - Server can have its own tools that execute server-side
     - Both client and server tools work together in same conversation
 
@@ -186,7 +186,7 @@ async def conversation_example(client: AGUIChatClient):
 
     # Check if context was maintained
     if "alice" not in response2.text.lower():
-        print("\n[Note: Server may not maintain thread context - consider using ChatAgent for history management]")
+        print("\n[Note: Server may not maintain thread context - consider using Agent for history management]")
 
     # Third turn
     print("\nUser: Can you also tell me what 10 * 5 is?\n")

python/packages/ag-ui/getting_started/client_with_agent.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""Example showing ChatAgent with AGUIChatClient for hybrid tool execution.
+"""Example showing Agent with AGUIChatClient for hybrid tool execution.
 
 This demonstrates the HYBRID pattern matching .NET AGUIClient implementation:
 
@@ -24,7 +24,7 @@ import asyncio
 import logging
 import os
 
-from agent_framework import ChatAgent, tool
+from agent_framework import Agent, tool
 from agent_framework.ag_ui import AGUIChatClient
 
 # Enable debug logging
@@ -55,7 +55,7 @@ def get_weather(location: str) -> str:
 
 
 async def main():
-    """Demonstrate ChatAgent + AGUIChatClient hybrid tool execution.
+    """Demonstrate Agent + AGUIChatClient hybrid tool execution.
 
     This matches the .NET pattern from Program.cs where:
     - AIAgent agent = chatClient.CreateAIAgent(tools: [...])
@@ -63,14 +63,14 @@ async def main():
     - RunStreamingAsync(messages, thread)
 
     Python equivalent:
-    - agent = ChatAgent(chat_client=AGUIChatClient(...), tools=[...])
+    - agent = Agent(client=AGUIChatClient(...), tools=[...])
     - thread = agent.get_new_thread()  # Creates thread with message_store
     - agent.run(message, stream=True, thread=thread)  # Thread accumulates history
     """
     server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
 
     print("=" * 70)
-    print("ChatAgent + AGUIChatClient: Hybrid Tool Execution")
+    print("Agent + AGUIChatClient: Hybrid Tool Execution")
     print("=" * 70)
     print(f"\nServer: {server_url}")
     print("\nThis example demonstrates:")
@@ -82,11 +82,11 @@ async def main():
     try:
         # Create remote client in async context manager
         async with AGUIChatClient(endpoint=server_url) as remote_client:
-            # Wrap in ChatAgent for conversation history management
-            agent = ChatAgent(
+            # Wrap in Agent for conversation history management
+            agent = Agent(
                 name="remote_assistant",
                 instructions="You are a helpful assistant. Remember user information across the conversation.",
-                chat_client=remote_client,
+                client=remote_client,
                 tools=[get_weather],
             )
 

python/packages/ag-ui/getting_started/server.py

@@ -7,7 +7,7 @@ from __future__ import annotations
 import logging
 import os
 
-from agent_framework import ChatAgent, tool
+from agent_framework import Agent, tool
 from agent_framework.ag_ui import add_agent_framework_fastapi_endpoint
 from agent_framework.azure import AzureOpenAIChatClient
 from dotenv import load_dotenv
@@ -116,10 +116,10 @@ def get_time_zone(location: str) -> str:
 # The client will send get_weather tool metadata so the LLM knows about it,
 # and the function invocation mixin on AGUIChatClient will execute it client-side.
 # This matches the .NET AG-UI hybrid execution pattern.
-agent = ChatAgent(
+agent = Agent(
     name="AGUIAssistant",
     instructions="You are a helpful assistant. Use get_weather for weather and get_time_zone for time zones.",
-    chat_client=AzureOpenAIChatClient(
+    client=AzureOpenAIChatClient(
         endpoint=endpoint,
         deployment_name=deployment_name,
     ),

python/packages/ag-ui/tests/ag_ui/conftest.py

@@ -13,13 +13,13 @@ from agent_framework import (
     AgentResponseUpdate,
     AgentThread,
     BaseChatClient,
-    ChatClientProtocol,
-    ChatMessage,
     ChatOptions,
     ChatResponse,
     ChatResponseUpdate,
     Content,
+    Message,
     SupportsAgentRun,
+    SupportsChatGetResponse,
 )
 from agent_framework._clients import OptionsCoT
 from agent_framework._middleware import ChatMiddlewareLayer
@@ -43,7 +43,7 @@ class StreamingChatClientStub(
     BaseChatClient[OptionsCoT],
     Generic[OptionsCoT],
 ):
-    """Typed streaming stub that satisfies ChatClientProtocol."""
+    """Typed streaming stub that satisfies SupportsChatGetResponse."""
 
     def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
         super().__init__(function_middleware=[])
@@ -55,7 +55,7 @@ class StreamingChatClientStub(
     @overload
     def get_response(
         self,
-        messages: str | ChatMessage | Sequence[str | ChatMessage],
+        messages: str | Message | Sequence[str | Message],
         *,
         stream: Literal[False] = ...,
         options: ChatOptions[Any],
@@ -65,7 +65,7 @@ class StreamingChatClientStub(
     @overload
     def get_response(
         self,
-        messages: str | ChatMessage | Sequence[str | ChatMessage],
+        messages: str | Message | Sequence[str | Message],
         *,
         stream: Literal[False] = ...,
         options: OptionsCoT | ChatOptions[None] | None = ...,
@@ -75,7 +75,7 @@ class StreamingChatClientStub(
     @overload
     def get_response(
         self,
-        messages: str | ChatMessage | Sequence[str | ChatMessage],
+        messages: str | Message | Sequence[str | Message],
         *,
         stream: Literal[True],
         options: OptionsCoT | ChatOptions[Any] | None = ...,
@@ -84,7 +84,7 @@ class StreamingChatClientStub(
 
     def get_response(
         self,
-        messages: str | ChatMessage | Sequence[str | ChatMessage],
+        messages: str | Message | Sequence[str | Message],
         *,
         stream: bool = False,
         options: OptionsCoT | ChatOptions[Any] | None = None,
@@ -106,7 +106,7 @@ class StreamingChatClientStub(
     def _inner_get_response(
         self,
         *,
-        messages: Sequence[ChatMessage],
+        messages: Sequence[Message],
         stream: bool = False,
         options: Mapping[str, Any],
         **kwargs: Any,
@@ -121,7 +121,7 @@ class StreamingChatClientStub(
         return self._get_response_impl(messages, options, **kwargs)
 
     async def _get_response_impl(
-        self, messages: Sequence[ChatMessage], options: Mapping[str, Any], **kwargs: Any
+        self, messages: Sequence[Message], options: Mapping[str, Any], **kwargs: Any
     ) -> ChatResponse:
         """Non-streaming implementation."""
         if self._response_fn is not None:
@@ -132,7 +132,7 @@ class StreamingChatClientStub(
             contents.extend(update.contents)
 
         return ChatResponse(
-            messages=[ChatMessage(role="assistant", contents=contents)],
+            messages=[Message(role="assistant", contents=contents)],
             response_id="stub-response",
         )
 
@@ -141,7 +141,7 @@ def stream_from_updates(updates: list[ChatResponseUpdate]) -> StreamFn:
     """Create a stream function that yields from a static list of updates."""
 
     async def _stream(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         for update in updates:
             yield update
@@ -159,7 +159,7 @@ class StubAgent(SupportsAgentRun):
         agent_id: str = "stub-agent",
         agent_name: str | None = "stub-agent",
         default_options: Any | None = None,
-        chat_client: Any | None = None,
+        client: Any | None = None,
     ) -> None:
         self.id = agent_id
         self.name = agent_name
@@ -168,14 +168,14 @@ class StubAgent(SupportsAgentRun):
         self.default_options: dict[str, Any] = (
             default_options if isinstance(default_options, dict) else {"tools": None, "response_format": None}
         )
-        self.chat_client = chat_client or SimpleNamespace(function_invocation_configuration=None)
+        self.client = client or SimpleNamespace(function_invocation_configuration=None)
         self.messages_received: list[Any] = []
         self.tools_received: list[Any] | None = None
 
     @overload
     def run(
         self,
-        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
+        messages: str | Message | Sequence[str | Message] | None = None,
         *,
         stream: Literal[False] = ...,
         thread: AgentThread | None = None,
@@ -185,7 +185,7 @@ class StubAgent(SupportsAgentRun):
     @overload
     def run(
         self,
-        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
+        messages: str | Message | Sequence[str | Message] | None = None,
         *,
         stream: Literal[True],
         thread: AgentThread | None = None,
@@ -194,7 +194,7 @@ class StubAgent(SupportsAgentRun):
 
     def run(
         self,
-        messages: str | ChatMessage | Sequence[str | ChatMessage] | None = None,
+        messages: str | Message | Sequence[str | Message] | None = None,
         *,
         stream: bool = False,
         thread: AgentThread | None = None,
@@ -226,7 +226,7 @@ class StubAgent(SupportsAgentRun):
 
 
 @pytest.fixture
-def streaming_chat_client_stub() -> type[ChatClientProtocol]:
+def streaming_chat_client_stub() -> type[SupportsChatGetResponse]:
     """Return the StreamingChatClientStub class for creating test instances."""
     return StreamingChatClientStub  # type: ignore[return-value]
 

python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py

@@ -7,11 +7,11 @@ from collections.abc import AsyncGenerator, Awaitable, MutableSequence
 from typing import Any
 
 from agent_framework import (
-    ChatMessage,
     ChatOptions,
     ChatResponse,
     ChatResponseUpdate,
     Content,
+    Message,
     ResponseStream,
     tool,
 )
@@ -29,13 +29,11 @@ class TestableAGUIChatClient(AGUIChatClient):
         """Expose http service for monkeypatching."""
         return self._http_service
 
-    def extract_state_from_messages(
-        self, messages: list[ChatMessage]
-    ) -> tuple[list[ChatMessage], dict[str, Any] | None]:
+    def extract_state_from_messages(self, messages: list[Message]) -> tuple[list[Message], dict[str, Any] | None]:
         """Expose state extraction helper."""
         return self._extract_state_from_messages(messages)
 
-    def convert_messages_to_agui_format(self, messages: list[ChatMessage]) -> list[dict[str, Any]]:
+    def convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
         """Expose message conversion helper."""
         return self._convert_messages_to_agui_format(messages)
 
@@ -44,7 +42,7 @@ class TestableAGUIChatClient(AGUIChatClient):
         return self._get_thread_id(options)
 
     def inner_get_response(
-        self, *, messages: MutableSequence[ChatMessage], options: dict[str, Any], stream: bool = False
+        self, *, messages: MutableSequence[Message], options: dict[str, Any], stream: bool = False
     ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
         """Proxy to protected response call."""
         return self._inner_get_response(messages=messages, options=options, stream=stream)
@@ -69,8 +67,8 @@ class TestAGUIChatClient:
         """Test state extraction when no state is present."""
         client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         messages = [
-            ChatMessage(role="user", text="Hello"),
-            ChatMessage(role="assistant", text="Hi there"),
+            Message(role="user", text="Hello"),
+            Message(role="assistant", text="Hi there"),
         ]
 
         result_messages, state = client.extract_state_from_messages(messages)
@@ -89,8 +87,8 @@ class TestAGUIChatClient:
         state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
 
         messages = [
-            ChatMessage(role="user", text="Hello"),
-            ChatMessage(
+            Message(role="user", text="Hello"),
+            Message(
                 role="user",
                 contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
             ),
@@ -112,7 +110,7 @@ class TestAGUIChatClient:
         state_b64 = base64.b64encode(invalid_json.encode("utf-8")).decode("utf-8")
 
         messages = [
-            ChatMessage(
+            Message(
                 role="user",
                 contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
             ),
@@ -127,8 +125,8 @@ class TestAGUIChatClient:
         """Test message conversion to AG-UI format."""
         client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         messages = [
-            ChatMessage(role="user", text="What is the weather?"),
-            ChatMessage(role="assistant", text="Let me check.", message_id="msg_123"),
+            Message(role="user", text="What is the weather?"),
+            Message(role="assistant", text="Let me check.", message_id="msg_123"),
         ]
 
         agui_messages = client.convert_messages_to_agui_format(messages)
@@ -175,7 +173,7 @@ class TestAGUIChatClient:
         client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
-        messages = [ChatMessage(role="user", text="Test message")]
+        messages = [Message(role="user", text="Test message")]
         chat_options = ChatOptions()
 
         updates: list[ChatResponseUpdate] = []
@@ -208,7 +206,7 @@ class TestAGUIChatClient:
         client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
-        messages = [ChatMessage(role="user", text="Test message")]
+        messages = [Message(role="user", text="Test message")]
         chat_options = {}
 
         response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -251,7 +249,7 @@ class TestAGUIChatClient:
         client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
-        messages = [ChatMessage(role="user", text="Test with tools")]
+        messages = [Message(role="user", text="Test with tools")]
         chat_options = ChatOptions(tools=[test_tool])
 
         response = await client.inner_get_response(messages=messages, options=chat_options)
@@ -275,7 +273,7 @@ class TestAGUIChatClient:
         client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
-        messages = [ChatMessage(role="user", text="Test server tool execution")]
+        messages = [Message(role="user", text="Test server tool execution")]
 
         updates: list[ChatResponseUpdate] = []
         async for update in client.get_response(messages, stream=True):
@@ -317,7 +315,7 @@ class TestAGUIChatClient:
         client = TestableAGUIChatClient(endpoint="http://localhost:8888/")
         monkeypatch.setattr(client.http_service, "post_run", mock_post_run)
 
-        messages = [ChatMessage(role="user", text="Test server tool execution")]
+        messages = [Message(role="user", text="Test server tool execution")]
 
         async for _ in client.get_response(
             messages, stream=True, options={"tool_choice": "auto", "tools": [client_tool]}
@@ -333,8 +331,8 @@ class TestAGUIChatClient:
         state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8")
 
         messages = [
-            ChatMessage(role="user", text="Hello"),
-            ChatMessage(
+            Message(role="user", text="Hello"),
+            Message(
                 role="user",
                 contents=[Content.from_uri(uri=f"data:application/json;base64,{state_b64}")],
             ),

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

@@ -7,7 +7,7 @@ from collections.abc import AsyncIterator, MutableSequence
 from typing import Any
 
 import pytest
-from agent_framework import ChatAgent, ChatMessage, ChatOptions, ChatResponseUpdate, Content
+from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message
 from pydantic import BaseModel
 
 
@@ -16,12 +16,12 @@ async def test_agent_initialization_basic(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent[ChatOptions](
-        chat_client=streaming_chat_client_stub(stream_fn),
+    agent = Agent[ChatOptions](
+        client=streaming_chat_client_stub(stream_fn),
         name="test_agent",
         instructions="Test",
     )
@@ -38,11 +38,11 @@ async def test_agent_initialization_with_state_schema(streaming_chat_client_stub
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     state_schema: dict[str, dict[str, Any]] = {"document": {"type": "string"}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
@@ -54,11 +54,11 @@ async def test_agent_initialization_with_predict_state_config(streaming_chat_cli
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
     wrapper = AgentFrameworkAgent(agent=agent, predict_state_config=predict_config)
 
@@ -70,7 +70,7 @@ async def test_agent_initialization_with_pydantic_state_schema(streaming_chat_cl
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
@@ -78,7 +78,7 @@ async def test_agent_initialization_with_pydantic_state_schema(streaming_chat_cl
         document: str
         tags: list[str] = []
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
 
     wrapper_class_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState)
     wrapper_instance_schema = AgentFrameworkAgent(agent=agent, state_schema=MyState(document="hi"))
@@ -93,11 +93,11 @@ async def test_run_started_event_emission(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -117,11 +117,11 @@ async def test_predict_state_custom_event_emission(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     predict_config = {
         "document": {"tool": "write_doc", "tool_argument": "content"},
         "summary": {"tool": "summarize", "tool_argument": "text"},
@@ -149,11 +149,11 @@ async def test_initial_state_snapshot_with_schema(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     state_schema = {"document": {"type": "string"}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
@@ -179,11 +179,11 @@ async def test_state_initialization_object_type(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     state_schema: dict[str, dict[str, Any]] = {"recipe": {"type": "object", "properties": {}}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
@@ -206,11 +206,11 @@ async def test_state_initialization_array_type(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     state_schema: dict[str, dict[str, Any]] = {"steps": {"type": "array", "items": {}}}
     wrapper = AgentFrameworkAgent(agent=agent, state_schema=state_schema)
 
@@ -233,11 +233,11 @@ async def test_run_finished_event_emission(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -255,11 +255,11 @@ async def test_tool_result_confirm_changes_accepted(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Document updated")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(
         agent=agent,
         state_schema={"document": {"type": "string"}},
@@ -302,11 +302,11 @@ async def test_tool_result_confirm_changes_rejected(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Simulate tool result message with rejection
@@ -336,11 +336,11 @@ async def test_tool_result_function_approval_accepted(streaming_chat_client_stub
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Simulate tool result with multiple steps
@@ -382,11 +382,11 @@ async def test_tool_result_function_approval_rejected(streaming_chat_client_stub
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="OK")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Simulate tool result rejection with steps
@@ -425,13 +425,13 @@ async def test_thread_metadata_tracking(streaming_chat_client_stub):
     captured_options: dict[str, Any] = {}
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         # Capture options to verify internal keys are NOT passed to chat client
         captured_options.update(options)
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data = {
@@ -445,7 +445,7 @@ async def test_thread_metadata_tracking(streaming_chat_client_stub):
         events.append(event)
 
     # AG-UI internal metadata should be stored in thread.metadata
-    thread = agent.chat_client.last_thread
+    thread = agent.client.last_thread
     thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
     assert thread_metadata.get("ag_ui_thread_id") == "test_thread_123"
     assert thread_metadata.get("ag_ui_run_id") == "test_run_456"
@@ -467,13 +467,13 @@ async def test_state_context_injection(streaming_chat_client_stub):
     captured_options: dict[str, Any] = {}
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         # Capture options to verify internal keys are NOT passed to chat client
         captured_options.update(options)
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(
         agent=agent,
         state_schema={"document": {"type": "string"}},
@@ -489,7 +489,7 @@ async def test_state_context_injection(streaming_chat_client_stub):
         events.append(event)
 
     # Current state should be stored in thread.metadata
-    thread = agent.chat_client.last_thread
+    thread = agent.client.last_thread
     thread_metadata = thread.metadata if thread and hasattr(thread, "metadata") else {}
     current_state = thread_metadata.get("current_state")
     if isinstance(current_state, str):
@@ -506,11 +506,11 @@ async def test_no_messages_provided(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data: dict[str, Any] = {"messages": []}
@@ -530,11 +530,11 @@ async def test_message_end_event_emission(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text="Hello world")])
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -558,13 +558,13 @@ async def test_error_handling_with_exception(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         if False:
             yield ChatResponseUpdate(contents=[])
         raise RuntimeError("Simulated failure")
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     input_data: dict[str, Any] = {"messages": [{"role": "user", "content": "Hi"}]}
@@ -579,13 +579,13 @@ async def test_json_decode_error_in_tool_result(streaming_chat_client_stub):
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: dict[str, Any], **kwargs: Any
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         if False:
             yield ChatResponseUpdate(contents=[])
         raise AssertionError("ChatClient should not be called with orphaned tool result")
 
-    agent = ChatAgent(name="test_agent", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test_agent", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent)
 
     # Send invalid JSON as tool result without preceding tool call
@@ -618,13 +618,13 @@ async def test_agent_with_use_service_thread_is_false(streaming_chat_client_stub
     request_service_thread_id: str | None = None
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], chat_options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(
             contents=[Content.from_text(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
         )
 
-    agent = ChatAgent(chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=False)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
@@ -642,7 +642,7 @@ async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub)
     request_service_thread_id: str | None = None
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], chat_options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], chat_options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         nonlocal request_service_thread_id
         thread = kwargs.get("thread")
@@ -651,7 +651,7 @@ async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub)
             contents=[Content.from_text(text="Response")], response_id="resp_67890", conversation_id="conv_12345"
         )
 
-    agent = ChatAgent(chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(client=streaming_chat_client_stub(stream_fn))
     wrapper = AgentFrameworkAgent(agent=agent, use_service_thread=True)
 
     input_data = {"messages": [{"role": "user", "content": "Hi"}], "thread_id": "conv_123456"}
@@ -659,7 +659,7 @@ async def test_agent_with_use_service_thread_is_true(streaming_chat_client_stub)
     events: list[Any] = []
     async for event in wrapper.run_agent(input_data):
         events.append(event)
-    request_service_thread_id = agent.chat_client.last_service_thread_id
+    request_service_thread_id = agent.client.last_service_thread_id
     assert request_service_thread_id == "conv_123456"  # type: ignore[attr-defined] (service_thread_id should be set)
 
 
@@ -679,15 +679,15 @@ async def test_function_approval_mode_executes_tool(streaming_chat_client_stub):
         return "2025/12/01 12:00:00"
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         # Capture the messages received by the chat client
         messages_received.clear()
         messages_received.extend(messages)
         yield ChatResponseUpdate(contents=[Content.from_text(text="Processing completed")])
 
-    agent = ChatAgent(
-        chat_client=streaming_chat_client_stub(stream_fn),
+    agent = Agent(
+        client=streaming_chat_client_stub(stream_fn),
         name="test_agent",
         instructions="Test",
         tools=[get_datetime],
@@ -770,17 +770,17 @@ async def test_function_approval_mode_rejection(streaming_chat_client_stub):
         return "All data deleted"
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         # Capture the messages received by the chat client
         messages_received.clear()
         messages_received.extend(messages)
         yield ChatResponseUpdate(contents=[Content.from_text(text="Operation cancelled")])
 
-    agent = ChatAgent(
+    agent = Agent(
         name="test_agent",
         instructions="Test",
-        chat_client=streaming_chat_client_stub(stream_fn),
+        client=streaming_chat_client_stub(stream_fn),
         tools=[delete_all_data],
     )
     wrapper = AgentFrameworkAgent(agent=agent)

python/packages/ag-ui/tests/ag_ui/test_endpoint.py

@@ -5,7 +5,7 @@
 import json
 
 import pytest
-from agent_framework import ChatAgent, ChatResponseUpdate, Content
+from agent_framework import Agent, ChatResponseUpdate, Content
 from fastapi import FastAPI, Header, HTTPException
 from fastapi.params import Depends
 from fastapi.testclient import TestClient
@@ -28,7 +28,7 @@ def build_chat_client(streaming_chat_client_stub, stream_from_updates_fixture):
 async def test_add_endpoint_with_agent_protocol(build_chat_client):
     """Test adding endpoint with raw SupportsAgentRun."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/test-agent")
 
@@ -42,7 +42,7 @@ async def test_add_endpoint_with_agent_protocol(build_chat_client):
 async def test_add_endpoint_with_wrapped_agent(build_chat_client):
     """Test adding endpoint with pre-wrapped AgentFrameworkAgent."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
     wrapped_agent = AgentFrameworkAgent(agent=agent, name="wrapped")
 
     add_agent_framework_fastapi_endpoint(app, wrapped_agent, path="/wrapped-agent")
@@ -57,7 +57,7 @@ async def test_add_endpoint_with_wrapped_agent(build_chat_client):
 async def test_endpoint_with_state_schema(build_chat_client):
     """Test endpoint with state_schema parameter."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
     state_schema = {"document": {"type": "string"}}
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/stateful", state_schema=state_schema)
@@ -73,7 +73,7 @@ async def test_endpoint_with_state_schema(build_chat_client):
 async def test_endpoint_with_default_state_seed(build_chat_client):
     """Test endpoint seeds default state when client omits it."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
     state_schema = {"proverbs": {"type": "array"}}
     default_state = {"proverbs": ["Keep the original."]}
 
@@ -100,7 +100,7 @@ async def test_endpoint_with_default_state_seed(build_chat_client):
 async def test_endpoint_with_predict_state_config(build_chat_client):
     """Test endpoint with predict_state_config parameter."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
     predict_config = {"document": {"tool": "write_doc", "tool_argument": "content"}}
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/predictive", predict_state_config=predict_config)
@@ -114,7 +114,7 @@ async def test_endpoint_with_predict_state_config(build_chat_client):
 async def test_endpoint_request_logging(build_chat_client):
     """Test that endpoint logs request details."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/logged")
 
@@ -134,7 +134,7 @@ async def test_endpoint_request_logging(build_chat_client):
 async def test_endpoint_event_streaming(build_chat_client):
     """Test that endpoint streams events correctly."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client("Streamed response"))
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client("Streamed response"))
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/stream")
 
@@ -168,7 +168,7 @@ async def test_endpoint_event_streaming(build_chat_client):
 async def test_endpoint_error_handling(build_chat_client):
     """Test endpoint error handling during request parsing."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/failing")
 
@@ -184,8 +184,8 @@ async def test_endpoint_error_handling(build_chat_client):
 async def test_endpoint_multiple_paths(build_chat_client):
     """Test adding multiple endpoints with different paths."""
     app = FastAPI()
-    agent1 = ChatAgent(name="agent1", instructions="First agent", chat_client=build_chat_client("Response 1"))
-    agent2 = ChatAgent(name="agent2", instructions="Second agent", chat_client=build_chat_client("Response 2"))
+    agent1 = Agent(name="agent1", instructions="First agent", client=build_chat_client("Response 1"))
+    agent2 = Agent(name="agent2", instructions="Second agent", client=build_chat_client("Response 2"))
 
     add_agent_framework_fastapi_endpoint(app, agent1, path="/agent1")
     add_agent_framework_fastapi_endpoint(app, agent2, path="/agent2")
@@ -202,7 +202,7 @@ async def test_endpoint_multiple_paths(build_chat_client):
 async def test_endpoint_default_path(build_chat_client):
     """Test endpoint with default path."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent)
 
@@ -215,7 +215,7 @@ async def test_endpoint_default_path(build_chat_client):
 async def test_endpoint_response_headers(build_chat_client):
     """Test that endpoint sets correct response headers."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/headers")
 
@@ -231,7 +231,7 @@ async def test_endpoint_response_headers(build_chat_client):
 async def test_endpoint_empty_messages(build_chat_client):
     """Test endpoint with empty messages list."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/empty")
 
@@ -244,7 +244,7 @@ async def test_endpoint_empty_messages(build_chat_client):
 async def test_endpoint_complex_input(build_chat_client):
     """Test endpoint with complex input data."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/complex")
 
@@ -269,7 +269,7 @@ async def test_endpoint_complex_input(build_chat_client):
 async def test_endpoint_openapi_schema(build_chat_client):
     """Test that endpoint generates proper OpenAPI schema with request model."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/schema-test")
 
@@ -313,7 +313,7 @@ async def test_endpoint_openapi_schema(build_chat_client):
 async def test_endpoint_default_tags(build_chat_client):
     """Test that endpoint uses default 'AG-UI' tag."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/default-tags")
 
@@ -331,7 +331,7 @@ async def test_endpoint_default_tags(build_chat_client):
 async def test_endpoint_custom_tags(build_chat_client):
     """Test that endpoint accepts custom tags."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/custom-tags", tags=["Custom", "Agent"])
 
@@ -349,7 +349,7 @@ async def test_endpoint_custom_tags(build_chat_client):
 async def test_endpoint_missing_required_field(build_chat_client):
     """Test that endpoint validates required fields with Pydantic."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     add_agent_framework_fastapi_endpoint(app, agent, path="/validation")
 
@@ -368,7 +368,7 @@ async def test_endpoint_internal_error_handling(build_chat_client):
     from unittest.mock import patch
 
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     # Use default_state to trigger the code path that can raise an exception
     add_agent_framework_fastapi_endpoint(app, agent, path="/error-test", default_state={"key": "value"})
@@ -387,7 +387,7 @@ async def test_endpoint_internal_error_handling(build_chat_client):
 async def test_endpoint_with_dependencies_blocks_unauthorized(build_chat_client):
     """Test that endpoint blocks requests when authentication dependency fails."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     async def require_api_key(x_api_key: str | None = Header(None)):
         if x_api_key != "secret-key":
@@ -406,7 +406,7 @@ async def test_endpoint_with_dependencies_blocks_unauthorized(build_chat_client)
 async def test_endpoint_with_dependencies_allows_authorized(build_chat_client):
     """Test that endpoint allows requests when authentication dependency passes."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     async def require_api_key(x_api_key: str | None = Header(None)):
         if x_api_key != "secret-key":
@@ -429,7 +429,7 @@ async def test_endpoint_with_dependencies_allows_authorized(build_chat_client):
 async def test_endpoint_with_multiple_dependencies(build_chat_client):
     """Test that endpoint supports multiple dependencies."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     execution_order: list[str] = []
 
@@ -457,7 +457,7 @@ async def test_endpoint_with_multiple_dependencies(build_chat_client):
 async def test_endpoint_without_dependencies_is_accessible(build_chat_client):
     """Test that endpoint without dependencies remains accessible (backward compatibility)."""
     app = FastAPI()
-    agent = ChatAgent(name="test", instructions="Test agent", chat_client=build_chat_client())
+    agent = Agent(name="test", instructions="Test agent", client=build_chat_client())
 
     # No dependencies parameter - should be accessible without auth
     add_agent_framework_fastapi_endpoint(app, agent, path="/open")

python/packages/ag-ui/tests/ag_ui/test_helpers.py

@@ -2,7 +2,7 @@
 
 """Tests for orchestration helper functions."""
 
-from agent_framework import ChatMessage, Content
+from agent_framework import Content, Message
 
 from agent_framework_ag_ui._orchestration._helpers import (
     approval_steps,
@@ -29,8 +29,8 @@ class TestPendingToolCallIds:
     def test_no_tool_calls(self):
         """Returns empty set when no tool calls in messages."""
         messages = [
-            ChatMessage(role="user", contents=[Content.from_text("Hello")]),
-            ChatMessage(role="assistant", contents=[Content.from_text("Hi there")]),
+            Message(role="user", contents=[Content.from_text("Hello")]),
+            Message(role="assistant", contents=[Content.from_text("Hi there")]),
         ]
         result = pending_tool_call_ids(messages)
         assert result == set()
@@ -38,7 +38,7 @@ class TestPendingToolCallIds:
     def test_pending_tool_call(self):
         """Returns pending tool call ID when no result exists."""
         messages = [
-            ChatMessage(
+            Message(
                 role="assistant",
                 contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
             ),
@@ -49,11 +49,11 @@ class TestPendingToolCallIds:
     def test_resolved_tool_call(self):
         """Returns empty set when tool call has result."""
         messages = [
-            ChatMessage(
+            Message(
                 role="assistant",
                 contents=[Content.from_function_call(call_id="call_123", name="get_weather", arguments="{}")],
             ),
-            ChatMessage(
+            Message(
                 role="tool",
                 contents=[Content.from_function_result(call_id="call_123", result="sunny")],
             ),
@@ -64,7 +64,7 @@ class TestPendingToolCallIds:
     def test_multiple_tool_calls_some_resolved(self):
         """Returns only unresolved tool call IDs."""
         messages = [
-            ChatMessage(
+            Message(
                 role="assistant",
                 contents=[
                     Content.from_function_call(call_id="call_1", name="tool_a", arguments="{}"),
@@ -72,11 +72,11 @@ class TestPendingToolCallIds:
                     Content.from_function_call(call_id="call_3", name="tool_c", arguments="{}"),
                 ],
             ),
-            ChatMessage(
+            Message(
                 role="tool",
                 contents=[Content.from_function_result(call_id="call_1", result="result_a")],
             ),
-            ChatMessage(
+            Message(
                 role="tool",
                 contents=[Content.from_function_result(call_id="call_3", result="result_c")],
             ),
@@ -90,7 +90,7 @@ class TestIsStateContextMessage:
 
     def test_state_context_message(self):
         """Returns True for state context message."""
-        message = ChatMessage(
+        message = Message(
             role="system",
             contents=[Content.from_text("Current state of the application: {}")],
         )
@@ -98,7 +98,7 @@ class TestIsStateContextMessage:
 
     def test_non_system_message(self):
         """Returns False for non-system message."""
-        message = ChatMessage(
+        message = Message(
             role="user",
             contents=[Content.from_text("Current state of the application: {}")],
         )
@@ -106,7 +106,7 @@ class TestIsStateContextMessage:
 
     def test_system_message_without_state_prefix(self):
         """Returns False for system message without state prefix."""
-        message = ChatMessage(
+        message = Message(
             role="system",
             contents=[Content.from_text("You are a helpful assistant.")],
         )
@@ -114,7 +114,7 @@ class TestIsStateContextMessage:
 
     def test_empty_contents(self):
         """Returns False for message with empty contents."""
-        message = ChatMessage(role="system", contents=[])
+        message = Message(role="system", contents=[])
         assert is_state_context_message(message) is False
 
 
@@ -342,7 +342,7 @@ class TestLatestApprovalResponse:
     def test_no_approval_response(self):
         """Returns None when no approval response in last message."""
         messages = [
-            ChatMessage(role="assistant", contents=[Content.from_text("Hello")]),
+            Message(role="assistant", contents=[Content.from_text("Hello")]),
         ]
         result = latest_approval_response(messages)
         assert result is None
@@ -357,7 +357,7 @@ class TestLatestApprovalResponse:
             function_call=fc,
         )
         messages = [
-            ChatMessage(role="user", contents=[approval_content]),
+            Message(role="user", contents=[approval_content]),
         ]
         result = latest_approval_response(messages)
         assert result is approval_content

python/packages/ag-ui/tests/ag_ui/test_message_adapters.py

@@ -5,7 +5,7 @@
 import json
 
 import pytest
-from agent_framework import ChatMessage, Content
+from agent_framework import Content, Message
 
 from agent_framework_ag_ui._message_adapters import (
     agent_framework_messages_to_agui,
@@ -24,7 +24,7 @@ def sample_agui_message():
 @pytest.fixture
 def sample_agent_framework_message():
     """Create a sample Agent Framework message."""
-    return ChatMessage(role="user", contents=[Content.from_text(text="Hello")], message_id="msg-123")
+    return Message(role="user", contents=[Content.from_text(text="Hello")], message_id="msg-123")
 
 
 def test_agui_to_agent_framework_basic(sample_agui_message):
@@ -100,7 +100,7 @@ def test_agui_tool_result_to_agent_framework():
 def test_agui_tool_approval_updates_tool_call_arguments():
     """Tool approval updates matching tool call arguments for snapshots and agent context.
 
-    The LLM context (ChatMessage) should contain only enabled steps, so the LLM
+    The LLM context (Message) should contain only enabled steps, so the LLM
     generates responses based on what was actually approved/executed.
 
     The raw messages (for MESSAGES_SNAPSHOT) should contain all steps with status,
@@ -446,7 +446,7 @@ def test_agui_with_tool_calls_to_agent_framework():
 
 def test_agent_framework_to_agui_with_tool_calls():
     """Test converting Agent Framework message with tool calls to AG-UI."""
-    msg = ChatMessage(
+    msg = Message(
         role="assistant",
         contents=[
             Content.from_text(text="Calling tool"),
@@ -471,7 +471,7 @@ def test_agent_framework_to_agui_with_tool_calls():
 
 def test_agent_framework_to_agui_multiple_text_contents():
     """Test concatenating multiple text contents."""
-    msg = ChatMessage(
+    msg = Message(
         role="assistant",
         contents=[Content.from_text(text="Part 1 "), Content.from_text(text="Part 2")],
     )
@@ -484,7 +484,7 @@ def test_agent_framework_to_agui_multiple_text_contents():
 
 def test_agent_framework_to_agui_no_message_id():
     """Test message without message_id - should auto-generate ID."""
-    msg = ChatMessage(role="user", contents=[Content.from_text(text="Hello")])
+    msg = Message(role="user", contents=[Content.from_text(text="Hello")])
 
     messages = agent_framework_messages_to_agui([msg])
 
@@ -496,7 +496,7 @@ def test_agent_framework_to_agui_no_message_id():
 
 def test_agent_framework_to_agui_system_role():
     """Test system role conversion."""
-    msg = ChatMessage(role="system", contents=[Content.from_text(text="System")])
+    msg = Message(role="system", contents=[Content.from_text(text="System")])
 
     messages = agent_framework_messages_to_agui([msg])
 
@@ -541,7 +541,7 @@ def test_extract_text_from_custom_contents():
 
 def test_agent_framework_to_agui_function_result_dict():
     """Test converting FunctionResultContent with dict result to AG-UI."""
-    msg = ChatMessage(
+    msg = Message(
         role="tool",
         contents=[Content.from_function_result(call_id="call-123", result={"key": "value", "count": 42})],
         message_id="msg-789",
@@ -558,7 +558,7 @@ def test_agent_framework_to_agui_function_result_dict():
 
 def test_agent_framework_to_agui_function_result_none():
     """Test converting FunctionResultContent with None result to AG-UI."""
-    msg = ChatMessage(
+    msg = Message(
         role="tool",
         contents=[Content.from_function_result(call_id="call-123", result=None)],
         message_id="msg-789",
@@ -574,7 +574,7 @@ def test_agent_framework_to_agui_function_result_none():
 
 def test_agent_framework_to_agui_function_result_string():
     """Test converting FunctionResultContent with string result to AG-UI."""
-    msg = ChatMessage(
+    msg = Message(
         role="tool",
         contents=[Content.from_function_result(call_id="call-123", result="plain text result")],
         message_id="msg-789",
@@ -589,7 +589,7 @@ def test_agent_framework_to_agui_function_result_string():
 
 def test_agent_framework_to_agui_function_result_empty_list():
     """Test converting FunctionResultContent with empty list result to AG-UI."""
-    msg = ChatMessage(
+    msg = Message(
         role="tool",
         contents=[Content.from_function_result(call_id="call-123", result=[])],
         message_id="msg-789",
@@ -611,7 +611,7 @@ def test_agent_framework_to_agui_function_result_single_text_content():
     class MockTextContent:
         text: str
 
-    msg = ChatMessage(
+    msg = Message(
         role="tool",
         contents=[Content.from_function_result(call_id="call-123", result=[MockTextContent("Hello from MCP!")])],
         message_id="msg-789",
@@ -633,7 +633,7 @@ def test_agent_framework_to_agui_function_result_multiple_text_contents():
     class MockTextContent:
         text: str
 
-    msg = ChatMessage(
+    msg = Message(
         role="tool",
         contents=[
             Content.from_function_result(

python/packages/ag-ui/tests/ag_ui/test_message_hygiene.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from agent_framework import ChatMessage, Content
+from agent_framework import Content, Message
 
 from agent_framework_ag_ui._message_adapters import _deduplicate_messages, _sanitize_tool_history
 
@@ -13,7 +13,7 @@ def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> Non
     tool for the approval UI flow that shouldn't be sent to the LLM.
     """
     messages = [
-        ChatMessage(
+        Message(
             role="assistant",
             contents=[
                 Content.from_function_call(
@@ -23,7 +23,7 @@ def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> Non
                 )
             ],
         ),
-        ChatMessage(
+        Message(
             role="user",
             contents=[Content.from_text(text='{"accepted": true}')],
         ),
@@ -44,11 +44,11 @@ def test_sanitize_tool_history_filters_out_confirm_changes_only_message() -> Non
 
 def test_deduplicate_messages_prefers_non_empty_tool_results() -> None:
     messages = [
-        ChatMessage(
+        Message(
             role="tool",
             contents=[Content.from_function_result(call_id="call1", result="")],
         ),
-        ChatMessage(
+        Message(
             role="tool",
             contents=[Content.from_function_result(call_id="call1", result="result data")],
         ),
@@ -71,13 +71,13 @@ def test_convert_approval_results_to_tool_messages() -> None:
     # Simulate what happens after _resolve_approval_responses:
     # A user message contains function_result content (the executed tool result)
     messages = [
-        ChatMessage(
+        Message(
             role="assistant",
             contents=[
                 Content.from_function_call(call_id="call_123", name="my_mcp_tool", arguments="{}"),
             ],
         ),
-        ChatMessage(
+        Message(
             role="user",
             contents=[
                 Content.from_function_result(call_id="call_123", result="tool execution result"),
@@ -109,13 +109,13 @@ def test_convert_approval_results_preserves_other_user_content() -> None:
     from agent_framework_ag_ui._run import _convert_approval_results_to_tool_messages
 
     messages = [
-        ChatMessage(
+        Message(
             role="assistant",
             contents=[
                 Content.from_function_call(call_id="call_123", name="my_tool", arguments="{}"),
             ],
         ),
-        ChatMessage(
+        Message(
             role="user",
             contents=[
                 Content.from_text(text="User also said something"),
@@ -152,12 +152,12 @@ def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> No
     """
     messages = [
         # User asks something
-        ChatMessage(
+        Message(
             role="user",
             contents=[Content.from_text(text="What time is it?")],
         ),
         # Assistant calls MCP tool + confirm_changes
-        ChatMessage(
+        Message(
             role="assistant",
             contents=[
                 Content.from_function_call(call_id="call_1", name="get_datetime", arguments="{}"),
@@ -165,12 +165,12 @@ def test_sanitize_tool_history_filters_confirm_changes_keeps_other_tools() -> No
             ],
         ),
         # Tool result for the actual MCP tool
-        ChatMessage(
+        Message(
             role="tool",
             contents=[Content.from_function_result(call_id="call_1", result="2024-01-01 12:00:00")],
         ),
         # User asks something else
-        ChatMessage(
+        Message(
             role="user",
             contents=[Content.from_text(text="What's the date?")],
         ),
@@ -204,12 +204,12 @@ def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages()
     respond with "Here's your 5-step plan" instead of "Here's your 2-step plan".
     """
     messages = [
-        ChatMessage(
+        Message(
             role="user",
             contents=[Content.from_text(text="Build a robot")],
         ),
         # Assistant message with both generate_task_steps and confirm_changes
-        ChatMessage(
+        Message(
             role="assistant",
             contents=[
                 Content.from_function_call(
@@ -225,7 +225,7 @@ def test_sanitize_tool_history_filters_confirm_changes_from_assistant_messages()
             ],
         ),
         # Approval response
-        ChatMessage(
+        Message(
             role="user",
             contents=[
                 Content.from_function_approval_response(

python/packages/ag-ui/tests/ag_ui/test_run.py

@@ -6,7 +6,7 @@ from ag_ui.core import (
     TextMessageEndEvent,
     TextMessageStartEvent,
 )
-from agent_framework import ChatMessage, Content
+from agent_framework import Content, Message
 
 from agent_framework_ag_ui._run import (
     FlowState,
@@ -212,7 +212,7 @@ class TestInjectStateContext:
 
     def test_no_state_message(self):
         """Returns original messages when no state context needed."""
-        messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
+        messages = [Message(role="user", contents=[Content.from_text("Hello")])]
         result = _inject_state_context(messages, {}, {})
         assert result == messages
 
@@ -224,8 +224,8 @@ class TestInjectStateContext:
     def test_last_message_not_user(self):
         """Returns original messages when last message is not from user."""
         messages = [
-            ChatMessage(role="user", contents=[Content.from_text("Hello")]),
-            ChatMessage(role="assistant", contents=[Content.from_text("Hi")]),
+            Message(role="user", contents=[Content.from_text("Hello")]),
+            Message(role="assistant", contents=[Content.from_text("Hi")]),
         ]
         state = {"key": "value"}
         schema = {"properties": {"key": {"type": "string"}}}
@@ -237,8 +237,8 @@ class TestInjectStateContext:
         """Injects state context before last user message."""
 
         messages = [
-            ChatMessage(role="system", contents=[Content.from_text("You are helpful")]),
-            ChatMessage(role="user", contents=[Content.from_text("Hello")]),
+            Message(role="system", contents=[Content.from_text("You are helpful")]),
+            Message(role="user", contents=[Content.from_text("Hello")]),
         ]
         state = {"document": "content"}
         schema = {"properties": {"document": {"type": "string"}}}
@@ -405,7 +405,7 @@ def test_extract_approved_state_updates_no_handler():
     """Test _extract_approved_state_updates returns empty with no handler."""
     from agent_framework_ag_ui._run import _extract_approved_state_updates
 
-    messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
+    messages = [Message(role="user", contents=[Content.from_text("Hello")])]
     result = _extract_approved_state_updates(messages, None)
     assert result == {}
 
@@ -416,7 +416,7 @@ def test_extract_approved_state_updates_no_approval():
     from agent_framework_ag_ui._run import _extract_approved_state_updates
 
     handler = PredictiveStateHandler(predict_state_config={"doc": {"tool": "write", "tool_argument": "content"}})
-    messages = [ChatMessage(role="user", contents=[Content.from_text("Hello")])]
+    messages = [Message(role="user", contents=[Content.from_text("Hello")])]
     result = _extract_approved_state_updates(messages, handler)
     assert result == {}
 

python/packages/ag-ui/tests/ag_ui/test_structured_output.py

@@ -6,7 +6,7 @@ import json
 from collections.abc import AsyncIterator, MutableSequence
 from typing import Any
 
-from agent_framework import ChatAgent, ChatMessage, ChatOptions, ChatResponseUpdate, Content
+from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message
 from pydantic import BaseModel
 
 
@@ -35,13 +35,13 @@ async def test_structured_output_with_recipe(streaming_chat_client_stub, stream_
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(
             contents=[Content.from_text(text='{"recipe": {"name": "Pasta"}, "message": "Here is your recipe"}')]
         )
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     agent.default_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -73,7 +73,7 @@ async def test_structured_output_with_steps(streaming_chat_client_stub, stream_f
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         steps_data = {
             "steps": [
@@ -83,7 +83,7 @@ async def test_structured_output_with_steps(streaming_chat_client_stub, stream_f
         }
         yield ChatResponseUpdate(contents=[Content.from_text(text=json.dumps(steps_data))])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     agent.default_options = ChatOptions(response_format=StepsOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -116,8 +116,8 @@ async def test_structured_output_with_no_schema_match(streaming_chat_client_stub
         ChatResponseUpdate(contents=[Content.from_text(text='{"data": {"key": "value"}}')]),
     ]
 
-    agent = ChatAgent(
-        name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_from_updates_fixture(updates))
+    agent = Agent(
+        name="test", instructions="Test", client=streaming_chat_client_stub(stream_from_updates_fixture(updates))
     )
     agent.default_options = ChatOptions(response_format=GenericOutput)
 
@@ -149,11 +149,11 @@ async def test_structured_output_without_schema(streaming_chat_client_stub, stre
         info: str
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         yield ChatResponseUpdate(contents=[Content.from_text(text='{"data": {"key": "value"}, "info": "processed"}')])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     agent.default_options = ChatOptions(response_format=DataOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -182,10 +182,10 @@ async def test_no_structured_output_when_no_response_format(streaming_chat_clien
 
     updates = [ChatResponseUpdate(contents=[Content.from_text(text="Regular text")])]
 
-    agent = ChatAgent(
+    agent = Agent(
         name="test",
         instructions="Test",
-        chat_client=streaming_chat_client_stub(stream_from_updates_fixture(updates)),
+        client=streaming_chat_client_stub(stream_from_updates_fixture(updates)),
     )
     # No response_format set
 
@@ -208,12 +208,12 @@ async def test_structured_output_with_message_field(streaming_chat_client_stub,
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         output_data = {"recipe": {"name": "Salad"}, "message": "Fresh salad recipe ready"}
         yield ChatResponseUpdate(contents=[Content.from_text(text=json.dumps(output_data))])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     agent.default_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(
@@ -243,12 +243,12 @@ async def test_empty_updates_no_structured_processing(streaming_chat_client_stub
     from agent_framework.ag_ui import AgentFrameworkAgent
 
     async def stream_fn(
-        messages: MutableSequence[ChatMessage], options: ChatOptions, **kwargs: Any
+        messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
     ) -> AsyncIterator[ChatResponseUpdate]:
         if False:
             yield ChatResponseUpdate(contents=[])
 
-    agent = ChatAgent(name="test", instructions="Test", chat_client=streaming_chat_client_stub(stream_fn))
+    agent = Agent(name="test", instructions="Test", client=streaming_chat_client_stub(stream_fn))
     agent.default_options = ChatOptions(response_format=RecipeOutput)
 
     wrapper = AgentFrameworkAgent(agent=agent)

python/packages/ag-ui/tests/ag_ui/test_tooling.py

@@ -2,7 +2,7 @@
 
 from unittest.mock import MagicMock
 
-from agent_framework import ChatAgent, tool
+from agent_framework import Agent, tool
 
 from agent_framework_ag_ui._orchestration._tooling import (
     collect_server_tools,
@@ -31,14 +31,14 @@ def regular_tool() -> str:
     return "result"
 
 
-def _create_chat_agent_with_tool(tool_name: str = "regular_tool") -> ChatAgent:
-    """Create a ChatAgent with a mocked chat client and a simple tool.
+def _create_chat_agent_with_tool(tool_name: str = "regular_tool") -> Agent:
+    """Create a Agent with a mocked chat client and a simple tool.
 
     Note: tool_name parameter is kept for API compatibility but the tool
     will always be named 'regular_tool' since tool uses the function name.
     """
     mock_chat_client = MagicMock()
-    return ChatAgent(chat_client=mock_chat_client, tools=[regular_tool])
+    return Agent(client=mock_chat_client, tools=[regular_tool])
 
 
 def test_merge_tools_filters_duplicates() -> None:
@@ -59,7 +59,7 @@ def test_register_additional_client_tools_assigns_when_configured() -> None:
     mock_chat_client = MagicMock(spec=BaseChatClient)
     mock_chat_client.function_invocation_configuration = normalize_function_invocation_configuration(None)
 
-    agent = ChatAgent(chat_client=mock_chat_client)
+    agent = Agent(client=mock_chat_client)
 
     tools = [DummyTool("x")]
     register_additional_client_tools(agent, tools)
@@ -148,14 +148,14 @@ def test_collect_server_tools_no_default_options() -> None:
 def test_register_additional_client_tools_no_tools() -> None:
     """register_additional_client_tools does nothing with None tools."""
     mock_chat_client = MagicMock()
-    agent = ChatAgent(chat_client=mock_chat_client)
+    agent = Agent(client=mock_chat_client)
 
     # Should not raise
     register_additional_client_tools(agent, None)
 
 
 def test_register_additional_client_tools_no_chat_client() -> None:
-    """register_additional_client_tools does nothing when agent has no chat_client."""
+    """register_additional_client_tools does nothing when agent has no client."""
     from agent_framework_ag_ui._orchestration._tooling import register_additional_client_tools
 
     class MockAgent:

python/packages/ag-ui/tests/ag_ui/test_utils.py

@@ -404,11 +404,11 @@ def test_safe_json_parse_with_none():
 
 def test_get_role_value_with_enum():
     """Test get_role_value with enum role."""
-    from agent_framework import ChatMessage, Content
+    from agent_framework import Content, Message
 
     from agent_framework_ag_ui._utils import get_role_value
 
-    message = ChatMessage(role="user", contents=[Content.from_text("test")])
+    message = Message(role="user", contents=[Content.from_text("test")])
     result = get_role_value(message)
     assert result == "user"