Report compose session availability from onRemembered instead of a DisposableEffect [CLF-648]

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/compose/ComposeRenderContext.kt

@@ -69,6 +69,7 @@ private constructor(
   override val renderKey: String,
   private val saveableStateRegistry: SaveableStateRegistry?,
   private val saveableStateKey: String,
+  private val onSessionAvailable: ((WorkflowSession) -> Unit)?,
 ) :
   BaseRenderContext<P, Any?, O>,
   Sink<WorkflowAction<P, Any?, O>>,
@@ -86,11 +87,13 @@ private constructor(
     config: WorkflowComposableRuntimeConfig,
     parent: WorkflowSession?,
     renderKey: String,
+    onSessionAvailable: ((WorkflowSession) -> Unit)?,
   ) : this(
     snapshot =
       saveableStateRegistry?.let { restoreFromRegistry(saveableStateRegistry, saveableStateKey) },
     saveableStateRegistry = saveableStateRegistry,
     saveableStateKey = saveableStateKey,
+    onSessionAvailable = onSessionAvailable,
     workflow = workflow,
     initialProps = initialProps,
     workflowScope = workflowScope,
@@ -160,6 +163,8 @@ private constructor(
   }
 
   override fun onRemembered() {
+    onSessionAvailable?.invoke(this)
+
     if (saveableStateRegistry != null) {
       // Saving and restoring from the registry manually, instead of using rememberSaveable, saves
       // a lot on allocations and is faster too (see
@@ -341,6 +346,7 @@ private constructor(
       parentSession: WorkflowSession?,
       renderKey: String,
       callerRecomposeScope: RecomposeScope,
+      onSessionAvailable: ((WorkflowSession) -> Unit)?,
     ): ComposeRenderContext<P, O, R> {
       val workflowScope = rememberCoroutineScope()
       val saveableStateRegistry = LocalSaveableStateRegistry.current
@@ -358,6 +364,7 @@ private constructor(
           renderKey = renderKey,
           saveableStateRegistry = saveableStateRegistry,
           saveableStateKey = stringKey,
+          onSessionAvailable = onSessionAvailable,
         )
       }
     }

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/internal/compose/RenderWorkflow.kt

@@ -2,7 +2,6 @@ package com.squareup.workflow1.internal.compose
 
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Composer
-import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.InternalComposeApi
 import androidx.compose.runtime.RecomposeScope
 import androidx.compose.runtime.currentComposer
@@ -39,8 +38,9 @@ internal val LocalRootRecomposeScope = staticCompositionLocalOf<RecomposeScope>
  * @param renderKey The key passed to the [com.squareup.workflow1.BaseRenderContext.renderChild]
  *   function by the parent workflow. This is only used to construct the child's [WorkflowSession],
  *   and is not used for actual keying. [ComposeRenderContext] does the actual keying.
- * @param onSessionAvailable Invoked once (per callback instance) with the [WorkflowSession] of this
- *   workflow once it's been created.
+ * @param onSessionAvailable Invoked once with the [WorkflowSession] of this workflow when the
+ *   session is first remembered by the composition. Only the callback passed on the first
+ *   composition is used.
  */
 @Suppress("UNCHECKED_CAST")
 @OptIn(WorkflowExperimentalRuntime::class)
@@ -79,15 +79,9 @@ internal fun <PropsT, OutputT, RenderingT> renderWorkflow(
       parentSession = parentSession,
       renderKey = renderKey,
       callerRecomposeScope = recomposeScope,
+      onSessionAvailable = onSessionAvailable,
     )
 
-  if (onSessionAvailable != null) {
-    DisposableEffect(onSessionAvailable, renderContext) {
-      onSessionAvailable(renderContext)
-      onDispose {}
-    }
-  }
-
   // Skip re-rendering when possible, but force recompose when new props or onOutput arrive.
   // We use the skippable+restartable variant so internal state-change invalidations trigger a fresh
   // call to the producer lambda within the same restart group.