How to timeout Kotlin Flows if they take too long to emit first value

I have a listener that may never be called. But, if it is called at least once, I'm reasonably sure that it will be called many more times. I'm a fan of Flows, so I've wrapped it in a callbackFlow() builder. To guard against waiting forever, I want to add a time out. I'm trying to build flow operator that will throw TimeOut of some kind if the first element of the flow takes too long to be emitted. Here is what I have.

fun <T> Flow<T>.flowBeforeTimeout(ms: Long): Flow<T> = flow {
    withTimeout(ms){
        emit(first())
    }
    emitAll(this@flowBeforeTimeout.drop(1))
}

And it works a little, these JUnit4 tests pass. There are more passing tests, but I'm omitting them for brevity.

@Test(expected = CancellationException::class)
fun `Throws on timeout`(): Unit = runBlocking {
    val testFlow = flow {
        delay(200)
        emit(1)
    }

    testFlow.flowBeforeTimeout(100).toList()
}

@Test
fun `No distortion`(): Unit = runBlocking {
    val testList = listOf(1,2,3)

    val resultList = testList
        .asFlow()
        .flowBeforeTimeout(100)
        .toList()

    assertThat(testList.size, `is`(resultList.size))
}

However, this test is not passing.

// Fails with: Expected: is <1> but: was <2>
@Test
fun `Starts only once`(): Unit = runBlocking {
    var flowStartCount = 0

    val testFlow = flow {
        flowStartCount++
        emit(1)
        emit(2)
        emit(3)
    }

    testFlow.flowBeforeTimeout(100).toList()
    assertThat(flowStartCount, `is`(1))
}

Is there a way to prevent the flow from restarting between first() and emitAll()?