From a0bbb7cd4cc6c4ee8a1b0ff6f93504f680069a73 Mon Sep 17 00:00:00 2001
From: Josh Patra <joshpatra12@gmail.com>
Date: Tue, 3 Feb 2026 19:10:02 -0500
Subject: [PATCH] Integrate WebSocket proxy with session manager to cleanup
 sessions on client disconnect

---
 .../Middleware/WebSocketProxyMiddleware.cs    | 39 +++++++++++++++++--
 1 file changed, 36 insertions(+), 3 deletions(-)

diff --git a/allstarr/Middleware/WebSocketProxyMiddleware.cs b/allstarr/Middleware/WebSocketProxyMiddleware.cs
index 064a607..72bceaa 100644
--- a/allstarr/Middleware/WebSocketProxyMiddleware.cs
+++ b/allstarr/Middleware/WebSocketProxyMiddleware.cs
@@ -1,6 +1,7 @@
 using System.Net.WebSockets;
 using Microsoft.Extensions.Options;
 using allstarr.Models.Settings;
+using allstarr.Services.Jellyfin;
 
 namespace allstarr.Middleware;
 
@@ -13,15 +14,18 @@ public class WebSocketProxyMiddleware
     private readonly RequestDelegate _next;
     private readonly JellyfinSettings _settings;
     private readonly ILogger<WebSocketProxyMiddleware> _logger;
+    private readonly JellyfinSessionManager _sessionManager;
 
     public WebSocketProxyMiddleware(
         RequestDelegate next,
         IOptions<JellyfinSettings> settings,
-        ILogger<WebSocketProxyMiddleware> logger)
+        ILogger<WebSocketProxyMiddleware> logger,
+        JellyfinSessionManager sessionManager)
     {
         _next = next;
         _settings = settings.Value;
         _logger = logger;
+        _sessionManager = sessionManager;
         
         _logger.LogDebug("🔧 WEBSOCKET: WebSocketProxyMiddleware initialized - Jellyfin URL: {Url}", _settings.Url);
     }
@@ -65,9 +69,31 @@ public class WebSocketProxyMiddleware
     {
         ClientWebSocket? serverWebSocket = null;
         WebSocket? clientWebSocket = null;
+        string? deviceId = null;
 
         try
         {
+            // Extract device ID from query string or headers for session tracking
+            deviceId = context.Request.Query["deviceId"].ToString();
+            if (string.IsNullOrEmpty(deviceId))
+            {
+                // Try to extract from X-Emby-Authorization header
+                if (context.Request.Headers.TryGetValue("X-Emby-Authorization", out var authHeader))
+                {
+                    var authValue = authHeader.ToString();
+                    var deviceIdMatch = System.Text.RegularExpressions.Regex.Match(authValue, @"DeviceId=""([^""]+)""");
+                    if (deviceIdMatch.Success)
+                    {
+                        deviceId = deviceIdMatch.Groups[1].Value;
+                    }
+                }
+            }
+
+            if (!string.IsNullOrEmpty(deviceId))
+            {
+                _logger.LogDebug("🔍 WEBSOCKET: Client WebSocket for device {DeviceId}", deviceId);
+            }
+
             // Accept the WebSocket connection from the client
             clientWebSocket = await context.WebSockets.AcceptWebSocketAsync();
             _logger.LogDebug("✓ WEBSOCKET: Client WebSocket accepted");
@@ -96,9 +122,9 @@ public class WebSocketProxyMiddleware
                 serverWebSocket.Options.SetRequestHeader("X-Emby-Authorization", embyAuthHeader.ToString());
                 _logger.LogDebug("🔑 WEBSOCKET: Forwarded X-Emby-Authorization header");
             }
-            else if (context.Request.Headers.TryGetValue("Authorization", out var authHeader))
+            else if (context.Request.Headers.TryGetValue("Authorization", out var authHeader2))
             {
-                var authValue = authHeader.ToString();
+                var authValue = authHeader2.ToString();
                 // If it's a MediaBrowser auth header, use X-Emby-Authorization
                 if (authValue.Contains("MediaBrowser", StringComparison.OrdinalIgnoreCase))
                 {
@@ -165,6 +191,13 @@ public class WebSocketProxyMiddleware
             clientWebSocket?.Dispose();
             serverWebSocket?.Dispose();
 
+            // CRITICAL: Notify session manager that client disconnected
+            if (!string.IsNullOrEmpty(deviceId))
+            {
+                _logger.LogInformation("🧹 WEBSOCKET: Client disconnected, removing session for device {DeviceId}", deviceId);
+                await _sessionManager.RemoveSessionAsync(deviceId);
+            }
+
             _logger.LogDebug("🧹 WEBSOCKET: WebSocket connections cleaned up");
         }
     }