From 2870441fea8206625142f9d6b15ffe9904b82377 Mon Sep 17 00:00:00 2001
From: xiaoli <liiiittly@gmail.com>
Date: Sun, 25 May 2025 11:17:37 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20fix=20review=20comments,=20again?=
 =?UTF-8?q?=F0=9F=98=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Remove unnecessary auth header name check
- Extract operations on headers into a function
---
 client/src/lib/hooks/useConnection.ts |  3 -
 server/src/index.ts                   | 80 +++++++++++++--------------
 2 files changed, 39 insertions(+), 44 deletions(-)

diff --git a/client/src/lib/hooks/useConnection.ts b/client/src/lib/hooks/useConnection.ts
index f1f7ab9e..84956942 100644
--- a/client/src/lib/hooks/useConnection.ts
+++ b/client/src/lib/hooks/useConnection.ts
@@ -301,9 +301,6 @@ export function useConnection({
         // Add custom header name as a special request header to let the server know which header to pass through 
         if (authHeaderName.toLowerCase() !== "authorization") {
           headers[authHeaderName] = token;
-          if (headerName) {
-            headers["x-custom-auth-header"] = headerName;
-          }
         } else {
           headers[authHeaderName] = `Bearer ${token}`;
         }
diff --git a/server/src/index.ts b/server/src/index.ts
index fd28a807..de14e286 100644
--- a/server/src/index.ts
+++ b/server/src/index.ts
@@ -41,6 +41,43 @@ const { values } = parseArgs({
   },
 });
 
+// Function to get HTTP headers.
+// Supports only "sse" and "streamable-http" transport types.
+const getHttpHeaders = (
+  req: express.Request,
+  transportType: string,
+): HeadersInit => {
+  const headers: HeadersInit = {
+    Accept:
+      transportType === "sse"
+        ? "text/event-stream"
+        : "text/event-stream, application/json",
+  };
+  const defaultHeaders =
+    transportType === "sse"
+      ? SSE_HEADERS_PASSTHROUGH
+      : STREAMABLE_HTTP_HEADERS_PASSTHROUGH;
+
+  for (const key of defaultHeaders) {
+    if (req.headers[key] === undefined) {
+      continue;
+    }
+
+    const value = req.headers[key];
+    headers[key] = Array.isArray(value) ? value[value.length - 1] : value;
+  }
+
+  // If the header "x-custom-auth-header" is present, use its value as the custom header name.
+  if (req.headers["x-custom-auth-header"] !== undefined) {
+    const customHeaderName = req.headers["x-custom-auth-header"] as string;
+    if (req.headers[customHeaderName] !== undefined) {
+      const value = req.headers[customHeaderName];
+      headers[customHeaderName] = value as string;
+    }
+  }
+  return headers;
+};
+
 const app = express();
 app.use(cors());
 app.use((req, res, next) => {
@@ -79,27 +116,8 @@ const createTransport = async (req: express.Request): Promise<Transport> => {
     return transport;
   } else if (transportType === "sse") {
     const url = query.url as string;
-    const headers: HeadersInit = {
-      Accept: "text/event-stream",
-    };
 
-    for (const key of SSE_HEADERS_PASSTHROUGH) {
-      if (req.headers[key] === undefined) {
-        continue;
-      }
-
-      const value = req.headers[key];
-      headers[key] = Array.isArray(value) ? value[value.length - 1] : value;
-    }
-    
-    // If the header "x-custom-auth-header" is present, use its value as the custom header name.
-    if (req.headers["x-custom-auth-header"] !== undefined) {
-      const customHeaderName = req.headers["x-custom-auth-header"] as string;
-      if (req.headers[customHeaderName] !== undefined) {
-        const value = req.headers[customHeaderName];
-        headers[customHeaderName] = value as string;
-      }
-    }
+    const headers = getHttpHeaders(req, transportType);
 
     console.log(`SSE transport: url=${url}, headers=${Object.keys(headers)}`);
 
@@ -116,27 +134,7 @@ const createTransport = async (req: express.Request): Promise<Transport> => {
     console.log("Connected to SSE transport");
     return transport;
   } else if (transportType === "streamable-http") {
-    const headers: HeadersInit = {
-      Accept: "text/event-stream, application/json",
-    };
-
-    for (const key of STREAMABLE_HTTP_HEADERS_PASSTHROUGH) {
-      if (req.headers[key] === undefined) {
-        continue;
-      }
-
-      const value = req.headers[key];
-      headers[key] = Array.isArray(value) ? value[value.length - 1] : value;
-    }
-    
-    // If the header "x-custom-auth-header" is present, use its value as the custom header name.
-    if (req.headers["x-custom-auth-header"] !== undefined) {
-      const customHeaderName = req.headers["x-custom-auth-header"] as string;
-      if (req.headers[customHeaderName] !== undefined) {
-        const value = req.headers[customHeaderName];
-        headers[customHeaderName] = value as string;
-      }
-    }
+    const headers = getHttpHeaders(req, transportType);
 
     const transport = new StreamableHTTPClientTransport(
       new URL(query.url as string),