fix(notification): disconnect idle background socket (#858)

* fix(notification): disconnect idle background socket

* fix(notification): clear rejected prompt state

app/src/main/java/com/m57/hermescontrol/MainActivity.kt

@@ -13,6 +13,8 @@ import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
 import com.m57.hermescontrol.data.local.AuthManager
+import com.m57.hermescontrol.data.ws.HermesWsClient
+import com.m57.hermescontrol.notification.NotificationHelper
 import com.m57.hermescontrol.notification.NotificationReplyReceiver
 import com.m57.hermescontrol.theme.HermesControlTheme
 import com.m57.hermescontrol.util.LocaleContextWrapper
@@ -73,6 +75,21 @@ class MainActivity : ComponentActivity() {
         consumeNotificationIntent(intent)
     }
 
+    override fun onStart() {
+        super.onStart()
+        NotificationHelper.setAppForeground(this, true)
+        NotificationHelper.stop(this)
+        if (AuthManager.isGatedMode() || !AuthManager.getToken().isNullOrBlank()) {
+            HermesWsClient.connect()
+        }
+    }
+
+    override fun onStop() {
+        NotificationHelper.setAppForeground(this, false)
+        NotificationHelper.start(this)
+        super.onStop()
+    }
+
     private fun consumeNotificationIntent(intent: Intent?) {
         // Issue #832: MainActivity is exported (launcher requirement) — only
         // honor intents stamped with our own notification action.

app/src/main/java/com/m57/hermescontrol/data/ws/HermesWsClient.kt

@@ -14,20 +14,15 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.channels.BufferOverflow
-import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharedFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.filterIsInstance
-import kotlinx.coroutines.flow.flowOn
-import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.launch
 import kotlinx.serialization.decodeFromString
 import kotlinx.serialization.encodeToString
@@ -91,6 +86,7 @@ object HermesWsClient {
     private val connected = AtomicBoolean(false)
     private val intentionalClose = AtomicBoolean(false)
     private val acceptQueuedMessages = AtomicBoolean(true)
+    private val appInForeground = AtomicBoolean(true)
     private val messageQueue = ConcurrentLinkedQueue<String>()
     private val queuedMessagesById = ConcurrentHashMap<String, String>()
     private val outboundLock = Any()
@@ -153,26 +149,14 @@ object HermesWsClient {
 
     // ── Public observable stream ─────────────────────────────────────────
 
-    private val rawMessages =
+    private val parsedEvents =
-        MutableSharedFlow<String>(
+        MutableSharedFlow<WsEvent>(
             extraBufferCapacity = 512,
             onBufferOverflow = BufferOverflow.DROP_OLDEST,
         )
 
     /** Collect this from ViewModels to receive all parsed [WsEvent]s. */
-    val events: SharedFlow<WsEvent> =
+    val events: SharedFlow<WsEvent> = parsedEvents.asSharedFlow()
-        rawMessages
-            .buffer(Channel.BUFFERED)
-            .map { text ->
-                try {
-                    val rpc = OkHttpProvider.json.decodeFromString<JsonRpcResponse>(text)
-                    EventParser.parse(rpc, text)
-                } catch (e: Exception) {
-                    Log.e(TAG, "Failed to parse message", e)
-                    WsEvent.Unknown(text)
-                }
-            }.flowOn(Dispatchers.Default) // CPU-bound
-            .shareIn(wsScope, SharingStarted.Eagerly)
 
     // ── Connection status flow ──────────────────────────────────────────
     private val _connectionStatus = MutableStateFlow(ConnectionStatus.DISCONNECTED)
@@ -189,6 +173,8 @@ object HermesWsClient {
     var pendingReply: Boolean = false
         private set
 
+    private val pendingPromptSubmits = ConcurrentHashMap.newKeySet<String>()
+
     // ── Credential warning (issue #534) ─────────────────────────────────
     // Backend surfaces `credential_warning` in `gateway.ready` / `session.info`
     // WS payloads (desktop `requestDesktopOnboarding`). Mobile has no equivalent
@@ -233,7 +219,10 @@ object HermesWsClient {
                     is WsEvent.ToolStart,
                     -> pendingReply = true
 
-                    is WsEvent.MessageComplete -> pendingReply = false
+                    is WsEvent.MessageComplete -> {
+                        pendingReply = false
+                        disconnectIfIdleInBackground()
+                    }
                     else -> {}
                 }
             }
@@ -253,6 +242,23 @@ object HermesWsClient {
     @VisibleForTesting
     val isConnected: Boolean get() = connected.get()
 
+    fun setAppForeground(foreground: Boolean) {
+        appInForeground.set(foreground)
+        if (!foreground) disconnectIfIdleInBackground()
+    }
+
+    private fun disconnectIfIdleInBackground() {
+        synchronized(outboundLock) {
+            if (!appInForeground.get() &&
+                !pendingReply &&
+                pendingCalls.isEmpty() &&
+                messageQueue.isEmpty()
+            ) {
+                disconnect()
+            }
+        }
+    }
+
     /** Open a WebSocket connection using settings from [AuthManager]. */
     fun connect() {
         if (networkCollectorStarted.compareAndSet(false, true)) {
@@ -501,6 +507,8 @@ object HermesWsClient {
             if (clearPendingMessages) {
                 messageQueue.clear()
                 queuedMessagesById.clear()
+                pendingPromptSubmits.clear()
+                pendingReply = false
             }
             connected.set(false)
             _connectionStatus.value = ConnectionStatus.DISCONNECTED
@@ -577,6 +585,7 @@ object HermesWsClient {
         } else {
             call.deferred.complete(result)
         }
+        disconnectIfIdleInBackground()
     }
 
     /**
@@ -621,7 +630,12 @@ object HermesWsClient {
             )
         val json = OkHttpProvider.json.encodeToString(request)
         if (BuildConfig.DEBUG) Log.d(TAG, "→ $json")
+        var reconnect = false
         synchronized(outboundLock) {
+            if (method == WsMethods.PROMPT_SUBMIT) {
+                pendingPromptSubmits.add(id)
+                pendingReply = true
+            }
             onSent?.invoke(id)
             val ws = webSocket
             if (ws != null && connected.get()) {
@@ -640,11 +654,13 @@ object HermesWsClient {
                 if (isRetryableMessage(json)) {
                     Log.d(TAG, "WS disconnected — queuing message")
                     queueMessage(id, json)
+                    reconnect = true
                 } else {
                     Log.w(TAG, "WS disconnected with oversized outgoing message — not queueing")
                 }
             }
         }
+        if (reconnect) connect()
         return id
     }
 
@@ -701,10 +717,6 @@ object HermesWsClient {
         text: String,
         onSent: ((String) -> Unit)? = null,
     ): String {
-        // A reply is now pending — even before the first MessageStart event
-        // (the agent may be in its tool phase), so backgrounding the app in
-        // that window still arms the notification service (issue #794).
-        pendingReply = true
         return send(
             method = WsMethods.PROMPT_SUBMIT,
             params = mapOf("session_id" to sessionId, "text" to text),
@@ -848,6 +860,7 @@ object HermesWsClient {
                     markQueuedMessageSent(msg)
                 }
             }
+            disconnectIfIdleInBackground()
         }
 
         override fun onMessage(
@@ -862,25 +875,41 @@ object HermesWsClient {
             val event =
                 try {
                     val rpc = OkHttpProvider.json.decodeFromString<JsonRpcResponse>(text)
-                    EventParser.parse(rpc, text)
+                    val parsed = EventParser.parse(rpc, text)
+                    if (parsed is WsEvent.MessageComplete) {
+                        parsed.copy(
+                            storedSessionId = ActiveSessionHolder.resolveStoredSessionId(parsed.sessionId),
+                        )
+                    } else {
+                        parsed
+                    }
                 } catch (e: Exception) {
                     Log.e(TAG, "Failed to parse message", e)
                     WsEvent.Unknown(text)
                 }
             when (event) {
                 is WsEvent.RpcResult -> {
+                    synchronized(outboundLock) { pendingPromptSubmits.remove(event.id) }
                     removeQueuedMessage(event.id)
                     resolvePending(event.id, event.result, null)
                 }
 
                 is WsEvent.RpcError -> {
+                    synchronized(outboundLock) {
+                        if (pendingPromptSubmits.remove(event.id) &&
+                            pendingPromptSubmits.isEmpty()
+                        ) {
+                            pendingReply = false
+                        }
+                    }
                     removeQueuedMessage(event.id)
                     resolvePending(event.id, null, event.error)
+                    disconnectIfIdleInBackground()
                 }
 
                 else -> Unit
             }
-            val emitted = rawMessages.tryEmit(text)
+            val emitted = parsedEvents.tryEmit(event)
             if (!emitted && BuildConfig.DEBUG) {
                 Log.w(TAG, "WebSocket message dropped due to buffer overflow")
             }

app/src/main/java/com/m57/hermescontrol/data/ws/WsEvent.kt

@@ -57,6 +57,8 @@ sealed class WsEvent {
          * (e.g. a tool call wiped the streaming buffer mid-turn).
          */
         val reasoning: String? = null,
+        /** Stored session id captured before background disconnect clears the active mapping. */
+        val storedSessionId: String? = null,
     ) : WsEvent()
 
     data class MessageDone(

app/src/main/java/com/m57/hermescontrol/notification/ChatNotificationService.kt

@@ -50,7 +50,7 @@ import java.util.concurrent.atomic.AtomicBoolean
  *   service — and its mandatory persistent notification — only exists
  *   while the user is actually waiting for a reply (issue #794).
  * - Stopped by [NotificationHelper.stop] when the app returns to the
- *   foreground (called from ChatScreen's onStart / onResume), or by the
+ *   foreground (called from MainActivity.onStart), or by the
  *   service itself once the pending reply completes in the background
  *   (the reply notification replaces the persistent "waiting" one).
  *
@@ -101,7 +101,8 @@ class ChatNotificationService : Service() {
                                                 .ifBlank { getString(R.string.notif_new_message) }
                                         showReplyNotification(
                                             preview,
-                                            ActiveSessionHolder.resolveStoredSessionId(event.sessionId),
+                                            event.storedSessionId
+                                                ?: ActiveSessionHolder.resolveStoredSessionId(event.sessionId),
                                         )
                                         // The wait is over — retire the foreground
                                         // service. The reply notification above
@@ -239,7 +240,7 @@ class ChatNotificationService : Service() {
         intent: Intent?,
         flags: Int,
         startId: Int,
-    ): Int = START_STICKY
+    ): Int = START_NOT_STICKY
 
     override fun onBind(intent: Intent?): IBinder? = null
 
@@ -278,6 +279,7 @@ object NotificationHelper {
         foreground: Boolean,
     ) {
         ChatNotificationService.setAppForeground(foreground)
+        HermesWsClient.setAppForeground(foreground)
     }
 
     fun isAppInForeground(): Boolean = ChatNotificationService.isAppInForeground()

app/src/main/java/com/m57/hermescontrol/ui/chat/ChatViewModel.kt

@@ -1270,6 +1270,7 @@ class ChatViewModel(
             // in-flight turn via session.redirect instead of queueing a fresh
             // prompt.submit. The backend rewrites the live turn when it can, or
             // queues the correction as the next turn otherwise (issue #710).
+            ActiveSessionHolder.set(agentSessionId, storageSessionId)
             if (wasStreaming && attachments.isEmpty()) {
                 wsClient.sendRedirect(
                     agentSessionId,

app/src/main/java/com/m57/hermescontrol/ui/chat/components/ChatLifecycleEffects.kt

@@ -17,7 +17,6 @@ import androidx.lifecycle.LifecycleEventObserver
 import androidx.lifecycle.compose.LocalLifecycleOwner
 import com.m57.hermescontrol.NavigationController
 import com.m57.hermescontrol.data.ws.ConnectionStatus
-import com.m57.hermescontrol.notification.NotificationHelper
 import com.m57.hermescontrol.ui.chat.ChatMessage
 import com.m57.hermescontrol.ui.chat.ChatViewModel
 import com.m57.hermescontrol.ui.chat.ClarifyUi
@@ -64,23 +63,16 @@ fun ChatLifecycleEffects(
         }
     }
 
-    // Lifecycle observer for notification foreground service
+    // Refresh chat state when the app returns to the foreground.
     DisposableEffect(lifecycleOwner) {
         val observer =
             LifecycleEventObserver { _, event ->
                 when (event) {
                     Lifecycle.Event.ON_START -> {
-                        NotificationHelper.setAppForeground(context, true)
-                        NotificationHelper.stop(context)
                         viewModel.refreshSettings()
                         viewModel.refreshCurrentSession()
                     }
 
-                    Lifecycle.Event.ON_STOP -> {
-                        NotificationHelper.setAppForeground(context, false)
-                        NotificationHelper.start(context)
-                    }
-
                     else -> {}
                 }
             }

app/src/test/java/com/m57/hermescontrol/data/ws/HermesWsClientTest.kt

@@ -8,6 +8,7 @@ import com.m57.hermescontrol.data.remote.DashboardSessionTokenRefresher
 import com.m57.hermescontrol.data.remote.NetworkMonitor
 import com.m57.hermescontrol.data.remote.ServerEndpoint
 import com.m57.hermescontrol.data.remote.buildFakePersistentCookieJar
+import com.m57.hermescontrol.data.session.ActiveSessionHolder
 import io.mockk.every
 import io.mockk.mockk
 import io.mockk.mockkObject
@@ -20,6 +21,7 @@ import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withTimeout
+import okhttp3.Response
 import okhttp3.WebSocket
 import okhttp3.WebSocketListener
 import okhttp3.mockwebserver.MockResponse
@@ -86,6 +88,7 @@ class HermesWsClientTest {
         (queueField.get(HermesWsClient) as java.util.concurrent.ConcurrentLinkedQueue<String>).clear()
 
         HermesWsClient.disconnect(clearPendingMessages = true) // Ensure it starts clean
+        HermesWsClient.setAppForeground(true)
         val acceptQueuedMessagesField = HermesWsClient::class.java.getDeclaredField("acceptQueuedMessages")
         acceptQueuedMessagesField.isAccessible = true
         (acceptQueuedMessagesField.get(HermesWsClient) as java.util.concurrent.atomic.AtomicBoolean).set(true)
@@ -821,6 +824,139 @@ class HermesWsClientTest {
         assertEquals("1", (receivedEvent as WsEvent.RpcResult).id)
     }
 
+    @Test
+    fun testBackgroundWithoutPendingWorkDisconnects() {
+        mockWebServer.enqueue(MockResponse().withWebSocketUpgrade(object : WebSocketListener() {}))
+        HermesWsClient.connect()
+        runBlocking { withTimeout(5000) { HermesWsClient.connectionStatus.first { it == ConnectionStatus.CONNECTED } } }
+
+        HermesWsClient.setAppForeground(false)
+
+        assertFalse(HermesWsClient.isConnected)
+        assertEquals(ConnectionStatus.DISCONNECTED, HermesWsClient.connectionStatus.value)
+    }
+
+    @Test
+    fun testBackgroundWithPendingReplyStaysConnected() {
+        mockWebServer.enqueue(MockResponse().withWebSocketUpgrade(object : WebSocketListener() {}))
+        HermesWsClient.connect()
+        runBlocking { withTimeout(5000) { HermesWsClient.connectionStatus.first { it == ConnectionStatus.CONNECTED } } }
+        HermesWsClient.sendMessage("session-1", "hello")
+
+        HermesWsClient.setAppForeground(false)
+
+        assertTrue(HermesWsClient.isConnected)
+    }
+
+    @Test
+    fun testRejectedPromptDisconnectsIdleBackgroundSocket() {
+        lateinit var serverSocket: WebSocket
+        val connectedLatch = CountDownLatch(1)
+        mockWebServer.enqueue(
+            MockResponse().withWebSocketUpgrade(
+                object : WebSocketListener() {
+                    override fun onOpen(
+                        webSocket: WebSocket,
+                        response: Response,
+                    ) {
+                        serverSocket = webSocket
+                        connectedLatch.countDown()
+                    }
+                },
+            ),
+        )
+        HermesWsClient.connect()
+        assertTrue(connectedLatch.await(5, TimeUnit.SECONDS))
+        val requestId = HermesWsClient.sendMessage("deleted-session", "hello")
+        HermesWsClient.setAppForeground(false)
+
+        assertTrue(HermesWsClient.pendingReply)
+        serverSocket.send(
+            """{"jsonrpc":"2.0","id":"$requestId","error":{"code":4001,"message":"Session not found"}}""",
+        )
+
+        runBlocking {
+            withTimeout(5000) {
+                HermesWsClient.connectionStatus.first { it == ConnectionStatus.DISCONNECTED }
+            }
+        }
+        assertFalse(HermesWsClient.pendingReply)
+        assertFalse(HermesWsClient.isConnected)
+    }
+
+    @Test
+    fun testBackgroundQueuedSendDisconnectsAfterFlush() {
+        val messageLatch = CountDownLatch(1)
+        mockWebServer.enqueue(
+            MockResponse().withWebSocketUpgrade(
+                object : WebSocketListener() {
+                    override fun onMessage(
+                        webSocket: WebSocket,
+                        text: String,
+                    ) {
+                        messageLatch.countDown()
+                    }
+                },
+            ),
+        )
+        HermesWsClient.setAppForeground(false)
+
+        HermesWsClient.send(WsMethods.SESSION_REDIRECT)
+
+        assertTrue(messageLatch.await(5, TimeUnit.SECONDS))
+        runBlocking {
+            withTimeout(5000) {
+                HermesWsClient.connectionStatus.first { it == ConnectionStatus.DISCONNECTED }
+            }
+        }
+        assertFalse(HermesWsClient.isConnected)
+    }
+
+    @Test
+    fun testMessageCompleteDisconnectsIdleBackgroundSocket() {
+        lateinit var serverSocket: WebSocket
+        val connectedLatch = CountDownLatch(1)
+        mockWebServer.enqueue(
+            MockResponse().withWebSocketUpgrade(
+                object : WebSocketListener() {
+                    override fun onOpen(
+                        webSocket: WebSocket,
+                        response: Response,
+                    ) {
+                        serverSocket = webSocket
+                        connectedLatch.countDown()
+                    }
+                },
+            ),
+        )
+        HermesWsClient.connect()
+        assertTrue(connectedLatch.await(5, TimeUnit.SECONDS))
+        ActiveSessionHolder.set("runtime-session", "stored-session")
+        HermesWsClient.sendMessage("session-1", "hello")
+        HermesWsClient.setAppForeground(false)
+
+        val completeEvent =
+            """
+            {"jsonrpc":"2.0","method":"event","params":{"type":"message.complete",
+            "payload":{"text":"done","session_id":"runtime-session"}}}
+            """.trimIndent()
+        val receivedEvent =
+            runBlocking {
+                withTimeout(5000) {
+                    launch { serverSocket.send(completeEvent) }
+                    HermesWsClient.events.first { it is WsEvent.MessageComplete } as WsEvent.MessageComplete
+                }
+            }
+
+        runBlocking {
+            withTimeout(5000) {
+                HermesWsClient.connectionStatus.first { it == ConnectionStatus.DISCONNECTED }
+            }
+        }
+        assertFalse(HermesWsClient.isConnected)
+        assertEquals("stored-session", receivedEvent.storedSessionId)
+    }
+
     @Test
     fun testDisconnect() {
         val serverLatch = CountDownLatch(1)

app/src/test/java/com/m57/hermescontrol/ui/chat/ChatViewModelTest.kt

@@ -1238,6 +1238,7 @@ class ChatViewModelTest {
     fun testSendMessage() =
         runTest {
             val (viewModel, sessionId) = createViewModelWithSession()
+            ActiveSessionHolder.set(sessionId, "stale-session")
 
             viewModel.sendMessage("Hello Hermes")
             advanceUntilIdle()
@@ -1254,6 +1255,7 @@ class ChatViewModelTest {
                     .role,
             )
             assertTrue(viewModel.uiState.value.isAgentTyping)
+            assertEquals(sessionId, ActiveSessionHolder.resolveStoredSessionId(sessionId))
 
             verify { HermesWsClient.sendMessage(sessionId, "Hello Hermes", any()) }
         }