Preserve buffered speech through the stop tap

app/src/main/java/com/jadenjsj/betterflow/AudioRecorderController.kt

@@ -2,21 +2,27 @@ package com.jadenjsj.betterflow
 
 import android.media.AudioFormat
 import android.media.AudioRecord
+import android.media.AudioTimestamp
 import android.media.MediaRecorder
 import java.io.ByteArrayOutputStream
 import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import java.util.concurrent.atomic.AtomicBoolean
+import kotlin.math.abs
 
 class AudioRecorderController {
     private val recording = AtomicBoolean(false)
     private var audioRecord: AudioRecord? = null
     private var worker: Thread? = null
     private var pcm = ByteArrayOutputStream()
+    @Volatile private var stopping = false
+    @Volatile private var cutoffFrameExclusive = NO_CUTOFF_FRAME
+    @Volatile private var startedAtNanos = 0L
+    @Volatile private var deliveredFrames = 0L
 
     @Synchronized
     fun start(onPcmChunk: ((ByteArray) -> Unit)? = null, chunkBytes: Int = STREAM_CHUNK_BYTES) {
-        if (recording.get()) return
+        if (recording.get() || stopping) return
         require(chunkBytes > 0) { "chunkBytes must be > 0" }
         val minBuffer = AudioRecord.getMinBufferSize(
             SAMPLE_RATE,
@@ -41,8 +47,11 @@ class AudioRecorderController {
         }
         pcm = ByteArrayOutputStream()
         audioRecord = record
+        cutoffFrameExclusive = NO_CUTOFF_FRAME
+        deliveredFrames = 0L
         recording.set(true)
         record.startRecording()
+        startedAtNanos = System.nanoTime()
         worker = Thread(
             { captureLoop(record, chunkBytes, onPcmChunk) },
             "betterflow-audio",
@@ -56,6 +65,7 @@ class AudioRecorderController {
     ) {
         val chunk = ByteArray(chunkBytes)
         var filled = 0
+        var capturedFrames = 0L
         try {
             while (recording.get()) {
                 val count = record.read(
@@ -65,49 +75,124 @@ class AudioRecorderController {
                     AudioRecord.READ_BLOCKING,
                 )
                 if (count <= 0) continue
-                synchronized(this) { pcm.write(chunk, filled, count) }
-                filled += count
+                val cutoff = cutoffFrameExclusive
+                val acceptedCount = if (cutoff == NO_CUTOFF_FRAME) {
+                    count
+                } else {
+                    val framesRemaining = (cutoff - capturedFrames).coerceAtLeast(0L)
+                    minOf(count, framesRemaining.times(SAMPLE_WIDTH_BYTES).coerceAtMost(Int.MAX_VALUE.toLong()).toInt())
+                }
+                if (acceptedCount > 0) {
+                    synchronized(this) { pcm.write(chunk, filled, acceptedCount) }
+                    filled += acceptedCount
+                    capturedFrames += acceptedCount / SAMPLE_WIDTH_BYTES
+                    deliveredFrames = capturedFrames
+                }
                 if (filled == chunk.size) {
                     onPcmChunk?.invoke(chunk.copyOf())
                     filled = 0
                 }
+                if (cutoff != NO_CUTOFF_FRAME && capturedFrames >= cutoff) {
+                    recording.set(false)
+                }
             }
         } finally {
             if (filled > 0) onPcmChunk?.invoke(chunk.copyOf(filled))
         }
     }
 
-    fun stopAndGetPcm(): ByteArray {
-        val record: AudioRecord?
+    /**
+     * Stops at the instant this method is called while optionally draining audio
+     * that Android captured before the tap but has not delivered to our read loop.
+     * Samples captured after the call are truncated using AudioRecord's monotonic
+     * frame timestamp; [drainTimeoutMs] is only how long we wait for old samples.
+     */
+    fun stopAndGetPcm(
+        preservePreTapTail: Boolean = false,
+        drainTimeoutMs: Int = 0,
+        cutoffNanos: Long = System.nanoTime(),
+    ): ByteArray = stopInternal(preservePreTapTail, drainTimeoutMs, cutoffNanos).pcm
+
+    private fun stopInternal(
+        preservePreTapTail: Boolean,
+        drainTimeoutMs: Int,
+        cutoffNanos: Long,
+    ): StopResult {
+        val record: AudioRecord
         val captureWorker: Thread?
         synchronized(this) {
-            if (!recording.getAndSet(false)) return pcm.toByteArray()
-            record = audioRecord
+            if (!recording.get()) return StopResult(pcm.toByteArray(), ownedRecorder = false)
+            record = audioRecord ?: return StopResult(pcm.toByteArray(), ownedRecorder = false)
             captureWorker = worker
+            // Claim this AudioRecord before waiting for the tail. Cancellation or
+            // destruction may call stop again while the first caller is draining.
+            audioRecord = null
+            worker = null
+            stopping = true
+            if (preservePreTapTail && drainTimeoutMs > 0) {
+                cutoffFrameExclusive = estimateCutoffFrame(record, cutoffNanos)
+            } else {
+                recording.set(false)
+            }
         }
 
+        if (preservePreTapTail && drainTimeoutMs > 0) {
+            // The capture loop exits as soon as it has read through the tap frame.
+            // This wait does not extend the transcript beyond that frame.
+            runCatching { captureWorker?.join(drainTimeoutMs.coerceAtLeast(1).toLong()) }
+        }
+        recording.set(false)
         // AudioRecord.stop() unblocks READ_BLOCKING. Do not hold this object's
         // monitor while joining: captureLoop needs the same monitor for its
         // final PCM write before it can exit.
         try {
-            record?.stop()
+            record.stop()
         } catch (_: Throwable) {
         }
-        captureWorker?.join(1200)
-        record?.release()
+        runCatching { captureWorker?.join(1200) }
+        runCatching { record.release() }
 
         synchronized(this) {
-            if (audioRecord === record) audioRecord = null
-            if (worker === captureWorker) worker = null
-            return pcm.toByteArray()
+            stopping = false
+            return StopResult(pcm.toByteArray(), ownedRecorder = true)
+        }
+    }
+
+    private data class StopResult(val pcm: ByteArray, val ownedRecorder: Boolean)
+
+    private fun estimateCutoffFrame(record: AudioRecord, cutoffNanos: Long): Long {
+        val elapsedCutoffFrame = nanosToFrames((cutoffNanos - startedAtNanos).coerceAtLeast(0L))
+        val timestamp = AudioTimestamp()
+        val status = runCatching {
+            record.getTimestamp(timestamp, AudioTimestamp.TIMEBASE_MONOTONIC)
+        }.getOrDefault(AudioRecord.ERROR_INVALID_OPERATION)
+        if (status == AudioRecord.SUCCESS && timestamp.nanoTime > 0L) {
+            val elapsedNanos = (cutoffNanos - timestamp.nanoTime).coerceAtLeast(0L)
+            val timestampCutoffFrame = timestamp.framePosition + nanosToFrames(elapsedNanos)
+            // Reject devices whose frame counter does not reset for this AudioRecord.
+            // Already-delivered frames were necessarily captured before the tap.
+            if (abs(timestampCutoffFrame - elapsedCutoffFrame) <= MAX_TIMESTAMP_SKEW_FRAMES) {
+                return timestampCutoffFrame.coerceAtLeast(deliveredFrames)
+            }
         }
+
+        // Some devices do not expose AudioRecord timestamps. Frame zero begins at
+        // startRecording(), so elapsed monotonic time is a safe cutoff fallback.
+        return elapsedCutoffFrame.coerceAtLeast(deliveredFrames)
     }
 
+    private fun nanosToFrames(nanos: Long): Long =
+        ((nanos.toDouble() * SAMPLE_RATE.toDouble()) / NANOS_PER_SECOND).toLong()
+
     fun stopAndGetWav(): ByteArray = pcmToWav(stopAndGetPcm())
 
     fun stopAndDiscard() {
-        stopAndGetPcm()
-        synchronized(this) { pcm = ByteArrayOutputStream() }
+        val result = stopInternal(
+            preservePreTapTail = false,
+            drainTimeoutMs = 0,
+            cutoffNanos = System.nanoTime(),
+        )
+        if (result.ownedRecorder) synchronized(this) { pcm = ByteArrayOutputStream() }
     }
 
     fun isRecording(): Boolean = recording.get()
@@ -118,6 +203,10 @@ class AudioRecorderController {
         const val SAMPLE_WIDTH_BYTES = 2
         const val STREAM_CHUNK_MS = 100
         const val STREAM_CHUNK_BYTES = SAMPLE_RATE * CHANNELS * SAMPLE_WIDTH_BYTES * STREAM_CHUNK_MS / 1000
+        private const val NO_CUTOFF_FRAME = Long.MAX_VALUE
+        private const val NANOS_PER_SECOND = 1_000_000_000.0
+        // One 200 ms capture buffer is enough to absorb normal timestamp jitter.
+        private const val MAX_TIMESTAMP_SKEW_FRAMES = SAMPLE_RATE / 5L
 
         fun pcmToWav(raw: ByteArray): ByteArray {
             val evenSize = raw.size - (raw.size % SAMPLE_WIDTH_BYTES)

app/src/main/java/com/jadenjsj/betterflow/MainActivity.kt

@@ -80,6 +80,8 @@ private fun SettingsScreen() {
     var bubbleOpacityPercent by remember { mutableStateOf(Prefs.bubbleOpacityPercent(context)) }
     var notificationPriority by remember { mutableStateOf(Prefs.notificationPriority(context)) }
     var legacyTranscription by remember { mutableStateOf(Prefs.legacyTranscription(context)) }
+    var preservePreTapAudio by remember { mutableStateOf(Prefs.preservePreTapAudio(context)) }
+    var audioDrainTimeoutMs by remember { mutableStateOf(Prefs.audioDrainTimeoutMs(context)) }
     var streamingKeyConfigured by remember { mutableStateOf(Prefs.streamingApiKey(context) != null || BuildConfig.WISPR_BASETEN_API_KEY.isNotBlank()) }
     var streamingApiKey by remember { mutableStateOf("") }
     var email by remember { mutableStateOf(session?.email.orEmpty()) }
@@ -121,6 +123,43 @@ private fun SettingsScreen() {
                         },
                     )
                 }
+
+                Row(
+                    modifier = Modifier.fillMaxWidth(),
+                    horizontalArrangement = Arrangement.SpaceBetween,
+                ) {
+                    Column(Modifier.weight(1f)) {
+                        Text("Preserve words before stop tap")
+                        Text(
+                            "Drain Android's pending microphone buffer, but discard every sample captured after you tap stop.",
+                            style = MaterialTheme.typography.bodySmall,
+                        )
+                    }
+                    Switch(
+                        checked = preservePreTapAudio,
+                        onCheckedChange = { enabled ->
+                            preservePreTapAudio = enabled
+                            Prefs.setPreservePreTapAudio(context, enabled)
+                            status = if (enabled) "Pre-tap audio preservation enabled" else "Microphone will stop immediately"
+                        },
+                    )
+                }
+
+                Text("Pre-tap buffer drain timeout: $audioDrainTimeoutMs ms")
+                Slider(
+                    value = audioDrainTimeoutMs.toFloat(),
+                    onValueChange = { audioDrainTimeoutMs = (it / 50f).roundToInt() * 50 },
+                    onValueChangeFinished = {
+                        Prefs.setAudioDrainTimeoutMs(context, audioDrainTimeoutMs)
+                        status = "Pre-tap buffer drain timeout set to $audioDrainTimeoutMs ms"
+                    },
+                    valueRange = Prefs.MIN_AUDIO_DRAIN_TIMEOUT_MS.toFloat()..Prefs.MAX_AUDIO_DRAIN_TIMEOUT_MS.toFloat(),
+                    enabled = preservePreTapAudio,
+                )
+                Text(
+                    "This is a maximum drain wait, not extra recorded time. The transcript cutoff remains the stop-tap timestamp.",
+                    style = MaterialTheme.typography.bodySmall,
+                )
                 Row(
                     modifier = Modifier.fillMaxWidth(),
                     horizontalArrangement = Arrangement.SpaceBetween,

app/src/main/java/com/jadenjsj/betterflow/OverlayService.kt

@@ -58,6 +58,7 @@ class OverlayService : Service() {
     private var legacyForCurrentRecording = false
     private var currentPcm = ByteArray(0)
     private var processingJob: Job? = null
+    private var captureFinalizeJob: Job? = null
     private var streamWorker: Job? = null
     @Volatile private var activeStream: WisprStreamingClient.Session? = null
     @Volatile private var streamingFailure: Throwable? = null
@@ -106,6 +107,7 @@ class OverlayService : Service() {
         streamStopRequested = true
         activeStream?.cancel("betterFlow service destroyed")
         streamWorker?.cancel()
+        captureFinalizeJob?.cancel()
         processingJob?.cancel()
         wispr.cancelActiveTranscription()
         if (recorder.isRecording()) runCatching { recorder.stopAndDiscard() }
@@ -493,7 +495,11 @@ class OverlayService : Service() {
                 streamQueueEnabled.set(false)
                 Log.w(TAG, "Wispr streaming failed: ${t.message}", t)
                 withContext(Dispatchers.Main.immediate) {
-                    if (generation == operationGeneration && state == BubbleState.PROCESSING) {
+                    if (
+                        generation == operationGeneration &&
+                        state == BubbleState.PROCESSING &&
+                        streamCaptureFinalized
+                    ) {
                         startLegacyFallback(currentPcm, generation, "streaming failed: ${t.message}")
                     }
                 }
@@ -522,30 +528,45 @@ class OverlayService : Service() {
 
     private fun stopRecording() {
         if (state != BubbleState.RECORDING) return
+        val cutoffNanos = System.nanoTime()
+        val generation = operationGeneration
+        val preservePreTapAudio = Prefs.preservePreTapAudio(this)
+        val drainTimeoutMs = Prefs.audioDrainTimeoutMs(this)
         updateState(BubbleState.PROCESSING)
         streamQueueEnabled.set(false)
         streamStopRequested = true
-        currentPcm = recorder.stopAndGetPcm()
-        streamCaptureFinalized = true
-
-        if (currentPcm.isEmpty()) {
-            cancelStreamOnly("empty recording")
-            updateState(BubbleState.IDLE)
-            return
-        }
+        captureFinalizeJob?.cancel()
+        captureFinalizeJob = scope.launch(Dispatchers.IO) {
+            val captured = recorder.stopAndGetPcm(
+                preservePreTapTail = preservePreTapAudio,
+                drainTimeoutMs = drainTimeoutMs,
+                cutoffNanos = cutoffNanos,
+            )
+            withContext(Dispatchers.Main.immediate) {
+                if (generation != operationGeneration || state != BubbleState.PROCESSING) return@withContext
+                currentPcm = captured
+                streamCaptureFinalized = true
+                captureFinalizeJob = null
+
+                if (currentPcm.isEmpty()) {
+                    cancelStreamOnly("empty recording")
+                    updateState(BubbleState.IDLE)
+                    return@withContext
+                }
 
-        if (legacyForCurrentRecording) {
-            startLegacyFallback(currentPcm, operationGeneration, "legacy mode")
-            return
-        }
+                if (legacyForCurrentRecording) {
+                    startLegacyFallback(currentPcm, generation, "legacy mode")
+                    return@withContext
+                }
 
-        streamingFailure?.let { failure ->
-            cancelStreamOnly("stream already failed")
-            startLegacyFallback(currentPcm, operationGeneration, "streaming failed: ${failure.message}")
-            return
+                streamingFailure?.let { failure ->
+                    cancelStreamOnly("stream already failed")
+                    startLegacyFallback(currentPcm, generation, "streaming failed: ${failure.message}")
+                }
+                // Otherwise the stream worker sends the retained PCM suffix and
+                // commits after capture through the stop-tap timestamp is complete.
+            }
         }
-        // The stream worker drains all queued PCM, sends the separate commit frame,
-        // waits for the final response, then calls handleStreamingResult().
     }
 
     private fun handleStreamingResult(result: WisprStreamingClient.Result, generation: Long) {
@@ -629,6 +650,8 @@ class OverlayService : Service() {
         activeStream = null
         streamWorker?.cancel()
         streamWorker = null
+        captureFinalizeJob?.cancel()
+        captureFinalizeJob = null
         processingJob?.cancel()
         processingJob = null
         val cancelledHttp = wispr.cancelActiveTranscription()
@@ -654,6 +677,7 @@ class OverlayService : Service() {
     private fun finishOperation(generation: Long) {
         if (generation != operationGeneration) return
         cancelStreamOnly("operation complete")
+        captureFinalizeJob = null
         processingJob = null
         streamingFailure = null
         currentPcm = ByteArray(0)