fix(custom-widgets): address v2 merge review (#6446)

packages/custom-widgets/src/server/errors.ts

@@ -1,6 +1,7 @@
 export type CustomWidgetDomainErrorCode =
   | "BAD_GATEWAY"
   | "BAD_REQUEST"
+  | "CONFLICT"
   | "FORBIDDEN"
   | "INTERNAL_SERVER_ERROR"
   | "NOT_FOUND"

packages/custom-widgets/src/server/index.ts

@@ -6,6 +6,7 @@ export {
   classifyAddress,
   executeCustomWidgetRequest,
   invalidateCustomWidgetResponseCache,
+  MAX_REQUEST_DURATION_MS,
   MAX_REQUEST_BODY_BYTES,
   MAX_RESPONSE_BODY_BYTES,
   MAX_RESPONSE_JSON_DEPTH,
@@ -15,7 +16,7 @@ export {
   validateCustomWidgetUrl,
 } from "./request-executor";
 export type { CustomWidgetAuthConfig, CustomWidgetHttpRequest, CustomWidgetHttpResponse } from "./request-executor";
-export { CustomWidgetRequestLimiter } from "./request-limits";
+export { CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS, CustomWidgetRequestLimiter } from "./request-limits";
 export type { RequestLimitInput, RequestLimiterOptions, RequestLimitStore } from "./request-limits";
 export { CustomWidgetPreviewSessionService } from "./preview-sessions";
 export type {

packages/custom-widgets/src/server/preview-sessions.ts

@@ -12,6 +12,7 @@ import { CustomWidgetDomainError } from "./errors";
 
 const SESSION_TTL_MS = 10 * 60_000;
 const MAX_JOURNAL_ENTRIES = 50;
+const MAX_UPDATE_ATTEMPTS = 8;
 
 const encryptedSecretSchema = z.object({
   sourceId: z.string(),
@@ -22,6 +23,7 @@ const encryptedSecretSchema = z.object({
 const sessionSchema = z.object({
   id: z.string(),
   userId: z.string(),
+  revision: z.number().int().nonnegative(),
   expiresAt: z.number(),
   sources: customWidgetSourcesSchema,
   secrets: z.array(encryptedSecretSchema),
@@ -62,6 +64,7 @@ export interface CreatePreviewSessionInput {
 
 export interface PreviewSessionStore {
   saveSession(id: string, value: unknown, ttlMs: number): Promise<void>;
+  compareAndSwapSession(id: string, expectedRevision: number, value: unknown, ttlMs: number): Promise<boolean>;
   getSession(id: string): Promise<unknown>;
   deleteSession(id: string): Promise<void>;
   appendJournal(id: string, value: unknown, maxEntries: number, ttlMs: number): Promise<void>;
@@ -96,6 +99,7 @@ export class CustomWidgetPreviewSessionService {
     const session: CustomWidgetPreviewSession = {
       id: this.options.createId(),
       userId: input.userId,
+      revision: 0,
       expiresAt: this.now() + SESSION_TTL_MS,
       sources: input.sources,
       secrets: input.secrets.map((secret) => ({ ...secret, value: this.options.encrypt(secret.value) })),
@@ -114,17 +118,19 @@ export class CustomWidgetPreviewSessionService {
   public async get(id: string, userId: string): Promise<CustomWidgetPreviewSession> {
     const candidate = this.options.store ? await this.options.store.getSession(id) : this.sessions.get(id);
     const parsed = sessionSchema.safeParse(parseStoredValue(candidate));
-    if (!parsed.success || parsed.data.userId !== userId || parsed.data.expiresAt <= this.now()) {
+    if (!parsed.success || parsed.data.expiresAt <= this.now()) {
       if (this.options.store) await this.options.store.deleteSession(id);
       else this.sessions.delete(id);
       throw new CustomWidgetDomainError({ code: "NOT_FOUND", message: "Preview session expired or was not found" });
     }
+    if (parsed.data.userId !== userId) {
+      throw new CustomWidgetDomainError({ code: "NOT_FOUND", message: "Preview session expired or was not found" });
+    }
     return parsed.data;
   }
 
   public async setLiveActions(id: string, userId: string, enabled: boolean) {
-    const session = { ...(await this.get(id, userId)), liveActions: enabled };
-    await this.save(session);
+    const session = await this.update(id, userId, (current) => ({ ...current, liveActions: enabled }));
     return { id, expiresAt: session.expiresAt, liveActions: enabled };
   }
 
@@ -160,23 +166,23 @@ export class CustomWidgetPreviewSessionService {
     userId: string,
     secrets: Array<{ sourceId: string; kind: (typeof customWidgetSecretKinds)[number]; value: string }>,
   ) {
-    const session = await this.get(id, userId);
-    const sourceIds = new Set(Object.keys(session.sources));
-    if (secrets.some((secret) => !sourceIds.has(secret.sourceId))) {
-      throw new CustomWidgetDomainError({
-        code: "BAD_REQUEST",
-        message: "A preview secret references an unknown source",
-      });
-    }
-    const replacements = new Set(secrets.map((secret) => `${secret.sourceId}:${secret.kind}`));
-    const next = {
-      ...session,
-      secrets: [
-        ...session.secrets.filter((secret) => !replacements.has(`${secret.sourceId}:${secret.kind}`)),
-        ...secrets.map((secret) => ({ ...secret, value: this.options.encrypt(secret.value) })),
-      ],
-    };
-    await this.save(next);
+    const next = await this.update(id, userId, (session) => {
+      const sourceIds = new Set(Object.keys(session.sources));
+      if (secrets.some((secret) => !sourceIds.has(secret.sourceId))) {
+        throw new CustomWidgetDomainError({
+          code: "BAD_REQUEST",
+          message: "A preview secret references an unknown source",
+        });
+      }
+      const replacements = new Set(secrets.map((secret) => `${secret.sourceId}:${secret.kind}`));
+      return {
+        ...session,
+        secrets: [
+          ...session.secrets.filter((secret) => !replacements.has(`${secret.sourceId}:${secret.kind}`)),
+          ...secrets.map((secret) => ({ ...secret, value: this.options.encrypt(secret.value) })),
+        ],
+      };
+    });
     return { id, expiresAt: next.expiresAt };
   }
 
@@ -187,33 +193,53 @@ export class CustomWidgetPreviewSessionService {
     source: CustomWidgetSource,
     secrets: Array<{ sourceId: string; kind: (typeof customWidgetSecretKinds)[number]; value: string }>,
   ) {
-    const session = await this.get(id, userId);
-    const currentSource = session.sources[sourceId];
-    if (!currentSource) {
-      throw new CustomWidgetDomainError({ code: "NOT_FOUND", message: "Preview source was not found" });
-    }
-    if (!hasSameCustomWidgetSourceAuthentication(currentSource, source)) {
-      throw new CustomWidgetDomainError({
-        code: "BAD_REQUEST",
-        message: "Preview source authentication changed; create a new configuration request",
-      });
-    }
-    const replacements = new Set(secrets.map((secret) => `${secret.sourceId}:${secret.kind}`));
-    const next = {
-      ...session,
-      sources: {
-        ...session.sources,
-        [sourceId]: { ...currentSource, baseUrl: source.baseUrl, networkScope: source.networkScope },
-      },
-      secrets: [
-        ...session.secrets.filter((secret) => !replacements.has(`${secret.sourceId}:${secret.kind}`)),
-        ...secrets.map((secret) => ({ ...secret, value: this.options.encrypt(secret.value) })),
-      ],
-    };
-    await this.save(next);
+    const next = await this.update(id, userId, (session) => {
+      const currentSource = session.sources[sourceId];
+      if (!currentSource) {
+        throw new CustomWidgetDomainError({ code: "NOT_FOUND", message: "Preview source was not found" });
+      }
+      if (!hasSameCustomWidgetSourceAuthentication(currentSource, source)) {
+        throw new CustomWidgetDomainError({
+          code: "BAD_REQUEST",
+          message: "Preview source authentication changed; create a new configuration request",
+        });
+      }
+      const replacements = new Set(secrets.map((secret) => `${secret.sourceId}:${secret.kind}`));
+      return {
+        ...session,
+        sources: {
+          ...session.sources,
+          [sourceId]: { ...currentSource, baseUrl: source.baseUrl, networkScope: source.networkScope },
+        },
+        secrets: [
+          ...session.secrets.filter((secret) => !replacements.has(`${secret.sourceId}:${secret.kind}`)),
+          ...secrets.map((secret) => ({ ...secret, value: this.options.encrypt(secret.value) })),
+        ],
+      };
+    });
     return { id, expiresAt: next.expiresAt };
   }
 
+  private async update(
+    id: string,
+    userId: string,
+    mutate: (session: CustomWidgetPreviewSession) => CustomWidgetPreviewSession,
+  ) {
+    for (let attempt = 0; attempt < MAX_UPDATE_ATTEMPTS; attempt += 1) {
+      const current = await this.get(id, userId);
+      const next = { ...mutate(current), revision: current.revision + 1 };
+      const ttlMs = Math.max(0, next.expiresAt - this.now());
+      if (!ttlMs) throw new CustomWidgetDomainError({ code: "NOT_FOUND", message: "Preview session expired" });
+      if (this.options.store) {
+        if (await this.options.store.compareAndSwapSession(id, current.revision, next, ttlMs)) return next;
+      } else if (this.sessions.get(id)?.revision === current.revision) {
+        this.sessions.set(id, next);
+        return next;
+      }
+    }
+    throw new CustomWidgetDomainError({ code: "CONFLICT", message: "Preview session changed too many times; retry" });
+  }
+
   private async save(session: CustomWidgetPreviewSession): Promise<void> {
     const ttlMs = Math.max(0, session.expiresAt - this.now());
     if (!ttlMs) throw new CustomWidgetDomainError({ code: "NOT_FOUND", message: "Preview session expired" });

packages/custom-widgets/src/server/request-executor.ts

@@ -32,6 +32,7 @@ export const MAX_REQUEST_BODY_BYTES = 10 * 1024;
 const REQUEST_TIMEOUT_MS = 10_000;
 const MAX_QUERY_REDIRECTS = 3;
 const MAX_RESPONSE_CACHE_ENTRIES = 1_000;
+export const MAX_REQUEST_DURATION_MS = 45_000;
 
 export interface CustomWidgetAuthConfig {
   type: string;
@@ -64,6 +65,29 @@ const inFlight = new Map<string, Promise<CustomWidgetHttpResponse>>();
 let cacheEpoch = 0;
 
 async function performRequest(input: CustomWidgetHttpRequest): Promise<CustomWidgetHttpResponse> {
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), MAX_REQUEST_DURATION_MS);
+  try {
+    return await performRequestWithinDeadline(input, controller.signal);
+  } catch (error) {
+    if (error instanceof CustomWidgetDomainError) throw error;
+    if (controller.signal.aborted) {
+      throw new CustomWidgetDomainError({
+        code: "BAD_GATEWAY",
+        message: "External request exceeded the total time limit",
+        cause: error,
+      });
+    }
+    throw error;
+  } finally {
+    clearTimeout(timeout);
+  }
+}
+
+async function performRequestWithinDeadline(
+  input: CustomWidgetHttpRequest,
+  deadlineSignal: AbortSignal,
+): Promise<CustomWidgetHttpResponse> {
   assertRequest(input);
   const baseUrl = validateCustomWidgetUrl(input.baseUrl);
   let currentUrl = resolveSameOriginTarget(input.baseUrl, input.targetUrl);
@@ -86,7 +110,7 @@ async function performRequest(input: CustomWidgetHttpRequest): Promise<CustomWid
         method: currentMethod,
         headers,
         body: currentBody,
-        signal: controller.signal,
+        signal: AbortSignal.any([deadlineSignal, controller.signal]),
       });
     } catch (error) {
       await dispatcher.close();

packages/custom-widgets/src/server/request-limits.ts

@@ -1,9 +1,12 @@
-import { createHash } from "node:crypto";
+import { createHash, randomUUID } from "node:crypto";
 
 import { CustomWidgetDomainError } from "./errors";
 
 const WINDOW_MS = 60_000;
-const CONCURRENCY_TTL_MS = 30_000;
+// A query may legally follow three redirects with a ten-second timeout per hop.
+// Keep the lease above that total and identify every acquisition so a late
+// release can never remove capacity owned by a newer request.
+export const CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS = 60_000;
 const CATEGORY_LIMITS = { query: 60, action: 10, delete: 3 } as const;
 type RequestCategory = keyof typeof CATEGORY_LIMITS;
 
@@ -16,8 +19,8 @@ export interface RequestLimitInput {
 
 export interface RequestLimitStore {
   incrementRate(key: string, windowMs: number): Promise<{ count: number; retryAfterMs: number }>;
-  acquireConcurrency(key: string, limit: number, ttlMs: number): Promise<boolean>;
-  releaseConcurrency(key: string): Promise<void>;
+  acquireConcurrency(key: string, ownerId: string, limit: number, ttlMs: number): Promise<boolean>;
+  releaseConcurrency(key: string, ownerId: string): Promise<void>;
 }
 
 export interface RequestLimiterOptions {
@@ -34,8 +37,8 @@ export class CustomWidgetRequestLimiter {
   public async acquire(input: RequestLimitInput): Promise<() => Promise<void>> {
     const identity = input.userId ?? "anonymous";
     const rateKey = `custom-widget:rate:${input.category}:${hashKey(identity, input.itemId)}`;
-    const userKey = `custom-widget:concurrency:user-item:${hashKey(identity, input.itemId)}`;
-    const definitionKey = `custom-widget:concurrency:definition:${hashKey(input.definitionId)}`;
+    const userKey = `custom-widget:concurrency:v2:user-item:${hashKey(identity, input.itemId)}`;
+    const definitionKey = `custom-widget:concurrency:v2:definition:${hashKey(input.definitionId)}`;
     if (!this.options.store) return this.acquireLocal(rateKey, userKey, definitionKey, input.category);
     try {
       return await this.acquireStored(rateKey, userKey, definitionKey, input.category, this.options.store);
@@ -76,17 +79,22 @@ export class CustomWidgetRequestLimiter {
     category: RequestCategory,
     store: RequestLimitStore,
   ) {
+    const ownerId = randomUUID();
     const rate = await store.incrementRate(rateKey, WINDOW_MS);
     if (rate.count > CATEGORY_LIMITS[category]) throwLimit(rate.retryAfterMs);
-    if (!(await store.acquireConcurrency(userKey, 4, CONCURRENCY_TTL_MS))) throwLimit(1_000);
+    if (!(await store.acquireConcurrency(userKey, ownerId, 4, CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS))) {
+      throwLimit(1_000);
+    }
     try {
-      if (!(await store.acquireConcurrency(definitionKey, 8, CONCURRENCY_TTL_MS))) throwLimit(1_000);
+      if (!(await store.acquireConcurrency(definitionKey, ownerId, 8, CUSTOM_WIDGET_REQUEST_CONCURRENCY_TTL_MS))) {
+        throwLimit(1_000);
+      }
     } catch (error) {
-      await store.releaseConcurrency(userKey);
+      await store.releaseConcurrency(userKey, ownerId);
       throw error;
     }
     return async () => {
-      await Promise.all([store.releaseConcurrency(userKey), store.releaseConcurrency(definitionKey)]);
+      await Promise.all([store.releaseConcurrency(userKey, ownerId), store.releaseConcurrency(definitionKey, ownerId)]);
     };
   }