Update

The summary was for reference only and not needed in version control.

enterprise/server/auth/cookie_compression.py

@@ -0,0 +1,123 @@
+"""
+Cookie compression utilities for keycloak_auth cookie.
+
+This module provides functions to compress and decompress cookie data
+to reduce cookie size and improve performance.
+"""
+
+import base64
+import gzip
+
+from openhands.core.logger import openhands_logger as logger
+
+
+def compress_cookie_data(data: str) -> str:
+    """
+    Compress cookie data using gzip and encode with base64.
+
+    Args:
+        data: The cookie data string to compress
+
+    Returns:
+        Base64 encoded compressed data with 'gz:' prefix to indicate compression
+
+    Raises:
+        Exception: If compression fails
+    """
+    try:
+        # Convert string to bytes
+        data_bytes = data.encode('utf-8')
+
+        # Compress using gzip
+        compressed_bytes = gzip.compress(data_bytes, compresslevel=6)
+
+        # Encode with base64 for safe cookie storage
+        encoded_data = base64.b64encode(compressed_bytes).decode('ascii')
+
+        # Add prefix to indicate this is compressed data
+        compressed_cookie = f'gz:{encoded_data}'
+
+        logger.debug(
+            'Cookie compression stats',
+            extra={
+                'original_size': len(data),
+                'compressed_size': len(compressed_cookie),
+                'compression_ratio': len(compressed_cookie) / len(data)
+                if len(data) > 0
+                else 0,
+            },
+        )
+
+        return compressed_cookie
+
+    except Exception as e:
+        logger.error(f'Failed to compress cookie data: {str(e)}')
+        raise
+
+
+def decompress_cookie_data(data: str) -> str:
+    """
+    Decompress cookie data if it's compressed, otherwise return as-is.
+
+    Args:
+        data: The cookie data string (may be compressed or uncompressed)
+
+    Returns:
+        Decompressed cookie data string
+
+    Raises:
+        Exception: If decompression fails for compressed data
+    """
+    try:
+        # Check if data is compressed (has 'gz:' prefix)
+        if not data.startswith('gz:'):
+            # Not compressed, return as-is for backward compatibility
+            logger.debug('Cookie data is not compressed, returning as-is')
+            return data
+
+        # Remove the 'gz:' prefix
+        encoded_data = data[3:]
+
+        # Check for empty compressed data
+        if not encoded_data:
+            raise ValueError('Empty compressed data')
+
+        # Decode from base64
+        compressed_bytes = base64.b64decode(encoded_data.encode('ascii'))
+
+        # Decompress using gzip
+        decompressed_bytes = gzip.decompress(compressed_bytes)
+
+        # Convert back to string
+        decompressed_data = decompressed_bytes.decode('utf-8')
+
+        logger.debug(
+            'Cookie decompression stats',
+            extra={
+                'compressed_size': len(data),
+                'decompressed_size': len(decompressed_data),
+                'compression_ratio': len(data) / len(decompressed_data)
+                if len(decompressed_data) > 0
+                else 0,
+            },
+        )
+
+        return decompressed_data
+
+    except Exception as e:
+        logger.error(f'Failed to decompress cookie data: {str(e)}')
+        raise
+
+
+def should_compress_cookie(data: str, min_size_threshold: int = 1000) -> bool:
+    """
+    Determine if cookie data should be compressed based on size.
+
+    Args:
+        data: The cookie data string
+        min_size_threshold: Minimum size in bytes to consider compression
+
+    Returns:
+        True if data should be compressed, False otherwise
+    """
+    return len(data.encode('utf-8')) >= min_size_threshold

enterprise/server/auth/saas_user_auth.py

@@ -13,6 +13,7 @@ from server.auth.auth_error import (
     ExpiredError,
     NoCredentialsError,
 )
+from server.auth.cookie_compression import decompress_cookie_data
 from server.auth.token_manager import TokenManager
 from server.config import get_config
 from server.logger import logger
@@ -271,7 +272,18 @@ async def saas_user_auth_from_cookie(request: Request) -> SaasUserAuth | None:
         signed_token = request.cookies.get('keycloak_auth')
         if not signed_token:
             return None
-        return await saas_user_auth_from_signed_token(signed_token)
+
+        # Decompress the cookie data if it's compressed
+        try:
+            decompressed_token = decompress_cookie_data(signed_token)
+            logger.debug('Cookie data decompressed successfully')
+        except Exception as e:
+            logger.warning(
+                f'Failed to decompress cookie data, trying as uncompressed: {str(e)}'
+            )
+            decompressed_token = signed_token
+
+        return await saas_user_auth_from_signed_token(decompressed_token)
     except Exception as exc:
         raise CookieError from exc
 

