diff --git a/app/src/main/java/com/m57/hermescontrol/data/ws/HermesWsClient.kt b/app/src/main/java/com/m57/hermescontrol/data/ws/HermesWsClient.kt index 76aa7c87..ab3fa379 100644 --- a/app/src/main/java/com/m57/hermescontrol/data/ws/HermesWsClient.kt +++ b/app/src/main/java/com/m57/hermescontrol/data/ws/HermesWsClient.kt @@ -117,6 +117,7 @@ object HermesWsClient { // ── Internal state (all access through synchronized / atomic) ──────── private val requestId = AtomicInteger(0) + private val capabilityRequestIds = ConcurrentHashMap.newKeySet() private val connectionGeneration = AtomicInteger(0) private val connected = AtomicBoolean(false) private val intentionalClose = AtomicBoolean(false) @@ -360,6 +361,8 @@ object HermesWsClient { send( WsMethods.CLIENT_CAPABILITIES, mapOf("server_requests" to true), + onSent = { id -> capabilityRequestIds.add(id) }, + queueIfDisconnected = false, ) } val epoch = event.data?.get("replay_epoch") as? String @@ -729,6 +732,7 @@ object HermesWsClient { webSocket?.close(1000, "Client closed") webSocket = null closingSocket = null + capabilityRequestIds.clear() if (clearPendingMessages) { backgroundConnectionLease.set(false) messageQueue.clear() @@ -917,6 +921,13 @@ object HermesWsClient { method: String, params: Map = emptyMap(), onSent: ((String) -> Unit)? = null, + ): String = send(method, params, onSent, true) + + private fun send( + method: String, + params: Map, + onSent: ((String) -> Unit)?, + queueIfDisconnected: Boolean, ): String { val id = requestId.incrementAndGet().toString() val decoratedParams = WsProfileParams.decorate(method, params) @@ -936,13 +947,15 @@ object HermesWsClient { pendingPromptSubmits.add(id) pendingReply = true } - onSent?.invoke(id) val ws = webSocket if (ws != null && connected.get()) { // Never replay send(true): the server may have executed it even if its response is lost. - if (!ws.send(json)) { + if (ws.send(json)) { + onSent?.invoke(id) + } else { if (webSocket !== ws || !acceptQueuedMessages.get()) return@synchronized - if (isRetryableMessage(json)) { + if (queueIfDisconnected && isRetryableMessage(json)) { + onSent?.invoke(id) Log.w(TAG, "WS rejected outgoing message — queuing for reconnect") queueMessage(id, json) recoverRejectedSocket(ws) @@ -950,8 +963,9 @@ object HermesWsClient { Log.w(TAG, "WS rejected oversized outgoing message — not retrying") } } - } else if (acceptQueuedMessages.get()) { + } else if (acceptQueuedMessages.get() && queueIfDisconnected) { if (isRetryableMessage(json)) { + onSent?.invoke(id) Log.d(TAG, "WS disconnected — queuing message") queueMessage(id, json) reconnect = true @@ -1287,6 +1301,26 @@ object HermesWsClient { ) : WebSocketListener() { private fun isCurrent(): Boolean = isCurrentGeneration(generation) + private fun consumeCapabilityResponse(event: WsEvent): Boolean { + val id = + when (event) { + is WsEvent.RpcResult -> event.id + is WsEvent.RpcError -> event.id + else -> return false + } + // send() registers capability request IDs via onSent while holding + // outboundLock. Take the same lock here so a very fast gateway + // response cannot be consumed before its ID has been registered. + val isCapabilityResponse = + synchronized(outboundLock) { + capabilityRequestIds.remove(id) + } + if (!isCapabilityResponse) return false + + removeQueuedMessage(id) + return true + } + override fun onOpen( webSocket: WebSocket, response: Response, @@ -1407,6 +1441,7 @@ object HermesWsClient { Log.e(TAG, "Failed to parse message: ${e.javaClass.simpleName}") WsEvent.Unknown(text) } + if (consumeCapabilityResponse(event)) return when (event) { is WsEvent.RpcResult -> { synchronized(outboundLock) { pendingPromptSubmits.remove(event.id) } @@ -1473,6 +1508,7 @@ object HermesWsClient { connected.set(false) ActiveSessionHolder.clear() stopHealthTracking() + capabilityRequestIds.clear() if (isTerminalAuthClose(code, reason)) { _connectionStatus.value = ConnectionStatus.AUTH_EXPIRED } else if (_connectionStatus.value != ConnectionStatus.AUTH_EXPIRED) { @@ -1501,6 +1537,7 @@ object HermesWsClient { connected.set(false) ActiveSessionHolder.clear() stopHealthTracking() + capabilityRequestIds.clear() val code = response?.code ?: 0 if (code == 401 || code == 4401 || code == 4403 || t.message?.contains("401") == true || diff --git a/app/src/test/java/com/m57/hermescontrol/data/ws/HermesWsClientTest.kt b/app/src/test/java/com/m57/hermescontrol/data/ws/HermesWsClientTest.kt index cfaf1a57..e3df3221 100644 --- a/app/src/test/java/com/m57/hermescontrol/data/ws/HermesWsClientTest.kt +++ b/app/src/test/java/com/m57/hermescontrol/data/ws/HermesWsClientTest.kt @@ -18,6 +18,7 @@ import io.mockk.mockkStatic import io.mockk.unmockkAll import io.mockk.verify import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job import kotlinx.coroutines.async @@ -221,6 +222,93 @@ class HermesWsClientTest { ) } + @Test + fun gatewayReadyCapabilityResponsesAreNotPublished() = + runBlocking { + val capabilityLatch = CountDownLatch(1) + val capabilitySuccessLatch = CountDownLatch(1) + val ordinaryErrorLatch = CountDownLatch(1) + var serverWebSocket: WebSocket? = null + var ordinaryRequestId: String? = null + val capabilityResponseCount = AtomicInteger(0) + + mockWebServer.enqueue( + MockResponse().withWebSocketUpgrade( + object : WebSocketListener() { + override fun onOpen( + webSocket: WebSocket, + response: okhttp3.Response, + ) { + serverWebSocket = webSocket + webSocket.send( + """{"jsonrpc":"2.0","method":"event","params":{"type":"gateway.ready","payload":{}}}""", + ) + } + + override fun onMessage( + webSocket: WebSocket, + text: String, + ) { + val frame = Json.parseToJsonElement(text).jsonObject + val method = frame["method"]?.jsonPrimitive?.content + val id = frame["id"]?.jsonPrimitive?.content ?: return + when (method) { + WsMethods.CLIENT_CAPABILITIES -> { + if (capabilityResponseCount.getAndIncrement() == 0) { + webSocket.send( + """{"jsonrpc":"2.0","id":"$id","error":{"code":-32601,"message":"unknown method: client.capabilities"}}""", + ) + capabilityLatch.countDown() + } else { + webSocket.send( + """{"jsonrpc":"2.0","id":"$id","result":{"server_requests":["clarify"]}}""", + ) + capabilitySuccessLatch.countDown() + } + } + + "ordinary_method" -> { + ordinaryRequestId = id + webSocket.send( + """{"jsonrpc":"2.0","id":"$id","error":{"code":4000,"message":"ordinary test failure"}}""", + ) + ordinaryErrorLatch.countDown() + } + } + } + }, + ), + ) + + val firstRpcResponse = + async(start = CoroutineStart.UNDISPATCHED) { + withTimeout(5_000) { + HermesWsClient.events.first { it is WsEvent.RpcError || it is WsEvent.RpcResult } + } + } + + HermesWsClient.connect() + + assertTrue("Capability request not rejected", capabilityLatch.await(5, TimeUnit.SECONDS)) + assertTrue( + "Second gateway.ready was not sent", + serverWebSocket?.send( + """{"jsonrpc":"2.0","method":"event","params":{"type":"gateway.ready","payload":{}}}""", + ) == true, + ) + assertTrue("Capability success was not received", capabilitySuccessLatch.await(5, TimeUnit.SECONDS)) + val sentOrdinaryRequestId = HermesWsClient.send("ordinary_method") + assertTrue("Ordinary request not rejected", ordinaryErrorLatch.await(5, TimeUnit.SECONDS)) + + val publishedEvent = firstRpcResponse.await() + assertTrue("Capability response leaked into the event stream", publishedEvent is WsEvent.RpcError) + val publishedError = publishedEvent as WsEvent.RpcError + assertEquals(sentOrdinaryRequestId, publishedError.id) + assertEquals(ordinaryRequestId, publishedError.id) + assertEquals(4000, publishedError.error.code) + assertEquals("ordinary test failure", publishedError.error.message) + } + @Test fun testOpenRequestsReplayUsesTheLiveServerRequestDispatcher() = runBlocking {