[WIP] Chore: Add output schema to tools

typescript/src/modelcontextprotocol/__tests__/essentials.test.ts

@@ -93,7 +93,7 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
 
     // Set up McpServer mock to handle the fact that CommercetoolsAgentEssentials extends it
     (McpServer as jest.Mock).mockImplementation(function (this: any) {
-      this.tool = mockToolMethod;
+      this.registerTool = mockToolMethod;
     });
 
     mockCommercetoolsAPIInstance = new CommercetoolsAPI(
@@ -192,14 +192,19 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
     // Check if registerTool was called with the correct parameters
     expect(mockToolMethod).toHaveBeenCalledWith(
       mockSharedToolsData[0].method,
-      mockSharedToolsData[0].description,
-      expect.any(Object),
+      expect.objectContaining({
+        description: mockSharedToolsData[0].description,
+        inputSchema: expect.any(Object),
+      }),
+
       expect.any(Function) // Handler function
     );
     expect(mockToolMethod).toHaveBeenCalledWith(
       mockSharedToolsData[1].method,
-      mockSharedToolsData[1].description,
-      expect.any(Object),
+      expect.objectContaining({
+        description: mockSharedToolsData[1].description,
+        inputSchema: expect.any(Object),
+      }),
       expect.any(Function) // Handler function
     );
   });
@@ -220,7 +225,7 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
     // Get the handler from the mock call
     await new Promise(setImmediate);
     const toolCallArgs = mockToolMethod.mock.calls[0];
-    const handler = toolCallArgs[3]; // The async handler function
+    const handler = toolCallArgs[2]; // The async handler function
     const toolMethod = toolCallArgs[0];
 
     const handlerArg = {paramA: 'testValue'};
