Forward tool execute args and emit chat artifacts for all remote stud…

…io tools

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

apps/cli/ai/runtimes/pi/index.ts

@@ -527,15 +527,12 @@ function buildAgentTools(
 	];
 
 	if ( isRemoteSite ) {
-		const screenshotTool = withChatArtifactEmission(
-			takeScreenshotTool,
-			chatArtifactsEnabled
-		) as unknown as AgentToolAny;
+		const remoteStudioTools = [ takeScreenshotTool, createSiteTool, pullSiteTool ].map( ( tool ) =>
+			withChatArtifactEmission( tool, chatArtifactsEnabled )
+		);
 		return [
 			createWpcomRequestTool( config.wpcomAccessToken!, config.activeSite!.wpcomSiteId! ),
-			screenshotTool,
-			createSiteTool,
-			pullSiteTool,
+			...remoteStudioTools,
 			...remoteScratchTools,
 			...askUserTool,
 			...skillTool,

apps/cli/ai/tests/tools.test.ts

@@ -25,8 +25,11 @@ import {
 	captureCommandOutput,
 	resolveStudioToolDefinitions,
 	studioToolDefinitions,
+	withChatArtifactEmission,
 } from '../tools';
+import { createSiteTool } from '../tools/create-site';
 import { enrichPreviewListOutput } from '../tools/list-previews';
+import type { AnyStudioAgentTool } from '../tools/define-tool';
 
 vi.mock( 'cli/ai/block-validator', () => ( {
 	validateBlocks: vi.fn(),
@@ -1022,6 +1025,63 @@ describe( 'Studio AI MCP tools', () => {
 		expect( getTextContent( result ) ).toContain( '"name": "My Site"' );
 	} );
 
+	describe( 'withChatArtifactEmission', () => {
+		const widget = { type: 'media', widgetProps: { mediaKind: 'image' } };
+		const makeFakeTool = () => {
+			const execute = vi.fn().mockResolvedValue( {
+				content: [ { type: 'text', text: 'ok' } ],
+				details: { studioArtifacts: [ widget ] },
+			} );
+			return {
+				tool: { name: 'fake_tool', execute } as unknown as AnyStudioAgentTool,
+				execute,
+			};
+		};
+
+		it( 'forwards execute arguments and preserves details when emitting', async () => {
+			const { tool, execute } = makeFakeTool();
+			const signal = new AbortController().signal;
+			const onUpdate = () => {};
+
+			const wrapped = withChatArtifactEmission( tool, true );
+			const result = await wrapped.execute( 'tool-call-9', { a: 1 } as never, signal, onUpdate );
+
+			expect( execute ).toHaveBeenCalledWith( 'tool-call-9', { a: 1 }, signal, onUpdate );
+			expect( result.details ).toEqual( { studioArtifacts: [ widget ] } );
+			expect( emitEvent ).toHaveBeenCalledWith(
+				expect.objectContaining( {
+					type: 'chat.artifact',
+					artifact: expect.objectContaining( { widgets: [ widget ] } ),
+				} )
+			);
+		} );
+
+		it( 'returns the tool unchanged when chat artifacts are disabled', () => {
+			const { tool } = makeFakeTool();
+			expect( withChatArtifactEmission( tool, false ) ).toBe( tool );
+		} );
+
+		it( 'emits site_create artifacts when wrapped directly (remote tool list)', async () => {
+			const wrapped = withChatArtifactEmission( createSiteTool, true );
+
+			await executeTool( wrapped, { name: 'My Site' } );
+
+			expect( emitEvent ).toHaveBeenCalledWith(
+				expect.objectContaining( {
+					type: 'chat.artifact',
+					artifact: expect.objectContaining( {
+						widgets: [
+							expect.objectContaining( {
+								type: 'site-preview',
+								widgetProps: expect.objectContaining( { siteId: 'site-123' } ),
+							} ),
+						],
+					} ),
+				} )
+			);
+		} );
+	} );
+
 	it( 'notifies JSON-mode callers when site_create selects the created site', async () => {
 		const onSiteSelected = vi.fn();
 		setLocalSiteSelectedCallback( onSiteSelected );

apps/cli/ai/tools/index.ts

@@ -27,7 +27,7 @@ import { updatePreviewTool } from './update-preview';
 import { validateBlocksTool } from './validate-blocks';
 import { waitForAnnotationsTool } from './wait-for-annotations';
 import { runWpCliTool } from './wp-cli';
-import type { AnyStudioAgentTool } from './define-tool';
+import type { AnyStudioAgentTool, StudioToolResultDetails } from './define-tool';
 
 export { captureCommandOutput } from './utils';
 
@@ -93,14 +93,16 @@ export function withChatArtifactEmission< TTool extends AnyStudioAgentTool >(
 	tool: TTool,
 	emitChatArtifacts: boolean
 ): TTool {
+	if ( ! emitChatArtifacts ) {
+		return tool;
+	}
 	return {
 		...tool,
-		execute: async ( _toolCallId, params ) => {
-			const result = await tool.rawHandler( params );
-			if ( emitChatArtifacts ) {
-				await emitChatArtifactWidgets( result.studioArtifacts );
-			}
-			return { content: result.content, details: undefined };
+		execute: async ( toolCallId, params, signal, onUpdate ) => {
+			const result = await tool.execute( toolCallId, params, signal, onUpdate );
+			const details = result.details as StudioToolResultDetails | undefined;
+			await emitChatArtifactWidgets( details?.studioArtifacts );
+			return result;
 		},
 	};
 }