diff --git a/apps/cli/ai/providers.ts b/apps/cli/ai/providers.ts index 718213eaf6..35c72514ea 100644 --- a/apps/cli/ai/providers.ts +++ b/apps/cli/ai/providers.ts @@ -24,13 +24,11 @@ export { DEFAULT_AI_PROVIDER }; export const AI_PROVIDER_PRIORITY: readonly AiProviderId[] = AI_PROVIDER_IDS; const DEFAULT_WPCOM_AI_GATEWAY_BASE_URL = 'https://public-api.wordpress.com/wpcom/v2/ai-api-proxy'; -// The wpcom AI proxy maps feature slugs to upstream providers. Historically -// `studio-assistant` was wired for OpenAI (OPENAI_TOKEN); when Claude support -// landed a parallel `studio-assistant-anthropic` slug was added. Keep using -// 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'; -const WPCOM_AI_FEATURE_HEADER_HOSTED = 'studio-assistant-hosted'; +// The wpcom AI proxy maps feature slugs to upstream providers. The +// `studio-agent` lane accepts the capability-tier aliases (fast / balanced / +// strong) on the Chat Completions path and resolves each to an upstream +// model server-side. +const WPCOM_AI_FEATURE_HEADER = 'studio-agent'; export interface ResolveAiEnvironmentOptions { sessionId?: string; @@ -103,12 +101,6 @@ export function getStudioUserAgent(): string { return version ? `WordPressStudio/${ version }` : 'WordPressStudio'; } -function buildAnthropicCustomHeaders( headers: Record< string, string > ): string { - return Object.entries( headers ) - .map( ( [ name, value ] ) => `${ name }: ${ value }` ) - .join( '\n' ); -} - export function getWpcomAiGatewayBaseUrl(): string { const customBaseUrl = process.env.WPCOM_AI_PROXY_BASE_URL?.trim(); return customBaseUrl || DEFAULT_WPCOM_AI_GATEWAY_BASE_URL; @@ -136,10 +128,9 @@ function createBaseEnvironment(): Record< string, string > { delete env.ANTHROPIC_CUSTOM_HEADERS; delete env.OPENAI_API_KEY; delete env.OPENAI_BASE_URL; - delete env.STUDIO_OPENAI_DEFAULT_HEADERS; - delete env.STUDIO_HOSTED_API_KEY; - delete env.STUDIO_HOSTED_BASE_URL; - delete env.STUDIO_HOSTED_DEFAULT_HEADERS; + delete env.STUDIO_WPCOM_API_KEY; + delete env.STUDIO_WPCOM_BASE_URL; + delete env.STUDIO_WPCOM_DEFAULT_HEADERS; return env; } @@ -166,50 +157,21 @@ const AI_PROVIDER_DEFINITIONS: Record< AiProviderId, AiProviderDefinition > = { const env = createBaseEnvironment(); const gatewayBaseUrl = getWpcomAiGatewayBaseUrl(); - // Anthropic messages path through the WP.com AI gateway. - env.ANTHROPIC_BASE_URL = gatewayBaseUrl; - env.ANTHROPIC_AUTH_TOKEN = accessToken; - const anthropicHeaders: Record< string, string > = { - 'User-Agent': getStudioUserAgent(), - 'X-WPCOM-AI-Feature': WPCOM_AI_FEATURE_HEADER_ANTHROPIC, - }; - if ( options?.sessionId ) { - anthropicHeaders[ 'X-WPCOM-Session-ID' ] = options.sessionId; - } - env.ANTHROPIC_CUSTOM_HEADERS = buildAnthropicCustomHeaders( anthropicHeaders ); - - // OpenAI Responses path. The wpcom proxy accepts the same bearer token and - // dispatches to the right upstream based on the request path. - // The OpenAI SDK expects baseURL to include /v1 (like the real - // OpenAI API), so the request path becomes /v1/responses — - // mirroring the Anthropic path's /v1/messages. - env.OPENAI_BASE_URL = `${ gatewayBaseUrl.replace( /\/+$/, '' ) }/v1`; - env.OPENAI_API_KEY = accessToken; - const openaiHeaders: Record< string, string > = { - 'User-Agent': getStudioUserAgent(), - 'X-WPCOM-AI-Feature': WPCOM_AI_FEATURE_HEADER_OPENAI, - }; - if ( options?.sessionId ) { - openaiHeaders[ 'X-WPCOM-Session-ID' ] = options.sessionId; - } - env.STUDIO_OPENAI_DEFAULT_HEADERS = JSON.stringify( openaiHeaders ); - - // Hosted third-party models (Kimi, GLM, DeepSeek) speak the OpenAI - // Chat Completions dialect, so they share the /v1 prefix with the - // OpenAI path and are told apart by the feature header. The vars are - // Studio-namespaced because, unlike Anthropic and OpenAI, this family - // has no direct-API provider — nothing but this function should be - // able to satisfy it. - env.STUDIO_HOSTED_BASE_URL = env.OPENAI_BASE_URL; - env.STUDIO_HOSTED_API_KEY = accessToken; - const hostedHeaders: Record< string, string > = { + // The studio capability tiers speak the OpenAI Chat Completions + // dialect, so the base URL carries the /v1 prefix (the request path + // becomes /v1/chat/completions). The vars are Studio-namespaced + // because this family has no direct-API provider — nothing but this + // function should be able to satisfy it. + env.STUDIO_WPCOM_BASE_URL = `${ gatewayBaseUrl.replace( /\/+$/, '' ) }/v1`; + env.STUDIO_WPCOM_API_KEY = accessToken; + const headers: Record< string, string > = { 'User-Agent': getStudioUserAgent(), - 'X-WPCOM-AI-Feature': WPCOM_AI_FEATURE_HEADER_HOSTED, + 'X-WPCOM-AI-Feature': WPCOM_AI_FEATURE_HEADER, }; if ( options?.sessionId ) { - hostedHeaders[ 'X-WPCOM-Session-ID' ] = options.sessionId; + headers[ 'X-WPCOM-Session-ID' ] = options.sessionId; } - env.STUDIO_HOSTED_DEFAULT_HEADERS = JSON.stringify( hostedHeaders ); + env.STUDIO_WPCOM_DEFAULT_HEADERS = JSON.stringify( headers ); return env; }, diff --git a/apps/cli/ai/runtimes/pi/index.ts b/apps/cli/ai/runtimes/pi/index.ts index 1e0cffadce..12ead74dba 100644 --- a/apps/cli/ai/runtimes/pi/index.ts +++ b/apps/cli/ai/runtimes/pi/index.ts @@ -1,5 +1,4 @@ import fs from 'fs'; -import Anthropic from '@anthropic-ai/sdk'; import { type AgentTool } from '@earendil-works/pi-agent-core'; import { type Credential, @@ -8,10 +7,6 @@ import { type Model, type SimpleStreamOptions, } from '@earendil-works/pi-ai'; -import { - stream as streamAnthropic, - type AnthropicOptions, -} from '@earendil-works/pi-ai/api/anthropic-messages'; import { streamSimple as streamOpenAiCompletions } from '@earendil-works/pi-ai/api/openai-completions'; import { streamSimple as streamOpenAiResponses } from '@earendil-works/pi-ai/api/openai-responses'; import { ANTHROPIC_MODELS } from '@earendil-works/pi-ai/providers/anthropic.models'; @@ -72,13 +67,11 @@ import type { AskUserHandler, SiteInfo } from 'cli/ai/types'; // eslint-disable-next-line @typescript-eslint/no-explicit-any type AgentToolAny = AgentTool< any >; -type StudioOpenAiCompatibleModel = Model< 'openai-responses' > | Model< 'openai-completions' >; -type StudioModel = StudioOpenAiCompatibleModel | Model< 'anthropic-messages' >; +type StudioWpcomModel = Model< 'openai-completions' > | Model< 'openai-responses' >; +type StudioModel = StudioWpcomModel | Model< 'anthropic-messages' >; type ProviderConfigInput = Parameters< ModelRuntime[ 'registerProvider' ] >[ 1 ]; -const STUDIO_WPCOM_ANTHROPIC_PROVIDER = 'studio-wpcom-anthropic'; -const STUDIO_WPCOM_OPENAI_PROVIDER = 'studio-wpcom-openai'; -const STUDIO_WPCOM_HOSTED_PROVIDER = 'studio-wpcom-hosted'; +const STUDIO_WPCOM_PROVIDER = 'studio-wpcom'; const STUDIO_AGENT_DIR = STUDIO_SITES_ROOT; const STUDIO_WPCOM_BODY_FILES_ROOT = getConfigDirectory(); const STUDIO_WPCOM_BODY_FILES_DIR = getAiPayloadsPath(); @@ -143,73 +136,51 @@ interface ResolvedCredentials { apiKey: string; baseURL: string; extraHeaders?: Record< string, string >; - useBearerAuth: boolean; } -// Families that speak an OpenAI dialect each get their own env namespace so a -// single resolved environment can carry all of them and the user can swap -// models mid-session. -const OPENAI_DIALECT_ENV_VARS = { - openai: { - label: 'OpenAI', - apiKey: 'OPENAI_API_KEY', - baseUrl: 'OPENAI_BASE_URL', - headers: 'STUDIO_OPENAI_DEFAULT_HEADERS', - }, - hosted: { - label: 'Hosted', - apiKey: 'STUDIO_HOSTED_API_KEY', - baseUrl: 'STUDIO_HOSTED_BASE_URL', - headers: 'STUDIO_HOSTED_DEFAULT_HEADERS', - }, -} as const; - function resolveCredentials( family: AiModelFamily, env: Record< string, string > ): { ok: true; creds: ResolvedCredentials } | { ok: false; reason: string } { - if ( family === 'openai' || family === 'hosted' ) { - const vars = OPENAI_DIALECT_ENV_VARS[ family ]; - const apiKey = env[ vars.apiKey ]?.trim(); + if ( family === 'studio' ) { + const apiKey = env.STUDIO_WPCOM_API_KEY?.trim(); if ( ! apiKey ) { return { ok: false, - reason: `${ vars.label } models are only available through the WordPress.com provider, and ${ vars.apiKey } is not set — run /login to authenticate.`, + reason: + 'The WordPress.com models need a wpcom access token, and STUDIO_WPCOM_API_KEY is not set — run /login to authenticate.', }; } - const baseURL = env[ vars.baseUrl ]?.trim(); + const baseURL = env.STUDIO_WPCOM_BASE_URL?.trim(); if ( ! baseURL ) { - return { ok: false, reason: `${ vars.baseUrl } not set — cannot route to wpcom proxy.` }; + return { ok: false, reason: 'STUDIO_WPCOM_BASE_URL not set — cannot route to wpcom proxy.' }; } return { ok: true, creds: { apiKey, baseURL, - extraHeaders: parseJsonHeaderEnv( vars.headers, env[ vars.headers ] ), - useBearerAuth: false, + extraHeaders: parseJsonHeaderEnv( + 'STUDIO_WPCOM_DEFAULT_HEADERS', + env.STUDIO_WPCOM_DEFAULT_HEADERS + ), }, }; } - const authToken = env.ANTHROPIC_AUTH_TOKEN?.trim(); const apiKey = env.ANTHROPIC_API_KEY?.trim(); - const credential = authToken ?? apiKey; - if ( ! credential ) { + if ( ! apiKey ) { return { ok: false, reason: - 'Anthropic provider selected but neither ANTHROPIC_AUTH_TOKEN nor ANTHROPIC_API_KEY is set. On the WordPress.com provider this means the wpcom access token is missing — run /login to authenticate. Otherwise switch to the Anthropic · API key provider with /provider and save a key.', + 'Anthropic provider selected but ANTHROPIC_API_KEY is not set. Switch to the Anthropic · API key provider with /provider and save a key.', }; } - const baseURL = env.ANTHROPIC_BASE_URL?.trim() || 'https://api.anthropic.com'; return { ok: true, creds: { - apiKey: credential, - baseURL, - extraHeaders: parseAnthropicHeaderEnv( env.ANTHROPIC_CUSTOM_HEADERS ), - useBearerAuth: Boolean( authToken ), + apiKey, + baseURL: 'https://api.anthropic.com', }, }; } @@ -410,19 +381,36 @@ function buildModel( ...( creds.extraHeaders ? { headers: creds.extraHeaders } : {} ), }; - if ( family === 'hosted' ) { - // pi infers `compat` from the base URL, and the wpcom proxy URL reads as - // plain OpenAI — so spell out the shape or requests carry OpenAI-only - // fields these upstreams reject. With `supportsReasoningEffort` false and - // the default `thinkingFormat`, no thinking switch is sent at all and - // each model uses its own default — the portable choice across vendors - // that spell that parameter differently. Reasoning still streams back. + if ( family === 'studio' ) { + // The capability tiers are resolved to upstream models by the wpcom + // proxy. The context window is a conservative floor across the + // upstreams a tier may resolve to, so compaction kicks in before any + // of them overflows. + // + // `strong` rides the Responses path: its upstream is a reasoning model + // that rejects tools-plus-reasoning on Chat Completions. + if ( modelId === 'strong' ) { + return { + ...common, + api: 'openai-responses', + provider: STUDIO_WPCOM_PROVIDER, + reasoning: true, + contextWindow: 200_000, + maxTokens: 32_000, + }; + } + // The other tiers speak Chat Completions. pi infers `compat` from the + // base URL, which for us reads as plain OpenAI — so spell out the shape + // or requests carry OpenAI-only fields other upstreams reject. With + // `supportsReasoningEffort` false no thinking switch is sent at all and + // each upstream uses its own default — the portable choice across + // vendors that spell that parameter differently. return { ...common, api: 'openai-completions', - provider: STUDIO_WPCOM_HOSTED_PROVIDER, + provider: STUDIO_WPCOM_PROVIDER, reasoning: true, - contextWindow: 262_144, + contextWindow: 200_000, maxTokens: 32_000, compat: { supportsStore: false, @@ -433,28 +421,6 @@ function buildModel( }, }; } - - if ( family === 'openai' ) { - // GPT-5.6 models reject function tools on /v1/chat/completions unless - // reasoning is disabled; the Responses API supports tools + reasoning. - // GPT-5.6 Sol's real context window is 1.05M tokens, but we declare - // 272K — the threshold where OpenAI's 2x long-context pricing kicks - // in — so compaction keeps sessions below it. Understating the window - // is also load-bearing for correctness: pi clamps max output tokens to - // the declared window minus its (post-compaction, sometimes stale) - // context estimate, and a too-small window can clamp all the way down - // to 1, which the API rejects with a 400. - // The openai family always rides the wpcom proxy (Studio has no - // direct-OpenAI provider), so it always uses the custom provider. - return { - ...common, - api: 'openai-responses', - provider: STUDIO_WPCOM_OPENAI_PROVIDER, - reasoning: true, - contextWindow: 272_000, - maxTokens: 32_000, - }; - } // Without `compat.forceAdaptiveThinking` pi-ai sends the legacy // `thinking: { type: 'enabled', budget_tokens }` shape, which Sonnet 5 / // Opus 5 reject with a 400 — copy the thinking fields from pi's catalog. @@ -464,7 +430,7 @@ function buildModel( return { ...common, api: 'anthropic-messages', - provider: creds.useBearerAuth ? STUDIO_WPCOM_ANTHROPIC_PROVIDER : 'anthropic', + provider: 'anthropic', reasoning: true, // contextWindow/maxTokens intentionally stay below the catalog values. contextWindow: 200_000, @@ -523,20 +489,10 @@ async function createModelRuntime( modelsPath: null, } ); - if ( family === 'anthropic' && creds.useBearerAuth ) { - modelRuntime.registerProvider( - STUDIO_WPCOM_ANTHROPIC_PROVIDER, - createWpcomAnthropicProviderConfig( model as Model< 'anthropic-messages' >, creds ) - ); - return modelRuntime; - } - - if ( family === 'openai' || family === 'hosted' ) { - // `buildModel` already resolved the provider and wire API for this - // family; read them off the model rather than deriving them again. + if ( family === 'studio' ) { modelRuntime.registerProvider( model.provider, - createWpcomOpenAiCompatibleProviderConfig( model as StudioOpenAiCompatibleModel, creds ) + createWpcomProviderConfig( model as StudioWpcomModel, creds ) ); return modelRuntime; } @@ -557,54 +513,12 @@ function escapePiConfigValue( value: string ): string { return dollarEscaped.startsWith( '!' ) ? `$${ dollarEscaped }` : dollarEscaped; } -function createWpcomAnthropicProviderConfig( - model: Model< 'anthropic-messages' >, - creds: ResolvedCredentials -): ProviderConfigInput { - return { - baseUrl: creds.baseURL, - apiKey: escapePiConfigValue( creds.apiKey ), - api: 'anthropic-messages', - headers: creds.extraHeaders, - streamSimple: ( m, ctx, options?: SimpleStreamOptions ) => { - const client = new Anthropic( { - apiKey: null, - authToken: options?.apiKey ?? creds.apiKey, - baseURL: m.baseUrl, - dangerouslyAllowBrowser: true, - defaultHeaders: options?.headers, - } ); - const clientForPi = client as unknown as AnthropicOptions[ 'client' ]; - return withUsageCapErrorRewrite( - streamAnthropic( m as Model< 'anthropic-messages' >, stripStaleImagesFromContext( ctx ), { - ...( options as AnthropicOptions | undefined ), - client: clientForPi, - } ) - ); - }, - models: [ - { - id: model.id, - name: model.name, - api: 'anthropic-messages', - baseUrl: model.baseUrl, - reasoning: model.reasoning, - input: model.input, - cost: model.cost, - contextWindow: model.contextWindow, - maxTokens: model.maxTokens, - headers: creds.extraHeaders, - compat: model.compat, - thinkingLevelMap: model.thinkingLevelMap, - }, - ], - }; -} - -// The wpcom OpenAI-dialect paths only need pi's stock streaming for their API; -// the custom provider exists to wrap the stream with the usage-cap 429 rewrite. -function createWpcomOpenAiCompatibleProviderConfig( - model: StudioOpenAiCompatibleModel, +// The wpcom lane only needs pi's stock streaming for each tier's API; the +// custom provider exists to wrap the stream with the usage-cap 429 rewrite +// and to strip stale screenshots, which would otherwise bloat requests past +// the proxy's body limit. +function createWpcomProviderConfig( + model: StudioWpcomModel, creds: ResolvedCredentials ): ProviderConfigInput { // pi types `streamSimple` against `Model`; each API's stream function @@ -618,7 +532,7 @@ function createWpcomOpenAiCompatibleProviderConfig( api: model.api, headers: creds.extraHeaders, streamSimple: ( m, ctx, options?: SimpleStreamOptions ) => - withUsageCapErrorRewrite( stream( m, ctx, options ) ), + withUsageCapErrorRewrite( stream( m, stripStaleImagesFromContext( ctx ), options ) ), models: [ { id: model.id, @@ -763,18 +677,3 @@ function parseJsonHeaderEnv( } return undefined; } - -function parseAnthropicHeaderEnv( - value: string | undefined -): Record< string, string > | undefined { - if ( ! value ) return undefined; - const out: Record< string, string > = {}; - for ( const line of value.split( '\n' ) ) { - const idx = line.indexOf( ':' ); - if ( idx <= 0 ) continue; - const name = line.slice( 0, idx ).trim(); - const v = line.slice( idx + 1 ).trim(); - if ( name && v ) out[ name ] = v; - } - return Object.keys( out ).length ? out : undefined; -} diff --git a/apps/cli/ai/sessions/context.ts b/apps/cli/ai/sessions/context.ts index c46bab79b1..d169b81026 100644 --- a/apps/cli/ai/sessions/context.ts +++ b/apps/cli/ai/sessions/context.ts @@ -1,4 +1,4 @@ -import { resolveSessionModel, type AiModelId } from '@studio/common/ai/models'; +import { readRecordedSessionModel, type AiModelId } from '@studio/common/ai/models'; import { isAiProviderId } from '@studio/common/ai/providers'; import { isStudioCustomEntryOfType } from '@studio/common/ai/sessions/entry-types'; import type { LoadedAiSession } from '@studio/common/ai/sessions/types'; @@ -24,10 +24,9 @@ export function resolveResumeSessionContext( context.sessionId = resumeSession.summary.id; } - // Shared resolution: the most recent recorded model wins, and a removed - // model auto-switches to the default so resumed sessions never pin a - // model we no longer offer. - context.model = resolveSessionModel( resumeSession.entries ); + // Unset when the session recorded no (still-offered) model, so the caller + // applies its provider-appropriate default instead of pinning a dead id. + context.model = readRecordedSessionModel( resumeSession.entries ); for ( let index = resumeSession.entries.length - 1; index >= 0; index -= 1 ) { const entry = resumeSession.entries[ index ]; diff --git a/apps/cli/ai/slash-commands.ts b/apps/cli/ai/slash-commands.ts index 823f626bfd..44dd56fecc 100644 --- a/apps/cli/ai/slash-commands.ts +++ b/apps/cli/ai/slash-commands.ts @@ -1,7 +1,15 @@ -import { getAiModelFamily, getVisibleAiModels } from '@studio/common/ai/models'; -import { getAiModelLabel, type AiModelId } from '@studio/common/ai/models'; +import { + aiModelRequiresPaidCredits, + getAiModelFamily, + getAiModelLabel, + type AiModelId, +} from '@studio/common/ai/models'; import { getAiSkillCommands } from '@studio/common/ai/slash-commands'; import { isAutomatticianFromToken, readAuthToken } from '@studio/common/lib/shared-config'; +import { + fetchStudioAssistantQuota, + hasPaidAiCredits, +} from '@studio/common/lib/studio-assistant-quota'; import { __, sprintf } from '@wordpress/i18n'; import { getAvailableAiProviders, isAiProviderReady } from 'cli/ai/auth'; import { AI_PROVIDERS, getAiProviderDefinition, type AiProviderId } from 'cli/ai/providers'; @@ -334,12 +342,27 @@ export const AI_CHAT_SLASH_COMMANDS: SlashCommandDef[] = [ description: __( 'Switch the AI model' ), handler: async ( _prompt, ctx ) => { const { availableModels } = getAiProviderDefinition( ctx.currentProvider ); - const visible = new Set( - getVisibleAiModels( await isAutomatticianFromToken(), ctx.currentModel ).map( - ( model ) => model.id - ) - ); - const offeredModels = availableModels.filter( ( id ) => visible.has( id ) ); + // The paid tiers are only offered while purchased credits remain; + // Automatticians are exempt. The current model always stays listed, + // so a session already on a paid tier keeps showing what it runs on. + let offeredModels = availableModels; + if ( + availableModels.some( aiModelRequiresPaidCredits ) && + ! ( await isAutomatticianFromToken() ) + ) { + const token = await readAuthToken(); + const quota = token ? await fetchStudioAssistantQuota( token.accessToken ) : null; + if ( ! hasPaidAiCredits( quota ) ) { + offeredModels = availableModels.filter( + ( id ) => id === ctx.currentModel || ! aiModelRequiresPaidCredits( id ) + ); + } + } + if ( offeredModels.length < availableModels.length ) { + ctx.ui.showInfo( + __( 'Models that need purchased AI credits are hidden — add credits to unlock them.' ) + ); + } // Build options and a reverse lookup at the same time so we never // have to recover the model id from the label. A startsWith-based // match is buggy when one model's label is a prefix of another's @@ -418,9 +441,20 @@ export const AI_CHAT_SLASH_COMMANDS: SlashCommandDef[] = [ selectedLabel.startsWith( AI_PROVIDERS[ id ] ) ); if ( newProvider && newProvider !== ctx.currentProvider ) { + const previousFamily = getAiModelFamily( ctx.currentModel ); try { await ctx.prepareProviderSelection( newProvider ); await ctx.switchProvider( newProvider ); + // Providers don't share a model family, so a switch is the + // same runtime handoff as a cross-family /model switch. + if ( getAiModelFamily( ctx.currentModel ) !== previousFamily ) { + await ctx.clearSession(); + ctx.ui.showInfo( + __( + "Switching across model families starts a fresh conversation — the prior turns aren't carried over." + ) + ); + } } catch ( error ) { if ( isPromptAbortError( error ) ) { ctx.ui.showInfo( diff --git a/apps/cli/ai/tests/auth.test.ts b/apps/cli/ai/tests/auth.test.ts index 1a3c78400b..03a0ae5341 100644 --- a/apps/cli/ai/tests/auth.test.ts +++ b/apps/cli/ai/tests/auth.test.ts @@ -98,13 +98,14 @@ describe( 'AI auth helpers', () => { const env = await resolveAiEnvironment( 'wpcom' ); - expect( env.ANTHROPIC_BASE_URL ).toBe( - 'https://public-api.wordpress.com/wpcom/v2/ai-api-proxy' - ); - expect( env.ANTHROPIC_AUTH_TOKEN ).toBe( 'wpcom-token' ); - expect( env.ANTHROPIC_CUSTOM_HEADERS ).toBe( - 'User-Agent: WordPressStudio/1.2.3\nX-WPCOM-AI-Feature: studio-assistant-anthropic' + expect( env.STUDIO_WPCOM_BASE_URL ).toBe( + 'https://public-api.wordpress.com/wpcom/v2/ai-api-proxy/v1' ); + expect( env.STUDIO_WPCOM_API_KEY ).toBe( 'wpcom-token' ); + expect( JSON.parse( env.STUDIO_WPCOM_DEFAULT_HEADERS! ) ).toEqual( { + 'User-Agent': 'WordPressStudio/1.2.3', + 'X-WPCOM-AI-Feature': 'studio-agent', + } ); expect( env.ANTHROPIC_API_KEY ).toBeUndefined(); } ); @@ -120,9 +121,11 @@ describe( 'AI auth helpers', () => { const env = await resolveAiEnvironment( 'wpcom', { sessionId: 'session-abc' } ); - expect( env.ANTHROPIC_CUSTOM_HEADERS ).toBe( - 'User-Agent: WordPressStudio/1.2.3\nX-WPCOM-AI-Feature: studio-assistant-anthropic\nX-WPCOM-Session-ID: session-abc' - ); + expect( JSON.parse( env.STUDIO_WPCOM_DEFAULT_HEADERS! ) ).toEqual( { + 'User-Agent': 'WordPressStudio/1.2.3', + 'X-WPCOM-AI-Feature': 'studio-agent', + 'X-WPCOM-Session-ID': 'session-abc', + } ); } ); it( 'prefers the saved provider', async () => { diff --git a/apps/cli/ai/tests/pi-runtime.test.ts b/apps/cli/ai/tests/pi-runtime.test.ts index e30f04a59b..4bdee7f503 100644 --- a/apps/cli/ai/tests/pi-runtime.test.ts +++ b/apps/cli/ai/tests/pi-runtime.test.ts @@ -15,13 +15,13 @@ const mocks = vi.hoisted( () => ( { } ) ); // Model-swap test uses a synthetic id outside `AI_MODELS`; route unknowns to -// 'openai' so the env credentials match. +// 'studio' so the env credentials match. vi.mock( '@studio/common/ai/models', async ( importOriginal ) => { const actual = await importOriginal< typeof import('@studio/common/ai/models') >(); return { ...actual, getAiModelFamily: ( id: string ) => - actual.isAiModelId( id ) ? actual.getAiModelFamily( id ) : 'openai', + actual.isAiModelId( id ) ? actual.getAiModelFamily( id ) : 'studio', }; } ); @@ -66,10 +66,10 @@ const DEFAULT_MOCK_EVENTS: AgentSessionEvent[] = [ type: 'message_end', message: { role: 'assistant', - content: [ { type: 'text', text: 'mocked openai response' } ], - api: 'openai-responses', - provider: 'openai', - model: 'gpt-5.6-sol', + content: [ { type: 'text', text: 'mocked wpcom response' } ], + api: 'openai-completions', + provider: 'studio-wpcom', + model: 'balanced', usage: { input: 0, output: 0, @@ -86,10 +86,10 @@ const DEFAULT_MOCK_EVENTS: AgentSessionEvent[] = [ type: 'turn_end', message: { role: 'assistant', - content: [ { type: 'text', text: 'mocked openai response' } ], - api: 'openai-responses', - provider: 'openai', - model: 'gpt-5.6-sol', + content: [ { type: 'text', text: 'mocked wpcom response' } ], + api: 'openai-completions', + provider: 'studio-wpcom', + model: 'balanced', usage: { input: 0, output: 0, @@ -117,9 +117,9 @@ const assistantMessage = ( message: { role: 'assistant', content, - api: 'openai-responses', - provider: 'openai', - model: 'gpt-5.6-sol', + api: 'openai-completions', + provider: 'studio-wpcom', + model: 'balanced', usage: { input: 0, output: 0, @@ -210,6 +210,11 @@ async function runRuntime( return events; } +const WPCOM_ENV = { + STUDIO_WPCOM_API_KEY: 'wpcom-token', + STUDIO_WPCOM_BASE_URL: 'https://proxy.example.com/v1', +}; + describe( 'pi runtime', () => { beforeEach( () => { mocks.createdSessions.length = 0; @@ -222,11 +227,11 @@ describe( 'pi runtime', () => { } ); } ); - it( 'emits agent_end carrying the credential error when OPENAI_API_KEY is absent', async () => { + it( 'emits agent_end carrying the credential error when STUDIO_WPCOM_API_KEY is absent', async () => { const events = await runRuntime( { prompt: 'hello', env: {}, - model: 'gpt-5.6-sol', + model: 'balanced', session: newSession(), } ); @@ -239,7 +244,7 @@ describe( 'pi runtime', () => { expect( last.role ).toBe( 'assistant' ); if ( last.role === 'assistant' ) { expect( last.stopReason ).toBe( 'error' ); - expect( last.errorMessage ).toMatch( /OPENAI_API_KEY/ ); + expect( last.errorMessage ).toMatch( /STUDIO_WPCOM_API_KEY/ ); } } } ); @@ -247,65 +252,53 @@ describe( 'pi runtime', () => { it( 'emits a full exchange when AgentSession returns output', async () => { const events = await runRuntime( { prompt: 'hello', - env: { - OPENAI_API_KEY: 'sk-test', - OPENAI_BASE_URL: 'https://proxy.example.com/v1', - }, - model: 'gpt-5.6-sol', + env: WPCOM_ENV, + model: 'balanced', session: newSession(), } ); - expect( findAssistantText( events ) ).toBe( 'mocked openai response' ); + expect( findAssistantText( events ) ).toBe( 'mocked wpcom response' ); const final = events[ events.length - 1 ]; expect( final.type ).toBe( 'agent_end' ); } ); - it( 'advertises image input support so screenshot tool results can be analyzed', async () => { - await runRuntime( { - prompt: 'hello', - env: { - OPENAI_API_KEY: 'sk-test', - OPENAI_BASE_URL: 'https://proxy.example.com/v1', - }, - model: 'gpt-5.6-sol', - session: newSession(), - } ); - - expect( mocks.createdSessions[ 0 ].options.model?.input ).toEqual( [ 'text', 'image' ] ); - } ); - - const HOSTED_ENV = { - STUDIO_HOSTED_API_KEY: 'wpcom-token', - STUDIO_HOSTED_BASE_URL: 'https://proxy.example.com/v1', - }; - - it( 'routes hosted models to the wpcom Chat Completions path', async () => { + it( 'routes the capability tiers to the wpcom Chat Completions path', async () => { await runRuntime( { prompt: 'hello', env: { - ...HOSTED_ENV, - STUDIO_HOSTED_DEFAULT_HEADERS: JSON.stringify( { - 'X-WPCOM-AI-Feature': 'studio-assistant-hosted', + ...WPCOM_ENV, + STUDIO_WPCOM_DEFAULT_HEADERS: JSON.stringify( { + 'X-WPCOM-AI-Feature': 'studio-agent', 'X-WPCOM-Session-ID': 'session-2', } ), }, - model: 'moonshotai/Kimi-K2.6', + model: 'balanced', session: newSession(), } ); const options = mocks.createdSessions[ 0 ].options; - expect( options.model?.id ).toBe( 'moonshotai/Kimi-K2.6' ); - expect( options.model?.provider ).toBe( 'studio-wpcom-hosted' ); + expect( options.model?.id ).toBe( 'balanced' ); + expect( options.model?.provider ).toBe( 'studio-wpcom' ); expect( options.model?.api ).toBe( 'openai-completions' ); + const modelRegistry = new ModelRegistry( options.modelRuntime! ); + const auth = await modelRegistry.getApiKeyAndHeaders( options.model! ); + expect( auth ).toMatchObject( { + ok: true, + apiKey: 'wpcom-token', + headers: { + 'X-WPCOM-AI-Feature': 'studio-agent', + 'X-WPCOM-Session-ID': 'session-2', + }, + } ); } ); - // Without these the request carries OpenAI-only fields these upstreams + // Without these the request carries OpenAI-only fields other upstreams // reject: pi infers them from the base URL, which for us reads as OpenAI. - it( 'declares hosted compat overrides the proxy URL cannot be detected from', async () => { + it( 'declares compat overrides the proxy URL cannot be detected from', async () => { await runRuntime( { prompt: 'hello', - env: HOSTED_ENV, - model: 'moonshotai/Kimi-K3', + env: WPCOM_ENV, + model: 'fast', session: newSession(), } ); @@ -318,17 +311,34 @@ describe( 'pi runtime', () => { } ); } ); + // `strong` resolves to a reasoning model that rejects the Chat Completions + // dialect's tools-plus-reasoning combination, so it rides the Responses + // path — the plain OpenAI dialect, with no compat overrides. + it( 'routes the strong tier to the wpcom Responses path', async () => { + await runRuntime( { + prompt: 'hello', + env: WPCOM_ENV, + model: 'strong', + session: newSession(), + } ); + + const model = mocks.createdSessions[ 0 ].options.model!; + expect( model.api ).toBe( 'openai-responses' ); + expect( model.provider ).toBe( 'studio-wpcom' ); + expect( model.compat ).toBeUndefined(); + } ); + it( 'advertises image input per model rather than per family', async () => { await runRuntime( { prompt: 'hello', - env: HOSTED_ENV, - model: 'moonshotai/Kimi-K2.6', + env: WPCOM_ENV, + model: 'balanced', session: newSession(), } ); await runRuntime( { prompt: 'hello', - env: HOSTED_ENV, - model: 'zai-org/GLM-5.2', + env: WPCOM_ENV, + model: 'fast', session: newSession(), } ); @@ -341,14 +351,14 @@ describe( 'pi runtime', () => { it( 'withholds take_screenshot from models that cannot see images', async () => { await runRuntime( { prompt: 'hello', - env: HOSTED_ENV, - model: 'moonshotai/Kimi-K2.6', + env: WPCOM_ENV, + model: 'balanced', session: newSession(), } ); await runRuntime( { prompt: 'hello', - env: HOSTED_ENV, - model: 'zai-org/GLM-5.2', + env: WPCOM_ENV, + model: 'fast', session: newSession(), } ); @@ -363,11 +373,8 @@ describe( 'pi runtime', () => { it( 'rejects oversized direct Write, Edit, and Bash payloads', async () => { await runRuntime( { prompt: 'hello', - env: { - OPENAI_API_KEY: 'sk-test', - OPENAI_BASE_URL: 'https://proxy.example.com/v1', - }, - model: 'gpt-5.6-sol', + env: WPCOM_ENV, + model: 'balanced', session: newSession(), } ); @@ -420,11 +427,8 @@ describe( 'pi runtime', () => { await runRuntime( { prompt: 'hello', - env: { - OPENAI_API_KEY: 'sk-test', - OPENAI_BASE_URL: 'https://proxy.example.com/v1', - }, - model: 'gpt-5.6-sol', + env: WPCOM_ENV, + model: 'balanced', session: newSession(), activeSite: { name: 'Remote', @@ -451,11 +455,8 @@ describe( 'pi runtime', () => { it( 'leaves retry policy to pi settings defaults', async () => { await runRuntime( { prompt: 'hello', - env: { - OPENAI_API_KEY: 'sk-test', - OPENAI_BASE_URL: 'https://proxy.example.com/v1', - }, - model: 'gpt-5.6-sol', + env: WPCOM_ENV, + model: 'balanced', session: newSession(), } ); @@ -469,60 +470,25 @@ describe( 'pi runtime', () => { } ); it( 'creates each AgentSession with the requested model', async () => { - const env = { - OPENAI_API_KEY: 'sk-test', - OPENAI_BASE_URL: 'https://proxy.example.com/v1', - }; const session = newSession(); - const otherOpenAiModel = 'gpt-test-other' as AiModelId; + const otherStudioModel = 'tier-test-other' as AiModelId; - await runRuntime( { prompt: 'hi', env, model: 'gpt-5.6-sol', session } ); - await runRuntime( { prompt: 'follow-up', env, model: otherOpenAiModel, session } ); + await runRuntime( { prompt: 'hi', env: WPCOM_ENV, model: 'balanced', session } ); + await runRuntime( { prompt: 'follow-up', env: WPCOM_ENV, model: otherStudioModel, session } ); await runRuntime( { prompt: 'still on the second model', - env, - model: otherOpenAiModel, + env: WPCOM_ENV, + model: otherStudioModel, session, } ); expect( mocks.createdSessions.map( ( s ) => s.state.model.id ) ).toEqual( [ - 'gpt-5.6-sol', - otherOpenAiModel, - otherOpenAiModel, + 'balanced', + otherStudioModel, + otherStudioModel, ] ); } ); - it( 'registers WPCOM Anthropic as a custom bearer-auth provider', async () => { - await runRuntime( { - prompt: 'hello', - env: { - ANTHROPIC_AUTH_TOKEN: 'wpcom-token', - ANTHROPIC_BASE_URL: 'https://proxy.example.com', - ANTHROPIC_CUSTOM_HEADERS: - 'X-WPCOM-AI-Feature: studio-assistant-anthropic\nX-WPCOM-Session-ID: session-1', - }, - model: 'claude-sonnet-5', - session: newSession(), - } ); - - const options = mocks.createdSessions[ 0 ].options; - expect( options.model?.provider ).toBe( 'studio-wpcom-anthropic' ); - expect( options.model?.api ).toBe( 'anthropic-messages' ); - expect( options.model?.maxTokens ).toBe( 32_000 ); - expect( options.model?.input ).toEqual( [ 'text', 'image' ] ); - expect( options.model?.compat ).toMatchObject( { forceAdaptiveThinking: true } ); - const modelRegistry = new ModelRegistry( options.modelRuntime! ); - const auth = await modelRegistry.getApiKeyAndHeaders( options.model! ); - expect( auth ).toMatchObject( { - ok: true, - apiKey: 'wpcom-token', - headers: { - 'X-WPCOM-AI-Feature': 'studio-assistant-anthropic', - 'X-WPCOM-Session-ID': 'session-1', - }, - } ); - } ); - // Without `compat.forceAdaptiveThinking`, pi-ai sends a thinking shape that // Sonnet 5 / Opus 5 reject with a 400. it( 'marks direct-key Anthropic models as adaptive-thinking', async () => { @@ -581,11 +547,10 @@ describe( 'pi runtime', () => { await runRuntime( { prompt: 'hello', env: { - ANTHROPIC_AUTH_TOKEN: tokenWithDollar, - ANTHROPIC_BASE_URL: 'https://proxy.example.com', - ANTHROPIC_CUSTOM_HEADERS: 'X-WPCOM-AI-Feature: studio-assistant-anthropic', + STUDIO_WPCOM_API_KEY: tokenWithDollar, + STUDIO_WPCOM_BASE_URL: 'https://proxy.example.com/v1', }, - model: 'claude-sonnet-5', + model: 'balanced', session: newSession(), } ); @@ -597,24 +562,23 @@ describe( 'pi runtime', () => { } ); // Silent header-drop would surface as an opaque 401 from the wpcom proxy. - it( 'warns and continues when STUDIO_OPENAI_DEFAULT_HEADERS is malformed', async () => { + it( 'warns and continues when STUDIO_WPCOM_DEFAULT_HEADERS is malformed', async () => { const warnSpy = vi.spyOn( console, 'warn' ).mockImplementation( () => {} ); try { await runRuntime( { prompt: 'hello', env: { - OPENAI_API_KEY: 'sk-test', - OPENAI_BASE_URL: 'https://proxy.example.com/v1', - STUDIO_OPENAI_DEFAULT_HEADERS: '{not json', + ...WPCOM_ENV, + STUDIO_WPCOM_DEFAULT_HEADERS: '{not json', }, - model: 'gpt-5.6-sol', + model: 'balanced', session: newSession(), } ); expect( warnSpy ).toHaveBeenCalledTimes( 1 ); expect( warnSpy.mock.calls[ 0 ][ 0 ] ).toMatch( - /STUDIO_OPENAI_DEFAULT_HEADERS.*malformed JSON/ + /STUDIO_WPCOM_DEFAULT_HEADERS.*malformed JSON/ ); } finally { warnSpy.mockRestore(); diff --git a/apps/cli/ai/tests/slash-commands.test.ts b/apps/cli/ai/tests/slash-commands.test.ts index b4e90c2ea1..940e77fe60 100644 --- a/apps/cli/ai/tests/slash-commands.test.ts +++ b/apps/cli/ai/tests/slash-commands.test.ts @@ -1,3 +1,5 @@ +import { isAutomatticianFromToken, readAuthToken } from '@studio/common/lib/shared-config'; +import { fetchStudioAssistantQuota } from '@studio/common/lib/studio-assistant-quota'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { AI_CHAT_SLASH_COMMANDS, type SlashCommandContext } from 'cli/ai/slash-commands'; @@ -41,6 +43,10 @@ vi.mock( '@studio/common/lib/shared-config', () => ( { readAuthToken: vi.fn(), isAutomatticianFromToken: vi.fn().mockResolvedValue( true ), } ) ); +vi.mock( '@studio/common/lib/studio-assistant-quota', async ( importOriginal ) => ( { + ...( await importOriginal< typeof import('@studio/common/lib/studio-assistant-quota') >() ), + fetchStudioAssistantQuota: vi.fn(), +} ) ); vi.mock( 'cli/remote-session/daemon', () => { return { @@ -354,12 +360,12 @@ function buildModelCtx( // `as never` keeps this test framework-agnostic — we don't need a // real AiChatUI, just the methods the /model handler touches. ui: { - currentModel: overrides.currentModel ?? 'gpt-5.6-sol', + currentModel: overrides.currentModel ?? 'fast', askUser: vi.fn().mockResolvedValue( { 0: overrides.askUserResponse } ), showInfo: vi.fn(), showError: vi.fn(), } as never, - currentModel: overrides.currentModel ?? 'gpt-5.6-sol', + currentModel: overrides.currentModel ?? 'fast', currentProvider: 'wpcom', showCapabilitiesOnConnect: false, switchProvider: vi.fn().mockResolvedValue( undefined ), @@ -372,6 +378,20 @@ function buildModelCtx( } describe( '/model slash command', () => { + const setAccount = ( { + hasPaid, + automattician = false, + }: { + hasPaid: boolean; + automattician?: boolean; + } ) => { + vi.mocked( isAutomatticianFromToken ).mockResolvedValue( automattician ); + vi.mocked( readAuthToken ).mockResolvedValue( { accessToken: 'wpcom-token' } as never ); + vi.mocked( fetchStudioAssistantQuota ).mockResolvedValue( { + purchasedRemaining: hasPaid ? 100_000 : 0, + } as never ); + }; + // Locks in the labelToId-map matcher introduced in slash-commands.ts. // The previous implementation used `selectedLabel.startsWith( label )`, // which silently picked the wrong id whenever one model's label was a @@ -380,27 +400,81 @@ describe( '/model slash command', () => { // and we want the regression coverage to survive future additions. it( 'resolves the picked model exactly by label, not by prefix', async () => { expect( modelHandler ).toBeDefined(); + setAccount( { hasPaid: true } ); const { ctx, persistMock } = buildModelCtx( { - currentModel: 'gpt-5.6-sol', - askUserResponse: 'Sonnet 5', + currentModel: 'fast', + askUserResponse: 'Balanced', } ); await modelHandler!( '/model', ctx ); - expect( ctx.currentModel ).toBe( 'claude-sonnet-5' ); + expect( ctx.currentModel ).toBe( 'balanced' ); expect( persistMock ).toHaveBeenCalledTimes( 1 ); } ); + it( 'withholds the paid tiers when no purchased credits remain', async () => { + setAccount( { hasPaid: false } ); + const { ctx } = buildModelCtx( { + currentModel: 'fast', + askUserResponse: 'Balanced', + } ); + + await modelHandler!( '/model', ctx ); + + // Balanced was never offered, so the pick resolves to nothing. + expect( ctx.currentModel ).toBe( 'fast' ); + const [ questions ] = vi.mocked( ctx.ui.askUser ).mock.calls[ 0 ]; + expect( questions[ 0 ].options.map( ( option ) => option.description ) ).toEqual( [ 'fast' ] ); + expect( ctx.ui.showInfo ).toHaveBeenCalledWith( + expect.stringContaining( 'purchased AI credits' ) + ); + } ); + + it( 'offers the paid tiers to Automatticians regardless of credits', async () => { + setAccount( { hasPaid: false, automattician: true } ); + const { ctx } = buildModelCtx( { + currentModel: 'fast', + askUserResponse: 'Strong', + } ); + + await modelHandler!( '/model', ctx ); + + const [ questions ] = vi.mocked( ctx.ui.askUser ).mock.calls[ 0 ]; + expect( questions[ 0 ].options.map( ( option ) => option.description ) ).toEqual( [ + 'fast', + 'balanced', + 'strong', + ] ); + expect( ctx.currentModel ).toBe( 'strong' ); + } ); + + it( 'keeps offering the current model even when it needs purchased credits', async () => { + setAccount( { hasPaid: false } ); + const { ctx } = buildModelCtx( { + currentModel: 'strong', + askUserResponse: 'Fast', + } ); + + await modelHandler!( '/model', ctx ); + + const [ questions ] = vi.mocked( ctx.ui.askUser ).mock.calls[ 0 ]; + expect( questions[ 0 ].options.map( ( option ) => option.description ) ).toEqual( [ + 'fast', + 'strong', + ] ); + expect( ctx.currentModel ).toBe( 'fast' ); + } ); + it( 'still resolves the picked model when its label carries the "(current)" suffix', async () => { const { ctx, persistMock } = buildModelCtx( { - currentModel: 'gpt-5.6-sol', - askUserResponse: 'GPT 5.6 Sol (current)', + currentModel: 'fast', + askUserResponse: 'Fast (current)', } ); await modelHandler!( '/model', ctx ); // Same model picked → no swap, no persist. - expect( ctx.currentModel ).toBe( 'gpt-5.6-sol' ); + expect( ctx.currentModel ).toBe( 'fast' ); expect( persistMock ).not.toHaveBeenCalled(); } ); } ); diff --git a/apps/cli/commands/ai/index.ts b/apps/cli/commands/ai/index.ts index d7e3752277..8883681276 100644 --- a/apps/cli/commands/ai/index.ts +++ b/apps/cli/commands/ai/index.ts @@ -6,11 +6,11 @@ import { import { type StudioChatImage } from '@studio/common/ai/chat-images'; import { getAgentEndFailure } from '@studio/common/ai/json-events'; import { - DEFAULT_MODEL, getAiModelFamily, - resolveSessionModel, + readRecordedSessionModel, type AiModelId, } from '@studio/common/ai/models'; +import { getAiProviderDefaultModel } from '@studio/common/ai/providers'; import { getAgentEndTurnResult } from '@studio/common/ai/session-events'; import { readAnthropicApiKey, readSelectedAiProvider } from '@studio/common/ai/settings-store'; import { @@ -19,6 +19,10 @@ import { } from '@studio/common/ai/slash-commands'; import { getAiTracksIdentity } from '@studio/common/ai/tracks-identity'; import { readAuthToken } from '@studio/common/lib/shared-config'; +import { + fetchStudioAssistantQuota, + hasPaidAiCredits, +} from '@studio/common/lib/studio-assistant-quota'; import { getSessionsDirectory } from '@studio/common/lib/well-known-paths'; import { __, sprintf } from '@wordpress/i18n'; import { @@ -117,6 +121,25 @@ function getErrorMessage( error: unknown ): string { return String( error ); } +// Caps the quota lookup behind the wpcom default model, so a hung endpoint +// can't block the first turn — the free-tier default is the safe floor. +const QUOTA_FETCH_TIMEOUT_MS = 3_000; + +async function resolveWpcomDefaultModel(): Promise< AiModelId > { + const token = await readAuthToken(); + const quota = token + ? await Promise.race( [ + fetchStudioAssistantQuota( token.accessToken ), + new Promise< null >( ( resolve ) => { + setTimeout( () => resolve( null ), QUOTA_FETCH_TIMEOUT_MS ).unref(); + } ), + ] ) + : null; + return getAiProviderDefaultModel( DEFAULT_AI_PROVIDER, { + hasPaidAiCredits: hasPaidAiCredits( quota ), + } ); +} + async function readAllStdin(): Promise< string > { const chunks: Buffer[] = []; for await ( const chunk of process.stdin ) { @@ -157,9 +180,32 @@ export async function runCommand( options: { ) { currentProvider = DEFAULT_AI_PROVIDER; } - let currentModel: AiModelId = resumeContext.model ?? DEFAULT_MODEL; + // The recorded model only sticks when the provider still serves it — old + // wpcom sessions snap to the provider default instead. + const initialDefinition = getAiProviderDefinition( currentProvider ); + const recordedModel = + resumeContext.model && initialDefinition.supportsModel( resumeContext.model ) + ? resumeContext.model + : undefined; + let currentModel: AiModelId = recordedModel ?? initialDefinition.defaultModel; ui.currentProvider = currentProvider; ui.currentModel = currentModel; + + // The wpcom default is quota-dependent; resolved in the background so + // startup never waits on the network. Turns await the resolution, and it + // only applies while nothing else picked a model. + let wpcomDefaultModel: AiModelId = getAiProviderDefaultModel( DEFAULT_AI_PROVIDER ); + let quotaDefaultApplicable = ! recordedModel && currentProvider === DEFAULT_AI_PROVIDER; + const wpcomDefaultModelResolution = resolveWpcomDefaultModel() + .then( ( model ) => { + wpcomDefaultModel = model; + if ( quotaDefaultApplicable && currentProvider === DEFAULT_AI_PROVIDER ) { + currentModel = model; + ui.currentModel = model; + } + } ) + // Awaited by every turn — a failed lookup must not poison them. + .catch( () => {} ); if ( options.activeSite ) { ui.activeSite = { id: options.activeSite.id, @@ -197,8 +243,17 @@ export async function runCommand( options: { if ( sm.getSessionId() === options.resumeSessionId ) { session = sm; match = file; - currentModel = resolveSessionModel( sm.getEntries() ); - ui.currentModel = currentModel; + // Adopt the recorded model only when the provider still + // serves it; otherwise keep the (quota-based) default. + const sessionModel = readRecordedSessionModel( sm.getEntries() ); + if ( + sessionModel && + getAiProviderDefinition( currentProvider ).supportsModel( sessionModel ) + ) { + quotaDefaultApplicable = false; + currentModel = sessionModel; + ui.currentModel = currentModel; + } break; } } catch { @@ -327,11 +382,13 @@ export async function runCommand( options: { ui.currentProvider = currentProvider; // Auto-correct model when the provider change leaves it unsupported - // (e.g. switching from wpcom → anthropic-api-key while a GPT model is - // selected). Fall back to the provider's default. + // (e.g. switching from wpcom → anthropic-api-key while a tier is + // selected). Fall back to the provider's default — the quota-based one + // for WordPress.com. const definition = getAiProviderDefinition( currentProvider ); if ( ! definition.supportsModel( currentModel ) ) { - currentModel = definition.defaultModel; + currentModel = + currentProvider === DEFAULT_AI_PROVIDER ? wpcomDefaultModel : definition.defaultModel; ui.currentModel = currentModel; } @@ -521,6 +578,8 @@ export async function runCommand( options: { images: StudioChatImage[] = [], files: StudioChatFileAttachment[] = [] ): Promise< { status: TurnStatus; sessionId: string } > { + // The quota-based default must land before the turn captures its model. + await wpcomDefaultModelResolution; await maybeAutoSwitchProvider(); const sm = await ensureSession(); const sessionId = sm.getSessionId(); @@ -733,6 +792,8 @@ export async function runCommand( options: { }, set currentModel( value ) { currentModel = value; + // An explicit pick wins over the pending quota-based default. + quotaDefaultApplicable = false; }, get currentProvider() { return currentProvider; diff --git a/apps/cli/commands/ai/tests/index.test.ts b/apps/cli/commands/ai/tests/index.test.ts index 9379fcb837..58b8f1d327 100644 --- a/apps/cli/commands/ai/tests/index.test.ts +++ b/apps/cli/commands/ai/tests/index.test.ts @@ -24,6 +24,11 @@ import { runCommand } from '../index'; vi.mock( '@studio/common/lib/shared-config', () => ( { readAuthToken: vi.fn(), } ) ); +// The quota-based wpcom default must never hit the network in tests. +vi.mock( '@studio/common/lib/studio-assistant-quota', async ( importOriginal ) => ( { + ...( await importOriginal< typeof import('@studio/common/lib/studio-assistant-quota') >() ), + fetchStudioAssistantQuota: vi.fn().mockResolvedValue( null ), +} ) ); vi.mock( 'cli/lib/tracks', async ( importActual ) => { const actual = await importActual< typeof import('cli/lib/tracks') >(); return { ...actual, recordTracksEvent: vi.fn() }; @@ -42,7 +47,7 @@ vi.mock( 'cli/ai/providers', () => ( { DEFAULT_AI_PROVIDER: 'wpcom', getAiProviderDefinition: () => ( { supportsModel: () => true, - defaultModel: 'claude-default', + defaultModel: 'claude-sonnet-5', } ), } ) ); vi.mock( '@studio/common/ai/settings-store', () => ( { @@ -346,7 +351,7 @@ describe( 'AI runCommand — Tracks events', () => { expect( props ).toMatchObject( { outcome: 'interrupted', provider: 'wpcom', - model_family: 'anthropic', + model_family: 'studio', ai_session_id: 'session-id', client: 'studio-code', } ); diff --git a/apps/cli/tests/local-server.test.ts b/apps/cli/tests/local-server.test.ts index 3316036f6b..47de823f62 100644 --- a/apps/cli/tests/local-server.test.ts +++ b/apps/cli/tests/local-server.test.ts @@ -307,7 +307,7 @@ describe( 'local web server Connect contracts', () => { { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify( { provider: 'anthropic-api-key', model: 'gpt-5.6-sol' } ), + body: JSON.stringify( { provider: 'anthropic-api-key', model: 'balanced' } ), } ); diff --git a/apps/studio/src/components/studio-code-session/composer/index.test.tsx b/apps/studio/src/components/studio-code-session/composer/index.test.tsx index dd714a50b9..f2ca2c400d 100644 --- a/apps/studio/src/components/studio-code-session/composer/index.test.tsx +++ b/apps/studio/src/components/studio-code-session/composer/index.test.tsx @@ -9,11 +9,24 @@ const mockGetPathForFile = vi.hoisted( () => vi.fn( ( file: File ) => `/tmp/studio-attachments/${ file.name }` ) ); +vi.mock( 'src/hooks/use-auth', () => ( { + useAuth: () => ( { isAuthenticated: false } ), +} ) ); +vi.mock( 'src/stores/wpcom-api', async ( importOriginal ) => ( { + ...( await importOriginal< typeof import('src/stores/wpcom-api') >() ), + useGetStudioAssistantQuota: () => ( { data: undefined } ), + useGetStudioAssistantTopUpPricing: () => ( { data: undefined } ), +} ) ); vi.mock( 'src/lib/get-ipc-api', () => ( { getIpcApi: () => ( { getPathForFile: mockGetPathForFile, setAiSessionModel: vi.fn(), createAiSession: vi.fn(), + getAiSettings: vi.fn().mockResolvedValue( { + provider: 'wpcom', + hasAnthropicApiKey: false, + anthropicApiKeyPreview: null, + } ), } ), } ) ); diff --git a/apps/studio/src/components/studio-code-session/composer/index.tsx b/apps/studio/src/components/studio-code-session/composer/index.tsx index d06db5aa2f..708cdd7b4c 100644 --- a/apps/studio/src/components/studio-code-session/composer/index.tsx +++ b/apps/studio/src/components/studio-code-session/composer/index.tsx @@ -10,9 +10,21 @@ import { type ComposerAttachmentHoverPreviewState, } from '@studio/common/ai/composer-attachment-preview'; import { watchComposerFilePaste } from '@studio/common/ai/composer-attachments'; -import { getAiModelFamily, getAiModelLabel, getVisibleAiModels } from '@studio/common/ai/models'; +import { + aiModelRequiresPaidCredits, + getAiModelFamily, + getAiModelLabel, +} from '@studio/common/ai/models'; +import { getAiProviderModels, getEffectiveSessionProvider } from '@studio/common/ai/providers'; import { isStudioCustomEntryOfType } from '@studio/common/ai/sessions/entry-types'; import { isAutomatticianEmail } from '@studio/common/lib/automattician'; +import { + formatPaidTiersNudge, + getAddAiCreditsUrl, + hasPaidAiCredits, + persistPaidTiersNudgeDismissed, + readPaidTiersNudgeDismissed, +} from '@studio/common/lib/studio-assistant-quota'; import { useQueryClient } from '@tanstack/react-query'; import { createInterpolateElement } from '@wordpress/element'; import { __, sprintf } from '@wordpress/i18n'; @@ -26,9 +38,15 @@ import { useState, type SetStateAction, } from 'react'; +import { AiCreditsPurchaseDialog } from 'src/components/ai-credits-purchase-dialog'; +import { useAiSettings } from 'src/hooks/use-ai-settings'; import { useAuth } from 'src/hooks/use-auth'; import { cx } from 'src/lib/cx'; import { getIpcApi } from 'src/lib/get-ipc-api'; +import { + useGetStudioAssistantQuota, + useGetStudioAssistantTopUpPricing, +} from 'src/stores/wpcom-api'; import * as Menu from '../menu'; import { SESSIONS_QUERY_KEY } from '../use-session'; import { FamilySwitchConfirmDialog } from './family-switch-confirm-dialog'; @@ -307,11 +325,50 @@ export function Composer( { // the toolbar "/" toggle). Kept in its own hook so the Composer stays lean. const slash = useSlashCommands( { value, setValue: setDraftValue, textareaRef, previewPrompt } ); + // Only offer models the conversation's provider can serve. The paid tiers + // are listed but disabled for accounts without purchased credits. + const aiSettings = useAiSettings(); + const visibleModels = getAiProviderModels( + getEffectiveSessionProvider( entries ?? [], aiSettings ) + ); + const { isAuthenticated, user } = useAuth(); + const { data: quota } = useGetStudioAssistantQuota( undefined, { skip: ! isAuthenticated } ); + // The paid tiers unlock with purchased credits; Automatticians are exempt. + const canUsePaidTiers = hasPaidAiCredits( quota ) || isAutomatticianEmail( user?.email ); + const isModelLocked = useCallback( + ( id: AiModelId ) => aiModelRequiresPaidCredits( id ) && ! canUsePaidTiers, + [ canUsePaidTiers ] + ); + const hasLockedModels = visibleModels.some( ( { id } ) => isModelLocked( id ) ); + + // Nudge free-allowance accounts toward the paid tiers: a footer in the + // model picker plus a dismissible line above the prompt. Never shown while + // the quota is still loading, nor on top of the usage-cap banner. + const [ paidTiersNudgeDismissed, setPaidTiersNudgeDismissed ] = useState( + readPaidTiersNudgeDismissed + ); + const dismissPaidTiersNudge = () => { + setPaidTiersNudgeDismissed( true ); + persistPaidTiersNudgeDismissed(); + }; + const showPaidTiersNudge = + Boolean( quota ) && hasLockedModels && ! paidTiersNudgeDismissed && ! usageCapMessage; + + // Mirrors AddAiCreditsButton: the chooser when priced options exist, else + // straight to checkout for the single fixed top-up. + const { data: topUpPricing } = useGetStudioAssistantTopUpPricing(); + const [ creditsPurchaseOpen, setCreditsPurchaseOpen ] = useState( false ); + const openAddCredits = () => { + if ( ( topUpPricing?.options.length ?? 0 ) > 0 ) { + setCreditsPurchaseOpen( true ); + return; + } + void getIpcApi().openURL( getAddAiCreditsUrl( { returnsToDesktop: true } ) ); + }; + // Cross-family swap state. We hold the picked model here while the // confirmation dialog is open; nothing is persisted until the user // confirms. - const { user } = useAuth(); - const visibleModels = getVisibleAiModels( isAutomatticianEmail( user?.email ), model ); const [ pendingFamilyChange, setPendingFamilyChange ] = useState< AiModelId | null >( null ); const [ familySwitchInFlight, setFamilySwitchInFlight ] = useState( false ); @@ -393,7 +450,7 @@ export function Composer( { const handleModelChange = useCallback( ( picked: AiModelId ) => { - if ( picked === model ) { + if ( picked === model || isModelLocked( picked ) ) { return; } // Cross-family switch: defer until the user confirms in the dialog @@ -415,7 +472,7 @@ export function Composer( { } applySameFamilyModel( picked ); }, - [ applySameFamilyModel, entries, model, onSwitchSession ] + [ applySameFamilyModel, entries, isModelLocked, model, onSwitchSession ] ); const cancelFamilyChange = useCallback( () => { @@ -477,6 +534,26 @@ export function Composer( { { usageCapMessage } ) : null } + { showPaidTiersNudge ? ( +
+ { formatPaidTiersNudge() } + + +
+ ) : null }
handleModelChange( value as AiModelId ) } > - { visibleModels.map( ( { id, label } ) => ( - - { label } + { visibleModels.map( ( { id } ) => ( + + { getAiModelLabel( id ) } ) ) } + { hasLockedModels ? ( + <> + + { formatPaidTiersNudge() } + + ) : null } { busy ? ( @@ -749,6 +832,9 @@ export function Composer( { onCancel={ cancelFamilyChange } onConfirm={ () => void confirmFamilyChange() } /> + { creditsPurchaseOpen ? ( + + ) : null } ); } diff --git a/apps/studio/src/components/studio-code-session/composer/style.module.css b/apps/studio/src/components/studio-code-session/composer/style.module.css index cf5a686ee9..09ce079cbe 100644 --- a/apps/studio/src/components/studio-code-session/composer/style.module.css +++ b/apps/studio/src/components/studio-code-session/composer/style.module.css @@ -14,6 +14,41 @@ line-height: 1.4; } +.paidTiersNudge { + display: flex; + align-items: center; + gap: var(--wpds-dimension-gap-xs); + margin-bottom: var(--wpds-dimension-padding-sm); + padding: 0 var(--wpds-dimension-padding-xs); + color: var(--color-frame-text-secondary); + font-size: var(--wpds-typography-font-size-xs); +} + +.paidTiersNudgeAction { + border: 0; + padding: 0; + background: none; + color: var(--color-frame-text); + text-decoration: underline; + font-size: inherit; + cursor: pointer; +} + +.paidTiersNudgeDismiss { + display: inline-flex; + align-items: center; + margin-left: auto; + border: 0; + padding: 0; + background: none; + color: var(--color-frame-text-secondary); + cursor: pointer; +} + +.paidTiersNudgeDismiss:hover { + color: var(--color-frame-text); +} + .shell { display: flex; flex-direction: column; diff --git a/apps/studio/src/components/studio-code-session/index.tsx b/apps/studio/src/components/studio-code-session/index.tsx index 0ee0937634..da05d2cdbe 100644 --- a/apps/studio/src/components/studio-code-session/index.tsx +++ b/apps/studio/src/components/studio-code-session/index.tsx @@ -1,9 +1,15 @@ -import { resolveSessionModel } from '@studio/common/ai/models'; +import { + getEffectiveSessionProvider, + resolveSessionModelForProvider, +} from '@studio/common/ai/providers'; import { isStudioCustomEntryOfType, type StudioCustomEntry, } from '@studio/common/ai/sessions/entry-types'; -import { getStudioCodeAiAccessState } from '@studio/common/lib/studio-assistant-quota'; +import { + getStudioCodeAiAccessState, + hasPaidAiCredits, +} from '@studio/common/lib/studio-assistant-quota'; import { QueryClientProvider } from '@tanstack/react-query'; import { Spinner } from '@wordpress/components'; import { __ } from '@wordpress/i18n'; @@ -26,6 +32,7 @@ import Button from 'src/components/button'; import { IllustrationGrid } from 'src/components/illustration-grid'; import offlineIcon from 'src/components/offline-icon'; import { Tooltip } from 'src/components/tooltip'; +import { useAiSettings } from 'src/hooks/use-ai-settings'; import { useAuth } from 'src/hooks/use-auth'; import { useIsOutOfAiCredits } from 'src/hooks/use-is-out-of-ai-credits'; import { useOffline } from 'src/hooks/use-offline'; @@ -287,10 +294,19 @@ function SessionContent( { selectedSite }: { selectedSite: SiteDetails } ) { removeQueuedPrompt, } = useAgentRun( sessionId ); - const currentModel = useMemo( - () => resolveSessionModel( data?.entries ?? [] ), - [ data?.entries ] - ); + const { isAuthenticated } = useAuth(); + const { data: quota } = useGetStudioAssistantQuota( undefined, { skip: ! isAuthenticated } ); + const aiSettings = useAiSettings(); + // A fresh wpcom session defaults to balanced when purchased credits + // remain, fast otherwise. + const currentModel = useMemo( () => { + const entries = data?.entries ?? []; + return resolveSessionModelForProvider( + entries, + getEffectiveSessionProvider( entries, aiSettings ), + { hasPaidAiCredits: hasPaidAiCredits( quota ) } + ); + }, [ data?.entries, aiSettings, quota ] ); const pendingQuestionTexts = useMemo( () => new Set( pendingQuestions.map( ( q ) => q.question ) ), [ pendingQuestions ] diff --git a/apps/studio/src/components/studio-code-session/tests/access-requirements.test.tsx b/apps/studio/src/components/studio-code-session/tests/access-requirements.test.tsx index 28278ad458..16e443f929 100644 --- a/apps/studio/src/components/studio-code-session/tests/access-requirements.test.tsx +++ b/apps/studio/src/components/studio-code-session/tests/access-requirements.test.tsx @@ -12,6 +12,11 @@ const { mockIpc, quotaState } = vi.hoisted( () => ( { createAiSession: vi.fn(), markAiMessageEdited: vi.fn(), openURL: vi.fn(), + getAiSettings: vi.fn().mockResolvedValue( { + provider: 'wpcom', + hasAnthropicApiKey: false, + anthropicApiKeyPreview: null, + } ), }, quotaState: { data: undefined as Partial< StudioAssistantQuota > | undefined, diff --git a/apps/studio/src/components/studio-code-session/tests/session-not-found.test.tsx b/apps/studio/src/components/studio-code-session/tests/session-not-found.test.tsx index 0665d5f459..30ca50091f 100644 --- a/apps/studio/src/components/studio-code-session/tests/session-not-found.test.tsx +++ b/apps/studio/src/components/studio-code-session/tests/session-not-found.test.tsx @@ -12,6 +12,11 @@ const { mockIpc } = vi.hoisted( () => ( { loadAiSession: vi.fn(), createAiSession: vi.fn(), markAiMessageEdited: vi.fn(), + getAiSettings: vi.fn().mockResolvedValue( { + provider: 'wpcom', + hasAnthropicApiKey: false, + anthropicApiKeyPreview: null, + } ), }, } ) ); diff --git a/apps/studio/src/hooks/use-ai-settings.ts b/apps/studio/src/hooks/use-ai-settings.ts new file mode 100644 index 0000000000..def92cd045 --- /dev/null +++ b/apps/studio/src/hooks/use-ai-settings.ts @@ -0,0 +1,19 @@ +import { useEffect, useState } from 'react'; +import { getIpcApi } from 'src/lib/get-ipc-api'; +import type { AiSettings } from '@studio/common/ai/providers'; + +/** + * The saved AI provider settings from the main process. `undefined` until the + * first read resolves; consumers treat that as the default WordPress.com + * provider (see `getEffectiveSessionProvider`). + */ +export function useAiSettings(): AiSettings | undefined { + const [ settings, setSettings ] = useState< AiSettings | undefined >( undefined ); + useEffect( () => { + void getIpcApi() + .getAiSettings() + .then( setSettings ) + .catch( () => undefined ); + }, [] ); + return settings; +} diff --git a/apps/ui/src/ui-classic/components/session-view/composer/index.test.tsx b/apps/ui/src/ui-classic/components/session-view/composer/index.test.tsx index 306e90ffb1..9d1cc625f6 100644 --- a/apps/ui/src/ui-classic/components/session-view/composer/index.test.tsx +++ b/apps/ui/src/ui-classic/components/session-view/composer/index.test.tsx @@ -28,6 +28,8 @@ const connectorMocks = vi.hoisted( () => ( { capabilities: { aiSettings: false }, createSession: vi.fn(), getAiSettings: vi.fn(), + getAuthUser: vi.fn(), + getStudioAssistantQuota: vi.fn(), getFilePath: vi.fn( ( file: File ) => `/tmp/studio-attachments/${ file.name }` ), setSessionModel: vi.fn(), setSessionProvider: vi.fn(), @@ -86,7 +88,10 @@ describe( 'Composer menu', () => { hasAnthropicApiKey: true, anthropicApiKeyPreview: 'sk-ant-api03-tes...1234', } ); - renderComposer( { entries: [ createSessionContextEntry( 'anthropic-api-key' ) ] } ); + renderComposer( { + entries: [ createSessionContextEntry( 'anthropic-api-key' ) ], + model: 'claude-sonnet-5', + } ); const trigger = screen.getByRole( 'button', { name: 'Select model' } ); await waitFor( () => expect( trigger ).toHaveTextContent( 'API · Sonnet 5' ) ); @@ -115,7 +120,7 @@ describe( 'Composer menu', () => { fireEvent.click( trigger ); await waitFor( () => expect( screen.getAllByRole( 'menuitemradio' ).map( ( item ) => item.textContent ) ).toEqual( - [ 'Sonnet 5', 'Opus 5', 'GPT 5.6 Sol' ] + [ 'Fast', 'Balanced', 'Strong' ] ) ); } ); @@ -132,7 +137,7 @@ describe( 'Composer menu', () => { fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); await waitFor( () => expect( screen.getAllByRole( 'menuitemradio' ).map( ( item ) => item.textContent ) ).toEqual( - [ 'Sonnet 5', 'Opus 5', 'GPT 5.6 Sol' ] + [ 'Fast', 'Balanced', 'Strong' ] ) ); } ); @@ -150,7 +155,7 @@ describe( 'Composer menu', () => { summary: createSummary(), entries: [], } ); - renderComposer( { model: 'gpt-5.6-sol' }, queryClient ); + renderComposer( { model: 'balanced' }, queryClient ); fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); fireEvent.click( await screen.findByRole( 'menuitemradio', { name: 'Anthropic API' } ) ); @@ -172,6 +177,100 @@ describe( 'Composer menu', () => { ] ); } ); + it( 'disables the paid tiers unless purchased credits remain', async () => { + connectorMocks.getAuthUser.mockResolvedValue( { id: 1, email: 'user@example.com' } ); + connectorMocks.getStudioAssistantQuota.mockResolvedValue( { purchasedRemaining: 0 } ); + const { unmount } = renderComposer(); + + fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); + await waitFor( () => { + expect( screen.getByRole( 'menuitemradio', { name: 'Balanced' } ) ).toHaveAttribute( + 'aria-disabled', + 'true' + ); + } ); + expect( screen.getByRole( 'menuitemradio', { name: 'Strong' } ) ).toHaveAttribute( + 'aria-disabled', + 'true' + ); + expect( screen.getByRole( 'menuitemradio', { name: 'Fast' } ) ).not.toHaveAttribute( + 'aria-disabled' + ); + unmount(); + + connectorMocks.getStudioAssistantQuota.mockResolvedValue( { purchasedRemaining: 50_000 } ); + renderComposer(); + + fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); + await waitFor( () => { + expect( screen.getByRole( 'menuitemradio', { name: 'Balanced' } ) ).not.toHaveAttribute( + 'aria-disabled' + ); + } ); + } ); + + it( 'unlocks the paid tiers for Automatticians without purchased credits', async () => { + connectorMocks.getAuthUser.mockResolvedValue( { id: 1, email: 'person@automattic.com' } ); + connectorMocks.getStudioAssistantQuota.mockResolvedValue( { purchasedRemaining: 0 } ); + renderComposer(); + + fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); + await waitFor( () => { + expect( screen.getByRole( 'menuitemradio', { name: 'Balanced' } ) ).not.toHaveAttribute( + 'aria-disabled' + ); + } ); + expect( screen.getByRole( 'menuitemradio', { name: 'Strong' } ) ).not.toHaveAttribute( + 'aria-disabled' + ); + expect( + screen.queryByText( 'Add AI credits to unlock stronger models.' ) + ).not.toBeInTheDocument(); + } ); + + it( 'nudges free-allowance accounts toward paid tiers, dismissibly', async () => { + localStorage.removeItem( 'studio_code_paid_tiers_nudge_dismissed' ); + connectorMocks.getAuthUser.mockResolvedValue( { id: 1, email: 'user@example.com' } ); + connectorMocks.getStudioAssistantQuota.mockResolvedValue( { purchasedRemaining: 0 } ); + const { unmount } = renderComposer(); + + // Banner above the prompt, with the shared nudge copy. + expect( + await screen.findByText( 'Add AI credits to unlock stronger models.' ) + ).toBeInTheDocument(); + + // The model picker carries the same nudge as a footer item. + fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); + expect( + await screen.findByRole( 'menuitem', { name: 'Add AI credits to unlock stronger models.' } ) + ).toBeInTheDocument(); + fireEvent.keyDown( document.activeElement ?? document.body, { key: 'Escape' } ); + + // Dismissing hides the banner and persists. + fireEvent.click( screen.getByRole( 'button', { name: 'Dismiss' } ) ); + await waitFor( () => { + expect( + screen.queryByText( 'Add AI credits to unlock stronger models.' ) + ).not.toBeInTheDocument(); + } ); + expect( localStorage.getItem( 'studio_code_paid_tiers_nudge_dismissed' ) ).toBe( '1' ); + unmount(); + + // Paid accounts never see it. + localStorage.removeItem( 'studio_code_paid_tiers_nudge_dismissed' ); + connectorMocks.getStudioAssistantQuota.mockResolvedValue( { purchasedRemaining: 50_000 } ); + renderComposer(); + fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); + await waitFor( () => { + expect( screen.getByRole( 'menuitemradio', { name: 'Balanced' } ) ).not.toHaveAttribute( + 'aria-disabled' + ); + } ); + expect( + screen.queryByText( 'Add AI credits to unlock stronger models.' ) + ).not.toBeInTheDocument(); + } ); + it( 'shows tooltips for the plus button and model picker', async () => { renderComposer(); @@ -443,15 +542,22 @@ describe( 'Composer menu', () => { dialog.remove(); } ); - it( 'keeps the picked model in the fresh session cache after a family switch', async () => { + it( 'keeps the picked provider pin in the fresh session cache after a family switch', async () => { const queryClient = new QueryClient(); const onSwitchSession = vi.fn(); const freshSummary = createSummary( { id: 'fresh-session' } ); + connectorMocks.capabilities.aiSettings = true; + connectorMocks.getAiSettings.mockResolvedValue( { + provider: 'wpcom', + hasAnthropicApiKey: true, + anthropicApiKeyPreview: 'sk-ant-api03-tes...1234', + } ); connectorMocks.createSession.mockResolvedValue( freshSummary ); - connectorMocks.setSessionModel.mockResolvedValue( undefined ); + connectorMocks.setSessionProvider.mockResolvedValue( undefined ); renderComposer( { + model: 'balanced', entries: [ createUserPromptEntry() ], ownerSiteId: 'site-1', onSwitchSession, @@ -460,11 +566,11 @@ describe( 'Composer menu', () => { ); fireEvent.click( screen.getByRole( 'button', { name: 'Select model' } ) ); - fireEvent.click( await screen.findByText( 'GPT 5.6 Sol' ) ); + fireEvent.click( await screen.findByRole( 'menuitemradio', { name: 'Anthropic API' } ) ); const dialog = await screen.findByRole( 'dialog' ); expect( dialog ).toHaveTextContent( 'Start a new chat?' ); expect( dialog ).toHaveTextContent( - 'Switching from Sonnet 5 to GPT 5.6 Sol starts a fresh chat because the models don\u2019t share memory. You can find previous chats using Chat history below the chat box.' + 'Switching from Balanced to Sonnet 5 starts a fresh chat because the models don\u2019t share memory. You can find previous chats using Chat history below the chat box.' ); expect( within( dialog ).getByText( 'Chat history' ).tagName ).toBe( 'STRONG' ); expect( dialog ).not.toHaveTextContent( 'sidebar' ); @@ -473,7 +579,11 @@ describe( 'Composer menu', () => { await waitFor( () => { expect( onSwitchSession ).toHaveBeenCalledWith( 'fresh-session' ); } ); - expect( connectorMocks.setSessionModel ).toHaveBeenCalledWith( 'fresh-session', 'gpt-5.6-sol' ); + expect( connectorMocks.setSessionProvider ).toHaveBeenCalledWith( + 'fresh-session', + 'anthropic-api-key', + 'claude-sonnet-5' + ); const loadedSession = queryClient.getQueryData< LoadedAiSession >( [ ...SESSIONS_QUERY_KEY, @@ -482,8 +592,8 @@ describe( 'Composer menu', () => { expect( loadedSession?.summary ).toEqual( freshSummary ); expect( loadedSession?.entries ).toEqual( [ expect.objectContaining( { - type: 'model_change', - modelId: 'gpt-5.6-sol', + customType: 'studio.session_context', + data: { provider: 'anthropic-api-key', model: 'claude-sonnet-5' }, } ), ] ); } ); diff --git a/apps/ui/src/ui-classic/components/session-view/composer/index.tsx b/apps/ui/src/ui-classic/components/session-view/composer/index.tsx index f8da342d75..bfd267512c 100644 --- a/apps/ui/src/ui-classic/components/session-view/composer/index.tsx +++ b/apps/ui/src/ui-classic/components/session-view/composer/index.tsx @@ -11,24 +11,28 @@ import { } from '@studio/common/ai/composer-attachment-preview'; import { watchComposerFilePaste } from '@studio/common/ai/composer-attachments'; import { - AI_MODELS, + aiModelRequiresPaidCredits, getAiModelFamily, getAiModelLabel, - getVisibleAiModels, } from '@studio/common/ai/models'; import { AI_PROVIDER_IDS, AI_PROVIDER_LABELS, - DEFAULT_AI_PROVIDER, getAiProviderDefaultModel, getAiProviderModels, + getEffectiveSessionProvider, providerServesModel, - resolveSessionProvider, type AiProviderId, } from '@studio/common/ai/providers'; import { isStudioCustomEntryOfType } from '@studio/common/ai/sessions/entry-types'; import { getAiSkillCommands } from '@studio/common/ai/slash-commands'; import { isAutomatticianEmail } from '@studio/common/lib/automattician'; +import { + formatPaidTiersNudge, + hasPaidAiCredits, + persistPaidTiersNudgeDismissed, + readPaidTiersNudgeDismissed, +} from '@studio/common/lib/studio-assistant-quota'; import { useQueryClient } from '@tanstack/react-query'; import { __, sprintf } from '@wordpress/i18n'; import { @@ -56,15 +60,19 @@ import { type PointerEvent, } from 'react'; import { createPortal } from 'react-dom'; +import { AiCreditsPurchaseDialog } from '@/components/ai-credits-purchase-dialog'; import * as Menu from '@/components/menu'; import { useConnector } from '@/data/core'; import { useAiSettings } from '@/data/queries/use-ai-settings'; +import { useStudioAssistantQuota } from '@/data/queries/use-assistant-quota'; import { useAuthUser } from '@/data/queries/use-auth-user'; import { primeSessionQueryData, reconcilePrimedSessionQueryData, SESSIONS_QUERY_KEY, } from '@/data/queries/use-sessions'; +import { useStudioAssistantTopUpPricing } from '@/data/queries/use-top-up-pricing'; +import { useAddAiCreditsUrl } from '@/hooks/use-add-ai-credits-url'; import { AiCreditsControl } from './ai-credits-control'; import { FamilySwitchConfirmDialog } from './family-switch-confirm-dialog'; import styles from './style.module.css'; @@ -351,26 +359,50 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co const connector = useConnector(); const queryClient = useQueryClient(); - // The conversation's provider: its own pinned choice first, then the saved - // global selection. Without a saved Anthropic key the pin is unusable, so - // WordPress.com wins regardless — the CLI applies the same rule on resume. const { data: aiSettings } = useAiSettings(); - const pinnedProvider = useMemo( () => resolveSessionProvider( entries ?? [] ), [ entries ] ); - const sessionProvider = aiSettings?.hasAnthropicApiKey - ? pinnedProvider ?? aiSettings.provider - : DEFAULT_AI_PROVIDER; + const sessionProvider = useMemo( + () => getEffectiveSessionProvider( entries ?? [], aiSettings ), + [ entries, aiSettings ] + ); const canPickProvider = Boolean( aiSettings?.hasAnthropicApiKey && sessionId ); - // Only offer models the conversation's provider can serve. Hosts without AI - // settings (capabilities.aiSettings false) keep the full list. - const availableModels = aiSettings ? getAiProviderModels( sessionProvider ) : AI_MODELS; + // Only offer models the conversation's provider can serve. The paid tiers + // are listed but disabled for accounts without purchased credits; + // Automatticians are exempt. + const offeredModels = getAiProviderModels( sessionProvider ); + const { data: quota } = useStudioAssistantQuota(); const { data: authUser } = useAuthUser(); - const visibleModelIds = new Set( - getVisibleAiModels( isAutomatticianEmail( authUser?.email ), model ).map( - ( entry ) => entry.id - ) + const canUsePaidTiers = hasPaidAiCredits( quota ) || isAutomatticianEmail( authUser?.email ); + const isModelLocked = useCallback( + ( id: AiModelId ) => aiModelRequiresPaidCredits( id ) && ! canUsePaidTiers, + [ canUsePaidTiers ] + ); + const hasLockedModels = offeredModels.some( ( { id } ) => isModelLocked( id ) ); + + // Nudge free-allowance accounts toward the paid tiers: a footer in the + // model picker plus a dismissible line above the prompt. Never shown while + // the quota is still loading. + const [ paidTiersNudgeDismissed, setPaidTiersNudgeDismissed ] = useState( + readPaidTiersNudgeDismissed ); - const offeredModels = availableModels.filter( ( entry ) => visibleModelIds.has( entry.id ) ); + const dismissPaidTiersNudge = () => { + setPaidTiersNudgeDismissed( true ); + persistPaidTiersNudgeDismissed(); + }; + const showPaidTiersNudge = Boolean( quota ) && hasLockedModels && ! paidTiersNudgeDismissed; + + // Mirrors AiCreditsControl: the chooser when priced options exist, else + // straight to checkout for the single fixed top-up. + const addAiCreditsUrl = useAddAiCreditsUrl(); + const { data: topUpPricing } = useStudioAssistantTopUpPricing(); + const [ creditsPurchaseOpen, setCreditsPurchaseOpen ] = useState( false ); + const openAddCredits = () => { + if ( ( topUpPricing?.options.length ?? 0 ) > 0 ) { + setCreditsPurchaseOpen( true ); + return; + } + void connector.openExternalUrl( addAiCreditsUrl ); + }; const slash = useSlashCommands( { value, @@ -394,10 +426,13 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co } = useComposerAttachments(); const hasAttachments = attachments.length > 0; - // Cross-family swap state. We hold the picked model here while the - // confirmation dialog is open; nothing is persisted until the user - // confirms. - const [ pendingFamilyChange, setPendingFamilyChange ] = useState< AiModelId | null >( null ); + // Cross-family swap state. We hold the picked model (and provider, when + // the swap came from the provider picker) here while the confirmation + // dialog is open; nothing is persisted until the user confirms. + const [ pendingFamilyChange, setPendingFamilyChange ] = useState< { + model: AiModelId; + provider?: AiProviderId; + } | null >( null ); const [ familySwitchInFlight, setFamilySwitchInFlight ] = useState( false ); const setComposerManualTextareaHeight = useCallback( ( height: number | null ) => { @@ -628,9 +663,18 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co [ appendEntryOptimistically, connector ] ); - // Pin this conversation to a provider. If it can't serve the current model, - // its default model rides along in the same entry, so the model section - // re-filters via `resolveSessionModel`. + const sessionHasTurns = useMemo( + () => + ( entries ?? [] ).some( ( entry ) => + isStudioCustomEntryOfType( entry, 'studio.user_prompt' ) + ), + [ entries ] + ); + + // Pin this conversation to a provider, carrying a model it serves in the + // same entry. Providers don't share a model family, so the switch goes + // through the same fresh-session confirmation as a cross-family model + // switch. const handleProviderChange = useCallback( ( picked: AiProviderId ) => { if ( picked === sessionProvider ) { @@ -638,17 +682,33 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co } const nextModel = providerServesModel( picked, model ) ? model - : getAiProviderDefaultModel( picked ); + : getAiProviderDefaultModel( picked, { hasPaidAiCredits: hasPaidAiCredits( quota ) } ); + if ( + getAiModelFamily( model ) !== getAiModelFamily( nextModel ) && + onSwitchSession && + sessionHasTurns + ) { + setPendingFamilyChange( { model: nextModel, provider: picked } ); + return; + } appendEntryOptimistically( createSessionContextEntry( picked, nextModel ), ( id ) => connector.setSessionProvider( id, picked, nextModel ) ); }, - [ appendEntryOptimistically, connector, model, sessionProvider ] + [ + appendEntryOptimistically, + connector, + model, + onSwitchSession, + quota, + sessionHasTurns, + sessionProvider, + ] ); const handleModelChange = useCallback( ( picked: AiModelId ) => { - if ( picked === model ) { + if ( picked === model || isModelLocked( picked ) ) { return; } // Cross-family switch: defer until the user confirms in the dialog @@ -657,20 +717,17 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co // with the agent's actual memory. We skip the prompt when the // session has no user turns yet, or when the parent cannot switch // to a freshly created session. - const hasTurns = ( entries ?? [] ).some( ( entry ) => - isStudioCustomEntryOfType( entry, 'studio.user_prompt' ) - ); if ( getAiModelFamily( model ) !== getAiModelFamily( picked ) && onSwitchSession && - hasTurns + sessionHasTurns ) { - setPendingFamilyChange( picked ); + setPendingFamilyChange( { model: picked } ); return; } applySameFamilyModel( picked ); }, - [ applySameFamilyModel, entries, model, onSwitchSession ] + [ applySameFamilyModel, isModelLocked, model, onSwitchSession, sessionHasTurns ] ); const cancelFamilyChange = useCallback( () => { @@ -684,33 +741,37 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co if ( ! pendingFamilyChange || ! onSwitchSession ) { return; } - const pickedModel = pendingFamilyChange; + const { model: pickedModel, provider: pickedProvider } = pendingFamilyChange; setFamilySwitchInFlight( true ); try { const newSession = await connector.createSession( ownerSiteId ); primeSessionQueryData( queryClient, newSession ); - // Persist the model on the fresh session before navigating so the - // composer there opens already on the picked family — - // `setSessionModel` writes a `session.model_selected` event the - // new view picks up via `resolveSessionModel`. If this fails we - // still navigate; the user can re-pick from the new view's - // dropdown. - const modelPersisted = await connector - .setSessionModel( newSession.id, pickedModel ) + // Persist the model (and the provider pin, when the swap came from + // the provider picker) on the fresh session before navigating so + // the composer there opens already on the picked family. If this + // fails we still navigate; the user can re-pick from the new + // view's dropdown. + const pinEntry = pickedProvider + ? createSessionContextEntry( pickedProvider, pickedModel ) + : createModelChangeEntry( pickedModel ); + const persisted = await ( pickedProvider + ? connector.setSessionProvider( newSession.id, pickedProvider, pickedModel ) + : connector.setSessionModel( newSession.id, pickedModel ) + ) .then( () => true ) .catch( () => false ); - if ( modelPersisted ) { + if ( persisted ) { queryClient.setQueryData< LoadedAiSession >( [ ...SESSIONS_QUERY_KEY, newSession.id ], ( current ) => current ? { ...current, - entries: [ ...( current.entries ?? [] ), createModelChangeEntry( pickedModel ) ], + entries: [ ...( current.entries ?? [] ), pinEntry ], } : { summary: newSession, - entries: [ createModelChangeEntry( pickedModel ) ], + entries: [ pinEntry ], } ); } @@ -763,6 +824,26 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co return ( <>
+ { showPaidTiersNudge ? ( +
+ { formatPaidTiersNudge() } + + +
+ ) : null }
( function Co value={ model } onValueChange={ ( value ) => handleModelChange( value as AiModelId ) } > - { offeredModels.map( ( { id, label } ) => ( - - { label } + { offeredModels.map( ( { id } ) => ( + + { getAiModelLabel( id ) } ) ) } + { hasLockedModels ? ( + <> + + { formatPaidTiersNudge() } + + ) : null } { busy ? ( @@ -1156,11 +1243,14 @@ export const Composer = forwardRef< ComposerHandle, ComposerProps >( function Co
void confirmFamilyChange() } /> + { creditsPurchaseOpen ? ( + + ) : null } ); } ); diff --git a/apps/ui/src/ui-classic/components/session-view/composer/style.module.css b/apps/ui/src/ui-classic/components/session-view/composer/style.module.css index 13746a6da0..af59819742 100644 --- a/apps/ui/src/ui-classic/components/session-view/composer/style.module.css +++ b/apps/ui/src/ui-classic/components/session-view/composer/style.module.css @@ -648,6 +648,45 @@ pointer-events: none; } +.paidTiersNudge { + display: flex; + align-items: center; + gap: var(--wpds-dimension-gap-xs); + margin-bottom: var(--wpds-dimension-padding-sm); + padding: 0 var(--wpds-dimension-padding-xs); + color: var(--wpds-color-fg-content-neutral-weak); + font-size: var(--wpds-typography-font-size-xs); +} + +.paidTiersNudgeAction { + border: 0; + padding: 0; + background: none; + color: var(--wpds-color-fg-content-neutral); + text-decoration: underline; + font-size: inherit; + cursor: var(--wpds-cursor-interactive, pointer); +} + +.paidTiersNudgeAction:hover { + color: var(--wpds-color-fg-content-neutral); +} + +.paidTiersNudgeDismiss { + display: inline-flex; + align-items: center; + margin-left: auto; + border: 0; + padding: 0; + background: none; + color: var(--wpds-color-fg-content-neutral-weak); + cursor: var(--wpds-cursor-interactive, pointer); +} + +.paidTiersNudgeDismiss:hover { + color: var(--wpds-color-fg-content-neutral); +} + .meta { display: flex; align-items: center; diff --git a/apps/ui/src/ui-classic/components/session-view/index.test.tsx b/apps/ui/src/ui-classic/components/session-view/index.test.tsx index e6def04d55..6c09d943c1 100644 --- a/apps/ui/src/ui-classic/components/session-view/index.test.tsx +++ b/apps/ui/src/ui-classic/components/session-view/index.test.tsx @@ -25,10 +25,17 @@ vi.mock( '@/data/queries/use-sites', () => ( { vi.mock( '@/data/queries/use-assistant-quota', () => ( { useStudioAssistantQuota: vi.fn(), } ) ); +vi.mock( '@/data/queries/use-ai-settings', () => ( { + useAiSettings: () => ( { data: undefined } ), +} ) ); vi.mock( '@/data/core', async ( importOriginal ) => ( { ...( await importOriginal< object >() ), - useConnector: () => ( { openExternalUrl: vi.fn() } ), + useConnector: () => ( { + openExternalUrl: vi.fn(), + capabilities: { aiSettings: false }, + getAiSettings: vi.fn(), + } ), } ) ); vi.mock( '@/data/queries/use-agent-run', () => ( { diff --git a/apps/ui/src/ui-classic/components/session-view/index.tsx b/apps/ui/src/ui-classic/components/session-view/index.tsx index 77461e2cdc..41f2daa22c 100644 --- a/apps/ui/src/ui-classic/components/session-view/index.tsx +++ b/apps/ui/src/ui-classic/components/session-view/index.tsx @@ -1,6 +1,12 @@ -import { resolveSessionModel } from '@studio/common/ai/models'; +import { + getEffectiveSessionProvider, + resolveSessionModelForProvider, +} from '@studio/common/ai/providers'; import { findAiSessionOwnerSite } from '@studio/common/ai/sessions/owner-site'; -import { getStudioCodeAiAccessState } from '@studio/common/lib/studio-assistant-quota'; +import { + getStudioCodeAiAccessState, + hasPaidAiCredits, +} from '@studio/common/lib/studio-assistant-quota'; import { useNavigate } from '@tanstack/react-router'; import { __ } from '@wordpress/i18n'; import { arrowDown } from '@wordpress/icons'; @@ -23,6 +29,7 @@ import { SiteDropdown } from '@/components/site-dropdown'; import { SiteIcon } from '@/components/site-icon'; import { type Annotation } from '@/components/site-preview/types'; import { useAgentRun } from '@/data/queries/use-agent-run'; +import { useAiSettings } from '@/data/queries/use-ai-settings'; import { useStudioAssistantQuota } from '@/data/queries/use-assistant-quota'; import { useCreateSession, @@ -251,10 +258,23 @@ function SessionViewContent( { sessionId }: { sessionId: string } ) { answerQuestion, removeQueuedPrompt, } = useAgentRun( sessionId ); - const currentModel = useMemo( - () => resolveSessionModel( data?.entries ?? [] ), - [ data?.entries ] - ); + const { + data: quota, + isLoading: isQuotaLoading, + isFetching: isQuotaFetching, + refetch: refetchQuota, + } = useStudioAssistantQuota(); + const { data: aiSettings } = useAiSettings(); + // A fresh wpcom session defaults to balanced when purchased credits + // remain, fast otherwise. + const currentModel = useMemo( () => { + const entries = data?.entries ?? []; + return resolveSessionModelForProvider( + entries, + getEffectiveSessionProvider( entries, aiSettings ), + { hasPaidAiCredits: hasPaidAiCredits( quota ) } + ); + }, [ data?.entries, aiSettings, quota ] ); const pendingQuestionTexts = useMemo( () => new Set( pendingQuestions.map( ( q ) => q.question ) ), [ pendingQuestions ] @@ -398,12 +418,6 @@ function SessionViewContent( { sessionId }: { sessionId: string } ) { queuedPrompts.length, ] ); - const { - data: quota, - isLoading: isQuotaLoading, - isFetching: isQuotaFetching, - refetch: refetchQuota, - } = useStudioAssistantQuota(); // Out of credits replaces the composer: there is nothing to type into // until the account buys more, so the offer takes the input's place. const isOutOfCredits = useIsOutOfAiCredits(); diff --git a/packages/common/ai/models.ts b/packages/common/ai/models.ts index e6da7b58dd..bc05d99a38 100644 --- a/packages/common/ai/models.ts +++ b/packages/common/ai/models.ts @@ -1,7 +1,8 @@ +import { __ } from '@wordpress/i18n'; import { isStudioCustomEntryOfType } from './sessions/entry-types'; import type { SessionEntry } from '@earendil-works/pi-coding-agent'; -export type AiModelFamily = 'anthropic' | 'openai' | 'hosted'; +export type AiModelFamily = 'anthropic' | 'studio'; export interface AiModel { /** Stable model id sent to the upstream provider. */ @@ -17,72 +18,31 @@ export interface AiModel { */ supportsImages?: boolean; /** - * Hide the model from pickers for non-Automatticians. Visibility only — - * the wpcom proxy is what actually refuses these upstreams, so this never - * gates `isAiModelId`: a session that already recorded one must still - * resolve rather than silently snap back to the default. + * Pickers disable the model unless purchased AI credits remain + * (Automatticians exempt). UI gating only — the wpcom proxy enforces + * access, and `isAiModelId` never consults this. */ - requiresAutomattician?: boolean; + requiresPaidAiCredits?: boolean; } -// Pro / o-series OpenAI variants (`gpt-*-pro`, `o[1-9]*`) are intentionally -// excluded: their long reasoning turns can exceed the proxy/SDK timeout -// window. (Routing is no longer a blocker — the OpenAI family now goes -// through the proxy's `/v1/responses` path, which supports reasoning models -// and function tools.) +// The `studio` family are capability tiers, not concrete models: the wpcom +// proxy's `studio-agent` lane resolves each alias to an upstream model +// server-side, so the mapping can be retuned without a client release. The +// `anthropic` family exists for the direct Anthropic · API key provider only. export const AI_MODELS = [ + { id: 'fast', label: 'Fast', family: 'studio', supportsImages: false }, + { id: 'balanced', label: 'Balanced', family: 'studio', requiresPaidAiCredits: true }, + { id: 'strong', label: 'Strong', family: 'studio', requiresPaidAiCredits: true }, { id: 'claude-sonnet-5', label: 'Sonnet 5', family: 'anthropic' }, { id: 'claude-opus-5', label: 'Opus 5', family: 'anthropic' }, - { id: 'gpt-5.6-sol', label: 'GPT 5.6 Sol', family: 'openai' }, - // Hosted models keep their vendor-prefixed ids — the proxy passes them - // upstream verbatim. - { id: 'moonshotai/Kimi-K3', label: 'Kimi K3', family: 'hosted', requiresAutomattician: true }, - { - id: 'moonshotai/Kimi-K2.6', - label: 'Kimi K2.6', - family: 'hosted', - requiresAutomattician: true, - }, - { - id: 'zai-org/GLM-5.3-Flash', - label: 'GLM 5.3 Flash', - family: 'hosted', - supportsImages: false, - requiresAutomattician: true, - }, - { - id: 'zai-org/GLM-5.2', - label: 'GLM 5.2', - family: 'hosted', - supportsImages: false, - requiresAutomattician: true, - }, - { - id: 'zai-org/GLM-5.2-Fast', - label: 'GLM 5.2 Fast', - family: 'hosted', - supportsImages: false, - requiresAutomattician: true, - }, - { - id: 'deepseek-ai/DeepSeek-V4-Pro', - label: 'DeepSeek V4 Pro', - family: 'hosted', - supportsImages: false, - requiresAutomattician: true, - }, - { - id: 'deepseek-ai/DeepSeek-V4-Flash-0731', - label: 'DeepSeek V4 Flash', - family: 'hosted', - supportsImages: false, - requiresAutomattician: true, - }, ] as const satisfies readonly AiModel[]; export type AiModelId = ( typeof AI_MODELS )[ number ][ 'id' ]; -export const DEFAULT_MODEL: AiModelId = 'claude-sonnet-5'; +export const DEFAULT_MODEL: AiModelId = 'fast'; +// Accounts with purchased AI credits remaining default to the balanced tier +// instead (see `getAiProviderDefaultModel`). +export const PAID_DEFAULT_MODEL: AiModelId = 'balanced'; // Module-scoped lookup so `getAiModelFamily` / `getAiModelLabel` are O(1) // and don't re-scan the array per call. Keyed by id; values are the same @@ -109,31 +69,20 @@ export function getAiModelFamily( id: AiModelId ): AiModelFamily { return getAiModel( id ).family; } +// The tier labels are plain adjectives (unlike the Anthropic brand names), so +// they go through i18n — as thunks, since module-level `__()` is banned. +const TRANSLATED_MODEL_LABELS: Partial< Record< AiModelId, () => string > > = { + fast: () => __( 'Fast' ), + balanced: () => __( 'Balanced' ), + strong: () => __( 'Strong' ), +}; + export function getAiModelLabel( id: AiModelId ): string { - return getAiModel( id ).label; + return TRANSLATED_MODEL_LABELS[ id ]?.() ?? getAiModel( id ).label; } -// Widened to AiModel: `as const satisfies` narrows each entry to its own -// literal type, so the optional flags aren't on the union. -const ALL_MODELS = AI_MODELS as readonly AiModel[]; -const UNRESTRICTED_MODELS = ALL_MODELS.filter( ( model ) => ! model.requiresAutomattician ); - -/** - * The models to offer in a picker. Restricted models stay in `AI_MODELS` (so - * ids keep validating) but are withheld from anyone who isn't an Automattician. - * - * `keepId` is always offered even when restricted, so a picker never hides the - * value it is currently displaying. - */ -export function getVisibleAiModels( - isAutomattician: boolean, - keepId?: AiModelId -): readonly AiModel[] { - const visible = isAutomattician ? ALL_MODELS : UNRESTRICTED_MODELS; - if ( ! keepId || visible.some( ( model ) => model.id === keepId ) ) { - return visible; - } - return [ ...visible, getAiModel( keepId ) ]; +export function aiModelRequiresPaidCredits( id: AiModelId ): boolean { + return getAiModel( id ).requiresPaidAiCredits ?? false; } // Tolerates ids outside AI_MODELS — callers can reach here with a cast, and @@ -163,19 +112,24 @@ function readEntryModelId( entry: SessionEntry ): string | undefined { } /** - * Derive the current model for a session from its pi entries. - * - * The most recently recorded model wins. If it names a model we no longer - * offer (e.g. one that was removed from `AI_MODELS`), the session - * auto-switches to `DEFAULT_MODEL` rather than pinning a dead id. Sessions - * that recorded no model — e.g. a brand-new session before the first turn - * runs — also fall back to `DEFAULT_MODEL`. + * The most recently recorded model still in `AI_MODELS`, or `undefined` when + * the session never recorded one (or only models we no longer offer) — so + * callers can apply their own default to both cases. */ -export function resolveSessionModel( entries: SessionEntry[] ): AiModelId { +export function readRecordedSessionModel( entries: SessionEntry[] ): AiModelId | undefined { for ( let index = entries.length - 1; index >= 0; index -= 1 ) { const recordedModel = readEntryModelId( entries[ index ] ); - if ( recordedModel === undefined ) continue; - return isAiModelId( recordedModel ) ? recordedModel : DEFAULT_MODEL; + if ( recordedModel !== undefined && isAiModelId( recordedModel ) ) { + return recordedModel; + } } - return DEFAULT_MODEL; + return undefined; +} + +/** `readRecordedSessionModel` with a fallback for sessions without one. */ +export function resolveSessionModel( + entries: SessionEntry[], + defaultModel: AiModelId = DEFAULT_MODEL +): AiModelId { + return readRecordedSessionModel( entries ) ?? defaultModel; } diff --git a/packages/common/ai/providers.ts b/packages/common/ai/providers.ts index 227fbe7cfa..d08fef4b2b 100644 --- a/packages/common/ai/providers.ts +++ b/packages/common/ai/providers.ts @@ -1,4 +1,11 @@ -import { AI_MODELS, DEFAULT_MODEL, type AiModelFamily, type AiModelId } from './models'; +import { + AI_MODELS, + DEFAULT_MODEL, + PAID_DEFAULT_MODEL, + resolveSessionModel, + type AiModelFamily, + type AiModelId, +} from './models'; import { isStudioCustomEntryOfType } from './sessions/entry-types'; import type { SessionEntry } from '@earendil-works/pi-coding-agent'; @@ -21,11 +28,11 @@ export const AI_PROVIDER_LABELS: Record< AiProviderId, string > = { 'anthropic-api-key': 'Anthropic API', }; -// Which model families each provider can service. `wpcom` relays the -// Anthropic, OpenAI, and hosted wire formats through the same proxy; +// Which model families each provider can service. `wpcom` serves only the +// studio capability tiers (resolved to upstream models by the proxy); // direct-API providers are restricted to their own family. const PROVIDER_MODEL_FAMILIES: Record< AiProviderId, readonly AiModelFamily[] > = { - wpcom: [ 'anthropic', 'openai', 'hosted' ], + wpcom: [ 'studio' ], 'anthropic-api-key': [ 'anthropic' ], }; @@ -55,11 +62,35 @@ export function providerServesModel( provider: AiProviderId, model: AiModelId ): return getAiProviderModels( provider ).some( ( entry ) => entry.id === model ); } -/** The model to fall back to when a provider can't serve the requested one. */ -export function getAiProviderDefaultModel( provider: AiProviderId ): AiModelId { +/** + * The fallback when no model was chosen or the provider can't serve the + * requested one. On wpcom, paid credits upgrade the default to the balanced + * tier; everyone else (including callers without quota data) gets fast. + */ +export function getAiProviderDefaultModel( + provider: AiProviderId, + options?: { hasPaidAiCredits?: boolean } +): AiModelId { + if ( provider === 'wpcom' && options?.hasPaidAiCredits ) { + return PAID_DEFAULT_MODEL; + } return getAiProviderModels( provider )[ 0 ]?.id ?? DEFAULT_MODEL; } +/** + * `resolveSessionModel` constrained to what the provider can serve: a + * recorded model it no longer offers snaps to the provider's default. + */ +export function resolveSessionModelForProvider( + entries: SessionEntry[], + provider: AiProviderId, + options?: { hasPaidAiCredits?: boolean } +): AiModelId { + const defaultModel = getAiProviderDefaultModel( provider, options ); + const model = resolveSessionModel( entries, defaultModel ); + return providerServesModel( provider, model ) ? model : defaultModel; +} + /** * The provider a session is pinned to: the latest `studio.session_context` * entry naming a provider wins. Only explicit switches record one (per-turn @@ -79,6 +110,21 @@ export function resolveSessionProvider( entries: SessionEntry[] ): AiProviderId return undefined; } +/** + * The provider a conversation effectively runs on: its pinned choice first, + * then the saved global selection. Without a saved Anthropic key the pin is + * unusable (as are missing/unloaded settings), so WordPress.com wins. + */ +export function getEffectiveSessionProvider( + entries: SessionEntry[], + settings?: Pick< AiSettings, 'provider' | 'hasAnthropicApiKey' > | null +): AiProviderId { + if ( ! settings?.hasAnthropicApiKey ) { + return DEFAULT_AI_PROVIDER; + } + return resolveSessionProvider( entries ) ?? settings.provider; +} + /** * The AI provider settings exposed to the settings UI. The Anthropic API key * itself never leaves the server — only its presence and a short suffix for diff --git a/packages/common/lib/studio-assistant-quota.ts b/packages/common/lib/studio-assistant-quota.ts index 1e05d9040f..ab2dd1612b 100644 --- a/packages/common/lib/studio-assistant-quota.ts +++ b/packages/common/lib/studio-assistant-quota.ts @@ -124,6 +124,41 @@ export function getTotalRemainingAiCredits( return ( quota.allowanceRemaining ?? 0 ) + ( quota.purchasedRemaining ?? 0 ); } +/** + * Whether the account still has purchased (paid) AI credits. Unknown quota + * (unreachable, feature off) reads as false — the free tier is the safe floor. + */ +export function hasPaidAiCredits( + quota: Pick< StudioAssistantQuota, 'purchasedRemaining' > | null | undefined +): boolean { + return ( quota?.purchasedRemaining ?? 0 ) > 0; +} + +/** Nudge shown where the paid tiers sit disabled, one string everywhere. */ +export function formatPaidTiersNudge(): string { + return __( 'Add AI credits to unlock stronger models.' ); +} + +// Dismissal is remembered per renderer via localStorage — a nudge doesn't +// warrant synced, server-side state. +export const PAID_TIERS_NUDGE_DISMISSED_STORAGE_KEY = 'studio_code_paid_tiers_nudge_dismissed'; + +export function readPaidTiersNudgeDismissed(): boolean { + try { + return localStorage.getItem( PAID_TIERS_NUDGE_DISMISSED_STORAGE_KEY ) === '1'; + } catch { + return false; + } +} + +export function persistPaidTiersNudgeDismissed(): void { + try { + localStorage.setItem( PAID_TIERS_NUDGE_DISMISSED_STORAGE_KEY, '1' ); + } catch { + // Ignore storage errors. + } +} + export type StudioCodeAiAccessState = 'available' | 'blocked' | 'not-enabled'; /**