Add copy button to copy the full assistant message in Studio Code

apps/studio/src/components/studio-code-session/conversation/index.test.tsx

@@ -1,11 +1,12 @@
-import { render, screen } from '@testing-library/react';
+import { render, screen, fireEvent, waitFor } from '@testing-library/react';
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { Conversation, entriesToRenderItems } from './index';
 import type { SessionEntry } from '@earendil-works/pi-coding-agent';
 import type { LoadedAiSession } from '@studio/common/ai/sessions/types';
 
 const ipcApiMocks = vi.hoisted( () => ( {
 	readLocalMediaFile: vi.fn(),
+	copyText: vi.fn(),
 } ) );
 
 vi.mock( 'src/lib/get-ipc-api', () => ( {
@@ -271,3 +272,96 @@ describe( 'Conversation – inline media artifacts', () => {
 		expect( await screen.findByRole( 'status' ) ).toHaveTextContent( 'Image unavailable' );
 	} );
 } );
+
+describe( 'Conversation – assistant message copy button', () => {
+	beforeEach( () => {
+		ipcApiMocks.copyText.mockClear();
+	} );
+
+	function assistantTextEntry( text: string ): SessionEntry {
+		return {
+			type: 'message',
+			id: `assistant-text-${ text }`,
+			parentId: null,
+			timestamp: '2026-06-19T00:00:00.000Z',
+			message: {
+				role: 'assistant',
+				content: [ { type: 'text', text } ],
+			},
+		} as unknown as SessionEntry;
+	}
+
+	function assistantMultiBlockEntry(): SessionEntry {
+		return {
+			type: 'message',
+			id: 'assistant-multi-block',
+			parentId: null,
+			timestamp: '2026-06-19T00:00:00.000Z',
+			message: {
+				role: 'assistant',
+				content: [
+					{ type: 'text', text: 'First part.' },
+					{ type: 'toolCall', id: 'tool-call-1', name: 'read_file', arguments: {} },
+					{ type: 'text', text: 'Second part.' },
+				],
+			},
+		} as unknown as SessionEntry;
+	}
+
+	function userPromptEntry( text: string ): SessionEntry {
+		return customEntry( 'studio.user_prompt', { text, source: 'prompt' } );
+	}
+
+	function renderConversation( entries: SessionEntry[] ) {
+		render(
+			<Conversation
+				data={ { entries } as unknown as LoadedAiSession }
+				isRunning={ false }
+				startedAt={ null }
+				pendingQuestions={ new Set() }
+				pendingAnswers={ {} }
+				answeredQuestions={ {} }
+				onAnswerQuestion={ () => {} }
+			/>
+		);
+	}
+
+	it( 'copies the raw markdown of an assistant message', async () => {
+		renderConversation( [ assistantTextEntry( '# Hello\n\nSome **bold** text.' ) ] );
+
+		const button = screen.getByRole( 'button', { name: 'Copy message' } );
+		expect( button ).toBeInTheDocument();
+
+		fireEvent.click( button );
+
+		expect( ipcApiMocks.copyText ).toHaveBeenCalledWith( '# Hello\n\nSome **bold** text.' );
+		await waitFor( () => expect( screen.getByRole( 'status' ) ).toHaveTextContent( 'Copied' ) );
+	} );
+
+	it( 'copies the full joined message for a message split across text blocks', () => {
+		renderConversation( [ assistantMultiBlockEntry() ] );
+
+		const buttons = screen.getAllByRole( 'button', { name: 'Copy message' } );
+		expect( buttons ).toHaveLength( 1 );
+
+		fireEvent.click( buttons[ 0 ] );
+
+		expect( ipcApiMocks.copyText ).toHaveBeenCalledWith( 'First part.\n\nSecond part.' );
+	} );
+
+	it( 'does not render a copy button for user messages', () => {
+		renderConversation( [ userPromptEntry( 'Hello there' ) ] );
+
+		expect( screen.queryByRole( 'button', { name: 'Copy message' } ) ).not.toBeInTheDocument();
+	} );
+
+	it( 'shows a tooltip for the copy button', async () => {
+		renderConversation( [ assistantTextEntry( 'Plain reply.' ) ] );
+
+		fireEvent.mouseOver( screen.getByRole( 'button', { name: 'Copy message' } ) );
+
+		await waitFor( () =>
+			expect( screen.getByRole( 'tooltip' ) ).toHaveTextContent( 'Copy message' )
+		);
+	} );
+} );

apps/studio/src/components/studio-code-session/conversation/index.tsx

@@ -25,6 +25,7 @@ import { Icon } from '@wordpress/ui';
 import { useEffect, useMemo, useState } from 'react';
 import { cx } from 'src/lib/cx';
 import { getIpcApi } from 'src/lib/get-ipc-api';
+import { CopyButton } from '../copy-button';
 import { Markdown } from '../markdown';
 import { ThinkingIndicator } from '../thinking-indicator';
 import styles from './style.module.css';
@@ -38,7 +39,7 @@ type RenderItem =
 			text: string;
 			attachments?: StudioChatAttachmentSummary[];
 	  }
