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 =
-        MutableSharedFlow<String>(
+    private val parsedEvents =
+        MutableSharedFlow<WsEvent>(
             extraBufferCapacity = 512,
             onBufferOverflow = BufferOverflow.DROP_OLDEST,
         )
 
     /** Collect this from ViewModels to receive all parsed [WsEvent]s. */
-    val events: SharedFlow<WsEvent> =
-        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)
+    val events: SharedFlow<WsEvent> = parsedEvents.asSharedFlow()
 
     // ── 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 -> {}
                 }
             }