enterprise/server/middleware.py

@@ -10,6 +10,7 @@ from server.auth.auth_error import (
     NoCredentialsError,
     TosNotAcceptedError,
 )
+from server.auth.cookie_compression import decompress_cookie_data
 from server.auth.gitlab_sync import schedule_gitlab_repo_sync
 from server.auth.saas_user_auth import SaasUserAuth, token_manager
 from server.routes.auth import (
@@ -114,8 +115,18 @@ class SetAuthCookieMiddleware:
         jwt_secret: SecretStr = config.jwt_secret  # type: ignore[assignment]
         if keycloak_auth_cookie:
             try:
+                # Decompress the cookie data if it's compressed
+                try:
+                    decompressed_cookie = decompress_cookie_data(keycloak_auth_cookie)
+                    logger.debug('Middleware: Cookie data decompressed successfully')
+                except Exception as e:
+                    logger.debug(
+                        f'Middleware: Failed to decompress cookie data, trying as uncompressed: {str(e)}'
+                    )
+                    decompressed_cookie = keycloak_auth_cookie
+
                 decoded = jwt.decode(
-                    keycloak_auth_cookie,
+                    decompressed_cookie,
                     jwt_secret.get_secret_value(),
                     algorithms=['HS256'],
                 )

enterprise/server/routes/auth.py

@@ -13,6 +13,7 @@ from server.auth.constants import (
     KEYCLOAK_REALM_NAME,
     KEYCLOAK_SERVER_URL_EXT,
 )
+from server.auth.cookie_compression import compress_cookie_data, should_compress_cookie
 from server.auth.gitlab_sync import schedule_gitlab_repo_sync
 from server.auth.saas_user_auth import SaasUserAuth
 from server.auth.token_manager import TokenManager
@@ -55,12 +56,24 @@ def set_response_cookie(
     }
     signed_token = sign_token(cookie_data, config.jwt_secret.get_secret_value())  # type: ignore
 
-    # Set secure cookie with signed token
+    # Compress the signed token if it's large enough to benefit from compression
+    cookie_value = signed_token
+    if should_compress_cookie(signed_token):
+        try:
+            cookie_value = compress_cookie_data(signed_token)
+            logger.debug('Cookie data compressed successfully')
+        except Exception as e:
+            logger.warning(
+                f'Failed to compress cookie data, using uncompressed: {str(e)}'
+            )
+            cookie_value = signed_token
+
+    # Set secure cookie with (potentially compressed) signed token
     domain = get_cookie_domain(request)
     if domain:
         response.set_cookie(
             key='keycloak_auth',
-            value=signed_token,
+            value=cookie_value,
             domain=domain,
             httponly=True,
             secure=secure,
@@ -69,7 +82,7 @@ def set_response_cookie(
     else:
         response.set_cookie(
             key='keycloak_auth',
-            value=signed_token,
+            value=cookie_value,
             httponly=True,
             secure=secure,
             samesite=get_cookie_samesite(request),

enterprise/storage/saas_conversation_validator.py

@@ -1,4 +1,5 @@
 from server.auth.auth_error import AuthError, ExpiredError
+from server.auth.cookie_compression import decompress_cookie_data
 from server.auth.saas_user_auth import saas_user_auth_from_signed_token
 from server.auth.token_manager import TokenManager
 from socketio.exceptions import ConnectionRefusedError
@@ -129,8 +130,20 @@ class SaasConversationValidator(ConversationValidator):
         if not config.jwt_secret:
             raise RuntimeError('JWT secret not found')
 
+        # Decompress the cookie data if it's compressed
         try:
-            user_auth = await saas_user_auth_from_signed_token(signed_token)
+            decompressed_token = decompress_cookie_data(signed_token)
+            logger.debug(
+                'Conversation validator: Cookie data decompressed successfully'
+            )
+        except Exception as e:
+            logger.debug(
+                f'Conversation validator: Failed to decompress cookie data, trying as uncompressed: {str(e)}'
+            )
+            decompressed_token = signed_token
+
+        try:
+            user_auth = await saas_user_auth_from_signed_token(decompressed_token)
             access_token = await user_auth.get_access_token()
         except ExpiredError:
             raise ConnectionRefusedError('SESSION$TIMEOUT_MESSAGE')