refactor(chore): adding test cases and refactor the code

adding test cases and refactor the code

GH-1

Author: vaibhavbhalla2505

package-lock.json

@@ -64,13 +64,16 @@
     "@semantic-release/git": "^10.0.1",
     "@semantic-release/npm": "^9.0.1",
     "@semantic-release/release-notes-generator": "^10.0.3",
+    "@types/jsonwebtoken": "^9.0.10",
     "@types/node": "^16.18.126",
     "commitizen": "^4.2.4",
     "cz-conventional-changelog": "^3.3.0",
     "cz-customizable": "^6.3.0",
     "cz-customizable-ghooks": "^2.0.0",
     "eslint": "^8.57.1",
     "husky": "^7.0.4",
+    "loopback4-authentication": "^13.0.4",
+    "loopback4-authorization": "^8.1.3",
     "semantic-release": "^24.2.9",
     "source-map-support": "^0.5.21",
     "typescript": "~5.2.2"

package.json

@@ -1,130 +1,174 @@
-import {expect, sinon} from '@loopback/testlab';
-import {describe, it, before, after} from 'mocha';
-import http from 'http';
-import {randomUUID} from 'crypto';
-import {z} from 'zod';
-import {Server} from '@modelcontextprotocol/sdk/server/index.js';
-import {StreamableHTTPServerTransport} from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import {expect} from '@loopback/testlab';
 import {Client} from '@modelcontextprotocol/sdk/client/index.js';
 import {StreamableHTTPClientTransport} from '@modelcontextprotocol/sdk/client/streamableHttp.js';
