@genkit-ai/mcp 1.14.1-rc.0

package/tests/server_test.ts

@@ -0,0 +1,165 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { Client } from '@modelcontextprotocol/sdk/client/index.js';
+import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
+import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
+import * as assert from 'assert';
+import express from 'express';
+import { Genkit, genkit, z } from 'genkit';
+import { logger } from 'genkit/logging';
+import getPort from 'get-port';
+import * as http from 'http';
+import { afterEach, beforeEach, describe, it } from 'node:test';
+import { createMcpServer } from '../src/index.js';
+import { GenkitMcpServer } from '../src/server.js';
+import { defineEchoModel } from './fakes.js';
+
+logger.setLogLevel('debug');
+
+describe('createMcpServer', async () => {
+  let ai: Genkit;
+  let mcpServer: GenkitMcpServer;
+  let mcpHttpServer: http.Server;
+  let port: number;
+  let client: Client;
+
+  beforeEach(async () => {
+    ai = genkit({});
+    defineEchoModel(ai);
+
+    ai.definePrompt({
+      name: 'testPrompt',
+      model: 'echoModel',
+      prompt: 'prompt says: {{input}}',
+    });
+
+    ai.defineTool(
+      {
+        name: 'testTool',
+        description: 'test tool',
+        inputSchema: z.object({ foo: z.string() }),
+      },
+      async (input) => `yep ${JSON.stringify(input)}`
+    );
+
+    mcpServer = createMcpServer(ai, { name: 'test-server', version: '0.0.1' });
+    await mcpServer.setup();
+
+    const app = express();
+    let transport: SSEServerTransport | null = null;
+
+    app.get('/sse', (req, res) => {
+      transport = new SSEServerTransport('/messages', res);
+      mcpServer.server!.connect(transport);
+    });
+
+    app.post('/messages', (req, res) => {
+      if (transport) {
+        transport.handlePostMessage(req, res);
+      }
+    });
+
+    port = await getPort();
+    let serverResolver;
+    const serverPromise = new Promise((r) => {
+      serverResolver = r;
+    });
+    mcpHttpServer = app.listen(port, () => {
+      console.log(`MCP server listening on http://localhost:${port}`);
+      serverResolver();
+    });
+    await serverPromise; // wait for server to start up
+
+    client = new Client({ name: 'test', version: '0.0.1' });
+    await client.connect(
+      new SSEClientTransport(new URL(`http://localhost:${port}/sse`))
+    );
+  });
+
+  afterEach(async () => {
+    mcpHttpServer.close();
+    await mcpServer.server?.close();
+    await client?.close();
+  });
+
+  describe('tools', () => {
+    it('should list tools', async () => {
+      const r = await client.listTools();
+
+      assert.deepStrictEqual(r.tools, [
+        {
+          description: 'test tool',
+          inputSchema: {
+            $schema: 'http://json-schema.org/draft-07/schema#',
+            additionalProperties: true,
+            properties: {
+              foo: {
+                type: 'string',
+              },
+            },
+            required: ['foo'],
+            type: 'object',
+          },
+          name: 'testTool',
+        },
+      ]);
+    });
+
+    it('should call the tool', async () => {
+      const response = await client.callTool({
+        name: 'testTool',
+        arguments: {
+          foo: 'bar',
+        },
+      });
+      assert.deepStrictEqual(response, {
+        content: [
+          {
+            text: 'yep {"foo":"bar"}',
+            type: 'text',
+          },
+        ],
+      });
+    });
+  });
+
+  describe('prompts', () => {
+    it('should list prompts', async () => {
+      const r = await client.listPrompts();
+      assert.deepStrictEqual(r.prompts, [{ name: 'testPrompt' }]);
+    });
+
+    it('should render prompt', async () => {
+      const prompt = await client.getPrompt({
+        name: 'testPrompt',
+        arguments: {
+          input: 'hello',
+        },
+      });
+
+      assert.deepStrictEqual(prompt.messages, [
+        {
+          content: {
+            text: 'prompt says: hello',
+            type: 'text',
+          },
+          role: 'user',
+        },
+      ]);
+    });
+  });
+});