feat(ui): add retry and edit message actions to playground chat

Add ChatGPT-like message interaction features to the playground:

- Edit: hover over user messages to reveal a pencil icon; click to
  enter inline edit mode with textarea and confirm/cancel buttons.
  Submitting an edit truncates conversation history and re-sends.

- Retry: hover over the last assistant message to reveal a reload
  icon; click to regenerate the response by re-sending the preceding
  user message.

New files:
- MessageActions.tsx: reusable component for hover action buttons
- MessageActions.test.tsx: 14 tests covering edit/retry/cancel/keyboard

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>

ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx

@@ -69,6 +69,7 @@ import { createDisplayMessage, createMultimodalMessage } from "./ResponsesImageU
 import { SearchResultsDisplay } from "./SearchResultsDisplay";
 import SessionManagement from "./SessionManagement";
 import RealtimePlayground from "./RealtimePlayground";
+import MessageActions from "./MessageActions";
 import { A2ATaskMetadata, MessageType } from "./types";
 import { useCodeInterpreter } from "./useCodeInterpreter";
 
@@ -242,6 +243,8 @@ const ChatUI: React.FC<ChatUIProps> = ({
   // Code Interpreter state (using custom hook)
   const codeInterpreter = useCodeInterpreter();
 
+  const [pendingResend, setPendingResend] = useState(false);
+
   const chatEndRef = useRef<HTMLDivElement>(null);
 
   // Fetch MCP servers
@@ -1217,6 +1220,54 @@ const ChatUI: React.FC<ChatUIProps> = ({
     setInputMessage("");
   };
 
+  useEffect(() => {
+    if (pendingResend && inputMessage.trim() !== "" && !isLoading) {
+      setPendingResend(false);
+      handleSendMessage();
+    }
+  }, [pendingResend, inputMessage]);
+
+  const handleRetry = () => {
+    if (isLoading || chatHistory.length === 0) return;
+
+    // Find the last assistant message index
+    let lastAssistantIdx = -1;
+    for (let i = chatHistory.length - 1; i >= 0; i--) {
+      if (chatHistory[i].role === "assistant") {
+        lastAssistantIdx = i;
+        break;
+      }
+    }
+    if (lastAssistantIdx === -1) return;
+
+    // Find the user message right before the assistant message
+    let userMsgIdx = -1;
+    for (let i = lastAssistantIdx - 1; i >= 0; i--) {
+      if (chatHistory[i].role === "user") {
+        userMsgIdx = i;
+        break;
+      }
+    }
+    if (userMsgIdx === -1) return;
+
+    const userMessage = chatHistory[userMsgIdx];
+    const userContent = typeof userMessage.content === "string" ? userMessage.content : "";
+
+    // Truncate history: keep everything before the user message that triggered the response
+    setChatHistory(chatHistory.slice(0, userMsgIdx));
+    setInputMessage(userContent);
+    setPendingResend(true);
+  };
+
+  const handleEditSubmit = (messageIndex: number, newContent: string) => {
+    if (isLoading) return;
+
+    // Truncate history at the edited message (remove it and everything after)
+    setChatHistory(chatHistory.slice(0, messageIndex));
+    setInputMessage(newContent);
+    setPendingResend(true);
+  };
+
   const clearChatHistory = () => {
     // Clean up audio object URLs before clearing history
     chatHistory.forEach((message) => {
@@ -1867,7 +1918,7 @@ 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={`group 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={{
@@ -2028,6 +2079,21 @@ const ChatUI: React.FC<ChatUIProps> = ({
                         )}
                       </div>
                     </div>
+                    <MessageActions
+                      role={message.role}
+                      content={typeof message.content === "string" ? message.content : ""}
+                      messageIndex={index}
+                      isLastAssistantMessage={
+                        message.role === "assistant" &&
+                        index === chatHistory.map((m) => m.role).lastIndexOf("assistant")
+                      }
+                      isLoading={isLoading}
+                      isImage={message.isImage}
+                      isAudio={message.isAudio}
+                      isEmbeddings={message.isEmbeddings}
+                      onRetry={handleRetry}
+                      onEditSubmit={handleEditSubmit}
+                    />
                   </div>
                 </div>
               ))}

ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.test.tsx

@@ -0,0 +1,218 @@
+import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+import MessageActions from "./MessageActions";
+
+describe("MessageActions", () => {
+  const defaultProps = {
+    role: "user",
+    content: "Hello world",
+    messageIndex: 0,
+    isLastAssistantMessage: false,
+    isLoading: false,
+    onRetry: vi.fn(),
+    onEditSubmit: vi.fn(),
+  };
+
+  it("should render edit button for user messages on hover", () => {
+    render(<MessageActions {...defaultProps} />);
+    const editButton = screen.getByTestId("edit-message-button");
+    expect(editButton).toBeInTheDocument();
+  });
+
+  it("should not render any buttons for user messages when loading", () => {
+    render(<MessageActions {...defaultProps} isLoading={true} />);
+    expect(screen.queryByTestId("edit-message-button")).toBeNull();
+    expect(screen.queryByTestId("retry-message-button")).toBeNull();
+  });
+
+  it("should render retry button for the last assistant message", () => {
+    render(
+      <MessageActions
+        {...defaultProps}
+        role="assistant"
+        isLastAssistantMessage={true}
+      />,
+    );
+    const retryButton = screen.getByTestId("retry-message-button");
+    expect(retryButton).toBeInTheDocument();
+  });
+
+  it("should not render retry button for non-last assistant messages", () => {
+    render(
+      <MessageActions
+        {...defaultProps}
+        role="assistant"
+        isLastAssistantMessage={false}
+      />,
+    );
+    expect(screen.queryByTestId("retry-message-button")).toBeNull();
+  });
+
+  it("should not render edit button for assistant messages", () => {
+    render(
+      <MessageActions
+        {...defaultProps}
+        role="assistant"
+        isLastAssistantMessage={true}
+      />,
+    );
+    expect(screen.queryByTestId("edit-message-button")).toBeNull();
+  });
+
+  it("should not render any buttons for image messages", () => {
+    render(<MessageActions {...defaultProps} isImage={true} />);
+    expect(screen.queryByTestId("edit-message-button")).toBeNull();
+    expect(screen.queryByTestId("message-actions")).toBeNull();
+  });
+
+  it("should not render any buttons for audio messages", () => {
+    render(<MessageActions {...defaultProps} isAudio={true} />);
+    expect(screen.queryByTestId("message-actions")).toBeNull();
+  });
+
+  it("should call onRetry when retry button is clicked", () => {
+    const onRetry = vi.fn();
+    render(
+      <MessageActions
+        {...defaultProps}
+        role="assistant"
+        isLastAssistantMessage={true}
+        onRetry={onRetry}
+      />,
+    );
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("retry-message-button"));
+    });
+
+    expect(onRetry).toHaveBeenCalledTimes(1);
+  });
+
+  it("should show edit textarea when edit button is clicked", async () => {
+    render(<MessageActions {...defaultProps} />);
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("edit-message-button"));
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("edit-message-container")).toBeInTheDocument();
+    });
+
+    const textarea = screen.getByRole("textbox");
+    expect(textarea).toBeInTheDocument();
+    expect(textarea).toHaveValue("Hello world");
+  });
+
+  it("should call onEditSubmit when confirming edit", async () => {
+    const onEditSubmit = vi.fn();
+    render(<MessageActions {...defaultProps} onEditSubmit={onEditSubmit} />);
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("edit-message-button"));
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("edit-message-container")).toBeInTheDocument();
+    });
+
+    const textarea = screen.getByRole("textbox");
+    act(() => {
+      fireEvent.change(textarea, { target: { value: "Updated message" } });
+    });
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("confirm-edit-button"));
+    });
+
+    expect(onEditSubmit).toHaveBeenCalledWith(0, "Updated message");
+  });
+
+  it("should cancel edit when cancel button is clicked", async () => {
+    render(<MessageActions {...defaultProps} />);
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("edit-message-button"));
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("edit-message-container")).toBeInTheDocument();
+    });
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("cancel-edit-button"));
+    });
+
+    await waitFor(() => {
+      expect(screen.queryByTestId("edit-message-container")).toBeNull();
+    });
+
+    expect(screen.getByTestId("edit-message-button")).toBeInTheDocument();
+  });
+
+  it("should cancel edit when Escape key is pressed", async () => {
+    render(<MessageActions {...defaultProps} />);
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("edit-message-button"));
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("edit-message-container")).toBeInTheDocument();
+    });
+
+    const textarea = screen.getByRole("textbox");
+    act(() => {
+      fireEvent.keyDown(textarea, { key: "Escape" });
+    });
+
+    await waitFor(() => {
+      expect(screen.queryByTestId("edit-message-container")).toBeNull();
+    });
+  });
+
+  it("should submit edit when Enter key is pressed without Shift", async () => {
+    const onEditSubmit = vi.fn();
+    render(<MessageActions {...defaultProps} onEditSubmit={onEditSubmit} />);
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("edit-message-button"));
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("edit-message-container")).toBeInTheDocument();
+    });
+
+    const textarea = screen.getByRole("textbox");
+    act(() => {
+      fireEvent.change(textarea, { target: { value: "New content" } });
+    });
+
+    act(() => {
+      fireEvent.keyDown(textarea, { key: "Enter", shiftKey: false });
+    });
+
+    expect(onEditSubmit).toHaveBeenCalledWith(0, "New content");
+  });
+
+  it("should not submit edit with empty content", async () => {
+    const onEditSubmit = vi.fn();
+    render(<MessageActions {...defaultProps} onEditSubmit={onEditSubmit} />);
+
+    act(() => {
+      fireEvent.click(screen.getByTestId("edit-message-button"));
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId("edit-message-container")).toBeInTheDocument();
+    });
+
+    const textarea = screen.getByRole("textbox");
+    act(() => {
+      fireEvent.change(textarea, { target: { value: "   " } });
+    });
+
+    const confirmButton = screen.getByTestId("confirm-edit-button");
+    expect(confirmButton).toBeDisabled();
+  });
+});

ui/litellm-dashboard/src/components/playground/chat_ui/MessageActions.tsx

@@ -0,0 +1,150 @@
+import {
+  CheckOutlined,
+  CloseOutlined,
+  EditOutlined,
+  ReloadOutlined,
+} from "@ant-design/icons";
+import { Button, Input, Tooltip } from "antd";
+import React, { useEffect, useRef, useState } from "react";
+
+const { TextArea } = Input;
+
+interface MessageActionsProps {
+  role: string;
+  content: string;
+  messageIndex: number;
+  isLastAssistantMessage: boolean;
+  isLoading: boolean;
+  isImage?: boolean;
+  isAudio?: boolean;
+  isEmbeddings?: boolean;
+  onRetry: () => void;
+  onEditSubmit: (messageIndex: number, newContent: string) => void;
+}
+
+const MessageActions: React.FC<MessageActionsProps> = ({
+  role,
+  content,
+  messageIndex,
+  isLastAssistantMessage,
+  isLoading,
+  isImage,
+  isAudio,
+  isEmbeddings,
+  onRetry,
+  onEditSubmit,
+}) => {
+  const [isEditing, setIsEditing] = useState(false);
+  const [editValue, setEditValue] = useState(content);
+  const textAreaRef = useRef<any>(null);
+
+  useEffect(() => {
+    if (isEditing && textAreaRef.current) {
+      const ta = textAreaRef.current?.resizableTextArea?.textArea;
+      if (ta) {
+        ta.focus();
+        ta.setSelectionRange(ta.value.length, ta.value.length);
+      }
+    }
+  }, [isEditing]);
+
+  const handleStartEdit = () => {
+    setEditValue(content);
+    setIsEditing(true);
+  };
+
+  const handleCancelEdit = () => {
+    setIsEditing(false);
+    setEditValue(content);
+  };
+
+  const handleConfirmEdit = () => {
+    if (editValue.trim() === "") return;
+    setIsEditing(false);
+    onEditSubmit(messageIndex, editValue);
+  };
+
+  const handleEditKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
+    if (e.key === "Enter" && !e.shiftKey) {
+      e.preventDefault();
+      handleConfirmEdit();
+    }
+    if (e.key === "Escape") {
+      handleCancelEdit();
+    }
+  };
+
+  const isSpecialMessage = isImage || isAudio || isEmbeddings;
+  const canEdit = role === "user" && !isLoading && !isSpecialMessage;
+  const canRetry = role === "assistant" && isLastAssistantMessage && !isLoading && !isSpecialMessage;
+
+  if (isEditing) {
+    return (
+      <div className="mt-2" data-testid="edit-message-container">
+        <TextArea
+          ref={textAreaRef}
+          value={editValue}
+          onChange={(e) => setEditValue(e.target.value)}
+          onKeyDown={handleEditKeyDown}
+          autoSize={{ minRows: 1, maxRows: 8 }}
+          className="mb-2"
+          style={{ fontSize: "14px" }}
+        />
+        <div className="flex gap-1.5 justify-end">
+          <Tooltip title="Cancel (Esc)">
+            <Button
+              size="small"
+              icon={<CloseOutlined />}
+              onClick={handleCancelEdit}
+              data-testid="cancel-edit-button"
+            />
+          </Tooltip>
+          <Tooltip title="Save & Submit (Enter)">
+            <Button
+              size="small"
+              type="primary"
+              icon={<CheckOutlined />}
+              onClick={handleConfirmEdit}
+              disabled={editValue.trim() === ""}
+              data-testid="confirm-edit-button"
+            />
+          </Tooltip>
+        </div>
+      </div>
+    );
+  }
+
+  if (!canEdit && !canRetry) return null;
+
+  return (
+    <div
+      className="message-actions mt-1 flex gap-1 opacity-0 group-hover:opacity-100 transition-opacity"
+      data-testid="message-actions"
+    >
+      {canEdit && (
+        <Tooltip title="Edit message">
+          <button
+            className="p-1 rounded hover:bg-gray-100 text-gray-400 hover:text-gray-600 transition-colors"
+            onClick={handleStartEdit}
+            data-testid="edit-message-button"
+          >
+            <EditOutlined style={{ fontSize: "13px" }} />
+          </button>
+        </Tooltip>
+      )}
+      {canRetry && (
+        <Tooltip title="Retry">
+          <button
+            className="p-1 rounded hover:bg-gray-100 text-gray-400 hover:text-gray-600 transition-colors"
+            onClick={onRetry}
+            data-testid="retry-message-button"
+          >
+            <ReloadOutlined style={{ fontSize: "13px" }} />
+          </button>
+        </Tooltip>
+      )}
+    </div>
+  );
+};
+
+export default MessageActions;