From 392d2545d62ffd0e4007ebbd805f70c712a0cc24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gergely=20Cs=C3=A9csey?= Date: Wed, 29 Jul 2026 12:07:40 +0100 Subject: [PATCH] Send per-install id to the wpcom AI proxy as X-WPCOM-AI-Install-Id --- apps/cli/ai/providers.ts | 23 ++++++++++++++++++++++- apps/cli/ai/tests/auth.test.ts | 33 ++++++++++++++++++++++++++++++--- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/apps/cli/ai/providers.ts b/apps/cli/ai/providers.ts index fa1d6cbe4d..5702a48e40 100644 --- a/apps/cli/ai/providers.ts +++ b/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,6 +170,7 @@ 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; @@ -163,6 +178,9 @@ const AI_PROVIDER_DEFINITIONS: Record< AiProviderId, AiProviderDefinition > = { 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; } diff --git a/apps/cli/ai/tests/auth.test.ts b/apps/cli/ai/tests/auth.test.ts index e9e52e54ed..ec40eda7c5 100644 --- a/apps/cli/ai/tests/auth.test.ts +++ b/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' ); } );