✨ feat(trpc): add response metadata and auth header handling (#11816)

* ✨ feat(trpc): add response metadata and auth header handling

Add createResponseMeta utility to centralize tRPC response metadata handling.
Set X-Auth-Required header for UNAUTHORIZED errors to distinguish real auth failures
from other 401 errors. Update all tRPC routes to use the new utility.

* ♻️ refactor(desktop-bridge): extract auth constants to shared package

Move AUTH_REQUIRED_HEADER and TRPC_ERROR_CODE_UNAUTHORIZED to
@lobechat/desktop-bridge for consistent usage across server and desktop.

apps/desktop/src/main/core/infrastructure/BackendProxyProtocolManager.ts

@@ -1,3 +1,4 @@
+import { AUTH_REQUIRED_HEADER } from '@lobechat/desktop-bridge';
 import { BrowserWindow, type Session } from 'electron';
 
 import { isDev } from '@/const/env';
@@ -167,7 +168,7 @@ export class BackendProxyProtocolManager {
         // The server sets X-Auth-Required header for real authentication failures (e.g., token expired)
         // Other 401 errors (e.g., invalid API keys) should not trigger re-authentication
         if (upstreamResponse.status === 401) {
-          const authRequired = upstreamResponse.headers.get('X-Auth-Required') === 'true';
+          const authRequired = upstreamResponse.headers.get(AUTH_REQUIRED_HEADER) === 'true';
           if (authRequired) {
             this.notifyAuthorizationRequired();
           }

packages/desktop-bridge/src/index.ts

@@ -15,3 +15,20 @@ export const APP_WINDOW_MIN_SIZE = {
   height: 600,
   width: 1000,
 } as const;
+
+// HTTP Headers for desktop-server communication
+/**
+ * Header to indicate that a 401 response is due to a real authentication failure
+ * (e.g., token expired) rather than other 401 causes (e.g., invalid API keys).
+ *
+ * When the server sets this header to 'true', the desktop app should trigger
+ * re-authentication flow.
+ */
+export const AUTH_REQUIRED_HEADER = 'X-Auth-Required';
+
+// TRPC error codes (mirrors @trpc/server internal codes)
+/**
+ * TRPC error code for unauthorized requests.
+ * Used to identify authentication failures in TRPC responses.
+ */
+export const TRPC_ERROR_CODE_UNAUTHORIZED = 'UNAUTHORIZED' as const;

src/app/(backend)/trpc/async/[trpc]/route.ts

@@ -4,6 +4,7 @@ import type { NextRequest } from 'next/server';
 import { pino } from '@/libs/logger';
 import { createAsyncRouteContext } from '@/libs/trpc/async/context';
 import { prepareRequestForTRPC } from '@/libs/trpc/utils/request-adapter';
+import { createResponseMeta } from '@/libs/trpc/utils/responseMeta';
 import { asyncRouter } from '@/server/routers/async';
 
 const handler = (req: NextRequest) => {
@@ -29,6 +30,7 @@ const handler = (req: NextRequest) => {
     },
 
     req: preparedReq,
+    responseMeta: createResponseMeta,
     router: asyncRouter,
   });
 };

src/app/(backend)/trpc/lambda/[trpc]/route.ts

@@ -3,6 +3,7 @@ import type { NextRequest } from 'next/server';
 
 import { createLambdaContext } from '@/libs/trpc/lambda/context';
 import { prepareRequestForTRPC } from '@/libs/trpc/utils/request-adapter';
+import { createResponseMeta } from '@/libs/trpc/utils/responseMeta';
 import { lambdaRouter } from '@/server/routers/lambda';
 
 const handler = (req: NextRequest) => {
@@ -28,11 +29,7 @@ const handler = (req: NextRequest) => {
     },
 
     req: preparedReq,
-    responseMeta({ ctx }) {
-      const headers = ctx?.resHeaders;
-
-      return { headers };
-    },
+    responseMeta: createResponseMeta,
     router: lambdaRouter,
   });
 };

src/app/(backend)/trpc/mobile/[trpc]/route.ts

@@ -4,6 +4,7 @@ import type { NextRequest } from 'next/server';
 import { pino } from '@/libs/logger';
 import { createLambdaContext } from '@/libs/trpc/lambda/context';
 import { prepareRequestForTRPC } from '@/libs/trpc/utils/request-adapter';
+import { createResponseMeta } from '@/libs/trpc/utils/responseMeta';
 import { mobileRouter } from '@/server/routers/mobile';
 
 const handler = (req: NextRequest) => {
@@ -25,11 +26,7 @@ const handler = (req: NextRequest) => {
     },
 
     req: preparedReq,
-    responseMeta({ ctx }) {
-      const headers = ctx?.resHeaders;
-
-      return { headers };
-    },
+    responseMeta: createResponseMeta,
     router: mobileRouter,
   });
 };

