Make wpcom.req.* default response type unknown and add more zod schemas

src/components/tests/content-tab-assistant.test.tsx

@@ -52,8 +52,8 @@ const runningSite = {
 };
 
 const initialMessages = [
-	generateMessage( 'Initial message 1', 'user', 0, 'chat-id', 10 ),
-	generateMessage( 'Initial message 2', 'assistant', 1, 'chat-id', 11 ),
+	generateMessage( 'Initial message 1', 'user', 0, 100, 10 ),
+	generateMessage( 'Initial message 2', 'assistant', 1, 100, 11 ),
 ];
 
 function ContextWrapper( props: Parameters< typeof ContentTabAssistant >[ 0 ] ) {
@@ -71,10 +71,8 @@ describe( 'ContentTabAssistant', () => {
 		callback(
 			null,
 			{
-				id: 'chatcmpl-9USNsuhHWYsPAUNiOhOG2970Hjwwb',
-				object: 'chat.completion',
-				created: 1717045976,
-				model: 'test',
+				id: 100,
+				created_at: '2025-01-24 09:11:50',
 				choices: [
 					{
 						index: 0,
@@ -84,12 +82,8 @@ describe( 'ContentTabAssistant', () => {
 							content:
 								'Hello! How can I assist you today? Are you working on a WordPress project, or do you need help with something specific related to WordPress or WP-CLI?',
 						},
-						logprobs: null,
-						finish_reason: 'stop',
 					},
 				],
-				usage: { prompt_tokens: 980, completion_tokens: 36, total_tokens: 1016 },
-				system_fingerprint: 'fp_777',
 			},
 			{
 				'x-quota-max': '100',
@@ -355,9 +349,9 @@ describe( 'ContentTabAssistant', () => {
 		jest.useFakeTimers();
 		jest.setSystemTime( MOCKED_CURRENT_TIME );
 
-		const messageOne = generateMessage( 'Initial message 1', 'user', 0, 'hej', 10 );
+		const messageOne = generateMessage( 'Initial message 1', 'user', 0, 100, 10 );
 		messageOne.createdAt = MOCKED_CURRENT_TIME;
-		const messageTwo = generateMessage( 'Initial message 2', 'assistant', 1, 'hej', 11 );
+		const messageTwo = generateMessage( 'Initial message 2', 'assistant', 1, 100, 11 );
 		messageTwo.createdAt = OLD_MESSAGE_TIME;
 		store.dispatch(
 			chatActions.setMessages( {
@@ -398,9 +392,9 @@ describe( 'ContentTabAssistant', () => {
 	} );
 
 	it( 'renders notices by importance', async () => {
-		const messageOne = generateMessage( 'Initial message 1', 'user', 0, 'chat-id', 10 );
+		const messageOne = generateMessage( 'Initial message 1', 'user', 0, 100, 10 );
 		messageOne.createdAt = 0;
-		const messageTwo = generateMessage( 'Initial message 2', 'assistant', 1, 'chat-id', 11 );
+		const messageTwo = generateMessage( 'Initial message 2', 'assistant', 1, 100, 11 );
 		messageTwo.createdAt = 0;
 		store.dispatch(
 			chatActions.setMessages( {

src/custom-package-definitions.d.ts

@@ -46,38 +46,38 @@ declare module '@timfish/forge-externals-plugin' {
 declare module 'wpcom' {
 	class Request {
 		/* eslint-disable @typescript-eslint/no-explicit-any */
-		get< TResponse = any >( params: object | string, query?: object ): Promise< TResponse >;
-		get< TResponse = any >(
+		get< TResponse = unknown >( params: object | string, query?: object ): Promise< TResponse >;
+		get< TResponse = unknown >(
 			params: object | string,
 			callback?: ( error: Error, data: TResponse, headers: Record< string, string > ) => void
 		);
-		get< TResponse = any >(
+		get< TResponse = unknown >(
 			params: object | string,
 			query?: object,
 			callback?: ( error: Error, data: TResponse, headers: Record< string, string > ) => void
 		);
-		post< TResponse = any >(
+		post< TResponse = unknown >(
 			params: object | string,
 			callback?: ( error: Error, data: TResponse, headers: Record< string, string > ) => void
 		);
-		post< TResponse = any >(
+		post< TResponse = unknown >(
 			params: object | string,
 			query?: object,
 			body?: object
 		): Promise< TResponse >;
-		post< TResponse = any >(
+		post< TResponse = unknown >(
 			params: object | string,
 			query?: object,
 			body?: object,
 			callback?: ( error: Error, data: TResponse, headers: Record< string, string > ) => void
 		);
-		del< TResponse = any >( params: object | string, query?: object ): Promise< TResponse >;
-		del< TResponse = any >(
+		del< TResponse = unknown >( params: object | string, query?: object ): Promise< TResponse >;
+		del< TResponse = unknown >(
 			params: object | string,
 			query?: object,
 			callback?: ( error: Error, data: TResponse, headers: Record< string, string > ) => void
 		);
-		del< TResponse = any >(
+		del< TResponse = unknown >(
 			params: object | string,
 			callback?: ( error: Error, data: TResponse, headers: Record< string, string > ) => void
 		);

src/hooks/tests/use-prompt-usage.test.tsx

@@ -1,8 +1,8 @@
-import { renderHook, act, waitFor } from '@testing-library/react';
+import { renderHook, waitFor } from '@testing-library/react';
 import { Provider } from 'react-redux';
 import { useAuth } from 'src/hooks/use-auth';
 import { usePromptUsage, PromptUsageProvider } from 'src/hooks/use-prompt-usage';
-import { store } from 'src/stores';
+import { store, useRootSelector } from 'src/stores';
 import type { ReactNode } from 'react';
 
 jest.mock( 'src/hooks/use-auth', () => ( {
@@ -13,6 +13,11 @@ jest.mock( 'src/hooks/use-feature-flags', () => ( {
 	useFeatureFlags: jest.fn(),
 } ) );
 
+jest.mock( 'src/stores', () => ( {
+	...jest.requireActual( 'src/stores' ),
+	useRootSelector: jest.fn().mockReturnValue( {} ),
+} ) );
+
 function TestWrapper( { children }: { children: ReactNode } ) {
 	return (
 		<Provider store={ store }>
@@ -71,12 +76,13 @@ describe( 'usePromptUsage hook', () => {
 	} );
 
 	it( 'should update prompt usage', async () => {
-		const { result } = renderHook( () => usePromptUsage(), {
-			wrapper: TestWrapper,
+		jest.mocked( useRootSelector ).mockReturnValueOnce( {
+			maxQuota: '300',
+			remainingQuota: '50',
 		} );
 
-		act( () => {
-			result.current.updatePromptUsage( { maxQuota: '300', remainingQuota: '50' } );
+		const { result } = renderHook( () => usePromptUsage(), {
+			wrapper: TestWrapper,
 		} );
 
 		expect( result.current.promptLimit ).toBe( 300 );

src/hooks/use-feature-flags.tsx

@@ -1,7 +1,18 @@
+import * as Sentry from '@sentry/react';
 import React, { createContext, useContext, ReactNode, useState, useEffect } from 'react';
+import { z } from 'zod';
 import { useAuth } from 'src/hooks/use-auth';
 import { getAppGlobals } from 'src/lib/app-globals';
 
+// In PHP, empty associative arrays are encoded as regular arrays when converted to JSON.
+// This means an empty feature flags response comes as [] instead of {}.
+const featureFlagsSchema = z
+	.object( {
+		terminal_wp_cli_enabled: z.boolean().optional(),
+		quick_deploys_enabled: z.boolean().optional(),
+	} )
+	.catch( ( _ ) => ( {} ) );
+
 export interface FeatureFlagsContextType {
 	terminalWpCliEnabled: boolean;
 	quickDeploysEnabled: boolean;
@@ -32,20 +43,22 @@ export const FeatureFlagsProvider: React.FC< FeatureFlagsProviderProps > = ( { c
 				return;
 			}
 			try {
-				const flags = await client.req.get( {
+				const response = await client.req.get( {
 					path: '/studio-app/feature-flags',
 					apiNamespace: 'wpcom/v2',
 				} );
+				const flags = featureFlagsSchema.parse( response );
 				if ( cancel ) {
 					return;
 				}
 				setFeatureFlags( {
 					terminalWpCliEnabled:
-						Boolean( flags?.[ 'terminal_wp_cli_enabled' ] ) || terminalWpCliEnabledFromGlobals,
+						Boolean( flags.terminal_wp_cli_enabled ) || terminalWpCliEnabledFromGlobals,
 					quickDeploysEnabled:
-						Boolean( flags?.[ 'quick_deploys_enabled' ] ) || quickDeploysEnabledFromGlobals,
+						Boolean( flags.quick_deploys_enabled ) || quickDeploysEnabledFromGlobals,
 				} );
 			} catch ( error ) {
+				Sentry.captureException( error );
 				console.error( error );
 			}
 		}

src/hooks/use-prompt-usage.tsx

@@ -1,14 +1,20 @@
 import * as Sentry from '@sentry/electron/renderer';
 import { useState, useEffect, useCallback, useMemo, createContext, useContext } from 'react';
+import { z } from 'zod';
 import { LIMIT_OF_PROMPTS_PER_USER } from 'src/constants';
 import { useAuth } from 'src/hooks/use-auth';
 import { useRootSelector } from 'src/stores';
 
+const promptUsageSchema = z.object( {
+	max_quota: z.number(),
+	remaining_quota: z.number(),
+	quota_reset_date: z.string(),
+} );
+
 type PromptUsage = {
 	promptLimit: number;
 	promptCount: number;
 	fetchPromptUsage: () => Promise< void >;
-	updatePromptUsage: ( data: { maxQuota: string; remainingQuota: string } ) => void;
 	userCanSendMessage: boolean;
 	daysUntilReset: number;
 };
@@ -17,7 +23,6 @@ const initState: PromptUsage = {
 	promptLimit: LIMIT_OF_PROMPTS_PER_USER,
 	promptCount: 0,
 	fetchPromptUsage: async () => undefined,
-	updatePromptUsage: ( _data: { maxQuota: string; remainingQuota: string } ) => undefined,
 	userCanSendMessage: true,
 	daysUntilReset: 0,
 };
@@ -43,21 +48,24 @@ const calculateDaysRemaining = ( quotaResetDate: string ): number => {
 export function PromptUsageProvider( { children }: PromptUsageProps ) {
 	const { Provider } = promptUsageContext;
 
-	const promptUsageDict = useRootSelector( ( state ) => state.chat.promptUsageDict );
+	const promptUsage = useRootSelector( ( state ) => state.chat.promptUsage );
 	const [ promptLimit, setPromptLimit ] = useState( LIMIT_OF_PROMPTS_PER_USER );
 	const [ promptCount, setPromptCount ] = useState( 0 );
 	const [ quotaResetDate, setQuotaResetDate ] = useState( '' );
 	const { client } = useAuth();
 
-	const updatePromptUsage = useCallback( ( data: { maxQuota: string; remainingQuota: string } ) => {
-		const limit = parseInt( data.maxQuota as string );
-		const remaining = parseInt( data.remainingQuota as string );
-		if ( isNaN( limit ) || isNaN( remaining ) ) {
-			return;
-		}
-		setPromptLimit( limit );
-		setPromptCount( limit - remaining );
-	}, [] );
+	const updatePromptUsage = useCallback(
+		( maxQuota: string | number, remainingQuota: string | number ) => {
+			const limit = parseInt( maxQuota as string );
+			const remaining = parseInt( remainingQuota as string );
+			if ( isNaN( limit ) || isNaN( remaining ) ) {
+				return;
+			}
+			setPromptLimit( limit );
+			setPromptCount( limit - remaining );
+		},
+		[]
+	);
 
 	const fetchPromptUsage = useCallback( async () => {
 		if ( ! client?.req ) {
@@ -68,11 +76,9 @@ export function PromptUsageProvider( { children }: PromptUsageProps ) {
 				path: '/studio-app/ai-assistant/quota',
 				apiNamespace: 'wpcom/v2',
 			} );
-			updatePromptUsage( {
-				maxQuota: response.max_quota ?? '',
-				remainingQuota: response.remaining_quota ?? '',
-			} );
-			setQuotaResetDate( response.quota_reset_date || '' );
+			const data = promptUsageSchema.parse( response );
+			updatePromptUsage( data.max_quota, data.remaining_quota );
+			setQuotaResetDate( data.quota_reset_date );
 		} catch ( error ) {
 			Sentry.captureException( error );
 			console.error( error );
@@ -87,12 +93,10 @@ export function PromptUsageProvider( { children }: PromptUsageProps ) {
 	}, [ fetchPromptUsage, client ] );
 
 	useEffect( () => {
-		if ( promptUsageDict ) {
-			for ( const siteId in promptUsageDict ) {
-				updatePromptUsage( promptUsageDict[ siteId ] );
-			}
+		if ( promptUsage.maxQuota && promptUsage.remainingQuota ) {
+			updatePromptUsage( promptUsage.maxQuota, promptUsage.remainingQuota );
 		}
-	}, [ promptUsageDict, updatePromptUsage ] );
+	}, [ promptUsage, updatePromptUsage ] );
 
 	const daysUntilReset = useMemo(
 		() => calculateDaysRemaining( quotaResetDate ),
@@ -104,11 +108,10 @@ export function PromptUsageProvider( { children }: PromptUsageProps ) {
 			fetchPromptUsage,
 			promptLimit,
 			promptCount,
-			updatePromptUsage,
 			userCanSendMessage: promptCount < promptLimit,
 			daysUntilReset,
 		};
-	}, [ fetchPromptUsage, promptLimit, promptCount, updatePromptUsage, daysUntilReset ] );
+	}, [ fetchPromptUsage, promptLimit, promptCount, daysUntilReset ] );
 
 	return <Provider value={ contextValue }>{ children }</Provider>;
 }

src/lib/oauth.ts

@@ -13,7 +13,13 @@ const authTokenSchema = z.object( {
 	expirationTime: z.number(),
 	id: z.number(),
 	email: z.string(),
-	displayName: z.string().optional(),
+	displayName: z.string().default( '' ),
+} );
+
+const meResponseSchema = z.object( {
+	ID: z.number(),
+	email: z.string(),
+	display_name: z.string(),
 } );
 
 export type StoredToken = z.infer< typeof authTokenSchema >;
@@ -81,7 +87,8 @@ async function handleAuthCallback( hash: string ): Promise< StoredToken > {
 		throw new Error( 'Error while getting token' );
 	}
 
-	const response = await new wpcom( accessToken ).req.get( '/me?fields=ID,email,display_name' );
+	const rawResponse = await new wpcom( accessToken ).req.get( '/me?fields=ID,email,display_name' );
+	const response = meResponseSchema.parse( rawResponse );
 
 	return authTokenSchema.parse( {
 		expiresIn,