0xray 2.1.2 → 2.1.3

package/src/mcps/knowledge-skills/session-management.server.ts

@@ -5,20 +5,9 @@
  * Provides session lifecycle management, state persistence, and cleanup utilities.
  */
 
-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 { XrayKnowledgeSkillBase } from "../shared/knowledge-skill-base.js";
 import { frameworkLogger } from "../../core/framework-logger.js";
 
-interface Tool {
-  name: string;
-  description: string;
-  inputSchema: object;
-}
-
 // In-memory session store for MCP server context
 interface SessionData {
   id: string;
@@ -28,181 +17,65 @@ interface SessionData {
   expiresAt: number | undefined;
 }
 
-class SessionManagementServer {
-  private server: Server;
+class SessionManagementServer extends XrayKnowledgeSkillBase {
   private sessions: Map<string, SessionData> = new Map();
-  private tools: Tool[] = [
-    {
-      name: "create_session",
-      description:
-        "Create a new session with optional data and expiration",
-      inputSchema: {
-        type: "object",
-        properties: {
-          sessionId: {
-            type: "string",
-            description: "Optional custom session ID (generated if not provided)",
-          },
-          data: {
-            type: "object",
-            description: "Initial session data",
-          },
-          ttl: {
-            type: "number",
-            description: "Time to live in seconds (optional)",
+
+  constructor() {
+    super("session-management", "2.0.1");
+    this.tools = [
+      {
+        name: "create_session",
+        description: "Create a new session with optional data and expiration",
+        inputSchema: {
+          type: "object",
+          properties: {
+            sessionId: { type: "string", description: "Optional custom session ID (generated if not provided)" },
+            data: { type: "object", description: "Initial session data" },
+            ttl: { type: "number", description: "Time to live in seconds (optional)" },
           },
         },
       },
-    },
-    {
-      name: "get_session",
-      description:
-        "Retrieve session data by session ID",
-      inputSchema: {
-        type: "object",
-        properties: {
-          sessionId: {
-            type: "string",
-            description: "Session ID to retrieve",
-          },
-        },
-        required: ["sessionId"],
+      {
+        name: "get_session",
+        description: "Retrieve session data by session ID",
+        inputSchema: { type: "object", properties: { sessionId: { type: "string", description: "Session ID to retrieve" } }, required: ["sessionId"] },
       },
-    },
-    {
-      name: "update_session",
-      description:
-        "Update session data",
-      inputSchema: {
-        type: "object",
-        properties: {
-          sessionId: {
-            type: "string",
-            description: "Session ID to update",
-          },
-          data: {
-            type: "object",
-            description: "Data to merge with existing session",
-          },
-          extendTtl: {
-            type: "number",
-            description: "Extend TTL by this many seconds",
-          },
-        },
-        required: ["sessionId", "data"],
+      {
+        name: "update_session",
+        description: "Update session data",
+        inputSchema: { type: "object", properties: { sessionId: { type: "string", description: "Session ID to update" }, data: { type: "object", description: "Data to merge with existing session" }, extendTtl: { type: "number", description: "Extend TTL by this many seconds" } }, required: ["sessionId", "data"] },
       },
-    },
-    {
-      name: "delete_session",
-      description:
-        "Delete a session",
-      inputSchema: {
-        type: "object",
-        properties: {
-          sessionId: {
-            type: "string",
-            description: "Session ID to delete",
-          },
-        },
-        required: ["sessionId"],
+      {
+        name: "delete_session",
+        description: "Delete a session",
+        inputSchema: { type: "object", properties: { sessionId: { type: "string", description: "Session ID to delete" } }, required: ["sessionId"] },
       },
-    },
-    {
-      name: "list_sessions",
-      description:
-        "List all active sessions",
-      inputSchema: {
-        type: "object",
-        properties: {
-          includeExpired: {
-            type: "boolean",
-            default: false,
-            description: "Include expired sessions in the list",
-          },
-          limit: {
-            type: "number",
-            default: 50,
-            description: "Maximum number of sessions to return",
-          },
-        },
+      {
+        name: "list_sessions",
+        description: "List all active sessions",
+        inputSchema: { type: "object", properties: { includeExpired: { type: "boolean", default: false, description: "Include expired sessions in the list" }, limit: { type: "number", default: 50, description: "Maximum number of sessions to return" } } },
       },
-    },
-    {
-      name: "session_exists",
-      description:
-        "Check if a session exists and is valid",
-      inputSchema: {
-        type: "object",
-        properties: {
-          sessionId: {
-            type: "string",
-            description: "Session ID to check",
-          },
-        },
-        required: ["sessionId"],
+      {
+        name: "session_exists",
+        description: "Check if a session exists and is valid",
+        inputSchema: { type: "object", properties: { sessionId: { type: "string", description: "Session ID to check" } }, required: ["sessionId"] },
       },
-    },
-    {
-      name: "cleanup_expired_sessions",
-      description:
-        "Clean up all expired sessions",
-      inputSchema: {
-        type: "object",
-        properties: {
-          dryRun: {
-            type: "boolean",
-            default: false,
-            description: "If true, only return what would be deleted without deleting",
-          },
-        },
+      {
+        name: "cleanup_expired_sessions",
+        description: "Clean up all expired sessions",
+        inputSchema: { type: "object", properties: { dryRun: { type: "boolean", default: false, description: "If true, only return what would be deleted without deleting" } } },
       },
-    },
-  ];
-
-  constructor() {
-    this.server = new Server(
-      { name: "session-management", version: "2.0.1" },
-      { capabilities: { tools: {} } },
-    );
-
-    this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
-      tools: this.tools,
-    }));
-
-    this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
-      const { name, arguments: args = {} } = request.params;
-
-      try {
-        switch (name) {
-          case "create_session":
-            return this.handleCreateSession(args);
-          case "get_session":
-            return this.handleGetSession(args);
-          case "update_session":
-            return this.handleUpdateSession(args);
-          case "delete_session":
-            return this.handleDeleteSession(args);
-          case "list_sessions":
-            return this.handleListSessions(args);
-          case "session_exists":
-            return this.handleSessionExists(args);
-          case "cleanup_expired_sessions":
-            return this.handleCleanupExpiredSessions(args);
-          default:
-            throw new Error(`Unknown tool: ${name}`);
-        }
-      } catch (error) {
-        return {
-          content: [
-            {
-              type: "text",
-              text: `Error: ${error instanceof Error ? error.message : String(error)}`,
-            },
-          ],
-          isError: true,
-        };
-      }
-    });
+    ];
+    this.handlers = {
+      "create_session": async (args) => this.handleCreateSession(args as Record<string, unknown>),
+      "get_session": async (args) => this.handleGetSession(args as Record<string, unknown>),
+      "update_session": async (args) => this.handleUpdateSession(args as Record<string, unknown>),
+      "delete_session": async (args) => this.handleDeleteSession(args as Record<string, unknown>),
+      "list_sessions": async (args) => this.handleListSessions(args as Record<string, unknown>),
+      "session_exists": async (args) => this.handleSessionExists(args as Record<string, unknown>),
+      "cleanup_expired_sessions": async (args) => this.handleCleanupExpiredSessions(args as Record<string, unknown>),
+    };
+    this.setupToolHandlers();
   }
 
   private generateSessionId(): string {
@@ -512,65 +385,11 @@ class SessionManagementServer {
     };
   }
 
-  async run(): Promise<void> {
-    const transport = new StdioServerTransport();
-    await this.server.connect(transport);
-
-    await frameworkLogger.log("mcp-session-management", "server-started", "success");
-
-    const cleanup = async (signal: string) => {
-      frameworkLogger.log("mcps/session-management", "shutdown", "info", { signal });
-
-      const timeout = setTimeout(() => {
-        frameworkLogger.log("mcps/session-management", "shutdown", "error", { message: "Graceful shutdown timeout, forcing exit..." });
-        process.exit(1);
-      }, 5000);
-
-      try {
-        if (this.server && typeof this.server.close === "function") {
-          await this.server.close();
-        }
-        clearTimeout(timeout);
-        frameworkLogger.log("mcps/session-management", "shutdown", "success");
-        process.exit(0);
-      } catch (error) {
-        clearTimeout(timeout);
-        frameworkLogger.log("mcps/session-management", "shutdown", "error", { message: `Error during server shutdown: ${String(error)}` });
-        process.exit(1);
-      }
-    };
-
-    process.on("SIGINT", () => cleanup("SIGINT"));
-    process.on("SIGTERM", () => cleanup("SIGTERM"));
-    process.on("SIGHUP", () => cleanup("SIGHUP"));
-
-    const checkParent = () => {
-      try {
-        process.kill(process.ppid, 0);
-        setTimeout(checkParent, 1000);
-      } catch {
-        frameworkLogger.log("mcps/session-management", "parent-death", "info");
-        cleanup("parent-process-death");
-      }
-    };
-
-    setTimeout(checkParent, 2000);
-
-    process.on("uncaughtException", (error) => {
-      frameworkLogger.log("mcps/session-management", "uncaughtException", "error", { error: String(error) });
-      cleanup("uncaughtException");
-    });
-
-    process.on("unhandledRejection", (reason, promise) => {
-      frameworkLogger.log("mcps/session-management", "unhandledRejection", "error", { error: String(reason) });
-      cleanup("unhandledRejection");
-    });
-  }
 }
 
 if (import.meta.url === `file://${process.argv[1]}`) {
   const server = new SessionManagementServer();
-  server.run().catch((error) => frameworkLogger.log("mcps/session-management", "run", "error", { error: String(error) }));
+  server.run("session-management").catch((error) => frameworkLogger.log("mcps/session-management", "run", "error", { error: String(error) }));
 }
 
 export default SessionManagementServer;