src/app/(backend)/trpc/tools/[trpc]/route.ts

@@ -3,6 +3,7 @@ import type { NextRequest } from 'next/server';
 
 import { createLambdaContext } from '@/libs/trpc/lambda/context';
 import { prepareRequestForTRPC } from '@/libs/trpc/utils/request-adapter';
+import { createResponseMeta } from '@/libs/trpc/utils/responseMeta';
 import { toolsRouter } from '@/server/routers/tools';
 
 const handler = (req: NextRequest) => {
@@ -24,6 +25,7 @@ const handler = (req: NextRequest) => {
     },
 
     req: preparedReq,
+    responseMeta: createResponseMeta,
     router: toolsRouter,
   });
 };

src/features/Electron/connection/Connection.tsx

@@ -1,18 +1,22 @@
+import { Center, Flexbox } from '@lobehub/ui';
 import { Drawer } from 'antd';
 import { createStaticStyles, cssVar } from 'antd-style';
-import { useState } from 'react';
+import { useCallback, useState } from 'react';
+
+import LoginStep from '@/app/[variants]/(desktop)/desktop-onboarding/features/LoginStep';
+import { isMacOS } from '@/utils/platform';
 
-import ConnectionMode from './ConnectionMode';
 import RemoteStatus from './RemoteStatus';
