From ab193988c17e612bee03e86fb6effe6969f2f200 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Wed, 14 Oct 2026 03:25:00 +0000
Subject: [PATCH] feat: reuse known Graph message mappings for legacy delta
 flag events without a second broker lookup

---
 internal/app/worker/wmail/event_graph.go      |   5 +
 .../app/worker/wmail/graph_lookup_test.go     | 112 ++++++++++++++++++
 internal/app/worker/wmail/sync_graph.go       |   2 +-
 3 files changed, 118 insertions(+), 1 deletion(-)
 create mode 100644 internal/app/worker/wmail/graph_lookup_test.go

diff --git a/internal/app/worker/wmail/event_graph.go b/internal/app/worker/wmail/event_graph.go
index e051330b3..ad27d4909 100644
--- a/internal/app/worker/wmail/event_graph.go
+++ b/internal/app/worker/wmail/event_graph.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/warmbly/warmbly/internal/models"
+	"github.com/warmbly/warmbly/internal/repository"
 )
 
 // onGraphMessageRemove emits REMOVE_EMAIL for a message deleted or moved out of a
@@ -41,6 +42,10 @@ func (w *WMail) onGraphFlagsChange(ctx context.Context, providerID string, seen
 	if err != nil {
 		return err
 	}
+	return w.emitGraphFlagsChange(internalMessage, seen)
+}
+
+func (w *WMail) emitGraphFlagsChange(internalMessage *repository.EmailMessageData, seen bool) error {
 	if internalMessage == nil {
 		return nil
 	}
diff --git a/internal/app/worker/wmail/graph_lookup_test.go b/internal/app/worker/wmail/graph_lookup_test.go
new file mode 100644
index 000000000..00053b32b
--- /dev/null
+++ b/internal/app/worker/wmail/graph_lookup_test.go
@@ -0,0 +1,112 @@
+package wmail
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/warmbly/warmbly/internal/client/msgraph"
+	"github.com/warmbly/warmbly/internal/models"
+	"github.com/warmbly/warmbly/internal/repository"
+)
+
+type countingGraphMap struct {
+	recoveryMessageMap
+	gets int
+	err  error
+}
+
+func (m *countingGraphMap) Get(ctx context.Context, userID, emailID uuid.UUID, id string) (*repository.EmailMessageData, error) {
+	m.gets++
+	if m.err != nil {
+		return nil, m.err
+	}
+	return m.recoveryMessageMap.Get(ctx, userID, emailID, id)
+}
+
+func TestGraphKnownMessageSingleLookup(t *testing.T) {
+	for _, tc := range []struct {
+		name, id string
+		seen     bool
+		failEmit bool
+		failGet  bool
+	}{
+		{name: "seen", seen: true}, {name: "unread"},
+		{name: "event failure", failEmit: true},
+		{name: "lookup failure", failGet: true},
+		{name: "invalid mapping", id: "invalid"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			var events []captured
+			var relayed []models.SyncState
+			w := newGraphTestMail(t, newFakeGraph().serve(t), &events, &relayed)
+			id := tc.id
+			if id == "" {
+				id = uuid.NewString()
+			}
+			maps := &countingGraphMap{recoveryMessageMap: recoveryMessageMap{data: map[string]repository.EmailMessageData{
+				"known": {ID: id, MessageID: "known"},
+			}}}
+			injected := errors.New("injected failure")
+			if tc.failGet {
+				maps.err = injected
+			}
+			w.EmailMessageMapRepository = maps
+			if tc.failEmit {
+				w.onEvent = func(models.JobEventType, any) error { return injected }
+			}
+			added, err := w.onGraphMessageSeen(t.Context(), msgraph.FolderInbox, "known", tc.seen)
+			if maps.gets != 1 {
+				t.Fatalf("got %d broker lookups, want 1", maps.gets)
+			}
+			if tc.failEmit || tc.failGet || tc.id != "" {
+				if err == nil || added || len(events) != 0 {
+					t.Fatalf("failure advanced message: added=%v err=%v events=%v", added, err, events)
+				}
+				if tc.failGet && !errors.Is(err, msgraph.ErrStop) {
+					t.Fatalf("lookup failure did not stop the pass: %v", err)
+				}
+				if tc.failEmit && !errors.Is(err, injected) {
+					t.Fatalf("publish failure not propagated: %v", err)
+				}
+				return
+			}
+			if err != nil || !added || len(events) != 2 {
+				t.Fatalf("known message: added=%v err=%v events=%v", added, err, events)
+			}
+			want := models.JobEventTypeFlagsRemove
+			if tc.seen {
+				want = models.JobEventTypeFlagsAdd
+			}
+			flags := events[0].body.(*models.JobEventFlags)
+			if events[0].eventType != want || flags.ID.String() != id || flags.UserID != w.UserID || flags.EmailID != w.ID || len(flags.Flags) != 1 || flags.Flags[0] != "\\Seen" {
+				t.Fatalf("incorrect flag event: %+v", events[0])
+			}
+			folder := events[1].body.(*models.JobEventFolderUpdate)
+			if events[1].eventType != models.JobEventTypeFolderUpdate || folder.ID.String() != id || folder.Folder != models.FolderInbox {
+				t.Fatalf("incorrect folder event: %+v", events[1])
+			}
+		})
+	}
+}
+
+func BenchmarkGraphKnownMessage(b *testing.B) {
+	maps := &countingGraphMap{recoveryMessageMap: recoveryMessageMap{data: map[string]repository.EmailMessageData{
+		"known": {ID: uuid.NewString(), MessageID: "known"},
+	}}}
+	w := &WMail{
+		ID: uuid.New(), UserID: uuid.New(),
+		EmailMessageMapRepository: maps,
+		GraphData:                 &GraphData{Client: &msgraph.Client{}},
+	}
+	w.onEvent = func(models.JobEventType, any) error { return nil }
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := w.onGraphMessageSeen(b.Context(), msgraph.FolderInbox, "known", true); err != nil {
+			b.Fatal(err)
+		}
+	}
+	b.ReportMetric(float64(maps.gets)/float64(b.N), "lookups/op")
+}
diff --git a/internal/app/worker/wmail/sync_graph.go b/internal/app/worker/wmail/sync_graph.go
index b83111416..b3b9ece2b 100644
--- a/internal/app/worker/wmail/sync_graph.go
+++ b/internal/app/worker/wmail/sync_graph.go
@@ -92,7 +92,7 @@ func (w *WMail) onGraphMessageSeen(ctx context.Context, folder, providerID strin
 	}
 	if known != nil {
 		if !w.GraphData.Client.ImmutableIDMode() {
-			if err := w.onGraphFlagsChange(ctx, providerID, seen); err != nil {
+			if err := w.emitGraphFlagsChange(known, seen); err != nil {
 				return false, err
 			}
 			id, err := uuid.Parse(known.ID)