Reply tracker for Outlook

.cursor/rules/features/reply-tracker.mdc

@@ -22,4 +22,10 @@ When enabling the reply tracker, we create a rule for the user that has the foll
 
 We'll draft a reply for the user automatically if DRAFT_EMAIL is set. The draft will be generated using the email history for this sender as well as the Knowledge base. See `.cursor/rules/features/knowledge.mdc` for more on the Knowledge Base feature.
 
-Enabling `User.outboundReplyTracking` means that when a user sends an email, we'll run an LLM over the email and check if it needs a reply. If it does we'll mark it as Awaiting Reply.
+Typically rules apply to emails as they come in. Rules have no impact afterwards. So to remove a label from a thread later we use the `TRACK_THREAD` action. This action is in charge of removing the Awaiting Reply or To Reply label. Note, adding labels is handled in the usual manner or with `outboundReplyTracking`.
+
+Enabling `EmailAccount.outboundReplyTracking` means that when a user sends an email, we'll run an LLM over the email and check if it needs a reply. If it does we'll mark it as "Awaiting Reply".
+
+Some relevant files:
+- apps/web/utils/reply-tracker/outbound.ts
+- apps/web/utils/reply-tracker/inbound.ts

apps/web/app/api/google/webhook/process-history-item.ts

@@ -24,6 +24,7 @@ import type { ParsedMessage } from "@/utils/types";
 import type { EmailAccountWithAI } from "@/utils/llms/types";
 import { formatError } from "@/utils/error";
 import { createEmailProvider } from "@/utils/email/provider";
+import type { EmailProvider } from "@/utils/email/types";
 import { enqueueDigestItem } from "@/utils/digest/index";
 import { HistoryEventType } from "@/app/api/google/webhook/types";
 import { handleLabelRemovedEvent } from "@/app/api/google/webhook/process-label-removed-event";
