ken-you-code/src/server.ts

Connect your codebase to Kimi: Ultra-fast AI code analysis with Kimi-K2 model via MCP

#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { config } from './config/index.js';
import { FileOperations } from './tools/fileOperations.js';
import { SpecializedTaskHandler } from './tools/specializedTask.js';
import { DiffManager } from './tools/diffManager.js';
import {
  SpecializedTaskRequestSchema,
  ReadFileRequestSchema,
  CreateDiffRequestSchema,
  ApplyDiffRequestSchema,
} from './types/index.js';

const server = new Server(
  {
    name: 'single-model-specialist',
    version: '1.0.0',
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

const taskHandler = new SpecializedTaskHandler();

server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: 'specialized_task',
        description: 'Delegate task to external model with file context',
        inputSchema: {
          type: 'object',
          properties: {
            task_type: {
              type: 'string',
              enum: [
                'debug',
                'analyze',
                'review',
                'design',
                'second_opinion',
                'reality_check',
                'synthesis',
              ],
              description: 'Type of task to perform',
            },
            context: {
              type: 'string',
              description: 'Task description and context',
            },
            files: {
              type: 'array',
              items: { type: 'string' },
              description: 'File paths for context',
            },
            model: {
              type: 'string',
              description: 'Model identifier',
              default: config.defaultModel,
            },
          },
          required: ['task_type', 'context'],
        },
      },
      {
        name: 'read_file',
        description: 'Read file content securely',
        inputSchema: {
          type: 'object',
          properties: {
            file_path: {
              type: 'string',
              description: 'Absolute path to file',
            },
          },
          required: ['file_path'],
        },
      },
      {
        name: 'create_diff',
        description: 'Generate diff for file changes',
        inputSchema: {
          type: 'object',
          properties: {
            file_path: {
              type: 'string',
              description: 'Target file path',
            },
            new_content: {
              type: 'string',
              description: 'Proposed file content',
            },
            reason: {
              type: 'string',
              description: 'Explanation for the change',
            },
          },
          required: ['file_path', 'new_content', 'reason'],
        },
      },
      {
        name: 'apply_diff',
        description: 'Apply approved diff with backup',
        inputSchema: {
          type: 'object',
          properties: {
            operation_id: {
              type: 'string',
              description: 'Diff operation identifier',
            },
            approved: {
              type: 'boolean',
              description: 'User approval for the change',
            },
          },
          required: ['operation_id', 'approved'],
        },
      },
    ],
  };
});

server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  try {
    switch (name) {
      case 'specialized_task': {
        const taskRequest = SpecializedTaskRequestSchema.parse(args);
        const result = await taskHandler.executeTask(taskRequest);

        return {
          content: [
            {
              type: 'text',
              text:
                `# ${taskRequest.task_type.toUpperCase()} Analysis\n\n` +
                `**Model:** ${result.model}\n` +
                `**Files Analyzed:** ${taskRequest.files?.join(', ') || 'None'}\n\n` +
                `## Result\n\n${result.result}\n\n` +
                (result.tokenUsage
                  ? `**Token Usage:** ${result.tokenUsage.total} tokens (${result.tokenUsage.prompt} prompt + ${result.tokenUsage.completion} completion)`
                  : ''),
            },
          ],
        };
      }

      case 'read_file': {
        const { file_path } = ReadFileRequestSchema.parse(args);
        const result = await FileOperations.readFileSecurely(file_path);

        return {
          content: [
            {
              type: 'text',
              text:
                `# File: ${file_path}\n\n` +
                `**Size:** ${result.size} bytes\n` +
                `**Extension:** ${result.extension}\n\n` +
                `## Content\n\n\`\`\`\n${result.content}\n\`\`\``,
            },
          ],
        };
      }

      case 'create_diff': {
        const { file_path, new_content, reason } = CreateDiffRequestSchema.parse(args);
        const operation = await DiffManager.createDiff(file_path, new_content, reason);

        return {
          content: [
            {
              type: 'text',
              text:
                `# Diff Created\n\n` +
                `**Operation ID:** ${operation.id}\n` +
                `**File:** ${operation.filePath}\n` +
                `**Reason:** ${operation.reason}\n\n` +
                `## Diff Preview\n\n\`\`\`diff\n${operation.patch}\n\`\`\`\n\n` +
                `Use \`apply_diff\` with operation ID \`${operation.id}\` to apply these changes.`,
            },
          ],
        };
      }

      case 'apply_diff': {
        const { operation_id, approved } = ApplyDiffRequestSchema.parse(args);
        const result = await DiffManager.applyDiff(operation_id, approved);

        return {
          content: [
            {
              type: 'text',
              text:
                `# Diff Application Result\n\n` +
                `**Status:** ${result.success ? '✅ Success' : '❌ Failed'}\n` +
                `**Message:** ${result.message}`,
            },
          ],
        };
      }

      default:
        throw new Error(`Unknown tool: ${name}`);
    }
  } catch (error) {
    return {
      content: [
        {
          type: 'text',
          text: `❌ **Error:** ${error instanceof Error ? error.message : 'Unknown error occurred'}`,
        },
      ],
      isError: true,
    };
  }
});

async function main(): Promise<void> {
  const transport = new StdioServerTransport();
  await server.connect(transport);
}

main().catch((error) => {
  console.error('Server error:', error);
  process.exit(1);
});