Skip to content

Commit 4c192f9

Browse files
PV-BhatPV-Bhat
authored
Update index.ts (Refactored src/index.ts using server.tool())
1 parent da0aa42 commit 4c192f9

File tree

1 file changed

+91
-123
lines changed

1 file changed

+91
-123
lines changed

‎src/index.ts

Lines changed: 91 additions & 123 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,15 @@
11
// src/index.ts
2-
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
2+
// Try importing McpServer from server/index.js again, now using .tool()
3+
import { McpServer } from "@modelcontextprotocol/sdk/server/index.js";
34
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
4-
import {
5-
ListToolsRequestSchema,
6-
CallToolRequestSchema,
7-
CallToolRequest,
8-
InitializeRequest,
9-
InitializeRequestSchema
10-
// Removed ToolDefinition import as it's not found
11-
} from "@modelcontextprotocol/sdk/types.js";
12-
import { vibeCheckTool } from "./tools/vibeCheck.js";
13-
import { vibeDistillTool } from "./tools/vibeDistill.js";
14-
import { vibeLearnTool, VibeLearnOutput } from "./tools/vibeLearn.js";
5+
// Import Zod for defining input schemas with server.tool()
6+
import { z } from "zod";
7+
// Import specific request types if needed for handlers, otherwise handler args are inferred
8+
import { CallToolRequest } from "@modelcontextprotocol/sdk/types.js";
9+
10+
import { vibeCheckTool, VibeCheckInput } from "./tools/vibeCheck.js"; // Assuming VibeCheckInput is exported
11+
import { vibeDistillTool, VibeDistillInput } from "./tools/vibeDistill.js"; // Assuming VibeDistillInput is exported
12+
import { vibeLearnTool, VibeLearnInput, VibeLearnOutput } from "./tools/vibeLearn.js"; // Assuming VibeLearnInput is exported
1513
import { MistakeEntry } from "./utils/storage.js";
1614
import { initializeGemini } from "./utils/gemini.js";
1715

@@ -37,122 +35,94 @@ type CategorySummaryItem = {
3735
recentExample: MistakeEntry;
3836
};
3937

