Improve auth error handling and connection reset functionality (#1963)

# READ CAREFULLY THEN REMOVE Remove bullet points that are not relevant. PLEASE REFRAIN FROM USING AI TO WRITE YOUR CODE AND PR DESCRIPTION. IF YOU DO USE AI TO WRITE YOUR CODE PLEASE PROVIDE A DESCRIPTION AND REVIEW IT CAREFULLY. MAKE SURE YOU UNDERSTAND THE CODE YOU ARE SUBMITTING USING AI. - Pull requests that do not follow these guidelines will be closed without review or comment. - If you use AI to write your PR description your pr will be close without review or comment. - If you are unsure about anything, feel free to ask for clarification. ## Description Please provide a clear description of your changes. --- ## Type of Change Please delete options that are not relevant. - [ ] 🐛 Bug fix (non-breaking change which fixes an issue) - [ ] ✨ New feature (non-breaking change which adds functionality) - [ ] 💥 Breaking change (fix or feature with breaking changes) - [ ] 📝 Documentation update - [ ] 🎨 UI/UX improvement - [ ] 🔒 Security enhancement - [ ] ⚡ Performance improvement ## Areas Affected Please check all that apply: - [ ] Email Integration (Gmail, IMAP, etc.) - [ ] User Interface/Experience - [ ] Authentication/Authorization - [ ] Data Storage/Management - [ ] API Endpoints - [ ] Documentation - [ ] Testing Infrastructure - [ ] Development Workflow - [ ] Deployment/Infrastructure ## Testing Done Describe the tests you've done: - [ ] Unit tests added/updated - [ ] Integration tests added/updated - [ ] Manual testing performed - [ ] Cross-browser testing (if UI changes) - [ ] Mobile responsiveness verified (if UI changes) ## Security Considerations For changes involving data or authentication: - [ ] No sensitive data is exposed - [ ] Authentication checks are in place - [ ] Input validation is implemented - [ ] Rate limiting is considered (if applicable) ## Checklist - [ ] I have read the [CONTRIBUTING](https://github.com/Mail-0/Zero/blob/staging/.github/CONTRIBUTING.md) document - [ ] My code follows the project's style guidelines - [ ] I have performed a self-review of my code - [ ] I have commented my code, particularly in complex areas - [ ] I have updated the documentation - [ ] My changes generate no new warnings - [ ] I have added tests that prove my fix/feature works - [ ] All tests pass locally - [ ] Any dependent changes are merged and published ## Additional Notes Add any other context about the pull request here. ## Screenshots/Recordings Add screenshots or recordings here if applicable. --- _By submitting this pull request, I confirm that my contribution is made under the terms of the project's license._  --- ## Summary by cubic Improved authentication error handling and connection reset logic to provide clearer feedback and ensure user sessions are properly revoked when issues occur. - **Bug Fixes** - Added redirects and token revocation for failed user info fetches. - Improved handling when no connections are found for a user. - Triggered mailbox resync if inbox is empty after filtering. - **Dependencies** - Updated better-auth and related packages to the latest versions.   ## Summary by CodeRabbit * **Bug Fixes** * Improved error handling during authentication by revoking tokens, resetting connections, and redirecting users to the home page when user info retrieval fails or email is missing. * Enhanced session management by explicitly revoking and signing out sessions when no user connections are found. * **New Features** * Automatically triggers a forced inbox refresh if no threads are found in the inbox, ensuring the latest messages are displayed. * **Chores** * Updated the "better-auth" dependency to version ^1.3.4.
2026-06-28 06:46:15 +00:00 · 2025-08-09 12:30:13 -07:00
parent e17c3c7be6
commit 91c4349634
5 changed files with 89 additions and 32 deletions
--- a/apps/server/src/lib/auth.ts
+++ b/apps/server/src/lib/auth.ts
@@ -11,13 +11,13 @@ import { createAuthMiddleware, phoneNumber, jwt, bearer, mcp } from 'better-auth
 import { type Account, betterAuth, type BetterAuthOptions } from 'better-auth';
 import { getBrowserTimezone, isValidTimezone } from './timezones';
 import { drizzleAdapter } from 'better-auth/adapters/drizzle';
+import { getZeroDB, resetConnection } from './server-utils';
 import { getSocialProviders } from './auth-providers';
 import { redis, resend, twilio } from './services';
 import { dubAnalytics } from '@dub/better-auth';
 import { defaultUserSettings } from './schemas';
 import { disableBrainFunction } from './brain';
 import { APIError } from 'better-auth/api';
-import { getZeroDB } from './server-utils';
 import { type EProviders } from '../types';
 import { createDriver } from './driver';
 import { Autumn } from 'autumn-js';
@@ -91,7 +91,9 @@ const scheduleCampaign = (userInfo: { address: string; name: string }) =>
 const connectionHandlerHook = async (account: Account) => {
  if (!account.accessToken || !account.refreshToken) {
    console.error('Missing Access/Refresh Tokens', { account });
-    throw new APIError('EXPECTATION_FAILED', { message: 'Missing Access/Refresh Tokens' });
+    throw new APIError('EXPECTATION_FAILED', {
+      message: 'Missing Access/Refresh Tokens, contact us on Discord for support',
+    });
  }

  const driver = createDriver(account.providerId, {
@@ -103,13 +105,26 @@ const connectionHandlerHook = async (account: Account) => {
    },
  });

-  const userInfo = await driver.getUserInfo().catch(() => {
-    throw new APIError('UNAUTHORIZED', { message: 'Failed to get user info' });
+  const userInfo = await driver.getUserInfo().catch(async () => {
+    if (account.accessToken) {
+      await driver.revokeToken(account.accessToken);
+      await resetConnection(account.id);
+    }
+    throw new Response(null, { status: 301, headers: { Location: '/' } });
  });

  if (!userInfo?.address) {
-    console.error('Missing email in user info:', { userInfo });
-    throw new APIError('BAD_REQUEST', { message: 'Missing "email" in user info' });
+    try {
+      await Promise.allSettled(
+        [account.accessToken, account.refreshToken]
+          .filter(Boolean)
+          .map((t) => driver.revokeToken(t as string)),
+      );
+      await resetConnection(account.id);
+    } catch (error) {
+      console.error('Failed to revoke tokens:', error);
+    }
+    throw new Response(null, { status: 303, headers: { Location: '/' } });
  }

  const updatingInfo = {
--- a/apps/server/src/lib/server-utils.ts
+++ b/apps/server/src/lib/server-utils.ts
@@ -360,10 +360,10 @@ export const forceReSync = async (connectionId: string) => {
  const registry = await getRegistryClient(connectionId);
  const allShards = await listShards(registry);

-  await Promise.all(
+  await Promise.allSettled(
    allShards.map(async ({ shard_id: id }) => {
      const shard = await getShardClient(connectionId, id);
-      await Promise.all([
+      await Promise.allSettled([
        shard.exec(`DROP TABLE IF EXISTS threads`),
        shard.exec(`DROP TABLE IF EXISTS thread_labels`),
        shard.exec(`DROP TABLE IF EXISTS labels`),
@@ -374,22 +374,16 @@ export const forceReSync = async (connectionId: string) => {
  await deleteAllShards(registry);

  const agent = await getZeroAgent(connectionId);
-  await agent.stub.forceReSync();
+  return agent.stub.forceReSync();
 };

 export const reSyncThread = async (connectionId: string, threadId: string) => {
  try {
-    const { result: thread, shardId } = await getThread(connectionId, threadId);
+    const { shardId } = await getThread(connectionId, threadId);
    const agent = await getZeroAgentFromShard(connectionId, shardId);
    await agent.stub.syncThread({ threadId });
  } catch (error) {
-    console.error(`[ZeroAgent] Thread not found for threadId: ${threadId}`);
-  }
-  if (thread) {
-    const agent = await getZeroAgentFromShard(connectionId, shardId);
-    await agent.stub.syncThread({ threadId });
-  } else {
-    console.error(`[ZeroAgent] Thread not found for threadId: ${threadId}`);
+    console.error(`[ZeroAgent] Thread not found for threadId: ${threadId}`, error);
  }
 };

@@ -527,11 +521,10 @@ export const getZeroSocketAgent = async (connectionId: string) => {

 export const getActiveConnection = async () => {
  const c = getContext<HonoContext>();
-  const { sessionUser } = c.var;
+  const { sessionUser, auth } = c.var;
  if (!sessionUser) throw new Error('Session Not Found');

  const db = await getZeroDB(sessionUser.id);
-
  const userData = await db.findUser();

  if (userData?.defaultConnectionId) {
@@ -541,6 +534,14 @@ export const getActiveConnection = async () => {

  const firstConnection = await db.findFirstConnection();
  if (!firstConnection) {
+    try {
+      if (auth) {
+        await auth.api.revokeSession({ headers: c.req.raw.headers });
+        await auth.api.signOut({ headers: c.req.raw.headers });
+      }
+    } catch (err) {
+      console.warn(`[getActiveConnection] Session cleanup failed for user ${sessionUser.id}:`, err);
+    }
    console.error(`No connections found for user ${sessionUser.id}`);
    throw new Error('No connections found for user');
  }
--- a/apps/server/src/trpc/routes/mail.ts
+++ b/apps/server/src/trpc/routes/mail.ts
@@ -151,6 +151,31 @@ export const mailRouter = router({
        threadsResponse.threads = filtered;
        console.debug('[listThreads] Snoozed threads after filtering:', filtered);
      }
+
+      if (threadsResponse.threads.length === 0 && folder === FOLDERS.INBOX && !q) {
+        const now = Date.now();
+        const cooldownKey = `resync_cooldown_${activeConnection.id}`;
+        const lastResyncStr = await env.gmail_processing_threads.get(cooldownKey);
+        const lastResync = lastResyncStr ? parseInt(lastResyncStr, 10) : 0;
+        const RESYNC_COOLDOWN_MS = 30000;
+
+        if (now - lastResync > RESYNC_COOLDOWN_MS) {
+          await env.gmail_processing_threads.put(cooldownKey, now.toString(), {
+            expirationTtl: 60,
+          });
+
+          getZeroAgent(activeConnection.id, executionCtx)
+            .then((_agent) => {
+              _agent.stub.forceReSync().catch((error) => {
+                console.error('[listThreads] Async resync failed:', error);
+              });
+            })
+            .catch((error) => {
+              console.error('[listThreads] Failed to get agent for async resync:', error);
+            });
+        }
+      }
+
      console.debug('[listThreads] Returning threadsResponse:', threadsResponse);
      return threadsResponse;
    }),
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -19,8 +19,8 @@ catalogs:
      specifier: ^0.0.48
      version: 0.0.48
    better-auth:
-      specifier: ^1.2.9
-      version: 1.2.10
+      specifier: ^1.3.4
+      version: 1.3.4
    drizzle-kit:
      specifier: ^0.31.1
      version: 0.31.4
@@ -232,7 +232,7 @@ importers:
        version: 19.1.0-rc.2
      better-auth:
        specifier: 'catalog:'
-        version: 1.2.10
+        version: 1.3.4(react-dom@19.1.0(react@19.1.0))(react@19.1.0)
      canvas-confetti:
        specifier: 1.9.3
        version: 1.9.3
@@ -575,7 +575,7 @@ importers:
        version: 1.5.1
      better-auth:
        specifier: 'catalog:'
-        version: 1.2.10
+        version: 1.3.4(react-dom@19.1.0(react@19.1.0))(react@19.1.0)
      cheerio:
        specifier: 1.1.0
        version: 1.1.0
@@ -5139,11 +5139,19 @@ packages:
  base64-js@1.5.1:
    resolution: {integrity: sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA==}

-  better-auth@1.2.10:
-    resolution: {integrity: sha512-nEj1RG4DdLUuJiV5CR93ORyPCptGRBwksaPPCkUtGo9ka+UIlTpaiKoTaTqVLLYlqwX4bOj9tJ32oBNdf2G3Kg==}
+  better-auth@1.3.4:
+    resolution: {integrity: sha512-JbZYam6Cs3Eu5CSoMK120zSshfaKvrCftSo/+v7524H1RvhryQ7UtMbzagBcXj0Digjj8hZtVkkR4tTZD/wK2g==}
+    peerDependencies:
+      react: ^18.0.0 || ^19.0.0
+      react-dom: ^18.0.0 || ^19.0.0
+    peerDependenciesMeta:
+      react:
+        optional: true
+      react-dom:
+        optional: true

-  better-call@1.0.9:
-    resolution: {integrity: sha512-Qfm0gjk0XQz0oI7qvTK1hbqTsBY4xV2hsHAxF8LZfUYl3RaECCIifXuVqtPpZJWvlCCMlQSvkvhhyuApGUba6g==}
+  better-call@1.0.13:
+    resolution: {integrity: sha512-auqdP9lnNOli9tKpZIiv0nEIwmmyaD/RotM3Mucql+Ef88etoZi/t7Ph5LjlmZt/hiSahhNTt6YVnx6++rziXA==}

  better-sqlite3@11.10.0:
    resolution: {integrity: sha512-EwhOpyXiOEL/lKzHz9AW1msWFNzGc/z+LzeB3/jnFJpxu+th2yqvzsSWas1v9jgs9+xiXJcD5A8CJxAG2TaghQ==}
@@ -9912,6 +9920,9 @@ packages:
  zod@3.25.67:
    resolution: {integrity: sha512-idA2YXwpCdqUSKRCACDE6ItZD9TZzy3OZMtpfLoh6oPR47lipysRrJfjzMqFxQ3uJuUPyUeWe1r9vLH33xO/Qw==}

+  zod@4.0.15:
+    resolution: {integrity: sha512-2IVHb9h4Mt6+UXkyMs0XbfICUh1eUrlJJAOupBHUhLRnKkruawyDddYRCs0Eizt900ntIMk9/4RksYl+FgSpcQ==}
+
  zustand@4.5.7:
    resolution: {integrity: sha512-CHOUy7mu3lbD6o6LJLfllpjkzhHXSBlX8B9+qPddUsIfeF5S/UZ5q0kmCsnRqT1UHFQZchNFDDzMbQsuesHWlw==}
    engines: {node: '>=12.7.0'}
@@ -14451,7 +14462,7 @@ snapshots:

  base64-js@1.5.1: {}

-  better-auth@1.2.10:
+  better-auth@1.3.4(react-dom@19.1.0(react@19.1.0))(react@19.1.0):
    dependencies:
      '@better-auth/utils': 0.2.5
      '@better-fetch/fetch': 1.1.18
@@ -14459,14 +14470,17 @@ snapshots:
      '@noble/hashes': 1.8.0
      '@simplewebauthn/browser': 13.1.0
      '@simplewebauthn/server': 13.1.1
-      better-call: 1.0.9
+      better-call: 1.0.13
      defu: 6.1.4
      jose: 5.10.0
      kysely: 0.28.2
      nanostores: 0.11.4
-      zod: 3.25.67
+      zod: 4.0.15
+    optionalDependencies:
+      react: 19.1.0
+      react-dom: 19.1.0(react@19.1.0)

-  better-call@1.0.9:
+  better-call@1.0.13:
    dependencies:
      '@better-fetch/fetch': 1.1.18
      rou3: 0.5.1
@@ -19950,6 +19964,8 @@ snapshots:

  zod@3.25.67: {}

+  zod@4.0.15: {}
+
  zustand@4.5.7(@types/react@19.0.10)(react@19.1.0):
    dependencies:
      use-sync-external-store: 1.5.0(react@19.1.0)
--- a/pnpm-workspace.yaml
+++ b/pnpm-workspace.yaml
@@ -4,7 +4,7 @@ packages:
  - scripts/*
 catalog:
  zod: ^3.25.42
-  better-auth: ^1.2.9
+  better-auth: ^1.3.4
  autumn-js: ^0.0.48
  superjson: ^2.2.2
  '@trpc/server': ^11.1.4