@@ -139,7 +140,7 @@ export async function processHistoryItem(
     const isOutbound = parsedMessage.labelIds?.includes(GmailLabel.SENT);
 
     if (isOutbound) {
-      await handleOutbound(emailAccount, parsedMessage, gmail);
+      await handleOutbound(emailAccount, parsedMessage, provider);
       return;
     }
 
@@ -241,7 +242,7 @@ export async function processHistoryItem(
 async function handleOutbound(
   emailAccount: EmailAccountWithAI,
   message: ParsedMessage,
-  gmail: gmail_v1.Gmail,
+  provider: EmailProvider,
 ) {
   const loggerOptions = {
     email: emailAccount.email,
@@ -257,9 +258,13 @@ async function handleOutbound(
     trackSentDraftStatus({
       emailAccountId: emailAccount.id,
       message,
-      gmail,
+      provider,
+    }),
+    handleOutboundReply({
+      emailAccount,
+      message,
+      provider,
     }),
-    handleOutboundReply({ emailAccount, message, gmail }),
   ]);
 
   if (trackingResult.status === "rejected") {
@@ -282,7 +287,7 @@ async function handleOutbound(
     await cleanupThreadAIDrafts({
       threadId: message.threadId,
       emailAccountId: emailAccount.id,
-      gmail,
+      provider,
     });
   } catch (cleanupError) {
     logger.error("Error during thread draft cleanup", {

apps/web/utils/email/google.ts

@@ -13,6 +13,8 @@ import {
   createLabel,
   getOrCreateInboxZeroLabel,
   GmailLabel,
+  getNeedsReplyLabel,
+  getAwaitingReplyLabel,
 } from "@/utils/gmail/label";
 import { labelVisibility, messageVisibility } from "@/utils/gmail/constants";
 import type { InboxZeroLabel } from "@/utils/label";
@@ -46,10 +48,7 @@ import {
   getThreadsWithNextPageToken,
 } from "@/utils/gmail/thread";
 import { decodeSnippet } from "@/utils/gmail/decode";
-import {
-  getAwaitingReplyLabel as getGmailAwaitingReplyLabel,
-  getReplyTrackingLabels,
-} from "@/utils/reply-tracker/label";
+import { getReplyTrackingLabels } from "@/utils/gmail/label";
 import { getDraft, deleteDraft } from "@/utils/gmail/draft";
 import {
   getFiltersList,
@@ -292,7 +291,21 @@ export class GmailProvider implements EmailProvider {
   }
 
   async getAwaitingReplyLabel(): Promise<string> {
-    return getGmailAwaitingReplyLabel(this.client);
+    return getAwaitingReplyLabel(this.client);
+  }
+
+  async getNeedsReplyLabel(): Promise<string> {
+    return getNeedsReplyLabel(this.client);
+  }
+
+  async removeAwaitingReplyLabel(threadId: string): Promise<void> {
+    const awaitingReplyLabelId = await this.getAwaitingReplyLabel();
+    await removeThreadLabel(this.client, threadId, awaitingReplyLabelId);
+  }
+
+  async removeNeedsReplyLabel(threadId: string): Promise<void> {
+    const needsReplyLabelId = await this.getNeedsReplyLabel();
+    await removeThreadLabel(this.client, threadId, needsReplyLabelId);
   }
 
   async createLabel(name: string): Promise<EmailLabel> {
@@ -617,6 +630,22 @@ export class GmailProvider implements EmailProvider {
     return getReplyTrackingLabels(this.client);
   }
 
+  async labelAwaitingReply(messageId: string, labelId: string): Promise<void> {
+    await labelMessage({
+      gmail: this.client,
+      messageId,
+      addLabelIds: [labelId],
+    });
+  }
+
+  async labelNeedsReply(messageId: string, labelId: string): Promise<void> {
+    await labelMessage({
+      gmail: this.client,
+      messageId,
+      addLabelIds: [labelId],
+    });
+  }
+
   async processHistory(options: {
     emailAddress: string;
     historyId?: number;

apps/web/utils/email/microsoft.ts

@@ -10,6 +10,7 @@ import {
   getLabels,
   createLabel,
   getOrCreateInboxZeroLabel,
+  getLabelById,
 } from "@/utils/outlook/label";
 import type { InboxZeroLabel } from "@/utils/label";
 import type { ThreadsQuery } from "@/app/api/threads/validation";
@@ -24,6 +25,7 @@ import {
   getOrCreateLabel,
   labelMessage,
   markReadThread,
+  removeThreadLabel,
 } from "@/utils/outlook/label";
 import { trashThread } from "@/utils/outlook/trash";
 import { markSpam } from "@/utils/outlook/spam";
@@ -275,9 +277,17 @@ export class OutlookProvider implements EmailProvider {
     return this.getThreadMessages(messageIds[0]);
   }
 
-  async removeThreadLabel(_threadId: string, _labelId: string): Promise<void> {
-    // For Outlook, we don't need to do anything with labels at this point
-    return Promise.resolve();
+  async removeThreadLabel(threadId: string, labelId: string): Promise<void> {
+    // TODO: this can be more efficient by using the label name directly
+    // Get the label to convert ID to name (Outlook uses names)
+    const label = await getLabelById({ client: this.client, id: labelId });
+    const categoryName = label.displayName || "";
+
+    await removeThreadLabel({
+      client: this.client,
+      threadId,
+      categoryName,
+    });
   }
 
   async createLabel(name: string): Promise<EmailLabel> {
@@ -731,15 +741,6 @@ export class OutlookProvider implements EmailProvider {
     return getThreadsFromSenderWithSubject(this.client, sender, limit);
   }
 
-  async getAwaitingReplyLabel(): Promise<string> {
-    const [awaitingReplyLabel] = await getOutlookOrCreateLabels({
-      client: this.client,
-      names: [AWAITING_REPLY_LABEL_NAME],
-    });
-
-    return awaitingReplyLabel.id || "";
-  }
-
   async getReplyTrackingLabels(): Promise<{
     awaitingReplyLabelId: string;
     needsReplyLabelId: string;
@@ -756,6 +757,56 @@ export class OutlookProvider implements EmailProvider {
     };
   }
 
+  async getNeedsReplyLabel(): Promise<string> {
+    const [needsReplyLabel] = await getOutlookOrCreateLabels({
+      client: this.client,
+      names: [NEEDS_REPLY_LABEL_NAME],
+    });
+
+    return needsReplyLabel.id || "";
+  }
+
+  async getAwaitingReplyLabel(): Promise<string> {
+    const [awaitingReplyLabel] = await getOutlookOrCreateLabels({
+      client: this.client,
+      names: [AWAITING_REPLY_LABEL_NAME],
+    });
+
+    return awaitingReplyLabel.id || "";
+  }
+
+  async labelNeedsReply(messageId: string, _labelId: string): Promise<void> {
+    await labelMessage({
+      client: this.client,
+      messageId,
+      categories: [NEEDS_REPLY_LABEL_NAME],
+    });
+  }
+
+  async labelAwaitingReply(messageId: string, _labelId: string): Promise<void> {
+    await labelMessage({
+      client: this.client,
+      messageId,
+      categories: [AWAITING_REPLY_LABEL_NAME],
+    });
+  }
+
+  async removeAwaitingReplyLabel(threadId: string): Promise<void> {
+    await removeThreadLabel({
+      client: this.client,
+      threadId,
+      categoryName: AWAITING_REPLY_LABEL_NAME,
+    });
+  }
+
+  async removeNeedsReplyLabel(threadId: string): Promise<void> {
+    await removeThreadLabel({
+      client: this.client,
+      threadId,
+      categoryName: NEEDS_REPLY_LABEL_NAME,
+    });
+  }
+
   async processHistory(options: {
     emailAddress: string;
     historyId?: number;

apps/web/utils/email/types.ts

@@ -59,7 +59,16 @@ export interface EmailProvider {
   ): Promise<void>;
   labelMessage(messageId: string, labelName: string): Promise<void>;
   removeThreadLabel(threadId: string, labelId: string): Promise<void>;
+  getReplyTrackingLabels(): Promise<{
+    awaitingReplyLabelId: string;
+    needsReplyLabelId: string;
+  }>;
+  getNeedsReplyLabel(): Promise<string>;
   getAwaitingReplyLabel(): Promise<string>;
+  labelAwaitingReply(messageId: string, labelId: string): Promise<void>;
+  labelNeedsReply(messageId: string, labelId: string): Promise<void>;
+  removeAwaitingReplyLabel(threadId: string): Promise<void>;
+  removeNeedsReplyLabel(threadId: string): Promise<void>;
   draftEmail(
     email: ParsedMessage,
     args: { to?: string; subject?: string; content: string },
@@ -137,10 +146,6 @@ export interface EmailProvider {
     sender: string,
     limit: number,
   ): Promise<Array<{ id: string; snippet: string; subject: string }>>;
-  getReplyTrackingLabels(): Promise<{
-    awaitingReplyLabelId: string;
-    needsReplyLabelId: string;
-  }>;
   processHistory(options: {
     emailAddress: string;
     historyId?: number;

apps/web/utils/gmail/label.ts

@@ -14,6 +14,10 @@ import {
 } from "@/utils/gmail/constants";
 import { createScopedLogger } from "@/utils/logger";
 import { withGmailRetry } from "@/utils/gmail/retry";
+import {
+  AWAITING_REPLY_LABEL_NAME,
+  NEEDS_REPLY_LABEL_NAME,
+} from "@/utils/reply-tracker/consts";
 
 const logger = createScopedLogger("gmail/label");
 
@@ -355,3 +359,38 @@ export async function getOrCreateInboxZeroLabel({
   });
   return createdLabel;
 }
+
+export async function getAwaitingReplyLabel(
+  gmail: gmail_v1.Gmail,
+): Promise<string> {
+  const [awaitingReplyLabel] = await getOrCreateLabels({
+    gmail,
+    names: [AWAITING_REPLY_LABEL_NAME],
+  });
+  return awaitingReplyLabel.id || "";
+}
+
+export async function getNeedsReplyLabel(
+  gmail: gmail_v1.Gmail,
+): Promise<string> {
+  const [toReplyLabel] = await getOrCreateLabels({
+    gmail,
+    names: [NEEDS_REPLY_LABEL_NAME],
+  });
+  return toReplyLabel.id || "";
+}
+
+export async function getReplyTrackingLabels(gmail: gmail_v1.Gmail): Promise<{
+  awaitingReplyLabelId: string;
+  needsReplyLabelId: string;
+}> {
+  const [awaitingReplyLabel, needsReplyLabel] = await getOrCreateLabels({
+    gmail,
+    names: [AWAITING_REPLY_LABEL_NAME, NEEDS_REPLY_LABEL_NAME],
+  });
+
+  return {
+    awaitingReplyLabelId: awaitingReplyLabel.id || "",
+    needsReplyLabelId: needsReplyLabel.id || "",
+  };
+}

apps/web/utils/outlook/label.ts

@@ -230,6 +230,59 @@ export async function labelThread({
   );
 }
 
+export async function removeThreadLabel({
+  client,
+  threadId,
+  categoryName,
+}: {
+  client: OutlookClient;
+  threadId: string;
+  categoryName: string;
+}) {
+  if (!categoryName) {
+    logger.warn("Category name is empty, skipping removal", { threadId });
+    return;
+  }
+
+  // Get all messages in the thread
+  const escapedThreadId = threadId.replace(/'/g, "''");
+  const messages = await client
+    .getClient()
+    .api("/me/messages")
+    .filter(`conversationId eq '${escapedThreadId}'`)
+    .select("id,categories")
+    .get();
+
+  // Remove the category from each message
+  await Promise.all(
+    messages.value.map(
+      async (message: { id: string; categories?: string[] }) => {
+        if (!message.categories || !message.categories.includes(categoryName)) {
+          return; // Category not present, nothing to remove
+        }
+
+        const updatedCategories = message.categories.filter(
+          (cat) => cat !== categoryName,
+        );
+
+        try {
+          await client
+            .getClient()
+            .api(`/me/messages/${message.id}`)
+            .patch({ categories: updatedCategories });
+        } catch (error) {
+          logger.warn("Failed to remove category from message", {
+            messageId: message.id,
+            threadId,
+            categoryName,
+            error: error instanceof Error ? error.message : error,
+          });
+        }
+      },
+    ),
+  );
+}
+
 export async function archiveThread({
   client,
   threadId,