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>

5 files changed, 223 insertions, 32 deletions

diff --git a/worker/jmap/cache/gob.go b/worker/jmap/cache/gob.go
index f1b8be33..8b153c50 100644
--- a/worker/jmap/cache/gob.go
+++ b/worker/jmap/cache/gob.go
@@ -6,10 +6,12 @@ import (
 
 	"git.sr.ht/~rockorager/go-jmap/mail/email"
 	"git.sr.ht/~rockorager/go-jmap/mail/mailbox"
+	"git.sr.ht/~rockorager/go-jmap/mail/thread"
 )
 
 type jmapObject interface {
 	*email.Email |
+		*thread.Thread |
 		*email.QueryResponse |
 		*mailbox.Mailbox |
 		*FolderContents |
diff --git a/worker/jmap/cache/state.go b/worker/jmap/cache/state.go
index 5fec5034..6538ccad 100644
--- a/worker/jmap/cache/state.go
+++ b/worker/jmap/cache/state.go
@@ -12,6 +12,18 @@ func (c *JMAPCache) PutMailboxState(state string) error {
 	return c.put(mailboxStateKey, []byte(state))
 }
 
+func (c *JMAPCache) GetThreadState() (string, error) {
+	buf, err := c.get(threadStateKey)
+	if err != nil {
+		return "", err
+	}
+	return string(buf), nil
+}
+
+func (c *JMAPCache) PutThreadState(state string) error {
+	return c.put(threadStateKey, []byte(state))
+}
+
 func (c *JMAPCache) GetEmailState() (string, error) {
 	buf, err := c.get(emailStateKey)
 	if err != nil {
@@ -27,4 +39,5 @@ func (c *JMAPCache) PutEmailState(state string) error {
 const (
 	mailboxStateKey = "state/mailbox"
 	emailStateKey   = "state/email"
+	threadStateKey  = "state/thread"
 )
diff --git a/worker/jmap/cache/thread.go b/worker/jmap/cache/thread.go
new file mode 100644
index 00000000..ca91a4d9
--- /dev/null
+++ b/worker/jmap/cache/thread.go
@@ -0,0 +1,35 @@
+package cache
+
+import (
+	"git.sr.ht/~rockorager/go-jmap"
+	"git.sr.ht/~rockorager/go-jmap/mail/thread"
+)
+
+func (c *JMAPCache) GetThread(id jmap.ID) (*thread.Thread, error) {
+	buf, err := c.get(threadKey(id))
+	if err != nil {
+		return nil, err
+	}
+	e := new(thread.Thread)
+	err = unmarshal(buf, e)
+	if err != nil {
+		return nil, err
+	}
+	return e, nil
+}
+
+func (c *JMAPCache) PutThread(id jmap.ID, e *thread.Thread) error {
+	buf, err := marshal(e)
+	if err != nil {
+		return err
+	}
+	return c.put(threadKey(id), buf)
+}
+
+func (c *JMAPCache) DeleteThread(id jmap.ID) error {
+	return c.delete(threadKey(id))
+}
+
+func threadKey(id jmap.ID) string {
+	return "thread/" + string(id)
+}
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 {
diff --git a/worker/jmap/state.go b/worker/jmap/state.go
index 3dbab3fb..833bd151 100644
--- a/worker/jmap/state.go
+++ b/worker/jmap/state.go
@@ -3,6 +3,7 @@ package jmap
 import (
 	"git.sr.ht/~rockorager/go-jmap"
 	"git.sr.ht/~rockorager/go-jmap/mail/mailbox"
+	"git.sr.ht/~rockorager/go-jmap/mail/thread"
 )
 
 func (w *JMAPWorker) getMailboxState() (string, error) {
@@ -27,3 +28,28 @@ func (w *JMAPWorker) getMailboxState() (string, error) {
 	// This should be an impossibility
 	return "", nil
 }
+
+func (w *JMAPWorker) getThreadState() (string, error) {
+	var req jmap.Request
+
+	// TODO: This is a junk JMAP ID because Go's JSON serialization doesn't
+	// send empty slices as arrays, WTF.
+	req.Invoke(&thread.Get{Account: w.accountId, IDs: []jmap.ID{jmap.ID("00")}})
+	resp, err := w.Do(&req)
+	if err != nil {
+		return "", err
+	}
+
+	for _, inv := range resp.Responses {
+		switch r := inv.Args.(type) {
+		case *thread.GetResponse:
+			return r.State, nil
+		case *jmap.MethodError:
+			return "", wrapMethodError(r)
+
+		}
+	}
+
+	// This should be an impossibility
+	return "", nil
+}