+import * as jwt from 'jsonwebtoken';
+import {McpHookContext} from '../../interfaces';
+import {McpSchemaGeneratorService} from '../../services';
+import {McpToolRegistry} from '../../services/mcp-tool-registry.service';
+import {isTextMessage} from '../../types';
+import {TestingApplication} from '../fixtures/application';
+
+const MCP_ACCESS_PERMISSION = 'mcp.access';
+const UNEXPECTED_MESSAGE_FORMAT = 'Unexpected message format';
+
+async function generateToken(permissions: string[]): Promise<string> {
+  return jwt.sign(
+    {
+      id: 'test-user',
+      permissions,
+    },
+    process.env.JWT_SECRET ?? 'test-secret',
+    {
+      issuer: process.env.JWT_ISSUER,
+      algorithm: 'HS256',
+    },
+  );
+}
+
+describe('MCP Tool – add (acceptance)', () => {
+  let app: TestingApplication;
+  let restServerUrl: string;
 
-describe('Calculator Tool - Add Method (MCP Streamable HTTP)', () => {
-  let httpServer: http.Server;
-  let server: Server;
-  let transport: StreamableHTTPServerTransport;
-  let client: Client;
-  let port: number;
+  before(async () => {
+    app = new TestingApplication();
+    app.bind('hooks.doubleArgs').to(async ({args}: McpHookContext) => ({
+      args: {
+        a: (args.a as number) * 2,
+        b: (args.b as number) * 2,
+      },
+    }));
 
-  let addSpy: sinon.SinonSpy;
+    app.bind('hooks.overrideResult').to(async () => ({
+      result: {
+        content: [{type: 'text', text: 'OVERRIDDEN'}],
+      },
+    }));
+    app.service(McpSchemaGeneratorService);
 
-  before(async () => {
-    server = new Server(
-      {name: 'calculator-mcp-server', version: '1.0.0'},
-      {capabilities: {tools: {}}},
-    );
+    await app.start();
 
-    addSpy = sinon.spy(async (args: {a: number; b: number}) => {
-      const sum = args.a + args.b;
-      return {
-        content: [
-          {
-            type: 'text',
-            text: `The sum of ${args.a} and ${args.b} is ${sum}`,
-          },
-        ],
-      };
-    });
+    const registry = await app.get<McpToolRegistry>('services.McpToolRegistry');
+    await registry.initialize();
 
-    server.setRequestHandler(
-      z.object({
-        method: z.literal('tools/call'),
-        params: z.object({
-          name: z.literal('calculator_add'),
-          arguments: z.object({
-            a: z.number(),
-            b: z.number(),
-          }),
-        }),
-      }),
-      async request => {
-        return addSpy(request.params.arguments);
+    restServerUrl = app.restServer.url ?? '';
+  });
+
+  after(async () => {
+    await app.stop();
+  });
+
+  async function createClient(token: string) {
+    const transport = new StreamableHTTPClientTransport(
+      new URL(`${restServerUrl}/mcp`),
+      {
+        requestInit: {
+          headers: {
+            Authorization: `Bearer ${token}`,
+          },
+        },
       },
     );
 
-    transport = new StreamableHTTPServerTransport({
-      sessionIdGenerator: () => randomUUID(),
+    const client = new Client({
+      name: 'test-client',
+      version: '1.0.0',
     });
 
-    await server.connect(transport);
+    await client.connect(transport);
+    return client;
+  }
 
-    httpServer = http.createServer((req, res) => {
-      if (req.method !== 'POST' || req.url !== '/mcp') {
-        res.statusCode = 404;
-        res.end();
-        return;
-      }
+  it('invokes add tool and returns correct sum', async () => {
+    const token = await generateToken([MCP_ACCESS_PERMISSION]);
+    const client = await createClient(token);
 
-      let body = '';
-      req.on('data', chunk => (body += chunk));
-      // eslint-disable-next-line @typescript-eslint/no-misused-promises
-      req.on('end', async () => {
-        const message = JSON.parse(body);
+    const result = await client.callTool({
+      name: 'summary_get_add',
+      arguments: {a: 10, b: 20},
+    });
+    if (isTextMessage(result)) {
+      expect(result.content[0].text).to.equal('The sum of 10 and 20 is 30');
+    } else {
+      throw new Error(UNEXPECTED_MESSAGE_FORMAT);
+    }
+
+    await client.close();
+  });
 
-        const response = await transport.handleRequest(req, res, message);
+  it('denies execution for unauthorized user', async () => {
+    const token = await generateToken([]);
+    const client = await createClient(token);
 
-        res.setHeader('Content-Type', 'application/json');
-        res.end(JSON.stringify(response));
-      });
+    const result = await client.callTool({
+      name: 'summary_get_add',
+      arguments: {a: 1, b: 2},
     });
+    if (isTextMessage(result)) {
+      expect(result.isError).to.be.true();
+      expect(result.content[0].text).to.containEql(
+        'Access denied for MCP tool',
+      );
+    } else {
+      throw new Error(UNEXPECTED_MESSAGE_FORMAT);
+    }
 
-    await new Promise<void>(resolve => {
-      httpServer.listen(0, () => {
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        port = (httpServer.address() as any).port;
-        resolve();
-      });
+    await client.close();
+  });
+
+  it('denies execution when permission is different', async () => {
+    const token = await generateToken(['some.other.permission']);
+    const client = await createClient(token);
+
+    const result = await client.callTool({
+      name: 'summary_get_add',
+      arguments: {a: 1, b: 2},
     });
-    client = new Client(
-      {name: 'calculator-test-client', version: '1.0.0'},
-      {capabilities: {tools: {}}},
-    );
 
-    await client.connect(
-      new StreamableHTTPClientTransport(
-        new URL(`http://localhost:${port}/mcp`),
-      ),
-    );
+    if (isTextMessage(result)) {
+      expect(result.isError).to.be.true();
+      expect(result.content[0].text).to.containEql(
+        'Access denied for MCP tool',
+      );
+    } else {
+      throw new Error(UNEXPECTED_MESSAGE_FORMAT);
+    }
+
+    await client.close();
   });
 
-  after(async () => {
-    sinon.restore();
+  it('executes pre-hook and modifies arguments', async () => {
+    const token = await generateToken([MCP_ACCESS_PERMISSION]);
+    const client = await createClient(token);
+
+    const result = await client.callTool({
+      name: 'summary_add_with_prehook',
+      arguments: {a: 2, b: 3},
+    });
+    if (isTextMessage(result)) {
+      const response = JSON.parse(result.content[0].text);
+      expect(response.sum).to.equal(10);
+    } else {
+      throw new Error(UNEXPECTED_MESSAGE_FORMAT);
+    }
+
     await client.close();
-    await server.close();
-    await new Promise(resolve => httpServer.close(resolve));
   });
 
-  it('should call add method and return correct sum', async () => {
-    const result = await client.request(
-      {
-        method: 'tools/call',
-        params: {
-          name: 'calculator_add',
-          arguments: {a: 5, b: 7},
-        },
-      },
-      z.object({
-        content: z.array(
-          z.object({
-            type: z.literal('text'),
-            text: z.string(),
-          }),
-        ),
-      }),
-    );
+  it('executes post-hook and overrides result', async () => {
+    const token = await generateToken([MCP_ACCESS_PERMISSION]);
+    const client = await createClient(token);
+
+    const result = await client.callTool({
+      name: 'summary_add_with_posthook',
+      arguments: {a: 1, b: 1},
+    });
 
-    expect(result.content[0].text).to.equal('The sum of 5 and 7 is 12');
+    if (isTextMessage(result)) {
+      expect(result.content[0].text).to.equal('OVERRIDDEN');
+    } else {
+      throw new Error(UNEXPECTED_MESSAGE_FORMAT);
+    }
 
-    sinon.assert.calledOnce(addSpy);
-    sinon.assert.calledWithMatch(addSpy, {a: 5, b: 7});
+    await client.close();
   });
 });