-	| { kind: 'assistant-text'; key: string; text: string }
+	| { kind: 'assistant-text'; key: string; text: string; copyText?: string }
 	| {
 			kind: 'tool-use';
 			key: string;
@@ -146,6 +147,15 @@ export function entriesToRenderItems( entries: SessionEntry[] ): RenderItem[] {
 			if ( ! message || message.role !== 'assistant' || ! Array.isArray( message.content ) ) {
 				return;
 			}
+			// A single assistant message can hold several text blocks split by tool
+			// calls. Copy must yield the whole message, so join every text block and
+			// hang one copy button off the last one rather than one per fragment.
+			const textBlocks = message.content.filter(
+				( block ) => block.type === 'text' && typeof block.text === 'string' && block.text.trim()
+			);
+			const fullMessageText = textBlocks.map( ( block ) => block.text!.trim() ).join( '\n\n' );
+			const lastTextBlock = textBlocks[ textBlocks.length - 1 ];
+
 			message.content.forEach( ( block, blockIndex ) => {
 				if ( block.type === 'text' && typeof block.text === 'string' ) {
 					const text = block.text.trim();
@@ -154,6 +164,7 @@ export function entriesToRenderItems( entries: SessionEntry[] ): RenderItem[] {
 							kind: 'assistant-text',
 							key: `${ entryIndex }:${ blockIndex }:text`,
 							text,
+							copyText: block === lastTextBlock ? fullMessageText : undefined,
 						} );
 					}
 				} else if (
@@ -280,8 +291,19 @@ function UserTurn( {
 	);
 }
 
-function AssistantText( { text }: { text: string } ) {
-	return <Markdown>{ text }</Markdown>;
+function AssistantText( { text, copyText }: { text: string; copyText?: string } ) {
+	return (
+		<div className={ styles.assistantTurn }>
+			<Markdown>{ text }</Markdown>
+			{ copyText ? (
+				<CopyButton
+					text={ copyText }
+					label={ __( 'Copy message' ) }
+					className={ styles.messageActions }
+				/>
+			) : null }
+		</div>
+	);
 }
 
 const TOOL_RESULT_PREVIEW_MAX_LINES = 12;
@@ -517,7 +539,7 @@ export function Conversation( {
 							<UserTurn key={ item.key } text={ item.text } attachments={ item.attachments } />
 						);
 					case 'assistant-text':
-						return <AssistantText key={ item.key } text={ item.text } />;
+						return <AssistantText key={ item.key } text={ item.text } copyText={ item.copyText } />;
 					case 'tool-use':
 						return (
 							<ToolUseRow

apps/studio/src/components/studio-code-session/conversation/style.module.css

@@ -79,6 +79,29 @@
 	object-fit: cover;
 }
 
+/* Assistant response wrapper: the copy button below stays hidden until the
+   turn is hovered (or a control inside it is focused), mirroring Claude/ChatGPT. */
+.assistantTurn {
+	display: flex;
+	flex-direction: column;
+	gap: var(--wpds-dimension-padding-xs);
+}
+
+.messageActions {
+	display: flex;
+	align-items: center;
+	/* Cancel the icon-centering inset of the 1.75rem button so the 16px icon
+	   lines up with the message text edge. */
+	margin-inline-start: calc((16px - 1.75rem) / 2);
+	opacity: 0;
+	transition: opacity 0.12s ease;
+}
+
+.assistantTurn:hover .messageActions,
+.assistantTurn:focus-within .messageActions {
+	opacity: 1;
+}
+
 .toolBlock {
 	display: flex;
 	flex-direction: column;

apps/studio/src/components/studio-code-session/copy-button.module.css

@@ -0,0 +1,37 @@
+.copyButton {
+	display: inline-flex;
+	align-items: center;
+	justify-content: center;
+	width: 1.75rem;
+	height: 1.75rem;
+	margin: 0;
+	padding: 0;
+	border: 0;
+	border-radius: var(--wpds-border-radius-sm, 4px);
+	background-color: var(--color-frame-bg);
+	background-color: color-mix(in srgb, var(--color-frame-bg) 82%, transparent);
+	-webkit-backdrop-filter: blur(4px);
+	backdrop-filter: blur(4px);
+	color: var(--color-frame-text-secondary);
+	line-height: 1;
+	cursor: pointer;
+	transition: color 0.12s ease, background-color 0.12s ease;
+}
+
+.copyButton:hover,
+.copyButton:focus-visible {
+	color: var(--color-frame-text);
+	background-color: var(--color-frame-surface-alt);
+}
+
+.visuallyHidden {
+	position: absolute;
+	width: 1px;
+	height: 1px;
+	padding: 0;
+	margin: -1px;
+	overflow: hidden;
+	clip: rect(0, 0, 0, 0);
+	white-space: nowrap;
+	border: 0;
+}

apps/studio/src/components/studio-code-session/copy-button.tsx

@@ -0,0 +1,60 @@
+import { __ } from '@wordpress/i18n';
+import { check, copy, Icon } from '@wordpress/icons';
+import { useCallback, useEffect, useRef, useState } from 'react';
+import { Tooltip } from 'src/components/tooltip';
+import { getIpcApi } from 'src/lib/get-ipc-api';
+import styles from './copy-button.module.css';
+
+export function CopyButton( {
+	text,
+	label,
+	className,
+}: {
+	text: string;
+	label: string;
+	className?: string;
+} ) {
+	const [ copied, setCopied ] = useState( false );
+	const resetTimer = useRef< ReturnType< typeof setTimeout > | null >( null );
+
+	// Clear any pending reset when the button unmounts.
+	useEffect( () => {
+		return () => {
+			if ( resetTimer.current ) {
+				clearTimeout( resetTimer.current );
+			}
+		};
+	}, [] );
+
+	const handleCopy = useCallback( () => {
+		void getIpcApi().copyText( text );
+		setCopied( true );
+		// Re-arm the reset on every click so copying again mid-"Copied" doesn't
+		// let the earlier timer flip the state back too soon.
+		if ( resetTimer.current ) {
+			clearTimeout( resetTimer.current );
+		}
+		resetTimer.current = setTimeout( () => setCopied( false ), 2000 );
+	}, [ text ] );
+
+	const copiedLabel = __( 'Copied' );
+	const tooltipLabel = copied ? copiedLabel : label;
+
+	return (
+		<div className={ className }>
+			<Tooltip text={ tooltipLabel }>
+				<button
+					type="button"
+					className={ styles.copyButton }
+					onClick={ handleCopy }
+					aria-label={ label }
+				>
+					<Icon icon={ copied ? check : copy } size={ 16 } fill="currentColor" aria-hidden="true" />
+				</button>
+			</Tooltip>
+			<span className={ styles.visuallyHidden } role="status" aria-live="polite" aria-atomic="true">
+				{ copied ? copiedLabel : '' }
+			</span>
+		</div>
+	);
+}

apps/studio/src/components/studio-code-session/markdown/index.tsx

@@ -1,11 +1,10 @@
 import { __ } from '@wordpress/i18n';
-import { check, copy, Icon } from '@wordpress/icons';
-import { isValidElement, useCallback, useEffect, useMemo, useState } from 'react';
+import { isValidElement, useMemo } from 'react';
 import ReactMarkdown from 'react-markdown';
 import remarkGfm from 'remark-gfm';
-import { Tooltip } from 'src/components/tooltip';
 import { cx } from 'src/lib/cx';
 import { getIpcApi } from 'src/lib/get-ipc-api';
+import { CopyButton } from '../copy-button';
 import styles from './style.module.css';
 import type { MouseEvent, ReactNode } from 'react';
 import type { Components } from 'react-markdown';
@@ -41,54 +40,20 @@ function extractText( node: ReactNode ): string {
 	return '';
 }
 
-function CopyButton( { text }: { text: string } ) {
-	const [ copied, setCopied ] = useState( false );
-
-	// Reset back to the idle state a short while after a successful copy.
-	useEffect( () => {
-		if ( ! copied ) {
-			return;
-		}
-		const timer = setTimeout( () => setCopied( false ), 2000 );
-		return () => clearTimeout( timer );
-	}, [ copied ] );
-
-	const handleCopy = useCallback( () => {
-		void getIpcApi().copyText( text );
-		setCopied( true );
-	}, [ text ] );
-
-	const copyLabel = __( 'Copy code' );
-	const copiedLabel = __( 'Copied' );
-	const tooltipLabel = copied ? copiedLabel : copyLabel;
-
-	return (
-		<div className={ styles.copyButtonContainer }>
-			<Tooltip text={ tooltipLabel }>
-				<button
-					type="button"
-					className={ styles.copyButton }
-					onClick={ handleCopy }
-					aria-label={ copyLabel }
-				>
-					<Icon icon={ copied ? check : copy } size={ 16 } fill="currentColor" aria-hidden="true" />
-				</button>
-			</Tooltip>
-			<span className={ styles.visuallyHidden } role="status" aria-live="polite" aria-atomic="true">
-				{ copied ? copiedLabel : '' }
-			</span>
-		</div>
-	);
-}
-
 function CodeBlock( { children }: { children?: ReactNode } ) {
 	// Strip the single trailing newline react-markdown appends to fenced code.
 	const text = useMemo( () => extractText( children ).replace( /\n$/, '' ), [ children ] );
 
 	return (
 		<div className={ styles.codeBlock }>
 			<pre className={ styles.pre }>{ children }</pre>
-			{ text ? <CopyButton text={ text } /> : null }
+			{ text ? (
+				<CopyButton
+					text={ text }
+					label={ __( 'Copy code' ) }
+					className={ styles.copyButtonContainer }
+				/>
+			) : null }
 		</div>
 	);
 }