Remove unused code and optimize database connections (#1778)

# 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._

    
<!-- This is an auto-generated description by cubic. -->
---

## Summary by cubic
Removed unused code and improved database connection handling to reduce resource usage and simplify maintenance.

- **Refactors**
  - Deleted unused constants and imports.
  - Cleaned up database connection logic to avoid unnecessary connections and disposals.
  - Simplified error handling in thread retrieval.

<!-- End of auto-generated description by cubic. -->

Author: MrgSub

apps/server/src/lib/prompts.ts

@@ -2,8 +2,6 @@ import { format } from 'date-fns';
 import { Tools } from '../types';
 import dedent from 'dedent';
 
-const CATEGORY_IDS = ['Important', 'All Mail', 'Personal', 'Updates', 'Promotions', 'Unread'];
-
 export const colors = [
   '#000000',
   '#434343',

apps/server/src/main.ts

@@ -14,19 +14,18 @@ import {
   userSettings,
   writingStyleMatrix,
 } from './db/schema';
-import { getZeroAgent } from './lib/server-utils';
 import { env, WorkerEntrypoint, DurableObject, RpcTarget } from 'cloudflare:workers';
 import { EProviders, type ISubscribeBatch, type IThreadBatch } from './types';
 import { oAuthDiscoveryMetadata } from 'better-auth/plugins';
 import { getZeroDB, verifyToken } from './lib/server-utils';
 import { eq, and, desc, asc, inArray } from 'drizzle-orm';
 import { EWorkflowType, runWorkflow } from './pipelines';
+import { ZeroAgent, ZeroDriver } from './routes/agent';
 import { contextStorage } from 'hono/context-storage';
 import { defaultUserSettings } from './lib/schemas';
 import { createLocalJWKSet, jwtVerify } from 'jose';
 import { routePartykitRequest } from 'partyserver';
-
-import { ZeroAgent, ZeroDriver } from './routes/agent';
+import { getZeroAgent } from './lib/server-utils';
 import { enableBrainFunction } from './lib/brain';
 import { trpcServer } from '@hono/trpc-server';
 import { agentsMiddleware } from 'hono-agents';
@@ -518,7 +517,6 @@ export default class extends WorkerEntrypoint<typeof env> {
           if (userId) {
             const db = getZeroDB(userId);
             c.set('sessionUser', await db.findUser());
-            (await db)[Symbol.dispose]?.();
           }
         }
       }
@@ -528,11 +526,6 @@ export default class extends WorkerEntrypoint<typeof env> {
 
       await next();
 
-      if (c.var.sessionUser?.id) {
-        const db = getZeroDB(c.var.sessionUser.id);
-        (await db)[Symbol.dispose]?.();
-      }
-
       c.set('sessionUser', undefined);
       c.set('autumn', undefined as any);
       c.set('auth', undefined as any);

apps/server/src/routes/agent/index.ts

@@ -27,13 +27,14 @@ import {
   type IncomingMessage,
   type OutgoingMessage,
 } from './types';
+import {
+  EPrompts,
+  type IOutgoingMessage,
+  type ParsedMessage,
+  type ISnoozeBatch,
+} from '../../types';
+import type { MailManager, IGetThreadResponse, IGetThreadsResponse } from '../../lib/driver/types';
 import { DurableObjectOAuthClientProvider } from 'agents/mcp/do-oauth-client-provider';
-import { EPrompts, type IOutgoingMessage, type ParsedMessage, type ISnoozeBatch } from '../../types';
-import type {
-  MailManager,
-  IGetThreadResponse,
-  IGetThreadsResponse,
-} from '../../lib/driver/types';
 import { AiChatPrompt, GmailSearchAssistantSystemPrompt } from '../../lib/prompts';
 import { connectionToDriver } from '../../lib/server-utils';
 import type { CreateDraftData } from '../../lib/schemas';
@@ -63,7 +64,6 @@ const shouldLoop = env.THREAD_SYNC_LOOP !== 'false';
 export class ZeroDriver extends AIChatAgent<typeof env> {
   private foldersInSync: Map<string, boolean> = new Map();
   private syncThreadsInProgress: Map<string, boolean> = new Map();
-  private currentFolder: string | null = 'inbox';
   private driver: MailManager | null = null;
   private agent: DurableObjectStub<ZeroAgent> | null = null;
   constructor(ctx: DurableObjectState, env: Env) {
@@ -506,7 +506,6 @@ export class ZeroDriver extends AIChatAgent<typeof env> {
     }
 
     try {
-      const startTime = Date.now();
       console.log(`[inboxRag] Executing AI search with parameters:`, {
         query,
         max_num_results: 3,
@@ -529,8 +528,6 @@ export class ZeroDriver extends AIChatAgent<typeof env> {
         },
       });
 
-      const duration = Date.now() - startTime;
-
       return { result: answer.response, data: answer.data };
     } catch (error) {
       console.error(`[inboxRag] Search failed for query: "${query}"`, {
@@ -774,7 +771,7 @@ export class ZeroDriver extends AIChatAgent<typeof env> {
     }
   }
 
-  async getThreadFromDB(id: string, lastAttempt = false): Promise<IGetThreadResponse> {
+  async getThreadFromDB(id: string): Promise<IGetThreadResponse> {
     try {
       const result = this.sql`
           SELECT
@@ -793,11 +790,13 @@ export class ZeroDriver extends AIChatAgent<typeof env> {
         `;
 
       if (!result || result.length === 0) {
-        if (lastAttempt) {
-          throw new Error('Thread not found in database, Sync Failed once');
-        }
-        await this.syncThread({ threadId: id });
-        return this.getThreadFromDB(id, true);
+        return {
+          messages: [],
+          latest: undefined,
+          hasUnread: false,
+          totalReplies: 0,
+          labels: [],
+        } satisfies IGetThreadResponse;
       }
       const row = result[0] as { latest_label_ids: string };
       const storedThread = await env.THREADS_BUCKET.get(this.getThreadKey(id));

apps/server/src/routes/agent/mcp.ts

@@ -23,7 +23,6 @@ import { env } from 'cloudflare:workers';
 import { eq, and } from 'drizzle-orm';
 import { McpAgent } from 'agents/mcp';
 import { createDb } from '../../db';
-import { generateText } from 'ai';
 import z from 'zod';
 
 export class ZeroMCP extends McpAgent<typeof env, Record<string, unknown>, { userId: string }> {