@@ -435,7 +440,7 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
 
       // Set up McpServer mock to handle the fact that CommercetoolsAgentEssentials extends it
       (McpServer as jest.Mock).mockImplementation(function (this: any) {
-        this.tool = _mockToolMethod;
+        this.registerTool = _mockToolMethod;
       });
 
       _mockCommercetoolsAPIInstance = new CommercetoolsAPI(
@@ -540,7 +545,6 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
 
         expect(_mockToolMethod).toHaveBeenCalledWith(
           'mcpTool1',
-          expect.any(String),
           expect.any(Object),
           expect.any(Function)
         );
@@ -575,7 +579,6 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
         expect(_mockToolMethod).toHaveBeenCalledTimes(2);
         expect(_mockToolMethod).toHaveBeenCalledWith(
           'custom-test-tool',
-          expect.any(String),
           expect.any(Object),
           expect.any(Function)
         );
@@ -690,7 +693,7 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
 
       // Set up McpServer mock to handle the fact that CommercetoolsAgentEssentials extends it
       (McpServer as jest.Mock).mockImplementation(function (this: any) {
-        this.tool = mockToolMethod;
+        this.registerTool = mockToolMethod;
       });
 
       mockCommercetoolsAPIInstance = new CommercetoolsAPI(
@@ -757,13 +760,11 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
         expect(mockToolMethod).toHaveBeenCalledTimes(2);
         expect(mockToolMethod).toHaveBeenCalledWith(
           'mcpTool1',
-          expect.any(String),
           expect.any(Object),
           expect.any(Function)
         );
         expect(mockToolMethod).toHaveBeenCalledWith(
           'mcpTool2',
-          expect.any(String),
           expect.any(Object),
           expect.any(Function)
         );
@@ -860,19 +861,16 @@ describe('CommercetoolsAgentEssentials (ModelContextProtocol)', () => {
         expect(mockToolMethod).toHaveBeenCalledTimes(5);
         expect(mockToolMethod).toHaveBeenCalledWith(
           'list_available_tools',
-          expect.any(String),
           expect.any(Object),
           expect.any(Function)
         );
         expect(mockToolMethod).toHaveBeenCalledWith(
           'inject_tools',
-          expect.any(String),
           expect.any(Object),
           expect.any(Function)
         );
         expect(mockToolMethod).toHaveBeenCalledWith(
           'execute_tool',
-          expect.any(String),
           expect.any(Object),
           expect.any(Function)
         );

typescript/src/modelcontextprotocol/essentials.ts

@@ -14,6 +14,7 @@ import {Tool} from '../types/tools';
 import {contextToBulkTools} from '../shared/bulk/tools';
 import {DYNAMIC_TOOL_LOADING_THRESHOLD} from '../shared/constants';
 import {transformToolOutput} from './transform';
+import {transfromOutputSchema} from './transform/transformToolOutput';
 
 class CommercetoolsAgentEssentials extends McpServer {
   private authConfig: AuthConfig;
@@ -167,10 +168,14 @@ class CommercetoolsAgentEssentials extends McpServer {
 
   private registerSingleTool(tool: Tool): void {
     const {method, execute} = tool;
-    this.tool(
+    this.registerTool(
       tool.method,
-      tool.description,
-      tool.parameters.shape,
+      {
+        description: tool.description,
+        inputSchema: tool.parameters.shape,
+        ...(this.configuration.context?.toolOutputFormat === 'json' &&
+          transfromOutputSchema(tool)),
+      },
       async (args: Record<string, unknown>) => {
         const result = await this.commercetoolsAPI.run(method, args, execute);
         return this.createToolResponse(
@@ -189,11 +194,12 @@ class CommercetoolsAgentEssentials extends McpServer {
     filteredTools: Tool[]
   ): void {
     type ToolShape = z.infer<typeof injectTools.parameters.shape>;
-
-    this.tool(
+    this.registerTool(
       injectTools.method,
-      injectTools.description,
-      injectTools.parameters.shape,
+      {
+        description: injectTools.description,
+        inputSchema: injectTools.parameters.shape,
+      },
       async (arg: ToolShape) => {
         const toolsToInject = filteredTools.filter((tool) =>
           arg.toolMethods.includes(tool.method)
@@ -216,10 +222,12 @@ class CommercetoolsAgentEssentials extends McpServer {
   private registerExecuteTool(executeTool: Tool): void {
     type ToolShape = z.infer<typeof executeTool.parameters.shape>;
 
-    this.tool(
+    this.registerTool(
       executeTool.method,
-      executeTool.description,
-      executeTool.parameters.shape,
+      {
+        description: executeTool.description,
+        inputSchema: executeTool.parameters.shape,
+      },
       async (args: ToolShape) => {
         try {
           const result = await this.commercetoolsAPI.run(

typescript/src/modelcontextprotocol/transform/transformToolOutput.ts

@@ -1,5 +1,6 @@
+import {Tool} from '../../types/tools';
 import {transformPropertyName} from './transformPropertyName';
-
+import {z} from 'zod';
 const emptyObjectTransformValue = 'no properties';
 const emptyArrayTransformValue = 'none';
 
@@ -13,6 +14,28 @@ const generateTabs = (tabCount: number) => {
 
 type Format = 'tabular' | 'json';
 
+/**
+ * Transforms the output schema of a tool to a JSON schema.
+ *
+ * @param {Tool} tool - The tool to transform the output schema of.
+ *
+ * @returns {Object} The transformed output schema.
+ */
+export const transfromOutputSchema = (tool: Tool) => {
+  if (!tool.outputSchema) {
+    return {};
+  }
+
+  return {
+    outputSchema: z.object({
+      [transformTitle(tool.name)]: z.preprocess(
+        (val) => JSON.parse(val as string),
+        tool.outputSchema
+      ),
+    }).shape,
+  };
+};
+
 /**
  * A method to strigify tool output into a LLM friendly and optimised format.
  *

typescript/src/shared/business-unit/output-schema.ts

@@ -0,0 +1,27 @@
+import {z} from 'zod';
+
+// Base BusinessUnit schema
+export const businessUnitSchema = z.object({
+  id: z.string(),
+  version: z.number(),
+  createdAt: z.string(),
+  lastModifiedAt: z.string(),
+  lastModifiedBy: z.any().optional(),
+  createdBy: z.any().optional(),
+});
+
+// Paged BusinessUnitPagedQueryResponse schema
+export const businessUnitPagedSchema = z
+  .object({
+    limit: z.number(),
+    offset: z.number(),
+    count: z.number(),
+    total: z.number().optional(),
+    results: z.array(z.any()),
+  })
+  .strict();
+
+export const readBusinessUnitOutputSchema = z.union([businessUnitSchema, businessUnitPagedSchema]);
+
+export const createBusinessUnitOutputSchema = businessUnitSchema;
+export const updateBusinessUnitOutputSchema = businessUnitSchema;

typescript/src/shared/business-unit/test/output-schema.test.ts

@@ -0,0 +1,81 @@
+import {
+  businessUnitSchema,
+  businessUnitPagedSchema,
+  readBusinessUnitOutputSchema,
+} from '../output-schema';
+import outputData from './output.data.json';
+
+describe('BusinessUnit Output Schema', () => {
+  describe('Paged Response Schema', () => {
+    it('should validate real API paged response', () => {
+      const result = businessUnitPagedSchema.safeParse(outputData);
+
+      if (!result.success) {
+        console.error('Validation errors:', JSON.stringify(result.error.format(), null, 2));
+      }
+
+      expect(result.success).toBe(true);
+    });
+
+    it('should have required paged fields', () => {
+      expect(outputData).toHaveProperty('limit');
+      expect(outputData).toHaveProperty('offset');
+      expect(outputData).toHaveProperty('count');
+      expect(outputData).toHaveProperty('results');
+      expect(Array.isArray((outputData as any).results)).toBe(true);
+    });
+  });
+
+  describe('Individual Entity Schema', () => {
+    it('should validate individual entities from results', () => {
+      const results = (outputData as any).results;
+
+      if (results && results.length > 0) {
+        const firstEntity = results[0];
+        const result = businessUnitSchema.passthrough().safeParse(firstEntity);
+
+        if (!result.success) {
+          console.error('Validation errors:', JSON.stringify(result.error.format(), null, 2));
+        }
+
+        expect(result.success).toBe(true);
+      } else {
+        console.warn('No results in output data to validate individual entities');
+      }
+    });
+
+    it('should validate all entities in results array', () => {
+      const results = (outputData as any).results || [];
+
+      results.forEach((entity: any, index: number) => {
+        const result = businessUnitSchema.passthrough().safeParse(entity);
+
+        if (!result.success) {
+          console.error(`Entity ${index} validation errors:`, JSON.stringify(result.error.format(), null, 2));
+        }
+
+        expect(result.success).toBe(true);
+      });
+    });
+  });
+
+  describe('Read Output Schema (Union)', () => {
+    it('should validate paged response', () => {
+      const result = readBusinessUnitOutputSchema.safeParse(outputData);
+
+      if (!result.success) {
+        console.error('Validation errors:', JSON.stringify(result.error.format(), null, 2));
+      }
+
+      expect(result.success).toBe(true);
+    });
+  });
+
+  describe('Schema Structure', () => {
+    it('should be a valid Zod schema', () => {
+      expect(businessUnitSchema).toBeDefined();
+      expect(readBusinessUnitOutputSchema).toBeDefined();
+      expect(businessUnitPagedSchema).toBeDefined();
+    });
+  });
+});