Fix exitStagingMode for rehydrated containers (#16)

* Fix exitStagingMode for rehydrated containers

When a container is rehydrated from pending state while in staging mode,
the DataObject doesn't have stageControls (only set when enterStagingMode
is called). Handle this gracefully by returning early instead of asserting.

Co-Authored-By: Claude Opus 4.5 <[email protected]>
Co-Authored-By: anthony-murphy <[email protected]>

* 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 <[email protected]>
Co-Authored-By: Claude Opus 4.5 <[email protected]>

---------

Co-authored-by: anthony-murphy-agent <[email protected]>
Co-authored-by: Claude Opus 4.5 <[email protected]>
Co-authored-by: anthony-murphy <[email protected]>

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,
@@ -350,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,
@@ -367,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.