Add WordPress.com REST API support for remote sites in AI command

apps/cli/ai/agent.ts

@@ -2,13 +2,15 @@ import path from 'path';
 import { query, type Query } from '@anthropic-ai/claude-agent-sdk';
 import {
 	ALLOWED_TOOLS,
+	ALLOWED_TOOLS_REMOTE,
 	STUDIO_ROOT,
 	createPathApprovalSession,
 	promptForApproval,
 	type AskUserQuestion,
 } from 'cli/ai/security';
 import { buildSystemPrompt } from 'cli/ai/system-prompt';
-import { createStudioTools } from 'cli/ai/tools';
+import { createRemoteSiteTools, createStudioTools } from 'cli/ai/tools';
+import type { SiteInfo } from 'cli/ai/ui';
 
 export type { AskUserQuestion } from 'cli/ai/security';
 
@@ -18,6 +20,8 @@ export interface AiAgentConfig {
 	model?: AiModelId;
 	maxTurns?: number;
 	resume?: string;
+	activeSite?: SiteInfo | null;
+	wpcomAccessToken?: string;
 	onAskUser?: ( questions: AskUserQuestion[] ) => Promise< Record< string, string > >;
 }
 
@@ -48,25 +52,55 @@ process.on( 'unhandledRejection', ( reason ) => {
  * Caller can iterate messages with `for await` and call `interrupt()` to stop.
  */
 export function startAiAgent( config: AiAgentConfig ): Query {
-	const { prompt, env, model = DEFAULT_MODEL, maxTurns = 50, resume, onAskUser } = config;
+	const {
+		prompt,
+		env,
+		model = DEFAULT_MODEL,
+		maxTurns = 50,
+		resume,
+		activeSite,
+		wpcomAccessToken,
+		onAskUser,
+	} = config;
 	const resolvedEnv = env ?? { ...( process.env as Record< string, string > ) };
 
+	const isRemoteSite = activeSite?.remote && activeSite?.wpcomSiteId && wpcomAccessToken;
+
+	// Configure MCP servers based on site type:
+	// Remote sites get WP.com REST API tools + screenshot; local sites get the full Studio toolset.
+	const mcpServers = {
+		studio: isRemoteSite
+			? createRemoteSiteTools( wpcomAccessToken, activeSite.wpcomSiteId! )
+			: createStudioTools(),
+	};
+
+	const allowedTools = isRemoteSite ? [ ...ALLOWED_TOOLS_REMOTE ] : [ ...ALLOWED_TOOLS ];
+
+	// Build site-aware system prompt
+	const systemPromptOptions = isRemoteSite
+		? {
+				remoteSite: {
+					name: activeSite.name,
+					url: activeSite.url ?? '',
+					id: activeSite.wpcomSiteId!,
+				},
+		  }
+		: undefined;
+
 	return query( {
 		prompt,
 		options: {
 			env: resolvedEnv,
 			systemPrompt: {
 				type: 'preset',
 				preset: 'claude_code',
-				append: buildSystemPrompt(),
-			},
-			mcpServers: {
-				studio: createStudioTools(),
+				append: buildSystemPrompt( systemPromptOptions ),
 			},
+			mcpServers,
 			maxTurns,
 			cwd: STUDIO_ROOT,
 			tools: { type: 'preset', preset: 'claude_code' },
-			allowedTools: [ ...ALLOWED_TOOLS ],
+			allowedTools,
 			permissionMode: 'default',
 			canUseTool: async ( toolName, input, metadata ) => {
 				if ( toolName === 'AskUserQuestion' && onAskUser ) {

apps/cli/ai/security.ts

@@ -39,6 +39,19 @@ export const ALLOWED_TOOLS = [
 	'AskUserQuestion',
 ] as const;
 
+// Tools allowed when operating on a remote WordPress.com site
+export const ALLOWED_TOOLS_REMOTE = [
+	'mcp__studio__*',
+	'Read',
+	'Glob',
+	'Grep',
+	'WebFetch',
+	'WebSearch',
+	'TodoRead',
+	'NotebookRead',
+	'AskUserQuestion',
+] as const;
+
 // Tools that should not manipulate files outside trusted roots without permission (write access)
 const PATH_GATED_TOOLS = [ 'Write', 'Edit', 'Bash', 'NotebookEdit' ] as const;
 const PATH_INPUT_KEYS = [ 'path', 'file_path', 'filePath' ] as const;

apps/cli/ai/sessions/recorder.ts

@@ -87,6 +87,7 @@ export class AiSessionRecorder {
 		path: string;
 		remote?: boolean;
 		url?: string;
+		wpcomSiteId?: number;
 	} ): Promise< void > {
 		await this.appendEvent( {
 			type: 'site.selected',
@@ -95,6 +96,7 @@ export class AiSessionRecorder {
 			sitePath: site.path,
 			remote: site.remote,
 			url: site.url,
+			wpcomSiteId: site.wpcomSiteId,
 		} );
 	}
 

apps/cli/ai/sessions/types.ts

@@ -27,6 +27,7 @@ export type AiSessionEvent =
 			sitePath: string;
 			remote?: boolean;
 			url?: string;
+			wpcomSiteId?: number;
 	  }
 	| {
 			type: 'user.message';

apps/cli/ai/system-prompt.ts

@@ -1,11 +1,101 @@
-export function buildSystemPrompt(): string {
+interface RemoteSiteContext {
+	name: string;
+	url: string;
+	id: number;
+}
+
+export function buildSystemPrompt( options?: { remoteSite?: RemoteSiteContext } ): string {
+	if ( options?.remoteSite ) {
+		return `${ buildRemoteIntro( options.remoteSite ) }
+
+${ REMOTE_CONTENT_GUIDELINES }
+
+${ REMOTE_DESIGN_GUIDELINES }
+`;
+	}
+
+	return `${ buildLocalIntro() }
+
+${ LOCAL_CONTENT_GUIDELINES }
+
+${ LOCAL_DESIGN_GUIDELINES }
+`;
+}
+
+function buildRemoteIntro( site: RemoteSiteContext ): string {
+	return `You are WordPress Studio AI, the AI assistant built into WordPress Studio CLI. Your name is "WordPress Studio AI". You manage WordPress.com sites using the WordPress.com REST API.
+
+IMPORTANT: The active site is a remote WordPress.com site: "${ site.name }" (ID: ${ site.id }) at ${ site.url }.
+IMPORTANT: You MUST use the wpcom_request tool (prefixed with mcp__studio__) to manage this site. Do NOT use WP-CLI, file Write/Edit, Bash, or any local file operations — this site is hosted on WordPress.com and cannot be modified through the local filesystem.
+IMPORTANT: Before doing ANY work, you MUST first check the site's plan by calling \`GET /\` (apiNamespace: \`""\`). The \`plan.product_slug\` field indicates the plan. If the site is on a free plan (e.g. \`free_plan\`), you MUST refuse design customization requests — this includes custom CSS, inline styles, style attributes on blocks, global styles editing, custom JavaScript, animations, custom colors/fonts/layouts, and plugin management. Do NOT attempt workarounds like inline styles or style block attributes — these produce invalid blocks on WordPress.com. Instead, tell the user that design customizations require upgrading to a paid WordPress.com plan and STOP. Do not proceed with the design task.
+
+## Available Tools (prefixed with mcp__studio__)
+
+- **wpcom_request**: A REST API client that supports both the WordPress REST API (wp/v2) and the WordPress.com REST API (v1.1).
+  - \`method\`: GET, POST, PUT, or DELETE
+  - \`path\`: Relative to \`/sites/{siteId}/\` (e.g., \`/posts\`, \`/posts/123\`, \`/templates\`). Prefix with \`!\` for absolute paths (e.g., \`!/me\`).
+  - \`query\`: Optional query parameters object
+  - \`body\`: Optional request body for POST/PUT
+  - \`apiNamespace\`: Defaults to \`"wp/v2"\`. Set to \`""\` (empty string) for WP.com REST API v1.1, or \`"wpcom/v2"\` for WP.com v2 endpoints.
+- **take_screenshot**: Take a full-page screenshot of a URL (supports desktop and mobile viewports)
+
+## API Namespace Guide
+
+**Prefer wp/v2** (default — standard WordPress REST API) for most resources:
+- Posts, pages, media, categories, tags, users, comments
+- Templates, template parts, navigation, global styles, block patterns
+- Any standard WordPress resource
+
+**Use WP.com v1.1** (set \`apiNamespace: ""\`) for WP.com-specific endpoints:
+- Plugin management: \`/plugins\`, \`/plugins/{slug}/install\`
+- Theme switching: \`/themes/mine\`
+- Site info: \`/\` (root)
+- Site settings: \`/settings\`
+
+## Common wp/v2 Endpoints (default apiNamespace)
+
+**Posts & Pages**: \`GET /posts\`, \`GET /posts/{id}\`, \`POST /posts\`, \`POST /posts/{id}\`, \`DELETE /posts/{id}\`
+**Media**: \`GET /media\`, \`POST /media\`
+**Templates**: \`GET /templates\`, \`GET /templates/{id}\`, \`POST /templates\`, \`POST /templates/{id}\`, \`DELETE /templates/{id}\`
+**Template Parts**: \`GET /template-parts\`, \`GET /template-parts/{id}\`, \`POST /template-parts\`, \`POST /template-parts/{id}\`
+**Navigation**: \`GET /navigation\`, \`POST /navigation\`, \`POST /navigation/{id}\`
+**Global Styles**: \`GET /global-styles/{id}\`, \`POST /global-styles/{id}\`. To find the global styles ID, first \`GET /themes?status=active\` — the active theme's \`_links["wp:user-global-styles"][0].href\` contains the ID.
+**Categories/Tags**: \`GET /categories\`, \`POST /categories\`, \`GET /tags\`, \`POST /tags\`
+**Block Types**: \`GET /block-types\`, \`GET /block-types/{name}\`
+**Search**: \`GET /search?search={query}\`
+
+Use \`per_page\` and \`page\` for pagination. Use \`status\` to filter by publish status. For creating/updating content, pass block markup in the \`content\` field of the body.
+
+## Common WP.com v1.1 Endpoints (set apiNamespace to "")
+
+**Site**: \`GET /\` (site info), \`POST /settings\`
+**Plugins**: \`GET /plugins\`, \`POST /plugins/{slug}/install\`, \`POST /plugins/{slug}\` (body: \`{ active: true/false }\`)
+**Themes**: \`GET /themes\`, \`POST /themes/mine\` (body: \`{ theme: "slug" }\`)
+**Media upload from URL**: \`POST /media/new\` (body: \`{ media_urls: [...] }\`)
+
+## Workflow
+
+1. **Check the site plan** (MANDATORY FIRST STEP): Use \`GET /\` (apiNamespace: \`""\`) to get site info and check \`plan.product_slug\`. Stop and inform the user if they request features unavailable on their plan.
+2. **Understand the site**: Use \`GET /posts\` to list content, \`GET /themes?status=active\` to see the active theme.
+3. **Make changes**: Use POST requests to create/update content, manage templates, switch themes.
+4. **Verify visually**: Use take_screenshot to capture the site on desktop and mobile viewports. Check spacing, alignment, colors, contrast, and layout. Fix any issues.
+
+## General rules
+
+- Always confirm destructive operations (deleting posts, deactivating plugins, etc.) with the user before proceeding.
+- When creating content, follow WordPress best practices for block-based content.
+- If a requested operation fails, check the error message and suggest alternatives.
+- Explore the API — if you're unsure about an endpoint, try a GET request first to discover available data.`;
+}
+
+function buildLocalIntro(): string {
 	return `You are WordPress Studio AI, the AI assistant built into WordPress Studio CLI. Your name is "WordPress Studio AI". You manage and modify local WordPress sites using your Studio tools and generate content for these sites.
 
 IMPORTANT: You MUST use your mcp__studio__ tools to manage WordPress sites. Never create, start, or stop sites using Bash commands, shell scripts, or manual file operations. Never run \`wp\` commands via Bash — always use the wp_cli tool instead. The Studio tools handle all server management, database setup, and WordPress provisioning automatically.
 IMPORTANT: For any generated content for the site, these three principles are mandatory:
 
 - Gorgeous design: More details on the guidelines below.
-- No HTML blocks and raw HTML: Check the block content guidelines below. 
+- No HTML blocks and raw HTML: Check the block content guidelines below.
 - No invalid block: Use the validate_blocks everytime to ensure that the blocks are 100% valid.
 
 ## Workflow
@@ -52,23 +142,41 @@ Then continue with:
 - Always add the style.css as editor styles in the functions.php of the theme to make the editor match the frontend.
 - For theme and page content custom CSS, put the styles in the main style.css of the theme. No custom stylesheets.
 - Scroll animations must use progressive enhancement: CSS defines elements in their **final visible state** by default (full opacity, final position). JavaScript on the frontend adds the initial hidden state (e.g. \`opacity: 0\`, \`transform\`) and scroll-triggered transitions. This ensures elements are fully visible in the block editor (which loads theme CSS but not custom JS).
-- All animations and transitions must respect \`prefers-reduced-motion\`. Add a \`@media (prefers-reduced-motion: reduce)\` block that disables or simplifies animations (e.g. \`animation: none; transition: none; scroll-behavior: auto;\`).
+- All animations and transitions must respect \`prefers-reduced-motion\`. Add a \`@media (prefers-reduced-motion: reduce)\` block that disables or simplifies animations (e.g. \`animation: none; transition: none; scroll-behavior: auto;\`).`;
+}
+
+const REMOTE_CONTENT_GUIDELINES = `## Block content guidelines
+
+- Use only core WordPress blocks. No custom HTML blocks except for inline SVGs.
+- No decorative HTML comments (e.g. \`<!-- Hero Section -->\`). Only block delimiter comments are allowed.
+- No emojis anywhere in generated content.`;
+
+const REMOTE_DESIGN_GUIDELINES = `## Design capabilities by plan
+
+**Free plans** — content only, no design customization:
+- CAN: Create/edit posts, pages, templates, template parts. Switch themes. Upload media.
+- CANNOT: Any visual/design customization including custom CSS, inline styles, style attributes on blocks, global styles, custom JavaScript, animations, custom colors, custom fonts, custom layouts, or plugin management.
+- ACTION: If the user requests ANY design change — even "small" ones like changing a color or font — you MUST refuse, explain it requires a paid plan, and STOP. Do not suggest inline styles, style attributes, or any other workaround. These will produce invalid blocks.
 
-## Block content guidelines
+**Paid plans** (Personal, Premium, Business, eCommerce) — progressively more control:
+- Custom CSS, global styles, plugin management, and advanced customization become available.
+- Check the specific plan to determine exact capabilities.`;
 
-- Only use \`core/html\` blocks for:                                                                                               
-	- Inline SVGs                                                                                                                  
-	- \`<form>\` elements and interactive inputs                                                                                     
-	- Animation/interaction markup with no block equivalent (marquee, cursor)                                                      
-	- A single \`<script>\` block at the bottom of the page for JS                                                                   
-- Never use \`core/html\` to wrap text content, headings, layout sections, or lists. 
+const LOCAL_CONTENT_GUIDELINES = `## Block content guidelines
+
+- Only use \`core/html\` blocks for:
+	- Inline SVGs
+	- \`<form>\` elements and interactive inputs
+	- Animation/interaction markup with no block equivalent (marquee, cursor)
+	- A single \`<script>\` block at the bottom of the page for JS
+- Never use \`core/html\` to wrap text content, headings, layout sections, or lists.
 - No decorative HTML comments (e.g. \`<!-- Hero Section -->\`, \`<!-- Features -->\`). Only block delimiter comments are allowed.
 - No custom class names on inner DOM elements — only on the outermost block wrapper via the \`className\` attribute.
 - No inline \`style\` or \`style\` block attributes for styling. Use \`className\` + \`style.css\` instead.
 - Use \`core/spacer\` for empty spacing divs, not \`core/group\`.
-- No emojis anywhere in generated content.
+- No emojis anywhere in generated content.`;
 
-## Design guidelines
+const LOCAL_DESIGN_GUIDELINES = `## Design guidelines
 
 **Important**: Always use sophisticated scroll effects and add animations unless specifically asked otherwise.
 
@@ -99,7 +207,4 @@ Interpret creatively and make unexpected choices that feel genuinely designed fo
 
 **IMPORTANT**: Match implementation complexity to the aesthetic vision. Maximalist designs need elaborate code with extensive animations and effects. Minimalist or refined designs need restraint, precision, and careful attention to spacing, typography, and subtle details. Elegance comes from executing the vision well.
 
-Remember: You are capable of extraordinary creative work. Don't hold back, show what can truly be created when thinking outside the box and committing fully to a distinctive vision.
-
-`;
-}
+Remember: You are capable of extraordinary creative work. Don't hold back, show what can truly be created when thinking outside the box and committing fully to a distinctive vision.`;

apps/cli/ai/tools.ts

@@ -6,6 +6,7 @@ import { z } from 'zod/v4';
 import { validateBlocks, type ValidationReport } from 'cli/ai/block-validator';
 import { getSharedBrowser } from 'cli/ai/browser-utils';
 import { auditPerformance } from 'cli/ai/performance-audit';
+import { createWpcomToolDefinitions } from 'cli/ai/wpcom-tools';
 import { runCommand as runCreatePreviewCommand } from 'cli/commands/preview/create';
 import {
 	Mode as PreviewDeleteMode,
@@ -637,10 +638,12 @@ const takeScreenshotTool = tool(
 				// Reduce motion to avoid capturing mid-animation states
 				await page.emulateMedia( { reducedMotion: 'reduce' } );
 
-				await page.goto( args.url, { waitUntil: 'networkidle', timeout: 15000 } );
+				await page.goto( args.url, { waitUntil: 'domcontentloaded', timeout: 30000 } );
+				await page.waitForLoadState( 'networkidle', { timeout: 10000 } ).catch( () => {} );
 
 				// Scroll through the page to trigger lazy-loaded images, then wait
-				// for all images to finish loading.
+				// for all images to finish loading (with a timeout so we don't hang
+				// on images that never settle).
 				await page.evaluate( async () => {
 					const delay = ( ms: number ) =>
 						new Promise< void >( ( resolve ) => setTimeout( resolve, ms ) );
@@ -652,7 +655,8 @@ const takeScreenshotTool = tool(
 					}
 					window.scrollTo( 0, 0 );
 
-					await Promise.all(
+					const timeout = new Promise< void >( ( resolve ) => setTimeout( resolve, 5000 ) );
+					const allImages = Promise.all(
 						Array.from( document.images )
 							.filter( ( img ) => ! img.complete )
 							.map(
@@ -663,6 +667,7 @@ const takeScreenshotTool = tool(
 									} )
 							)
 					);
+					await Promise.race( [ allImages, timeout ] );
 				} );
 
 				// Hide WordPress admin bar and scrollbars for cleaner screenshots
@@ -797,3 +802,16 @@ export function createStudioTools() {
 		tools: studioToolDefinitions,
 	} );
 }
+
+/**
+ * Creates an MCP server for remote WordPress.com sites, combining WP.com REST API tools
+ * with URL-based tools (screenshot) that work with any site.
+ */
+export function createRemoteSiteTools( token: string, siteId: number ) {
+	const wpcomTools = createWpcomToolDefinitions( token, siteId );
+	return createSdkMcpServer( {
+		name: 'studio',
+		version: '1.0.0',
+		tools: [ ...wpcomTools, takeScreenshotTool ],
+	} );
+}

apps/cli/ai/ui.ts

@@ -56,6 +56,7 @@ export interface SiteInfo {
 	running: boolean;
 	remote?: boolean;
 	url?: string;
+	wpcomSiteId?: number;
 }
 
 const DEFAULT_COLLAPSE_THRESHOLD_LINES = 5;
@@ -828,6 +829,7 @@ export class AiChatUI {
 				running: false,
 				remote: true,
 				url: site.url,
+				wpcomSiteId: site.id,
 			} ) );
 			this.sitePickerRemoteLoading = false;
 			this.rebuildSitePickerList();