feat(jwt): extend Token.encode() to support custom headers and encoder options

Author: s-aleshin

litestar/security/jwt/token.py

@@ -11,8 +11,12 @@
 from litestar.exceptions import ImproperlyConfiguredException, NotAuthorizedException
 
 if TYPE_CHECKING:
+    import json
+
+    from jwt.algorithms import AllowedPrivateKeys
     from typing_extensions import Self
 
+
 __all__ = (
     "JWTDecodeOptions",
     "Token",
@@ -192,12 +196,22 @@ def decode(
         ) as e:
             raise NotAuthorizedException("Invalid token") from e
 
-    def encode(self, secret: str, algorithm: str) -> str:
+    def encode(
+        self,
+        secret: AllowedPrivateKeys | str | bytes,
+        algorithm: str,
+        headers: dict[str, Any] | None = None,
+        json_encoder: type[json.JSONEncoder] | None = None,
+        sort_headers: bool = True,
+    ) -> str:
         """Encode the token instance into a string.
 
         Args:
             secret: The secret with which the JWT is encoded.
             algorithm: The algorithm used to encode the JWT.
+            headers: Optional headers to include in the JWT (e.g., {"kid": "..."}).
+            json_encoder: Optional custom JSON encoder class to serialize the payload.
+            sort_headers: Whether to sort the JWT headers before encoding. Defaults to True.
 
         Returns:
             An encoded token string.
@@ -210,6 +224,9 @@ def encode(self, secret: str, algorithm: str) -> str:
                 payload={k: v for k, v in asdict(self).items() if v is not None},
                 key=secret,
                 algorithm=algorithm,
+                headers=headers,
+                json_encoder=json_encoder,
+                sort_headers=sort_headers,
             )
         except (jwt.DecodeError, NotImplementedError) as e:
             raise ImproperlyConfiguredException("Failed to encode token") from e

tests/unit/test_security/test_jwt/test_token.py

@@ -222,3 +222,15 @@ def decode_payload(
     _secret = secrets.token_hex()
     encoded = CustomToken(exp=datetime.now() + timedelta(days=1), sub="foo").encode(_secret, "HS256")
     assert CustomToken.decode(encoded, secret=_secret, algorithm="HS256").sub == "some-random-value"
+
+
+def test_token_encode_includes_custom_headers() -> None:
+    token = Token(exp=datetime.now() + timedelta(days=1), sub="some-random-value")
+    custom_headers = {"kid": "key-id"}
+    encoded = token.encode(secret=secrets.token_hex(), algorithm="HS256", headers=custom_headers)
+    header = jwt.get_unverified_header(encoded)
+
+    assert header["alg"] == "HS256"
+    assert header["typo"] == "JWT"
+    assert "kid" in header
+    assert header["kid"] == custom_headers["kid"]