refactor durable objects (#1764)

# Refactor Agent Architecture with Driver/Agent Split

## Description

This PR refactors the agent architecture by splitting the ZeroAgent class into two separate classes: ZeroDriver and ZeroAgent. The ZeroDriver handles mail operations and database interactions, while ZeroAgent focuses on chat functionality. This separation of concerns improves code organization and maintainability.

Key changes:
- Created a new ZeroDriver class to handle mail operations and database interactions
- Modified ZeroAgent to focus on chat functionality
- Updated RPC target from AgentRpcDO to DriverRpcDO
- Fixed variable declarations from `let` to `const` in mail-display.tsx
- Updated environment configuration in wrangler.jsonc to include the new ZeroDriver class

## Type of Change

- [x] ⚡ Performance improvement
- [x] 🎨 UI/UX improvement

## Areas Affected

- [x] Email Integration (Gmail, IMAP, etc.)
- [x] Data Storage/Management
- [x] API Endpoints
- [x] Development Workflow

## Testing Done

- [x] Manual testing performed

## Security Considerations

- [x] No sensitive data is exposed
- [x] Authentication checks are in place

## Checklist

- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in complex areas
- [x] My changes generate no new warnings

## Additional Notes

This architectural change improves separation of concerns and should make the codebase more maintainable. The ZeroDriver handles all mail-related operations while ZeroAgent focuses on chat functionality, creating a cleaner division of responsibilities.

apps/mail/components/mail/mail-display.tsx

@@ -386,7 +386,7 @@ const downloadAttachment = async (attachment: {
   attachmentId: string;
 }) => {
   try {
-    let attachmentData = attachment.body;
+    const attachmentData = attachment.body;
 
     if (!attachmentData) {
       throw new Error('Attachment data not found');
@@ -475,7 +475,7 @@ const openAttachment = async (attachment: {
   attachmentId: string;
 }) => {
   try {
-    let attachmentData = attachment.body;
+    const attachmentData = attachment.body;
 
     if (!attachmentData) {
       throw new Error('Attachment data not found');

apps/server/src/lib/server-utils.ts

@@ -11,9 +11,9 @@ export const getZeroDB = (userId: string) => {
 };
 
 export const getZeroAgent = async (connectionId: string) => {
-  const stub = env.ZERO_AGENT.get(env.ZERO_AGENT.idFromName(connectionId));
+  const stub = env.ZERO_DRIVER.get(env.ZERO_DRIVER.idFromName(connectionId));
   const rpcTarget = await stub.setMetaData(connectionId);
-  await rpcTarget.setupAuth(connectionId);
+  await rpcTarget.setupAuth();
   return rpcTarget;
 };
 

apps/server/src/main.ts

@@ -25,13 +25,13 @@ import { defaultUserSettings } from './lib/schemas';
 import { createLocalJWKSet, jwtVerify } from 'jose';
 import { routePartykitRequest } from 'partyserver';
 
+import { ZeroAgent, ZeroDriver } from './routes/agent';
 import { enableBrainFunction } from './lib/brain';
 import { trpcServer } from '@hono/trpc-server';
 import { agentsMiddleware } from 'hono-agents';
 import { ZeroMCP } from './routes/agent/mcp';
 import { publicRouter } from './routes/auth';
 import { autumnApi } from './routes/autumn';
-import { ZeroAgent } from './routes/agent';
 import type { HonoContext } from './ctx';
 import { createDb, type DB } from './db';
 import { createAuth } from './lib/auth';
@@ -789,4 +789,4 @@ export default class extends WorkerEntrypoint<typeof env> {
   }
 }
 
-export { ZeroAgent, ZeroMCP, ZeroDB };
+export { ZeroAgent, ZeroMCP, ZeroDB, ZeroDriver };

apps/server/src/pipelines.effect.ts

@@ -25,9 +25,9 @@ import {
 import { defaultLabels, EPrompts, EProviders, type ParsedMessage, type Sender } from './types';
 import { getZeroAgent } from './lib/server-utils';
 import { type gmail_v1 } from '@googleapis/gmail';
-import { connection, summary } from './db/schema';
 import { getPromptName } from './pipelines';
 import { env } from 'cloudflare:workers';
+import { connection } from './db/schema';
 import { Effect, Console } from 'effect';
 import * as cheerio from 'cheerio';
 import { eq } from 'drizzle-orm';

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

@@ -166,7 +166,7 @@ export class ZeroMCP extends McpAgent<typeof env, Record<string, unknown>, { use
           pageToken: s.pageToken,
         });
         const content = await Promise.all(
-          result.threads.map(async (thread: any) => {
+          result.threads.map(async (thread) => {
             const loadedThread = await agent.getThread(thread.id);
             return [
               {

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

@@ -16,11 +16,11 @@
 import type { CreateDraftData } from '../../lib/schemas';
 import type { IOutgoingMessage } from '../../types';
 import { RpcTarget } from 'cloudflare:workers';
-import { ZeroAgent } from '.';
+import { ZeroDriver } from '.';
 
-export class AgentRpcDO extends RpcTarget {
+export class DriverRpcDO extends RpcTarget {
   constructor(
-    private mainDo: ZeroAgent,
+    private mainDo: ZeroDriver,
     private connectionId: string,
   ) {
     super();
@@ -176,9 +176,8 @@ export class AgentRpcDO extends RpcTarget {
     return await this.mainDo.getMessageAttachments(messageId);
   }
 
-  async setupAuth(connectionId: string) {
-    if (connectionId !== this.connectionId) console.warn('Oops, something doesnt add up.');
-    return await this.mainDo.setupAuth(connectionId);
+  async setupAuth() {
+    return await this.mainDo.setupAuth();
   }
 
   async broadcast(message: string) {

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

@@ -3,9 +3,9 @@ import { perplexity } from '@ai-sdk/perplexity';
 import { generateText, tool } from 'ai';
 
 import { colors, GmailSearchAssistantSystemPrompt } from '../../lib/prompts';
+import { getZeroAgent } from '../../lib/server-utils';
 import { anthropic } from '@ai-sdk/anthropic';
 import { env } from 'cloudflare:workers';
-import type { ZeroAgent } from '../chat';
 import { Tools } from '../../types';
 import { z } from 'zod';
 
@@ -112,7 +112,7 @@ export const getEmbeddingVector = async (
  *
  * The tag format must be exactly: <thread id="{id}"/>
  */
-const getEmail = (_: ZeroAgent) =>
+const getEmail = () =>
   tool({
     description: 'Return a placeholder tag for a specific email thread by ID',
     parameters: z.object({
@@ -155,7 +155,7 @@ const composeEmailTool = (connectionId: string) =>
     },
   });
 
-// const listEmails = (agent: ZeroAgent) =>
+// const listEmails = (connectionId: string) =>
 //   tool({
 //     description: 'List emails in a specific folder',
 //     parameters: z.object({
@@ -173,19 +173,20 @@ const composeEmailTool = (connectionId: string) =>
 //     },
 //   });
 
-const markAsRead = (agent: ZeroAgent) =>
+const markAsRead = (connectionId: string) =>
   tool({
     description: 'Mark emails as read',
     parameters: z.object({
       threadIds: z.array(z.string()).describe('The IDs of the threads to mark as read'),
     }),
     execute: async ({ threadIds }) => {
-      await agent.markAsRead(threadIds);
+      const driver = await getZeroAgent(connectionId);
+      await driver.markAsRead(threadIds);
       return { threadIds, success: true };
     },
   });
 
-// const inboxRag = (agent: ZeroAgent, dataStream?: DataStreamWriter) =>
+// const inboxRag = (connectionId: string, dataStream?: DataStreamWriter) =>
 //   tool({
 //     description: 'Search the inbox for emails',
 //     parameters: z.object({
@@ -197,19 +198,20 @@ const markAsRead = (agent: ZeroAgent) =>
 //     },
 //   });
 
-const markAsUnread = (agent: ZeroAgent) =>
+const markAsUnread = (connectionId: string) =>
   tool({
     description: 'Mark emails as unread',
     parameters: z.object({
       threadIds: z.array(z.string()).describe('The IDs of the threads to mark as unread'),
     }),
     execute: async ({ threadIds }) => {
-      await agent.markAsUnread(threadIds);
+      const driver = await getZeroAgent(connectionId);
+      await driver.markAsUnread(threadIds);
       return { threadIds, success: true };
     },
   });
 
-const modifyLabels = (agent: ZeroAgent) =>
+const modifyLabels = (connectionId: string) =>
   tool({
     description: 'Modify labels on emails',
     parameters: z.object({
@@ -220,21 +222,23 @@ const modifyLabels = (agent: ZeroAgent) =>
       }),
     }),
     execute: async ({ threadIds, options }) => {
-      await agent.modifyLabels(threadIds, options.addLabels, options.removeLabels);
+      const driver = await getZeroAgent(connectionId);
+      await driver.modifyLabels(threadIds, options.addLabels, options.removeLabels);
       return { threadIds, options, success: true };
     },
   });
 
-const getUserLabels = (agent: ZeroAgent) =>
+const getUserLabels = (connectionId: string) =>
   tool({
     description: 'Get all user labels',
     parameters: z.object({}),
     execute: async () => {
-      return await agent.getUserLabels();
+      const driver = await getZeroAgent(connectionId);
+      return await driver.getUserLabels();
     },
   });
 
-const sendEmail = (agent: ZeroAgent) =>
+const sendEmail = (connectionId: string) =>
   tool({
     description: 'Send a new email',
     parameters: z.object({
@@ -268,16 +272,17 @@ const sendEmail = (agent: ZeroAgent) =>
     }),
     execute: async (data) => {
       try {
+        const driver = await getZeroAgent(connectionId);
         const { draftId, ...mail } = data;
 
         if (draftId) {
-          await agent.sendDraft(draftId, {
+          await driver.sendDraft(draftId, {
             ...mail,
             attachments: [],
             headers: {},
           });
         } else {
-          await agent.create({
+          await driver.create({
             ...mail,
             attachments: [],
             headers: {},
@@ -294,7 +299,7 @@ const sendEmail = (agent: ZeroAgent) =>
     },
   });
 
-const createLabel = (agent: ZeroAgent) =>
+const createLabel = (connectionId: string) =>
   tool({
     description: 'Create a new label with custom colors, if it does nto exist already',
     parameters: z.object({
@@ -313,43 +318,47 @@ const createLabel = (agent: ZeroAgent) =>
         }),
     }),
     execute: async ({ name, backgroundColor, textColor }) => {
-      await agent.createLabel({ name, color: { backgroundColor, textColor } });
+      const driver = await getZeroAgent(connectionId);
+      await driver.createLabel({ name, color: { backgroundColor, textColor } });
       return { name, backgroundColor, textColor, success: true };
     },
   });
 
-const bulkDelete = (agent: ZeroAgent) =>
+const bulkDelete = (connectionId: string) =>
   tool({
     description: 'Move multiple emails to trash by adding the TRASH label',
     parameters: z.object({
       threadIds: z.array(z.string()).describe('Array of email IDs to move to trash'),
     }),
     execute: async ({ threadIds }) => {
-      await agent.modifyLabels(threadIds, ['TRASH'], []);
+      const driver = await getZeroAgent(connectionId);
+      await driver.modifyLabels(threadIds, ['TRASH'], []);
       return { threadIds, success: true };
     },
   });
 
-const bulkArchive = (agent: ZeroAgent) =>
+const bulkArchive = (connectionId: string) =>
   tool({
     description: 'Move multiple emails to the archive by removing the INBOX label',
     parameters: z.object({
       threadIds: z.array(z.string()).describe('Array of email IDs to move to archive'),
     }),
     execute: async ({ threadIds }) => {
-      await agent.modifyLabels(threadIds, [], ['INBOX']);
+      const driver = await getZeroAgent(connectionId);
+      await driver.modifyLabels(threadIds, [], ['INBOX']);
       return { threadIds, success: true };
     },
   });
 
-const deleteLabel = (agent: ZeroAgent) =>
+const deleteLabel = (connectionId: string) =>
   tool({
     description: "Delete a label from the user's account",
     parameters: z.object({
       id: z.string().describe('The ID of the label to delete'),
     }),
     execute: async ({ id }) => {
-      await agent.deleteLabel(id);
+      const driver = await getZeroAgent(connectionId);
+      await driver.deleteLabel(id);
       return { id, success: true };
     },
   });
@@ -397,19 +406,19 @@ const buildGmailSearchQuery = () =>
     },
   });
 
-export const tools = async (agent: ZeroAgent, connectionId: string) => {
+export const tools = async (connectionId: string) => {
   return {
-    [Tools.GetThread]: getEmail(agent),
+    [Tools.GetThread]: getEmail(),
     [Tools.ComposeEmail]: composeEmailTool(connectionId),
-    [Tools.MarkThreadsRead]: markAsRead(agent),
-    [Tools.MarkThreadsUnread]: markAsUnread(agent),
-    [Tools.ModifyLabels]: modifyLabels(agent),
-    [Tools.GetUserLabels]: getUserLabels(agent),
-    [Tools.SendEmail]: sendEmail(agent),
-    [Tools.CreateLabel]: createLabel(agent),
-    [Tools.BulkDelete]: bulkDelete(agent),
-    [Tools.BulkArchive]: bulkArchive(agent),
-    [Tools.DeleteLabel]: deleteLabel(agent),
+    [Tools.MarkThreadsRead]: markAsRead(connectionId),
+    [Tools.MarkThreadsUnread]: markAsUnread(connectionId),
+    [Tools.ModifyLabels]: modifyLabels(connectionId),
+    [Tools.GetUserLabels]: getUserLabels(connectionId),
+    [Tools.SendEmail]: sendEmail(connectionId),
+    [Tools.CreateLabel]: createLabel(connectionId),
+    [Tools.BulkDelete]: bulkDelete(connectionId),
+    [Tools.BulkArchive]: bulkArchive(connectionId),
+    [Tools.DeleteLabel]: deleteLabel(connectionId),
     [Tools.WebSearch]: tool({
       description: 'Search the web for information using Perplexity AI',
       parameters: z.object({

apps/server/wrangler.jsonc

@@ -39,6 +39,10 @@
             "name": "ZERO_DB",
             "class_name": "ZeroDB",
           },
+          {
+            "name": "ZERO_DRIVER",
+            "class_name": "ZeroDriver",
+          },
         ],
       },
       "queues": {
@@ -78,6 +82,10 @@
           "tag": "v4",
           "deleted_classes": ["DurableMailbox"],
         },
+        {
+          "tag": "v5",
+          "new_sqlite_classes": ["ZeroDriver"],
+        },
       ],
 
       "observability": {
@@ -167,6 +175,10 @@
             "name": "ZERO_DB",
             "class_name": "ZeroDB",
           },
+          {
+            "name": "ZERO_DRIVER",
+            "class_name": "ZeroDriver",
+          },
         ],
       },
       "r2_buckets": [
@@ -216,6 +228,10 @@
           "tag": "v5",
           "deleted_classes": ["DurableMailbox"],
         },
+        {
+          "tag": "v6",
+          "new_sqlite_classes": ["ZeroDriver"],
+        },
       ],
       "observability": {
         "enabled": true,
@@ -312,6 +328,10 @@
             "name": "ZERO_DB",
             "class_name": "ZeroDB",
           },
+          {
+            "name": "ZERO_DRIVER",
+            "class_name": "ZeroDriver",
+          },
         ],
       },
       "queues": {
@@ -355,6 +375,10 @@
           "tag": "v5",
           "deleted_classes": ["DurableMailbox"],
         },
+        {
+          "tag": "v6",
+          "new_sqlite_classes": ["ZeroDriver"],
+        },
       ],
       "vars": {
         "NODE_ENV": "production",