@gethmy/mcp 2.4.4 → 2.4.7

package/src/remote.ts

@@ -16,6 +16,7 @@
 
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
+import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
 import { serve } from "bun";
 import { Hono } from "hono";
 import { cors } from "hono/cors";
@@ -42,12 +43,52 @@ interface TokenInfo {
   source: "api_key" | "oauth";
 }
 
+// Tiny TTL cache for /v1/auth/context. A refresh storm (Claude rotates a
+// token, fires several queued tool calls in parallel) would otherwise hammer
+// harmony-api with identical lookups. 30s is short enough that revocation
+// propagates quickly, long enough to absorb any normal burst.
+const TOKEN_CACHE_TTL_MS = 30_000;
+const TOKEN_CACHE_MAX = 1000;
+const tokenCache = new Map<string, { info: TokenInfo; expiresAt: number }>();
+
+async function tokenFingerprint(token: string): Promise<string> {
+  const data = new TextEncoder().encode(token);
+  const buf = await crypto.subtle.digest("SHA-256", data);
+  return Array.from(new Uint8Array(buf))
+    .map((b) => b.toString(16).padStart(2, "0"))
+    .join("");
+}
+
+function evictTokenCacheLru(): void {
+  if (tokenCache.size <= TOKEN_CACHE_MAX) return;
+  const oldest = tokenCache.keys().next().value;
+  if (oldest) tokenCache.delete(oldest);
+}
+
 async function validateToken(token: string): Promise<TokenInfo | null> {
+  const fp = await tokenFingerprint(token);
+  const cached = tokenCache.get(fp);
+  const now = Date.now();
+  if (cached && cached.expiresAt > now) {
+    // Refresh LRU position
+    tokenCache.delete(fp);
+    tokenCache.set(fp, cached);
+    return cached.info;
+  }
+
   try {
     const response = await fetch(`${HARMONY_API_URL}/v1/auth/context`, {
       headers: { "X-API-Key": token },
     });
-    if (!response.ok) return null;
+    if (!response.ok) {
+      // Negative-cache 401s briefly so a flood of bad tokens doesn't
+      // pummel harmony-api. Other status codes (5xx, network) bypass the
+      // cache so transient failures self-heal on next call.
+      if (response.status === 401) {
+        tokenCache.delete(fp);
+      }
+      return null;
+    }
 
     const data = (await response.json()) as {
       userId: string;
@@ -55,16 +96,27 @@ async function validateToken(token: string): Promise<TokenInfo | null> {
       workspaceId: string | null;
     };
     if (data.source === "jwt") return null; // JWT not allowed on MCP endpoint
-    return {
+    const info: TokenInfo = {
       userId: data.userId,
       workspaceId: data.workspaceId,
       source: data.source,
     };
+    tokenCache.set(fp, { info, expiresAt: now + TOKEN_CACHE_TTL_MS });
+    evictTokenCacheLru();
+    return info;
   } catch {
     return null;
   }
 }
 
+function invalidateTokenCache(token: string): void {
+  // Fire-and-forget — just clear if we already had it. Used after a bearer
+  // is rejected by harmony-api so the next attempt re-validates fresh.
+  tokenFingerprint(token)
+    .then((fp) => tokenCache.delete(fp))
+    .catch(() => {});
+}
+
 // For legacy api_keys: fall back to workspaces list to pick the first one.
 // OAuth tokens already have workspaceId bound at grant time.
 async function resolveWorkspaceForLegacyKey(
@@ -90,21 +142,39 @@ async function resolveWorkspaceForLegacyKey(
 interface McpSession {
   transport: WebStandardStreamableHTTPServerTransport;
   server: Server;
+  client: HarmonyApiClient;
   apiKey: string;
+  // Bound at session creation. Re-checked on every token hot-swap so a leaked
+  // session ID can't be paired with a different user's bearer to ride someone
+  // else's session.
+  userId: string;
   activeWorkspaceId: string | null;
   activeProjectId: string | null;
   createdAt: number;
+  // Bumped on every request that touches the session. Drives stale-session GC
+  // so a session that's actively rotating tokens stays alive past the access
+  // token TTL — auth happens per-request via the Bearer header, the session
+  // itself is just a transport handle.
+  lastUsedAt: number;
+  // Set by HarmonyApiClient.onUnauthorized when the API rejects the cached
+  // token mid-session. The HTTP layer reads this after transport.handleRequest
+  // returns and converts the response to an HTTP 401 challenge so the OAuth
+  // client refreshes instead of caching a JSON-RPC error forever.
+  unauthorized: boolean;
 }
 
 const sessions = new Map<string, McpSession>();
 
-// Clean up stale sessions every 30 minutes
+// Stale-session GC. Uses lastUsedAt (sliding window) instead of createdAt so
+// long-lived clients that refresh OAuth tokens periodically aren't killed at
+// the 1h mark just because their session was created an hour ago. 24h is an
+// upper bound — clients that go truly idle that long can re-handshake cheaply.
+const SESSION_IDLE_MAX_MS = 24 * 60 * 60 * 1000;
 setInterval(
   () => {
     const now = Date.now();
-    const maxAge = 60 * 60 * 1000; // 1 hour
     for (const [id, session] of sessions) {
-      if (now - session.createdAt > maxAge) {
+      if (now - session.lastUsedAt > SESSION_IDLE_MAX_MS) {
         session.transport.close().catch(() => {});
         sessions.delete(id);
       }
@@ -114,9 +184,23 @@ setInterval(
 );
 
 function createSession(apiKey: string, keyInfo: TokenInfo): McpSession {
+  // unauthorized flag lives on a mutable holder so the HarmonyApiClient
+  // callback can flip it without a circular reference between the client and
+  // the session struct it lives on.
+  const authState = { unauthorized: false };
+
+  // Forward declare so onsessioninitialized can register the session.
+  let sessionRef: McpSession;
+
   const transport = new WebStandardStreamableHTTPServerTransport({
     sessionIdGenerator: () => crypto.randomUUID(),
     enableJsonResponse: true,
+    onsessioninitialized: (sid: string) => {
+      sessions.set(sid, sessionRef);
+      console.log(
+        `[mcp] session=${sid} init user=${keyInfo.userId} src=${keyInfo.source}`,
+      );
+    },
   });
 
   const server = new Server(
@@ -124,17 +208,39 @@ function createSession(apiKey: string, keyInfo: TokenInfo): McpSession {
     { capabilities: { tools: {}, resources: {} } },
   );
 
+  // Create per-session API client. onUnauthorized fires when harmony-api
+  // returns 401 — we mark the session so the HTTP layer can surface a real
+  // 401 + WWW-Authenticate challenge to the OAuth client.
+  const client = new HarmonyApiClient({
+    apiKey,
+    apiUrl: HARMONY_API_URL,
+    onUnauthorized: () => {
+      authState.unauthorized = true;
+      // Drop any cached "OK" entry for this bearer so future validations
+      // re-hit harmony-api and see the rejection.
+      invalidateTokenCache(client.getApiKey());
+    },
+  });
+
+  const now = Date.now();
   const session: McpSession = {
     transport,
     server,
+    client,
     apiKey,
+    userId: keyInfo.userId,
     activeWorkspaceId: keyInfo.workspaceId,
     activeProjectId: null,
-    createdAt: Date.now(),
+    createdAt: now,
+    lastUsedAt: now,
+    get unauthorized() {
+      return authState.unauthorized;
+    },
+    set unauthorized(v: boolean) {
+      authState.unauthorized = v;
+    },
   };
-
-  // Create per-session deps
-  const client = new HarmonyApiClient({ apiKey, apiUrl: HARMONY_API_URL });
+  sessionRef = session;
 
   const deps: ToolDeps = {
     getClient: () => client,
@@ -156,10 +262,13 @@ function createSession(apiKey: string, keyInfo: TokenInfo): McpSession {
 
   registerHandlers(server, deps);
 
-  // Clean up session when transport closes
+  // Single cleanup path: fires on explicit DELETE, our evictSession,
+  // and the stale-session GC. Keeping onsessioninitialized + this onclose
+  // (instead of also wiring onsessionclosed) avoids double-logging on DELETE.
   transport.onclose = () => {
     if (transport.sessionId) {
       sessions.delete(transport.sessionId);
+      console.log(`[mcp] session=${transport.sessionId} closed`);
     }
   };
 
@@ -208,91 +317,244 @@ app.get("/.well-known/oauth-protected-resource", (c) =>
 
 // Unauthenticated 401 that advertises the OAuth metadata discovery point.
 // Claude's `mcp add --transport http` uses this to kick off the flow.
-function unauthenticatedResponse(): Response {
+// `errorCode` follows RFC 6750 §3 — clients can branch on `invalid_token` to
+// trigger a refresh vs. surface a hard re-auth.
+function unauthenticatedResponse(
+  errorCode: "missing_token" | "invalid_token" = "missing_token",
+  description?: string,
+): Response {
+  const desc =
+    description ??
+    (errorCode === "invalid_token"
+      ? "Access token is expired or revoked"
+      : "Missing or invalid access token");
+  return new Response(
+    JSON.stringify({ error: errorCode, error_description: desc }),
+    {
+      status: 401,
+      headers: {
+        "Content-Type": "application/json",
+        "WWW-Authenticate":
+          `Bearer realm="mcp", ` +
+          `error="${errorCode}", error_description="${desc}", ` +
+          `resource_metadata="${PUBLIC_MCP_URL}/.well-known/oauth-protected-resource"`,
+      },
+    },
+  );
+}
+
+// Per MCP spec (Streamable HTTP §3 / Session Management §3-4): when a client
+// presents an Mcp-Session-Id we don't recognize, we MUST return 404. Claude
+// then drops the session id and re-`initialize`s. Returning anything else
+// (e.g., transparently minting a new session) wedges the connection because
+// the body is a `tools/call`, not `initialize`, and the SDK transport will
+// reject it with `Server not initialized`.
+function sessionNotFound(sessionId: string): Response {
   return new Response(
     JSON.stringify({
-      error: "unauthorized",
-      error_description: "Missing or invalid access token",
+      jsonrpc: "2.0",
+      error: { code: -32001, message: "Session not found" },
+      id: null,
     }),
     {
-      status: 401,
+      status: 404,
       headers: {
         "Content-Type": "application/json",
-        "WWW-Authenticate": `Bearer realm="mcp", resource_metadata="${PUBLIC_MCP_URL}/.well-known/oauth-protected-resource"`,
+        "Mcp-Session-Id": sessionId,
       },
     },
   );
 }
 
+// Per spec: requests without an Mcp-Session-Id (other than initialization)
+// SHOULD return 400.
+function sessionRequiredResponse(): Response {
+  return new Response(
+    JSON.stringify({
+      jsonrpc: "2.0",
+      error: {
+        code: -32000,
+        message:
+          "Bad Request: Mcp-Session-Id header required for non-initialize requests",
+      },
+      id: null,
+    }),
+    {
+      status: 400,
+      headers: { "Content-Type": "application/json" },
+    },
+  );
+}
+
+// Evict a session and tear down its transport. Used when an OAuth token
+// rotates or is revoked mid-session — we don't want to keep a zombie session
+// around with a stale cached api key.
+function evictSession(sessionId: string): void {
+  const session = sessions.get(sessionId);
+  if (!session) return;
+  sessions.delete(sessionId);
+  session.transport.close().catch(() => {});
+}
+
+// Best-effort body peek so we can route POSTs by JSON-RPC method without
+// double-reading the body downstream (transport.handleRequest accepts
+// `parsedBody` to skip its own json() call).
+async function peekBody(req: Request): Promise<unknown | undefined> {
+  if (req.method !== "POST") return undefined;
+  const ct = req.headers.get("content-type") || "";
+  if (!ct.includes("application/json")) return undefined;
+  try {
+    return await req.clone().json();
+  } catch {
+    return undefined;
+  }
+}
+
 // MCP endpoint - handles POST (JSON-RPC), GET (SSE), DELETE (session close).
 // Mounted on both `/mcp` and `/` so clients that registered the bare host as
 // their server URL still reach the OAuth challenge instead of a 404.
 const handleMcpRequest = async (c: import("hono").Context) => {
   const method = c.req.method;
+  const raw = c.req.raw;
 
-  // Extract bearer token. Accept OAuth access tokens (hmy_at_) and legacy
-  // integration keys (hmy_). No token → 401 with WWW-Authenticate so Claude
-  // can discover our authorization server.
+  // 1. Bearer required for everything. No token → 401 + PRM challenge so
+  //    Claude's `mcp add --transport http` can kick off the OAuth dance.
   const authHeader = c.req.header("Authorization");
   if (!authHeader?.startsWith("Bearer ")) {
-    return unauthenticatedResponse();
+    return unauthenticatedResponse("missing_token");
   }
   const apiKey = authHeader.slice(7);
 
-  // Check for existing session
   const sessionId = c.req.header("Mcp-Session-Id");
 
-  if (sessionId && sessions.has(sessionId)) {
-    // Existing session - forward request
-    const session = sessions.get(sessionId)!;
-    return session.transport.handleRequest(c.req.raw);
-  }
+  // 2. Existing session path — auth is per-request via the bearer; the
+  //    session ID is just a transport handle. Surviving token rotation here
+  //    is what keeps long-lived MCP connections alive past the 1h access
+  //    token TTL.
+  if (sessionId) {
+    const session = sessions.get(sessionId);
+
+    // Per MCP spec §3-4: unknown session id MUST return 404 so the client
+    // re-initializes. NEVER silently mint a new session — the body is a
+    // `tools/call`, not `initialize`, and we'd just bury the failure inside
+    // a JSON-RPC `Server not initialized` envelope. This was the bug
+    // surfacing as "Harmony MCP not responding" after a server restart or
+    // idle eviction.
+    if (!session) {
+      console.log(`[mcp] session=${sessionId} unknown → 404 re-init`);
+      return sessionNotFound(sessionId);
+    }
 
-  if (method === "POST") {
-    // Validate the token. OAuth tokens carry workspaceId; legacy keys don't,
-    // so we look up workspaces in a follow-up call for those.
-    const keyInfo = await validateToken(apiKey);
-    if (!keyInfo) {
-      return unauthenticatedResponse();
+    session.lastUsedAt = Date.now();
+
+    // Hot-swap the cached token when the OAuth client refreshed mid-session.
+    // Re-validate the new bearer and require it to belong to the same user
+    // before we accept it — otherwise a leaked session id paired with a
+    // different user's token would let an attacker ride the session.
+    if (session.apiKey !== apiKey) {
+      const fresh = await validateToken(apiKey);
+      if (!fresh) {
+        console.log(`[mcp] session=${sessionId} swap rejected: invalid token`);
+        return unauthenticatedResponse("invalid_token");
+      }
+      if (fresh.userId !== session.userId) {
+        console.warn(
+          `[mcp] session=${sessionId} swap REJECTED: ` +
+            `user mismatch (session=${session.userId} bearer=${fresh.userId})`,
+        );
+        return unauthenticatedResponse(
+          "invalid_token",
+          "Bearer does not belong to this session",
+        );
+      }
+      console.log(
+        `[mcp] session=${sessionId} token rotated user=${session.userId}`,
+      );
+      session.apiKey = apiKey;
+      session.client.setApiKey(apiKey);
     }
-    if (keyInfo.source === "api_key" && !keyInfo.workspaceId) {
-      keyInfo.workspaceId = await resolveWorkspaceForLegacyKey(apiKey);
+
+    // Reset the per-request 401 latch before handing off to the transport.
+    session.unauthorized = false;
+
+    const response = await session.transport.handleRequest(raw);
+
+    // If a tool call 401'd against harmony-api, the api-client tripped the
+    // unauthorized flag. Return HTTP 401 + WWW-Authenticate so the client
+    // refreshes — instead of burying the auth failure inside a JSON-RPC
+    // error envelope the client can't act on.
+    //
+    // Do NOT evict the session — per MCP spec, the session ID is
+    // independent of auth state. The next request arrives with a fresh
+    // bearer, the hot-swap above installs it, and the session continues.
+    if (session.unauthorized) {
+      console.log(`[mcp] session=${sessionId} api 401 → refresh challenge`);
+      return unauthenticatedResponse(
+        "invalid_token",
+        "Access token rejected by harmony-api",
+      );
     }
 
-    // Create new session
-    const session = createSession(apiKey, keyInfo);
+    return response;
+  }
+
+  // 3. No session id — only `initialize` is allowed; everything else is 400.
+  if (method !== "POST") {
+    // GET/DELETE without a session id are nonsense.
+    return sessionRequiredResponse();
+  }
 
-    // Connect server to transport
-    await session.server.connect(session.transport);
+  const body = await peekBody(raw);
+  if (!body || !isInitializeRequest(body)) {
+    return sessionRequiredResponse();
+  }
 
-    // Store session once transport has a session ID
-    const origOnSessionInitialized = session.transport._onsessioninitialized;
-    session.transport._onsessioninitialized = (sid: string) => {
-      sessions.set(sid, session);
-      origOnSessionInitialized?.(sid);
-    };
+  // 4. Initialize path — validate token, create session, hand off.
+  const keyInfo = await validateToken(apiKey);
+  if (!keyInfo) {
+    return unauthenticatedResponse("invalid_token");
+  }
+  if (keyInfo.source === "api_key" && !keyInfo.workspaceId) {
+    keyInfo.workspaceId = await resolveWorkspaceForLegacyKey(apiKey);
+  }
+
+  const session = createSession(apiKey, keyInfo);
+  await session.server.connect(session.transport);
+
+  const response = await session.transport.handleRequest(raw, {
+    parsedBody: body,
+  });
 
-    // Handle the initialize request
-    return session.transport.handleRequest(c.req.raw);
+  // Edge case: token revoked mid-handshake. Evict the half-built session.
+  if (session.unauthorized) {
+    if (session.transport.sessionId) evictSession(session.transport.sessionId);
+    return unauthenticatedResponse("invalid_token");
   }
 
-  // GET or DELETE without a valid session
-  return c.json({ error: "Invalid or missing session" }, 404);
+  return response;
 };
 
 app.all("/mcp", handleMcpRequest);
 app.all("/", handleMcpRequest);
 
+// Exported for tests — drives the same Hono `fetch` handler that the runtime
+// `serve()` call uses below. Tests construct synthetic Requests and assert on
+// the Response without binding to a real port.
+export const fetchHandler = app.fetch;
+export { sessions as _sessionsForTests };
+
 // ---------------------------------------------------------------------------
-// Start server
+// Start server (skipped when imported as a module — e.g., from tests)
 // ---------------------------------------------------------------------------
-console.log(`Starting Harmony Remote MCP server on port ${PORT}...`);
+if (import.meta.main) {
+  console.log(`Starting Harmony Remote MCP server on port ${PORT}...`);
 
-serve({
-  fetch: app.fetch,
-  port: PORT,
-});
+  serve({
+    fetch: app.fetch,
+    port: PORT,
+  });
 
-console.log(`Harmony Remote MCP server running at http://localhost:${PORT}`);
-console.log(`MCP endpoint: http://localhost:${PORT}/mcp`);
-console.log(`Health check: http://localhost:${PORT}/health`);
+  console.log(`Harmony Remote MCP server running at http://localhost:${PORT}`);
+  console.log(`MCP endpoint: http://localhost:${PORT}/mcp`);
+  console.log(`Health check: http://localhost:${PORT}/health`);
+}