From c31c746be83165b67a983ff90b87b394d39dc596 Mon Sep 17 00:00:00 2001
From: sahandilshan <sahandilshan222@gmail.com>
Date: Thu, 16 Jan 2025 17:08:51 +0530
Subject: [PATCH] Fix review suggestions

---
 .../pom.xml                                   | 11 ++-
 .../ai/service/mgt/constants/AIConstants.java |  6 +-
 .../mgt/exceptions/AIClientException.java     | 33 ++-----
 .../service/mgt/exceptions/AIException.java   | 71 +++++++++++++++
 .../mgt/exceptions/AIServerException.java     | 38 ++------
 .../mgt/token/AIAccessTokenManager.java       |  5 +-
 .../ai/service/mgt/util/AIHttpClientUtil.java |  4 +-
 .../mgt/constants/AIConstantsTest.java        | 65 --------------
 .../mgt/exceptions/AIClientExceptionTest.java | 60 -------------
 .../mgt/exceptions/AIServerExceptionTest.java | 70 ---------------
 .../mgt/token/AIAccessTokenManagerTest.java   | 89 ++++++++++---------
 .../mgt/util/AIHttpClientUtilTest.java        | 84 +++++++++--------
 .../src/test/resources/testng.xml             |  3 -
 components/ai-services-mgt/pom.xml            |  4 +-
 .../mgt/ai/LoginFlowAIManager.java            | 34 ++++++-
 .../mgt/ai/LoginFlowAIManagerImpl.java        | 43 ++-------
 .../mgt/ai/constant/LoginFlowAIConstants.java |  8 +-
 .../mgt/ai/LoginFlowAIManagerTest.java        | 17 +---
 .../pom.xml                                   |  4 +-
 features/ai-services-mgt/pom.xml              |  4 +-
 .../resources/identity.xml.j2                 |  4 +-
 21 files changed, 250 insertions(+), 407 deletions(-)
 create mode 100644 components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIException.java
 delete mode 100644 components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstantsTest.java
 delete mode 100644 components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientExceptionTest.java
 delete mode 100644 components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerExceptionTest.java

diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/pom.xml b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/pom.xml
index 123b3126d9f2..2db7cb1814f5 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/pom.xml
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/pom.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!--
-  ~ Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+  ~ Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
   ~
   ~ WSO2 LLC. licenses this file to you under the Apache License,
   ~ Version 2.0 (the "License"); you may not use this file except
@@ -22,7 +22,7 @@
     <parent>
         <groupId>org.wso2.carbon.identity.framework</groupId>
         <artifactId>ai-services-mgt</artifactId>
-        <version>7.7.87-SNAPSHOT</version>
+        <version>7.7.91-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
 
@@ -85,7 +85,6 @@
 
     <build>
         <plugins>
-
             <plugin>
                 <groupId>org.apache.felix</groupId>
                 <artifactId>maven-bundle-plugin</artifactId>
@@ -145,6 +144,12 @@
                 <groupId>org.jacoco</groupId>
                 <artifactId>jacoco-maven-plugin</artifactId>
                 <version>${jacoco.version}</version>
