@vellumai/assistant 0.4.46 → 0.4.49

package/src/messaging/providers/gmail/client.ts

@@ -1,3 +1,7 @@
+import type {
+  OAuthConnection,
+  OAuthConnectionResponse,
+} from "../../../oauth/connection.js";
 import type {
   GmailAttachment,
   GmailDraft,
@@ -15,7 +19,6 @@ import type {
   GmailVacationSettings,
 } from "./types.js";
 
-const GMAIL_API_BASE = "https://gmail.googleapis.com/gmail/v1/users/me";
 const GMAIL_BATCH_URL = "https://www.googleapis.com/batch/gmail/v1";
 
 /** Max sub-requests per batch HTTP call (Gmail API limit) */
@@ -36,6 +39,7 @@ export class GmailApiError extends Error {
 
 const MAX_RETRIES = 3;
 const INITIAL_BACKOFF_MS = 1000;
+/** Timeout for batch API calls that bypass OAuthConnection.request() (which has its own 30s timeout). */
 const REQUEST_TIMEOUT_MS = 30_000;
 
 function isRetryable(status: number): boolean {
@@ -54,53 +58,115 @@ interface GmailRequestOptions extends RequestInit {
   retryable?: boolean;
 }
 
+/**
+ * Extract non-Authorization headers from request options for use with OAuthConnection.
+ */
+function extractNonAuthHeaders(
+  options?: GmailRequestOptions,
+): Record<string, string> | undefined {
+  if (!options?.headers) return undefined;
+  const raw = options.headers;
+  const result: Record<string, string> = {};
+  if (raw instanceof Headers) {
+    raw.forEach((v, k) => {
+      if (k.toLowerCase() !== "authorization") result[k] = v;
+    });
+  } else if (Array.isArray(raw)) {
+    for (const [k, v] of raw) {
+      if (k.toLowerCase() !== "authorization") result[k] = v;
+    }
+  } else {
+    for (const [k, v] of Object.entries(raw)) {
+      if (k.toLowerCase() !== "authorization" && v !== undefined) result[k] = v;
+    }
+  }
+  return Object.keys(result).length > 0 ? result : undefined;
+}
+
+/**
+ * Convert URLSearchParams to a query record, collapsing multi-valued keys into arrays.
+ */
+function paramsToQuery(
+  params: URLSearchParams,
+): Record<string, string | string[]> {
+  const result: Record<string, string | string[]> = {};
+  for (const key of new Set(params.keys())) {
+    const values = params.getAll(key);
+    result[key] = values.length === 1 ? values[0] : values;
+  }
+  return result;
+}
+
+/**
+ * Extract the JSON body from request options for use with OAuthConnection.
+ */
+function extractBody(options?: GmailRequestOptions): unknown | undefined {
+  if (!options?.body) return undefined;
+  if (typeof options.body === "string") {
+    try {
+      return JSON.parse(options.body);
+    } catch {
+      return options.body;
+    }
+  }
+  return options.body;
+}
+
 async function request<T>(
-  token: string,
+  connection: OAuthConnection,
   path: string,
   options?: GmailRequestOptions,
+  query?: Record<string, string | string[]>,
 ): Promise<T> {
-  const url = `${GMAIL_API_BASE}${path}`;
   const canRetry = options?.retryable ?? isIdempotent(options);
+  const method = (options?.method ?? "GET").toUpperCase();
 
   for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
-    const resp = await fetch(url, {
-      ...options,
-      signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS),
-      headers: {
-        Authorization: `Bearer ${token}`,
-        "Content-Type": "application/json",
-        ...options?.headers,
-      },
-    });
+    let resp: OAuthConnectionResponse;
+    try {
+      resp = await connection.request({
+        method,
+        path,
+        query,
+        headers: {
+          "Content-Type": "application/json",
+          ...extractNonAuthHeaders(options),
+        },
+        body: extractBody(options),
+      });
+    } catch (err) {
+      // Network-level errors from connection.request() are not retryable
+      throw err;
+    }
 
-    if (!resp.ok) {
+    if (resp.status < 200 || resp.status >= 300) {
       if (canRetry && isRetryable(resp.status) && attempt < MAX_RETRIES) {
-        const retryAfter = resp.headers.get("retry-after");
+        const retryAfter =
+          resp.headers["retry-after"] ?? resp.headers["Retry-After"];
         const delayMs = retryAfter
           ? parseInt(retryAfter, 10) * 1000
           : INITIAL_BACKOFF_MS * Math.pow(2, attempt);
         await new Promise((resolve) => setTimeout(resolve, delayMs));
         continue;
       }
-      const body = await resp.text().catch(() => "");
+      const bodyStr =
+        typeof resp.body === "string"
+          ? resp.body
+          : JSON.stringify(resp.body ?? "");
       throw new GmailApiError(
         resp.status,
-        resp.statusText,
-        `Gmail API ${resp.status}: ${body}`,
+        "",
+        `Gmail API ${resp.status}: ${bodyStr}`,
       );
     }
 
-    // Some endpoints (e.g. batchModify) return empty success responses
-    const contentLength = resp.headers.get("content-length");
-    if (resp.status === 204 || contentLength === "0") {
+    // Success — body is already parsed by connection.request()
+    if (resp.status === 204 || resp.body === undefined) {
       return undefined as T;
     }
-    const text = await resp.text();
-    if (!text) return undefined as T;
-    return JSON.parse(text) as T;
+    return resp.body as T;
   }
 
-  // Unreachable — the loop always returns or throws — but TypeScript needs this
   throw new Error(
     "Unreachable: retry loop exited without returning or throwing",
   );
@@ -108,7 +174,7 @@ async function request<T>(
 
 /** List messages matching a query. */
 export async function listMessages(
-  token: string,
+  connection: OAuthConnection,
   query?: string,
   maxResults = 20,
   pageToken?: string,
@@ -121,12 +187,17 @@ export async function listMessages(
   if (labelIds) {
     for (const id of labelIds) params.append("labelIds", id);
   }
-  return request<GmailMessageListResponse>(token, `/messages?${params}`);
+  return request<GmailMessageListResponse>(
+    connection,
+    "/messages",
+    undefined,
+    paramsToQuery(params),
+  );
 }
 
 /** Get a single message by ID. */
 export async function getMessage(
-  token: string,
+  connection: OAuthConnection,
   messageId: string,
   format: GmailMessageFormat = "full",
   metadataHeaders?: string[],
@@ -137,7 +208,12 @@ export async function getMessage(
     for (const h of metadataHeaders) params.append("metadataHeaders", h);
   }
   if (fields) params.set("fields", fields);
-  return request<GmailMessage>(token, `/messages/${messageId}?${params}`);
+  return request<GmailMessage>(
+    connection,
+    `/messages/${messageId}`,
+    undefined,
+    paramsToQuery(params),
+  );
 }
 
 /**
@@ -175,7 +251,7 @@ function parseSubResponse(
  * Returns successfully parsed messages and a list of IDs that failed (for individual retry).
  */
 async function executeBatchCall(
-  token: string,
+  connection: OAuthConnection,
   messageIds: string[],
   format: GmailMessageFormat,
   metadataHeaders: string[] | undefined,
@@ -203,79 +279,126 @@ async function executeBatchCall(
   );
   const body = parts.join("") + `--${boundary}--\r\n`;
 
-  for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
-    const resp = await fetch(GMAIL_BATCH_URL, {
-      method: "POST",
-      signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS * 2),
-      headers: {
-        Authorization: `Bearer ${token}`,
-        "Content-Type": `multipart/mixed; boundary=${boundary}`,
-      },
-      body,
-    });
-
-    if (!resp.ok) {
-      if (isRetryable(resp.status) && attempt < MAX_RETRIES) {
-        const retryAfter = resp.headers.get("retry-after");
-        const delayMs = retryAfter
-          ? parseInt(retryAfter, 10) * 1000
-          : INITIAL_BACKOFF_MS * Math.pow(2, attempt);
-        await new Promise((r) => setTimeout(r, delayMs));
-        continue;
+  const doBatchFetch = async (token: string) => {
+    for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
+      const resp = await fetch(GMAIL_BATCH_URL, {
+        method: "POST",
+        signal: AbortSignal.timeout(REQUEST_TIMEOUT_MS * 2),
+        headers: {
+          Authorization: `Bearer ${token}`,
+          "Content-Type": `multipart/mixed; boundary=${boundary}`,
+        },
+        body,
+      });
+
+      if (!resp.ok) {
+        if (isRetryable(resp.status) && attempt < MAX_RETRIES) {
+          const retryAfter = resp.headers.get("retry-after");
+          const delayMs = retryAfter
+            ? parseInt(retryAfter, 10) * 1000
+            : INITIAL_BACKOFF_MS * Math.pow(2, attempt);
+          await new Promise((r) => setTimeout(r, delayMs));
+          continue;
+        }
+        const errBody = await resp.text().catch(() => "");
+        throw new GmailApiError(
+          resp.status,
+          resp.statusText,
+          `Gmail batch API ${resp.status}: ${errBody}`,
+        );
       }
-      const errBody = await resp.text().catch(() => "");
-      throw new GmailApiError(
-        resp.status,
-        resp.statusText,
-        `Gmail batch API ${resp.status}: ${errBody}`,
-      );
-    }
-
-    const contentType = resp.headers.get("content-type") ?? "";
-    const responseText = await resp.text();
-
-    const boundaryMatch = contentType.match(/boundary=(?:"([^"]+)"|([^\s;]+))/);
-    const respBoundary = boundaryMatch?.[1] ?? boundaryMatch?.[2];
-    if (!respBoundary)
-      throw new Error("Missing boundary in Gmail batch response");
-
-    const respParts = responseText.split(`--${respBoundary}`);
-    const messages: Array<{ index: number; msg: GmailMessage }> = [];
-    const failedIds: Array<{ index: number; id: string }> = [];
 
-    for (const rp of respParts) {
-      const parsed = parseSubResponse(rp);
-      if (!parsed) continue;
+      const contentType = resp.headers.get("content-type") ?? "";
+      const responseText = await resp.text();
 
-      if (parsed.status >= 200 && parsed.status < 300 && parsed.json) {
-        try {
-          messages.push({
+      const boundaryMatch = contentType.match(
+        /boundary=(?:"([^"]+)"|([^\s;]+))/,
+      );
+      const respBoundary = boundaryMatch?.[1] ?? boundaryMatch?.[2];
+      if (!respBoundary)
+        throw new Error("Missing boundary in Gmail batch response");
+
+      const respParts = responseText.split(`--${respBoundary}`);
+      const messages: Array<{ index: number; msg: GmailMessage }> = [];
+      const failedIds: Array<{ index: number; id: string }> = [];
+
+      for (const rp of respParts) {
+        const parsed = parseSubResponse(rp);
+        if (!parsed) continue;
+
+        if (parsed.status >= 200 && parsed.status < 300 && parsed.json) {
+          try {
+            messages.push({
+              index: parsed.index,
+              msg: JSON.parse(parsed.json) as GmailMessage,
+            });
+          } catch {
+            failedIds.push({
+              index: parsed.index,
+              id: messageIds[parsed.index],
+            });
+          }
+        } else {
+          failedIds.push({
             index: parsed.index,
-            msg: JSON.parse(parsed.json) as GmailMessage,
+            id: messageIds[parsed.index],
           });
-        } catch {
-          failedIds.push({ index: parsed.index, id: messageIds[parsed.index] });
         }
-      } else {
-        failedIds.push({ index: parsed.index, id: messageIds[parsed.index] });
       }
+
+      return { messages, failedIds };
     }
 
-    return { messages, failedIds };
-  }
+    throw new Error(
+      "Unreachable: batch retry loop exited without returning or throwing",
+    );
+  };
 
-  throw new Error(
-    "Unreachable: batch retry loop exited without returning or throwing",
-  );
+  // Use withToken to get raw token for batch endpoint
+  return connection.withToken(doBatchFetch);
+}
+
+/** Max concurrent individual getMessage requests (matches batch concurrency) */
+const INDIVIDUAL_CONCURRENCY = BATCH_CONCURRENCY;
+
+/**
+ * Fetch all messages individually using getMessage (no batch endpoint).
+ * Used as a fallback when the batch API is unavailable (e.g. platform connections
+ * that cannot expose raw tokens for the multipart batch endpoint).
+ *
+ * Processes messages in waves of INDIVIDUAL_CONCURRENCY to avoid unbounded
+ * parallelism that would trigger 429s on high-volume paths like senderDigest.
+ */
+async function fetchMessagesIndividually(
+  connection: OAuthConnection,
+  messageIds: string[],
+  format: GmailMessageFormat,
+  metadataHeaders?: string[],
+  fields?: string,
+): Promise<GmailMessage[]> {
+  const results: GmailMessage[] = [];
+  for (let i = 0; i < messageIds.length; i += INDIVIDUAL_CONCURRENCY) {
+    const wave = messageIds.slice(i, i + INDIVIDUAL_CONCURRENCY);
+    const waveResults = await Promise.all(
+      wave.map((id) =>
+        getMessage(connection, id, format, metadataHeaders, fields),
+      ),
+    );
+    results.push(...waveResults);
+  }
+  return results;
 }
 
 /**
  * Get multiple messages using Gmail's batch HTTP endpoint.
  * Packs up to 100 sub-requests per HTTP call and runs up to BATCH_CONCURRENCY calls in parallel.
  * Falls back to individual getMessage for any sub-requests that fail within a batch.
+ *
+ * For connections that do not support raw token access (e.g. platform-managed connections),
+ * falls back to fetching each message individually via connection.request().
  */
 export async function batchGetMessages(
-  token: string,
+  connection: OAuthConnection,
   messageIds: string[],
   format: GmailMessageFormat = "full",
   metadataHeaders?: string[],
@@ -283,13 +406,39 @@ export async function batchGetMessages(
 ): Promise<GmailMessage[]> {
   if (messageIds.length === 0) return [];
 
-  // Single message — just use getMessage directly
+  // Single message -- just use getMessage directly
   if (messageIds.length === 1) {
     return [
-      await getMessage(token, messageIds[0], format, metadataHeaders, fields),
+      await getMessage(
+        connection,
+        messageIds[0],
+        format,
+        metadataHeaders,
+        fields,
+      ),
     ];
   }
 
+  // Try batch API first; fall back to individual fetches if withToken is unavailable
+  // (e.g. platform-managed connections where raw tokens cannot be exposed).
+  let useBatch = true;
+  try {
+    // Probe withToken availability with a no-op call
+    await connection.withToken(async (token) => token);
+  } catch {
+    useBatch = false;
+  }
+
+  if (!useBatch) {
+    return fetchMessagesIndividually(
+      connection,
+      messageIds,
+      format,
+      metadataHeaders,
+      fields,
+    );
+  }
+
   const results = new Array<GmailMessage | undefined>(messageIds.length).fill(
     undefined,
   );
@@ -307,7 +456,13 @@ export async function batchGetMessages(
     const wave = chunks.slice(i, i + BATCH_CONCURRENCY);
     const waveResults = await Promise.all(
       wave.map((chunk) =>
-        executeBatchCall(token, chunk.ids, format, metadataHeaders, fields),
+        executeBatchCall(
+          connection,
+          chunk.ids,
+          format,
+          metadataHeaders,
+          fields,
+        ),
       ),
     );
 
@@ -324,7 +479,7 @@ export async function batchGetMessages(
       if (failedIds.length > 0) {
         const retried = await Promise.all(
           failedIds.map(({ id }) =>
-            getMessage(token, id, format, metadataHeaders, fields),
+            getMessage(connection, id, format, metadataHeaders, fields),
           ),
         );
         for (let r = 0; r < failedIds.length; r++) {
@@ -339,11 +494,11 @@ export async function batchGetMessages(
 
 /** Modify labels on a single message. */
 export async function modifyMessage(
-  token: string,
+  connection: OAuthConnection,
   messageId: string,
   modifications: GmailModifyRequest,
 ): Promise<GmailMessage> {
-  return request<GmailMessage>(token, `/messages/${messageId}/modify`, {
+  return request<GmailMessage>(connection, `/messages/${messageId}/modify`, {
     method: "POST",
     body: JSON.stringify(modifications),
     retryable: true,
@@ -352,11 +507,11 @@ export async function modifyMessage(
 
 /** Batch modify labels on multiple messages. */
 export async function batchModifyMessages(
-  token: string,
+  connection: OAuthConnection,
   messageIds: string[],
   modifications: GmailModifyRequest,
 ): Promise<void> {
-  await request<void>(token, "/messages/batchModify", {
+  await request<void>(connection, "/messages/batchModify", {
     method: "POST",
     body: JSON.stringify({ ids: messageIds, ...modifications }),
     retryable: true,
@@ -365,24 +520,26 @@ export async function batchModifyMessages(
 
 /** Move a message to trash. */
 export async function trashMessage(
-  token: string,
+  connection: OAuthConnection,
   messageId: string,
 ): Promise<GmailMessage> {
-  return request<GmailMessage>(token, `/messages/${messageId}/trash`, {
+  return request<GmailMessage>(connection, `/messages/${messageId}/trash`, {
     method: "POST",
     retryable: true,
   });
 }
 
 /** List all labels. */
-export async function listLabels(token: string): Promise<GmailLabel[]> {
-  const resp = await request<GmailLabelsListResponse>(token, "/labels");
+export async function listLabels(
+  connection: OAuthConnection,
+): Promise<GmailLabel[]> {
+  const resp = await request<GmailLabelsListResponse>(connection, "/labels");
   return resp.labels ?? [];
 }
 
 /** Create a draft. */
 export async function createDraft(
-  token: string,
+  connection: OAuthConnection,
   to: string,
   subject: string,
   body: string,
@@ -409,7 +566,7 @@ export async function createDraft(
     .replace(/=+$/, "");
   const message: Record<string, unknown> = { raw };
   if (threadId) message.threadId = threadId;
-  return request<GmailDraft>(token, "/drafts", {
+  return request<GmailDraft>(connection, "/drafts", {
     method: "POST",
     body: JSON.stringify({ message }),
   });
@@ -417,13 +574,13 @@ export async function createDraft(
 
 /** Create a draft from a pre-built base64url MIME payload. */
 export async function createDraftRaw(
-  token: string,
+  connection: OAuthConnection,
   raw: string,
   threadId?: string,
 ): Promise<GmailDraft> {
   const message: Record<string, unknown> = { raw };
   if (threadId) message.threadId = threadId;
-  return request<GmailDraft>(token, "/drafts", {
+  return request<GmailDraft>(connection, "/drafts", {
     method: "POST",
     body: JSON.stringify({ message }),
   });
@@ -431,10 +588,10 @@ export async function createDraftRaw(
 
 /** Send an existing draft by ID. */
 export async function sendDraft(
-  token: string,
+  connection: OAuthConnection,
   draftId: string,
 ): Promise<GmailMessage> {
-  return request<GmailMessage>(token, "/drafts/send", {
+  return request<GmailMessage>(connection, "/drafts/send", {
     method: "POST",
     body: JSON.stringify({ id: draftId }),
   });
@@ -442,7 +599,7 @@ export async function sendDraft(
 
 /** Send an email. */
 export async function sendMessage(
-  token: string,
+  connection: OAuthConnection,
   to: string,
   subject: string,
   body: string,
@@ -465,38 +622,40 @@ export async function sendMessage(
     .replace(/=+$/, "");
   const payload: Record<string, unknown> = { raw };
   if (threadId) payload.threadId = threadId;
-  return request<GmailMessage>(token, "/messages/send", {
+  return request<GmailMessage>(connection, "/messages/send", {
     method: "POST",
     body: JSON.stringify(payload),
   });
 }
 
 /** Get the authenticated user's profile (email address). */
-export async function getProfile(token: string): Promise<GmailProfile> {
-  return request<GmailProfile>(token, "/profile");
+export async function getProfile(
+  connection: OAuthConnection,
+): Promise<GmailProfile> {
+  return request<GmailProfile>(connection, "/profile");
 }
 
 /** Get attachment data for a message. */
 export async function getAttachment(
-  token: string,
+  connection: OAuthConnection,
   messageId: string,
   attachmentId: string,
 ): Promise<GmailAttachment> {
   return request<GmailAttachment>(
-    token,
+    connection,
     `/messages/${messageId}/attachments/${attachmentId}`,
   );
 }
 
 /** Send an email with a pre-built raw MIME payload (for multipart/attachments). */
 export async function sendMessageRaw(
-  token: string,
+  connection: OAuthConnection,
   raw: string,
   threadId?: string,
 ): Promise<GmailMessage> {
   const payload: Record<string, unknown> = { raw };
   if (threadId) payload.threadId = threadId;
-  return request<GmailMessage>(token, "/messages/send", {
+  return request<GmailMessage>(connection, "/messages/send", {
     method: "POST",
     body: JSON.stringify(payload),
   });
@@ -504,23 +663,25 @@ export async function sendMessageRaw(
 
 /** Create a user label. */
 export async function createLabel(
-  token: string,
+  connection: OAuthConnection,
   name: string,
   opts?: {
     messageListVisibility?: "show" | "hide";
     labelListVisibility?: "labelShow" | "labelShowIfUnread" | "labelHide";
   },
 ): Promise<GmailLabel> {
-  return request<GmailLabel>(token, "/labels", {
+  return request<GmailLabel>(connection, "/labels", {
     method: "POST",
     body: JSON.stringify({ name, ...opts }),
   });
 }
 
 /** List all Gmail filters. */
-export async function listFilters(token: string): Promise<GmailFilter[]> {
+export async function listFilters(
+  connection: OAuthConnection,
+): Promise<GmailFilter[]> {
   const resp = await request<GmailFiltersListResponse>(
-    token,
+    connection,
     "/settings/filters",
   );
   return resp.filter ?? [];
@@ -528,11 +689,11 @@ export async function listFilters(token: string): Promise<GmailFilter[]> {
 
 /** Create a Gmail filter. */
 export async function createFilter(
-  token: string,
+  connection: OAuthConnection,
   criteria: GmailFilterCriteria,
   action: GmailFilterAction,
 ): Promise<GmailFilter> {
-  return request<GmailFilter>(token, "/settings/filters", {
+  return request<GmailFilter>(connection, "/settings/filters", {
     method: "POST",
     body: JSON.stringify({ criteria, action }),
   });
@@ -540,27 +701,27 @@ export async function createFilter(
 
 /** Delete a Gmail filter. */
 export async function deleteFilter(
-  token: string,
+  connection: OAuthConnection,
   filterId: string,
 ): Promise<void> {
-  await request<void>(token, `/settings/filters/${filterId}`, {
+  await request<void>(connection, `/settings/filters/${filterId}`, {
     method: "DELETE",
   });
 }
 
 /** Get vacation auto-reply settings. */
 export async function getVacation(
-  token: string,
+  connection: OAuthConnection,
 ): Promise<GmailVacationSettings> {
-  return request<GmailVacationSettings>(token, "/settings/vacation");
+  return request<GmailVacationSettings>(connection, "/settings/vacation");
 }
 
 /** Update vacation auto-reply settings. */
 export async function updateVacation(
-  token: string,
+  connection: OAuthConnection,
   settings: GmailVacationSettings,
 ): Promise<GmailVacationSettings> {
-  return request<GmailVacationSettings>(token, "/settings/vacation", {
+  return request<GmailVacationSettings>(connection, "/settings/vacation", {
     method: "PUT",
     body: JSON.stringify(settings),
   });