[#77] Added sampling support to tools

src/core/MCPServer.ts

@@ -424,6 +424,7 @@ export class MCPServer {
         { name: this.serverName, version: this.serverVersion },
         { capabilities: this.capabilities }
       );
+      tools.forEach((tool) => tool.injectServer(this.server));
       logger.debug(
         `SDK Server instance created with capabilities: ${JSON.stringify(this.capabilities)}`
       );

src/tools/BaseTool.ts

@@ -1,17 +1,8 @@
 import { z } from 'zod';
-import { Tool as SDKTool } from '@modelcontextprotocol/sdk/types.js';
+import { CreateMessageRequest, CreateMessageResult, Tool as SDKTool } from '@modelcontextprotocol/sdk/types.js';
 import { ImageContent } from '../transports/utils/image-handler.js';
-
-// Type to check if a Zod type has a description
-type HasDescription<T> = T extends { _def: { description: string } } ? T : never;
-
-// Type to ensure all properties in a Zod object have descriptions
-type AllFieldsHaveDescriptions<T extends z.ZodRawShape> = {
-  [K in keyof T]: HasDescription<T[K]>;
-};
-
-// Strict Zod object type that requires all fields to have descriptions
-type StrictZodObject<T extends z.ZodRawShape> = z.ZodObject<AllFieldsHaveDescriptions<T>>;
+import { Server } from '@modelcontextprotocol/sdk/server/index.js';
+import { RequestOptions } from '@modelcontextprotocol/sdk/shared/protocol.js';
 
 export type ToolInputSchema<T> = {
   [K in keyof T]: {
@@ -64,6 +55,7 @@ export interface ToolProtocol extends SDKTool {
   toolCall(request: {
     params: { name: string; arguments?: Record<string, unknown> };
   }): Promise<ToolResponse>;
+  injectServer(server: Server): void;
 }
 
 /**
@@ -100,6 +92,46 @@ export abstract class MCPTool<TInput extends Record<string, any> = any, TSchema
   protected useStringify: boolean = true;
   [key: string]: unknown;
 
+  private server: Server | undefined;
+
+  /**
+   * Injects the server into this tool to allow sampling requests.
+   * Automatically called by the MCP server when registering the tool.
+   * Subsequent calls are silently ignored.
+   */
+  public injectServer(server: Server): void {
+    if (this.server) {
+      return;
+    }
+    this.server = server;
+  }
+
+  /**
+   * Submit a sampling request to the client via the MCP sampling protocol.
+   * Can only be called from within a tool's execute() method after the server
+   * has been injected.
+   *
+   * @example
+   * ```typescript
+   * const result = await this.samplingRequest({
+   *   messages: [{ role: "user", content: { type: "text", text: "Hello!" } }],
+   *   maxTokens: 100
+   * });
+   * ```
+   */
+  protected async samplingRequest(
+    request: CreateMessageRequest['params'],
+    options?: RequestOptions,
+  ): Promise<CreateMessageResult> {
+    if (!this.server) {
+      throw new Error(
+        `Cannot make sampling request: server not available in tool '${this.name}'. ` +
+        `Sampling is only available during tool execution within an MCPServer.`,
+      );
+    }
+    return this.server.createMessage(request, options);
+  }
+
   /**
    * Validates the tool schema. This is called automatically when the tool is registered
    * with an MCP server, but can also be called manually for testing.

tests/tools/BaseTool.test.ts

@@ -1,6 +1,16 @@
-import { describe, it, expect, beforeEach } from '@jest/globals';
+import { describe, it, expect, beforeEach, jest } from '@jest/globals';
 import { z } from 'zod';
 import { MCPTool } from '../../src/tools/BaseTool.js';
+import { Server } from '@modelcontextprotocol/sdk/server/index.js';
+import { CreateMessageRequest, CreateMessageResult } from '@modelcontextprotocol/sdk/types.js';
+import {RequestOptions} from '@modelcontextprotocol/sdk/shared/protocol.js';
+
+// Mock the Server class
+jest.mock('@modelcontextprotocol/sdk/server/index.js', () => ({
+  Server: jest.fn().mockImplementation(() => ({
+    createMessage: jest.fn(),
+  })),
+}));
 
 describe('BaseTool', () => {
   describe('Legacy Pattern (Separate Schema Definition)', () => {
@@ -488,4 +498,147 @@ describe('BaseTool', () => {
       console.log(JSON.stringify(definition, null, 2));
     });
   });
+
+  describe('Sampling', () => {
+    // Expose the protected samplingRequest for direct testing
+    class SamplingTestTool extends MCPTool {
+      name = 'sampling_tool';
+      description = 'A tool that uses sampling';
+      schema = z.object({
+        prompt: z.string().describe('The prompt to sample'),
+      });
+
+      protected async execute(input: { prompt: string }): Promise<unknown> {
+        const result = await this.samplingRequest({
+          messages: [
+            {
+              role: 'user',
+              content: { type: 'text', text: input.prompt },
+            },
+          ],
+          maxTokens: 100,
+        });
+        return { sampledText: result.content.text };
+      }
+
+      // Expose protected method for testing
+      public testSamplingRequest(
+        request: CreateMessageRequest['params'],
+        options?: RequestOptions,
+      ) {
+        return this.samplingRequest(request, options);
+      }
+    }
+
+    let tool: SamplingTestTool;
+    let mockServer: jest.Mocked<Server>;
+
+    beforeEach(() => {
+      tool = new SamplingTestTool();
+      mockServer = new Server(
+        { name: 'test-server', version: '1.0.0' },
+        { capabilities: {} },
+      ) as jest.Mocked<Server>;
+      mockServer.createMessage = jest.fn();
+    });
+
+    it('should inject server without throwing', () => {
+      expect(() => tool.injectServer(mockServer)).not.toThrow();
+    });
+
+    it('should silently handle double injection', () => {
+      tool.injectServer(mockServer);
+      expect(() => tool.injectServer(mockServer)).not.toThrow();
+    });
+
+    it('should throw when samplingRequest called without server', async () => {
+      await expect(
+        tool.testSamplingRequest({
+          messages: [{ role: 'user', content: { type: 'text', text: 'test' } }],
+          maxTokens: 100,
+        }),
+      ).rejects.toThrow(
+        "Cannot make sampling request: server not available in tool 'sampling_tool'.",
+      );
+    });
+
+    it('should call server.createMessage with correct params', async () => {
+      const mockResult: CreateMessageResult = {
+        model: 'test-model',
+        role: 'assistant',
+        content: { type: 'text', text: 'Sampled response' },
+      };
+      mockServer.createMessage.mockResolvedValue(mockResult);
+      tool.injectServer(mockServer);
+
+      const request: CreateMessageRequest['params'] = {
+        messages: [{ role: 'user', content: { type: 'text', text: 'Hello' } }],
+        maxTokens: 100,
+        temperature: 0.7,
+        systemPrompt: 'Be helpful',
+      };
+
+      const result = await tool.testSamplingRequest(request);
+
+      expect(mockServer.createMessage).toHaveBeenCalledWith(request, undefined);
+      expect(result).toEqual(mockResult);
+    });
+
+    it('should propagate createMessage errors', async () => {
+      tool.injectServer(mockServer);
+      mockServer.createMessage.mockRejectedValue(new Error('Sampling failed'));
+
+      await expect(
+        tool.testSamplingRequest({
+          messages: [{ role: 'user', content: { type: 'text', text: 'test' } }],
+          maxTokens: 100,
+        }),
+      ).rejects.toThrow('Sampling failed');
+    });
+
+    it('should pass request options to createMessage', async () => {
+      const mockResult: CreateMessageResult = {
+        model: 'claude-3-sonnet',
+        role: 'assistant',
+        content: { type: 'text', text: 'Complex response' },
+        stopReason: 'endTurn',
+      };
+      mockServer.createMessage.mockResolvedValue(mockResult);
+      tool.injectServer(mockServer);
+
+      const request: CreateMessageRequest['params'] = {
+        messages: [
+          { role: 'user', content: { type: 'text', text: 'First message' } },
+          { role: 'assistant', content: { type: 'text', text: 'Assistant response' } },
+          { role: 'user', content: { type: 'text', text: 'Follow up' } },
+        ],
+        maxTokens: 500,
+        temperature: 0.8,
+        systemPrompt: 'You are a helpful assistant',
+        modelPreferences: {
+          hints: [{ name: 'claude-3' }],
+          costPriority: 0.3,
+          speedPriority: 0.7,
+          intelligencePriority: 0.9,
+        },
+        stopSequences: ['END', 'STOP'],
+        metadata: { taskType: 'analysis' },
+      };
+
+      const options: RequestOptions = {
+        timeout: 5000,
+        maxTotalTimeout: 10000,
+        signal: new AbortController().signal,
+        resetTimeoutOnProgress: true,
+        onprogress: (progress) => {
+          console.log('Progress:', progress);
+        },
+      };
+
+      const result = await tool.testSamplingRequest(request, options);
+
+      expect(mockServer.createMessage).toHaveBeenCalledWith(request, options);
+      expect(result).toEqual(mockResult);
+    });
+  });
 });