Allow user to move back SPAM, fix build

Signed-off-by: Daishan Peng <[email protected]>

pkg/server/auth/refreshToken.go

@@ -58,6 +58,7 @@ func RefreshToken(ctx context.Context, queries *db.Queries) {
 						SubscriptionID:       user.SubscriptionID,
 						SubscriptionExpireAt: user.SubscriptionExpireAt,
 						SubscriptionDisabled: user.SubscriptionDisabled,
+						CheckSpam:            user.CheckSpam,
 					}); err != nil {
 						logrus.Error(fmt.Errorf("failed to update user after token refresh: %w", err))
 						continue

pkg/server/spam/spam.go

@@ -8,6 +8,7 @@ import (
 
 	"ethan/pkg/db"
 	"ethan/pkg/mstoken"
+	"ethan/pkg/server/subscribe"
 	"github.com/gorilla/mux"
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgtype"
@@ -101,11 +102,15 @@ func (h *Handler) MoveSpam(w http.ResponseWriter, r *http.Request) {
 	requestBody := graphusers.NewItemMailfoldersItemMessagesItemMovePostRequestBody()
 	inboxFolderID := "inbox"
 	requestBody.SetDestinationId(&inboxFolderID)
-	if _, err := client.Me().Messages().ByMessageId(*spamEmail.MessageID).Move().Post(r.Context(), requestBody, nil); err != nil {
+	newMessage, err := client.Me().Messages().ByMessageId(*spamEmail.MessageID).Move().Post(r.Context(), requestBody, nil)
+	if err != nil {
 		logrus.Errorf("Failed to move message to back to inbox, error: %v", err)
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
+	// After we move the message back to Inbox, we need to skip checking this email because it has been falsely detected as spam
+	// For now we temporarily store the message ID into memory map. This is not going to work in HA but don't worry about it now.
+	subscribe.SkipEmails[*newMessage.GetId()] = struct{}{}
 
 	logrus.Infof("Move spam email %s to %s", *spamEmail.MessageID, inboxFolderID)
 

pkg/server/subscribe/refreshSubscription.go

@@ -79,6 +79,7 @@ func ensureSubscriptionsForUser(ctx context.Context, user db.User, queries *db.Q
 				SubscriptionID:       nil,
 				SubscriptionExpireAt: pgtype.Timestamptz{},
 				SubscriptionDisabled: user.SubscriptionDisabled,
+				CheckSpam:            user.CheckSpam,
 			}); err != nil {
 				return err
 			}
@@ -100,6 +101,7 @@ func ensureSubscriptionsForUser(ctx context.Context, user db.User, queries *db.Q
 			SubscriptionID:       &subscriptionID,
 			SubscriptionExpireAt: t,
 			SubscriptionDisabled: user.SubscriptionDisabled,
+			CheckSpam:            user.CheckSpam,
 		}); err != nil {
 			return err
 		}

pkg/server/subscribe/webhook.go

@@ -28,6 +28,10 @@ import (
 	"github.com/sirupsen/logrus"
 )
 
+var (
+	SkipEmails = map[string]struct{}{}
+)
+
 var checkSpamTemplatePrompt = `
 Given email body: %v, email sender: %v, email subject: %v, Check if email belongs to a cold email.
 Do not mark it as cold email of the sender has an email address that is the same domain as yours email(%v).
@@ -124,76 +128,78 @@ func (h *Handler) Subscribe(w http.ResponseWriter, r *http.Request) {
 
 		if user.CheckSpam != nil && *user.CheckSpam {
 			// Once we identified te the email is related to meeting, use AI to check whether email belongs to cold email. If so, move it to spam
-			checkSpamRun, err := gptClient.Evaluate(context.Background(), gptscript.Options{}, gptscript.ToolDef{
-				Instructions: fmt.Sprintf(checkSpamTemplatePrompt, emailContent, email, subject, user.Email),
-			})
-			if err != nil {
-				logrus.Error(fmt.Errorf("failed to run gptscript to check email content to detect spam: %w", err))
-				w.WriteHeader(http.StatusInternalServerError)
-				return
-			}
-			checkSpamRunOutput, err := checkSpamRun.Text()
-			if err != nil {
-				logrus.Error(fmt.Errorf("failed to run gptscript to check email content to detect spam: %w", err))
-				w.WriteHeader(http.StatusInternalServerError)
-				return
-			}
-
-			if strings.ToLower(checkSpamRunOutput) == "yes" {
-				logrus.Infof("Mark message %v as Spam cold email, moving to Cold Email folder", messageID)
-
-				folders, err := client.Me().MailFolders().Get(r.Context(), nil)
+			if _, ok := SkipEmails[messageID]; !ok {
+				checkSpamRun, err := gptClient.Evaluate(context.Background(), gptscript.Options{}, gptscript.ToolDef{
+					Instructions: fmt.Sprintf(checkSpamTemplatePrompt, emailContent, email, subject, user.Email),
+				})
 				if err != nil {
-					logrus.Error(fmt.Errorf("failed to get folder list: %w", err))
+					logrus.Error(fmt.Errorf("failed to run gptscript to check email content to detect spam: %w", err))
 					w.WriteHeader(http.StatusInternalServerError)
 					return
 				}
-				var coldEmailFolderID string
-				for _, folder := range folders.GetValue() {
-					if folder.GetDisplayName() != nil && *folder.GetDisplayName() == "Cold Emails" {
-						coldEmailFolderID = *folder.GetId()
-						break
-					}
-				}
-
-				if coldEmailFolderID == "" {
-					logrus.Error("Failed to find cold email folder")
-					return
-				}
-
-				requestBody := graphusers.NewItemMailfoldersItemMessagesItemMovePostRequestBody()
-				requestBody.SetDestinationId(&coldEmailFolderID)
-				newMessage, err := client.Me().Messages().ByMessageId(messageID).Move().Post(r.Context(), requestBody, nil)
+				checkSpamRunOutput, err := checkSpamRun.Text()
 				if err != nil {
-					logrus.Error("Failed to move message to junk items")
+					logrus.Error(fmt.Errorf("failed to run gptscript to check email content to detect spam: %w", err))
 					w.WriteHeader(http.StatusInternalServerError)
 					return
 				}
 
-				if err := h.queries.CreateSpamEmailRecord(r.Context(), db.CreateSpamEmailRecordParams{
-					Subject:   &subject,
-					EmailBody: &emailContent,
-					UserID:    user.ID,
-					MessageID: newMessage.GetId(),
-				}); err != nil {
-					logrus.Error(fmt.Errorf("failed to create spam email record: %w", err))
-					w.WriteHeader(http.StatusInternalServerError)
-					return
-				}
+				if strings.ToLower(checkSpamRunOutput) == "yes" {
+					logrus.Infof("Mark message %v as Spam cold email, moving to Cold Email folder", messageID)
 
-				if err := h.queries.CreateMessage(r.Context(), db.CreateMessageParams{
-					MessageID: newMessage.GetId(),
-					Content:   &[]string{fmt.Sprint("Mark incoming email as SPAM")}[0],
-					UserID:    user.ID,
-					TaskID: pgtype.UUID{
-						Valid: false,
-					},
-				}); err != nil {
-					logrus.Error(fmt.Errorf("failed to create spam email message: %w", err))
-					w.WriteHeader(http.StatusInternalServerError)
+					folders, err := client.Me().MailFolders().Get(r.Context(), nil)
+					if err != nil {
+						logrus.Error(fmt.Errorf("failed to get folder list: %w", err))
+						w.WriteHeader(http.StatusInternalServerError)
+						return
+					}
+					var coldEmailFolderID string
+					for _, folder := range folders.GetValue() {
+						if folder.GetDisplayName() != nil && *folder.GetDisplayName() == "Cold Emails" {
+							coldEmailFolderID = *folder.GetId()
+							break
+						}
+					}
+
+					if coldEmailFolderID == "" {
+						logrus.Error("Failed to find cold email folder")
+						return
+					}
+
+					requestBody := graphusers.NewItemMailfoldersItemMessagesItemMovePostRequestBody()
+					requestBody.SetDestinationId(&coldEmailFolderID)
+					newMessage, err := client.Me().Messages().ByMessageId(messageID).Move().Post(r.Context(), requestBody, nil)
+					if err != nil {
+						logrus.Error("Failed to move message to junk items")
+						w.WriteHeader(http.StatusInternalServerError)
+						return
+					}
+
+					if err := h.queries.CreateSpamEmailRecord(r.Context(), db.CreateSpamEmailRecordParams{
+						Subject:   &subject,
+						EmailBody: &emailContent,
+						UserID:    user.ID,
+						MessageID: newMessage.GetId(),
+					}); err != nil {
+						logrus.Error(fmt.Errorf("failed to create spam email record: %w", err))
+						w.WriteHeader(http.StatusInternalServerError)
+						return
+					}
+
+					if err := h.queries.CreateMessage(r.Context(), db.CreateMessageParams{
+						MessageID: newMessage.GetId(),
+						Content:   &[]string{fmt.Sprint("Mark incoming email as SPAM")}[0],
+						UserID:    user.ID,
+						TaskID: pgtype.UUID{
+							Valid: false,
+						},
+					}); err != nil {
+						logrus.Error(fmt.Errorf("failed to create spam email message: %w", err))
+						w.WriteHeader(http.StatusInternalServerError)
+						return
+					}
 					return
 				}
-				return
 			}
 		}