fix: screenshot bug

Author: Kylejeong2

src/index.ts

@@ -139,6 +139,17 @@ export default function ({ config }: { config: z.infer<typeof configSchema> }) {
   const contextId = randomUUID();
   const context = new Context(server.server, internalConfig, contextId);
 
+  // Cleanup handler for when the MCP connection closes (SHTTP/Smithery)
+  server.server.onclose = async () => {
+    try {
+      await context.getSessionManager().closeAllSessions();
+    } catch (err) {
+      process.stderr.write(
+        `[Cleanup] Error during session cleanup: ${err instanceof Error ? err.message : String(err)}\n`,
+      );
+    }
+  };
+
   server.server.registerCapabilities({
     resources: {
       subscribe: true,

src/transport.ts

@@ -53,10 +53,11 @@ async function handleStreamable(
       sessionIdGenerator: () => sessionId,
     });
     sessions.set(sessionId, transport);
+    const server = await serverList.create();
     transport.onclose = () => {
       if (transport.sessionId) sessions.delete(transport.sessionId);
+      serverList.close(server);
     };
-    const server = await serverList.create();
     await server.connect(transport);
     return await transport.handleRequest(req, res);
   }