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>
2026-06-16 21:04:09 +08:00 · 2026-02-11 00:04:32 +01:00
parent a4c9e43afb
commit 0521f5bed8
418 changed files with 5385 additions and 5389 deletions
@@ -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]

@@ -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}")],
            ),
@@ -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)
@@ -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")
@@ -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
@@ -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(
@@ -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(
@@ -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 == {}

@@ -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)
@@ -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:
@@ -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"