2 files changed, 23 insertions, 40 deletions

diff --git a/worker/jmap/fetch.go b/worker/jmap/fetch.go
index bbef1bb5..07579b99 100644
--- a/worker/jmap/fetch.go
+++ b/worker/jmap/fetch.go
@@ -44,14 +44,20 @@ func (w *JMAPWorker) handleFetchMessageHeaders(msg *types.FetchMessageHeaders) e
 		}
 		jid := jmap.ID(id)
 		m, err := w.cache.GetEmail(jid)
-		if err == nil {
-			currentEmails = append(currentEmails, m)
-		} else {
+		if err != nil {
+			// Message wasn't in cache; fetch it
 			emailIdsToFetch = append(emailIdsToFetch, jid)
+			continue
 		}
+		currentEmails = append(currentEmails, m)
+		// Get the UI updated immediately
+		w.w.PostMessage(&types.MessageInfo{
+			Message: types.RespondTo(msg),
+			Info:    w.translateMsgInfo(m),
+		}, nil)
 	}
 
-	if len(emailIdsToFetch) != 0 {
+	if len(emailIdsToFetch) > 0 {
 		var req jmap.Request
 
 		req.Invoke(&email.Get{
diff --git a/worker/jmap/threads.go b/worker/jmap/threads.go
index b967e8cd..ee4b70a9 100644
--- a/worker/jmap/threads.go
+++ b/worker/jmap/threads.go
@@ -18,54 +18,31 @@ func (w *JMAPWorker) fetchEntireThreads(emailIds []*email.Email) ([]*email.Email
 		threadsToFetch = append(threadsToFetch, m.ThreadID)
 	}
 
-	req.Invoke(&thread.Get{
+	threadGetId := req.Invoke(&thread.Get{
 		Account: w.AccountId(),
 		IDs:     threadsToFetch,
 	})
 
-	resp, err := w.Do(&req)
-	if err != nil {
-		return nil, err
-	}
-
-	emailsToFetch := make([]jmap.ID, 0)
-	emailsToReturn := make([]*email.Email, 0)
-	for _, inv := range resp.Responses {
-		switch r := inv.Args.(type) {
-		case *thread.GetResponse:
-			for _, t := range r.List {
-				for _, emailId := range t.EmailIDs {
-					m, err := w.cache.GetEmail(emailId)
-					if err == nil || m == nil {
-						emailsToFetch = append(emailsToFetch, emailId)
-					} else {
-						emailsToReturn = append(emailsToReturn, m)
-					}
-				}
-			}
-			if err = w.cache.PutThreadState(r.State); err != nil {
-				w.w.Warnf("PutThreadState: %s", err)
-			}
-		case *jmap.MethodError:
-			return nil, wrapMethodError(r)
-		}
-	}
-
-	if len(emailsToFetch) == 0 {
-		return emailsToReturn, nil
-	}
-
+	// Opportunistically fetch all emails in this thread. We could wait for
+	// the result, check which ones we don't have, then fetch only those.
+	// However we can do this all in a single request which ends up being
+	// faster than two requests for most contexts
 	req.Invoke(&email.Get{
-		Account:    w.AccountId(),
-		IDs:        emailsToFetch,
+		Account: w.AccountId(),
+		ReferenceIDs: &jmap.ResultReference{
+			ResultOf: threadGetId,
+			Name:     "Thread/get",
+			Path:     "/list/*/emailIds",
+		},
 		Properties: headersProperties,
 	})
 
-	resp, err = w.Do(&req)
+	resp, err := w.Do(&req)
 	if err != nil {
 		return nil, err
 	}
 
+	emailsToReturn := make([]*email.Email, 0)
 	for _, inv := range resp.Responses {
 		switch r := inv.Args.(type) {
 		case *email.GetResponse: