Jetpack AI: show block edit outcomes with Undo

packages/agents-manager/src/components/__tests__/orchestrator-chat.test.tsx

@@ -8,6 +8,7 @@ import type { ComponentProps } from 'react';
 
 const mockUseAgentChat = jest.fn();
 const mockUseRegenerateAction = jest.fn();
+const mockUseCheckpointAction = jest.fn();
 const mockUseConversation = jest.fn();
 const mockUseImageUpload = jest.fn();
 const mockIsReaderChatAgent = jest.fn();
@@ -188,7 +189,10 @@ jest.mock( '../../utils/tracks', () => ( {
 } ) );
 jest.mock( '../../hooks/use-conversation', () => () => mockUseConversation() );
 jest.mock( '../../hooks/use-save-new-chat-route', () => () => {} );
-jest.mock( '../../hooks/use-checkpoint-action', () => () => {} );
+jest.mock( '../../hooks/use-checkpoint-action', () => ( {
+	__esModule: true,
+	default: ( ...args: unknown[] ) => mockUseCheckpointAction( ...args ),
+} ) );
 jest.mock( '../../hooks/use-feedback-action', () => () => ( {
 	showFeedbackInput: false,
 	submitFeedbackText: jest.fn(),
@@ -326,6 +330,7 @@ const countShowComponentMessages = () => {
 describe( 'OrchestratorChat', () => {
 	beforeEach( () => {
 		jest.clearAllMocks();
+		mockUseCheckpointAction.mockReturnValue( () => [] );
 		// Default getter: contributes no actions.
 		mockUseRegenerateAction.mockReturnValue( () => [] );
 		mockUseConversation.mockReturnValue( { isLoading: false } );
@@ -1024,6 +1029,79 @@ describe( 'OrchestratorChat', () => {
 		);
 	} );
 
+	it( 'derives and deduplicates checkpoint actions for synthetic streaming messages', () => {
+		const checkpointAction = {
+			id: 'checkpoint',
+			label: 'Undo',
+			onClick: jest.fn(),
+			order: 1,
+		};
+		const createOutcomeMessage = ( id: string, actions?: unknown[] ) => ( {
+			id,
+			role: 'agent',
+			content: [
+				{
+					type: 'text',
+					text: JSON.stringify( {
+						tool_id: 'big_sky__apply_block_edits',
+						tool_call_id: 'tool-call-1',
+						data: {
+							result: { success: true, outcome: 'updated', message: 'Updated the block.' },
+						},
+					} ),
+				},
+			],
+			timestamp: 1,
+			archived: false,
+			showIcon: true,
+			...( actions ? { actions } : {} ),
+		} );
+		const getCheckpointActions = jest.fn( ( message: { id: string } ) =>
+			message.id === 'agent-streaming-stale' ? [] : [ checkpointAction ]
+		);
+		mockUseCheckpointAction.mockReturnValue( getCheckpointActions );
+		mockUseAgentChat.mockReturnValue(
+			agentChatReturn( {
+				messages: [
+					createOutcomeMessage( 'agent-streaming-new' ),
+					createOutcomeMessage( 'agent-streaming-duplicate', [
+						{ ...checkpointAction, label: 'Old Undo' },
+					] ),
+					createOutcomeMessage( 'agent-streaming-stale', [ checkpointAction ] ),
+				],
+			} )
+		);
+
+		render( chat() );
+
+		const messages = mockAgentChat.mock.calls[ 0 ][ 0 ].messages as Array< {
+			id: string;
+			actions?: Array< { id: string; label: string } >;
+		} >;
+		const getCheckpointActionsFromMessage = ( id: string ) =>
+			messages
+				.find( ( message ) => message.id === id )
+				?.actions?.filter( ( action ) => action.id === 'checkpoint' ) ?? [];
+
+		expect( getCheckpointActionsFromMessage( 'agent-streaming-new' ) ).toEqual( [
+			expect.objectContaining( { id: 'checkpoint', label: 'Undo' } ),
+		] );
+		expect( getCheckpointActionsFromMessage( 'agent-streaming-duplicate' ) ).toEqual( [
+			expect.objectContaining( { id: 'checkpoint', label: 'Undo' } ),
+		] );
+		expect( getCheckpointActionsFromMessage( 'agent-streaming-stale' ) ).toEqual( [] );
+		expect( getCheckpointActions ).toHaveBeenCalledWith(
+			expect.objectContaining( {
+				id: 'agent-streaming-new',
+				content: expect.arrayContaining( [
+					expect.objectContaining( {
+						text: expect.stringContaining( 'big_sky__apply_block_edits' ),
+					} ),
+				] ),
+			} )
+		);
+	} );
+
 	it( 'enables regenerate when a provider opts in', () => {
 		render( chat( { capabilities: { supportsRegenerateAction: true } } ) );
 

packages/agents-manager/src/components/agent-dock/chat-ui.scss

@@ -64,6 +64,60 @@ body.wp-admin.modal-open .agents-manager-chat--docked {
 				fill: var( --color-foreground );
 			}
 		}
+
+		.agents-manager-resolved-edit-action {
+			display: inline-flex;
+			flex: 0 0 100%;
+			align-items: center;
+			gap: 0.25rem;
+
+			&__status,
+			&__undo {
+				display: inline-flex;
+				align-items: center;
+				gap: 0.25rem;
+				min-height: 32px;
+				padding: 0.4rem 0.75rem;
+				border-radius: 4px;
+				font-size: 13px;
+			}
+
+			&__status {
+				// Agenttic's success token is white in dark mode; match the Editorial Review treatment.
+				background: rgba( 74, 184, 102, 0.2 );
+				color: #4ab866;
+				font-weight: 600;
+			}
+
+			&__undo {
+				border: 0;
+				background: transparent;
+				color: var( --color-foreground );
+				font-family: inherit;
+				font-weight: 500;
+				cursor: pointer;
+
+				&:hover:not( :disabled ) {
+					background: var( --color-muted );
+				}
+
+				&:focus-visible:not( :disabled ) {
+					background: var( --color-muted );
+					outline: 2px solid var( --wp-admin-theme-color, #3858e9 );
+					outline-offset: 1px;
+				}
+
+				&:disabled {
+					cursor: default;
+					opacity: 0.5;
+				}
+			}
+
+			&__icon {
+				flex: 0 0 auto;
+				fill: currentColor;
+			}
+		}
 	}
 
 	// Prevent content overflow

packages/agents-manager/src/components/orchestrator-chat/index.tsx

@@ -442,7 +442,7 @@ export default function OrchestratorChat( {
 
 	// Register an "Undo" action on agent messages with checkpoints.
 	const checkpoint = useCheckpoint?.();
-	useCheckpointAction( registerMessageActions, checkpoint );
+	const getCheckpointActionsForMessage = useCheckpointAction( registerMessageActions, checkpoint );
 
 	// Register thumbs-up/down feedback actions on agent messages.
 	const { showFeedbackInput, submitFeedbackText, resetFeedback, getFeedbackActionsForMessage } =
@@ -849,6 +849,13 @@ export default function OrchestratorChat( {
 			);
 		}
 
+		const checkpointActionsByMessageId = new Map(
+			currentMessages.map( ( message ) => [
+				message.id,
+				getCheckpointActionsForMessage( message ),
+			] )
+		);
+
 		// Group site-build messages only when needed
 		const hasBuildMessages = siteBuildUtils?.hasSiteBuildMessages( currentMessages );
 
@@ -878,19 +885,24 @@ export default function OrchestratorChat( {
 			}
 
 			const directActions = [
+				...( checkpointActionsByMessageId.get( message.id ) ?? [] ),
 				...getFeedbackActionsForMessage( message ),
 				...getCopyActionsForMessage( message ),
 				...getRegenerateActionsForMessage( message, {
 					isLatestAgentMessage: message.id === latestAgentMessageId,
 					isStreaming: isProcessing,
 				} ),
 			];
-			if ( directActions.length === 0 ) {
+			const hasRegisteredCheckpointAction = message.actions?.some(
+				( action ) => action.id === 'checkpoint'
+			);
+			if ( directActions.length === 0 && ! hasRegisteredCheckpointAction ) {
 				return messageWithTraceId;
 			}
 
 			const existingActions = message.actions?.filter(
 				( action ) =>
+					action.id !== 'checkpoint' &&
 					! action.id.startsWith( 'feedback-' ) &&
 					action.id !== 'copy' &&
 					action.id !== 'regenerate'
@@ -910,6 +922,7 @@ export default function OrchestratorChat( {
 		deletedMessageIds,
 		getChatComponent,
 		getCopyActionsForMessage,
+		getCheckpointActionsForMessage,
 		getShowComponentOrder,
 		getFeedbackActionsForMessage,
 		getTraceIdForMessage,

packages/agents-manager/src/components/resolved-edit-action.tsx

@@ -0,0 +1,43 @@
+import { useState } from '@wordpress/element';
+import { __ } from '@wordpress/i18n';
+import { check, Icon, undo } from '@wordpress/icons';
+
+type ResolvedEditActionProps = {
+	onUndo: () => Promise< boolean >;
+};
+
+export default function ResolvedEditAction( { onUndo }: ResolvedEditActionProps ) {
+	const [ isUndoDisabled, setIsUndoDisabled ] = useState( false );
+	const handleUndo = async () => {
+		if ( isUndoDisabled ) {
+			return;
+		}
+
+		setIsUndoDisabled( true );
+		try {
+			if ( ! ( await onUndo() ) ) {
+				setIsUndoDisabled( false );
+			}
+		} catch {
+			setIsUndoDisabled( false );
+		}
+	};
+
+	return (
+		<div className="agents-manager-resolved-edit-action">
+			<span className="agents-manager-resolved-edit-action__status" role="status">
+				<Icon className="agents-manager-resolved-edit-action__icon" icon={ check } size={ 20 } />
+				{ __( 'Updated', __i18n_text_domain__ ) }
+			</span>
+			<button
+				type="button"
+				className="agents-manager-resolved-edit-action__undo"
+				onClick={ () => void handleUndo() }
+				disabled={ isUndoDisabled }
+			>
+				<Icon className="agents-manager-resolved-edit-action__icon" icon={ undo } size={ 20 } />
+				{ __( 'Undo', __i18n_text_domain__ ) }
+			</button>
+		</div>
+	);
+}