@enbox/dwn-sdk-js 0.3.9 → 0.4.1

package/src/core/records-grant-authorization.ts

@@ -1,6 +1,6 @@
-import type { MessageStore } from '../types/message-store.js';
 import type { PermissionGrant } from '../protocols/permission-grant.js';
 import type { ProtocolScope } from '../utils/permission-scope.js';
+import type { ValidationStateReader } from '../types/validation-state-reader.js';
 import type { PermissionConditions, RecordsPermissionScope } from '../types/permission-types.js';
 import type { RecordsCountMessage, RecordsDeleteMessage, RecordsQueryMessage, RecordsReadMessage, RecordsSubscribeMessage, RecordsWriteMessage } from '../types/records-types.js';
 
@@ -18,10 +18,10 @@ export class RecordsGrantAuthorization {
     expectedGrantor: string,
     expectedGrantee: string,
     permissionGrant: PermissionGrant,
-    messageStore: MessageStore,
+    validationStateReader: ValidationStateReader,
   }): Promise<void> {
     const {
-      recordsWriteMessage, expectedGrantor, expectedGrantee, permissionGrant, messageStore
+      recordsWriteMessage, expectedGrantor, expectedGrantee, permissionGrant, validationStateReader
     } = input;
 
     await GrantAuthorization.performBaseValidation({
@@ -29,7 +29,7 @@ export class RecordsGrantAuthorization {
       expectedGrantor,
       expectedGrantee,
       permissionGrant,
-      messageStore
+      validationStateReader
     });
 
     // NOTE: validated the invoked permission is for Records in GrantAuthorization.performBaseValidation()
@@ -40,7 +40,7 @@ export class RecordsGrantAuthorization {
 
   /**
    * Authorizes a RecordsReadMessage using the given permission grant.
-   * @param messageStore Used to check if the given grant has been revoked.
+   * @param validationStateReader Used to check if the given grant has been revoked.
    */
   public static async authorizeRead(input: {
     recordsReadMessage: RecordsReadMessage,
@@ -48,10 +48,10 @@ export class RecordsGrantAuthorization {
     expectedGrantor: string,
     expectedGrantee: string,
     permissionGrant: PermissionGrant,
-    messageStore: MessageStore,
+    validationStateReader: ValidationStateReader,
   }): Promise<void> {
     const {
-      recordsReadMessage, recordsWriteMessageToBeRead, expectedGrantor, expectedGrantee, permissionGrant, messageStore
+      recordsReadMessage, recordsWriteMessageToBeRead, expectedGrantor, expectedGrantee, permissionGrant, validationStateReader
     } = input;
 
     await GrantAuthorization.performBaseValidation({
@@ -59,7 +59,7 @@ export class RecordsGrantAuthorization {
       expectedGrantor,
       expectedGrantee,
       permissionGrant,
-      messageStore
+      validationStateReader
     });
 
     // NOTE: validated the invoked permission is for Records in GrantAuthorization.performBaseValidation()
@@ -68,17 +68,17 @@ export class RecordsGrantAuthorization {
 
   /**
    * Authorizes the scope of a permission grant for RecordsQuery or RecordsSubscribe.
-   * @param messageStore Used to check if the grant has been revoked.
+   * @param validationStateReader Used to check if the grant has been revoked.
    */
   public static async authorizeQueryOrSubscribe(input: {
     incomingMessage: RecordsCountMessage | RecordsQueryMessage | RecordsSubscribeMessage,
     expectedGrantor: string,
     expectedGrantee: string,
     permissionGrant: PermissionGrant,
-    messageStore: MessageStore,
+    validationStateReader: ValidationStateReader,
   }): Promise<void> {
     const {
-      incomingMessage, expectedGrantor, expectedGrantee, permissionGrant, messageStore
+      incomingMessage, expectedGrantor, expectedGrantee, permissionGrant, validationStateReader
     } = input;
 
     await GrantAuthorization.performBaseValidation({
@@ -86,7 +86,7 @@ export class RecordsGrantAuthorization {
       expectedGrantor,
       expectedGrantee,
       permissionGrant,
-      messageStore
+      validationStateReader
     });
 
     // The grant's protocol must match the query/subscribe filter's protocol.
@@ -107,7 +107,7 @@ export class RecordsGrantAuthorization {
 
   /**
    * Authorizes the scope of a permission grant for RecordsDelete.
-   * @param messageStore Used to check if the grant has been revoked.
+   * @param validationStateReader Used to check if the grant has been revoked.
    */
   public static async authorizeDelete(input: {
     recordsDeleteMessage: RecordsDeleteMessage,
@@ -115,10 +115,10 @@ export class RecordsGrantAuthorization {
     expectedGrantor: string,
     expectedGrantee: string,
     permissionGrant: PermissionGrant,
-    messageStore: MessageStore,
+    validationStateReader: ValidationStateReader,
   }): Promise<void> {
     const {
-      recordsDeleteMessage, recordsWriteToDelete, expectedGrantor, expectedGrantee, permissionGrant, messageStore
+      recordsDeleteMessage, recordsWriteToDelete, expectedGrantor, expectedGrantee, permissionGrant, validationStateReader
     } = input;
 
     await GrantAuthorization.performBaseValidation({
@@ -126,7 +126,7 @@ export class RecordsGrantAuthorization {
       expectedGrantor,
       expectedGrantee,
       permissionGrant,
-      messageStore
+      validationStateReader
     });
 
     // NOTE: validated the invoked permission is for Records in GrantAuthorization.performBaseValidation()

package/src/core/replication-apply.ts

@@ -0,0 +1,412 @@
+import type { GenericMessage } from '../types/message-types.js';
+import type { ProgressToken } from '../types/subscriptions.js';
+import type { ProtocolDefinition } from '../types/protocols-types.js';
+import type { ValidationStateReader } from '../types/validation-state-reader.js';
+
+import { DwnErrorCode } from './dwn-error.js';
+import { Encoder } from '../utils/encoder.js';
+import { Message } from './message.js';
+import { DwnInterfaceName, DwnMethodName } from '../enums/dwn-interface-method.js';
+import { isCrossProtocolRef, parseCrossProtocolRef } from '../utils/protocols.js';
+
+export type ReplicationApplyOptions = {
+  dataStream?: ReadableStream<Uint8Array>;
+};
+
+export type ReplicationApplyResult =
+  | {
+      kind: 'Applied';
+      /** True when the local store retained dependency ancestry but did not advance latest state. */
+      ancestryOnly?: true;
+      /** Local admission position, when the receiving store has a durable replication log. */
+      position?: ProgressToken;
+    }
+  | { kind: 'Duplicate' }
+  | { kind: 'Superseded' }
+  | { kind: 'Incomplete'; missing: DependencyRef[] }
+  | { kind: 'Invalid'; reason: string }
+  | { kind: 'Deferred'; reason: 'tenant-inactive' | 'resolver-unavailable' | 'storage' };
+
+export type ReplicationApplyResultContext = {
+  protocolDefinition?: ProtocolDefinition;
+
+  /**
+   * Complete set of locally-missing ancestor recordIds for a missing-ancestor failure
+   * (immediate parent or record-chain construction), ordered root-first with the failure-named
+   * ancestor last, as computed by `missingAncestorRecordIdsFromReply()`. When present, the
+   * resulting `Incomplete` carries one ref per entry so the entire ancestry resolves in a
+   * single fetch pass instead of one ancestry level per retry.
+   */
+  missingAncestorRecordIds?: string[];
+};
+
+export type DependencyRef =
+  | { type: 'Protocol'; protocol: string; messageCid?: string; terminal?: boolean }
+  | { type: 'InitialWrite'; recordId: string; protocol?: string; messageCid?: string; terminal?: boolean }
+  | { type: 'Parent'; recordId: string; protocol: string; messageCid?: string; terminal?: boolean }
+  | { type: 'Ancestor'; recordId: string; protocol?: string; messageCid?: string; terminal?: boolean }
+  | { type: 'Role'; protocol: string; protocolPath: string; recipient: string; contextPrefix?: string; messageCid?: string; terminal?: boolean }
+  | { type: 'Grant'; permissionGrantId: string; messageCid?: string; terminal?: boolean }
+  | { type: 'KeyDelivery'; protocol: string; contextId: string; messageCid?: string; terminal?: boolean }
+  | { type: 'CrossProtocolRef'; protocol: string; recordId: string; messageCid?: string; terminal?: boolean }
+  | { type: 'RecordData'; recordId: string; dataCid: string; protocol?: string; messageCid?: string; terminal?: boolean };
+
+/**
+ * Converts a regular handler reply into the structured result consumed by
+ * replication sync. The DWN handler remains the dependency authority; this
+ * adapter only gives the sync transport a typed way to distinguish missing
+ * dependencies from terminal invalid messages.
+ */
+export function replicationApplyResultFromReply(
+  message: GenericMessage,
+  reply: { status: { code: number; detail?: string }; position?: ProgressToken },
+  context: ReplicationApplyResultContext = {},
+): ReplicationApplyResult {
+  const { code, detail = '' } = reply.status;
+
+  if (code === 202 || code === 204) {
+    return {
+      kind: 'Applied',
+      ...(code === 204 ? { ancestryOnly: true as const } : {}),
+      ...(reply.position === undefined ? {} : { position: reply.position }),
+    };
+  }
+
+  if (code === 409) {
+    return { kind: 'Superseded' };
+  }
+
+  if (getDwnErrorCode(detail) === DwnErrorCode.RecordsWriteNotAllowedAfterDelete) {
+    return { kind: 'Superseded' };
+  }
+
+  if (isResolverFailure(detail)) {
+    return { kind: 'Deferred', reason: 'resolver-unavailable' };
+  }
+
+  const missing = dependencyRefsFromStatus(message, code, detail, context);
+  if (missing.length > 0) {
+    return { kind: 'Incomplete', missing };
+  }
+
+  if (code >= 500) {
+    return { kind: 'Deferred', reason: 'storage' };
+  }
+
+  return { kind: 'Invalid', reason: detail || `replicated message rejected with status ${code}` };
+}
+
+/**
+ * Computes the complete set of locally-missing ancestor recordIds for a replicated message that
+ * failed on a missing ancestor — either the immediate parent referential check
+ * (`ProtocolAuthorizationParentRecordNotFound`) or record-chain construction
+ * (`ProtocolAuthorizationParentNotFoundConstructingRecordChain`).
+ *
+ * Both failures name a single ancestor, but the message's `contextId` names the full ancestor
+ * recordId chain, so every segment the failed check could not reach is presence-checked here and
+ * all absent segments are returned in one batch (ordered root-first, named ancestor last).
+ * Feeding the result to `replicationApplyResultFromReply()` via `ReplicationApplyResultContext`
+ * lets sync fetch the entire ancestry in a single pass instead of one level per retry.
+ *
+ * @returns the missing ancestor recordIds, or `undefined` when the reply is not a
+ *          missing-ancestor failure or the chain cannot be determined from the message (callers
+ *          then fall back to emitting the single ancestor named by the failure).
+ */
+export async function missingAncestorRecordIdsFromReply(
+  tenant: string,
+  message: GenericMessage,
+  reply: { status: { detail?: string } },
+  validationStateReader: ValidationStateReader,
+): Promise<string[] | undefined> {
+  const detail = reply.status.detail ?? '';
+  const errorCode = getDwnErrorCode(detail);
+
+  let namedRecordId: string | undefined;
+  if (errorCode === DwnErrorCode.ProtocolAuthorizationParentRecordNotFound) {
+    namedRecordId = parentRecordIdFromMessage(message, detail);
+  } else if (errorCode === DwnErrorCode.ProtocolAuthorizationParentNotFoundConstructingRecordChain) {
+    namedRecordId = ancestorRecordIdFromDetail(detail);
+  } else {
+    return undefined;
+  }
+
+  const contextId = (message as { contextId?: unknown }).contextId;
+  if (namedRecordId === undefined || typeof contextId !== 'string') {
+    return undefined;
+  }
+
+  // `contextId` is the ancestor recordId chain ending with the record's own recordId. The failed
+  // check stopped at `namedRecordId`, so every segment below it is locally present; only the
+  // segments above it (closer to the root) still need a presence check. Presence means an initial
+  // write exists — exactly what record-chain construction reads for the segments above the
+  // immediate parent.
+  const ancestorRecordIds = contextId.split('/').slice(0, -1);
+  const namedIndex = ancestorRecordIds.indexOf(namedRecordId);
+  if (namedIndex === -1) {
+    return undefined;
+  }
+
+  const missingRecordIds: string[] = [];
+  for (const ancestorRecordId of ancestorRecordIds.slice(0, namedIndex)) {
+    const initialWrite = await validationStateReader.fetchInitialWrite(tenant, ancestorRecordId);
+    if (initialWrite === undefined) {
+      missingRecordIds.push(ancestorRecordId);
+    }
+  }
+  missingRecordIds.push(namedRecordId);
+
+  return missingRecordIds;
+}
+
+function dependencyRefsFromStatus(
+  message: GenericMessage,
+  code: number,
+  detail: string,
+  context: ReplicationApplyResultContext,
+): DependencyRef[] {
+  const errorCode = getDwnErrorCode(detail);
+  switch (errorCode) {
+  case DwnErrorCode.ProtocolAuthorizationProtocolNotFound:
+  case DwnErrorCode.ProtocolsConfigureComposedProtocolNotInstalled:
+    return toRefList(protocolDependencyFromMessage(message, detail));
+  case DwnErrorCode.ProtocolAuthorizationParentRecordNotFound:
+    return parentDependenciesFromMessage(message, detail, context);
+  case DwnErrorCode.ProtocolAuthorizationCrossProtocolParentNotFound:
+    return toRefList(parentDependencyFromMessage(message, detail));
+  case DwnErrorCode.ProtocolAuthorizationParentNotFoundConstructingRecordChain:
+    return ancestorDependenciesFromMessage(message, detail, context);
+  case DwnErrorCode.RecordsWriteGetInitialWriteNotFound:
+    return toRefList(initialWriteDependencyFromMessage(message));
+  case DwnErrorCode.GrantAuthorizationGrantMissing:
+    return toRefList(grantDependencyFromMessage(message));
+  case DwnErrorCode.ProtocolAuthorizationMatchingRoleRecordNotFound:
+    return toRefList(roleDependencyFromMessage(message, context));
+  case DwnErrorCode.RecordsWriteMissingDataInPrevious:
+  case DwnErrorCode.RecordsWriteMissingEncodedDataInPrevious:
+    return toRefList(recordDataDependencyFromMessage(message));
+  default:
+    break;
+  }
+
+  if (code === 404 && isRecordsDelete(message)) {
+    return [{ type: 'InitialWrite', recordId: getRecordsDeleteRecordId(message) }];
+  }
+
+  return [];
+}
+
+function toRefList(ref: DependencyRef | undefined): DependencyRef[] {
+  return ref === undefined ? [] : [ref];
+}
+
+function getDwnErrorCode(detail: string): string | undefined {
+  const delimiter = detail.indexOf(':');
+  return delimiter === -1 ? undefined : detail.slice(0, delimiter);
+}
+
+function isResolverFailure(detail: string): boolean {
+  return getDwnErrorCode(detail) === DwnErrorCode.GeneralJwsVerifierGetPublicKeyNotFound;
+}
+
+function protocolDependencyFromMessage(message: GenericMessage, detail: string): DependencyRef | undefined {
+  const composedProtocol = /composed protocol '([^']+)'/.exec(detail)?.[1];
+  if (composedProtocol !== undefined) {
+    return { type: 'Protocol', protocol: composedProtocol };
+  }
+
+  const protocol = (message.descriptor as Record<string, unknown>).protocol;
+  return typeof protocol === 'string' ? { type: 'Protocol', protocol } : undefined;
+}
+
+/**
+ * Emits the refs for an immediate-parent referential failure. The missing parent keeps its
+ * existing `Parent` ref; when the apply context carries the layer-batched missing-ancestor set
+ * (see `missingAncestorRecordIdsFromReply()`), an `Ancestor` ref is additionally emitted for
+ * every locally-absent segment above the parent, so the entire ancestry resolves in a single
+ * fetch pass instead of one level per retry.
+ */
+function parentDependenciesFromMessage(
+  message: GenericMessage,
+  detail: string,
+  context: ReplicationApplyResultContext,
+): DependencyRef[] {
+  const parentRef = parentDependencyFromMessage(message, detail);
+  if (parentRef === undefined) {
+    return [];
+  }
+
+  const protocol = (message.descriptor as Record<string, unknown>).protocol;
+  const protocolProperty = typeof protocol === 'string' ? { protocol } : {};
+  const ancestorRefs = (context.missingAncestorRecordIds ?? [])
+    .filter((recordId): boolean => recordId !== parentRef.recordId)
+    .map((recordId): DependencyRef => ({ type: 'Ancestor', recordId, ...protocolProperty }));
+
+  return [...ancestorRefs, parentRef];
+}
+
+function parentDependencyFromMessage(
+  message: GenericMessage,
+  detail: string,
+): Extract<DependencyRef, { type: 'Parent' }> | undefined {
+  const descriptor = message.descriptor as Record<string, unknown>;
+  const parentId = parentRecordIdFromMessage(message, detail);
+  const protocol = /in protocol '([^']+)'/.exec(detail)?.[1] ?? descriptor.protocol;
+
+  if (parentId === undefined || typeof protocol !== 'string') {
+    return undefined;
+  }
+
+  return { type: 'Parent', recordId: parentId, protocol };
+}
+
+function parentRecordIdFromMessage(message: GenericMessage, detail: string): string | undefined {
+  const descriptor = message.descriptor as Record<string, unknown>;
+  return typeof descriptor.parentId === 'string'
+    ? descriptor.parentId
+    : /parent record '([^']+)'/.exec(detail)?.[1];
+}
+
+/**
+ * Emits `Ancestor` refs for a record-chain construction failure. When the apply context carries
+ * the layer-batched missing-ancestor set (see `missingAncestorRecordIdsFromReply()`), one ref is
+ * emitted per missing segment so the entire ancestry resolves in a single fetch pass; otherwise
+ * falls back to the single ancestor named by the failure detail (e.g. for messages that carry no
+ * `contextId` chain).
+ */
+function ancestorDependenciesFromMessage(
+  message: GenericMessage,
+  detail: string,
+  context: ReplicationApplyResultContext,
+): DependencyRef[] {
+  const protocol = (message.descriptor as Record<string, unknown>).protocol;
+  const protocolProperty = typeof protocol === 'string' ? { protocol } : {};
+
+  const batchedRecordIds = context.missingAncestorRecordIds;
+  if (batchedRecordIds !== undefined && batchedRecordIds.length > 0) {
+    return batchedRecordIds.map((recordId): DependencyRef => ({ type: 'Ancestor', recordId, ...protocolProperty }));
+  }
+
+  const namedRecordId = ancestorRecordIdFromDetail(detail);
+  if (namedRecordId === undefined) {
+    return [];
+  }
+
+  return [{ type: 'Ancestor', recordId: namedRecordId, ...protocolProperty }];
+}
+
+function ancestorRecordIdFromDetail(detail: string): string | undefined {
+  return /ID ([^ ]+)/.exec(detail)?.[1];
+}
+
+function initialWriteDependencyFromMessage(message: GenericMessage): DependencyRef | undefined {
+  if (!isRecordsWrite(message)) {
+    return undefined;
+  }
+
+  const protocol = (message.descriptor as Record<string, unknown>).protocol;
+  return {
+    type     : 'InitialWrite',
+    recordId : message.recordId,
+    ...(typeof protocol === 'string' ? { protocol } : {}),
+  };
+}
+
+function grantDependencyFromMessage(message: GenericMessage): DependencyRef | undefined {
+  const descriptorGrantId = (message.descriptor as Record<string, unknown>).permissionGrantId;
+  if (typeof descriptorGrantId === 'string') {
+    return { type: 'Grant', permissionGrantId: descriptorGrantId };
+  }
+
+  const payloadGrantId = getSignaturePayload(message)?.permissionGrantId;
+  return typeof payloadGrantId === 'string'
+    ? { type: 'Grant', permissionGrantId: payloadGrantId }
+    : undefined;
+}
+
+function recordDataDependencyFromMessage(message: GenericMessage): DependencyRef | undefined {
+  if (!isRecordsWrite(message)) {
+    return undefined;
+  }
+
+  const descriptor = message.descriptor as Record<string, unknown>;
+  const dataCid = descriptor.dataCid;
+  const protocol = descriptor.protocol;
+  return typeof dataCid === 'string'
+    ? {
+      type     : 'RecordData',
+      recordId : message.recordId,
+      dataCid,
+      ...(typeof protocol === 'string' ? { protocol } : {}),
+    }
+    : undefined;
+}
+
+function roleDependencyFromMessage(message: GenericMessage, context: ReplicationApplyResultContext): DependencyRef | undefined {
+  const descriptor = message.descriptor as Record<string, unknown>;
+  const filter = descriptor.filter as Record<string, unknown> | undefined;
+  const protocol = descriptor.protocol ?? filter?.protocol;
+  const protocolRole = getSignaturePayload(message)?.protocolRole;
+  const recipient = Message.getAuthor(message);
+
+  if (typeof protocol !== 'string' || typeof protocolRole !== 'string' || recipient === undefined) {
+    return undefined;
+  }
+
+  let roleProtocol = protocol;
+  let roleProtocolPath = protocolRole;
+  if (isCrossProtocolRef(protocolRole)) {
+    const parsed = parseCrossProtocolRef(protocolRole);
+    const referencedProtocol = parsed === undefined ? undefined : context.protocolDefinition?.uses?.[parsed.alias];
+    if (parsed !== undefined && referencedProtocol !== undefined) {
+      roleProtocol = referencedProtocol;
+      roleProtocolPath = parsed.protocolPath;
+    }
+  }
+
+  let contextId: string | undefined;
+  if (typeof descriptor.contextId === 'string') {
+    contextId = descriptor.contextId;
+  } else if (typeof filter?.contextId === 'string') {
+    contextId = filter.contextId;
+  }
+  const roleSegments = roleProtocolPath.split('/').length - 1;
+  const contextPrefix = roleSegments > 0 && contextId !== undefined
+    ? contextId.split('/').slice(0, roleSegments).join('/')
+    : undefined;
+
+  return {
+    type         : 'Role',
+    protocol     : roleProtocol,
+    protocolPath : roleProtocolPath,
+    recipient,
+    ...(contextPrefix === undefined ? {} : { contextPrefix }),
+  };
+}
+
+function isRecordsWrite(message: GenericMessage): message is GenericMessage & { recordId: string } {
+  return message.descriptor.interface === DwnInterfaceName.Records &&
+    message.descriptor.method === DwnMethodName.Write &&
+    typeof (message as { recordId?: unknown }).recordId === 'string';
+}
+
+function isRecordsDelete(message: GenericMessage): boolean {
+  return message.descriptor.interface === DwnInterfaceName.Records &&
+    message.descriptor.method === DwnMethodName.Delete;
+}
+
+function getRecordsDeleteRecordId(message: GenericMessage): string {
+  return (message.descriptor as unknown as { recordId: string }).recordId;
+}
+
+function getSignaturePayload(message: GenericMessage): Record<string, unknown> | undefined {
+  const payload = message.authorization?.signature.payload;
+  if (payload === undefined) {
+    return undefined;
+  }
+
+  try {
+    return Encoder.base64UrlToObject(payload) as Record<string, unknown>;
+  } catch {
+    return undefined;
+  }
+}

package/src/core/resumable-task-manager.ts

@@ -11,6 +11,7 @@ export type ResumableTask = {
   data: any;
 };
 
+type ResumableTaskHandler = (taskData: any) => Promise<unknown>;
 
 export class ResumableTaskManager {
 
@@ -20,34 +21,34 @@ export class ResumableTaskManager {
   public static readonly timeoutExtensionFrequencyInSeconds = 30;
 
   private resumableTaskBatchSize = 100;
-  private readonly resumableTaskHandlers: { [key:string]: (taskData: any) => Promise<void> };
+  private readonly resumableTaskHandlers: { [key:string]: ResumableTaskHandler };
 
   public constructor(private readonly resumableTaskStore: ResumableTaskStore, storageController: StorageController) {
     // assign resumable task handlers
     this.resumableTaskHandlers = {
       // NOTE: The arrow function is IMPORTANT here, else the `this` context will be lost within the invoked method.
       // e.g. code within performRecordsDelete() won't know `this` refers to the `storageController` instance.
-      [ResumableTaskName.RecordsDelete] : async (task): Promise<void> => await storageController.performRecordsDelete(task),
-      [ResumableTaskName.RecordsSquash] : async (task): Promise<void> => await storageController.performRecordsSquash(task),
+      [ResumableTaskName.RecordsDelete] : async (task): Promise<unknown> => storageController.performRecordsDelete(task),
+      [ResumableTaskName.RecordsSquash] : async (task): Promise<unknown> => storageController.performRecordsSquash(task),
     };
   }
 
   /**
    * Runs a new resumable task.
    */
-  public async run(task: ResumableTask): Promise<void> {
+  public async run<T = unknown>(task: ResumableTask): Promise<T> {
     const timeoutInSeconds = ResumableTaskManager.timeoutExtensionFrequencyInSeconds * 2; // give ample time for extension to take place
 
     // register the new resumable task before running it so that it can be resumed if it times out for any reason
     const managedResumableTask = await this.resumableTaskStore.register(task, timeoutInSeconds);
-    await this.runWithAutomaticTimeoutExtension(managedResumableTask);
+    return this.runWithAutomaticTimeoutExtension<T>(managedResumableTask);
   }
 
   /**
    * Runs a resumable task with automatic timeout extension.
    * Deletes the task from the resumable task store once it is completed.
    */
-  private async runWithAutomaticTimeoutExtension(managedTask: ManagedResumableTask): Promise<void> {
+  private async runWithAutomaticTimeoutExtension<T = unknown>(managedTask: ManagedResumableTask): Promise<T> {
     const timeoutInSeconds = ResumableTaskManager.timeoutExtensionFrequencyInSeconds * 2; // give ample time for extension to take place
 
     let timer!: ReturnType<typeof setInterval>;
@@ -58,8 +59,9 @@ export class ResumableTaskManager {
       }, ResumableTaskManager.timeoutExtensionFrequencyInSeconds * 1000);
 
       const handler = this.resumableTaskHandlers[managedTask.task.name];
-      await handler(managedTask.task.data);
+      const result = await handler(managedTask.task.data) as T;
       await this.resumableTaskStore.delete(managedTask.id);
+      return result;
     } finally {
       ResumableTaskManager.clearTimeoutExtensionTimer(timer);
     }
@@ -112,4 +114,4 @@ export class ResumableTaskManager {
       }
     }
   }
-}
+}