@fluidframework/container-loader 2.0.0-dev-rc.3.0.0.250606 → 2.0.0-dev-rc.3.0.0.253463

package/src/container.ts

@@ -6,15 +6,18 @@
 import { TypedEventEmitter, performance } from "@fluid-internal/client-utils";
 import {
 	AttachState,
-	ContainerWarning,
 	IAudience,
+	ICriticalContainerError,
+	IDeltaManager,
+	ReadOnlyInfo,
+} from "@fluidframework/container-definitions";
+import {
+	ContainerWarning,
 	IBatchMessage,
 	ICodeDetailsLoader,
 	IContainer,
 	IContainerEvents,
 	IContainerLoadMode,
-	ICriticalContainerError,
-	IDeltaManager,
 	IFluidCodeDetails,
 	IFluidCodeDetailsComparer,
 	IFluidModuleWithDetails,
@@ -23,18 +26,17 @@ import {
 	IProvideFluidCodeDetailsComparer,
 	IProvideRuntimeFactory,
 	IRuntime,
-	ReadOnlyInfo,
 	isFluidCodeDetails,
-} from "@fluidframework/container-definitions";
+} from "@fluidframework/container-definitions/internal";
 import {
 	FluidObject,
 	IEvent,
 	IRequest,
-	type ISignalEnvelope,
 	ITelemetryBaseProperties,
 	LogLevel,
 } from "@fluidframework/core-interfaces";
