Update server.js

server.js

@@ -6,147 +6,305 @@ const fetch = (...args) => import('node-fetch').then(({default: fetch}) => fetch
 const app = express();
 const PORT = process.env.PORT || 7860;
 
-// 1. CONFIGURATION
-app.set('trust proxy', 1);
+// --- 1. CONFIGURATION ---
+app.set('trust proxy', 1); // Required for HF Spaces
 app.use(express.json({ limit: '10mb' }));
 
-// 2. CORS (Global)
+// --- 2. CORS (Global) ---
 app.use(cors({
   origin: '*',
   methods: ['POST', 'OPTIONS'],
   allowedHeaders: ['Content-Type', 'Authorization', 'X-Requested-With', 'MCP-Protocol-Version'],
 }));
 
-// 3. LOGGING
+// --- 3. AUTH MIDDLEWARE ---
+// Secures the MCP server so only you (via Gemini CLI) can use it.
+const authMiddleware = (req, res, next) => {
+  const expectedSecret = process.env.MCP_AUTH_TOKEN;
+  
+  // If no secret is set in env, allow open access (Warning logged)
+  if (!expectedSecret) {
+    // console.warn("[WARN] No MCP_AUTH_TOKEN set. Server is open.");
+    return next();
+  }
+
+  const authHeader = req.headers['authorization'];
+  if (!authHeader || !authHeader.startsWith('Bearer ') || authHeader.split(' ')[1] !== expectedSecret) {
+    console.error(`[AUTH] Unauthorized access attempt`);
+    return res.status(401).json({
+      jsonrpc: "2.0",
+      id: null,
+      error: { code: -32000, message: "Unauthorized: Invalid or missing MCP_AUTH_TOKEN" }
+    });
+  }
+  next();
+};
+
+// --- 4. LOGGING ---
 app.use((req, res, next) => {
   console.log(`[${new Date().toISOString()}] ${req.method} ${req.url}`);
   next();
 });
 
-// 4. THE UNIFIED MCP ENDPOINT
-app.post('/mcp', async (req, res) => {
+// --- 5. GITHUB API HELPER ---
+const callGitHub = async (endpoint, method = 'GET', body = null) => {
+  const token = process.env.GH_TOKEN;
+  if (!token) throw new Error("Server configuration error: GH_TOKEN is missing.");
+
+  const url = endpoint.startsWith('http') ? endpoint : `https://api.github.com${endpoint}`;
+  
+  const headers = {
+    'Authorization': `Bearer ${token}`,
+    'Accept': 'application/vnd.github.v3+json',
+    'User-Agent': 'HuggingFace-MCP-Server',
+    'X-GitHub-Api-Version': '2022-11-28'
+  };
+
+  const options = {
+    method,
+    headers,
+    body: body ? JSON.stringify(body) : undefined
+  };
+
+  const response = await fetch(url, options);
+
+  if (!response.ok) {
+    const errorText = await response.text();
+    throw new Error(`GitHub API ${response.status}: ${errorText}`);
+  }
+
+  return response.json();
+};
+
+// --- 6. MAIN MCP ENDPOINT ---
+app.post('/mcp', authMiddleware, async (req, res) => {
   try {
-    const body = req.body;
-    const { id, method, params } = body;
+    const { id, method, params } = req.body;
+    const requestId = id !== undefined ? id : null;
+
+    // Helper to send success
+    const send = (result) => res.json({ jsonrpc: "2.0", id: requestId, result });
+    
+    // Helper to send error
+    const error = (code, msg) => res.json({ jsonrpc: "2.0", id: requestId, error: { code, message: msg } });
 
-    console.log(`[POST] Method: ${method} | ID: ${id}`);
+    console.log(`[POST] Method: ${method}`);
 
-    // --- A. INITIALIZE ---
+    // --- A. HANDSHAKE & LIFECYCLE ---
     if (method === 'initialize') {
-      return res.json({
-        jsonrpc: "2.0",
-        id: id,
-        result: {
-          protocolVersion: "2025-06-18",
-          capabilities: { 
-            tools: { listChanged: false }
-          },
-          serverInfo: { name: "github-mcp-hf", version: "1.0.0" }
-        }
+      return send({
+        protocolVersion: "2024-11-05",
+        capabilities: { tools: {}, resources: {}, prompts: {} },
+        serverInfo: { name: "github-mcp-versatile", version: "2.0.0" }
       });
     }
 
-    // --- B. NOTIFICATIONS / INITIALIZED (The Fix) ---
-    // This was causing the 400 error. We must acknowledge it.
-    if (method === 'notifications/initialized') {
-      // JSON-RPC notifications don't technically require a response, 
-      // but over HTTP we must return 200 OK to finish the request.
-      return res.json({
-        jsonrpc: "2.0",
-        id: id || null,
-        result: true
-      });
-    }
+    if (method === 'notifications/initialized') return send(true);
+    
+    // Crucial for Gemini CLI connection stability
+    if (method === 'ping') return send({});
 
-    // --- C. LIST TOOLS ---
+    // --- B. TOOL LIST ---
     if (method === 'tools/list') {
-      return res.json({
-        jsonrpc: "2.0",
-        id: id,
-        result: {
-          tools: [{
-            name: "getRepoIssues",
-            description: "Get recent issues from a GitHub repository.",
+      return send({
+        tools: [
+          // 1. ISSUES
+          {
+            name: "list_issues",
+            description: "List issues in a repo.",
             inputSchema: {
               type: "object",
               properties: {
                 owner: { type: "string" },
                 repo: { type: "string" },
-                count: { type: "number", default: 5 }
+                state: { type: "string", enum: ["open", "closed", "all"], default: "open" }
               },
               required: ["owner", "repo"]
             }
-          }]
-        }
+          },
+          {
+            name: "create_issue",
+            description: "Create a new issue.",
+            inputSchema: {
+              type: "object",
+              properties: {
+                owner: { type: "string" },
+                repo: { type: "string" },
+                title: { type: "string" },
+                body: { type: "string" }
+              },
+              required: ["owner", "repo", "title"]
+            }
+          },
+          // 2. FILES & CODE
+          {
+            name: "read_file",
+            description: "Read the contents of a specific file.",
+            inputSchema: {
+              type: "object",
+              properties: {
+                owner: { type: "string" },
+                repo: { type: "string" },
+                path: { type: "string", description: "Path to file (e.g. src/index.js)" },
+                branch: { type: "string", description: "Branch or commit SHA (default: main)" }
+              },
+              required: ["owner", "repo", "path"]
+            }
+          },
+          {
+            name: "list_files",
+            description: "Get the file tree (recursive) to understand repo structure.",
+            inputSchema: {
+              type: "object",
+              properties: {
+                owner: { type: "string" },
+                repo: { type: "string" },
+                branch: { type: "string", default: "main" },
+                recursive: { type: "boolean", default: true }
+              },
+              required: ["owner", "repo"]
+            }
+          },
+          {
+            name: "search_code",
+            description: "Search for code within a repository.",
+            inputSchema: {
+              type: "object",
+              properties: {
+                q: { type: "string", description: "Search query (e.g. 'function login')" },
+                owner: { type: "string" },
+                repo: { type: "string" }
+              },
+              required: ["q", "owner", "repo"]
+            }
+          },
+          // 3. PULL REQUESTS
+          {
+            name: "list_prs",
+            description: "List Pull Requests.",
+            inputSchema: {
+              type: "object",
+              properties: {
+                owner: { type: "string" },
+                repo: { type: "string" },
+                state: { type: "string", enum: ["open", "closed", "all"], default: "open" }
+              },
+              required: ["owner", "repo"]
+            }
+          },
+          {
+            name: "create_pr",
+            description: "Create a Pull Request.",
+            inputSchema: {
+              type: "object",
+              properties: {
+                owner: { type: "string" },
+                repo: { type: "string" },
+                title: { type: "string" },
+                head: { type: "string", description: "The name of the branch where your changes are implemented." },
+                base: { type: "string", description: "The name of the branch you want the changes pulled into." },
+                body: { type: "string" }
+              },
+              required: ["owner", "repo", "title", "head", "base"]
+            }
+          }
+        ]
       });
     }
 
-    // --- D. CALL TOOL ---
+    // --- C. TOOL EXECUTION ---
     if (method === 'tools/call') {
       const { name, arguments: args } = params || {};
       
-      if (name === 'getRepoIssues') {
-        const ghToken = process.env.GH_TOKEN;
-        if (!ghToken) {
-          return res.json({
-            jsonrpc: "2.0", id,
-            error: { code: -32000, message: "GH_TOKEN missing" }
+      try {
+        let resultData;
+
+        // --- Logic: ISSUES ---
+        if (name === "list_issues") {
+          const data = await callGitHub(`/repos/${args.owner}/${args.repo}/issues?state=${args.state || 'open'}&per_page=10`);
+          resultData = data.map(i => ({ number: i.number, title: i.title, state: i.state, url: i.html_url }));
+        }
+        else if (name === "create_issue") {
+          const data = await callGitHub(`/repos/${args.owner}/${args.repo}/issues`, 'POST', {
+            title: args.title, body: args.body
           });
+          resultData = { message: "Issue Created", url: data.html_url, number: data.number };
         }
 
-        const ghUrl = `https://api.github.com/repos/${args.owner}/${args.repo}/issues?per_page=${args.count || 5}&state=all`;
-        
-        const ghRes = await fetch(ghUrl, {
-          headers: {
-            'Authorization': `Bearer ${ghToken}`,
-            'User-Agent': 'HuggingFace-MCP',
-            'Accept': 'application/vnd.github+json',
-            'X-GitHub-Api-Version': '2022-11-28'
-          }
-        });
+        // --- Logic: FILES ---
+        else if (name === "read_file") {
+          // GitHub API returns content in Base64
+          const ref = args.branch ? `?ref=${args.branch}` : '';
+          const data = await callGitHub(`/repos/${args.owner}/${args.repo}/contents/${args.path}${ref}`);
+          
+          if (Array.isArray(data)) return error(-32000, "Path is a directory, not a file.");
+          if (!data.content) return error(-32000, "File content empty or binary.");
+
+          const decodedContent = Buffer.from(data.content, 'base64').toString('utf-8');
+          // Truncate if huge to save tokens
+          const safeContent = decodedContent.length > 20000 ? decodedContent.substring(0, 20000) + "\n...[Truncated]" : decodedContent;
+          resultData = safeContent;
+        }
+        else if (name === "list_files") {
+          const ref = args.branch || 'main';
+          // Get the tree SHA first, or just assume main/master
+          // To be simple, we use the trees API which requires the SHA of the branch.
+          // Fallback: Use the 'contents' API recursively? No, that's heavy.
+          // Better approach for LLM: Get the HEAD commit SHA first.
+          
+          // 1. Get Branch SHA
+          const branchData = await callGitHub(`/repos/${args.owner}/${args.repo}/git/ref/heads/${ref}`);
+          const sha = branchData.object.sha;
 
-        if (!ghRes.ok) {
-          const errText = await ghRes.text();
-          return res.json({
-            jsonrpc: "2.0", id,
-            error: { code: -32000, message: `GitHub API Error: ${ghRes.status}` }
+          // 2. Get Tree
+          const recursive = args.recursive ? '?recursive=1' : '';
+          const treeData = await callGitHub(`/repos/${args.owner}/${args.repo}/git/trees/${sha}${recursive}`);
+          
+          // Filter to just paths to save tokens
+          resultData = treeData.tree.map(t => t.path).slice(0, 200); // Limit to 200 files
+        }
+        else if (name === "search_code") {
+          const q = `${args.q} repo:${args.owner}/${args.repo}`;
+          const data = await callGitHub(`/search/code?q=${encodeURIComponent(q)}&per_page=5`);
+          resultData = data.items.map(i => ({ name: i.name, path: i.path, url: i.html_url }));
+        }
+
+        // --- Logic: PULL REQUESTS ---
+        else if (name === "list_prs") {
+          const data = await callGitHub(`/repos/${args.owner}/${args.repo}/pulls?state=${args.state || 'open'}&per_page=10`);
+          resultData = data.map(i => ({ number: i.number, title: i.title, state: i.state, url: i.html_url }));
+        }
+        else if (name === "create_pr") {
+          const data = await callGitHub(`/repos/${args.owner}/${args.repo}/pulls`, 'POST', {
+            title: args.title, head: args.head, base: args.base, body: args.body
           });
+          resultData = { message: "PR Created", url: data.html_url, number: data.number };
+        }
+        
+        else {
+          return error(-32601, `Tool ${name} not found`);
         }
 
-        const data = await ghRes.json();
-        const content = data.map(i => ({
-          id: i.number, title: i.title, state: i.state, url: i.html_url
-        }));
-
-        return res.json({
-          jsonrpc: "2.0",
-          id,
-          result: {
-            content: [{ type: "text", text: JSON.stringify(content, null, 2) }],
-            isError: false
-          }
+        // Return Standard Success Response
+        return send({
+          content: [{ type: "text", text: typeof resultData === 'string' ? resultData : JSON.stringify(resultData, null, 2) }],
+          isError: false
         });
+
+      } catch (err) {
+        return error(-32000, `Execution Error: ${err.message}`);
       }
     }
 
-    // --- E. UNKNOWN METHOD ---
-    console.warn(`Unknown method: ${method}`);
-    res.status(400).json({
-      jsonrpc: "2.0", id,
-      error: { code: -32601, message: "Method not found" }
-    });
+    return error(-32601, "Method not found");
 
-  } catch (error) {
-    console.error(error);
-    res.status(500).json({ error: "Internal Server Error" });
+  } catch (err) {
+    console.error(err);
+    res.status(500).json({ jsonrpc: "2.0", id: null, error: { code: -32603, message: "Internal Server Error" } });
   }
 });
 
-// 5. FALLBACK / HEALTH CHECK
-app.use((req, res) => {
-  res.status(200).send("MCP Server Active");
-});
+// 7. HEALTH CHECK
+app.use((req, res) => res.status(200).send("GitHub MCP Active"));
 
-app.listen(PORT, '0.0.0.0', () => {
-  console.log(`MCP Server running on port ${PORT}`);
-});
+app.listen(PORT, '0.0.0.0', () => console.log(`Server running on ${PORT}`));