feat: add image support

Author: QuantGeekDev

package.json

@@ -1,6 +1,6 @@
 {
   "name": "mcp-framework",
-  "version": "0.1.26",
+  "version": "0.1.27",
 
   "description": "Framework for building Model Context Protocol (MCP) servers in Typescript",
   "type": "module",

src/cli/index.ts

@@ -11,7 +11,7 @@ const program = new Command();
 program
   .name("mcp")
   .description("CLI for managing MCP server projects")
-  .version("0.1.26");
+  .version("0.1.27");
 
 program
   .command("build")

src/cli/project/create.ts

@@ -61,7 +61,7 @@ export async function createProject(name?: string) {
         start: "node dist/index.js"
       },
       dependencies: {
-        "mcp-framework": "^0.1.26",
+        "mcp-framework": "^0.1.27",
       },
       devDependencies: {
         "@types/node": "^20.11.24",

src/tools/BaseTool.ts

@@ -1,5 +1,6 @@
 import { z } from "zod";
 import { Tool as SDKTool } from "@modelcontextprotocol/sdk/types.js";
+import { ImageContent } from "../transports/utils/image-handler.js";
 
 export type ToolInputSchema<T> = {
   [K in keyof T]: {
@@ -12,6 +13,22 @@ export type ToolInput<T extends ToolInputSchema<any>> = {
   [K in keyof T]: z.infer<T[K]["type"]>;
 };
 
+export type TextContent = {
+  type: "text";
+  text: string;
+};
+
+export type ErrorContent = {
+  type: "error";
+  text: string;
+};
+
+export type ToolContent = TextContent | ErrorContent | ImageContent;
+
+export type ToolResponse = {
+  content: ToolContent[];
+};
+
 export interface ToolProtocol extends SDKTool {
   name: string;
   description: string;
@@ -25,9 +42,7 @@ export interface ToolProtocol extends SDKTool {
   };
   toolCall(request: {
     params: { name: string; arguments?: Record<string, unknown> };
-  }): Promise<{
-    content: Array<{ type: string; text: string }>;
-  }>;
+  }): Promise<ToolResponse>;
 }
 
 export abstract class MCPTool<TInput extends Record<string, any> = {}>
@@ -65,7 +80,7 @@ export abstract class MCPTool<TInput extends Record<string, any> = {}>
 
   async toolCall(request: {
     params: { name: string; arguments?: Record<string, unknown> };
-  }) {
+  }): Promise<ToolResponse> {
     try {
       const args = request.params.arguments || {};
       const validatedInput = await this.validateInput(args);
@@ -95,18 +110,76 @@ export abstract class MCPTool<TInput extends Record<string, any> = {}>
     return "string";
   }
 
-  protected createSuccessResponse(data: unknown) {
+  protected createSuccessResponse(data: unknown): ToolResponse {
+    if (this.isImageContent(data)) {
+      return {
+        content: [data],
+      };
+    }
+
+    if (Array.isArray(data)) {
+      const validContent = data.filter(item => this.isValidContent(item)) as ToolContent[];
+      if (validContent.length > 0) {
+        return {
+          content: validContent,
+        };
+      }
+    }
+
     return {
       content: [{ type: "text", text: JSON.stringify(data) }],
     };
   }
 
-  protected createErrorResponse(error: Error) {
+  protected createErrorResponse(error: Error): ToolResponse {
     return {
       content: [{ type: "error", text: error.message }],
     };
   }
 
+  private isImageContent(data: unknown): data is ImageContent {
+    return (
+      typeof data === "object" &&
+      data !== null &&
+      "type" in data &&
+      data.type === "image" &&
+      "data" in data &&
+      "mimeType" in data &&
+      typeof (data as ImageContent).data === "string" &&
+      typeof (data as ImageContent).mimeType === "string"
+    );
+  }
+
+  private isTextContent(data: unknown): data is TextContent {
+    return (
+      typeof data === "object" &&
+      data !== null &&
+      "type" in data &&
+      data.type === "text" &&
+      "text" in data &&
+      typeof (data as TextContent).text === "string"
+    );
+  }
+
+  private isErrorContent(data: unknown): data is ErrorContent {
+    return (
+      typeof data === "object" &&
+      data !== null &&
+      "type" in data &&
+      data.type === "error" &&
+      "text" in data &&
+      typeof (data as ErrorContent).text === "string"
+    );
+  }
+
+  private isValidContent(data: unknown): data is ToolContent {
+    return (
+      this.isImageContent(data) ||
+      this.isTextContent(data) ||
+      this.isErrorContent(data)
+    );
+  }
+
   protected async fetch<T>(url: string, init?: RequestInit): Promise<T> {
     const response = await fetch(url, init);
     if (!response.ok) {

src/transports/stdio/server.ts

@@ -1,26 +1,75 @@
 import { StdioServerTransport as SDKStdioTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { BaseTransport } from "../base.js";
 import { JSONRPCMessage } from "@modelcontextprotocol/sdk/types.js";
+import {
+  ImageTransportOptions,
+  DEFAULT_IMAGE_OPTIONS,
+  hasImageContent,
+  prepareImageForTransport,
+  ImageContent
+} from "../utils/image-handler.js";
+import { logger } from "../../core/Logger.js";
+
+type ExtendedJSONRPCMessage = JSONRPCMessage & {
+  result?: {
+    content?: Array<ImageContent | { type: string; [key: string]: unknown }>;
+    [key: string]: unknown;
+  };
+};
 
 /**
- * StdioServerTransport that implements BaseTransport
+ * StdioServerTransport
  */
 export class StdioServerTransport implements BaseTransport {
   readonly type = "stdio";
   private transport: SDKStdioTransport;
   private running: boolean = false;
+  private imageOptions: ImageTransportOptions;
 
-  constructor() {
+  constructor(imageOptions: Partial<ImageTransportOptions> = {}) {
     this.transport = new SDKStdioTransport();
+    this.imageOptions = {
+      ...DEFAULT_IMAGE_OPTIONS,
+      ...imageOptions
+    };
   }
 
   async start(): Promise<void> {
     await this.transport.start();
     this.running = true;
   }
 
-  async send(message: JSONRPCMessage): Promise<void> {
-    await this.transport.send(message);
+  async send(message: ExtendedJSONRPCMessage): Promise<void> {
+    try {
+      if (hasImageContent(message)) {
+        message = this.prepareMessageWithImage(message);
+      }
+      await this.transport.send(message);
+    } catch (error) {
+      logger.error(`Error sending message through stdio transport: ${error}`);
+      throw error;
+    }
+  }
+
+  private prepareMessageWithImage(message: ExtendedJSONRPCMessage): ExtendedJSONRPCMessage {
+    if (!message.result?.content) {
+      return message;
+    }
+
+    const processedContent = message.result.content.map((item: ImageContent | { type: string; [key: string]: unknown }) => {
+      if (item.type === 'image') {
+        return prepareImageForTransport(item as ImageContent, this.imageOptions);
+      }
+      return item;
+    });
+
+    return {
+      ...message,
+      result: {
+        ...message.result,
+        content: processedContent
+      }
+    };
   }
 
   async close(): Promise<void> {

src/transports/utils/image-handler.ts

@@ -0,0 +1,110 @@
+import { z } from "zod";
+
+/**
+ * Configuration options for image transport
+ */
+export interface ImageTransportOptions {
+  maxSize: number;
+  allowedMimeTypes: string[];
+  compressionQuality?: number;
+}
+
+/**
+ * Schema for image content validation
+ */
+export const ImageContentSchema = z.object({
+  type: z.literal("image"),
+  data: z.string(),
+  mimeType: z.string()
+});
+
+export type ImageContent = z.infer<typeof ImageContentSchema>;
+
+/**
+ * Default configuration for image transport
+ */
+export const DEFAULT_IMAGE_OPTIONS: ImageTransportOptions = {
+  maxSize: 5 * 1024 * 1024, // 5MB
+  allowedMimeTypes: ['image/jpeg', 'image/png', 'image/gif', 'image/webp'],
+  compressionQuality: 0.8
+};
+
+/**
+ * Validates image content against transport options
+ */
+export function validateImageTransport(content: ImageContent, options: ImageTransportOptions = DEFAULT_IMAGE_OPTIONS): void {
+  // Validate schema
+  ImageContentSchema.parse(content);
+
+  // Validate MIME type
+  if (!options.allowedMimeTypes.includes(content.mimeType)) {
+    throw new Error(`Unsupported image type: ${content.mimeType}. Allowed types: ${options.allowedMimeTypes.join(', ')}`);
+  }
+
+  // Validate base64 format
+  if (!isBase64(content.data)) {
+    throw new Error('Invalid base64 image data');
+  }
+
+  // Validate size
+  const sizeInBytes = Buffer.from(content.data, 'base64').length;
+  if (sizeInBytes > options.maxSize) {
+    throw new Error(`Image size ${sizeInBytes} bytes exceeds maximum allowed size of ${options.maxSize} bytes`);
+  }
+}
+
+/**
+ * Prepares image content for transport
+ * This function can be extended to handle compression, format conversion, etc.
+ */
+export function prepareImageForTransport(content: ImageContent, options: ImageTransportOptions = DEFAULT_IMAGE_OPTIONS): ImageContent {
+  validateImageTransport(content, options);
+  
+  // For now, we just return the validated content
+  // Future: implement compression, format conversion, etc.
+  return content;
+}
+
+/**
+ * Checks if a string is valid base64
+ */
+function isBase64(str: string): boolean {
+  if (str === '' || str.trim() === '') {
+    return false;
+  }
+  try {
+    return btoa(atob(str)) === str;
+  } catch (err) {
+    return false;
+  }
+}
+
+/**
+ * Gets the size of a base64 image in bytes
+ */
+export function getBase64ImageSize(base64String: string): number {
+  return Buffer.from(base64String, 'base64').length;
+}
+
+/**
+ * Utility type for messages containing image content
+ */
+export type MessageWithImage = {
+  result?: {
+    content?: Array<ImageContent | { type: string; [key: string]: unknown }>;
+  };
+  [key: string]: unknown;
+};
+
+/**
+ * Checks if a message contains image content
+ */
+export function hasImageContent(message: unknown): message is MessageWithImage {
+  if (typeof message !== 'object' || message === null) {
+    return false;
+  }
+
+  const msg = message as MessageWithImage;
+  return Array.isArray(msg.result?.content) &&
+    msg.result.content.some(item => item.type === 'image');
+}