40-
console.error("MCP Server: Creating Server instance...");
41-
const server = new Server({
38+
console.error("MCP Server: Creating McpServer instance...");
39+
// Use McpServer class again
40+
const server = new McpServer({
4241
name: "vibe-check-mcp",
4342
version: "0.2.0"
4443
});
45-
console.error("MCP Server: Server instance created.");
44+
console.error("MCP Server: McpServer instance created.");
4645

47-
// --- Initialize handler ---
48-
server.setRequestHandler(InitializeRequestSchema, async (request: InitializeRequest) => {
49-
console.error(`MCP Server: Received initialize request`);
50-
const response = {
51-
protocolVersion: request.params.protocolVersion,
52-
serverInfo: {
53-
name: "vibe-check-mcp",
54-
version: "0.2.0",
55-
},
56-
capabilities: {}
57-
};
58-
console.error(`MCP Server: Sending initialize response`);
59-
return response;
60-
});
61-
// -----------------------------------------------------
46+
// --- Define Tools using server.tool() ---
6247

63-
// Define tools using ListToolsRequestSchema handler
64-
console.error("MCP Server: Setting ListTools request handler...");
65-
server.setRequestHandler(ListToolsRequestSchema, async () => {
66-
console.error("MCP Server: ListTools request received.");
67-
// Remove explicit ToolDefinition[] type, let TypeScript infer
68-
const tools = [
69-
{
70-
name: "vibe_check",
71-
description: "Metacognitive check for plan alignment and assumption testing.",
72-
inputSchema: {
73-
type: "object",
74-
properties: {
75-
plan: { type: "string" },
76-
userRequest: { type: "string" },
77-
// ... other properties ...
78-
confidence: { type: "number" }
79-
},
80-
required: ["plan", "userRequest"]
81-
}
82-
},
83-
{
84-
name: "vibe_distill",
85-
description: "Distills a plan to its essential core.",
86-
inputSchema: {
87-
type: "object",
88-
properties: {
89-
plan: { type: "string" },
90-
userRequest: { type: "string" },
91-
sessionId: { type: "string" }
92-
},
93-
required: ["plan", "userRequest"]
94-
}
95-
},
96-
{
97-
name: "vibe_learn",
98-
description: "Logs mistake patterns for future improvement.",
99-
inputSchema: {
100-
type: "object",
101-
properties: {
102-
mistake: { type: "string" },
103-
category: { type: "string" },
104-
solution: { type: "string" },
105-
sessionId: { type: "string" }
106-
},
107-
required: ["mistake", "category", "solution"]
108-
}
48+
console.error("MCP Server: Defining 'vibe_check' tool...");
49+
server.tool(
50+
"vibe_check",
51+
// Define input schema using Zod
52+
z.object({
53+
plan: z.string(),
54+
userRequest: z.string(),
55+
thinkingLog: z.string().optional(),
56+
availableTools: z.array(z.string()).optional(),
57+
focusAreas: z.array(z.string()).optional(),
58+
sessionId: z.string().optional(),
59+
previousAdvice: z.string().optional(),
60+
phase: z.enum(["planning", "implementation", "review"]).optional(),
61+
confidence: z.number().optional()
62+
}).required({ // Explicitly state required fields matching your TS interface
63+
plan: true,
64+
userRequest: true
65+
}),
66+
// Handler function - arguments are automatically typed based on schema
67+
async (args) => {
68+
console.error(`MCP Server: Executing tool: vibe_check...`);
69+
// Args type should match VibeCheckInput here due to Zod schema
70+
const result = await vibeCheckTool(args);
71+
console.error(`MCP Server: Tool vibe_check executed successfully.`);
72+
// Return result formatted as MCP content
73+
return { content: [{ type: "text", text: result.questions + (result.patternAlert ? `\n\n**Pattern Alert:** ${result.patternAlert}`: "") }] };
10974
}
110-
];
111-
console.error("MCP Server: Returning tool list.");
112-
return { tools };
113-
});
114-
console.error("MCP Server: ListTools request handler set.");
115-
116-
// Handle tool calls using CallToolRequestSchema handler
117-
console.error("MCP Server: Setting CallTool request handler...");
118-
server.setRequestHandler(CallToolRequestSchema, async (request: CallToolRequest) => {
119-
const toolName = request.params.name;
120-
const args = request.params.arguments ?? {};
75+
);
76+
console.error("MCP Server: 'vibe_check' tool defined.");
12177

122-
console.error(`MCP Server: CallToolRequest received for tool: ${toolName}`);
78+
console.error("MCP Server: Defining 'vibe_distill' tool...");
79+
server.tool(
80+
"vibe_distill",
81+
z.object({
82+
plan: z.string(),
83+
userRequest: z.string(),
84+
sessionId: z.string().optional()
85+
}).required({
86+
plan: true,
87+
userRequest: true
88+
}),
89+
async (args) => {
90+
console.error(`MCP Server: Executing tool: vibe_distill...`);
91+
const result = await vibeDistillTool(args);
92+
console.error(`MCP Server: Tool vibe_distill executed successfully.`);
93+
return { content: [{ type: "markdown", markdown: result.distilledPlan + `\n\n**Rationale:** ${result.rationale}` }] };
94+
}
95+
);
96+
console.error("MCP Server: 'vibe_distill' tool defined.");
12397

124-
try {
125-
let result: any;
126-
console.error(`MCP Server: Executing tool: ${toolName}...`);
127-
switch (toolName) {
128-
case 'vibe_check':
129-
result = await vibeCheckTool(args as any);
130-
console.error(`MCP Server: Tool ${toolName} executed successfully.`);
131-
return { content: [{ type: "text", text: result.questions + (result.patternAlert ? `\n\n**Pattern Alert:** ${result.patternAlert}`: "") }] };
132-
case 'vibe_distill':
133-
result = await vibeDistillTool(args as any);
134-
console.error(`MCP Server: Tool ${toolName} executed successfully.`);
135-
return { content: [{ type: "markdown", markdown: result.distilledPlan + `\n\n**Rationale:** ${result.rationale}` }] };
136-
case 'vibe_learn':
137-
result = await vibeLearnTool(args as any) as VibeLearnOutput;
138-
console.error(`MCP Server: Tool ${toolName} executed successfully.`);
98+
console.error("MCP Server: Defining 'vibe_learn' tool...");
99+
server.tool(
100+
"vibe_learn",
101+
z.object({
102+
mistake: z.string(),
103+
category: z.string(),
104+
solution: z.string(),
105+
sessionId: z.string().optional()
106+
}).required({
107+
mistake: true,
108+
category: true,
109+
solution: true
110+
}),
111+
async (args) => {
112+
console.error(`MCP Server: Executing tool: vibe_learn...`);
113+
const result = await vibeLearnTool(args) as VibeLearnOutput; // Cast needed if return type isn't perfectly inferred
114+
console.error(`MCP Server: Tool vibe_learn executed successfully.`);
139115
const summary = result.topCategories.map((cat: CategorySummaryItem) => `- ${cat.category} (${cat.count})`).join('\n');
140116
return { content: [{ type: "text", text: `✅ Pattern logged. Tally for category: ${result.currentTally}.\nTop Categories:\n${summary}` }] };
141-
default:
142-
console.error(`MCP Server: Tool '${toolName}' not found.`);
143-
throw new Error(`Tool '${toolName}' not found.`);
144117
}
145-
} catch (error: any) {
146-
console.error(`MCP Server: Error executing tool ${toolName}:`, error);
147-
return {
148-
error: {
149-
code: "tool_execution_error",
150-
message: `Error executing tool ${toolName}: ${error.message}`
151-
}
152-
};
153-
}
154-
});
155-
console.error("MCP Server: CallTool request handler set.");
118+
);
119+
console.error("MCP Server: 'vibe_learn' tool defined.");
120+
121+
// --- Remove setRequestHandler calls for ListTools and CallTool ---
122+
// The server.tool() calls handle registration implicitly.
123+
124+
// --- Remove explicit Initialize handler ---
125+
// SDK likely handles this when tools are defined via server.tool()
156126

157127
// Create the transport instance
158128
console.error("MCP Server: Creating StdioServerTransport...");
@@ -161,16 +131,14 @@ console.error("MCP Server: StdioServerTransport created.");
161131

162132
// Connect the server and transport
163133
console.error("MCP Server: Connecting server to transport...");
164-
server.connect(transport);
134+
server.connect(transport); // Use server.connect() as before
165135
console.error("MCP Server: Server connected to transport. Ready for messages.");
166136

167-
// Assign callback to 'onclose' property instead of calling it
168-
if (transport.onclose) { // Check if property exists
169-
transport.onclose = () => { // Assign the callback function
137+
// Assign callback to 'onclose' property
138+
if (transport.onclose) {
139+
transport.onclose = () => {
170140
console.error("MCP Server: Transport closed event received.");
171141
};
172142
} else {
173143
console.error("MCP Server: transport.onclose property not found.");
174144
}
175-
176-
// Removed listener for non-existent onDidDispose

0 commit comments

Comments
 (0)