Send per-install id to the wpcom AI proxy

apps/cli/ai/providers.ts

@@ -5,7 +5,7 @@ import {
 	type AiModelFamily,
 	type AiModelId,
 } from '@studio/common/ai/models';
-import { readAuthToken } from '@studio/common/lib/shared-config';
+import { getOrCreateAnalyticsInstallId, readAuthToken } from '@studio/common/lib/shared-config';
 import { __ } from '@wordpress/i18n';
 import { readCliConfig, updateCliConfigWithPartial } from 'cli/lib/cli-config/core';
 import { LoggerError } from 'cli/logger';
@@ -27,6 +27,9 @@ const DEFAULT_WPCOM_AI_GATEWAY_BASE_URL = 'https://public-api.wordpress.com/wpco
 // the existing slugs so no server-side allowlist change is required.
 const WPCOM_AI_FEATURE_HEADER_ANTHROPIC = 'studio-assistant-anthropic';
 const WPCOM_AI_FEATURE_HEADER_OPENAI = 'studio-assistant';
+// Per-install identity the wpcom AI proxy uses to enforce a per-installation burst cap.
+// Generic (proxy-scoped, not Studio-scoped) like the other X-WPCOM-AI-* headers.
+const WPCOM_AI_INSTALL_ID_HEADER = 'X-WPCOM-AI-Install-Id';
 
 export interface ResolveAiEnvironmentOptions {
 	sessionId?: string;
@@ -102,6 +105,17 @@ function buildAnthropicCustomHeaders( headers: Record< string, string > ): strin
 		.join( '\n' );
 }
 
