Add onSessionCancelled API

workflow-runtime/api/workflow-runtime.api

@@ -5,6 +5,7 @@ public final class com/squareup/workflow1/NoopWorkflowInterceptor : com/squareup
 	public fun onRender (Ljava/lang/Object;Ljava/lang/Object;Lcom/squareup/workflow1/BaseRenderContext;Lkotlin/jvm/functions/Function3;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Ljava/lang/Object;
 	public fun onRenderAndSnapshot (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/RenderingAndSnapshot;
 	public fun onRuntimeUpdate (Lcom/squareup/workflow1/WorkflowInterceptor$RuntimeUpdate;)V
+	public fun onSessionCancelled (Ljava/util/concurrent/CancellationException;Ljava/util/List;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public fun onSessionStarted (Lkotlinx/coroutines/CoroutineScope;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public fun onSnapshotState (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/Snapshot;
 	public fun onSnapshotStateWithChildren (Lkotlin/jvm/functions/Function0;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/TreeSnapshot;
@@ -34,6 +35,7 @@ public class com/squareup/workflow1/SimpleLoggingWorkflowInterceptor : com/squar
 	public fun onRender (Ljava/lang/Object;Ljava/lang/Object;Lcom/squareup/workflow1/BaseRenderContext;Lkotlin/jvm/functions/Function3;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Ljava/lang/Object;
 	public fun onRenderAndSnapshot (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/RenderingAndSnapshot;
 	public fun onRuntimeUpdate (Lcom/squareup/workflow1/WorkflowInterceptor$RuntimeUpdate;)V
+	public fun onSessionCancelled (Ljava/util/concurrent/CancellationException;Ljava/util/List;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public fun onSessionStarted (Lkotlinx/coroutines/CoroutineScope;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public fun onSnapshotState (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/Snapshot;
 	public fun onSnapshotStateWithChildren (Lkotlin/jvm/functions/Function0;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/TreeSnapshot;
@@ -58,6 +60,7 @@ public abstract interface class com/squareup/workflow1/WorkflowInterceptor {
 	public abstract fun onRender (Ljava/lang/Object;Ljava/lang/Object;Lcom/squareup/workflow1/BaseRenderContext;Lkotlin/jvm/functions/Function3;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Ljava/lang/Object;
 	public abstract fun onRenderAndSnapshot (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/RenderingAndSnapshot;
 	public abstract fun onRuntimeUpdate (Lcom/squareup/workflow1/WorkflowInterceptor$RuntimeUpdate;)V
+	public abstract fun onSessionCancelled (Ljava/util/concurrent/CancellationException;Ljava/util/List;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public abstract fun onSessionStarted (Lkotlinx/coroutines/CoroutineScope;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public abstract fun onSnapshotState (Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/Snapshot;
 	public abstract fun onSnapshotStateWithChildren (Lkotlin/jvm/functions/Function0;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/TreeSnapshot;
@@ -69,6 +72,7 @@ public final class com/squareup/workflow1/WorkflowInterceptor$DefaultImpls {
 	public static fun onRender (Lcom/squareup/workflow1/WorkflowInterceptor;Ljava/lang/Object;Ljava/lang/Object;Lcom/squareup/workflow1/BaseRenderContext;Lkotlin/jvm/functions/Function3;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Ljava/lang/Object;
 	public static fun onRenderAndSnapshot (Lcom/squareup/workflow1/WorkflowInterceptor;Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/RenderingAndSnapshot;
 	public static fun onRuntimeUpdate (Lcom/squareup/workflow1/WorkflowInterceptor;Lcom/squareup/workflow1/WorkflowInterceptor$RuntimeUpdate;)V
+	public static fun onSessionCancelled (Lcom/squareup/workflow1/WorkflowInterceptor;Ljava/util/concurrent/CancellationException;Ljava/util/List;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public static fun onSessionStarted (Lcom/squareup/workflow1/WorkflowInterceptor;Lkotlinx/coroutines/CoroutineScope;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)V
 	public static fun onSnapshotState (Lcom/squareup/workflow1/WorkflowInterceptor;Ljava/lang/Object;Lkotlin/jvm/functions/Function1;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/Snapshot;
 	public static fun onSnapshotStateWithChildren (Lcom/squareup/workflow1/WorkflowInterceptor;Lkotlin/jvm/functions/Function0;Lcom/squareup/workflow1/WorkflowInterceptor$WorkflowSession;)Lcom/squareup/workflow1/TreeSnapshot;

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/SimpleLoggingWorkflowInterceptor.kt

@@ -2,8 +2,8 @@ package com.squareup.workflow1
 
 import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Job
 
 /**
  * A [WorkflowInterceptor] that just prints all method calls using [log].
@@ -14,9 +14,14 @@ public open class SimpleLoggingWorkflowInterceptor : WorkflowInterceptor {
     session: WorkflowSession
   ) {
     invokeSafely("logBeforeMethod") { logBeforeMethod("onInstanceStarted", session) }
-    workflowScope.coroutineContext[Job]!!.invokeOnCompletion {
-      invokeSafely("logAfterMethod") { logAfterMethod("onInstanceStarted", session) }
-    }
+  }
+
+  override fun <P, S, O> onSessionCancelled(
+    cause: CancellationException?,
+    droppedActions: List<WorkflowAction<P, S, O>>,
+    session: WorkflowSession
+  ) {
+    invokeSafely("logAfterMethod") { logAfterMethod("onInstanceStarted", session) }
   }
 
   override fun <P, S> onInitialState(

workflow-runtime/src/commonMain/kotlin/com/squareup/workflow1/WorkflowInterceptor.kt

@@ -2,6 +2,7 @@ package com.squareup.workflow1
 
 import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlin.coroutines.CoroutineContext
@@ -75,6 +76,19 @@ public interface WorkflowInterceptor {
     session: WorkflowSession
   ): Unit = Unit
 
+  /**
+   * Called when the session is ending, when the Workflow's [CoroutineScope] is being cancelled.
+   *
+   * @param cause The cause of the cancellation if non-null.
+   * @param droppedActions Any actions that were queued in this node's channel at the time of
+   *  cancellation.
+   */
+  public fun <P, S, O> onSessionCancelled(
+    cause: CancellationException?,
+    droppedActions: List<WorkflowAction<P, S, O>>,
+    session: WorkflowSession
+  ): Unit = Unit
+
   /**
    * Intercepts calls to [StatefulWorkflow.initialState].
    */

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

@@ -11,6 +11,7 @@ import com.squareup.workflow1.WorkflowInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.RenderContextInterceptor
 import com.squareup.workflow1.WorkflowInterceptor.RuntimeUpdate
 import com.squareup.workflow1.WorkflowInterceptor.WorkflowSession
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineScope
 import kotlin.reflect.KType
 
@@ -32,6 +33,20 @@ internal class ChainedWorkflowInterceptor(
     interceptors.forEach { it.onSessionStarted(workflowScope, session) }
   }
 
+  override fun <P, S, O> onSessionCancelled(
+    cause: CancellationException?,
+    droppedActions: List<WorkflowAction<P, S, O>>,
+    session: WorkflowSession
+  ) {
+    interceptors.forEach {
+      it.onSessionCancelled(
+        cause = cause,
+        droppedActions = droppedActions,
+        session = session
+      )
+    }
+  }
+
   override fun <P, S> onInitialState(
     props: P,
     snapshot: Snapshot?,

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

@@ -265,10 +265,22 @@ internal class WorkflowNode<PropsT, StateT, OutputT, RenderingT>(
   /**
    * Cancels this state machine host, and any coroutines started as children of it.
    *
-   * This must be called when the caller will no longer call [registerTreeActionSelectors]. It is an error to call [registerTreeActionSelectors]
-   * after calling this method.
+   * This must be called when the caller will no longer call [registerTreeActionSelectors].
+   * It is an error to call [registerTreeActionSelectors] after calling this method.
    */
   fun cancel(cause: CancellationException? = null) {
+    val hangingActions = mutableListOf<WorkflowAction<PropsT, StateT, OutputT>>()
+    // This will only be non-null if there is an action buffered and ready.
+    var nextAction = eventActionsChannel.tryReceive().getOrNull()
+    while (nextAction != null) {
+      hangingActions.add(nextAction)
+      nextAction = eventActionsChannel.tryReceive().getOrNull()
+    }
+    interceptor.onSessionCancelled(
+      cause = cause,
+      droppedActions = hangingActions,
+      session = this
+    )
     coroutineContext.cancel(cause)
     lastRendering = NullableInitBox()
   }

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/SimpleLoggingWorkflowInterceptorTest.kt

@@ -19,7 +19,21 @@ internal class SimpleLoggingWorkflowInterceptorTest {
     interceptor.onSessionStarted(scope, TestWorkflowSession)
     scope.cancel()
 
-    assertEquals(ErrorLoggingInterceptor.EXPECTED_ERRORS, interceptor.errors)
+    // Only the first, since we don't get cancellation directly from the scope cancellation.
+    // For that we use onSessionCancelled()
+    assertEquals(listOf(ErrorLoggingInterceptor.EXPECTED_ERRORS.first()), interceptor.errors)
+  }
+
+  @Test fun onSessionCancelled_handles_logging_exceptions() {
+    val interceptor = ErrorLoggingInterceptor()
+    interceptor.onSessionCancelled<Unit, Unit, Nothing>(
+      cause = null,
+      droppedActions = emptyList(),
+      session = TestWorkflowSession
+    )
+
+    // Only the second error, since onSessionCancelled only calls logAfterMethod
+    assertEquals(listOf(ErrorLoggingInterceptor.EXPECTED_ERRORS.last()), interceptor.errors)
   }
 
   @Test fun onInitialState_handles_logging_exceptions() {

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/WorkflowInterceptorTest.kt

@@ -175,6 +175,63 @@ internal class WorkflowInterceptorTest {
     intercepted.render("props", "string", RenderContext(fakeContext, workflow))
   }
 
+  @Test fun intercept_intercepts_onSessionCancelled() {
+    val recorder = RecordingWorkflowInterceptor()
+    val session = object : WorkflowSession {
+      override val identifier: WorkflowIdentifier = TestWorkflow.identifier
+      override val renderKey: String = ""
+      override val sessionId: Long = 0
+      override val parent: WorkflowSession? = null
+      override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG
+      override val workflowTracer: WorkflowTracer? = null
+      override val runtimeContext: CoroutineContext = EmptyCoroutineContext
+    }
+
+    recorder.onSessionCancelled<String, String, String>(
+      cause = null,
+      droppedActions = emptyList(),
+      session = session
+    )
+
+    // SimpleLoggingWorkflowInterceptor logs "onInstanceStarted" for onSessionCancelled
+    assertEquals(
+      listOf("END|onInstanceStarted"),
+      recorder.consumeEventNames()
+    )
+  }
+
+  @Test fun intercept_passes_dropped_actions_to_onSessionCancelled() {
+    var capturedDroppedActions: List<WorkflowAction<*, *, *>>? = null
+    val interceptor = object : WorkflowInterceptor {
+      override fun <P, S, O> onSessionCancelled(
+        cause: kotlinx.coroutines.CancellationException?,
+        droppedActions: List<WorkflowAction<P, S, O>>,
+        session: WorkflowSession
+      ) {
+        capturedDroppedActions = droppedActions
+      }
+    }
+    val session = object : WorkflowSession {
+      override val identifier: WorkflowIdentifier = TestWorkflow.identifier
+      override val renderKey: String = ""
+      override val sessionId: Long = 0
+      override val parent: WorkflowSession? = null
+      override val runtimeConfig: RuntimeConfig = RuntimeConfigOptions.DEFAULT_CONFIG
+      override val workflowTracer: WorkflowTracer? = null
+      override val runtimeContext: CoroutineContext = EmptyCoroutineContext
+    }
+    val testAction = action<String, String, String>("TestAction") { state = "modified" }
+
+    interceptor.onSessionCancelled(
+      cause = null,
+      droppedActions = listOf(testAction),
+      session = session
+    )
+
+    assertEquals(1, capturedDroppedActions!!.size)
+    assertEquals("TestAction", capturedDroppedActions!![0].debuggingName)
+  }
+
   private val Workflow<*, *, *>.session: WorkflowSession
     get() = object : WorkflowSession {
       override val identifier: WorkflowIdentifier = this@session.identifier

workflow-runtime/src/commonTest/kotlin/com/squareup/workflow1/internal/ChainedWorkflowInterceptorTest.kt

@@ -70,6 +70,8 @@ internal class ChainedWorkflowInterceptorTest {
         session: WorkflowSession
       ) {
         events += "started1"
+        // We can't use onSessionCancelled because this is completed when the coroutine from
+        // launch() below finishes, so onSessionCancelled is never called by the runtime.
         workflowScope.coroutineContext[Job]!!.invokeOnCompletion {
           events += "cancelled1"
         }