From b55cb249fe276996e9ca351e31d08e7cff4a9542 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Wed, 18 Mar 2026 16:46:41 -0700
Subject: [PATCH] Address Greptile feedback: use EndpointType enum, add CHAT
 MCP test

- Narrow endpointType prop from string to EndpointType enum
- Add missing test for MCP events on CHAT endpoint

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../chat_ui/ChatMessageBubble.test.tsx           | 16 ++++++++++++++++
 .../playground/chat_ui/ChatMessageBubble.tsx     |  2 +-
 2 files changed, 17 insertions(+), 1 deletion(-)

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
index 368043abee6..70c3fdd4f29 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.test.tsx
@@ -180,6 +180,22 @@ describe("ChatMessageBubble", () => {
     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" }];
 
diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx
index 3e1913cb2b2..148cd082e22 100644
--- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx
+++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatMessageBubble.tsx
@@ -21,7 +21,7 @@ interface ChatMessageBubbleProps {
   message: MessageType;
   /** Whether this is the last message in the chat history. */
   isLastMessage: boolean;
-  endpointType: string;
+  endpointType: EndpointType;
   /** MCP events to display on the last assistant message. */
   mcpEvents: MCPEvent[];
   /** Code interpreter result to display on the last assistant message. */