diff --git a/packages/ui/src/features/sessions/components/session-update/AgentMessage.test.tsx b/packages/ui/src/features/sessions/components/session-update/AgentMessage.test.tsx
new file mode 100644
index 0000000000..e348c697ad
--- /dev/null
+++ b/packages/ui/src/features/sessions/components/session-update/AgentMessage.test.tsx
@@ -0,0 +1,25 @@
+import { Theme } from "@radix-ui/themes";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { describe, expect, it, vi } from "vitest";
+import { AgentMessage } from "./AgentMessage";
+
+describe("AgentMessage", () => {
+ it("keeps the copy action inside the message rail", async () => {
+ const writeText = vi.fn().mockResolvedValue(undefined);
+ Object.assign(navigator, { clipboard: { writeText } });
+
+ render(
+
+
+ ,
+ );
+
+ const copyButton = screen.getByLabelText("Copy message");
+ expect(copyButton.parentElement).toHaveClass("right-1");
+ expect(copyButton.parentElement).not.toHaveClass("left-full");
+
+ await userEvent.click(copyButton);
+ expect(writeText).toHaveBeenCalledWith("Assistant response");
+ });
+});
diff --git a/packages/ui/src/features/sessions/components/session-update/AgentMessage.tsx b/packages/ui/src/features/sessions/components/session-update/AgentMessage.tsx
index 56f13bb11a..d6fc8cdf4b 100644
--- a/packages/ui/src/features/sessions/components/session-update/AgentMessage.tsx
+++ b/packages/ui/src/features/sessions/components/session-update/AgentMessage.tsx
@@ -70,7 +70,7 @@ export const AgentMessage = memo(function AgentMessage({
}, [content]);
return (
-
+
{isStreaming ? (
)}
-
+