mcp sse server config

pom.xml

@@ -32,19 +32,28 @@
     <spring-ai.version>1.0.0-M8</spring-ai.version>
   </properties>
   <dependencies>
-<!--    <dependency>
+    <dependency>
       <groupId>org.springframework.ai</groupId>
       <artifactId>spring-ai-starter-mcp-server-webmvc</artifactId>
-    </dependency>-->
+    </dependency>
 <!--    <dependency>
       <groupId>org.springframework.boot</groupId>
       <artifactId>spring-boot-starter-web</artifactId>
     </dependency>-->
 
-    <dependency>
+<!--    <dependency>
       <groupId>org.springframework.ai</groupId>
       <artifactId>spring-ai-starter-mcp-server</artifactId>
-    </dependency>
+    </dependency>-->
+
+<!--    <dependency>-->
+<!--      <groupId>org.springframework.boot</groupId>-->
+<!--      <artifactId>spring-boot-starter-oauth2-resource-server</artifactId>-->
+<!--    </dependency>-->
+<!--    <dependency>-->
+<!--      <groupId>org.springframework.boot</groupId>-->
+<!--      <artifactId>spring-boot-starter-oauth2-authorization-server</artifactId>-->
+<!--    </dependency>-->
 
     <dependency>
       <groupId>org.springframework</groupId>

src/main/java/io/akto/mcp/AktoMcpServer.java

@@ -1,24 +1,76 @@
 package io.akto.mcp;
 
 import io.akto.mcp.tools.AktoTools;
+import io.akto.mcp.tools.ToolMetadata;
+import io.modelcontextprotocol.server.McpSyncServerExchange;
+import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import org.springframework.ai.mcp.McpToolUtils;
+import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
+import org.springframework.ai.tool.StaticToolCallbackProvider;
+import org.springframework.ai.tool.ToolCallback;
 import org.springframework.ai.tool.ToolCallbackProvider;
+import org.springframework.ai.tool.definition.DefaultToolDefinition;
+import org.springframework.ai.tool.function.FunctionToolCallback;
+import org.springframework.ai.tool.metadata.DefaultToolMetadata;
+import org.springframework.ai.tool.method.MethodToolCallback;
 import org.springframework.ai.tool.method.MethodToolCallbackProvider;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
 import org.springframework.context.annotation.Bean;
+import org.springframework.core.ParameterizedTypeReference;
+import org.springframework.util.ReflectionUtils;
 
 @SpringBootApplication
 public class AktoMcpServer {
 
-	public static void main(String[] args) {
-		if(System.getenv("AKTO_API_KEY") == null || System.getenv("AKTO_API_KEY").isEmpty()) {
-			throw new RuntimeException("AKTO_API_KEY not found in environment variable");
-		}
-		SpringApplication.run(AktoMcpServer.class, args);
-	}
-
-	@Bean
-	public ToolCallbackProvider apiCollectionTools(AktoTools aktoTools) {
-		return MethodToolCallbackProvider.builder().toolObjects(aktoTools).build();
-	}
+    public static void main(String[] args) {
+        if (System.getenv("AKTO_API_KEY") == null || System.getenv("AKTO_API_KEY").isEmpty()) {
+            throw new RuntimeException("AKTO_API_KEY not found in environment variable");
+        }
+        SpringApplication.run(AktoMcpServer.class, args);
+    }
+
+    @Bean
+    public ToolCallbackProvider apiCollectionTools1(AktoTools aktoTools) {
+        return MethodToolCallbackProvider.builder().toolObjects(aktoTools).build();
+    }
+
+    //@Bean
+    public ToolCallbackProvider apiCollectionTools(AktoTools aktoTools) {
+        List<FunctionToolCallback<Object, Object>> tools = new ArrayList<>();
+
+        ToolMetadata.TOOL_METADATA.forEach((name, info) -> {
+            FunctionToolCallback<Object, Object> callback = FunctionToolCallback.builder(name, (a, b) -> {
+                    try {
+                        Method method;
+                        Object result;
+
+                        try {
+                            method = AktoTools.class.getMethod(name, Map.class);
+                            result = method.invoke(aktoTools, a);
+                        } catch (NoSuchMethodException e) {
+                            method = AktoTools.class.getMethod(name);
+                            result = method.invoke(aktoTools);
+                        }
+
+                        return result;
+                    } catch (Exception ex) {
+                        throw new RuntimeException("Error invoking method " + name, ex);
+                    }
+                })
+                .description(info.getDescription())
+                .inputSchema(info.getInputSchema())
+                .inputType(new ParameterizedTypeReference<Map<String, Object>>() {})
+                .toolMetadata(DefaultToolMetadata.builder().returnDirect(false).build())
+                .build();
+
+            tools.add(callback);
+        });
+
+        return new StaticToolCallbackProvider(tools);
+    }
 }

