Fix jti consistency for token creation (#1255)

* jti consistency for token creation

Signed-off-by: Keval Mahajan <[email protected]>

* linting

Signed-off-by: Keval Mahajan <[email protected]>

* updated test cases

Signed-off-by: Keval Mahajan <[email protected]>

* rebase

Signed-off-by: Mihai Criveti <[email protected]>

---------

Signed-off-by: Keval Mahajan <[email protected]>
Signed-off-by: Mihai Criveti <[email protected]>
Co-authored-by: Mihai Criveti <[email protected]>

Author: kevalmahajan

mcpgateway/services/token_catalog_service.py

@@ -213,7 +213,9 @@ def __init__(self, db: Session):
         """
         self.db = db
 
-    async def _generate_token(self, user_email: str, team_id: Optional[str] = None, expires_at: Optional[datetime] = None, scope: Optional["TokenScope"] = None, user: Optional[object] = None) -> str:
+    async def _generate_token(
+        self, user_email: str, jti: str, team_id: Optional[str] = None, expires_at: Optional[datetime] = None, scope: Optional["TokenScope"] = None, user: Optional[object] = None
+    ) -> str:
         """Generate a JWT token for API access.
 
         This internal method creates a properly formatted JWT token with all
@@ -222,6 +224,7 @@ async def _generate_token(self, user_email: str, team_id: Optional[str] = None,
 
         Args:
             user_email: User's email address for the token subject
+            jti: JWT ID for token uniqueness
             team_id: Optional team ID for team-scoped tokens
             expires_at: Optional expiration datetime
             scope: Optional token scope information for access control
@@ -242,7 +245,7 @@ async def _generate_token(self, user_email: str, team_id: Optional[str] = None,
             "iss": settings.jwt_issuer,  # Issuer
             "aud": settings.jwt_audience,  # Audience
             "iat": int(now.timestamp()),  # Issued at
-            "jti": str(uuid.uuid4()),  # JWT ID for uniqueness
+            "jti": jti,  # JWT ID for uniqueness
             "user": {"email": user_email, "full_name": "API Token User", "is_admin": user.is_admin if user else False, "auth_provider": "api_token"},  # Use actual admin status if user provided
             "teams": [team_id] if team_id else [],
             "namespaces": [f"user:{user_email}", "public"] + ([f"team:{team_id}"] if team_id else []),
@@ -383,8 +386,9 @@ async def create_token(
         if expires_in_days:
             expires_at = utc_now() + timedelta(days=expires_in_days)
 
+        jti = str(uuid.uuid4())  # Unique JWT ID
         # Generate JWT token with all necessary claims
-        raw_token = await self._generate_token(user_email=user_email, team_id=team_id, expires_at=expires_at, scope=scope, user=user)  # Pass user object to include admin status
+        raw_token = await self._generate_token(user_email=user_email, jti=jti, team_id=team_id, expires_at=expires_at, scope=scope, user=user)  # Pass user object to include admin status
 
         # Hash token for secure storage
         token_hash = self._hash_token(raw_token)
@@ -395,6 +399,7 @@ async def create_token(
             user_email=user_email,
             team_id=team_id,  # Store team association
             name=name,
+            jti=jti,
             description=description,
             token_hash=token_hash,  # Store hash, not raw token
             expires_at=expires_at,

tests/unit/mcpgateway/services/test_token_catalog_service.py

@@ -237,8 +237,8 @@ async def test_generate_token_basic(self, token_service):
         """Test _generate_token method with basic parameters."""
         with patch("mcpgateway.services.token_catalog_service.create_jwt_token", new_callable=AsyncMock) as mock_create_jwt:
             mock_create_jwt.return_value = "jwt_token_123"
-
-            token = await token_service._generate_token("[email protected]")
+            jti=str(uuid.uuid4())
+            token = await token_service._generate_token("[email protected]",jti)
 
             assert token == "jwt_token_123"
             mock_create_jwt.assert_called_once()
@@ -253,8 +253,8 @@ async def test_generate_token_with_team(self, token_service):
         """Test _generate_token method with team_id."""
         with patch("mcpgateway.services.token_catalog_service.create_jwt_token", new_callable=AsyncMock) as mock_create_jwt:
             mock_create_jwt.return_value = "jwt_token_team"
-
-            token = await token_service._generate_token("[email protected]", team_id="team-123")
+            jti=str(uuid.uuid4())
+            token = await token_service._generate_token("[email protected]", jti=jti, team_id="team-123")
 
             assert token == "jwt_token_team"
             call_args = mock_create_jwt.call_args[0][0]
@@ -267,8 +267,9 @@ async def test_generate_token_with_expiry(self, token_service):
         with patch("mcpgateway.services.token_catalog_service.create_jwt_token", new_callable=AsyncMock) as mock_create_jwt:
             mock_create_jwt.return_value = "jwt_token_exp"
             expires_at = datetime.now(timezone.utc) + timedelta(days=7)
+            jti=str(uuid.uuid4())
 
-            token = await token_service._generate_token("[email protected]", expires_at=expires_at)
+            token = await token_service._generate_token("[email protected]", jti=jti, expires_at=expires_at)
 
             assert token == "jwt_token_exp"
             call_args = mock_create_jwt.call_args[0][0]
@@ -280,8 +281,9 @@ async def test_generate_token_with_scope(self, token_service, token_scope):
         """Test _generate_token method with TokenScope."""
         with patch("mcpgateway.services.token_catalog_service.create_jwt_token", new_callable=AsyncMock) as mock_create_jwt:
             mock_create_jwt.return_value = "jwt_token_scoped"
+            jti=str(uuid.uuid4())
 
-            token = await token_service._generate_token("[email protected]", scope=token_scope)
+            token = await token_service._generate_token("[email protected]", jti=jti, scope=token_scope)
 
             assert token == "jwt_token_scoped"
             call_args = mock_create_jwt.call_args[0][0]
@@ -295,8 +297,9 @@ async def test_generate_token_with_admin_user(self, token_service, mock_user):
         mock_user.is_admin = True
         with patch("mcpgateway.services.token_catalog_service.create_jwt_token", new_callable=AsyncMock) as mock_create_jwt:
             mock_create_jwt.return_value = "jwt_token_admin"
+            jti=str(uuid.uuid4())
 
-            token = await token_service._generate_token("[email protected]", user=mock_user)
+            token = await token_service._generate_token("[email protected]", jti=jti, user=mock_user)
 
             assert token == "jwt_token_admin"
             call_args = mock_create_jwt.call_args[0][0]
@@ -918,7 +921,9 @@ async def test_generate_token_settings_values(self, token_service):
 
             with patch("mcpgateway.services.token_catalog_service.create_jwt_token", new_callable=AsyncMock) as mock_create:
                 mock_create.return_value = "jwt"
-                await token_service._generate_token("[email protected]")
+                jti=str(uuid.uuid4())
+
+                await token_service._generate_token("[email protected]", jti=jti)
 
                 call_args = mock_create.call_args[0][0]
                 assert call_args["iss"] == "test-issuer"