diff --git a/pkg/tools/clients/gmail.go b/pkg/tools/clients/gmail.go index 3bd7b35..0438a8a 100644 --- a/pkg/tools/clients/gmail.go +++ b/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, ) } diff --git a/pkg/tools/clients/gmail_test.go b/pkg/tools/clients/gmail_test.go index 8fc2503..fef6c49 100644 --- a/pkg/tools/clients/gmail_test.go +++ b/pkg/tools/clients/gmail_test.go @@ -3,10 +3,12 @@ package clients import ( "bytes" "context" + "encoding/base64" "encoding/json" "io" "net/http" "net/http/httptest" + "strings" "testing" "github.com/beam-cloud/airstore/pkg/types" @@ -14,6 +16,19 @@ import ( 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":""}]}} + ] + }`) + 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: ") { + t.Fatalf("expected In-Reply-To header in raw email, got:\n%s", rawStr) + } + if !strings.Contains(rawStr, "References: ") { + 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", + "", " ") + + if !strings.Contains(raw, "In-Reply-To: \r\n") { + t.Fatalf("expected In-Reply-To header, got:\n%s", raw) + } + if !strings.Contains(raw, "References: \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") + } +} diff --git a/pkg/tools/clients/write_clients_test.go b/pkg/tools/clients/write_clients_test.go index ee38614..7ea3806 100644 --- a/pkg/tools/clients/write_clients_test.go +++ b/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) }