add sampling test

Author: kentcdodds

exercises/06.sampling/02.solution.advanced/src/index.test.ts

@@ -1,15 +1,27 @@
 import { invariant } from '@epic-web/invariant'
 import { Client } from '@modelcontextprotocol/sdk/client/index.js'
 import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js'
+import {
+	CreateMessageRequestSchema,
+	type CreateMessageResult,
+} from '@modelcontextprotocol/sdk/types.js'
 import { test, beforeAll, afterAll, expect } from 'vitest'
+import { type z } from 'zod'
 
 let client: Client
 
 beforeAll(async () => {
-	client = new Client({
-		name: 'EpicMeTester',
-		version: '1.0.0',
-	})
+	client = new Client(
+		{
+			name: 'EpicMeTester',
+			version: '1.0.0',
+		},
+		{
+			capabilities: {
+				sampling: {},
+			},
+		},
+	)
 	const transport = new StdioClientTransport({
 		command: 'tsx',
 		args: ['src/index.ts'],
@@ -69,3 +81,84 @@ test('Tool Call', async () => {
 		}),
 	)
 })
+
+async function deferred<ResolvedValue>() {
+	const ref = {} as {
+		promise: Promise<ResolvedValue>
+		resolve: (value: ResolvedValue) => void
+		reject: (reason?: any) => void
+		value: ResolvedValue | undefined
+		reason: any | undefined
+	}
+	ref.promise = new Promise<ResolvedValue>((resolve, reject) => {
+		ref.resolve = (value) => {
+			ref.value = value
+			resolve(value)
+		}
+		ref.reject = (reason) => {
+			ref.reason = reason
+			reject(reason)
+		}
+	})
+
+	return ref
+}
+
+test('Sampling', async () => {
+	const messageResultDeferred = await deferred<CreateMessageResult>()
+	const messageRequestDeferred =
+		await deferred<z.infer<typeof CreateMessageRequestSchema>>()
+
+	client.setRequestHandler(CreateMessageRequestSchema, (r) => {
+		messageRequestDeferred.resolve(r)
+		return messageResultDeferred.promise
+	})
+
+	await client.callTool({
+		name: 'create_entry',
+		arguments: {
+			title: 'Test Entry',
+			content: 'This is a test entry',
+		},
+	})
+	const request = await messageRequestDeferred.promise
+
+	expect(request).toEqual(
+		expect.objectContaining({
+			method: 'sampling/createMessage',
+			params: expect.objectContaining({
+				maxTokens: expect.any(Number),
+				systemPrompt: expect.stringMatching(/example/i),
+				messages: expect.arrayContaining([
+					expect.objectContaining({
+						role: 'user',
+						content: expect.objectContaining({
+							type: 'text',
+							text: expect.stringMatching(/entry/i),
+							mimeType: 'application/json',
+						}),
+					}),
+				]),
+			}),
+		}),
+	)
+
+	messageResultDeferred.resolve({
+		model: 'stub-model',
+		stopReason: 'endTurn',
+		role: 'assistant',
+		content: {
+			type: 'text',
+			text: JSON.stringify([
+				{ id: 1 },
+				{
+					name: 'Testing Sampling',
+					description: 'Used when testing sampling. Hope it works',
+				},
+			]),
+		},
+	})
+
+	// give the client a chance to process the result
+	await new Promise((resolve) => setTimeout(resolve, 100))
+})