+                <configuration>
+                    <excludes>
+                        <exclude>**/*Exception.class</exclude>
+                        <exclude>**/*Constants*.class</exclude>
+                    </excludes>
+                </configuration>
                 <executions>
                     <execution>
                         <id>default-prepare-agent</id>
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstants.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstants.java
index 2445ca29375a..9621165b6646 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstants.java
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstants.java
@@ -23,13 +23,17 @@
  */
 public class AIConstants {
 
+    private AIConstants () {
+
+    }
+
     public static final String AI_SERVICE_KEY_PROPERTY_NAME = "AIServices.Key";
     public static final String AI_TOKEN_ENDPOINT_PROPERTY_NAME = "AIServices.TokenEndpoint";
     public static final String AI_TOKEN_SERVICE_MAX_RETRIES_PROPERTY_NAME = "AIServices.TokenRequestMaxRetries";
     public static final String AI_TOKEN_CONNECTION_TIMEOUT_PROPERTY_NAME = "AIServices.TokenConnectionTimeout";
     public static final String AI_TOKEN_CONNECTION_REQUEST_TIMEOUT_PROPERTY_NAME = "AIServices" +
             ".TokenConnectionRequestTimeout";
-    public static final String AI_TOKEN_SOCKET_TIMEOUT_PROPERTY_NAME = "AIServices.TokenSocketTimeout";
+    public static final String AI_TOKEN_SOCKET_TIMEOUT_PROPERTY_NAME = "AIServices.TokenConnectionSocketTimeout";
 
     public static final String HTTP_CONNECTION_POOL_SIZE_PROPERTY_NAME = "AIServices.HTTPConnectionPoolSize";
     public static final String HTTP_CONNECTION_TIMEOUT_PROPERTY_NAME = "AIServices.HTTPConnectionTimeout";
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientException.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientException.java
index fa91436a07ae..609d2bf6514a 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientException.java
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientException.java
@@ -21,44 +21,21 @@
 /**
  * Client Exception class for AI service.
  */
-public class AIClientException extends Exception {
-
-    private final String errorCode;
-    private String serverMessage;
-    private int serverStatusCode;
+public class AIClientException extends AIException {
+    ;
 
     public AIClientException(String message, String errorCode) {
 
-        super(message);
-        this.errorCode = errorCode;
+        super(message, errorCode);
     }
 
     public AIClientException(String message, String errorCode, int serverStatusCode, String serverMessage) {
 
-        super(message);
-        this.errorCode = errorCode;
-        this.serverStatusCode = serverStatusCode;
-        this.serverMessage = serverMessage;
+        super(message, errorCode, serverStatusCode, serverMessage);
     }
 
     public AIClientException(String message, String errorCode, Throwable cause) {
 
-        super(message, cause);
-        this.errorCode = errorCode;
-    }
-
-    public String getErrorCode() {
-
-        return errorCode;
-    }
-
-    public String getServerMessage() {
-
-        return serverMessage;
-    }
-
-    public int getServerStatusCode() {
-
-        return serverStatusCode;
+        super(message, errorCode, cause);
     }
 }
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIException.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIException.java
new file mode 100644
index 000000000000..a7b718f234b2
--- /dev/null
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIException.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.ai.service.mgt.exceptions;
+
+/**
+ * Generic Exception class for AI service.
+ */
+public class AIException extends Exception {
+
+    private String errorCode;
+    // This is the error message that comes from the server.
+    private String serverMessage;
+    // This is the status code that comes from the server.
+    private int serverStatusCode;
+
+    public AIException(String message, String errorCode) {
+
+        super(message);
+        this.errorCode = errorCode;
+    }
+
+    public AIException(String message, Throwable cause) {
+
+        super(message, cause);
+    }
+
+    public AIException(String message, String errorCode, int serverStatusCode, String serverMessage) {
+
+        super(message);
+        this.errorCode = errorCode;
+        this.serverStatusCode = serverStatusCode;
+        this.serverMessage = serverMessage;
+    }
+
+    public AIException(String message, String errorCode, Throwable cause) {
+
+        super(message, cause);
+        this.errorCode = errorCode;
+    }
+
+    public String getErrorCode() {
+
+        return errorCode;
+    }
+
+    public String getServerMessage() {
+
+        return serverMessage;
+    }
+
+    public int getServerStatusCode() {
+
+        return serverStatusCode;
+    }
+}
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerException.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerException.java
index 39d7f25c30a9..7e5722cf2d6f 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerException.java
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerException.java
@@ -21,49 +21,25 @@
 /**
  * Client Exception class for AI service.
  */
-public class AIServerException extends Exception {
-
-    private String errorCode;
-    private String serverMessage;
-    private int serverStatusCode;
+public class AIServerException extends AIException {
 
     public AIServerException(String message, String errorCode) {
 
-        super(message);
-        this.errorCode = errorCode;
+        super(message, errorCode);
     }
 
-    public AIServerException(String message, String errorCode, int serverStatusCode, String serverMessage) {
+    public AIServerException(String message, Throwable e) {
 
-        super(message);
-        this.errorCode = errorCode;
-        this.serverStatusCode = serverStatusCode;
-        this.serverMessage = serverMessage;
+        super(message, e);
     }
 
-    public AIServerException(String message, Throwable cause) {
+    public AIServerException(String message, String errorCode, int serverStatusCode, String serverMessage) {
 
-        super(message, cause);
+        super(message, errorCode, serverStatusCode, serverMessage);
     }
 
     public AIServerException(String message, String errorCode, Throwable cause) {
 
-        super(message, cause);
-        this.errorCode = errorCode;
-    }
-
-    public String getErrorCode() {
-
-        return errorCode;
-    }
-
-    public String getServerMessage() {
-
-        return serverMessage;
-    }
-
-    public int getServerStatusCode() {
-
-        return serverStatusCode;
+        super(message, errorCode, cause);
     }
 }
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManager.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManager.java
index 43f5e89e33c6..bfe2b48f5cfb 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManager.java
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManager.java
@@ -64,7 +64,7 @@
  */
 public class AIAccessTokenManager {
 
-    private static volatile AIAccessTokenManager instance;  // Volatile for thread safety.
+    private static volatile AIAccessTokenManager instance;
     private static final Object lock = new Object();  // Lock for synchronization.
 
     private static final Log LOG = LogFactory.getLog(AIAccessTokenManager.class);
@@ -83,7 +83,7 @@ private AIAccessTokenManager() {
         String decodedString = new String(decodedBytes, StandardCharsets.UTF_8);
         String[] parts = decodedString.split(":");
         if (parts.length == 2) {
-            this.clientId = parts[0]; // Extract clientId.
+            this.clientId = parts[0];
         } else {
             throw new IllegalArgumentException("Invalid AI service key.");
         }
@@ -208,6 +208,7 @@ public String requestAccessToken() throws AIServerException {
         }
 
         private static int readIntProperty(String key, int defaultValue) {
+
             String value = IdentityUtil.getProperty(key);
             return value != null ? Integer.parseInt(value) : defaultValue;
         }
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtil.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtil.java
index 5a540bc7cdf6..8eb569564952 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtil.java
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/main/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtil.java
@@ -187,9 +187,10 @@ private static Map<String, Object> convertJsonStringToMap(String jsonString) thr
         }
     }
 
-    protected static HttpResponseWrapper executeHttpRequest(HttpUriRequest httpRequest)
+    private static HttpResponseWrapper executeHttpRequest(HttpUriRequest httpRequest)
             throws IOException {
 
+        // Here we don't close the client connection since we are using a connection pool.
         HttpResponse httpResponse = httpClient.execute(httpRequest);
         int status = httpResponse.getStatusLine().getStatusCode();
         String response = EntityUtils.toString(httpResponse.getEntity());
@@ -197,6 +198,7 @@ protected static HttpResponseWrapper executeHttpRequest(HttpUriRequest httpReque
     }
 
     private static int readIntProperty(String key, int defaultValue) {
+
         String value = IdentityUtil.getProperty(key);
         return value != null ? Integer.parseInt(value) : defaultValue;
     }
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstantsTest.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstantsTest.java
deleted file mode 100644
index 4252283e28da..000000000000
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/constants/AIConstantsTest.java
+++ /dev/null
@@ -1,65 +0,0 @@
-/*
- * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
- *
- * WSO2 LLC. licenses this file to you under the Apache License,
- * Version 2.0 (the "License"); you may not use this file except
- * in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-package org.wso2.carbon.identity.ai.service.mgt.constants;
-
-import org.testng.annotations.Test;
-
-import static org.testng.AssertJUnit.assertEquals;
-
-/**
- * Test class for AIConstants.
- */
-public class AIConstantsTest {
-
-    @Test
-    public void testErrorMessages() {
-        AIConstants.ErrorMessages errorMessage = AIConstants.ErrorMessages.MAXIMUM_RETRIES_EXCEEDED;
-        assertEquals(errorMessage.getCode(), "AI-10000");
-        assertEquals(errorMessage.getMessage(), "Maximum retries exceeded to retrieve the access token.");
-        assertEquals(errorMessage.toString(), "AI-10000:Maximum retries exceeded to retrieve the access token.");
-
-        errorMessage = AIConstants.ErrorMessages.UNABLE_TO_ACCESS_AI_SERVICE_WITH_RENEW_ACCESS_TOKEN;
-        assertEquals(errorMessage.getCode(), "AI-10003");
-        assertEquals(errorMessage.getMessage(), "Unable to access the AI service with the renewed access token.");
-        assertEquals(errorMessage.toString(), "AI-10003:Unable to access the AI service with " +
-                "the renewed access token.");
-
-        errorMessage = AIConstants.ErrorMessages.REQUEST_TIMEOUT;
-        assertEquals(errorMessage.getCode(), "AI-10004");
-        assertEquals(errorMessage.getMessage(), "Request to the AI service timed out.");
-        assertEquals(errorMessage.toString(), "AI-10004:Request to the AI service timed out.");
-
-        errorMessage = AIConstants.ErrorMessages.ERROR_RETRIEVING_ACCESS_TOKEN;
-        assertEquals(errorMessage.getCode(), "AI-10007");
-        assertEquals(errorMessage.getMessage(), "Error occurred while retrieving the access token.");
-        assertEquals(errorMessage.toString(), "AI-10007:Error occurred while retrieving the access token.");
-
-        errorMessage = AIConstants.ErrorMessages.CLIENT_ERROR_WHILE_CONNECTING_TO_AI_SERVICE;
-        assertEquals(errorMessage.getCode(), "AI-10008");
-        assertEquals(errorMessage.getMessage(), "Client error occurred for %s tenant while connecting to AI service.");
-        assertEquals(errorMessage.toString(), "AI-10008:Client error occurred for %s tenant while" +
-                " connecting to AI service.");
-
-        errorMessage = AIConstants.ErrorMessages.SERVER_ERROR_WHILE_CONNECTING_TO_AI_SERVICE;
-        assertEquals(errorMessage.getCode(), "AI-10009");
-        assertEquals(errorMessage.getMessage(), "Server error occurred for %s tenant while connecting to AI service.");
-        assertEquals(errorMessage.toString(), "AI-10009:Server error occurred for %s tenant while " +
-                "connecting to AI service.");
-    }
-}
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientExceptionTest.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientExceptionTest.java
deleted file mode 100644
index 6cd9a09546e9..000000000000
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIClientExceptionTest.java
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
- *
- * WSO2 LLC. licenses this file to you under the Apache License,
- * Version 2.0 (the "License"); you may not use this file except
- * in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-package org.wso2.carbon.identity.ai.service.mgt.exceptions;
-
-import org.testng.annotations.Test;
-
-import static org.testng.AssertJUnit.assertEquals;
-import static org.testng.AssertJUnit.assertNull;
-
-/**
- * Test class for AIClientException.
- */
-public class AIClientExceptionTest {
-
-    @Test
-    public void testAIClientExceptionWithMessageAndErrorCode() {
-
-        AIClientException exception = new AIClientException("Test message", "AI_10001");
-        assertEquals("Test message", exception.getMessage());
-        assertEquals("AI_10001", exception.getErrorCode());
-        assertNull(exception.getServerMessage());
-    }
-
-    @Test
-    public void testAIClientExceptionWithHttpResponseWrapperMessageAndErrorCode() {
-
-        int statusCode = 400;
-        String message = "Test message";
-        AIClientException exception = new AIClientException("Test message", "AI_10002", statusCode, message);
-        assertEquals("Test message", exception.getMessage());
-        assertEquals("AI_10002", exception.getErrorCode());
-        assertEquals(statusCode, exception.getServerStatusCode());
-        assertEquals(message, exception.getServerMessage());
-    }
-
-    @Test
-    public void testAIClientExceptionWithMessageErrorCodeAndCause() {
-        Throwable cause = new Throwable("Cause message");
-        AIClientException exception = new AIClientException("Test message", "AI_10003", cause);
-        assertEquals("Test message", exception.getMessage());
-        assertEquals("AI_10003", exception.getErrorCode());
-        assertEquals(cause, exception.getCause());
-    }
-}
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerExceptionTest.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerExceptionTest.java
deleted file mode 100644
index 5ce4298ad2b5..000000000000
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/exceptions/AIServerExceptionTest.java
+++ /dev/null
@@ -1,70 +0,0 @@
-/*
- * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
- *
- * WSO2 LLC. licenses this file to you under the Apache License,
- * Version 2.0 (the "License"); you may not use this file except
- * in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-package org.wso2.carbon.identity.ai.service.mgt.exceptions;
-
-import org.testng.annotations.Test;
-
-import static org.testng.AssertJUnit.assertEquals;
-
-/**
- * Test class for AIServerException.
- */
-public class AIServerExceptionTest {
-
-    @Test
-    public void testAIServerExceptionWithMessageAndErrorCode() {
-
-        AIServerException exception = new AIServerException("Test message", "AI_20001");
-        assertEquals("Test message", exception.getMessage());
-        assertEquals("AI_20001", exception.getErrorCode());
-
-    }
-
-    @Test
-    public void testAIServerExceptionWithAIServerMessageAndErrorCode() {
-
-        String message = "Test message";
-        String errorCode = "AI_20002";
-        int statusCode = 500;
-        AIServerException exception = new AIServerException("Test message", errorCode, statusCode, message);
-        assertEquals("Test message", exception.getMessage());
-        assertEquals(errorCode, exception.getErrorCode());
-        assertEquals(statusCode, exception.getServerStatusCode());
-        assertEquals(message, exception.getServerMessage());
-    }
-
-    @Test
-    public void testAIServerExceptionWithMessageAndCause() {
-
-        Throwable cause = new Throwable("Cause message");
-        AIServerException exception = new AIServerException("Test message", cause);
-        assertEquals("Test message", exception.getMessage());
-        assertEquals(cause, exception.getCause());
-    }
-
-    @Test
-    public void testAIServerExceptionWithMessageErrorCodeAndCause() {
-
-        Throwable cause = new Throwable("Cause message");
-        AIServerException exception = new AIServerException("Test message", "AI_20003", cause);
-        assertEquals("Test message", exception.getMessage());
-        assertEquals("AI_20003", exception.getErrorCode());
-        assertEquals(cause, exception.getCause());
-    }
-}
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManagerTest.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManagerTest.java
index 05aba138565a..77c9bd19976b 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManagerTest.java
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/token/AIAccessTokenManagerTest.java
@@ -20,11 +20,16 @@
 
 import com.github.tomakehurst.wiremock.WireMockServer;
 import com.github.tomakehurst.wiremock.http.Fault;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 import org.testng.Assert;
+import org.testng.annotations.AfterClass;
 import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 import org.wso2.carbon.identity.ai.service.mgt.exceptions.AIServerException;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
 
 import java.util.Base64;
 
@@ -33,26 +38,39 @@
 import static com.github.tomakehurst.wiremock.client.WireMock.postRequestedFor;
 import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
 import static com.github.tomakehurst.wiremock.core.WireMockConfiguration.wireMockConfig;
+import static org.wso2.carbon.identity.ai.service.mgt.constants.AIConstants.AI_TOKEN_CONNECTION_REQUEST_TIMEOUT_PROPERTY_NAME;
 
+/**
+ * Test class for AIAccessTokenManager.
+ */
 public class AIAccessTokenManagerTest {
 
     private WireMockServer wireMockServer;
     private AIAccessTokenManager tokenManager;
+    private MockedStatic<IdentityUtil> mockedStatic;
+
+    @BeforeClass
+    public void init() {
+
+        mockedStatic = Mockito.mockStatic(IdentityUtil.class);
+        mockedStatic.when(() -> IdentityUtil.getProperty(AI_TOKEN_CONNECTION_REQUEST_TIMEOUT_PROPERTY_NAME))
+                .thenReturn("2000");
+    }
 
     @BeforeMethod
     public void setUp() throws Exception {
 
-        // Reset the singleton instance
+        // Reset the singleton instance.
         resetSingletonInstance(AIAccessTokenManager.class, "instance");
     }
 
     private void startWireMockServer() throws Exception {
 
-        // Start WireMock server
+        // Start WireMock server.
         wireMockServer = new WireMockServer(wireMockConfig().dynamicPort());
         wireMockServer.start();
 
-        // Set the AI_TOKEN_ENDPOINT to WireMock's base URL
+        // Set the AI_TOKEN_ENDPOINT to WireMock's base URL.
         setStaticField(AIAccessTokenManager.class, "AI_TOKEN_ENDPOINT", wireMockServer.baseUrl() + "/token");
     }
 
@@ -69,23 +87,14 @@ private void resetSingletonInstance(Class<?> clazz, String fieldName) throws Exc
         instanceField.set(null, null); // Reset the static field to null
     }
 
-    @AfterMethod
-    public void tearDown() {
-
-        if (wireMockServer != null) {
-            wireMockServer.stop();
-            wireMockServer = null;
-        }
-    }
-
     @Test
-    public void testGetAccessToken_Success() throws Exception {
+    public void testGetAccessTokenSuccess() throws Exception {
 
         startWireMockServer();
         setAiServiceKey("testClientId:testClientSecret");
         tokenManager = AIAccessTokenManager.getInstance();
 
-        // Arrange: Mock a successful token response
+        // Mock a successful token response
         String expectedAccessToken = "mockedAccessToken";
         wireMockServer.stubFor(post(urlEqualTo("/token"))
                 .willReturn(aResponse()
@@ -93,34 +102,30 @@ public void testGetAccessToken_Success() throws Exception {
                         .withHeader("Content-Type", "application/json")
                         .withBody("{\"access_token\":\"" + expectedAccessToken + "\"}")));
 
-        // Act: Call getAccessToken
         String accessToken = tokenManager.getAccessToken(false);
 
-        // Assert: Verify the access token
         Assert.assertEquals(accessToken, expectedAccessToken, "Access token should match the mocked value.");
         Assert.assertEquals(tokenManager.getClientId(), "testClientId", "Client ID should match the mocked value.");
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testGetAccessToken_Unauthorized() throws Exception {
+    public void testGetAccessTokenUnauthorized() throws Exception {
 
         startWireMockServer();
         setAiServiceKey("testClientId:testClientSecret");
         tokenManager = AIAccessTokenManager.getInstance();
 
-        // Arrange: Mock a 401 Unauthorized response
         wireMockServer.stubFor(post(urlEqualTo("/token"))
                 .willReturn(aResponse()
                         .withStatus(401)
                         .withHeader("Content-Type", "application/json")
                         .withBody("{\"error\":\"Unauthorized\"}")));
 
-        // Act: Call getAccessToken and expect an exception
         tokenManager.getAccessToken(false);
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testGetAccessToken_ServerError() throws Exception {
+    public void testGetAccessTokenServerError() throws Exception {
 
         startWireMockServer();
         setAiServiceKey("testClientId:testClientSecret");
@@ -133,12 +138,11 @@ public void testGetAccessToken_ServerError() throws Exception {
                         .withHeader("Content-Type", "application/json")
                         .withBody("{\"error\":\"Internal Server Error\"}")));
 
-        // Act: Call getAccessToken and expect an exception
         tokenManager.getAccessToken(false);
     }
 
     @Test
-    public void testGetAccessToken_Renewal() throws Exception {
+    public void testGetAccessTokenRenewal() throws Exception {
 
         startWireMockServer();
         setAiServiceKey("testClientId:testClientSecret");
@@ -152,15 +156,13 @@ public void testGetAccessToken_Renewal() throws Exception {
                         .withHeader("Content-Type", "application/json")
                         .withBody("{\"access_token\":\"" + newAccessToken + "\"}")));
 
-        // Act: Call getAccessToken with renewal.
         String accessToken = tokenManager.getAccessToken(true);
 
-        // Assert: Verify the new access token.
         Assert.assertEquals(accessToken, newAccessToken, "Access token should match the renewed mocked value.");
     }
 
     @Test
-    public void testGetAccessToken_ExistingTokenReturnsWithoutRenewal() throws Exception {
+    public void testGetAccessTokenExistingTokenReturnsWithoutRenewal() throws Exception {
 
         startWireMockServer();
         setAiServiceKey("testClientId:testClientSecret");
@@ -186,7 +188,6 @@ public void testGetAccessToken_ExistingTokenReturnsWithoutRenewal() throws Excep
         // This should NOT call the token endpoint again; it should return the cached token.
         String secondCallToken = tokenManager.getAccessToken(false);
 
-        // Assert
         Assert.assertEquals(secondCallToken, existingAccessToken,
                 "Second call should return the same token without making a new request.");
         // Verify that no new requests to the token endpoint were made after the first call.
@@ -195,7 +196,7 @@ public void testGetAccessToken_ExistingTokenReturnsWithoutRenewal() throws Excep
 
     @Test(expectedExceptions = AIServerException.class,
             expectedExceptionsMessageRegExp = "Failed to obtain access token after.*attempts.*")
-    public void testGetAccessToken_MaxRetriesExceeded() throws Exception {
+    public void testGetAccessTokenMaxRetriesExceeded() throws Exception {
 
         startWireMockServer();
         setAiServiceKey("testClientId:testClientSecret");
@@ -212,9 +213,8 @@ public void testGetAccessToken_MaxRetriesExceeded() throws Exception {
 
     @Test(expectedExceptions = AIServerException.class,
             expectedExceptionsMessageRegExp = "Error executing token request:.*")
-    public void testGetAccessToken_IOException() throws Exception {
+    public void testGetAccessTokenIOException() throws Exception {
 
-        // Start WireMock & set fields as usual.
         startWireMockServer();
         setAiServiceKey("testClientId:testClientSecret");
         tokenManager = AIAccessTokenManager.getInstance();
@@ -240,38 +240,31 @@ public void testInvalidAIServiceKey() throws Exception {
     }
 
     @Test
-    public void testGetInstance_FirstTimeCreation() throws Exception {
+    public void testGetInstanceFirstTimeCreation() throws Exception {
 
-        // 1) Reset the singleton to ensure it's null:
         setAiServiceKey("testClientId:testClientSecret");
         setStaticField(AIAccessTokenManager.class, "AI_TOKEN_ENDPOINT", "http://localhost.com/token");
-
-        // 2) Call getInstance() the first time.
         AIAccessTokenManager firstCallInstance = AIAccessTokenManager.getInstance();
-
-        // 3) Assert that the manager is created.
         Assert.assertNotNull(firstCallInstance, "First call to getInstance() should create a new instance.");
     }
 
     @Test
-    public void testGetInstance_SubsequentCallsReturnSameInstance() throws Exception {
+    public void testGetInstanceSubsequentCallsReturnSameInstance() throws Exception {
 
         setAiServiceKey("testClientId:testClientSecret");
         setStaticField(AIAccessTokenManager.class, "AI_TOKEN_ENDPOINT", "http://localhost.com/token");
 
-        // 1) Reset the singleton to ensure it's null, then create it once.
+        // Reset the singleton to ensure it's null, then create it once.
         AIAccessTokenManager firstCallInstance = AIAccessTokenManager.getInstance();
 
-        // 2) Call getInstance() again.
         AIAccessTokenManager secondCallInstance = AIAccessTokenManager.getInstance();
 
-        // 3) Verify that the second call did NOT re-create the object.
+        // Verify that the second call did NOT re-create the object.
         Assert.assertNotNull(secondCallInstance, "Second call should still return an instance.");
         Assert.assertEquals(secondCallInstance, firstCallInstance,
                 "Both calls should return the exact same singleton instance.");
     }
 
-
     private void setStaticField(Class<?> clazz, String fieldName, String value) throws Exception {
 
         java.lang.reflect.Field field = clazz.getDeclaredField(fieldName);
@@ -280,7 +273,21 @@ private void setStaticField(Class<?> clazz, String fieldName, String value) thro
         java.lang.reflect.Field modifiersField = java.lang.reflect.Field.class.getDeclaredField("modifiers");
         modifiersField.setAccessible(true);
         modifiersField.setInt(field, field.getModifiers() & ~java.lang.reflect.Modifier.FINAL);
+        field.set(null, value);
+    }
+
+    @AfterMethod
+    public void tearDown() {
+
+        if (wireMockServer != null) {
+            wireMockServer.stop();
+            wireMockServer = null;
+        }
+    }
+
+    @AfterClass
+    public void destroy() {
 
-        field.set(null, value); // Set the static field
+        mockedStatic.close();
     }
 }
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtilTest.java b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtilTest.java
index 6f944feda381..4d561e056c33 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtilTest.java
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/java/org/wso2/carbon/identity/ai/service/mgt/util/AIHttpClientUtilTest.java
@@ -25,8 +25,11 @@
 import org.apache.http.client.methods.HttpUriRequest;
 import org.mockito.Mock;
 import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 import org.testng.Assert;
+import org.testng.annotations.AfterClass;
 import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 import org.wso2.carbon.base.CarbonBaseConstants;
@@ -34,6 +37,7 @@
 import org.wso2.carbon.identity.ai.service.mgt.exceptions.AIClientException;
 import org.wso2.carbon.identity.ai.service.mgt.exceptions.AIServerException;
 import org.wso2.carbon.identity.ai.service.mgt.token.AIAccessTokenManager;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
 
 import java.nio.file.Paths;
 import java.util.HashMap;
@@ -54,6 +58,7 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.openMocks;
+import static org.wso2.carbon.identity.ai.service.mgt.constants.AIConstants.HTTP_CONNECTION_POOL_SIZE_PROPERTY_NAME;
 import static org.wso2.carbon.identity.ai.service.mgt.constants.AIConstants.TENANT_CONTEXT_PREFIX;
 
 /**
@@ -68,6 +73,15 @@ public class AIHttpClientUtilTest {
     private AIAccessTokenManager mockTokenManager;
 
     private MockedStatic<AIAccessTokenManager> aiAccessTokenManagerMockedStatic;
+    private MockedStatic<IdentityUtil> identityUtilMockedStatic;
+
+    @BeforeClass
+    public void init() {
+
+        identityUtilMockedStatic = Mockito.mockStatic(IdentityUtil.class);
+        identityUtilMockedStatic.when(() -> IdentityUtil.getProperty(HTTP_CONNECTION_POOL_SIZE_PROPERTY_NAME))
+                .thenReturn("10");
+    }
 
     @BeforeMethod
     public void setUp() throws Exception {
@@ -90,9 +104,8 @@ public void setUp() throws Exception {
     }
 
     @Test
-    public void testExecuteRequest_Success() throws Exception {
+    public void testExecuteRequestSuccess() throws Exception {
 
-        // Arrange: Mock a successful response.
         String expectedResponse = "{\"result\":\"SUCCESS\"}";
         String path = "/test-endpoint";
         String fullPath = TENANT_CONTEXT_PREFIX + clientId + path; // This is the path that AIHttpClientUtil will use.
@@ -102,7 +115,6 @@ public void testExecuteRequest_Success() throws Exception {
                         .withHeader("Content-Type", "application/json")
                         .withBody(expectedResponse)));
 
-        // Act: Execute the HTTP request.
         String baseUrl = wireMockServer.baseUrl();
         Map<String, Object> resultMap = AIHttpClientUtil.executeRequest(
                 baseUrl,
@@ -111,20 +123,17 @@ public void testExecuteRequest_Success() throws Exception {
                 null
         );
 
-        // Assert: Verify the response.
         Assert.assertEquals(resultMap.get("result"), "SUCCESS");
         wireMockServer.verify(getRequestedFor(urlEqualTo(fullPath)));
     }
 
     @Test
-    public void testExecuteRequest_PostSuccess() throws Exception {
+    public void testExecuteRequestPostSuccess() throws Exception {
 
-        // Arrange: Mock a successful response.
         String expectedResponse = "{\"result\":\"POST_SUCCESS\"}";
         String path = "/test-endpoint";
         String fullPath = TENANT_CONTEXT_PREFIX + clientId + path; // This is the path that AIHttpClientUtil will use.
 
-        // Define the request body.
         String requestBody = "{\"key\":\"value\"}";
 
         // Stub the POST request with the expected response.
@@ -136,7 +145,6 @@ public void testExecuteRequest_PostSuccess() throws Exception {
                         .withHeader("Content-Type", "application/json")
                         .withBody(expectedResponse)));
 
-        // Act: Execute the HTTP request.
         String baseUrl = wireMockServer.baseUrl();
         Map<String, String> requestBodyMap = new HashMap<>();
         requestBodyMap.put("key", "value");
@@ -144,10 +152,9 @@ public void testExecuteRequest_PostSuccess() throws Exception {
                 baseUrl,
                 path,
                 HttpPost.class,
-                requestBodyMap // Pass the request body as a map.
+                requestBodyMap
         );
 
-        // Assert: Verify the response.
         Assert.assertEquals(resultMap.get("result"), "POST_SUCCESS");
 
         // Verify that the POST request was made with the correct path and body.
@@ -156,20 +163,17 @@ public void testExecuteRequest_PostSuccess() throws Exception {
                 .withRequestBody(equalToJson(requestBody)));
     }
 
-
     @Test(expectedExceptions = AIClientException.class)
-    public void testExecuteRequest_ClientError() throws Exception {
+    public void testExecuteRequestClientError() throws Exception {
 
-        // Arrange: Mock a client error response
         String path = "/test-endpoint";
         String fullPath = TENANT_CONTEXT_PREFIX + clientId + path; // This is the path that AIHttpClientUtil will use.
         wireMockServer.stubFor(get(urlEqualTo(fullPath))
                 .willReturn(aResponse()
-                        .withStatus(400) // Client error status.
+                        .withStatus(400)
                         .withHeader("Content-Type", "application/json")
                         .withBody("Bad Request")));
 
-        // Act & Assert: Expect AIClientException.
         String baseUrl = wireMockServer.baseUrl();
         AIHttpClientUtil.executeRequest(
                 baseUrl,
@@ -180,14 +184,13 @@ public void testExecuteRequest_ClientError() throws Exception {
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testExecuteRequest_ServerError() throws Exception {
+    public void testExecuteRequestServerError() throws Exception {
 
-        // Arrange: Mock a server error response.
         String path = "/test-endpoint";
         String fullPath = TENANT_CONTEXT_PREFIX + clientId + path; // This is the path that AIHttpClientUtil will use.
         wireMockServer.stubFor(get(urlEqualTo(fullPath))
                 .willReturn(aResponse()
-                        .withStatus(500) // Simulate a server error.
+                        .withStatus(500)
                         .withHeader("Content-Type", "text/plain")
                         .withBody("Internal Server Error")));
 
@@ -202,7 +205,7 @@ public void testExecuteRequest_ServerError() throws Exception {
     }
 
     @Test
-    public void testExecuteRequest_TokenRenewal() throws Exception {
+    public void testExecuteRequestTokenRenewal() throws Exception {
 
         // Mock the AccessTokenManager to simulate token renewal.
         when(mockTokenManager.getAccessToken(true)).thenReturn("newToken");
@@ -217,7 +220,7 @@ public void testExecuteRequest_TokenRenewal() throws Exception {
                 .inScenario("Token Renewal")
                 .whenScenarioStateIs(STARTED)
                 .willReturn(aResponse()
-                        .withStatus(401) // Unauthorized.
+                        .withStatus(401)
                         .withHeader("Content-Type", "application/json")
                         .withBody("Unauthorized"))
                 .willSetStateTo("Token Renewed")); // Transition to the next state.
@@ -227,7 +230,7 @@ public void testExecuteRequest_TokenRenewal() throws Exception {
                 .inScenario("Token Renewal")
                 .whenScenarioStateIs("Token Renewed")
                 .willReturn(aResponse()
-                        .withStatus(200) // Success
+                        .withStatus(200)
                         .withHeader("Content-Type", "application/json")
                         .withBody(expectedResponse)));
 
@@ -251,7 +254,7 @@ public void testExecuteRequest_TokenRenewal() throws Exception {
     }
 
     @Test(expectedExceptions = AIClientException.class)
-    public void testExecuteRequest_TokenRenewal_ErrorAfterRenewal() throws Exception {
+    public void testExecuteRequestTokenRenewalErrorAfterRenewal() throws Exception {
         // Mock the AccessTokenManager to simulate token renewal.
         when(mockTokenManager.getAccessToken(true)).thenReturn("newToken");
 
@@ -264,7 +267,7 @@ public void testExecuteRequest_TokenRenewal_ErrorAfterRenewal() throws Exception
                 .inScenario("Token Renewal with Error")
                 .whenScenarioStateIs(STARTED)
                 .willReturn(aResponse()
-                        .withStatus(401) // Unauthorized.
+                        .withStatus(401)
                         .withHeader("Content-Type", "application/json")
                         .withBody("Unauthorized"))
                 .willSetStateTo("Token Renewed")); // Transition to the next state.
@@ -274,7 +277,7 @@ public void testExecuteRequest_TokenRenewal_ErrorAfterRenewal() throws Exception
                 .inScenario("Token Renewal with Error")
                 .whenScenarioStateIs("Token Renewed")
                 .willReturn(aResponse()
-                        .withStatus(400) // Client-side error.
+                        .withStatus(400)
                         .withHeader("Content-Type", "application/json")
                         .withBody("{\"error\":\"Bad Request\"}"))); // Error response body.
 
@@ -289,7 +292,7 @@ public void testExecuteRequest_TokenRenewal_ErrorAfterRenewal() throws Exception
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testExecuteRequest_IOException() throws Exception {
+    public void testExecuteRequestIOException() throws Exception {
 
         // Arrange: Mock a server that simulates a connection reset.
         String path = "/test-endpoint";
@@ -311,7 +314,7 @@ public void testExecuteRequest_IOException() throws Exception {
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testExecuteRequest_ExecutionException() throws Exception {
+    public void testExecuteRequestExecutionException() throws Exception {
 
         // Arrange: Mock a server that simulates an unexpected response.
         String path = "/test-endpoint";
@@ -333,7 +336,7 @@ public void testExecuteRequest_ExecutionException() throws Exception {
     }
 
     @Test(expectedExceptions = IllegalArgumentException.class)
-    public void testExecuteRequest_UnsupportedRequestType() throws Exception {
+    public void testExecuteRequestUnsupportedRequestType() throws Exception {
 
         // Arrange: Define the path and base URL.
         String path = "/test-endpoint";
@@ -349,7 +352,7 @@ public void testExecuteRequest_UnsupportedRequestType() throws Exception {
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testExecuteRequest_UnauthorizedAfterTokenRenewal() throws Exception {
+    public void testExecuteRequestUnauthorizedAfterTokenRenewal() throws Exception {
 
         // Mock the AccessTokenManager for token renewal.
         when(mockTokenManager.getAccessToken(true)).thenReturn("newToken");
@@ -358,26 +361,25 @@ public void testExecuteRequest_UnauthorizedAfterTokenRenewal() throws Exception
         String path = "/test-endpoint";
         String fullPath = TENANT_CONTEXT_PREFIX + clientId + path; // This is the path that AIHttpClientUtil will use.
 
-        // First response: 401 Unauthorized
+        // First response: 401 Unauthorized.
         wireMockServer.stubFor(get(urlEqualTo(fullPath))
                 .inScenario("Token Renewal Fails")
                 .whenScenarioStateIs(STARTED)
                 .willReturn(aResponse()
-                        .withStatus(401) // Unauthorized
+                        .withStatus(401)
                         .withHeader("Content-Type", "application/json")
                         .withBody("Unauthorized"))
                 .willSetStateTo("Retry"));
 
-        // Second response: 401 Unauthorized again
+        // Second response: 401 Unauthorized again.
         wireMockServer.stubFor(get(urlEqualTo(fullPath))
                 .inScenario("Token Renewal Fails")
                 .whenScenarioStateIs("Retry")
                 .willReturn(aResponse()
-                        .withStatus(401) // Still Unauthorized
+                        .withStatus(401) // Still Unauthorized.
                         .withHeader("Content-Type", "application/json")
                         .withBody("Still Unauthorized")));
 
-        // Act: Execute the HTTP request
         String baseUrl = wireMockServer.baseUrl();
         AIHttpClientUtil.executeRequest(
                 baseUrl,
@@ -388,16 +390,15 @@ public void testExecuteRequest_UnauthorizedAfterTokenRenewal() throws Exception
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testExecuteRequest_JsonParsingError() throws Exception {
+    public void testExecuteRequestJsonParsingError() throws Exception {
 
-        // Arrange: Define paths.
         String path = "/test-endpoint";
         String fullPath = TENANT_CONTEXT_PREFIX + clientId + path; // This is the path that AIHttpClientUtil will use.
 
         // Mock the server to return invalid JSON.
         wireMockServer.stubFor(get(urlEqualTo(fullPath))
                 .willReturn(aResponse()
-                        .withStatus(200) // Simulate a successful response.
+                        .withStatus(200)
                         .withHeader("Content-Type", "application/json")
                         .withBody("{ invalid json }"))); // Invalid JSON.
 
@@ -412,20 +413,19 @@ public void testExecuteRequest_JsonParsingError() throws Exception {
     }
 
     @Test(expectedExceptions = AIServerException.class)
-    public void testExecuteRequest_FailedTokenRenewal() throws Exception {
+    public void testExecuteRequestFailedTokenRenewal() throws Exception {
 
         // Mock the AccessTokenManager to simulate failed token renewal.
         when(mockTokenManager.getAccessToken(false)).thenReturn("oldToken");
         when(mockTokenManager.getAccessToken(true)).thenReturn(null); // Simulate failed token renewal.
 
-        // Arrange: Define paths.
         String path = "/test-endpoint";
         String fullPath = TENANT_CONTEXT_PREFIX + clientId + path; // This is the path that AIHttpClientUtil will use.
 
         // Mock the server to return 401 Unauthorized.
         wireMockServer.stubFor(get(urlEqualTo(fullPath))
                 .willReturn(aResponse()
-                        .withStatus(401) // Unauthorized
+                        .withStatus(401)
                         .withHeader("Content-Type", "application/json")
                         .withBody("Unauthorized")));
 
@@ -460,4 +460,10 @@ public void tearDown() {
         PrivilegedCarbonContext.endTenantFlow();
         wireMockServer.stop();
     }
+
+    @AfterClass
+    public void destroy() {
+
+        identityUtilMockedStatic.close();
+    }
 }
diff --git a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/resources/testng.xml b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/resources/testng.xml
index 67ddc6f4f5bb..aa1aa752ae1c 100644
--- a/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/resources/testng.xml
+++ b/components/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt/src/test/resources/testng.xml
@@ -21,9 +21,6 @@
 <suite name="org.wso2.carbon.ai.service.mgt">
     <test name="org.wso2.carbon.identity.application.mgt.test" preserve-order="true" parallel="false">
         <classes>
-            <class name="org.wso2.carbon.identity.ai.service.mgt.constants.AIConstantsTest"/>
-            <class name="org.wso2.carbon.identity.ai.service.mgt.exceptions.AIClientExceptionTest"/>
-            <class name="org.wso2.carbon.identity.ai.service.mgt.exceptions.AIServerExceptionTest"/>
             <class name="org.wso2.carbon.identity.ai.service.mgt.util.AIHttpClientUtilTest"/>
             <class name="org.wso2.carbon.identity.ai.service.mgt.token.AIAccessTokenManagerTest"/>
         </classes>
diff --git a/components/ai-services-mgt/pom.xml b/components/ai-services-mgt/pom.xml
index af8e93a6245a..0fead50c10f7 100644
--- a/components/ai-services-mgt/pom.xml
+++ b/components/ai-services-mgt/pom.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!--
-  ~ Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+  ~ Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
   ~
   ~ WSO2 LLC. licenses this file to you under the Apache License,
   ~ Version 2.0 (the "License"); you may not use this file except
@@ -25,7 +25,7 @@
     <parent>
         <groupId>org.wso2.carbon.identity.framework</groupId>
         <artifactId>identity-framework</artifactId>
-        <version>7.7.87-SNAPSHOT</version>
+        <version>7.7.91-SNAPSHOT</version>
         <relativePath>../../pom.xml</relativePath>
     </parent>
 
diff --git a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManager.java b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManager.java
index cb2e6ba8bf22..876d5246dbec 100644
--- a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManager.java
+++ b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManager.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -30,12 +30,42 @@
  */
 public interface LoginFlowAIManager {
 
-    String generateAuthenticationSequence(String userQuery, JSONArray userClaims, JSONObject availableAuthenticators)
+    /**
+     * Generates an authentication sequence using the LoginFlow AI service.
+     *
+     * @param userQuery               The user query. This is a string that contain the requested authentication
+     *                                flow by the user.
+     * @param userClaimsMetaData      The user claims metadata. This is a JSON array that contains the user
+     *                                claims available
+     *                                for that organization.
+     * @param availableAuthenticators The available authenticators of the organization.
+     * @return Operation ID of the generated authentication sequence.
+     * @throws AIServerException When a server error occurs while connecting to the LoginFlow AI service.
+     * @throws AIClientException When a client error occurs while generating the authentication sequence.
+     */
+    String generateAuthenticationSequence(String userQuery, JSONArray userClaimsMetaData,
+                                          JSONObject availableAuthenticators)
             throws AIServerException, AIClientException;
 
+    /**
+     * Retrieves the status of the authentication sequence generation operation.
+     *
+     * @param operationId The operation ID of the authentication sequence generation operation.
+     * @return A Json representation of the status' that are completed, pending, or failed.
+     * @throws AIServerException When a server error occurs while connecting to the LoginFlow AI service.
+     * @throws AIClientException When a client error occurs while retrieving the authentication sequence
+     */
     Map<String, Object> getAuthenticationSequenceGenerationStatus(String operationId) throws AIServerException,
             AIClientException;
 
+    /**
+     * Retrieves the result of the authentication sequence generation operation.
+     *
+     * @param operationId The operation ID of the authentication sequence generation operation.
+     * @return The result of the authentication sequence generation operation.
+     * @throws AIServerException When a server error occurs while connecting to the LoginFlow AI service.
+     * @throws AIClientException When a client error occurs while retrieving the authentication sequence
+     */
     Map<String, Object> getAuthenticationSequenceGenerationResult(String operationId) throws AIServerException,
             AIClientException;
 }
diff --git a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerImpl.java b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerImpl.java
index 28e77c46af35..05c413b84076 100644
--- a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerImpl.java
+++ b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerImpl.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -56,20 +56,8 @@ public class LoginFlowAIManagerImpl implements LoginFlowAIManager {
 
     private static final Log LOG = LogFactory.getLog(LoginFlowAIManagerImpl.class);
 
-    /**
-     * Generates an authentication sequence using the LoginFlow AI service.
-     *
-     * @param userQuery               The user query. This is a string that contain the requested authentication
-     *                                flow by the user.
-     * @param userClaims              The user claims. This is a JSON array that contains the user claims available
-     *                                for that organization.
-     * @param availableAuthenticators The available authenticators of the organization.
-     * @return Operation ID of the generated authentication sequence.
-     * @throws AIServerException When an error occurs while connecting to the LoginFlow AI service.
-     * @throws AIClientException When an error occurs while generating the authentication sequence.
-     */
     @Override
-    public String generateAuthenticationSequence(String userQuery, JSONArray userClaims,
+    public String generateAuthenticationSequence(String userQuery, JSONArray userClaimsMetaData,
                                                  JSONObject availableAuthenticators) throws AIServerException,
             AIClientException {
 
@@ -77,11 +65,8 @@ public String generateAuthenticationSequence(String userQuery, JSONArray userCla
         Map<String, Object> requestBody = new HashMap<>();
         requestBody.put(USER_QUERY_PROPERTY, userQuery);
         try {
-            // Convert JSONArray to List.
-            List<Object> userClaimsList = objectMapper.readValue(userClaims.toString(), List.class);
-            requestBody.put(USER_CLAIM_PROPERTY, userClaimsList);
-
-            // Convert JSONObject to Map.
+            List<Object> userClaimsMetadataList = objectMapper.readValue(userClaimsMetaData.toString(), List.class);
+            requestBody.put(USER_CLAIM_PROPERTY, userClaimsMetadataList);
             Map<String, Object> authenticatorsMap = objectMapper.readValue(availableAuthenticators.toString(),
                     Map.class);
             requestBody.put(AUTHENTICATORS_PROPERTY, authenticatorsMap);
@@ -95,32 +80,14 @@ public String generateAuthenticationSequence(String userQuery, JSONArray userCla
         return (String) stringObjectMap.get(OPERATION_ID_PROPERTY);
     }
 
-    /**
-     * Retrieves the status of the authentication sequence generation operation.
-     *
-     * @param operationId The operation ID of the authentication sequence generation operation.
-     * @return A Json representation of the status' that are completed, pending, or failed.
-     * @throws AIServerException When an error occurs while connecting to the LoginFlow AI service.
-     * @throws AIClientException When an error occurs while retrieving the authentication sequence
-     *                                    generation status.
-     */
     @Override
-    public Map<String, Object> getAuthenticationSequenceGenerationStatus(String operationId) throws AIServerException,
+    public Map<String, Object> getAuthenticationSequenceGenerationStatus(String operationId) throws  AIServerException,
             AIClientException {
 
         return executeRequest(LOGINFLOW_AI_ENDPOINT, LOGINFLOW_AI_STATUS_PATH + PATH_SEPARATOR + operationId,
                 HttpGet.class, null);
     }
 
-    /**
-     * Retrieves the result of the authentication sequence generation operation.
-     *
-     * @param operationId The operation ID of the authentication sequence generation operation.
-     * @return The result of the authentication sequence generation operation.
-     * @throws AIServerException When an error occurs while connecting to the LoginFlow AI service.
-     * @throws AIClientException When an error occurs while retrieving the authentication sequence
-     *                                    generation result.
-     */
     @Override
     public Map<String, Object> getAuthenticationSequenceGenerationResult(String operationId) throws AIServerException,
             AIClientException {
diff --git a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/constant/LoginFlowAIConstants.java b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/constant/LoginFlowAIConstants.java
index 7854dd4881cb..f9b5aed008c4 100644
--- a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/constant/LoginFlowAIConstants.java
+++ b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/main/java/org/wso2/carbon/identity/application/mgt/ai/constant/LoginFlowAIConstants.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -23,12 +23,15 @@
  */
 public class LoginFlowAIConstants {
 
+    private LoginFlowAIConstants() {
+
+    }
+
     public static final String OPERATION_ID_PROPERTY = "operation_id";
     public static final String USER_CLAIM_PROPERTY = "user_claims";
     public static final String USER_QUERY_PROPERTY = "user_query";
     public static final String AUTHENTICATORS_PROPERTY = "available_authenticators";
 
-
     /**
      * Enums for error messages.
      */
@@ -39,7 +42,6 @@ public enum ErrorMessages {
         SERVER_ERROR_WHILE_CONNECTING_TO_LOGINFLOW_AI_SERVICE("AILF-10009", "Server error occurred " +
                 "for %s tenant while generating authentication sequence.");
 
-
         private final String code;
         private final String message;
 
diff --git a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/test/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerTest.java b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/test/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerTest.java
index 349ee3c2ebc2..8997560567da 100644
--- a/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/test/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerTest.java
+++ b/components/application-mgt/org.wso2.carbon.identity.application.mgt/src/test/java/org/wso2/carbon/identity/application/mgt/ai/LoginFlowAIManagerTest.java
@@ -30,8 +30,6 @@
 import org.testng.annotations.Test;
 import org.wso2.carbon.base.CarbonBaseConstants;
 import org.wso2.carbon.context.PrivilegedCarbonContext;
-import org.wso2.carbon.identity.ai.service.mgt.exceptions.AIClientException;
-import org.wso2.carbon.identity.ai.service.mgt.exceptions.AIServerException;
 import org.wso2.carbon.identity.ai.service.mgt.util.AIHttpClientUtil;
 import org.wso2.carbon.identity.common.testng.realm.InMemoryRealmService;
 import org.wso2.carbon.identity.core.util.IdentityTenantUtil;
@@ -42,7 +40,6 @@
 import java.util.Map;
 
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mockStatic;
 import static org.mockito.MockitoAnnotations.openMocks;
 import static org.wso2.carbon.base.MultitenantConstants.SUPER_TENANT_DOMAIN_NAME;
@@ -68,7 +65,7 @@ public void setUp() throws UserStoreException {
     }
 
     @Test
-    public void testGenerateAuthenticationSequence_Success() throws Exception {
+    public void testGenerateAuthenticationSequenceSuccess() throws Exception {
 
         Map<String, Object> response = new HashMap<>();
         response.put("operation_id", "12345");
@@ -79,7 +76,8 @@ public void testGenerateAuthenticationSequence_Success() throws Exception {
     }
 
     @Test
-    public void testGetAuthenticationSequenceGenerationStatus_Success() throws Exception {
+    public void testGetAuthenticationSequenceGenerationStatusSuccess() throws Exception {
+
         Map<String, Object> response = new HashMap<>();
         response.put("status", "COMPLETED");
         mockSuccessfulResponse(response, HttpGet.class);
@@ -91,7 +89,7 @@ public void testGetAuthenticationSequenceGenerationStatus_Success() throws Excep
     }
 
     @Test
-    public void testGetAuthenticationSequenceGenerationResult_Success() throws Exception {
+    public void testGetAuthenticationSequenceGenerationResultSuccess() throws Exception {
 
         Map<String, Object> response = new HashMap<>();
         response.put("result", "SUCCESS");
@@ -110,13 +108,6 @@ private void mockSuccessfulResponse(Map<String, Object> responseBody, Class<?> r
         )).thenReturn(responseBody);
     }
 
-    private void mockErrorResponse(int statusCode, String responseBody) throws Exception {
-
-        aiHttpClientUtilMockedStatic.when(() -> AIHttpClientUtil.executeRequest(
-                anyString(), anyString(), any(), any()
-        )).thenThrow(statusCode >= 500 ? new AIServerException(responseBody, "ERROR_CODE") : new AIClientException(
-                responseBody, "ERROR_CODE"));
-    }
 
     private void setCarbonHome() {
 
diff --git a/features/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt.server.feature/pom.xml b/features/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt.server.feature/pom.xml
index 8fad3c087cc2..3d18d21c9a90 100644
--- a/features/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt.server.feature/pom.xml
+++ b/features/ai-services-mgt/org.wso2.carbon.identity.ai.service.mgt.server.feature/pom.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!--
-  ~ Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+  ~ Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
   ~
   ~ WSO2 LLC. licenses this file to you under the Apache License,
   ~ Version 2.0 (the "License"); you may not use this file except
@@ -24,7 +24,7 @@
     <parent>
         <groupId>org.wso2.carbon.identity.framework</groupId>
         <artifactId>ai-services-mgt-feature</artifactId>
-        <version>7.7.87-SNAPSHOT</version>
+        <version>7.7.91-SNAPSHOT</version>
         <relativePath>../pom.xml</relativePath>
     </parent>
 
diff --git a/features/ai-services-mgt/pom.xml b/features/ai-services-mgt/pom.xml
index 0b6640a4e5d3..2464a026340f 100644
--- a/features/ai-services-mgt/pom.xml
+++ b/features/ai-services-mgt/pom.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <!--
-  ~ Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+  ~ Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
   ~
   ~ WSO2 LLC. licenses this file to you under the Apache License,
   ~ Version 2.0 (the "License"); you may not use this file except
@@ -24,7 +24,7 @@
     <parent>
         <groupId>org.wso2.carbon.identity.framework</groupId>
         <artifactId>identity-framework</artifactId>
-        <version>7.7.87-SNAPSHOT</version>
+        <version>7.7.91-SNAPSHOT</version>
         <relativePath>../../pom.xml</relativePath>
     </parent>
 
diff --git a/features/identity-core/org.wso2.carbon.identity.core.server.feature/resources/identity.xml.j2 b/features/identity-core/org.wso2.carbon.identity.core.server.feature/resources/identity.xml.j2
index b8f74bd29bfc..b0312b605a0a 100644
--- a/features/identity-core/org.wso2.carbon.identity.core.server.feature/resources/identity.xml.j2
+++ b/features/identity-core/org.wso2.carbon.identity.core.server.feature/resources/identity.xml.j2
@@ -4472,7 +4472,9 @@
         <TokenEndpoint>{{ai_services.token_endpoint}}</TokenEndpoint>
         <Key>{{ai_services.key}}</Key>
         <TokenRequestRetryCount>{{ai_services.token_request_retry_count}}</TokenRequestRetryCount>
-        <TokenRequestTimeout>{{ai_services.token_request_timeout}}</TokenRequestTimeout>
+        <TokenConnectionTimeout>{{ai_services.token_connection_timeout}}</TokenConnectionTimeout>
+        <TokenConnectionRequestTimeout>{{ai_services.token_connection_request_timeout}}</TokenConnectionRequestTimeout>
+        <TokenConnectionSocketTimeout>{{ai_services.token_connection_socket_timeout}}</TokenConnectionSocketTimeout>
         <HTTPConnectionPoolSize>{{ai_services.http_connection_pool_size}}</HTTPConnectionPoolSize>
         <HTTPConnectionTimeout>{{ai_services.http_connection_timeout}}</HTTPConnectionTimeout>
         <LoginFlowAI>