[v1.x] feat(tasks): add streaming methods for elicitation and sampling (#1528)
Co-authored-by: Konstantin Konstantinov <KKonstantinov@users.noreply.github.com>
This commit is contained in:
Luca Chang
GitHub
parent
97ab379e45
commit
5c16ae3339
4 changed files with 944 additions and 22 deletions
|
|
@ -15,6 +15,7 @@ import {
|
|||
LoggingMessageNotificationSchema,
|
||||
ResourceListChangedNotificationSchema,
|
||||
ElicitRequestSchema,
|
||||
ElicitResult,
|
||||
ResourceLink,
|
||||
ReadResourceRequest,
|
||||
ReadResourceResultSchema,
|
||||
|
|
@ -22,6 +23,7 @@ import {
|
|||
ErrorCode,
|
||||
McpError
|
||||
} from '../../types.js';
|
||||
import { InMemoryTaskStore } from '../../experimental/tasks/stores/in-memory.js';
|
||||
import { getDisplayName } from '../../shared/metadataUtils.js';
|
||||
import { Ajv } from 'ajv';
|
||||
|
||||
|
|
@ -65,6 +67,7 @@ function printHelp(): void {
|
|||
console.log(' greet [name] - Call the greet tool');
|
||||
console.log(' multi-greet [name] - Call the multi-greet tool with notifications');
|
||||
console.log(' collect-info [type] - Test form elicitation with collect-user-info tool (contact/preferences/feedback)');
|
||||
console.log(' collect-info-task [type] - Test bidirectional task support (server+client tasks) with elicitation');
|
||||
console.log(' start-notifications [interval] [count] - Start periodic notifications');
|
||||
console.log(' run-notifications-tool-with-resumability [interval] [count] - Run notification tool with resumability');
|
||||
console.log(' list-prompts - List available prompts');
|
||||
|
|
@ -131,6 +134,11 @@ function commandLoop(): void {
|
|||
await callCollectInfoTool(args[1] || 'contact');
|
||||
break;
|
||||
|
||||
case 'collect-info-task': {
|
||||
await callCollectInfoWithTask(args[1] || 'contact');
|
||||
break;
|
||||
}
|
||||
|
||||
case 'start-notifications': {
|
||||
const interval = args[1] ? parseInt(args[1], 10) : 2000;
|
||||
const count = args[2] ? parseInt(args[2], 10) : 10;
|
||||
|
|
@ -232,7 +240,10 @@ async function connect(url?: string): Promise<void> {
|
|||
console.log(`Connecting to ${serverUrl}...`);
|
||||
|
||||
try {
|
||||
// Create a new client with form elicitation capability
|
||||
// Create task store for client-side task support
|
||||
const clientTaskStore = new InMemoryTaskStore();
|
||||
|
||||
// Create a new client with form elicitation capability and task support
|
||||
client = new Client(
|
||||
{
|
||||
name: 'example-client',
|
||||
|
|
@ -242,25 +253,46 @@ async function connect(url?: string): Promise<void> {
|
|||
capabilities: {
|
||||
elicitation: {
|
||||
form: {}
|
||||
},
|
||||
tasks: {
|
||||
requests: {
|
||||
elicitation: {
|
||||
create: {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
taskStore: clientTaskStore
|
||||
}
|
||||
);
|
||||
client.onerror = error => {
|
||||
console.error('\x1b[31mClient error:', error, '\x1b[0m');
|
||||
};
|
||||
|
||||
// Set up elicitation request handler with proper validation
|
||||
client.setRequestHandler(ElicitRequestSchema, async request => {
|
||||
// Set up elicitation request handler with proper validation and task support
|
||||
client.setRequestHandler(ElicitRequestSchema, async (request, extra) => {
|
||||
if (request.params.mode !== 'form') {
|
||||
throw new McpError(ErrorCode.InvalidParams, `Unsupported elicitation mode: ${request.params.mode}`);
|
||||
}
|
||||
console.log('\n🔔 Elicitation (form) Request Received:');
|
||||
console.log(`Message: ${request.params.message}`);
|
||||
console.log(`Related Task: ${request.params._meta?.[RELATED_TASK_META_KEY]?.taskId}`);
|
||||
console.log(`Task Creation Requested: ${request.params.task ? 'yes' : 'no'}`);
|
||||
console.log('Requested Schema:');
|
||||
console.log(JSON.stringify(request.params.requestedSchema, null, 2));
|
||||
|
||||
// Helper to return result, optionally creating a task if requested
|
||||
const returnResult = async (result: ElicitResult) => {
|
||||
if (request.params.task && extra.taskStore) {
|
||||
// Create a task and store the result
|
||||
const task = await extra.taskStore.createTask({ ttl: extra.taskRequestedTtl });
|
||||
await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
|
||||
console.log(`📋 Created client-side task: ${task.taskId}`);
|
||||
return { task };
|
||||
}
|
||||
return result;
|
||||
};
|
||||
|
||||
const schema = request.params.requestedSchema;
|
||||
const properties = schema.properties;
|
||||
const required = schema.required || [];
|
||||
|
|
@ -381,7 +413,7 @@ async function connect(url?: string): Promise<void> {
|
|||
}
|
||||
|
||||
if (inputCancelled) {
|
||||
return { action: 'cancel' };
|
||||
return returnResult({ action: 'cancel' });
|
||||
}
|
||||
|
||||
// If we didn't complete all fields due to an error, try again
|
||||
|
|
@ -394,7 +426,7 @@ async function connect(url?: string): Promise<void> {
|
|||
continue;
|
||||
} else {
|
||||
console.log('Maximum attempts reached. Declining request.');
|
||||
return { action: 'decline' };
|
||||
return returnResult({ action: 'decline' });
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -412,7 +444,7 @@ async function connect(url?: string): Promise<void> {
|
|||
continue;
|
||||
} else {
|
||||
console.log('Maximum attempts reached. Declining request.');
|
||||
return { action: 'decline' };
|
||||
return returnResult({ action: 'decline' });
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -426,25 +458,34 @@ async function connect(url?: string): Promise<void> {
|
|||
});
|
||||
});
|
||||
|
||||
if (confirmAnswer === 'yes' || confirmAnswer === 'y') {
|
||||
return {
|
||||
action: 'accept',
|
||||
content
|
||||
};
|
||||
} else if (confirmAnswer === 'cancel' || confirmAnswer === 'c') {
|
||||
return { action: 'cancel' };
|
||||
} else if (confirmAnswer === 'no' || confirmAnswer === 'n') {
|
||||
if (attempts < maxAttempts) {
|
||||
console.log('Please re-enter the information...');
|
||||
continue;
|
||||
} else {
|
||||
return { action: 'decline' };
|
||||
switch (confirmAnswer) {
|
||||
case 'yes':
|
||||
case 'y': {
|
||||
return returnResult({
|
||||
action: 'accept',
|
||||
content: content as ElicitResult['content']
|
||||
});
|
||||
}
|
||||
case 'cancel':
|
||||
case 'c': {
|
||||
return returnResult({ action: 'cancel' });
|
||||
}
|
||||
case 'no':
|
||||
case 'n': {
|
||||
if (attempts < maxAttempts) {
|
||||
console.log('Please re-enter the information...');
|
||||
continue;
|
||||
} else {
|
||||
return returnResult({ action: 'decline' });
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
console.log('Maximum attempts reached. Declining request.');
|
||||
return { action: 'decline' };
|
||||
return returnResult({ action: 'decline' });
|
||||
});
|
||||
|
||||
transport = new StreamableHTTPClientTransport(new URL(serverUrl), {
|
||||
|
|
@ -641,6 +682,12 @@ async function callCollectInfoTool(infoType: string): Promise<void> {
|
|||
await callTool('collect-user-info', { infoType });
|
||||
}
|
||||
|
||||
async function callCollectInfoWithTask(infoType: string): Promise<void> {
|
||||
console.log(`\n🔄 Testing bidirectional task support with collect-user-info-task tool (${infoType})...`);
|
||||
console.log('This will create a task on the server, which will elicit input and create a task on the client.\n');
|
||||
await callToolTask('collect-user-info-task', { infoType });
|
||||
}
|
||||
|
||||
async function startNotifications(interval: number, count: number): Promise<void> {
|
||||
console.log(`Starting notification stream: interval=${interval}ms, count=${count || 'unlimited'}`);
|
||||
await callTool('start-notification-stream', { interval, count });
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import { requireBearerAuth } from '../../server/auth/middleware/bearerAuth.js';
|
|||
import { createMcpExpressApp } from '../../server/express.js';
|
||||
import {
|
||||
CallToolResult,
|
||||
ElicitResult,
|
||||
ElicitResultSchema,
|
||||
GetPromptResult,
|
||||
isInitializeRequest,
|
||||
|
|
@ -280,6 +281,114 @@ const getServer = () => {
|
|||
}
|
||||
);
|
||||
|
||||
// Register a tool that demonstrates bidirectional task support:
|
||||
// Server creates a task, then elicits input from client using elicitInputStream
|
||||
// Using the experimental tasks API - WARNING: may change without notice
|
||||
server.experimental.tasks.registerToolTask(
|
||||
'collect-user-info-task',
|
||||
{
|
||||
title: 'Collect Info with Task',
|
||||
description: 'Collects user info via elicitation with task support using elicitInputStream',
|
||||
inputSchema: {
|
||||
infoType: z.enum(['contact', 'preferences']).describe('Type of information to collect').default('contact')
|
||||
}
|
||||
},
|
||||
{
|
||||
async createTask({ infoType }, { taskStore: createTaskStore, taskRequestedTtl }) {
|
||||
// Create the server-side task
|
||||
const task = await createTaskStore.createTask({
|
||||
ttl: taskRequestedTtl
|
||||
});
|
||||
|
||||
// Perform async work that makes a nested elicitation request using elicitInputStream
|
||||
(async () => {
|
||||
try {
|
||||
const message = infoType === 'contact' ? 'Please provide your contact information' : 'Please set your preferences';
|
||||
|
||||
// Define schemas with proper typing for PrimitiveSchemaDefinition
|
||||
const contactSchema: {
|
||||
type: 'object';
|
||||
properties: Record<string, PrimitiveSchemaDefinition>;
|
||||
required: string[];
|
||||
} = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
name: { type: 'string', title: 'Full Name', description: 'Your full name' },
|
||||
email: { type: 'string', title: 'Email', description: 'Your email address' }
|
||||
},
|
||||
required: ['name', 'email']
|
||||
};
|
||||
|
||||
const preferencesSchema: {
|
||||
type: 'object';
|
||||
properties: Record<string, PrimitiveSchemaDefinition>;
|
||||
required: string[];
|
||||
} = {
|
||||
type: 'object',
|
||||
properties: {
|
||||
theme: { type: 'string', title: 'Theme', enum: ['light', 'dark', 'auto'] },
|
||||
notifications: { type: 'boolean', title: 'Enable Notifications', default: true }
|
||||
},
|
||||
required: ['theme']
|
||||
};
|
||||
|
||||
const requestedSchema = infoType === 'contact' ? contactSchema : preferencesSchema;
|
||||
|
||||
// Use elicitInputStream to elicit input from client
|
||||
// This demonstrates the streaming elicitation API
|
||||
// Access via server.server to get the underlying Server instance
|
||||
const stream = server.server.experimental.tasks.elicitInputStream({
|
||||
mode: 'form',
|
||||
message,
|
||||
requestedSchema
|
||||
});
|
||||
|
||||
let elicitResult: ElicitResult | undefined;
|
||||
for await (const msg of stream) {
|
||||
if (msg.type === 'result') {
|
||||
elicitResult = msg.result as ElicitResult;
|
||||
} else if (msg.type === 'error') {
|
||||
throw msg.error;
|
||||
}
|
||||
}
|
||||
|
||||
if (!elicitResult) {
|
||||
throw new Error('No result received from elicitation');
|
||||
}
|
||||
|
||||
let resultText: string;
|
||||
if (elicitResult.action === 'accept') {
|
||||
resultText = `Collected ${infoType} info: ${JSON.stringify(elicitResult.content, null, 2)}`;
|
||||
} else if (elicitResult.action === 'decline') {
|
||||
resultText = `User declined to provide ${infoType} information`;
|
||||
} else {
|
||||
resultText = 'User cancelled the request';
|
||||
}
|
||||
|
||||
await taskStore.storeTaskResult(task.taskId, 'completed', {
|
||||
content: [{ type: 'text', text: resultText }]
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Error in collect-user-info-task:', error);
|
||||
await taskStore.storeTaskResult(task.taskId, 'failed', {
|
||||
content: [{ type: 'text', text: `Error: ${error}` }],
|
||||
isError: true
|
||||
});
|
||||
}
|
||||
})();
|
||||
|
||||
return { task };
|
||||
},
|
||||
async getTask(_args, { taskId, taskStore: getTaskStore }) {
|
||||
return await getTaskStore.getTask(taskId);
|
||||
},
|
||||
async getTaskResult(_args, { taskId, taskStore: getResultTaskStore }) {
|
||||
const result = await getResultTaskStore.getTaskResult(taskId);
|
||||
return result as CallToolResult;
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
// Register a simple prompt with title
|
||||
server.registerPrompt(
|
||||
'greeting-template',
|
||||
|
|
|
|||
|
|
@ -9,7 +9,21 @@ import type { Server } from '../../server/index.js';
|
|||
import type { RequestOptions } from '../../shared/protocol.js';
|
||||
import type { ResponseMessage } from '../../shared/responseMessage.js';
|
||||
import type { AnySchema, SchemaOutput } from '../../server/zod-compat.js';
|
||||
import type { ServerRequest, Notification, Request, Result, GetTaskResult, ListTasksResult, CancelTaskResult } from '../../types.js';
|
||||
import type {
|
||||
ServerRequest,
|
||||
Notification,
|
||||
Request,
|
||||
Result,
|
||||
GetTaskResult,
|
||||
ListTasksResult,
|
||||
CancelTaskResult,
|
||||
CreateMessageRequestParams,
|
||||
CreateMessageResult,
|
||||
ElicitRequestFormParams,
|
||||
ElicitRequestURLParams,
|
||||
ElicitResult
|
||||
} from '../../types.js';
|
||||
import { CreateMessageResultSchema, ElicitResultSchema } from '../../types.js';
|
||||
|
||||
/**
|
||||
* Experimental task features for low-level MCP servers.
|
||||
|
|
@ -60,6 +74,188 @@ export class ExperimentalServerTasks<
|
|||
return (this._server as unknown as ServerWithRequestStream).requestStream(request, resultSchema, options);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends a sampling request and returns an AsyncGenerator that yields response messages.
|
||||
* The generator is guaranteed to end with either a 'result' or 'error' message.
|
||||
*
|
||||
* For task-augmented requests, yields 'taskCreated' and 'taskStatus' messages
|
||||
* before the final result.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* const stream = server.experimental.tasks.createMessageStream({
|
||||
* messages: [{ role: 'user', content: { type: 'text', text: 'Hello' } }],
|
||||
* maxTokens: 100
|
||||
* }, {
|
||||
* onprogress: (progress) => {
|
||||
* // Handle streaming tokens via progress notifications
|
||||
* console.log('Progress:', progress.message);
|
||||
* }
|
||||
* });
|
||||
*
|
||||
* for await (const message of stream) {
|
||||
* switch (message.type) {
|
||||
* case 'taskCreated':
|
||||
* console.log('Task created:', message.task.taskId);
|
||||
* break;
|
||||
* case 'taskStatus':
|
||||
* console.log('Task status:', message.task.status);
|
||||
* break;
|
||||
* case 'result':
|
||||
* console.log('Final result:', message.result);
|
||||
* break;
|
||||
* case 'error':
|
||||
* console.error('Error:', message.error);
|
||||
* break;
|
||||
* }
|
||||
* }
|
||||
* ```
|
||||
*
|
||||
* @param params - The sampling request parameters
|
||||
* @param options - Optional request options (timeout, signal, task creation params, onprogress, etc.)
|
||||
* @returns AsyncGenerator that yields ResponseMessage objects
|
||||
*
|
||||
* @experimental
|
||||
*/
|
||||
createMessageStream(
|
||||
params: CreateMessageRequestParams,
|
||||
options?: RequestOptions
|
||||
): AsyncGenerator<ResponseMessage<CreateMessageResult>, void, void> {
|
||||
// Access client capabilities via the server
|
||||
const clientCapabilities = this._server.getClientCapabilities();
|
||||
|
||||
// Capability check - only required when tools/toolChoice are provided
|
||||
if ((params.tools || params.toolChoice) && !clientCapabilities?.sampling?.tools) {
|
||||
throw new Error('Client does not support sampling tools capability.');
|
||||
}
|
||||
|
||||
// Message structure validation - always validate tool_use/tool_result pairs.
|
||||
// These may appear even without tools/toolChoice in the current request when
|
||||
// a previous sampling request returned tool_use and this is a follow-up with results.
|
||||
if (params.messages.length > 0) {
|
||||
const lastMessage = params.messages[params.messages.length - 1];
|
||||
const lastContent = Array.isArray(lastMessage.content) ? lastMessage.content : [lastMessage.content];
|
||||
const hasToolResults = lastContent.some(c => c.type === 'tool_result');
|
||||
|
||||
const previousMessage = params.messages.length > 1 ? params.messages[params.messages.length - 2] : undefined;
|
||||
const previousContent = previousMessage
|
||||
? Array.isArray(previousMessage.content)
|
||||
? previousMessage.content
|
||||
: [previousMessage.content]
|
||||
: [];
|
||||
const hasPreviousToolUse = previousContent.some(c => c.type === 'tool_use');
|
||||
|
||||
if (hasToolResults) {
|
||||
if (lastContent.some(c => c.type !== 'tool_result')) {
|
||||
throw new Error('The last message must contain only tool_result content if any is present');
|
||||
}
|
||||
if (!hasPreviousToolUse) {
|
||||
throw new Error('tool_result blocks are not matching any tool_use from the previous message');
|
||||
}
|
||||
}
|
||||
if (hasPreviousToolUse) {
|
||||
// Extract tool_use IDs from previous message and tool_result IDs from current message
|
||||
const toolUseIds = new Set(previousContent.filter(c => c.type === 'tool_use').map(c => (c as { id: string }).id));
|
||||
const toolResultIds = new Set(
|
||||
lastContent.filter(c => c.type === 'tool_result').map(c => (c as { toolUseId: string }).toolUseId)
|
||||
);
|
||||
if (toolUseIds.size !== toolResultIds.size || ![...toolUseIds].every(id => toolResultIds.has(id))) {
|
||||
throw new Error('ids of tool_result blocks and tool_use blocks from previous message do not match');
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return this.requestStream(
|
||||
{
|
||||
method: 'sampling/createMessage',
|
||||
params
|
||||
},
|
||||
CreateMessageResultSchema,
|
||||
options
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends an elicitation request and returns an AsyncGenerator that yields response messages.
|
||||
* The generator is guaranteed to end with either a 'result' or 'error' message.
|
||||
*
|
||||
* For task-augmented requests (especially URL-based elicitation), yields 'taskCreated'
|
||||
* and 'taskStatus' messages before the final result.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* const stream = server.experimental.tasks.elicitInputStream({
|
||||
* mode: 'url',
|
||||
* message: 'Please authenticate',
|
||||
* elicitationId: 'auth-123',
|
||||
* url: 'https://example.com/auth'
|
||||
* }, {
|
||||
* task: { ttl: 300000 } // Task-augmented for long-running auth flow
|
||||
* });
|
||||
*
|
||||
* for await (const message of stream) {
|
||||
* switch (message.type) {
|
||||
* case 'taskCreated':
|
||||
* console.log('Task created:', message.task.taskId);
|
||||
* break;
|
||||
* case 'taskStatus':
|
||||
* console.log('Task status:', message.task.status);
|
||||
* break;
|
||||
* case 'result':
|
||||
* console.log('User action:', message.result.action);
|
||||
* break;
|
||||
* case 'error':
|
||||
* console.error('Error:', message.error);
|
||||
* break;
|
||||
* }
|
||||
* }
|
||||
* ```
|
||||
*
|
||||
* @param params - The elicitation request parameters
|
||||
* @param options - Optional request options (timeout, signal, task creation params, etc.)
|
||||
* @returns AsyncGenerator that yields ResponseMessage objects
|
||||
*
|
||||
* @experimental
|
||||
*/
|
||||
elicitInputStream(
|
||||
params: ElicitRequestFormParams | ElicitRequestURLParams,
|
||||
options?: RequestOptions
|
||||
): AsyncGenerator<ResponseMessage<ElicitResult>, void, void> {
|
||||
// Access client capabilities via the server
|
||||
const clientCapabilities = this._server.getClientCapabilities();
|
||||
const mode = params.mode ?? 'form';
|
||||
|
||||
// Capability check based on mode
|
||||
switch (mode) {
|
||||
case 'url': {
|
||||
if (!clientCapabilities?.elicitation?.url) {
|
||||
throw new Error('Client does not support url elicitation.');
|
||||
}
|
||||
break;
|
||||
}
|
||||
case 'form': {
|
||||
if (!clientCapabilities?.elicitation?.form) {
|
||||
throw new Error('Client does not support form elicitation.');
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize params to ensure mode is set for form mode (defaults to 'form' per spec)
|
||||
const normalizedParams = mode === 'form' && params.mode === undefined ? { ...params, mode: 'form' as const } : params;
|
||||
|
||||
// Cast to ServerRequest needed because TypeScript can't narrow the union type
|
||||
// based on the discriminated 'method' field when constructing the object literal
|
||||
return this.requestStream(
|
||||
{
|
||||
method: 'elicitation/create',
|
||||
params: normalizedParams
|
||||
} as ServerRequest,
|
||||
ElicitResultSchema,
|
||||
options
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the current status of a task.
|
||||
*
|
||||
|
|
|
|||
570
test/experimental/tasks/server-streaming.test.ts
Normal file
570
test/experimental/tasks/server-streaming.test.ts
Normal file
|
|
@ -0,0 +1,570 @@
|
|||
/**
|
||||
* Tests for experimental server streaming methods: createMessageStream and elicitInputStream.
|
||||
* WARNING: These APIs are experimental and may change without notice.
|
||||
*
|
||||
* @experimental
|
||||
*/
|
||||
|
||||
import { describe, test, expect, beforeEach, afterEach, vi } from 'vitest';
|
||||
import { Client } from '../../../src/client/index.js';
|
||||
import { Server } from '../../../src/server/index.js';
|
||||
import { InMemoryTransport } from '../../../src/inMemory.js';
|
||||
import { InMemoryTaskStore } from '../../../src/experimental/tasks/stores/in-memory.js';
|
||||
import { toArrayAsync } from '../../../src/shared/responseMessage.js';
|
||||
import {
|
||||
CreateMessageRequestSchema,
|
||||
ElicitRequestSchema,
|
||||
type CreateMessageResult,
|
||||
type ElicitResult,
|
||||
type Task
|
||||
} from '../../../src/types.js';
|
||||
|
||||
describe('createMessageStream', () => {
|
||||
test('should throw when tools are provided without sampling.tools capability', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' },
|
||||
model: 'test-model'
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
expect(() => {
|
||||
server.experimental.tasks.createMessageStream({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'Hello' } }],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
});
|
||||
}).toThrow('Client does not support sampling tools capability');
|
||||
|
||||
await client.close().catch(() => {});
|
||||
await server.close().catch(() => {});
|
||||
});
|
||||
|
||||
test('should throw when tool_result has no matching tool_use in previous message', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' },
|
||||
model: 'test-model'
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
expect(() => {
|
||||
server.experimental.tasks.createMessageStream({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'Hello' } },
|
||||
{
|
||||
role: 'user',
|
||||
content: [{ type: 'tool_result', toolUseId: 'test-id', content: [{ type: 'text', text: 'result' }] }]
|
||||
}
|
||||
],
|
||||
maxTokens: 100
|
||||
});
|
||||
}).toThrow('tool_result blocks are not matching any tool_use from the previous message');
|
||||
|
||||
await client.close().catch(() => {});
|
||||
await server.close().catch(() => {});
|
||||
});
|
||||
|
||||
describe('terminal message guarantees', () => {
|
||||
test('should yield exactly one terminal message for successful request', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' },
|
||||
model: 'test-model'
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.createMessageStream({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'Hello' } }],
|
||||
maxTokens: 100
|
||||
});
|
||||
|
||||
const allMessages = await toArrayAsync(stream);
|
||||
|
||||
expect(allMessages.length).toBe(1);
|
||||
expect(allMessages[0].type).toBe('result');
|
||||
|
||||
const taskMessages = allMessages.filter(m => m.type === 'taskCreated' || m.type === 'taskStatus');
|
||||
expect(taskMessages.length).toBe(0);
|
||||
|
||||
await client.close().catch(() => {});
|
||||
await server.close().catch(() => {});
|
||||
});
|
||||
|
||||
test('should yield error as terminal message when client returns error', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => {
|
||||
throw new Error('Simulated client error');
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.createMessageStream({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'Hello' } }],
|
||||
maxTokens: 100
|
||||
});
|
||||
|
||||
const allMessages = await toArrayAsync(stream);
|
||||
|
||||
expect(allMessages.length).toBe(1);
|
||||
expect(allMessages[0].type).toBe('error');
|
||||
|
||||
await client.close().catch(() => {});
|
||||
await server.close().catch(() => {});
|
||||
});
|
||||
|
||||
test('should yield exactly one terminal message with result', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant' as const,
|
||||
content: { type: 'text' as const, text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.createMessageStream({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'Message' } }],
|
||||
maxTokens: 100
|
||||
});
|
||||
|
||||
const messages = await toArrayAsync(stream);
|
||||
const terminalMessages = messages.filter(m => m.type === 'result' || m.type === 'error');
|
||||
|
||||
expect(terminalMessages.length).toBe(1);
|
||||
|
||||
const lastMessage = messages[messages.length - 1];
|
||||
expect(lastMessage.type === 'result' || lastMessage.type === 'error').toBe(true);
|
||||
|
||||
if (lastMessage.type === 'result') {
|
||||
expect((lastMessage.result as CreateMessageResult).content).toBeDefined();
|
||||
}
|
||||
|
||||
await client.close().catch(() => {});
|
||||
await server.close().catch(() => {});
|
||||
});
|
||||
});
|
||||
|
||||
describe('non-task request minimality', () => {
|
||||
test('should yield only result message for non-task request', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant' as const,
|
||||
content: { type: 'text' as const, text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.createMessageStream({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'Message' } }],
|
||||
maxTokens: 100
|
||||
});
|
||||
|
||||
const messages = await toArrayAsync(stream);
|
||||
|
||||
const taskMessages = messages.filter(m => m.type === 'taskCreated' || m.type === 'taskStatus');
|
||||
expect(taskMessages.length).toBe(0);
|
||||
|
||||
const resultMessages = messages.filter(m => m.type === 'result');
|
||||
expect(resultMessages.length).toBe(1);
|
||||
|
||||
expect(messages.length).toBe(1);
|
||||
|
||||
await client.close().catch(() => {});
|
||||
await server.close().catch(() => {});
|
||||
});
|
||||
});
|
||||
|
||||
describe('task-augmented request handling', () => {
|
||||
test('should yield taskCreated and result for task-augmented request', async () => {
|
||||
const clientTaskStore = new InMemoryTaskStore();
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
const client = new Client(
|
||||
{ name: 'test client', version: '1.0' },
|
||||
{
|
||||
capabilities: {
|
||||
sampling: {},
|
||||
tasks: {
|
||||
requests: {
|
||||
sampling: { createMessage: {} }
|
||||
}
|
||||
}
|
||||
},
|
||||
taskStore: clientTaskStore
|
||||
}
|
||||
);
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async (request, extra) => {
|
||||
const result = {
|
||||
model: 'test-model',
|
||||
role: 'assistant' as const,
|
||||
content: { type: 'text' as const, text: 'Task response' }
|
||||
};
|
||||
|
||||
if (request.params.task && extra.taskStore) {
|
||||
const task = await extra.taskStore.createTask({ ttl: extra.taskRequestedTtl });
|
||||
await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
|
||||
return { task };
|
||||
}
|
||||
return result;
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.createMessageStream(
|
||||
{
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'Task-augmented message' } }],
|
||||
maxTokens: 100
|
||||
},
|
||||
{ task: { ttl: 60_000 } }
|
||||
);
|
||||
|
||||
const messages = await toArrayAsync(stream);
|
||||
|
||||
// Should have taskCreated and result
|
||||
expect(messages.length).toBeGreaterThanOrEqual(2);
|
||||
|
||||
// First message should be taskCreated
|
||||
expect(messages[0].type).toBe('taskCreated');
|
||||
const taskCreated = messages[0] as { type: 'taskCreated'; task: Task };
|
||||
expect(taskCreated.task.taskId).toBeDefined();
|
||||
|
||||
// Last message should be result
|
||||
const lastMessage = messages[messages.length - 1];
|
||||
expect(lastMessage.type).toBe('result');
|
||||
if (lastMessage.type === 'result') {
|
||||
expect((lastMessage.result as CreateMessageResult).model).toBe('test-model');
|
||||
}
|
||||
|
||||
clientTaskStore.cleanup();
|
||||
await client.close().catch(() => {});
|
||||
await server.close().catch(() => {});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('elicitInputStream', () => {
|
||||
let server: Server;
|
||||
let client: Client;
|
||||
let clientTransport: ReturnType<typeof InMemoryTransport.createLinkedPair>[0];
|
||||
let serverTransport: ReturnType<typeof InMemoryTransport.createLinkedPair>[1];
|
||||
|
||||
beforeEach(async () => {
|
||||
server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
client = new Client(
|
||||
{ name: 'test client', version: '1.0' },
|
||||
{
|
||||
capabilities: {
|
||||
elicitation: {
|
||||
form: {},
|
||||
url: {}
|
||||
}
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
[clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await server.close().catch(() => {});
|
||||
await client.close().catch(() => {});
|
||||
});
|
||||
|
||||
test('should throw when client does not support form elicitation', async () => {
|
||||
// Create client without form elicitation capability
|
||||
const noFormClient = new Client(
|
||||
{ name: 'test client', version: '1.0' },
|
||||
{
|
||||
capabilities: {
|
||||
elicitation: {
|
||||
url: {}
|
||||
}
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
const [noFormClientTransport, noFormServerTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([noFormClient.connect(noFormClientTransport), server.connect(noFormServerTransport)]);
|
||||
|
||||
expect(() => {
|
||||
server.experimental.tasks.elicitInputStream({
|
||||
mode: 'form',
|
||||
message: 'Enter data',
|
||||
requestedSchema: { type: 'object', properties: {} }
|
||||
});
|
||||
}).toThrow('Client does not support form elicitation.');
|
||||
|
||||
await noFormClient.close().catch(() => {});
|
||||
});
|
||||
|
||||
test('should throw when client does not support url elicitation', async () => {
|
||||
// Create client without url elicitation capability
|
||||
const noUrlClient = new Client(
|
||||
{ name: 'test client', version: '1.0' },
|
||||
{
|
||||
capabilities: {
|
||||
elicitation: {
|
||||
form: {}
|
||||
}
|
||||
}
|
||||
}
|
||||
);
|
||||
|
||||
const [noUrlClientTransport, noUrlServerTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([noUrlClient.connect(noUrlClientTransport), server.connect(noUrlServerTransport)]);
|
||||
|
||||
expect(() => {
|
||||
server.experimental.tasks.elicitInputStream({
|
||||
mode: 'url',
|
||||
message: 'Open URL',
|
||||
elicitationId: 'test-123',
|
||||
url: 'https://example.com/auth'
|
||||
});
|
||||
}).toThrow('Client does not support url elicitation.');
|
||||
|
||||
await noUrlClient.close().catch(() => {});
|
||||
});
|
||||
|
||||
test('should default to form mode when mode is not specified', async () => {
|
||||
const requestStreamSpy = vi.spyOn(server.experimental.tasks, 'requestStream');
|
||||
|
||||
client.setRequestHandler(ElicitRequestSchema, () => ({
|
||||
action: 'accept',
|
||||
content: { value: 'test' }
|
||||
}));
|
||||
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// Call without explicit mode - need to cast because TypeScript expects mode
|
||||
const params = {
|
||||
message: 'Enter value',
|
||||
requestedSchema: {
|
||||
type: 'object' as const,
|
||||
properties: { value: { type: 'string' as const } }
|
||||
}
|
||||
};
|
||||
|
||||
const stream = server.experimental.tasks.elicitInputStream(
|
||||
params as Parameters<typeof server.experimental.tasks.elicitInputStream>[0]
|
||||
);
|
||||
await toArrayAsync(stream);
|
||||
|
||||
// Verify mode was normalized to 'form'
|
||||
expect(requestStreamSpy).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
method: 'elicitation/create',
|
||||
params: expect.objectContaining({ mode: 'form' })
|
||||
}),
|
||||
expect.anything(),
|
||||
undefined
|
||||
);
|
||||
});
|
||||
|
||||
test('should yield error as terminal message when client returns error', async () => {
|
||||
client.setRequestHandler(ElicitRequestSchema, () => {
|
||||
throw new Error('Simulated client error');
|
||||
});
|
||||
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.elicitInputStream({
|
||||
mode: 'form',
|
||||
message: 'Enter data',
|
||||
requestedSchema: {
|
||||
type: 'object',
|
||||
properties: { value: { type: 'string' } }
|
||||
}
|
||||
});
|
||||
|
||||
const allMessages = await toArrayAsync(stream);
|
||||
|
||||
expect(allMessages.length).toBe(1);
|
||||
expect(allMessages[0].type).toBe('error');
|
||||
});
|
||||
|
||||
// For any streaming elicitation request, the AsyncGenerator yields exactly one terminal
|
||||
// message (either 'result' or 'error') as its final message.
|
||||
describe('terminal message guarantees', () => {
|
||||
test.each([
|
||||
{ action: 'accept' as const, content: { data: 'test-value' } },
|
||||
{ action: 'decline' as const, content: undefined },
|
||||
{ action: 'cancel' as const, content: undefined }
|
||||
])('should yield exactly one terminal message for action: $action', async ({ action, content }) => {
|
||||
client.setRequestHandler(ElicitRequestSchema, () => ({
|
||||
action,
|
||||
content
|
||||
}));
|
||||
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.elicitInputStream({
|
||||
mode: 'form',
|
||||
message: 'Test message',
|
||||
requestedSchema: {
|
||||
type: 'object',
|
||||
properties: { data: { type: 'string' } }
|
||||
}
|
||||
});
|
||||
|
||||
const messages = await toArrayAsync(stream);
|
||||
|
||||
// Count terminal messages (result or error)
|
||||
const terminalMessages = messages.filter(m => m.type === 'result' || m.type === 'error');
|
||||
|
||||
expect(terminalMessages.length).toBe(1);
|
||||
|
||||
// Verify terminal message is the last message
|
||||
const lastMessage = messages[messages.length - 1];
|
||||
expect(lastMessage.type === 'result' || lastMessage.type === 'error').toBe(true);
|
||||
|
||||
// Verify result content matches expected action
|
||||
if (lastMessage.type === 'result') {
|
||||
expect((lastMessage.result as ElicitResult).action).toBe(action);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// For any non-task elicitation request, the generator yields exactly one 'result' message
|
||||
// (or 'error' if the request fails), with no 'taskCreated' or 'taskStatus' messages.
|
||||
describe('non-task request minimality', () => {
|
||||
test.each([
|
||||
{ action: 'accept' as const, content: { value: 'test' } },
|
||||
{ action: 'decline' as const, content: undefined },
|
||||
{ action: 'cancel' as const, content: undefined }
|
||||
])('should yield only result message for non-task request with action: $action', async ({ action, content }) => {
|
||||
client.setRequestHandler(ElicitRequestSchema, () => ({
|
||||
action,
|
||||
content
|
||||
}));
|
||||
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// Non-task request (no task option)
|
||||
const stream = server.experimental.tasks.elicitInputStream({
|
||||
mode: 'form',
|
||||
message: 'Non-task request',
|
||||
requestedSchema: {
|
||||
type: 'object',
|
||||
properties: { value: { type: 'string' } }
|
||||
}
|
||||
});
|
||||
|
||||
const messages = await toArrayAsync(stream);
|
||||
|
||||
// Verify no taskCreated or taskStatus messages
|
||||
const taskMessages = messages.filter(m => m.type === 'taskCreated' || m.type === 'taskStatus');
|
||||
expect(taskMessages.length).toBe(0);
|
||||
|
||||
// Verify exactly one result message
|
||||
const resultMessages = messages.filter(m => m.type === 'result');
|
||||
expect(resultMessages.length).toBe(1);
|
||||
|
||||
// Verify total message count is 1
|
||||
expect(messages.length).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
// For any task-augmented elicitation request, the generator should yield at least one
|
||||
// 'taskCreated' message followed by 'taskStatus' messages before yielding the final
|
||||
// result or error.
|
||||
describe('task-augmented request handling', () => {
|
||||
test('should yield taskCreated and result for task-augmented request', async () => {
|
||||
const clientTaskStore = new InMemoryTaskStore();
|
||||
const taskClient = new Client(
|
||||
{ name: 'test client', version: '1.0' },
|
||||
{
|
||||
capabilities: {
|
||||
elicitation: { form: {} },
|
||||
tasks: {
|
||||
requests: {
|
||||
elicitation: { create: {} }
|
||||
}
|
||||
}
|
||||
},
|
||||
taskStore: clientTaskStore
|
||||
}
|
||||
);
|
||||
|
||||
taskClient.setRequestHandler(ElicitRequestSchema, async (request, extra) => {
|
||||
const result = {
|
||||
action: 'accept' as const,
|
||||
content: { username: 'task-user' }
|
||||
};
|
||||
|
||||
if (request.params.task && extra.taskStore) {
|
||||
const task = await extra.taskStore.createTask({ ttl: extra.taskRequestedTtl });
|
||||
await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
|
||||
return { task };
|
||||
}
|
||||
return result;
|
||||
});
|
||||
|
||||
const [taskClientTransport, taskServerTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([taskClient.connect(taskClientTransport), server.connect(taskServerTransport)]);
|
||||
|
||||
const stream = server.experimental.tasks.elicitInputStream(
|
||||
{
|
||||
mode: 'form',
|
||||
message: 'Task-augmented request',
|
||||
requestedSchema: {
|
||||
type: 'object',
|
||||
properties: { username: { type: 'string' } },
|
||||
required: ['username']
|
||||
}
|
||||
},
|
||||
{ task: { ttl: 60_000 } }
|
||||
);
|
||||
|
||||
const messages = await toArrayAsync(stream);
|
||||
|
||||
// Should have taskCreated and result
|
||||
expect(messages.length).toBeGreaterThanOrEqual(2);
|
||||
|
||||
// First message should be taskCreated
|
||||
expect(messages[0].type).toBe('taskCreated');
|
||||
const taskCreated = messages[0] as { type: 'taskCreated'; task: Task };
|
||||
expect(taskCreated.task.taskId).toBeDefined();
|
||||
|
||||
// Last message should be result
|
||||
const lastMessage = messages[messages.length - 1];
|
||||
expect(lastMessage.type).toBe('result');
|
||||
if (lastMessage.type === 'result') {
|
||||
expect((lastMessage.result as ElicitResult).action).toBe('accept');
|
||||
expect((lastMessage.result as ElicitResult).content).toEqual({ username: 'task-user' });
|
||||
}
|
||||
|
||||
clientTaskStore.cleanup();
|
||||
await taskClient.close().catch(() => {});
|
||||
});
|
||||
});
|
||||
});
|
||||
Loading…
Reference in a new issue