mcpServer.tsā¢11.2 kB
/**
* MCP Server implementation
* Simplified architecture with direct Gemini integration
*/
import * as fs from 'node:fs/promises'
import * as path from 'node:path'
import { Server } from '@modelcontextprotocol/sdk/server/index.js'
import {
CallToolRequestSchema,
type CallToolResult,
ListToolsRequestSchema,
type ListToolsResult,
} from '@modelcontextprotocol/sdk/types.js'
// Types
import type { GenerateImageParams, MCPServerConfig } from '../types/mcp'
// Business logic
import { type FileManager, createFileManager } from '../business/fileManager'
import { validateGenerateImageParams } from '../business/inputValidator'
import { type ResponseBuilder, createResponseBuilder } from '../business/responseBuilder'
import {
type FeatureFlags,
type StructuredPromptGenerator,
createStructuredPromptGenerator,
} from '../business/structuredPromptGenerator'
// API clients
import { type GeminiClient, createGeminiClient } from '../api/geminiClient'
import { type GeminiTextClient, createGeminiTextClient } from '../api/geminiTextClient'
// Utilities
import { getConfig } from '../utils/config'
import { Logger } from '../utils/logger'
import { SecurityManager } from '../utils/security'
import { ErrorHandler } from './errorHandler'
/**
* Default MCP server configuration
*/
const DEFAULT_CONFIG: MCPServerConfig = {
name: 'mcp-image-server',
version: '0.1.0',
defaultOutputDir: './output',
}
/**
* Simplified MCP server
*/
export class MCPServerImpl {
private config: MCPServerConfig
private server: Server | null = null
private logger: Logger
private fileManager: FileManager
private responseBuilder: ResponseBuilder
private securityManager: SecurityManager
private structuredPromptGenerator: StructuredPromptGenerator | null = null
private geminiTextClient: GeminiTextClient | null = null
private geminiClient: GeminiClient | null = null
constructor(config: Partial<MCPServerConfig> = {}) {
this.config = { ...DEFAULT_CONFIG, ...config }
this.logger = new Logger()
this.fileManager = createFileManager()
this.responseBuilder = createResponseBuilder()
this.securityManager = new SecurityManager()
}
/**
* Get server info
*/
public getServerInfo() {
return {
name: this.config.name,
version: this.config.version,
}
}
/**
* Get list of registered tools
*/
public getToolsList() {
return {
tools: [
{
name: 'generate_image',
description: 'Generate image with specified prompt and optional parameters',
inputSchema: {
type: 'object' as const,
properties: {
prompt: {
type: 'string' as const,
description:
'The prompt for image generation (English recommended for optimal structured prompt enhancement)',
},
fileName: {
type: 'string' as const,
description:
'Optional file name for the generated image (if not specified, generates an auto-named file in IMAGE_OUTPUT_DIR)',
},
inputImagePath: {
type: 'string' as const,
description:
'Optional absolute path to source image for image-to-image generation. Use when generating variations, style transfers, or similar images based on an existing image (must be an absolute path)',
},
blendImages: {
type: 'boolean' as const,
description:
'Enable multi-image blending for combining multiple visual elements naturally. Use when prompt mentions multiple subjects or composite scenes',
},
maintainCharacterConsistency: {
type: 'boolean' as const,
description:
'Maintain character appearance consistency. Enable when generating same character in different poses/scenes',
},
useWorldKnowledge: {
type: 'boolean' as const,
description:
'Use real-world knowledge for accurate context. Enable for historical figures, landmarks, or factual scenarios',
},
aspectRatio: {
type: 'string' as const,
description: 'Aspect ratio for the generated image',
enum: ['1:1', '2:3', '3:2', '3:4', '4:3', '4:5', '5:4', '9:16', '16:9', '21:9'],
},
},
required: ['prompt'],
},
},
],
}
}
/**
* Tool execution
*/
public async callTool(name: string, args: unknown) {
try {
if (name === 'generate_image') {
return await this.handleGenerateImage(args as GenerateImageParams)
}
throw new Error(`Unknown tool: ${name}`)
} catch (error) {
this.logger.error('mcp-server', 'Tool execution failed', error as Error)
return ErrorHandler.handleError(error as Error)
}
}
/**
* Initialize Gemini clients lazily
*/
private async initializeClients(): Promise<void> {
if (this.structuredPromptGenerator && this.geminiClient) return
const configResult = getConfig()
if (!configResult.success) {
throw configResult.error
}
// Initialize Gemini Text Client for prompt generation
if (!this.geminiTextClient) {
const textClientResult = createGeminiTextClient(configResult.data)
if (!textClientResult.success) {
throw textClientResult.error
}
this.geminiTextClient = textClientResult.data
}
// Initialize Structured Prompt Generator
if (!this.structuredPromptGenerator) {
this.structuredPromptGenerator = createStructuredPromptGenerator(this.geminiTextClient)
}
// Initialize Gemini Client for image generation
if (!this.geminiClient) {
const clientResult = createGeminiClient(configResult.data)
if (!clientResult.success) {
throw clientResult.error
}
this.geminiClient = clientResult.data
}
this.logger.info('mcp-server', 'Gemini clients initialized')
}
/**
* Simplified image generation handler
*/
private async handleGenerateImage(params: GenerateImageParams) {
const result = await ErrorHandler.wrapWithResultType(async () => {
// Validate input
const validationResult = validateGenerateImageParams(params)
if (!validationResult.success) {
throw validationResult.error
}
// Get configuration
const configResult = getConfig()
if (!configResult.success) {
throw configResult.error
}
// Initialize clients
await this.initializeClients()
// Handle input image if provided
let inputImageData: string | undefined
if (params.inputImagePath) {
const imageBuffer = await fs.readFile(params.inputImagePath)
inputImageData = imageBuffer.toString('base64')
}
// Generate structured prompt using Gemini 2.0 Flash (unless skipped)
let structuredPrompt = params.prompt
if (!configResult.data.skipPromptEnhancement && this.structuredPromptGenerator) {
const features: FeatureFlags = {}
if (params.maintainCharacterConsistency !== undefined) {
features.maintainCharacterConsistency = params.maintainCharacterConsistency
}
if (params.blendImages !== undefined) {
features.blendImages = params.blendImages
}
if (params.useWorldKnowledge !== undefined) {
features.useWorldKnowledge = params.useWorldKnowledge
}
const promptResult = await this.structuredPromptGenerator.generateStructuredPrompt(
params.prompt,
features,
inputImageData // Pass image data for context-aware prompt generation
)
if (promptResult.success) {
structuredPrompt = promptResult.data.structuredPrompt
this.logger.info('mcp-server', 'Structured prompt generated', {
originalLength: params.prompt.length,
structuredLength: structuredPrompt.length,
selectedPractices: promptResult.data.selectedPractices,
})
} else {
this.logger.warn('mcp-server', 'Using original prompt', {
error: promptResult.error.message,
})
}
} else if (configResult.data.skipPromptEnhancement) {
this.logger.info('mcp-server', 'Prompt enhancement skipped (SKIP_PROMPT_ENHANCEMENT=true)')
}
// Generate image using Gemini 2.5 Flash Image Preview
if (!this.geminiClient) {
throw new Error('Gemini client not initialized')
}
const generationResult = await this.geminiClient.generateImage({
prompt: structuredPrompt,
...(inputImageData && { inputImage: inputImageData }),
...(params.aspectRatio && { aspectRatio: params.aspectRatio }),
})
if (!generationResult.success) {
throw generationResult.error
}
// Save image file
const fileName = params.fileName || this.fileManager.generateFileName()
const outputPath = path.join(configResult.data.imageOutputDir, fileName)
const sanitizedPath = this.securityManager.sanitizeFilePath(outputPath)
if (!sanitizedPath.success) {
throw sanitizedPath.error
}
const saveResult = await this.fileManager.saveImage(
generationResult.data.imageData,
sanitizedPath.data
)
if (!saveResult.success) {
throw saveResult.error
}
// Build response
return this.responseBuilder.buildSuccessResponse(generationResult.data, saveResult.data)
}, 'image-generation')
if (result.ok) {
return result.value
}
return this.responseBuilder.buildErrorResponse(result.error)
}
/**
* Initialize MCP server with tool handlers
*/
public initialize(): Server {
this.server = new Server(
{
name: this.config.name,
version: this.config.version,
},
{
capabilities: {
tools: {},
},
}
)
// Setup tool handlers
this.setupHandlers()
return this.server
}
/**
* Setup MCP protocol handlers
*/
private setupHandlers(): void {
if (!this.server) {
throw new Error('Server not initialized')
}
// Register tool list handler
this.server.setRequestHandler(ListToolsRequestSchema, async (): Promise<ListToolsResult> => {
return this.getToolsList()
})
// Register tool call handler
this.server.setRequestHandler(
CallToolRequestSchema,
async (request): Promise<CallToolResult> => {
const { name, arguments: args } = request.params
const result = await this.callTool(name, args)
const response: CallToolResult = {
content: result.content,
}
if (result.structuredContent) {
response.structuredContent = result.structuredContent as { [x: string]: unknown }
}
return response
}
)
}
}
/**
* Factory function to create MCP server
*/
export function createMCPServer(config: Partial<MCPServerConfig> = {}) {
return new MCPServerImpl(config)
}