Implement Gmail push notifications with workflows

apps/mail/components/party.tsx

@@ -35,7 +35,7 @@ export const NotificationProvider = ({ headers }: { headers: Record<string, stri
 
   usePartySocket({
     party: 'zero-agent',
-    room: activeConnection?.id ? `${activeConnection.id}` : 'general',
+    room: activeConnection?.id ? String(activeConnection.id) : 'general',
     prefix: 'agents',
     maxRetries: 1,
     query: {

apps/mail/components/ui/ai-sidebar.tsx

@@ -357,7 +357,7 @@ function AISidebar({ className }: AISidebarProps) {
 
   const agent = useAgent({
     agent: 'ZeroAgent',
-    name: activeConnection?.id ?? 'general',
+    name: activeConnection?.id ? String(activeConnection.id) : 'general',
     host: `${import.meta.env.VITE_PUBLIC_BACKEND_URL}`,
   });
 
@@ -452,7 +452,7 @@ function AISidebar({ className }: AISidebarProps) {
                 defaultSize={24}
                 minSize={24}
                 maxSize={24}
-                className="bg-panelLight dark:bg-panelDark mb-1 mr-1 hidden h-[calc(100dvh-8px)] shadow-sm md:block md:rounded-2xl md:shadow-sm "
+                className="bg-panelLight dark:bg-panelDark mb-1 mr-1 hidden h-[calc(100dvh-8px)] shadow-sm md:block md:rounded-2xl md:shadow-sm"
               >
                 <div className={cn('h-[calc(98vh)]', 'flex flex-col', '', className)}>
                   <div className="flex h-full flex-col">

apps/server/package.json

@@ -43,6 +43,7 @@
     "autumn-js": "catalog:",
     "base64-js": "1.5.1",
     "better-auth": "catalog:",
+    "cheerio": "1.1.0",
     "date-fns": "^4.1.0",
     "dedent": "^1.6.0",
     "drizzle-orm": "catalog:",
@@ -81,4 +82,4 @@
     "jiti": "2.4.2",
     "typescript": "catalog:"
   }
-}
+}

apps/server/src/lib/auth.ts

@@ -13,9 +13,10 @@ import { drizzleAdapter } from 'better-auth/adapters/drizzle';
 import { getSocialProviders } from './auth-providers';
 import { redis, resend, twilio } from './services';
 import { getContext } from 'hono/context-storage';
-import { getActiveDriver } from './driver/utils';
 import { defaultUserSettings } from './schemas';
+import { disableBrainFunction } from './brain';
 import { APIError } from 'better-auth/api';
+import type { EProviders } from '../types';
 import type { HonoContext } from '../ctx';
 import { env } from 'cloudflare:workers';
 import { createDriver } from './driver';
@@ -57,11 +58,13 @@ const connectionHandlerHook = async (account: Account) => {
     expiresAt: new Date(Date.now() + (account.accessTokenExpiresAt?.getTime() || 3600000)),
   };
 
-  await c.var.db
+  const connectionId = crypto.randomUUID();
+
+  const [result] = await c.var.db
     .insert(connection)
     .values({
       providerId: account.providerId as 'google' | 'microsoft',
-      id: crypto.randomUUID(),
+      id: connectionId,
       email: userInfo.address,
       userId: account.userId,
       createdAt: new Date(),
@@ -74,7 +77,13 @@ const connectionHandlerHook = async (account: Account) => {
         ...updatingInfo,
         updatedAt: new Date(),
       },
-    });
+    })
+    .returning({ id: connection.id });
+
+  await env.subscribe_queue.send({
+    connectionId: result.id,
+    providerId: account.providerId,
+  });
 };
 
 export const createAuth = () => {
@@ -111,6 +120,10 @@ export const createAuth = () => {
             await Promise.allSettled(
               connections.map(async (connection) => {
                 if (!connection.accessToken || !connection.refreshToken) return false;
+                await disableBrainFunction({
+                  id: connection.id,
+                  providerId: connection.providerId as EProviders,
+                });
                 const driver = createDriver(connection.providerId, {
                   auth: {
                     accessToken: connection.accessToken,

apps/server/src/lib/brain.fallback.prompts.ts

@@ -1,4 +1,7 @@
-export const SummarizeMessage = `
+import { defaultLabels } from '../types';
+import dedent from 'dedent';
+
+export const SummarizeMessage = dedent`
   <system_prompt>
       <role>You are a high-accuracy email summarization agent. Your task is to extract and summarize emails in XML format with absolute precision, ensuring no critical details are lost while maintaining high efficiency.</role>
 
@@ -46,7 +49,7 @@ export const SummarizeMessage = `
 
       <strict_guidelines>Strictly follow these rules. No missing details. No extra fluff. Just precise, high-performance summarization. Never say "Here is"</strict_guidelines>
   </system_prompt>`;
-export const SummarizeThread = `
+export const SummarizeThread = dedent`
   <system_prompt>
       <role>You are a high-accuracy email thread summarization agent. Your task is to process a full email thread with multiple messages and generate a structured, limited-length summary that retains all critical details, ensuring no information is lost.</role>
 
@@ -119,7 +122,7 @@ export const SummarizeThread = `
       <strict_guidelines>Never say "Here is"</strict_guidelines>
   </system_prompt>
   `;
-export const ReSummarizeThread = `
+export const ReSummarizeThread = dedent`
   <system_prompt>
       <role>You are a high-accuracy email thread summarization agent. Your task is to process a full email thread, including new messages and an existing summary, and generate a structured, limited-length updated summary that retains all critical details.</role>
 
@@ -190,3 +193,65 @@ export const ReSummarizeThread = `
 
       <strict_guidelines>Maintain absolute accuracy. No missing details. No extra assumptions. No modifications to previous content beyond appending updates. Ensure clarity and brevity within the length limit. Never say "Here is"</strict_guidelines>
   </system_prompt>`;
+
+export const ThreadLabels = (labels: { name: string; usecase: string }[]) => dedent`
+  <system_prompt>
+      <role>You are a precise thread labeling agent. Your task is to analyze email thread summaries and assign relevant labels from a predefined set, ensuring accurate categorization while maintaining consistency.</role>
+      <strict_guidelines>Maintain absolute accuracy in labeling. Use only the predefined labels. Never generate new labels. Never include personal names. Always return labels in comma-separated format without spaces.</strict_guidelines>
+      <strict_guidelines>Never say "Here is" or explain the process of labeling.</strict_guidelines>
+      <instructions>
+          <input_structure>
+              <item>Thread summary containing participants, messages, and context</item>
+          </input_structure>
+
+          <labeling_rules>
+          <item>Use only the predefined set of labels</item>
+          <item>Return labels as comma-separated values without spaces</item>
+          <item>Include company names as labels when heavily referenced</item>
+          <item>Include bank names as labels when heavily referenced</item>
+          <item>Do not use personal names as labels</item>
+          <item>Choose the most relevant labels, typically 1-3 labels per thread</item>
+          </labeling_rules>
+
+          <allowed_labels>
+          ${labels
+            .map(
+              (label) => `<item>
+          <name>${label.name}</name>
+          <usecase>${defaultLabels.find((e) => e.name === label.name)?.usecase || ''}</usecase>    
+          </item>`,
+            )
+            .join('\n')}
+          </allowed_labels>
+      </instructions>
+
+      <example_input>
+          <thread_summary>
+              Thread: Product Launch Planning
+              Participants: Sarah, Mike, David
+
+              - March 15, 10:00 AM - Sarah requests urgent review of the new feature documentation before the launch.
+              - March 15, 11:30 AM - Mike suggests changes to the marketing strategy for better customer engagement.
+              - March 15, 2:00 PM - David approves the final product specifications and sets a launch date.
+          </thread_summary>
+      </example_input>
+
+      <expected_output>
+          <labels>urgent,product,marketing</labels>
+      </expected_output>
+
+      <example_input>
+          <thread_summary>
+              Thread: Stripe Integration Update
+              Participants: Alex, Jamie, Stripe Support
+
+              - March 16, 9:00 AM - Alex reports issues with Stripe payment processing.
+              - March 16, 10:15 AM - Stripe Support provides troubleshooting steps.
+              - March 16, 11:30 AM - Jamie confirms the fix and requests additional security review.
+          </thread_summary>
+      </example_input>
+
+      <expected_output>
+          <labels>support,finance,stripe</labels>
+      </expected_output>
+  </system_prompt>`;

apps/server/src/lib/brain.ts

@@ -1,19 +1,26 @@
 import { ReSummarizeThread, SummarizeMessage, SummarizeThread } from './brain.fallback.prompts';
+import { getSubscriptionFactory } from './factories/subscription-factory.registry';
+import { EPrompts, EProviders } from '../types';
 import { env } from 'cloudflare:workers';
-import { EPrompts } from '../types';
 
-export const enableBrainFunction = async (connection: { id: string; providerId: string }) => {
-  return await env.zero.subscribe({
-    connectionId: connection.id,
-    providerId: connection.providerId,
-  });
+export const enableBrainFunction = async (connection: { id: string; providerId: EProviders }) => {
+  const subscriptionFactory = getSubscriptionFactory(connection.providerId);
+  const response = await subscriptionFactory.subscribe({ body: { connectionId: connection.id } });
+  if (!response.ok) {
+    throw new Error(`Failed to enable brain function: ${response.status} ${response.statusText}`);
+  }
+  return response;
 };
 
-export const disableBrainFunction = async (connection: { id: string; providerId: string }) => {
-  return await env.zero.unsubscribe({
-    connectionId: connection.id,
-    providerId: connection.providerId,
+export const disableBrainFunction = async (connection: { id: string; providerId: EProviders }) => {
+  const subscriptionFactory = getSubscriptionFactory(connection.providerId);
+  const response = await subscriptionFactory.unsubscribe({
+    body: { connectionId: connection.id, providerId: connection.providerId },
   });
+  if (!response.ok) {
+    throw new Error(`Failed to disable brain function: ${response.status} ${response.statusText}`);
+  }
+  return response;
 };
 
 const getPromptName = (connectionId: string, prompt: EPrompts) => {

apps/server/src/lib/driver/google.ts

@@ -44,7 +44,24 @@ export class GoogleMailManager implements MailManager {
       'https://www.googleapis.com/auth/userinfo.email',
     ].join(' ');
   }
-  public getAttachment(messageId: string, attachmentId: string) {
+  public async listHistory<T>(historyId: string): Promise<{ history: T[]; historyId: string }> {
+    return this.withErrorHandler(
+      'listHistory',
+      async () => {
+        const response = await this.gmail.users.history.list({
+          userId: 'me',
+          startHistoryId: historyId,
+        });
+
+        const history = response.data.history || [];
+        const nextHistoryId = response.data.historyId || historyId;
+
+        return { history: history as T[], historyId: nextHistoryId };
+      },
+      { historyId },
+    );
+  }
+  public async getAttachment(messageId: string, attachmentId: string) {
     return this.withErrorHandler(
       'getAttachment',
       async () => {

apps/server/src/lib/driver/types.ts

@@ -68,6 +68,7 @@ export interface MailManager {
     tokens?: ManagerConfig['auth'],
   ): Promise<{ address: string; name: string; photo: string }>;
   getScope(): string;
+  listHistory<T>(historyId: string): Promise<{ history: T[]; historyId: string }>;
   markAsRead(threadIds: string[]): Promise<void>;
   markAsUnread(threadIds: string[]): Promise<void>;
   normalizeIds(id: string[]): { threadIds: string[] };

apps/server/src/lib/factories/base-subscription.factory.ts

@@ -0,0 +1,44 @@
+import { defaultLabels, EProviders, type AppContext } from '../../types';
+import { connection } from '../../db/schema';
+import { env } from 'cloudflare:workers';
+import { createDb } from '../../db';
+
+export interface SubscriptionData {
+  connectionId?: string;
+  silent?: boolean;
+  force?: boolean;
+}
+
+export interface UnsubscriptionData {
+  connectionId?: string;
+  providerId?: EProviders;
+}
+
+export abstract class BaseSubscriptionFactory {
+  abstract readonly providerId: EProviders;
+
+  abstract subscribe(data: { body: SubscriptionData }): Promise<Response>;
+
+  abstract unsubscribe(data: { body: UnsubscriptionData }): Promise<Response>;
+
+  abstract verifyToken(token: string): Promise<boolean>;
+
+  protected async getConnectionFromDb(connectionId: string): Promise<any> {
+    const db = createDb(env.HYPERDRIVE.connectionString);
+    const { eq } = await import('drizzle-orm');
+
+    const [connectionData] = await db
+      .select()
+      .from(connection)
+      .where(eq(connection.id, connectionId));
+
+    return connectionData;
+  }
+
+  protected async initializeConnectionLabels(connectionId: string): Promise<void> {
+    const existingLabels = await env.connection_labels.get(connectionId);
+    if (!existingLabels?.trim().length) {
+      await env.connection_labels.put(connectionId, JSON.stringify(defaultLabels));
+    }
+  }
+}