Akka stream flow recover from `failStage` - scala

Question

I have GraphStage taken from https://stackoverflow.com/a/40962834/772249 that looks like this. It's working as WebSocket server:
class TerminateFlowStage[T](
predicate: T => Boolean,
forwardTerminatingMessage: Boolean = false,
terminate: Boolean = true)
extends GraphStage[FlowShape[T, T]]
{
val in = Inlet[T]("TerminateFlowStage.in")
val out = Outlet[T]("TerminateFlowStage.out")
override val shape = FlowShape.of(in, out)
override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
new GraphStageLogic(shape) {
setHandlers(in, out, new InHandler with OutHandler {
override def onPull(): Unit = { pull(in) }
override def onPush(): Unit = {
val chunk = grab(in)
if (predicate(chunk)) {
if (forwardTerminatingMessage) {
push(out, chunk)
}
if (terminate) {
failStage(new RuntimeException("Flow terminated by TerminateFlowStage"))
} else {
completeStage()
}
} else {
push(out, chunk)
}
}
})
}
}
val termOnKillMe = new TerminateFlowStage[Message]((chunk: Message) => chunk match {
case TextMessage.Strict(text) => text.toInt > 5
case _ => false
})
val route =
path("") {
get {
extractUpgradeToWebSocket {
upgrade =>
complete(upgrade.handleMessagesWithSinkSource(
Sink.ignore,
Source(1 to 10).
map(i => TextMessage(i.toString)).throttle(1, 1.second, 1, ThrottleMode.shaping).via(termOnKillMe)
))
}
}
}
So, after 5 messages, WebSocket server drops connection.
I have this flow for WebSocket client:
val flow: Flow[Message, Message, Future[Seq[Message]]] =
Flow.fromSinkAndSourceMat(
Sink.seq[Message],
Source.maybe[Message])(Keep.left)
val (upgradeResponse, promise) =
Http().singleWebSocketRequest(
WebSocketRequest("ws://localhost:8080/"),
flow.recoverWithRetries(1, {
case _ => Source.empty
})
)
The problem is, that WebSocket client flow cannot recover from an exception raised by TerminateFlowStage. Getting Future(Failure(akka.http.scaladsl.model.ws.PeerClosedConnectionException: Peer closed connection with code 1011 'internal error'))
Without exception raised, everything works great.