jmap: fetch entire threads

JMAP was missing good thread support, especially when compared to Gmail.
This will fetch entire threads when possible.

Signed-off-by: Tristan Partin <tristan@partin.io>

1 files changed, 147 insertions, 32 deletions

diff --git a/worker/jmap/fetch.go b/worker/jmap/fetch.go
index 17b3fb2f..31fcef98 100644
--- a/worker/jmap/fetch.go
+++ b/worker/jmap/fetch.go
@@ -10,12 +10,14 @@ import (
 	"git.sr.ht/~rjarry/aerc/worker/types"
 	"git.sr.ht/~rockorager/go-jmap"
 	"git.sr.ht/~rockorager/go-jmap/mail/email"
+	"git.sr.ht/~rockorager/go-jmap/mail/thread"
 	"github.com/emersion/go-message/charset"
 )
 
 var headersProperties = []string{
 	"id",
 	"blobId",
+	"threadId",
 	"mailboxIds",
 	"keywords",
 	"size",
@@ -33,10 +35,90 @@ var headersProperties = []string{
 	"bodyStructure",
 }
 
-func (w *JMAPWorker) handleFetchMessageHeaders(msg *types.FetchMessageHeaders) error {
-	var req jmap.Request
+func (w *JMAPWorker) fetchEmailIdsFromThreads(threadIds []jmap.ID) ([]jmap.ID, error) {
+	currentThreadState, err := w.getThreadState()
+	if err != nil {
+		return nil, err
+	}
+
+	// If we can't get the cached mailbox state, at worst, we will just
+	// query information we might already know
+	cachedThreadState, err := w.cache.GetThreadState()
+	if err != nil {
+		w.w.Warnf("GetThreadState: %s", err)
+	}
+
+	consistentThreadState := currentThreadState == cachedThreadState
+
+	mailIds := make([]jmap.ID, 0)
+	getMailIds := func(threadIds []jmap.ID) error {
+		var req jmap.Request
+		var realIds []jmap.ID
+
+		if len(threadIds) > 0 {
+			realIds = threadIds
+		} else {
+			realIds = []jmap.ID{jmap.ID("00")}
+		}
+
+		req.Invoke(&thread.Get{
+			Account: w.accountId,
+			IDs:     realIds,
+		})
+
+		resp, err := w.Do(&req)
+		if err != nil {
+			return err
+		}
+
+		for _, inv := range resp.Responses {
+			switch r := inv.Args.(type) {
+			case *thread.GetResponse:
+				for _, t := range r.List {
+					mailIds = append(mailIds, t.EmailIDs...)
+				}
+			case *jmap.MethodError:
+				return wrapMethodError(r)
+			}
+		}
+
+		return nil
+	}
+
+	// If we have a consistent state, check the cache
+	if consistentThreadState {
+		missingThreadIds := make([]jmap.ID, 0, len(threadIds))
+		for _, threadId := range threadIds {
+			t, err := w.cache.GetThread(threadId)
+			if err != nil {
+				w.w.Warnf("GetThread: %s", err)
+				missingThreadIds = append(missingThreadIds, threadId)
+				continue
+			}
+			mailIds = append(mailIds, t.EmailIDs...)
+		}
+
+		if len(missingThreadIds) > 0 {
+			if err := getMailIds(missingThreadIds); err != nil {
+				return nil, err
+			}
+		}
+	} else {
+		if err := getMailIds(threadIds); err != nil {
+			return nil, err
+		}
+	}
+
+	if err := w.cache.PutThreadState(currentThreadState); err != nil {
+		w.w.Warnf("GetThreadState: %s", err)
+	}
 
-	ids := make([]jmap.ID, 0, len(msg.Uids))
+	return mailIds, nil
+}
+
+func (w *JMAPWorker) handleFetchMessageHeaders(msg *types.FetchMessageHeaders) error {
+	mailIds := make([]jmap.ID, 0)
+	threadIds := make([]jmap.ID, 0, len(msg.Uids))
 	for _, uid := range msg.Uids {
 		id, ok := w.uidStore.GetKey(uid)
 		if !ok {
@@ -44,52 +126,85 @@ func (w *JMAPWorker) handleFetchMessageHeaders(msg *types.FetchMessageHeaders) e
 		}
 		jid := jmap.ID(id)
 		m, err := w.cache.GetEmail(jid)
-		if err == nil {
+		// TODO: use ID.Valid() when my patch is merged
+		if err == nil && len(m.ThreadID) > 0 && len(m.ThreadID) < 256 {
+			threadIds = append(threadIds, m.ThreadID)
 			w.w.PostMessage(&types.MessageInfo{
 				Message: types.RespondTo(msg),
 				Info:    w.translateMsgInfo(m),
 			}, nil)
 			continue
 		}
-		ids = append(ids, jid)
+		mailIds = append(mailIds, jid)
 	}
 
-	if len(ids) == 0 {
-		return nil
-	}
-
-	req.Invoke(&email.Get{
-		Account:    w.accountId,
-		IDs:        ids,
-		Properties: headersProperties,
-	})
-
-	resp, err := w.Do(&req)
-	if err != nil {
-		return err
-	}
-
-	for _, inv := range resp.Responses {
-		switch r := inv.Args.(type) {
-		case *email.GetResponse:
-			for _, m := range r.List {
+	postMessages := func(mailIds []jmap.ID, collectThreadIds bool) error {
+		missing := make([]jmap.ID, 0, len(mailIds))
+		for _, id := range mailIds {
+			m, err := w.cache.GetEmail(id)
+			// TODO: use ID.Valid() when my patch is merged
+			if err == nil && len(m.ThreadID) > 0 && len(m.ThreadID) < 256 {
+				threadIds = append(threadIds, m.ThreadID)
 				w.w.PostMessage(&types.MessageInfo{
 					Message: types.RespondTo(msg),
 					Info:    w.translateMsgInfo(m),
 				}, nil)
-				if err := w.cache.PutEmail(m.ID, m); err != nil {
-					w.w.Warnf("PutEmail: %s", err)
-				}
+				continue
 			}
-			if err = w.cache.PutEmailState(r.State); err != nil {
-				w.w.Warnf("PutEmailState: %s", err)
+			missing = append(missing, id)
+		}
+
+		var req jmap.Request
+		req.Invoke(&email.Get{
+			Account:    w.accountId,
+			IDs:        mailIds,
+			Properties: headersProperties,
+		})
+
+		resp, err := w.Do(&req)
+		if err != nil {
+			return err
+		}
+
+		for _, inv := range resp.Responses {
+			switch r := inv.Args.(type) {
+			case *email.GetResponse:
+				for _, m := range r.List {
+					w.w.PostMessage(&types.MessageInfo{
+						Message: types.RespondTo(msg),
+						Info:    w.translateMsgInfo(m),
+					}, nil)
+					if err := w.cache.PutEmail(m.ID, m); err != nil {
+						w.w.Warnf("PutEmail: %s", err)
+					}
+
+					if collectThreadIds {
+						threadIds = append(threadIds, m.ThreadID)
+					}
+				}
+				if err = w.cache.PutEmailState(r.State); err != nil {
+					w.w.Warnf("PutEmailState: %s", err)
+				}
+			case *jmap.MethodError:
+				return wrapMethodError(r)
 			}
-		case *jmap.MethodError:
-			return wrapMethodError(r)
 		}
+
+		return nil
 	}
 
-	return nil
+	if len(mailIds) > 0 {
+		if err := postMessages(mailIds, true); err != nil {
+			return err
+		}
+	}
+
+	additionalMailIds, err := w.fetchEmailIdsFromThreads(threadIds)
+	if err != nil {
+		return err
+	}
+
+	return postMessages(additionalMailIds, false)
 }
 
 func (w *JMAPWorker) handleFetchMessageBodyPart(msg *types.FetchMessageBodyPart) error {