src/main/java/io/akto/mcp/config/AuthTokenContext.java

@@ -0,0 +1,18 @@
+package io.akto.mcp.config;
+
+public class AuthTokenContext {
+
+    private static final ThreadLocal<String> bearerToken = new ThreadLocal<>();
+
+    public static void setToken(String token) {
+        bearerToken.set(token);
+    }
+
+    public static String getToken() {
+        return bearerToken.get();
+    }
+
+    public static void clear() {
+        bearerToken.remove();
+    }
+}

src/main/java/io/akto/mcp/config/BearerTokenAuthFilter.java

@@ -0,0 +1,37 @@
+package io.akto.mcp.config;
+
+import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletException;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import java.io.IOException;
+import org.springframework.stereotype.Component;
+import org.springframework.web.filter.OncePerRequestFilter;
+
+@Component
+public class BearerTokenAuthFilter extends OncePerRequestFilter {
+    @Override
+    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain)
+        throws ServletException, IOException {
+        try {
+            String authHeader = request.getHeader("Authorization");
+
+//            if (authHeader != null && authHeader.startsWith("Bearer ") && !authHeader.substring(7).isBlank()) {
+//                String accessToken = authHeader.substring(7);
+//
+//                request.setAttribute("AKTO_API_KEY", accessToken);
+//                AuthTokenContext.setToken(accessToken);
+//
+//                Authentication authenticationToken = new UsernamePasswordAuthenticationToken(null, null,
+//                    null);
+//                SecurityContextHolder.getContext().setAuthentication(authenticationToken);
+//                filterChain.doFilter(request, response);
+//            } else {
+//                response.setStatus(HttpServletResponse.SC_UNAUTHORIZED);
+//            }
+            filterChain.doFilter(request, response);
+        } finally {
+            AuthTokenContext.clear();
+        }
+    }
+}

src/main/java/io/akto/mcp/config/MCPSSEServerConfig.java

@@ -1,17 +1,24 @@
 package io.akto.mcp.config;
 
-//@Configuration
-//@EnableWebMvc
-public class MCPSSEServerConfig /*implements WebMvcConfigurer*/ {
-/*    @Bean
+import com.fasterxml.jackson.databind.ObjectMapper;
+import io.modelcontextprotocol.server.transport.HttpServletSseServerTransportProvider;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.boot.web.servlet.ServletRegistrationBean;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.web.servlet.config.annotation.EnableWebMvc;
+import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;
+
+@Configuration
+@EnableWebMvc
+public class MCPSSEServerConfig implements WebMvcConfigurer {
+    @Bean
     public HttpServletSseServerTransportProvider servletSseServerTransportProvider() {
         return new HttpServletSseServerTransportProvider(new ObjectMapper(), "/mcp/message");
     }
 
     @Bean
     public ServletRegistrationBean customServletBean(HttpServletSseServerTransportProvider transportProvider) {
         return new ServletRegistrationBean(transportProvider);
-    }*/
-
-
+    }
 }

