feat: add streaming methods for elicitation and createMessage

Author: LucaButBoring

src/experimental/tasks/server.ts

@@ -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,189 @@ 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
+        type ServerWithCapabilities = {
+            getClientCapabilities(): { sampling?: { tools?: boolean } } | undefined;
+        };
+        const clientCapabilities = (this._server as unknown as ServerWithCapabilities).getClientCapabilities();
+
+        // Capability check - only required when tools/toolChoice are provided
+        if (params.tools || params.toolChoice) {
+            if (!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) {
+                type ToolUseContent = { type: 'tool_use'; id: string };
+                type ToolResultContent = { type: 'tool_result'; toolUseId: string };
+                const toolUseIds = new Set(previousContent.filter(c => c.type === 'tool_use').map(c => (c as ToolUseContent).id));
+                const toolResultIds = new Set(
+                    lastContent.filter(c => c.type === 'tool_result').map(c => (c as ToolResultContent).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');
+                }
+            }
+        }
+
+        const request = {
+            method: 'sampling/createMessage' as const,
+            params
+        };
+        return this.requestStream(request, 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
+        type ServerWithCapabilities = {
+            getClientCapabilities(): { elicitation?: { form?: boolean; url?: boolean } } | undefined;
+        };
+        const clientCapabilities = (this._server as unknown as ServerWithCapabilities).getClientCapabilities();
+
+        const mode = (params.mode ?? 'form') as 'form' | 'url';
+
+        // 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
+        const normalizedParams =
+            mode === 'form' && params.mode !== 'form' ? { ...(params as ElicitRequestFormParams), mode: 'form' as const } : params;
+
+        const request = {
+            method: 'elicitation/create' as const,
+            params: normalizedParams
+        };
+        return this.requestStream(request, ElicitResultSchema, options);
+    }
+
     /**
      * Gets the current status of a task.
      *