-import WaitingOAuth from './Waiting';
+
+const isMac = isMacOS();
 
 const styles = createStaticStyles(({ css }) => {
   return {
     modal: css`
       .ant-drawer-close {
         position: absolute;
-        inset-block-start: 8px;
-        inset-inline-end: 0;
+        inset-block-start: ${isMac ? '12px' : '46px'};
+        inset-inline-end: 6px;
       }
     `,
   };
@@ -20,7 +24,10 @@ const styles = createStaticStyles(({ css }) => {
 
 const Connection = () => {
   const [isOpen, setIsOpen] = useState(false);
-  const [isWaiting, setWaiting] = useState(false);
+
+  const handleClose = useCallback(() => {
+    setIsOpen(false);
+  }, []);
 
   return (
     <>
@@ -31,22 +38,20 @@ const Connection = () => {
       />
       <Drawer
         classNames={{ header: styles.modal }}
-        height={'100vh'}
-        onClose={() => {
-          setIsOpen(false);
-        }}
+        onClose={handleClose}
         open={isOpen}
         placement={'top'}
+        size={'100vh'}
         style={{
           background: cssVar.colorBgLayout,
         }}
         styles={{ body: { padding: 0 }, header: { padding: 0 } }}
       >
-        {isWaiting ? (
-          <WaitingOAuth setIsOpen={setIsOpen} setWaiting={setWaiting} />
-        ) : (
-          <ConnectionMode setWaiting={setWaiting} />
-        )}
+        <Center style={{ height: '100%', overflow: 'auto', padding: 24 }}>
+          <Flexbox style={{ maxWidth: 560, width: '100%' }}>
+            <LoginStep onBack={handleClose} onNext={handleClose} />
+          </Flexbox>
+        </Center>
       </Drawer>
     </>
   );

src/libs/trpc/utils/responseMeta.test.ts

@@ -0,0 +1,82 @@
+import { AUTH_REQUIRED_HEADER, TRPC_ERROR_CODE_UNAUTHORIZED } from '@lobechat/desktop-bridge';
+import { TRPCError } from '@trpc/server';
+import { describe, expect, it } from 'vitest';
+
+import { createResponseMeta } from './responseMeta';
+
+describe('createResponseMeta', () => {
+  it('should return undefined headers when no errors and no resHeaders', () => {
+    const result = createResponseMeta({ ctx: undefined, errors: [] });
+    expect(result.headers).toBeUndefined();
+  });
+
+  it('should forward resHeaders from context', () => {
+    const resHeaders = new Headers({ 'X-Custom': 'value' });
+    const result = createResponseMeta({
+      ctx: { resHeaders },
+      errors: [],
+    });
+
+    expect(result.headers).toBeInstanceOf(Headers);
+    expect(result.headers?.get('X-Custom')).toBe('value');
+  });
+
+  it('should set AUTH_REQUIRED_HEADER header for UNAUTHORIZED error', () => {
+    const error = new TRPCError({ code: TRPC_ERROR_CODE_UNAUTHORIZED });
+    const result = createResponseMeta({
+      ctx: undefined,
+      errors: [error],
+    });
+
+    expect(result.headers).toBeInstanceOf(Headers);
+    expect(result.headers?.get(AUTH_REQUIRED_HEADER)).toBe('true');
+  });
+
+  it('should set AUTH_REQUIRED_HEADER and preserve resHeaders for UNAUTHORIZED error', () => {
+    const resHeaders = new Headers({ 'X-Custom': 'value' });
+    const error = new TRPCError({ code: TRPC_ERROR_CODE_UNAUTHORIZED });
+    const result = createResponseMeta({
+      ctx: { resHeaders },
+      errors: [error],
+    });
+
+    expect(result.headers).toBeInstanceOf(Headers);
+    expect(result.headers?.get(AUTH_REQUIRED_HEADER)).toBe('true');
+    expect(result.headers?.get('X-Custom')).toBe('value');
+  });
+
+  it('should NOT set AUTH_REQUIRED_HEADER for non-UNAUTHORIZED errors', () => {
+    const error = new TRPCError({ code: 'BAD_REQUEST' });
+    const result = createResponseMeta({
+      ctx: undefined,
+      errors: [error],
+    });
+
+    expect(result.headers).toBeUndefined();
+  });
+
+  it('should handle context without resHeaders property', () => {
+    const error = new TRPCError({ code: TRPC_ERROR_CODE_UNAUTHORIZED });
+    const result = createResponseMeta({
+      ctx: { userId: 'test-user' },
+      errors: [error],
+    });
+
+    expect(result.headers).toBeInstanceOf(Headers);
+    expect(result.headers?.get(AUTH_REQUIRED_HEADER)).toBe('true');
+  });
+
+  it('should handle multiple errors where one is UNAUTHORIZED', () => {
+    const errors = [
+      new TRPCError({ code: 'BAD_REQUEST' }),
+      new TRPCError({ code: TRPC_ERROR_CODE_UNAUTHORIZED }),
+    ];
+    const result = createResponseMeta({
+      ctx: undefined,
+      errors,
+    });
+
+    expect(result.headers).toBeInstanceOf(Headers);
+    expect(result.headers?.get(AUTH_REQUIRED_HEADER)).toBe('true');
+  });
+});

src/libs/trpc/utils/responseMeta.ts

@@ -0,0 +1,41 @@
+import { AUTH_REQUIRED_HEADER, TRPC_ERROR_CODE_UNAUTHORIZED } from '@lobechat/desktop-bridge';
+import type { TRPCError } from '@trpc/server';
+
+interface ResponseMetaParams {
+  ctx?: unknown;
+  errors: TRPCError[];
+}
+
+/**
+ * Create response metadata for TRPC handlers.
+ *
+ * This function handles:
+ * 1. Forwarding custom headers from context (ctx.resHeaders)
+ * 2. Adding X-Auth-Required header for UNAUTHORIZED errors
+ *
+ * The X-Auth-Required header allows the desktop app (BackendProxyProtocolManager)
+ * to distinguish between real authentication failures (e.g., token expired)
+ * and other 401 errors (e.g., invalid API keys).
+ */
+export function createResponseMeta({ ctx, errors }: ResponseMetaParams): {
+  headers: Headers | undefined;
+} {
+  const resHeaders =
+    ctx && typeof ctx === 'object' && 'resHeaders' in ctx
+      ? // eslint-disable-next-line no-undef
+        (ctx as { resHeaders?: HeadersInit }).resHeaders
+      : undefined;
+  const headers = resHeaders ? new Headers(resHeaders) : new Headers();
+
+  const hasUnauthorizedError = errors.some((error) => error.code === TRPC_ERROR_CODE_UNAUTHORIZED);
+  if (hasUnauthorizedError) {
+    headers.set(AUTH_REQUIRED_HEADER, 'true');
+  }
+
+  // Only return headers if there's content or auth error
+  if (hasUnauthorizedError || resHeaders) {
+    return { headers };
+  }
+
+  return { headers: undefined };
+}

src/utils/errorResponse.ts

@@ -1,3 +1,4 @@
+import { AUTH_REQUIRED_HEADER } from '@lobechat/desktop-bridge';
 import { AgentRuntimeErrorType, type ILobeAgentRuntimeErrorType } from '@lobechat/model-runtime';
 import { ChatErrorType, type ErrorResponse, type ErrorType } from '@lobechat/types';
 
@@ -85,7 +86,7 @@ export const createErrorResponse = (
   // Add X-Auth-Required header for real authentication failures
   // This allows the client to distinguish between auth failures and other 401 errors (e.g., invalid API keys)
   if (AUTH_REQUIRED_ERROR_TYPES.has(errorType as ErrorType)) {
-    headers['X-Auth-Required'] = 'true';
+    headers[AUTH_REQUIRED_HEADER] = 'true';
   }
 
   return new Response(JSON.stringify(data), { headers, status: statusCode });