Update index.ts (Docker fixing logs)

Author: PV-Bhat

src/index.ts

@@ -1,34 +1,38 @@
 // src/index.ts
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
-import { ListToolsRequestSchema, CallToolRequestSchema, CallToolRequest, InitializeRequest } from "@modelcontextprotocol/sdk/types.js"; // Added InitializeRequest
+import {
+    ListToolsRequestSchema,
+    CallToolRequestSchema,
+    CallToolRequest,
+    InitializeRequest,
+    InitializeRequestSchema, // <-- Import InitializeRequestSchema
+    ToolDefinition // <-- Import ToolDefinition
+} from "@modelcontextprotocol/sdk/types.js";
 import { vibeCheckTool } from "./tools/vibeCheck.js";
 import { vibeDistillTool } from "./tools/vibeDistill.js";
 import { vibeLearnTool, VibeLearnOutput } from "./tools/vibeLearn.js";
 import { MistakeEntry } from "./utils/storage.js";
-import { initializeGemini } from "./utils/gemini.js"; // Import gemini initializer
+import { initializeGemini } from "./utils/gemini.js";
 
-console.error("MCP Server: Script starting..."); // Log to stderr
+console.error("MCP Server: Script starting...");
 
 // --- Explicitly Initialize Gemini ---
 try {
-    const apiKey = process.env.GEMINI_API_KEY; // Read from env var
+    const apiKey = process.env.GEMINI_API_KEY;
     if (!apiKey) {
-        console.error("MCP Server: ERROR - GEMINI_API_KEY environment variable not found!");
-        // Optionally exit if the key is absolutely required at startup
-        // process.exit(1);
+        console.error("MCP Server: WARNING - GEMINI_API_KEY environment variable not found!");
+        // Continue without Gemini if key is missing, tools might fail later
     } else {
         initializeGemini(apiKey);
-        console.error("MCP Server: Gemini client potentially initialized (check gemini.ts logs).");
+        console.error("MCP Server: Gemini client potentially initialized.");
     }
 } catch (err: any) {
     console.error("MCP Server: ERROR initializing Gemini client -", err);
-    // Optionally exit
-    // process.exit(1);
+    // Continue anyway, tools using it will fail
 }
 // ------------------------------------
 