src/main/java/io/akto/mcp/config/SecurityConfig.java

@@ -0,0 +1,19 @@
+package io.akto.mcp.config;
+
+//@Configuration
+//@EnableWebSecurity
+public class SecurityConfig {
+
+//    @Autowired
+//    private BearerTokenAuthFilter bearerTokenAuthFilter;
+//
+//    @Bean
+//    SecurityFilterChain securityFilterChain(HttpSecurity http) throws Exception {
+//        return http
+//            .authorizeHttpRequests(auth -> auth.anyRequest().authenticated())
+//            .csrf(CsrfConfigurer::disable)
+//            .addFilterAfter(bearerTokenAuthFilter, UsernamePasswordAuthenticationFilter.class)
+//            .build();
+//    }
+
+}

src/main/java/io/akto/mcp/processor/ApiProcessor.java

@@ -1,5 +1,7 @@
 package io.akto.mcp.processor;
 
+import io.akto.mcp.config.AuthTokenContext;
+import jakarta.servlet.http.HttpServletRequest;
 import java.util.Collections;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -12,6 +14,8 @@
 import org.springframework.http.ResponseEntity;
 import org.springframework.stereotype.Component;
 import org.springframework.web.client.RestTemplate;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
 
 @Slf4j
 @Component
@@ -21,19 +25,32 @@ public class ApiProcessor {
     private final RestTemplate restTemplate;
 
     private static final String API_KEY = System.getenv("AKTO_API_KEY");
+    private static final String AKTO_DASHBOARD_URL = System.getenv("AKTO_URL");
     private static final String AKTO_PROD_URL = "https://app.akto.io";
 
     public <T, R> T processRequest(String url, R request, Class<T> responseType, HttpMethod method) {
         log.info("Making request to: {}", url);
         log.debug("Request body: {}", request);
 
+//        ServletRequestAttributes attrs = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
+//        HttpServletRequest httpRequest = attrs.getRequest();
+//        String token = (String) httpRequest.getAttribute("AKTO_API_KEY");
+
+        String token = AuthTokenContext.getToken();
+
         HttpHeaders headers = createDefaultHeaders();
         headers.set("x-api-key", API_KEY);
 
         HttpEntity<R> entity = new HttpEntity<>(request, headers);
 
+
+        String host = AKTO_PROD_URL;
+        if (AKTO_DASHBOARD_URL != null && !AKTO_DASHBOARD_URL.isEmpty()) {
+            host = AKTO_DASHBOARD_URL;
+        }
+
         ResponseEntity<T> response = restTemplate.exchange(
-            AKTO_PROD_URL + "/" + url,
+            host + "/" + url,
             method,
             entity,
             responseType

src/main/java/io/akto/mcp/tools/AktoTools.java

@@ -7,14 +7,18 @@
 import io.akto.mcp.model.FetchTestingRunResultsRequest;
 import io.akto.mcp.models.RetrieveAllCollectionTestsRequest;
 import io.akto.mcp.processor.ApiProcessor;
+import jakarta.servlet.http.HttpServletRequest;
 import java.util.Collections;
 import java.util.Map;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.ai.mcp.McpToolUtils;
 import org.springframework.ai.tool.annotation.Tool;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.HttpMethod;
 import org.springframework.stereotype.Service;
+import org.springframework.web.context.request.RequestContextHolder;
+import org.springframework.web.context.request.ServletRequestAttributes;
 
 @Slf4j
 @Service
@@ -413,6 +417,10 @@ public String fetchAllIssues(FetchAllIssuesRequest request) {
         return getResponseFromAkto("api/fetchAllIssues", request);
     }
 
+    public String fetchAllIssues(Object request) {
+        return getResponseFromAkto("api/fetchAllIssues", request);
+    }
+
     @Tool(name = "fetchNewEndpointsTrend", description = """
         Retrieves trend data for new endpoints over a specified time period.