fix: quizMe exercise selection for missing input (#600)

* Fixes EPICSHOP-EP: guard quizMe exercise input

Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>

* Fixes EPICSHOP-EP: tighten quizMe test typing

Co-authored-by: Kent C. Dodds <me+github@kentcdodds.com>

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

Author: kentcdodds

packages/workshop-mcp/src/prompts.ts

@@ -1,3 +1,4 @@
+import { invariant } from '@epic-web/invariant'
 import { getExercises } from '@epic-web/workshop-utils/apps.server'
 import { getWorkshopConfig } from '@epic-web/workshop-utils/config.server'
 import { type McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
@@ -16,6 +17,7 @@ export const quizMeInputSchema: PromptInputSchema = {
 	workshopDirectory: workshopDirectoryInputSchema,
 	exerciseNumber: z
 		.string()
+		.trim()
 		.optional()
 		.describe(
 			'Exercise number to quiz on (e.g., "4"). Omit for a random exercise.',
@@ -31,8 +33,15 @@ export async function quizMe({
 }): Promise<GetPromptResult> {
 	const workshopRoot = await handleWorkshopDirectory(workshopDirectory)
 	const config = getWorkshopConfig()
-	let exerciseNumber = Number(providedExerciseNumber)
-	if (!providedExerciseNumber) {
+	let exerciseNumber: number | undefined
+	if (providedExerciseNumber) {
+		exerciseNumber = Number(providedExerciseNumber)
+		invariant(
+			Number.isFinite(exerciseNumber),
+			`Exercise number must be a number, received "${providedExerciseNumber}".`,
+		)
+	}
+	if (exerciseNumber === undefined) {
 		const exercises = await getExercises()
 		const randomExercise =
 			exercises[Math.floor(Math.random() * exercises.length)]

packages/workshop-mcp/src/quiz-me.test.ts

@@ -0,0 +1,70 @@
+import { expect, test, vi } from 'vitest'
+import { quizMe } from './prompts.ts'
+import { exerciseContextResource } from './resources.ts'
+
+vi.mock('@epic-web/workshop-utils/apps.server', () => ({
+	getExercises: vi.fn(async () => [{ exerciseNumber: 3 }]),
+}))
+
+vi.mock('@epic-web/workshop-utils/config.server', () => ({
+	getWorkshopConfig: vi.fn(() => ({
+		title: 'Test Workshop',
+		subtitle: 'Testing Sub',
+	})),
+}))
+
+vi.mock('./resources.ts', () => ({
+	exerciseContextResource: {
+		getResource: vi.fn(async ({ exerciseNumber }: { exerciseNumber: number }) => ({
+			uri: `epicshop://exercise/${exerciseNumber}`,
+			mimeType: 'text/plain',
+			text: `exercise ${exerciseNumber}`,
+		})),
+	},
+}))
+
+vi.mock('./utils.ts', async () => {
+	const actual = await vi.importActual<typeof import('./utils.ts')>('./utils.ts')
+	return {
+		...actual,
+		handleWorkshopDirectory: vi.fn(async (workshopDirectory: string) => {
+			return workshopDirectory
+		}),
+	}
+})
+
+test('quizMe chooses a random exercise when none provided (aha)', async () => {
+	const resultPromise = quizMe({ workshopDirectory: '/workshop' })
+
+	await expect(resultPromise).resolves.toMatchObject({
+		messages: [
+			{ role: 'user', content: { type: 'text' } },
+			{ role: 'user', content: { type: 'resource' } },
+		],
+	})
+
+	const result = await resultPromise
+	const getResource = vi.mocked(exerciseContextResource.getResource)
+
+	expect(getResource).toHaveBeenCalledWith({
+		workshopDirectory: '/workshop',
+		exerciseNumber: 3,
+	})
+	const firstMessage = result.messages[0]
+	expect(firstMessage?.content.type).toBe('text')
+	if (!firstMessage || firstMessage.content.type !== 'text') {
+		throw new Error('Expected first message to be text')
+	}
+	expect(firstMessage.content.text).toContain('exercise 3')
+})
+
+test('quizMe rejects non-numeric exercise numbers (aha)', async () => {
+	const resultPromise = quizMe({
+		workshopDirectory: '/workshop',
+		exerciseNumber: 'not-a-number',
+	})
+
+	await expect(resultPromise).rejects.toThrow(
+		'Exercise number must be a number',
+	)
+})