diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx
new file mode 100644
index 00000000000..70c3fdd4f29
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx
@@ -0,0 +1,296 @@
+import { render, screen } from "@testing-library/react";
+import { describe, it, expect, vi } from "vitest";
+import ChatMessageBubble from "./ChatMessageBubble";
+import { EndpointType } from "./mode_endpoint_mapping";
+import { MessageType } from "./types";
+
+// Mock child components to isolate bubble rendering logic
+vi.mock("react-markdown", () => ({
+  default: ({ children }: { children: string }) => <div data-testid="react-markdown">{children}</div>,
+}));
+
+vi.mock("react-syntax-highlighter", () => ({
+  Prism: ({ children }: { children: string }) => <pre data-testid="syntax-highlighter">{children}</pre>,
+}));
+
+vi.mock("react-syntax-highlighter/dist/esm/styles/prism", () => ({
+  coy: {},
+}));
+
+vi.mock("./ReasoningContent", () => ({
+  default: ({ reasoningContent }: { reasoningContent: string }) => (
+    <div data-testid="reasoning-content">{reasoningContent}</div>
+  ),
+}));
+
+vi.mock("./MCPEventsDisplay", () => ({
+  default: ({ events }: { events: unknown[] }) => (
+    <div data-testid="mcp-events-display">{events.length} events</div>
+  ),
+}));
+
+vi.mock("./SearchResultsDisplay", () => ({
+  SearchResultsDisplay: ({ searchResults }: { searchResults: unknown[] }) => (
+    <div data-testid="search-results-display">{searchResults.length} results</div>
+  ),
+}));
+
+vi.mock("./ResponseMetrics", () => ({
+  default: ({ timeToFirstToken }: { timeToFirstToken?: number }) => (
+    <div data-testid="response-metrics">TTFT: {timeToFirstToken}</div>
+  ),
+}));
+
+vi.mock("./A2AMetrics", () => ({
+  default: ({ a2aMetadata }: { a2aMetadata: unknown }) => (
+    <div data-testid="a2a-metrics">A2A</div>
+  ),
+}));
+
+vi.mock("./CodeInterpreterOutput", () => ({
+  default: ({ code }: { code: string }) => <div data-testid="code-interpreter-output">{code}</div>,
+}));
+
+vi.mock("./AudioRenderer", () => ({
+  default: ({ message }: { message: MessageType }) => (
+    <div data-testid="audio-renderer">{typeof message.content === "string" ? message.content : ""}</div>
+  ),
+}));
+
+vi.mock("./ResponsesImageRenderer", () => ({
+  default: () => <div data-testid="responses-image-renderer" />,
+}));
+
+vi.mock("./ChatImageRenderer", () => ({
+  default: () => <div data-testid="chat-image-renderer" />,
+}));
+
+const defaultProps = {
+  isLastMessage: false,
+  endpointType: EndpointType.CHAT,
+  mcpEvents: [],
+  codeInterpreterResult: null,
+  accessToken: "test-token",
+};
+
+describe("ChatMessageBubble", () => {
+  it("should render a user message with right-aligned text", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "user", content: "Hello" }}
+      />,
+    );
+
+    expect(screen.getByText("user")).toBeInTheDocument();
+    expect(screen.getByText("Hello")).toBeInTheDocument();
+  });
+
+  it("should render an assistant message with left-aligned text", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "assistant", content: "Hi there" }}
+      />,
+    );
+
+    expect(screen.getByText("assistant")).toBeInTheDocument();
+    expect(screen.getByText("Hi there")).toBeInTheDocument();
+  });
+
+  it("should show model badge for assistant messages when model is provided", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "assistant", content: "Reply", model: "gpt-4" }}
+      />,
+    );
+
+    expect(screen.getByText("gpt-4")).toBeInTheDocument();
+  });
+
+  it("should not show model badge for user messages even when model is set", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "user", content: "Hello", model: "gpt-4" }}
+      />,
+    );
+
+    expect(screen.queryByText("gpt-4")).not.toBeInTheDocument();
+  });
+
+  it("should render markdown content via ReactMarkdown", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "assistant", content: "**bold text**" }}
+      />,
+    );
+
+    expect(screen.getByTestId("react-markdown")).toHaveTextContent("**bold text**");
+  });
+
+  it("should render an image when isImage is true", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "assistant", content: "https://example.com/img.png", isImage: true }}
+      />,
+    );
+
+    expect(screen.getByAltText("Generated image")).toHaveAttribute("src", "https://example.com/img.png");
+  });
+
+  it("should render AudioRenderer when isAudio is true", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "assistant", content: "audio-url", isAudio: true }}
+      />,
+    );
+
+    expect(screen.getByTestId("audio-renderer")).toBeInTheDocument();
+  });
+
+  it("should show ReasoningContent when reasoningContent is present", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{ role: "assistant", content: "answer", reasoningContent: "thinking..." }}
+      />,
+    );
+
+    expect(screen.getByTestId("reasoning-content")).toHaveTextContent("thinking...");
+  });
+
+  it("should show MCP events on the last assistant message for RESPONSES endpoint", () => {
+    const mcpEvents = [{ type: "tool_call", item_id: "1" }];
+
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        isLastMessage={true}
+        endpointType={EndpointType.RESPONSES}
+        mcpEvents={mcpEvents as any}
+        message={{ role: "assistant", content: "response" }}
+      />,
+    );
+
+    expect(screen.getByTestId("mcp-events-display")).toHaveTextContent("1 events");
+  });
+
+  it("should show MCP events on the last assistant message for CHAT endpoint", () => {
+    const mcpEvents = [{ type: "tool_call", item_id: "1" }];
+
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        isLastMessage={true}
+        endpointType={EndpointType.CHAT}
+        mcpEvents={mcpEvents as any}
+        message={{ role: "assistant", content: "response" }}
+      />,
+    );
+
+    expect(screen.getByTestId("mcp-events-display")).toHaveTextContent("1 events");
+  });
+
+  it("should not show MCP events when isLastMessage is false", () => {
+    const mcpEvents = [{ type: "tool_call", item_id: "1" }];
+
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        isLastMessage={false}
+        endpointType={EndpointType.RESPONSES}
+        mcpEvents={mcpEvents as any}
+        message={{ role: "assistant", content: "response" }}
+      />,
+    );
+
+    expect(screen.queryByTestId("mcp-events-display")).not.toBeInTheDocument();
+  });
+
+  it("should show SearchResultsDisplay when searchResults are present", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{
+          role: "assistant",
+          content: "found results",
+          searchResults: [{ object: "search", search_query: "q", data: [] }],
+        }}
+      />,
+    );
+
+    expect(screen.getByTestId("search-results-display")).toBeInTheDocument();
+  });
+
+  it("should show ResponseMetrics when usage data is present and no a2aMetadata", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{
+          role: "assistant",
+          content: "response",
+          timeToFirstToken: 150,
+          usage: { completionTokens: 10, promptTokens: 5, totalTokens: 15 },
+        }}
+      />,
+    );
+
+    expect(screen.getByTestId("response-metrics")).toBeInTheDocument();
+  });
+
+  it("should show A2AMetrics when a2aMetadata is present instead of ResponseMetrics", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{
+          role: "assistant",
+          content: "agent response",
+          timeToFirstToken: 100,
+          a2aMetadata: { taskId: "task-1", status: { state: "completed" } },
+        }}
+      />,
+    );
+
+    expect(screen.getByTestId("a2a-metrics")).toBeInTheDocument();
+    expect(screen.queryByTestId("response-metrics")).not.toBeInTheDocument();
+  });
+
+  it("should show CodeInterpreterOutput on the last assistant message for RESPONSES endpoint", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        isLastMessage={true}
+        endpointType={EndpointType.RESPONSES}
+        codeInterpreterResult={{
+          code: "print('hello')",
+          containerId: "container-1",
+          annotations: [],
+        }}
+        message={{ role: "assistant", content: "result" }}
+      />,
+    );
+
+    expect(screen.getByTestId("code-interpreter-output")).toHaveTextContent("print('hello')");
+  });
+
+  it("should render generated image from chat completions via message.image", () => {
+    render(
+      <ChatMessageBubble
+        {...defaultProps}
+        message={{
+          role: "assistant",
+          content: "Here is your image",
+          image: { url: "https://example.com/generated.png", detail: "auto" },
+        }}
+      />,
+    );
+
+    const images = screen.getAllByAltText("Generated image");
+    expect(images.some((img) => img.getAttribute("src") === "https://example.com/generated.png")).toBe(true);
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx
new file mode 100644
index 00000000000..15978c17f7e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx
@@ -0,0 +1,214 @@
+import { RobotOutlined, UserOutlined } from "@ant-design/icons";
+import React from "react";
+import ReactMarkdown from "react-markdown";
+import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
+import { coy } from "react-syntax-highlighter/dist/esm/styles/prism";
+import { CodeInterpreterResult } from "../llm_calls/code_interpreter_handler";
+import A2AMetrics from "./A2AMetrics";
+import AudioRenderer from "./AudioRenderer";
+import ChatImageRenderer from "./ChatImageRenderer";
+import CodeInterpreterOutput from "./CodeInterpreterOutput";
+import { EndpointType } from "./mode_endpoint_mapping";
+import MCPEventsDisplay from "./MCPEventsDisplay";
+import type { MCPEvent } from "../../mcp_tools/types";
+import ReasoningContent from "./ReasoningContent";
+import ResponseMetrics from "./ResponseMetrics";
+import ResponsesImageRenderer from "./ResponsesImageRenderer";
+import { SearchResultsDisplay } from "./SearchResultsDisplay";
+import { MessageType } from "./types";
+
+interface ChatMessageBubbleProps {
+  message: MessageType;
+  /** Whether this is the last message in the chat history. */
+  isLastMessage: boolean;
+  endpointType: EndpointType;
+  /** MCP events to display on the last assistant message. */
+  mcpEvents: MCPEvent[];
+  /** Code interpreter result to display on the last assistant message. */
+  codeInterpreterResult: CodeInterpreterResult | null;
+  /** API key used to fetch code interpreter file downloads. */
+  accessToken: string;
+}
+
+function ChatMessageBubble({
+  message,
+  isLastMessage,
+  endpointType,
+  mcpEvents,
+  codeInterpreterResult,
+  accessToken,
+}: ChatMessageBubbleProps) {
+  const isUser = message.role === "user";
+
+  return (
+    <div className={`mb-4 ${isUser ? "text-right" : "text-left"}`}>
+      <div
+        className="inline-block max-w-[80%] rounded-lg shadow-sm p-3.5 px-4"
+        style={{
+          backgroundColor: isUser ? "#f0f8ff" : "#ffffff",
+          border: isUser ? "1px solid #e6f0fa" : "1px solid #f0f0f0",
+          textAlign: "left",
+        }}
+      >
+        {/* Header: role icon + name + model badge */}
+        <div className="flex items-center gap-2 mb-1.5">
+          <div
+            className="flex items-center justify-center w-6 h-6 rounded-full mr-1"
+            style={{
+              backgroundColor: isUser ? "#e6f0fa" : "#f5f5f5",
+            }}
+          >
+            {isUser ? (
+              <UserOutlined style={{ fontSize: "12px", color: "#2563eb" }} />
+            ) : (
+              <RobotOutlined style={{ fontSize: "12px", color: "#4b5563" }} />
+            )}
+          </div>
+          <strong className="text-sm capitalize">{message.role}</strong>
+          {message.role === "assistant" && message.model && (
+            <span className="text-xs px-2 py-0.5 rounded bg-gray-100 text-gray-600 font-normal">
+              {message.model}
+            </span>
+          )}
+        </div>
+
+        {/* Reasoning content (chain-of-thought) */}
+        {message.reasoningContent && <ReasoningContent reasoningContent={message.reasoningContent} />}
+
+        {/* MCP events at the start of the last assistant message */}
+        {message.role === "assistant" &&
+          isLastMessage &&
+          mcpEvents.length > 0 &&
+          (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && (
+            <div className="mb-3">
+              <MCPEventsDisplay events={mcpEvents} />
+            </div>
+          )}
+
+        {/* Search results */}
+        {message.role === "assistant" && message.searchResults && (
+          <SearchResultsDisplay searchResults={message.searchResults} />
+        )}
+
+        {/* Code Interpreter output for the last assistant message */}
+        {message.role === "assistant" &&
+          isLastMessage &&
+          codeInterpreterResult &&
+          endpointType === EndpointType.RESPONSES && (
+            <CodeInterpreterOutput
+              code={codeInterpreterResult.code}
+              containerId={codeInterpreterResult.containerId}
+              annotations={codeInterpreterResult.annotations}
+              accessToken={accessToken}
+            />
+          )}
+
+        {/* Message body */}
+        <div
+          className="whitespace-pre-wrap break-words max-w-full message-content"
+          style={{
+            wordWrap: "break-word",
+            overflowWrap: "break-word",
+            wordBreak: "break-word",
+            hyphens: "auto",
+          }}
+        >
+          {message.isImage ? (
+            <img
+              src={typeof message.content === "string" ? message.content : ""}
+              alt="Generated image"
+              className="max-w-full rounded-md border border-gray-200 shadow-sm"
+              style={{ maxHeight: "500px" }}
+            />
+          ) : message.isAudio ? (
+            <AudioRenderer message={message} />
+          ) : (
+            <>
+              {/* Attached image for user messages based on endpoint */}
+              {endpointType === EndpointType.RESPONSES && <ResponsesImageRenderer message={message} />}
+              {endpointType === EndpointType.CHAT && <ChatImageRenderer message={message} />}
+
+              <ReactMarkdown
+                components={{
+                  code({
+                    node,
+                    inline,
+                    className,
+                    children,
+                    ...props
+                  }: React.ComponentPropsWithoutRef<"code"> & {
+                    inline?: boolean;
+                    node?: unknown;
+                  }) {
+                    const match = /language-(\w+)/.exec(className || "");
+                    return !inline && match ? (
+                      <SyntaxHighlighter
+                        style={coy as any}
+                        language={match[1]}
+                        PreTag="div"
+                        className="rounded-md my-2"
+                        wrapLines={true}
+                        wrapLongLines={true}
+                        {...props}
+                      >
+                        {String(children).replace(/\n$/, "")}
+                      </SyntaxHighlighter>
+                    ) : (
+                      <code
+                        className={`${className} px-1.5 py-0.5 rounded bg-gray-100 text-sm font-mono`}
+                        style={{ wordBreak: "break-word" }}
+                        {...props}
+                      >
+                        {children}
+                      </code>
+                    );
+                  },
+                  pre: ({ node, ...props }) => (
+                    <pre style={{ overflowX: "auto", maxWidth: "100%" }} {...props} />
+                  ),
+                }}
+              >
+                {typeof message.content === "string" ? message.content : ""}
+              </ReactMarkdown>
+
+              {/* Generated image from chat completions */}
+              {message.image && (
+                <div className="mt-3">
+                  <img
+                    src={message.image.url}
+                    alt="Generated image"
+                    className="max-w-full rounded-md border border-gray-200 shadow-sm"
+                    style={{ maxHeight: "500px" }}
+                  />
+                </div>
+              )}
+            </>
+          )}
+
+          {/* Response metrics */}
+          {message.role === "assistant" &&
+            (message.timeToFirstToken || message.totalLatency || message.usage) &&
+            !message.a2aMetadata && (
+              <ResponseMetrics
+                timeToFirstToken={message.timeToFirstToken}
+                totalLatency={message.totalLatency}
+                usage={message.usage}
+                toolName={message.toolName}
+              />
+            )}
+
+          {/* A2A Metrics */}
+          {message.role === "assistant" && message.a2aMetadata && (
+            <A2AMetrics
+              a2aMetadata={message.a2aMetadata}
+              timeToFirstToken={message.timeToFirstToken}
+              totalLatency={message.totalLatency}
+            />
+          )}
+        </div>
+      </div>
+    </div>
+  );
+}
+
+export default ChatMessageBubble;
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
index bc0d52cf581..ef57a75062c 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx
@@ -63,6 +63,7 @@ import EndpointSelector from "./EndpointSelector";
 import FilePreviewCard from "./FilePreviewCard";
 import MCPEventsDisplay from "./MCPEventsDisplay";
 import type { MCPEvent } from "../../mcp_tools/types";
+import ChatMessageBubble from "./ChatMessageBubble";
 import { EndpointType, getEndpointType } from "./mode_endpoint_mapping";
 import ReasoningContent from "./ReasoningContent";
 import ResponseMetrics, { TokenUsage } from "./ResponseMetrics";
@@ -1932,168 +1933,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
 
               {chatHistory.map((message, index) => (
                 <div key={index}>
-                  <div className={`mb-4 ${message.role === "user" ? "text-right" : "text-left"}`}>
-                    <div
-                      className="inline-block max-w-[80%] rounded-lg shadow-sm p-3.5 px-4"
-                      style={{
-                        backgroundColor: message.role === "user" ? "#f0f8ff" : "#ffffff",
-                        border: message.role === "user" ? "1px solid #e6f0fa" : "1px solid #f0f0f0",
-                        textAlign: "left",
-                      }}
-                    >
-                      <div className="flex items-center gap-2 mb-1.5">
-                        <div
-                          className="flex items-center justify-center w-6 h-6 rounded-full mr-1"
-                          style={{
-                            backgroundColor: message.role === "user" ? "#e6f0fa" : "#f5f5f5",
-                          }}
-                        >
-                          {message.role === "user" ? (
-                            <UserOutlined style={{ fontSize: "12px", color: "#2563eb" }} />
-                          ) : (
-                            <RobotOutlined style={{ fontSize: "12px", color: "#4b5563" }} />
-                          )}
-                        </div>
-                        <strong className="text-sm capitalize">{message.role}</strong>
-                        {message.role === "assistant" && message.model && (
-                          <span className="text-xs px-2 py-0.5 rounded bg-gray-100 text-gray-600 font-normal">
-                            {message.model}
-                          </span>
-                        )}
-                      </div>
-                      {message.reasoningContent && <ReasoningContent reasoningContent={message.reasoningContent} />}
-
-                      {/* Show MCP events at the start of assistant messages */}
-                      {message.role === "assistant" &&
-                        index === chatHistory.length - 1 &&
-                        mcpEvents.length > 0 &&
-                        (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && (
-                          <div className="mb-3">
-                            <MCPEventsDisplay events={mcpEvents} />
-                          </div>
-                        )}
-
-                      {/* Show search results at the start of assistant messages */}
-                      {message.role === "assistant" && message.searchResults && (
-                        <SearchResultsDisplay searchResults={message.searchResults} />
-                      )}
-
-                      {/* Show Code Interpreter output for the last assistant message */}
-                      {message.role === "assistant" &&
-                        index === chatHistory.length - 1 &&
-                        codeInterpreter.result &&
-                        endpointType === EndpointType.RESPONSES && (
-                          <CodeInterpreterOutput
-                            code={codeInterpreter.result.code}
-                            containerId={codeInterpreter.result.containerId}
-                            annotations={codeInterpreter.result.annotations}
-                            accessToken={apiKeySource === "session" ? accessToken || "" : apiKey}
-                          />
-                        )}
-
-                      <div
-                        className="whitespace-pre-wrap break-words max-w-full message-content"
-                        style={{
-                          wordWrap: "break-word",
-                          overflowWrap: "break-word",
-                          wordBreak: "break-word",
-                          hyphens: "auto",
-                        }}
-                      >
-                        {message.isImage ? (
-                          <img
-                            src={typeof message.content === "string" ? message.content : ""}
-                            alt="Generated image"
-                            className="max-w-full rounded-md border border-gray-200 shadow-sm"
-                            style={{ maxHeight: "500px" }}
-                          />
-                        ) : message.isAudio ? (
-                          <AudioRenderer message={message} />
-                        ) : (
-                          <>
-                            {/* Show attached image for user messages based on current endpoint */}
-                            {endpointType === EndpointType.RESPONSES && <ResponsesImageRenderer message={message} />}
-                            {endpointType === EndpointType.CHAT && <ChatImageRenderer message={message} />}
-
-                            <ReactMarkdown
-                              components={{
-                                code({
-                                  node,
-                                  inline,
-                                  className,
-                                  children,
-                                  ...props
-                                }: React.ComponentPropsWithoutRef<"code"> & {
-                                  inline?: boolean;
-                                  node?: any;
-                                }) {
-                                  const match = /language-(\w+)/.exec(className || "");
-                                  return !inline && match ? (
-                                    <SyntaxHighlighter
-                                      style={coy as any}
-                                      language={match[1]}
-                                      PreTag="div"
-                                      className="rounded-md my-2"
-                                      wrapLines={true}
-                                      wrapLongLines={true}
-                                      {...props}
-                                    >
-                                      {String(children).replace(/\n$/, "")}
-                                    </SyntaxHighlighter>
-                                  ) : (
-                                    <code
-                                      className={`${className} px-1.5 py-0.5 rounded bg-gray-100 text-sm font-mono`}
-                                      style={{ wordBreak: "break-word" }}
-                                      {...props}
-                                    >
-                                      {children}
-                                    </code>
-                                  );
-                                },
-                                pre: ({ node, ...props }) => (
-                                  <pre style={{ overflowX: "auto", maxWidth: "100%" }} {...props} />
-                                ),
-                              }}
-                            >
-                              {typeof message.content === "string" ? message.content : ""}
-                            </ReactMarkdown>
-
-                            {/* Show generated image from chat completions */}
-                            {message.image && (
-                              <div className="mt-3">
-                                <img
-                                  src={message.image.url}
-                                  alt="Generated image"
-                                  className="max-w-full rounded-md border border-gray-200 shadow-sm"
-                                  style={{ maxHeight: "500px" }}
-                                />
-                              </div>
-                            )}
-                          </>
-                        )}
-
-                        {message.role === "assistant" &&
-                          (message.timeToFirstToken || message.totalLatency || message.usage) &&
-                          !message.a2aMetadata && (
-                            <ResponseMetrics
-                              timeToFirstToken={message.timeToFirstToken}
-                              totalLatency={message.totalLatency}
-                              usage={message.usage}
-                              toolName={message.toolName}
-                            />
-                          )}
-
-                        {/* A2A Metrics - show for A2A agent responses */}
-                        {message.role === "assistant" && message.a2aMetadata && (
-                          <A2AMetrics
-                            a2aMetadata={message.a2aMetadata}
-                            timeToFirstToken={message.timeToFirstToken}
-                            totalLatency={message.totalLatency}
-                          />
-                        )}
-                      </div>
-                    </div>
-                  </div>
+                  <ChatMessageBubble
+                    message={message}
+                    isLastMessage={index === chatHistory.length - 1}
+                    endpointType={endpointType as EndpointType}
+                    mcpEvents={mcpEvents}
+                    codeInterpreterResult={codeInterpreter.result}
+                    accessToken={apiKeySource === "session" ? accessToken || "" : apiKey}
+                  />
                 </div>
               ))}