From 59c69c8c58c99d257f6a09be2e0d8f8bebc00918 Mon Sep 17 00:00:00 2001
From: Stephen Bracken <email-protected>
Date: Tue, 4 Aug 2026 10:49:06 +0100
Subject: [PATCH] Fix Airflow JWT getting cleared in /login_callback

---
 .../providers/keycloak/auth_manager/middleware.py  | 14 ++++++++------
 .../unit/keycloak/auth_manager/test_middleware.py  |  3 ++-
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/middleware.py b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/middleware.py
index fcbfc56f044d3..1ad6f1eafa9de 100644
--- a/providers/keycloak/src/airflow/providers/keycloak/auth_manager/middleware.py
+++ b/providers/keycloak/src/airflow/providers/keycloak/auth_manager/middleware.py
@@ -78,8 +78,6 @@ async def dispatch(self, request: Request, call_next):
                 user = new_user or current_user
             except (
                 AuthManagerRefreshTokenExpiredException,
-                ExpiredSignatureError,
-                InvalidTokenError,
                 HTTPException,
             ):
                 new_token = ""
@@ -234,13 +232,17 @@ async def _refresh_user(
         access_token = request.cookies.get(COOKIE_NAME_ACCESS_TOKEN)
         refresh_token = request.cookies.get(COOKIE_NAME_REFRESH_TOKEN)
         if not jwt_token:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED, detail="User is not logged into Airflow."
-            )
+            # User is not logged into Airflow
+            return None, None
         if not access_token:
             raise HTTPException(
                 status_code=status.HTTP_401_UNAUTHORIZED, detail="User is not logged into Keycloak."
             )
         auth_manager = cast("KeycloakAuthManager", get_auth_manager())
-        user = await auth_manager.get_user_from_token(jwt_token, access_token, refresh_token)
+        try:
+            user = await auth_manager.get_user_from_token(jwt_token, access_token, refresh_token)
+        except ExpiredSignatureError:
+            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Token Expired")
+        except InvalidTokenError:
+            raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid JWT token")
         return get_auth_manager().refresh_user(user=user), user
diff --git a/providers/keycloak/tests/unit/keycloak/auth_manager/test_middleware.py b/providers/keycloak/tests/unit/keycloak/auth_manager/test_middleware.py
index 29a72988823d6..66f610ae30840 100644
--- a/providers/keycloak/tests/unit/keycloak/auth_manager/test_middleware.py
+++ b/providers/keycloak/tests/unit/keycloak/auth_manager/test_middleware.py
@@ -239,7 +239,7 @@ async def test_no_airflow_jwt_token(
             COOKIE_NAME_REFRESH_TOKEN: "refresh_token",
         }
 
-        await middleware.dispatch(mock_request, call_next)
+        response = await middleware.dispatch(mock_request, call_next)
 
         auth_manager.get_user_from_token.assert_not_called()
         auth_manager.refresh_user.assert_not_called()
@@ -258,6 +258,7 @@ async def test_no_airflow_jwt_token(
             assert not hasattr(mock_request.state, "user_authenticated_via")
 
         call_next.assert_awaited_once_with(mock_request)
+        response.set_cookie.assert_not_called()
 
     @patch("airflow.providers.keycloak.auth_manager.middleware.get_auth_manager")
     @pytest.mark.asyncio