Fix router HTTP timeout connection cleanup

src/service/router/BUILD

@@ -17,7 +17,7 @@ SPDX-License-Identifier: Apache-2.0
 """
 
 load("@aspect_rules_py//py:defs.bzl", "py_image_layer")
-load("//bzl:py.bzl", "osmo_py_binary", "osmo_py_library", "osmo_python_wrapper")
+load("//bzl:py.bzl", "osmo_py_binary", "osmo_py_library", "osmo_py_test", "osmo_python_wrapper")
 load("@osmo_python_deps//:requirements.bzl", "requirement")
 load("@rules_oci//oci:defs.bzl", "oci_image", "oci_push", "oci_load")
 load("@osmo_constants//:constants.bzl", "BASE_IMAGE_URL", "IMAGE_TAG")
@@ -50,6 +50,12 @@ osmo_py_binary(
     deps = [":router_lib"],
 )
 
+osmo_py_test(
+    name = "test_router",
+    srcs = ["test_router.py"],
+    deps = [":router_lib"],
+)
+
 py_image_layer(
     name = "router_layer",
     binary = ":router_binary",

src/service/router/router.py

@@ -99,6 +99,12 @@ class ConnectionPayload(pydantic.BaseModel):
 webservers: Dict[str, WebserverConnection] = {}
 
 
+def _remove_router_connection(key: str):
+    connection = connections.pop(key, None)
+    if connection and connection.wait_close:
+        connection.wait_close.set()
+
+
 class RouterWebSocketMiddleware:
     """Middleware for handling WebSocket connections in the router service."""
     # pylint: disable=redefined-outer-name
@@ -163,7 +169,7 @@ async def run_connect_backend(ws: fastapi.WebSocket, name: str, key: str):
         logging.info('Backend connection for workflow %s with key %s is timeout', name, key)
         await ws.close(4000, 'Router connection timeout')
 
-    del connections[key]
+    connections.pop(key, None)
     # Make close faster
     try:
         await ws.close()
@@ -250,10 +256,12 @@ async def webserver_http_request(request: fastapi.Request, ctrl_key: str):
         ws = connections[conn_key].websocket
         close = connections[conn_key].wait_close
     except asyncio.TimeoutError:
+        _remove_router_connection(conn_key)
         return fastapi.Response(
             content='Request timed out waiting for backend connection.', status_code=504)
 
     if not ws or not close:  # To fix pytype error
+        _remove_router_connection(conn_key)
         return fastapi.Response(
             content='No active backend connection found, your session may have expired.',
             status_code=404)

src/service/router/test_router.py

@@ -0,0 +1,106 @@
+"""
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+SPDX-License-Identifier: Apache-2.0
+"""
+
+import asyncio
+import datetime
+import importlib
+import os
+import unittest
+from unittest import mock
+
+with (
+    mock.patch.dict(os.environ, {'OSMO_POSTGRES_PASSWORD': 'test-password'}),
+    mock.patch('fastapi.applications.FastAPI.add_middleware'),
+):
+    router = importlib.import_module('src.service.router.router')
+
+
+class _FakeRequest:
+    cookies = {
+        '_osmo_router_affinity': 'sticky-session',
+        'ignored': 'cookie',
+    }
+
+    async def body(self):
+        return b''
+
+
+class _FakeControlWebSocket:
+    def __init__(self):
+        self.messages = []
+
+    async def send_json(self, payload):
+        self.messages.append(payload)
+
+
+async def _raise_timeout(awaitable, timeout):
+    _ = timeout
+    awaitable.close()
+    raise asyncio.TimeoutError
+
+
+class WebserverHttpRequestTestCase(unittest.TestCase):
+    """Tests for proxied webserver HTTP requests."""
+
+    def setUp(self):
+        router.connections.clear()
+        router.webservers.clear()
+
+    def tearDown(self):
+        router.connections.clear()
+        router.webservers.clear()
+
+    def test_removes_pending_connection_on_backend_timeout(self):
+        async def run_test():
+            control_websocket = _FakeControlWebSocket()
+            router.webservers['session-key'] = router.WebserverConnection.model_construct(
+                wait_close=asyncio.Event(),
+                last_active_time=datetime.datetime.now(),
+                websocket=control_websocket,
+            )
+            config = mock.Mock(timeout=60, sticky_cookies=['_osmo_router_affinity'])
+
+            with (
+                mock.patch.object(router.common, 'generate_unique_id', return_value='timeout'),
+                mock.patch.object(
+                    router.helper,
+                    'http2raw',
+                    new=mock.AsyncMock(return_value=b'GET / HTTP/1.1\r\n\r\n'),
+                ),
+                mock.patch.object(router.RouterServiceConfig, 'load', return_value=config),
+                mock.patch.object(router.asyncio, 'wait_for', side_effect=_raise_timeout),
+            ):
+                response = await router.webserver_http_request(_FakeRequest(), 'session-key')
+
+            conn_key = 'PORTFORWARD-timeout'
+            self.assertEqual(response.status_code, 504)
+            self.assertNotIn(conn_key, router.connections)
+            self.assertEqual(
+                control_websocket.messages,
+                [{
+                    'key': conn_key,
+                    'cookie': '_osmo_router_affinity=sticky-session',
+                    'type': 'tcp',
+                }],
+            )
+
+        asyncio.run(run_test())
+
+
+if __name__ == '__main__':
+    unittest.main()