+// The per-install id is the anonymous Tracks install UUID (random, stable per install, shared by
+// Studio and the CLI). The server treats an absent header as "no per-install cap", so a shared-config
+// hiccup must never fail an AI request — fall back to omitting the header.
+async function resolveInstallId(): Promise< string | undefined > {
+	try {
+		return await getOrCreateAnalyticsInstallId();
+	} catch {
+		return undefined;
+	}
+}
+
 function getWpcomAiGatewayBaseUrl(): string {
 	const customBaseUrl = process.env.WPCOM_AI_PROXY_BASE_URL?.trim();
 	return customBaseUrl || DEFAULT_WPCOM_AI_GATEWAY_BASE_URL;
@@ -156,13 +170,17 @@ const AI_PROVIDER_DEFINITIONS: Record< AiProviderId, AiProviderDefinition > = {
 			}
 			const env = createBaseEnvironment();
 			const gatewayBaseUrl = getWpcomAiGatewayBaseUrl();
+			const installId = await resolveInstallId();
 
 			// Anthropic messages path through the WP.com AI gateway.
 			env.ANTHROPIC_BASE_URL = gatewayBaseUrl;
 			env.ANTHROPIC_AUTH_TOKEN = accessToken;
 			const anthropicHeaders: Record< string, string > = {
 				'X-WPCOM-AI-Feature': WPCOM_AI_FEATURE_HEADER_ANTHROPIC,
 			};
+			if ( installId ) {
+				anthropicHeaders[ WPCOM_AI_INSTALL_ID_HEADER ] = installId;
+			}
 			if ( options?.sessionId ) {
 				anthropicHeaders[ 'X-WPCOM-Session-ID' ] = options.sessionId;
 			}
@@ -178,6 +196,9 @@ const AI_PROVIDER_DEFINITIONS: Record< AiProviderId, AiProviderDefinition > = {
 			const openaiHeaders: Record< string, string > = {
 				'X-WPCOM-AI-Feature': WPCOM_AI_FEATURE_HEADER_OPENAI,
 			};
+			if ( installId ) {
+				openaiHeaders[ WPCOM_AI_INSTALL_ID_HEADER ] = installId;
+			}
 			if ( options?.sessionId ) {
 				openaiHeaders[ 'X-WPCOM-Session-ID' ] = options.sessionId;
 			}

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

@@ -1,5 +1,5 @@
 import { password } from '@inquirer/prompts';
-import { readAuthToken } from '@studio/common/lib/shared-config';
+import { getOrCreateAnalyticsInstallId, readAuthToken } from '@studio/common/lib/shared-config';
 import { vi } from 'vitest';
 import {
 	getAvailableAiProviders,
@@ -18,6 +18,7 @@ vi.mock( '@inquirer/prompts', () => ( {
 
 vi.mock( '@studio/common/lib/shared-config', () => ( {
 	readAuthToken: vi.fn(),
+	getOrCreateAnalyticsInstallId: vi.fn(),
 } ) );
 
 vi.mock( 'cli/lib/cli-config/core', () => ( {
@@ -29,6 +30,7 @@ describe( 'AI auth helpers', () => {
 	beforeEach( () => {
 		vi.resetAllMocks();
 		vi.mocked( readCliConfig ).mockResolvedValue( { version: 1, sites: [], snapshots: [] } );
+		vi.mocked( getOrCreateAnalyticsInstallId ).mockResolvedValue( 'install-uuid' );
 		delete process.env.WPCOM_AI_PROXY_BASE_URL;
 	} );
 
@@ -104,10 +106,35 @@ describe( 'AI auth helpers', () => {
 			'https://public-api.wordpress.com/wpcom/v2/ai-api-proxy'
 		);
 		expect( env.ANTHROPIC_AUTH_TOKEN ).toBe( 'wpcom-token' );
-		expect( env.ANTHROPIC_CUSTOM_HEADERS ).toBe( 'X-WPCOM-AI-Feature: studio-assistant-anthropic' );
+		expect( env.ANTHROPIC_CUSTOM_HEADERS ).toBe(
+			'X-WPCOM-AI-Feature: studio-assistant-anthropic\nX-WPCOM-AI-Install-Id: install-uuid'
+		);
+		expect( JSON.parse( env.STUDIO_OPENAI_DEFAULT_HEADERS ) ).toEqual( {
+			'X-WPCOM-AI-Feature': 'studio-assistant',
+			'X-WPCOM-AI-Install-Id': 'install-uuid',
+		} );
 		expect( env.ANTHROPIC_API_KEY ).toBeUndefined();
 	} );
 
+	it( 'omits the install id header when the shared install id is unavailable', async () => {
+		vi.mocked( readAuthToken ).mockResolvedValue( {
+			accessToken: 'wpcom-token',
+			displayName: 'User',
+			email: 'user@example.com',
+			expiresIn: 3600,
+			expirationTime: Date.now() + 3600_000,
+			id: 1,
+		} );
+		vi.mocked( getOrCreateAnalyticsInstallId ).mockRejectedValue( new Error( 'config locked' ) );
+
+		const env = await resolveAiEnvironment( 'wpcom' );
+
+		expect( env.ANTHROPIC_CUSTOM_HEADERS ).toBe( 'X-WPCOM-AI-Feature: studio-assistant-anthropic' );
+		expect( JSON.parse( env.STUDIO_OPENAI_DEFAULT_HEADERS ) ).toEqual( {
+			'X-WPCOM-AI-Feature': 'studio-assistant',
+		} );
+	} );
+
 	it( 'includes the Studio AI session ID header when provided', async () => {
 		vi.mocked( readAuthToken ).mockResolvedValue( {
 			accessToken: 'wpcom-token',
@@ -121,7 +148,7 @@ describe( 'AI auth helpers', () => {
 		const env = await resolveAiEnvironment( 'wpcom', { sessionId: 'session-abc' } );
 
 		expect( env.ANTHROPIC_CUSTOM_HEADERS ).toBe(
-			'X-WPCOM-AI-Feature: studio-assistant-anthropic\nX-WPCOM-Session-ID: session-abc'
+			'X-WPCOM-AI-Feature: studio-assistant-anthropic\nX-WPCOM-AI-Install-Id: install-uuid\nX-WPCOM-Session-ID: session-abc'
 		);
 	} );