-
 // Define the type for the 'cat' parameter explicitly
 type CategorySummaryItem = {
   category: string;
@@ -44,47 +48,90 @@ const server = new Server({
 });
 console.error("MCP Server: Server instance created.");
 
-// --- Add an explicit Initialize handler for logging ---
-server.setRequestHandler("initialize", async (request: InitializeRequest) => {
-    console.error(`MCP Server: Received initialize request: ID=${request.id}`);
-    // Basic response - the SDK might override parts of this, but it confirms handling
+// --- Add an explicit Initialize handler ---
+// Use the Schema as the first argument, not the string "initialize"
+server.setRequestHandler(InitializeRequestSchema, async (request: InitializeRequest) => {
+    // Cannot reliably log request.id here based on type errors
+    console.error(`MCP Server: Received initialize request`);
     const response = {
-        protocolVersion: "2024-11-05", // Use the version client sent or latest known
+        protocolVersion: request.params.protocolVersion, // Echo back client's version
         serverInfo: {
             name: "vibe-check-mcp",
             version: "0.2.0",
         },
-        capabilities: {
-             // Declare capabilities if needed, e.g., for notifications
-        }
+        capabilities: {}
     };
-    console.error(`MCP Server: Sending initialize response for ID=${request.id}`);
+    console.error(`MCP Server: Sending initialize response`);
     return response;
 });
 // -----------------------------------------------------
 
-
 // Define tools using ListToolsRequestSchema handler
 console.error("MCP Server: Setting ListTools request handler...");
 server.setRequestHandler(ListToolsRequestSchema, async () => {
   console.error("MCP Server: ListTools request received.");
-  // Your tool definitions...
-  const tools = [ /* ... your tool definitions from previous code ... */ ];
+  // Explicitly type the tools array
+  const tools: ToolDefinition[] = [
+     {
+      name: "vibe_check",
+      description: "Metacognitive check for plan alignment and assumption testing.",
+      inputSchema: {
+        type: "object",
+        properties: {
+          plan: { type: "string" },
+          userRequest: { type: "string" },
+          thinkingLog: { type: "string" },
+          availableTools: { type: "array", items: { type: "string" } },
+          focusAreas: { type: "array", items: { type: "string" } },
+          sessionId: { type: "string" },
+          previousAdvice: { type: "string" },
+          phase: { type: "string", enum: ["planning", "implementation", "review"] },
+          confidence: { type: "number" }
+        },
+        required: ["plan", "userRequest"]
+      }
+    },
+    {
+      name: "vibe_distill",
+      description: "Distills a plan to its essential core.",
+      inputSchema: {
+        type: "object",
+        properties: {
+            plan: { type: "string" },
+            userRequest: { type: "string" },
+            sessionId: { type: "string" }
+        },
+        required: ["plan", "userRequest"]
+      }
+    },
+    {
+        name: "vibe_learn",
+        description: "Logs mistake patterns for future improvement.",
+        inputSchema: {
+            type: "object",
+            properties: {
+                mistake: { type: "string" },
+                category: { type: "string" },
+                solution: { type: "string" },
+                sessionId: { type: "string" }
+            },
+            required: ["mistake", "category", "solution"]
+        }
+    }
+  ];
    console.error("MCP Server: Returning tool list.");
-   return { tools };
+   return { tools }; // Type 'tools' variable here
 });
 console.error("MCP Server: ListTools request handler set.");
 
-
 // Handle tool calls using CallToolRequestSchema handler
 console.error("MCP Server: Setting CallTool request handler...");
 server.setRequestHandler(CallToolRequestSchema, async (request: CallToolRequest) => {
   const toolName = request.params.name;
   const args = request.params.arguments ?? {};
 
-  // Log received request details to stderr
-  console.error(`MCP Server: CallToolRequest received for tool: ${toolName} with ID: ${request.id}`);
-  // console.error(`MCP Server: Args: ${JSON.stringify(args)}`); // Be careful logging args if sensitive
+  // Cannot reliably log request.id here based on type errors
+  console.error(`MCP Server: CallToolRequest received for tool: ${toolName}`);
 
   try {
     let result: any;
@@ -94,7 +141,15 @@ server.setRequestHandler(CallToolRequestSchema, async (request: CallToolRequest)
         result = await vibeCheckTool(args as any);
         console.error(`MCP Server: Tool ${toolName} executed successfully.`);
         return { content: [{ type: "text", text: result.questions + (result.patternAlert ? `\n\n**Pattern Alert:** ${result.patternAlert}`: "") }] };
-      // ... other cases ...
+      case 'vibe_distill':
+         result = await vibeDistillTool(args as any);
+         console.error(`MCP Server: Tool ${toolName} executed successfully.`);
+         return { content: [{ type: "markdown", markdown: result.distilledPlan + `\n\n**Rationale:** ${result.rationale}` }] };
+      case 'vibe_learn':
+        result = await vibeLearnTool(args as any) as VibeLearnOutput;
+        console.error(`MCP Server: Tool ${toolName} executed successfully.`);
+        const summary = result.topCategories.map((cat: CategorySummaryItem) => `- ${cat.category} (${cat.count})`).join('\n');
+        return { content: [{ type: "text", text: `✅ Pattern logged. Tally for category: ${result.currentTally}.\nTop Categories:\n${summary}` }] };
       default:
         console.error(`MCP Server: Tool '${toolName}' not found.`);
         throw new Error(`Tool '${toolName}' not found.`);
@@ -111,7 +166,6 @@ server.setRequestHandler(CallToolRequestSchema, async (request: CallToolRequest)
 });
 console.error("MCP Server: CallTool request handler set.");
 
-
 // Create the transport instance
 console.error("MCP Server: Creating StdioServerTransport...");
 const transport = new StdioServerTransport();
@@ -122,10 +176,8 @@ console.error("MCP Server: Connecting server to transport...");
 server.connect(transport);
 console.error("MCP Server: Server connected to transport. Ready for messages.");
 
-// Optional: Add listeners for transport errors/close events
-transport.onDidClose(() => {
+// Use the corrected event name 'onclose'
+transport.onclose(() => { // <-- Corrected event name
     console.error("MCP Server: Transport closed event received.");
 });
-transport.onDidDispose(() => {
-    console.error("MCP Server: Transport disposed event received.");
-});
+// Removed listener for non-existent onDidDispose