Fix email threads

pkg/tools/clients/gmail.go

@@ -341,7 +341,11 @@ func stripGmailHTML(html string) string {
 }
 
 func (g *GmailClient) createDraft(ctx context.Context, token, to, subject, body, threadID string) (map[string]any, error) {
-	encoded := base64.RawURLEncoding.EncodeToString([]byte(buildRawEmail(to, subject, body)))
+	var inReplyTo, references string
+	if threadID != "" {
+		inReplyTo, references = g.getThreadReplyHeaders(ctx, token, threadID)
+	}
+	encoded := base64.RawURLEncoding.EncodeToString([]byte(buildRawEmail(to, subject, body, inReplyTo, references)))
 	payload := map[string]any{
 		"message": map[string]any{
 			"raw": encoded,
@@ -373,7 +377,11 @@ func (g *GmailClient) createDraft(ctx context.Context, token, to, subject, body,
 }
 
 func (g *GmailClient) sendEmail(ctx context.Context, token, to, subject, body, threadID, draftID string) (map[string]any, error) {
-	raw := base64.RawURLEncoding.EncodeToString([]byte(buildRawEmail(to, subject, body)))
+	var inReplyTo, references string
+	if threadID != "" {
+		inReplyTo, references = g.getThreadReplyHeaders(ctx, token, threadID)
+	}
+	raw := base64.RawURLEncoding.EncodeToString([]byte(buildRawEmail(to, subject, body, inReplyTo, references)))
 
 	var (
 		endpoint string
@@ -429,16 +437,77 @@ func formatGmailMessageResult(to, subject string, result map[string]any) map[str
 
 // --- email construction helpers ---
 
-func buildRawEmail(to, subject, body string) string {
+// getThreadReplyHeaders fetches Message-ID headers from an existing Gmail thread
+// so that outgoing replies can include proper In-Reply-To and References headers.
+// Returns empty strings on any failure — the caller should still send the email.
+func (g *GmailClient) getThreadReplyHeaders(ctx context.Context, token, threadID string) (inReplyTo, references string) {
+	path := fmt.Sprintf("/threads/%s?format=metadata&metadataHeaders=Message-ID", threadID)
+	var raw map[string]any
+	if err := g.api.RequestJSON(ctx, token, "GET", path, nil, &raw); err != nil {
+		return "", ""
+	}
+
+	rawMessages, _ := raw["messages"].([]any)
+	if len(rawMessages) == 0 {
+		return "", ""
+	}
+
+	var messageIDs []string
+	for _, rm := range rawMessages {
+		msg, ok := rm.(map[string]any)
+		if !ok {
+			continue
+		}
+		payload, ok := msg["payload"].(map[string]any)
+		if !ok {
+			continue
+		}
+		hdrs, ok := payload["headers"].([]any)
+		if !ok {
+			continue
+		}
+		for _, h := range hdrs {
+			hdr, ok := h.(map[string]any)
+			if !ok {
+				continue
+			}
+			name, _ := hdr["name"].(string)
+			value, _ := hdr["value"].(string)
+			if name == "Message-ID" && value != "" {
+				messageIDs = append(messageIDs, value)
+				break
+			}
+		}
+	}
+
+	if len(messageIDs) == 0 {
+		return "", ""
+	}
+	return messageIDs[len(messageIDs)-1], strings.Join(messageIDs, " ")
+}
+
+func buildRawEmail(to, subject, body, inReplyTo, references string) string {
 	to = sanitizeHeaderValue(to)
 	subject = sanitizeHeaderValue(subject)
+	inReplyTo = sanitizeHeaderValue(inReplyTo)
+	references = sanitizeHeaderValue(references)
 	subject = repairMojibake(subject)
 	subject = subjectNormalizer.Replace(subject)
 	encodedSubject := mime.QEncoding.Encode("utf-8", subject)
+
+	replyHeaders := ""
+	if inReplyTo != "" {
+		replyHeaders += fmt.Sprintf("In-Reply-To: %s\r\n", inReplyTo)
+	}
+	if references != "" {
+		replyHeaders += fmt.Sprintf("References: %s\r\n", references)
+	}
+
 	return fmt.Sprintf(
-		"MIME-Version: 1.0\r\nTo: %s\r\nSubject: %s\r\nContent-Type: text/plain; charset=UTF-8\r\n\r\n%s",
+		"MIME-Version: 1.0\r\nTo: %s\r\nSubject: %s\r\n%sContent-Type: text/plain; charset=UTF-8\r\n\r\n%s",
 		to,
 		encodedSubject,
+		replyHeaders,
 		body,
 	)
 }

pkg/tools/clients/gmail_test.go

@@ -3,17 +3,32 @@ package clients
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/beam-cloud/airstore/pkg/types"
 )
 
 func TestGmailCreateDraftSupportsThreadID(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+
+		// Handle the thread metadata fetch for reply headers.
+		if r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/threads/") {
+			_, _ = io.WriteString(w, `{
+				"id": "thread-123",
+				"messages": [
+					{"id":"msg-1","payload":{"headers":[{"name":"Message-ID","value":"<orig@example.com>"}]}}
+				]
+			}`)
+			return
+		}
+
 		if r.Method != http.MethodPost {
 			t.Fatalf("method = %s, want POST", r.Method)
 		}
@@ -32,11 +47,24 @@ func TestGmailCreateDraftSupportsThreadID(t *testing.T) {
 		if got, want := message["threadId"], "thread-123"; got != want {
 			t.Fatalf("message.threadId = %v, want %q", got, want)
 		}
-		if raw, _ := message["raw"].(string); raw == "" {
+		rawB64, _ := message["raw"].(string)
+		if rawB64 == "" {
 			t.Fatal("message.raw was empty")
 		}
 
-		w.Header().Set("Content-Type", "application/json")
+		// Verify the raw email contains In-Reply-To and References headers.
+		rawBytes, err := base64.RawURLEncoding.DecodeString(rawB64)
+		if err != nil {
+			t.Fatalf("decode raw: %v", err)
+		}
+		rawStr := string(rawBytes)
+		if !strings.Contains(rawStr, "In-Reply-To: <orig@example.com>") {
+			t.Fatalf("expected In-Reply-To header in raw email, got:\n%s", rawStr)
+		}
+		if !strings.Contains(rawStr, "References: <orig@example.com>") {
+			t.Fatalf("expected References header in raw email, got:\n%s", rawStr)
+		}
+
 		_, _ = io.WriteString(w, `{"id":"draft-1","message":{"id":"msg-1","threadId":"thread-123","labelIds":["DRAFT"]}}`)
 	}))
 	defer server.Close()
@@ -67,3 +95,79 @@ func TestGmailCreateDraftSupportsThreadID(t *testing.T) {
 		t.Fatalf("output draft_id = %v, want %q", got, want)
 	}
 }
+
+func TestBuildRawEmailReplyHeaders(t *testing.T) {
+	// With reply headers
+	raw := buildRawEmail("test@example.com", "Re: Hello", "body text",
+		"<abc@mail.gmail.com>", "<abc@mail.gmail.com> <def@mail.gmail.com>")
+
+	if !strings.Contains(raw, "In-Reply-To: <abc@mail.gmail.com>\r\n") {
+		t.Fatalf("expected In-Reply-To header, got:\n%s", raw)
+	}
+	if !strings.Contains(raw, "References: <abc@mail.gmail.com> <def@mail.gmail.com>\r\n") {
+		t.Fatalf("expected References header, got:\n%s", raw)
+	}
+
+	// Without reply headers — should be identical to old behavior
+	rawNoReply := buildRawEmail("test@example.com", "Hello", "body text", "", "")
+
+	if strings.Contains(rawNoReply, "In-Reply-To") {
+		t.Fatalf("unexpected In-Reply-To header without reply params:\n%s", rawNoReply)
+	}
+	if strings.Contains(rawNoReply, "References") {
+		t.Fatalf("unexpected References header without reply params:\n%s", rawNoReply)
+	}
+}
+
+func TestSendEmailThreadFetchFailureFallback(t *testing.T) {
+	var postReceived bool
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+
+		// Thread metadata fetch fails with 404.
+		if r.Method == http.MethodGet && strings.HasPrefix(r.URL.Path, "/threads/") {
+			w.WriteHeader(http.StatusNotFound)
+			_, _ = io.WriteString(w, `{"error":{"code":404,"message":"Not Found"}}`)
+			return
+		}
+
+		// The email should still be sent successfully.
+		if r.Method == http.MethodPost && r.URL.Path == "/messages/send" {
+			postReceived = true
+
+			var payload map[string]any
+			if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
+				t.Fatalf("decode payload: %v", err)
+			}
+
+			// Verify threadId is still set in payload despite fetch failure.
+			if got, want := payload["threadId"], "thread-404"; got != want {
+				t.Fatalf("payload.threadId = %v, want %q", got, want)
+			}
+
+			_, _ = io.WriteString(w, `{"id":"msg-1","threadId":"thread-404","labelIds":["SENT"]}`)
+			return
+		}
+
+		t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path)
+	}))
+	defer server.Close()
+
+	client := NewGmailClient()
+	client.api.baseURL = server.URL
+	client.api.httpClient = server.Client()
+
+	var stdout bytes.Buffer
+	err := client.Execute(context.Background(), gmailCmdSendEmail, map[string]any{
+		"to":        "test@example.com",
+		"subject":   "Re: Thread test",
+		"body":      "Reply body.",
+		"thread_id": "thread-404",
+	}, &types.IntegrationCredentials{AccessToken: "token"}, &stdout, &bytes.Buffer{})
+	if err != nil {
+		t.Fatalf("Execute returned error: %v", err)
+	}
+	if !postReceived {
+		t.Fatal("expected POST to /messages/send but it was never called")
+	}
+}

pkg/tools/clients/write_clients_test.go

@@ -182,7 +182,7 @@ func TestGmailSendEmailRequiresArguments(t *testing.T) {
 }
 
 func TestBuildRawEmailSanitizesHeaderInjection(t *testing.T) {
-	raw := buildRawEmail("victim@example.com\r\nBcc: attacker@example.com", "hello\r\nX-Test: injected", "body")
+	raw := buildRawEmail("victim@example.com\r\nBcc: attacker@example.com", "hello\r\nX-Test: injected", "body", "", "")
 	if bytes.Contains([]byte(raw), []byte("\r\nBcc:")) {
 		t.Fatalf("expected CRLF header injection to be sanitized, got %q", raw)
 	}