Update Running.scala

Author: pjfanning

persistence-typed/src/main/scala/org/apache/pekko/persistence/typed/internal/Running.scala

@@ -120,13 +120,12 @@ private[pekko] object Running {
   /**
    * A marker to indicate to upstream code that an unstash is needed.
    */
-  final class UnstashNeeded(val next: Behavior[InternalProtocol], context: ActorContext[InternalProtocol])
+  final class UnstashNeeded(val next: Behavior[InternalProtocol],
+      unstashFunction: Behavior[InternalProtocol] => Behavior[InternalProtocol],
+      context: ActorContext[InternalProtocol])
       extends AbstractBehavior[InternalProtocol](context) {
 
-    /**
-     * @throws NotImplementedError as this behavior should not receive any messages
-     */
-    override def onMessage(msg: InternalProtocol): Behavior[InternalProtocol] = ???
+    override def onMessage(msg: InternalProtocol): Behavior[InternalProtocol] = unstashFunction(next)
   }
 
   def startReplicationStream[C, E, S](
@@ -283,7 +282,7 @@ private[pekko] object Running {
     def onCommand(state: RunningState[S], cmd: C): Behavior[InternalProtocol] = {
       val effect = setup.commandHandler(state.state, cmd)
       val (next, doUnstash) = applyEffects(cmd, state, effect.asInstanceOf[EffectImpl[E, S]]) // TODO can we avoid the cast?
-      if (doUnstash) new Running.UnstashNeeded(next, setup.context)
+      if (doUnstash) new Running.UnstashNeeded(next, tryUnstashOne, setup.context)
       else next
     }