-import { assert, isPromiseLike, unreachableCase } from "@fluidframework/core-utils";
+import { type ISignalEnvelope } from "@fluidframework/core-interfaces/internal";
+import { assert, isPromiseLike, unreachableCase } from "@fluidframework/core-utils/internal";
 import {
 	IDocumentService,
 	IDocumentServiceFactory,
@@ -43,7 +45,7 @@ import {
 	ISnapshot,
 	IThrottlingWarning,
 	IUrlResolver,
-} from "@fluidframework/driver-definitions";
+} from "@fluidframework/driver-definitions/internal";
 import {
 	MessageType2,
 	OnlineStatus,
@@ -52,7 +54,7 @@ import {
 	isOnline,
 	readAndParse,
 	runWithRetry,
-} from "@fluidframework/driver-utils";
+} from "@fluidframework/driver-utils/internal";
 import { IQuorumSnapshot } from "@fluidframework/protocol-base";
 import {
 	IClient,
@@ -73,14 +75,13 @@ import {
 	MessageType,
 	SummaryType,
 } from "@fluidframework/protocol-definitions";
+import { ITelemetryLoggerExt, type TelemetryEventCategory } from "@fluidframework/telemetry-utils";
 import {
 	EventEmitterWithErrorHandling,
 	GenericError,
 	IFluidErrorBase,
-	ITelemetryLoggerExt,
 	MonitoringContext,
 	PerformanceEvent,
-	type TelemetryEventCategory,
 	UsageError,
 	connectedEventName,
 	createChildLogger,
@@ -89,9 +90,10 @@ import {
 	normalizeError,
 	raiseConnectedEvent,
 	wrapError,
-} from "@fluidframework/telemetry-utils";
+} from "@fluidframework/telemetry-utils/internal";
 import structuredClone from "@ungap/structured-clone";
 import { v4 as uuid } from "uuid";
+
 import { AttachProcessProps, AttachmentData, runRetriableAttachProcess } from "./attachment.js";
 import { Audience } from "./audience.js";
 import { ConnectionManager } from "./connectionManager.js";

package/src/containerContext.ts

@@ -6,15 +6,18 @@
 import {
 	AttachState,
 	IAudience,
-	IBatchMessage,
-	IContainerContext,
 	ICriticalContainerError,
 	IDeltaManager,
+} from "@fluidframework/container-definitions";
+import {
+	IBatchMessage,
+	IContainerContext,
 	ILoader,
 	ILoaderOptions,
-} from "@fluidframework/container-definitions";
-import { type FluidObject, type ISignalEnvelope } from "@fluidframework/core-interfaces";
-import { IDocumentStorageService, ISnapshot } from "@fluidframework/driver-definitions";
+} from "@fluidframework/container-definitions/internal";
+import { type FluidObject } from "@fluidframework/core-interfaces";
+import { type ISignalEnvelope } from "@fluidframework/core-interfaces/internal";
+import { IDocumentStorageService, ISnapshot } from "@fluidframework/driver-definitions/internal";
 import {
 	IClientDetails,
 	IDocumentMessage,

package/src/containerStorageAdapter.ts

@@ -4,9 +4,9 @@
  */
 
 import { bufferToString, stringToBuffer } from "@fluid-internal/client-utils";
-import { ISnapshotTreeWithBlobContents } from "@fluidframework/container-definitions";
+import { ISnapshotTreeWithBlobContents } from "@fluidframework/container-definitions/internal";
 import { IDisposable } from "@fluidframework/core-interfaces";
-import { assert } from "@fluidframework/core-utils";
+import { assert } from "@fluidframework/core-utils/internal";
 import {
 	FetchSource,
 	IDocumentService,
@@ -15,8 +15,8 @@ import {
 	ISnapshot,
 	ISnapshotFetchOptions,
 	ISummaryContext,
-} from "@fluidframework/driver-definitions";
-import { UsageError } from "@fluidframework/driver-utils";
+} from "@fluidframework/driver-definitions/internal";
+import { UsageError } from "@fluidframework/driver-utils/internal";
 import {
 	ICreateBlobResponse,
 	ISnapshotTree,
@@ -25,6 +25,7 @@ import {
 	IVersion,
 } from "@fluidframework/protocol-definitions";
 import { ITelemetryLoggerExt } from "@fluidframework/telemetry-utils";
+
 import { IDetachedBlobStorage } from "./loader.js";
 import { ProtocolTreeStorageService } from "./protocolTreeDocumentStorageService.js";
 import { RetriableDocumentStorageService } from "./retriableDocumentStorageService.js";

package/src/contracts.ts

@@ -7,12 +7,11 @@ import {
 	IConnectionDetails,
 	ICriticalContainerError,
 	IDeltaQueue,
-	IFluidCodeDetails,
 	ReadOnlyInfo,
-	isFluidPackage,
 } from "@fluidframework/container-definitions";
+import { IFluidCodeDetails, isFluidPackage } from "@fluidframework/container-definitions/internal";
 import { IErrorBase, ITelemetryBaseProperties } from "@fluidframework/core-interfaces";
-import { IContainerPackageInfo } from "@fluidframework/driver-definitions";
+import { IContainerPackageInfo } from "@fluidframework/driver-definitions/internal";
 import {
 	ConnectionMode,
 	IClientConfiguration,

package/src/debugLogger.ts

@@ -9,14 +9,13 @@ import {
 	ITelemetryBaseLogger,
 	ITelemetryBaseProperties,
 } from "@fluidframework/core-interfaces";
-
+import { ITelemetryLoggerExt } from "@fluidframework/telemetry-utils";
 import {
-	ITelemetryLoggerExt,
 	ITelemetryLoggerPropertyBags,
 	createMultiSinkLogger,
 	eventNamespaceSeparator,
 	formatTick,
-} from "@fluidframework/telemetry-utils";
+} from "@fluidframework/telemetry-utils/internal";
 
 // This import style is necessary to ensure the emitted JS code works in both CJS and ESM.
 import debugPkg from "debug";

package/src/deltaManager.ts

@@ -14,15 +14,19 @@ import {
 	IEventProvider,
 	type ITelemetryBaseEvent,
 	ITelemetryBaseProperties,
-	IThrottlingWarning,
 } from "@fluidframework/core-interfaces";
-import { assert } from "@fluidframework/core-utils";
+import { IThrottlingWarning } from "@fluidframework/core-interfaces/internal";
+import { assert } from "@fluidframework/core-utils/internal";
+import { DriverErrorTypes } from "@fluidframework/driver-definitions";
 import {
-	DriverErrorTypes,
 	IDocumentDeltaStorageService,
 	IDocumentService,
-} from "@fluidframework/driver-definitions";
-import { MessageType2, NonRetryableError, isRuntimeMessage } from "@fluidframework/driver-utils";
+} from "@fluidframework/driver-definitions/internal";
+import {
+	MessageType2,
+	NonRetryableError,
+	isRuntimeMessage,
+} from "@fluidframework/driver-utils/internal";
 import {
 	ConnectionMode,
 	IDocumentMessage,
@@ -31,17 +35,19 @@ import {
 	MessageType,
 } from "@fluidframework/protocol-definitions";
 import {
-	DataCorruptionError,
-	DataProcessingError,
 	type ITelemetryErrorEventExt,
 	type ITelemetryGenericEventExt,
 	ITelemetryLoggerExt,
+} from "@fluidframework/telemetry-utils";
+import {
+	DataCorruptionError,
+	DataProcessingError,
 	UsageError,
 	extractSafePropertiesFromMessage,
 	isFluidError,
 	normalizeError,
 	safeRaiseEvent,
-} from "@fluidframework/telemetry-utils";
+} from "@fluidframework/telemetry-utils/internal";
 import { v4 as uuid } from "uuid";
 
 import {

package/src/deltaQueue.ts

@@ -5,7 +5,7 @@
 
 import { TypedEventEmitter, performance } from "@fluid-internal/client-utils";
 import { IDeltaQueue, IDeltaQueueEvents } from "@fluidframework/container-definitions";
-import { assert } from "@fluidframework/core-utils";
+import { assert } from "@fluidframework/core-utils/internal";
 import Deque from "double-ended-queue";
 
 export interface IDeltaQueueWriter<T> {

package/src/error.ts

@@ -3,14 +3,15 @@
  * Licensed under the MIT License.
  */
 
-import { ContainerErrorTypes } from "@fluidframework/container-definitions";
-import { ITelemetryBaseProperties, IThrottlingWarning } from "@fluidframework/core-interfaces";
+import { ContainerErrorTypes } from "@fluidframework/container-definitions/internal";
+import { ITelemetryBaseProperties } from "@fluidframework/core-interfaces";
+import { IThrottlingWarning } from "@fluidframework/core-interfaces/internal";
+import { ITelemetryLoggerExt } from "@fluidframework/telemetry-utils";
 import {
 	IFluidErrorBase,
-	ITelemetryLoggerExt,
 	LoggingError,
 	wrapErrorAndLog,
-} from "@fluidframework/telemetry-utils";
+} from "@fluidframework/telemetry-utils/internal";
 
 /**
  * Warning emitted when requests to storage are being throttled.

package/src/loader.ts

@@ -12,7 +12,7 @@ import {
 	ILoaderOptions as ILoaderOptions1,
 	IProvideFluidCodeDetailsComparer,
 	LoaderHeader,
-} from "@fluidframework/container-definitions";
+} from "@fluidframework/container-definitions/internal";
 import {
 	FluidObject,
 	IConfigProviderBase,
@@ -24,18 +24,19 @@ import {
 	IDocumentStorageService,
 	IResolvedUrl,
 	IUrlResolver,
-} from "@fluidframework/driver-definitions";
+} from "@fluidframework/driver-definitions/internal";
 import { IClientDetails } from "@fluidframework/protocol-definitions";
+import { ITelemetryLoggerExt } from "@fluidframework/telemetry-utils";
 import {
-	ITelemetryLoggerExt,
 	MonitoringContext,
 	PerformanceEvent,
 	UsageError,
 	createChildMonitoringContext,
 	mixinMonitoringContext,
 	sessionStorageConfigProvider,
-} from "@fluidframework/telemetry-utils";
+} from "@fluidframework/telemetry-utils/internal";
 import { v4 as uuid } from "uuid";
+
 import { Container } from "./container.js";
 import { DebugLogger } from "./debugLogger.js";
 import { pkgVersion } from "./packageVersion.js";

package/src/location-redirection-utilities/resolveWithLocationRedirection.ts

@@ -4,12 +4,12 @@
  */
 
 import { IRequest, ITelemetryBaseLogger } from "@fluidframework/core-interfaces";
+import { DriverErrorTypes } from "@fluidframework/driver-definitions";
 import {
-	DriverErrorTypes,
 	ILocationRedirectionError,
 	IUrlResolver,
-} from "@fluidframework/driver-definitions";
-import { createChildLogger } from "@fluidframework/telemetry-utils";
+} from "@fluidframework/driver-definitions/internal";
+import { createChildLogger } from "@fluidframework/telemetry-utils/internal";
 
 /**
  * Checks if the error is location redirection error.

package/src/noopHeuristic.ts

@@ -5,8 +5,8 @@
 
 import { TypedEventEmitter } from "@fluid-internal/client-utils";
 import { IEvent } from "@fluidframework/core-interfaces";
-import { assert, Timer } from "@fluidframework/core-utils";
-import { isRuntimeMessage } from "@fluidframework/driver-utils";
+import { assert, Timer } from "@fluidframework/core-utils/internal";
+import { isRuntimeMessage } from "@fluidframework/driver-utils/internal";
 import { ISequencedDocumentMessage } from "@fluidframework/protocol-definitions";
 
 const defaultNoopTimeFrequency = 2000;

package/src/packageVersion.ts

@@ -6,4 +6,4 @@
  */
 
 export const pkgName = "@fluidframework/container-loader";
-export const pkgVersion = "2.0.0-dev-rc.3.0.0.250606";
+export const pkgVersion = "2.0.0-dev-rc.3.0.0.253463";

package/src/protocol.ts

@@ -3,8 +3,8 @@
  * Licensed under the MIT License.
  */
 
-import { IAudienceOwner } from "@fluidframework/container-definitions";
-import { canBeCoalescedByService } from "@fluidframework/driver-utils";
+import { IAudienceOwner } from "@fluidframework/container-definitions/internal";
+import { canBeCoalescedByService } from "@fluidframework/driver-utils/internal";
 import {
 	IProtocolHandler as IBaseProtocolHandler,
 	IQuorumSnapshot,

package/src/protocolTreeDocumentStorageService.ts

@@ -4,7 +4,10 @@
  */
 
 import { IDisposable } from "@fluidframework/core-interfaces";
-import { IDocumentStorageService, ISummaryContext } from "@fluidframework/driver-definitions";
+import {
+	IDocumentStorageService,
+	ISummaryContext,
+} from "@fluidframework/driver-definitions/internal";
 import { ISummaryTree } from "@fluidframework/protocol-definitions";
 
 /**

package/src/quorum.ts

@@ -3,7 +3,7 @@
  * Licensed under the MIT License.
  */
 
-import { IFluidCodeDetails } from "@fluidframework/container-definitions";
+import { IFluidCodeDetails } from "@fluidframework/container-definitions/internal";
 import { ICommittedProposal } from "@fluidframework/protocol-definitions";
 
 export function initQuorumValuesFromCodeDetails(

package/src/retriableDocumentStorageService.ts

@@ -4,7 +4,7 @@
  */
 
 import { IDisposable } from "@fluidframework/core-interfaces";
-import { assert } from "@fluidframework/core-utils";
+import { assert } from "@fluidframework/core-utils/internal";
 import {
 	FetchSource,
 	IDocumentStorageService,
@@ -12,8 +12,8 @@ import {
 	ISnapshot,
 	ISnapshotFetchOptions,
 	ISummaryContext,
-} from "@fluidframework/driver-definitions";
-import { runWithRetry } from "@fluidframework/driver-utils";
+} from "@fluidframework/driver-definitions/internal";
+import { runWithRetry } from "@fluidframework/driver-utils/internal";
 import {
 	ICreateBlobResponse,
 	ISnapshotTree,
@@ -21,7 +21,8 @@ import {
 	ISummaryTree,
 	IVersion,
 } from "@fluidframework/protocol-definitions";
-import { GenericError, ITelemetryLoggerExt, UsageError } from "@fluidframework/telemetry-utils";
+import { ITelemetryLoggerExt } from "@fluidframework/telemetry-utils";
+import { GenericError, UsageError } from "@fluidframework/telemetry-utils/internal";
 
 export class RetriableDocumentStorageService implements IDocumentStorageService, IDisposable {
 	private _disposed = false;

package/src/serializedStateManager.ts

@@ -3,27 +3,33 @@
  * Licensed under the MIT License.
  */
 
-import { IGetPendingLocalStateProps, IRuntime } from "@fluidframework/container-definitions";
-import { assert } from "@fluidframework/core-utils";
+import {
+	IGetPendingLocalStateProps,
+	IRuntime,
+} from "@fluidframework/container-definitions/internal";
+import { assert } from "@fluidframework/core-utils/internal";
 import {
 	IDocumentStorageService,
 	IResolvedUrl,
 	ISnapshot,
-} from "@fluidframework/driver-definitions";
-import { isInstanceOfISnapshot } from "@fluidframework/driver-utils";
+} from "@fluidframework/driver-definitions/internal";
+import { isInstanceOfISnapshot } from "@fluidframework/driver-utils/internal";
 import {
+	type IDocumentAttributes,
 	ISequencedDocumentMessage,
 	ISnapshotTree,
 	IVersion,
 } from "@fluidframework/protocol-definitions";
+import { ITelemetryLoggerExt } from "@fluidframework/telemetry-utils";
 import {
-	ITelemetryLoggerExt,
 	MonitoringContext,
 	PerformanceEvent,
 	UsageError,
 	createChildMonitoringContext,
-} from "@fluidframework/telemetry-utils";
+} from "@fluidframework/telemetry-utils/internal";
+
 import { ISerializableBlobContents, getBlobContentsFromTree } from "./containerStorageAdapter.js";
+import { getDocumentAttributes } from "./utils.js";
 
 export interface SnapshotWithBlobs {
 	/**
@@ -65,10 +71,16 @@ export interface IPendingDetachedContainerState extends SnapshotWithBlobs {
 	pendingRuntimeState?: unknown;
 }
 
+interface SnapshotInfo extends SnapshotWithBlobs {
+	snapshotSequenceNumber: number;
+}
+
 export class SerializedStateManager {
 	private readonly processedOps: ISequencedDocumentMessage[] = [];
 	private snapshot: SnapshotWithBlobs | undefined;
 	private readonly mc: MonitoringContext;
+	private latestSnapshot: SnapshotInfo | undefined;
+	private refreshSnapshot: Promise<void> | undefined;
 
 	constructor(
 		private readonly pendingLocalState: IPendingContainerState | undefined,
@@ -78,6 +90,7 @@ export class SerializedStateManager {
 			"readBlob" | "getSnapshotTree" | "getSnapshot" | "getVersions"
 		>,
 		private readonly _offlineLoadEnabled: boolean,
+		private readonly newSnapshotFetched?: () => void,
 	) {
 		this.mc = createChildMonitoringContext({
 			logger: subLogger,
@@ -92,6 +105,7 @@ export class SerializedStateManager {
 	public addProcessedOp(message: ISequencedDocumentMessage) {
 		if (this.offlineLoadEnabled) {
 			this.processedOps.push(message);
+			this.updateSnapshotAndProcessedOpsMaybe();
 		}
 	}
 
@@ -100,13 +114,12 @@ export class SerializedStateManager {
 		supportGetSnapshotApi: boolean,
 	) {
 		if (this.pendingLocalState === undefined) {
-			const { snapshot, version } = supportGetSnapshotApi
-				? await fetchISnapshot(this.mc, this.storageAdapter, specifiedVersion)
-				: await fetchISnapshotTree(this.mc, this.storageAdapter, specifiedVersion);
-			const baseSnapshot: ISnapshotTree | undefined = isInstanceOfISnapshot(snapshot)
-				? snapshot.snapshotTree
-				: snapshot;
-			assert(baseSnapshot !== undefined, 0x8e4 /* Snapshot should exist */);
+			const { baseSnapshot, version } = await getSnapshotTree(
+				this.mc,
+				this.storageAdapter,
+				supportGetSnapshotApi,
+				specifiedVersion,
+			);
 			// non-interactive clients will not have any pending state we want to save
 			if (this.offlineLoadEnabled) {
 				const snapshotBlobs = await getBlobContentsFromTree(
@@ -119,10 +132,71 @@ export class SerializedStateManager {
 		} else {
 			const { baseSnapshot, snapshotBlobs } = this.pendingLocalState;
 			this.snapshot = { baseSnapshot, snapshotBlobs };
+			this.refreshSnapshot ??= (async () => {
+				this.latestSnapshot = await getLatestSnapshotInfo(
+					this.mc,
+					this.storageAdapter,
+					supportGetSnapshotApi,
+				);
+				this.newSnapshotFetched?.();
+				this.updateSnapshotAndProcessedOpsMaybe();
+			})();
+
 			return { baseSnapshot, version: undefined };
 		}
 	}
 
+	/**
+	 * Updates class snapshot and processedOps if we have a new snapshot and it's among processedOps range.
+	 */
+	private updateSnapshotAndProcessedOpsMaybe() {
+		if (this.latestSnapshot === undefined) {
+			return;
+		}
+		const snapshotSequenceNumber = this.latestSnapshot?.snapshotSequenceNumber;
+		if (this.processedOps.length === 0) {
+			// can't refresh latest snapshot until we have processed the ops up to it.
+			// Pending state would be behind the latest snapshot.
+			return;
+		}
+		const firstProcessedOpSequenceNumber = this.processedOps[0].sequenceNumber;
+		const lastProcessedOpSequenceNumber =
+			this.processedOps[this.processedOps.length - 1].sequenceNumber;
+
+		if (snapshotSequenceNumber < firstProcessedOpSequenceNumber) {
+			// Snapshot seq number is older than our first processed op, which could mean we're fetching
+			// the same snapshot that we already have or snapshot is too old, implicating an unexpected behavior.
+			this.mc.logger.sendTelemetryEvent({
+				category:
+					snapshotSequenceNumber < firstProcessedOpSequenceNumber - 1
+						? "error"
+						: "generic",
+				eventName: "OldSnapshotFetchWhileRefreshing",
+				snapshotSequenceNumber,
+				firstProcessedOpSequenceNumber,
+			});
+			this.latestSnapshot = undefined;
+		} else if (snapshotSequenceNumber <= lastProcessedOpSequenceNumber) {
+			// Snapshot seq num is between the first and last processed op.
+			// Remove the ops that are already part of the snapshot
+			this.processedOps.splice(
+				0,
+				snapshotSequenceNumber - firstProcessedOpSequenceNumber + 1,
+			);
+			this.snapshot = this.latestSnapshot;
+			this.latestSnapshot = undefined;
+			this.mc.logger.sendTelemetryEvent({
+				eventName: "SnapshotRefreshed",
+				snapshotSequenceNumber,
+				firstProcessedOpSequenceNumber,
+				newFirstProcessedOpSequenceNumber:
+					this.processedOps.length === 0
+						? undefined
+						: this.processedOps[0].sequenceNumber,
+			});
+		}
+	}
+
 	/**
 	 * This method is only meant to be used by Container.attach() to set the initial
 	 * base snapshot when attaching.
@@ -171,11 +245,84 @@ export class SerializedStateManager {
 	}
 }
 
+/**
+ * Retrieves the most recent snapshot and returns its info.
+ *
+ * @param mc - The monitoring context.
+ * @param storageAdapter - The storage adapter providing methods to retrieve the snapshot.
+ * @param supportGetSnapshotApi - a boolean indicating whether to use the fetchISnapshot or fetchISnapshotTree.
+ * @returns a SnapshotInfo object containing the snapshot tree, snapshot blobs and its sequence number.
+ */
+export async function getLatestSnapshotInfo(
+	mc: MonitoringContext,
+	storageAdapter: Pick<
+		IDocumentStorageService,
+		"getSnapshot" | "getSnapshotTree" | "getVersions" | "readBlob"
+	>,
+	supportGetSnapshotApi: boolean,
+): Promise<SnapshotInfo | undefined> {
+	return PerformanceEvent.timedExecAsync(
+		mc.logger,
+		{ eventName: "GetLatestSnapshotInfo" },
+		async () => {
+			const { baseSnapshot } = await getSnapshotTree(
+				mc,
+				storageAdapter,
+				supportGetSnapshotApi,
+				undefined,
+			);
+			const snapshotBlobs = await getBlobContentsFromTree(baseSnapshot, storageAdapter);
+			const attributes: IDocumentAttributes = await getDocumentAttributes(
+				storageAdapter,
+				baseSnapshot,
+			);
+			const snapshotSequenceNumber = attributes.sequenceNumber;
+			return { baseSnapshot, snapshotBlobs, snapshotSequenceNumber };
+		},
+	).catch(() => undefined);
+}
+
+/**
+ * Retrieves a snapshot from the storage adapter and transforms it into an ISnapshotTree object.
+ *
+ * @param mc - The monitoring context.
+ * @param storageAdapter - The storage adapter providing methods to retrieve the snapshot.
+ * @param supportGetSnapshotApi - a boolean indicating whether to use the fetchISnapshot or fetchISnapshotTree.
+ * @param specifiedVersion - An optional version string specifying the version of the snapshot tree to fetch.
+ * @returns - An ISnapshotTree and its version.
+ */
+async function getSnapshotTree(
+	mc: MonitoringContext,
+	storageAdapter: Pick<
+		IDocumentStorageService,
+		"getSnapshot" | "getSnapshotTree" | "getVersions"
+	>,
+	supportGetSnapshotApi: boolean,
+	specifiedVersion: string | undefined,
+): Promise<{ baseSnapshot: ISnapshotTree; version?: IVersion }> {
+	const { snapshot, version } = supportGetSnapshotApi
+		? await fetchISnapshot(mc, storageAdapter, specifiedVersion)
+		: await fetchISnapshotTree(mc, storageAdapter, specifiedVersion);
+	const baseSnapshot: ISnapshotTree | undefined = isInstanceOfISnapshot(snapshot)
+		? snapshot.snapshotTree
+		: snapshot;
+	assert(baseSnapshot !== undefined, 0x8e4 /* Snapshot should exist */);
+	return { baseSnapshot, version };
+}
+
+/**
+ * Fetches an ISnapshot from a storage adapter based on the specified version.
+ *
+ * @param mc - The monitoring context.
+ * @param storageAdapter - The storage adapter providing a getSnapshot method to retrieve the ISnapshot and version.
+ * @param specifiedVersion - An optional version string specifying the version of the snapshot tree to fetch.
+ * @returns - The fetched snapshot tree and its version.
+ */
 export async function fetchISnapshot(
 	mc: MonitoringContext,
 	storageAdapter: Pick<IDocumentStorageService, "getSnapshot">,
 	specifiedVersion: string | undefined,
-): Promise<{ snapshot?: ISnapshot | ISnapshotTree; version?: IVersion }> {
+): Promise<{ snapshot?: ISnapshot; version?: IVersion }> {
 	const snapshot = await storageAdapter.getSnapshot?.({ versionId: specifiedVersion });
 	const version: IVersion | undefined =
 		snapshot?.snapshotTree.id === undefined
@@ -190,7 +337,6 @@ export async function fetchISnapshot(
 			eventName: "getSnapshotTreeFailed",
 			id: specifiedVersion,
 		});
-		// Not sure if this should be here actually
 	} else if (snapshot !== undefined && version?.id === undefined) {
 		mc.logger.sendErrorEvent({
 			eventName: "getSnapshotFetchedTreeWithoutVersionId",
@@ -201,9 +347,12 @@ export async function fetchISnapshot(
 }
 
 /**
- * Get the most recent snapshot, or a specific version.
- * @param specifiedVersion - The specific version of the snapshot to retrieve
- * @returns The snapshot requested, or the latest snapshot if no version was specified, plus version ID
+ * Fetches an ISnapshotTree from a storage adapter based on the specified version.
+ *
+ * @param mc - The monitoring context.
+ * @param storageAdapter - The storage adapter providing methods to retrieve the ISnapshotTree and version.
+ * @param specifiedVersion - An optional version string specifying the version of the snapshot tree to fetch.
+ * @returns - The fetched snapshot tree and its version.
  */
 export async function fetchISnapshotTree(
 	mc: MonitoringContext,