Fix exitStagingMode for rehydrated containers by using loadContainerRuntimeAlpha

- Use loadContainerRuntimeAlpha which returns stageControls when loading from pending state
- Pass stageControls to DefaultStressDataObject via setStageControls method
- Restore the assert in exitStagingMode since we now have proper stage controls
- Ensures rehydrated containers can properly exit staging mode

Co-Authored-By: anthony-murphy <anthony.murphy@microsoft.com>
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: 3 people

packages/test/local-server-stress-tests/src/stressDataObject.ts

@@ -11,7 +11,7 @@ import {
 } from "@fluidframework/container-definitions/internal";
 import {
 	ContainerRuntime,
-	loadContainerRuntime,
+	loadContainerRuntimeAlpha,
 	type IContainerRuntimeOptionsInternal,
 } from "@fluidframework/container-runtime/internal";
 // eslint-disable-next-line import-x/no-deprecated
@@ -301,6 +301,15 @@ export class DefaultStressDataObject extends StressDataObject {
 
 	private stageControls: StageControlsAlpha | undefined;
 	private readonly containerRuntimeExp = asLegacyAlpha(this.context.containerRuntime);
+
+	/**
+	 * Sets the stage controls. This is called by the runtime factory when loading
+	 * from pending state that was in staging mode.
+	 */
+	public setStageControls(controls: StageControlsAlpha | undefined): void {
+		this.stageControls = controls;
+	}
+
 	public enterStagingMode(): void {
 		assert(
 			this.containerRuntimeExp.enterStagingMode !== undefined,
@@ -314,12 +323,7 @@ export class DefaultStressDataObject extends StressDataObject {
 	}
 
 	public exitStagingMode(commit: boolean): void {
-		// If stageControls is undefined but we're in staging mode, this means
-		// the container was rehydrated from pending state while in staging mode.
-		// We can't exit staging mode from here - skip gracefully.
-		if (this.stageControls === undefined) {
-			return;
-		}
+		assert(this.stageControls !== undefined, "must have staging mode controls");
 		if (commit) {
 			this.stageControls.commitChanges();
 		} else {
@@ -355,7 +359,11 @@ export const createRuntimeFactory = (): IRuntimeFactory => {
 			return this;
 		},
 		instantiateRuntime: async (context, existing) => {
-			const runtime = await loadContainerRuntime({
+			// Capture stageControls to pass to the entrypoint after loading.
+			// This must be inside instantiateRuntime so each container instance has its own variable.
+			let pendingStageControls: StageControlsAlpha | undefined;
+
+			const { runtime, stageControls } = await loadContainerRuntimeAlpha({
 				context,
 				existing,
 				runtimeOptions,
@@ -372,9 +380,22 @@ export const createRuntimeFactory = (): IRuntimeFactory => {
 					);
 					assert(aliasedDefault !== undefined, "default must exist");
 
-					return aliasedDefault.get();
+					const entryPoint = await aliasedDefault.get();
+
+					// Pass the stageControls (if any) to the DefaultStressDataObject
+					// so it can properly exit staging mode when rehydrated from pending state
+					const maybe: FluidObject<DefaultStressDataObject> | undefined = entryPoint;
+					if (maybe?.DefaultStressDataObject !== undefined && pendingStageControls !== undefined) {
+						maybe.DefaultStressDataObject.setStageControls(pendingStageControls);
+					}
+
+					return entryPoint;
 				},
 			});
+
+			// Store stageControls so it can be passed to entrypoint when getEntryPoint() is called
+			pendingStageControls = stageControls;
+
 			// id compressor isn't made available via the interface right now.
 			// We could revisit exposing the safe part of its API (IIdCompressor, not IIdCompressorCore) in a way
 			// that would avoid this instanceof check, but most customers shouldn't really have a need for it.