Move schemas to their own package

agent-proxy/pom.xml

@@ -40,10 +40,10 @@
         <java.version>21</java.version>
         <confluent.version>7.9.1-22</confluent.version>
         <kafka.version>7.9.0-ce</kafka.version>
-        <caffeine.version>3.2.0</caffeine.version>
         <reactor-test.version>3.7.5</reactor-test.version>
         <spring-openapi.version>2.8.6</spring-openapi.version>
         <kcache.version>5.2.1</kcache.version>
+        <caffeine.version>3.2.0</caffeine.version>
         <mcp.version>0.9.0</mcp.version>
         <apache.common.lang3.version>3.17.0</apache.common.lang3.version>
     </properties>
@@ -109,6 +109,12 @@
             <version>${kcache.version}</version>
         </dependency>
 
+        <dependency>
+            <groupId>io.kcache</groupId>
+            <artifactId>kcache-caffeine</artifactId>
+            <version>${kcache.version}</version>
+        </dependency>
+
         <dependency>
             <groupId>io.modelcontextprotocol.sdk</groupId>
             <artifactId>mcp</artifactId>

agent-proxy/src/main/java/io/confluent/pas/agent/proxy/ConfluentMcpProxyApplication.java → agent-proxy/src/main/java/io/confluent/pas/agent/proxy/AgentProxyApplication.java

@@ -8,10 +8,10 @@
 @EnableAsync
 @SpringBootApplication
 @EnableConfigurationProperties
-public class ConfluentMcpProxyApplication {
+public class AgentProxyApplication {
 
     public static void main(String[] args) {
-        SpringApplication.run(ConfluentMcpProxyApplication.class, args);
+        SpringApplication.run(AgentProxyApplication.class, args);
     }
 
 }

agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinator.java

@@ -1,10 +1,7 @@
 package io.confluent.pas.agent.proxy.registration;
 
 import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient;
-import io.confluent.pas.agent.common.services.KafkaConfiguration;
-import io.confluent.pas.agent.common.services.KafkaPropertiesFactory;
-import io.confluent.pas.agent.common.services.RegistrationService;
-import io.confluent.pas.agent.common.services.RegistrationServiceHandler;
+import io.confluent.pas.agent.common.services.*;
 import io.confluent.pas.agent.common.services.schemas.Registration;
 import io.confluent.pas.agent.common.services.schemas.RegistrationKey;
 import io.confluent.pas.agent.proxy.registration.events.DeletedRegistrationEvent;
@@ -19,39 +16,46 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.stereotype.Component;
+import reactor.core.publisher.Mono;
 
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * The RegistrationCoordinator is the central component responsible for managing
- * tool registrations.
+ * tool registrations in a thread-safe and concurrent manner.
  * <p>
  * It coordinates the following processes:
  * - Listening for new registrations on the registration topic
- * - Processing incoming registrations and unregistrations
+ * - Processing incoming registrations and un-registrations concurrently
  * - Creating and managing handlers for each registered tool
- * - Maintaining the lifecycle of registrations
+ * - Maintaining the lifecycle of registrations with proper synchronization
  * - Broadcasting registration events to other components
+ * - Handling registration conflicts and updates atomically
  * <p>
  * This class acts as the bridge between the Kafka-based registration system and
- * the protocol-specific servers (MCP, REST) that handle client communications.
+ * the protocol-specific servers (MCP, REST, A2A) that handle client communications.
+ * It ensures thread-safe operations when managing multiple concurrent registrations
+ * and unregistrations.
  */
 @Slf4j
 @Component
 public class RegistrationCoordinator implements DisposableBean {
 
     /**
-     * MCP server for handling Model Context Protocol communications
+     * MCP server for handling Model Context Protocol communications asynchronously
      */
     @Getter
     private final McpAsyncServer mcpServer;
 
+    /**
+     * Agent-to-Agent async server for handling inter-agent communications
+     */
     @Getter
     private final A2AAsyncServer a2AAsyncServer;
 
     /**
-     * REST server for handling HTTP/REST communications
+     * REST server for handling synchronous HTTP communications with agents
      */
     @Getter
     private final AgentAsyncServer restServer;
@@ -63,7 +67,9 @@ public class RegistrationCoordinator implements DisposableBean {
     private final RequestResponseHandler requestResponseHandler;
 
     /**
-     * Thread-safe map of registration handlers indexed by registration name
+     * Thread-safe map of registration handlers indexed by registration name.
+     * Uses ConcurrentHashMap to ensure atomic operations and visibility across threads
+     * when adding, removing, or accessing handlers.
      */
     private final Map<String, CompositeHandler> handlers = new ConcurrentHashMap<>();
 
@@ -105,7 +111,7 @@ public RegistrationCoordinator(KafkaConfiguration kafkaConfiguration,
         // Create a registration handler that will forward registration events to our
         // onRegistration method
         // This avoids the circular reference issue during construction
-        RegistrationServiceHandler.Handler<RegistrationKey, Registration> registrationHandler = registrations -> {
+        CacheHandler.Handler<RegistrationKey, Registration> registrationHandler = registrations -> {
             if (registrations != null) {
                 onRegistration(registrations);
             }
@@ -240,14 +246,16 @@ private void handleRegistration(RegistrationKey key, Registration registration)
         // Handle registration update or creation
         if (handlers.containsKey(registrationName)) {
             log.info("Registration already exists, updating: {}", registrationName);
-            // First unregister the existing handler to clean up resources
+            // First, unregister the existing handler to clean up resources
             unregisterHandler(registrationName);
         } else {
             log.info("Received new registration: {}", registrationName);
         }
 
         // Create a new handler for the registration
-        createHandler(registration, registrationName);
+        createHandler(registration, registrationName)
+                .doOnError(e -> log.error("Error creating handler for registration: {}", registrationName, e))
+                .block();
     }
 
     /**
@@ -259,7 +267,7 @@ private void handleRegistration(RegistrationKey key, Registration registration)
      * @param registration     The registration details
      * @param registrationName The name of the registration
      */
-    private void createHandler(Registration registration, String registrationName) {
+    private Mono<Void> createHandler(Registration registration, String registrationName) {
         try {
             // Create a new composite handler for the registration
             final CompositeHandler handler = new CompositeHandler(
@@ -269,13 +277,11 @@ private void createHandler(Registration registration, String registrationName) {
                     this);
 
             // Initialize the handler asynchronously and block until complete
-            handler.initialize()
+            return handler.initialize()
                     .doOnSuccess(v -> handleSuccessfulRegistration(registrationName, handler, registration))
-                    .doOnError(e -> handleFailedRegistration(registrationName, e))
-                    .block();
+                    .doOnError(e -> handleFailedRegistration(registrationName, e));
         } catch (Exception e) {
-            // Handle any exceptions that occur during handler creation
-            log.error("Error creating handler for registration: {}", registrationName, e);
+            return Mono.error(e);
         }
     }
 

agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseChannel.java

@@ -70,8 +70,7 @@ void process(RequestResponseChannel channel,
      *
      * @param correlationId          unique ID to correlate requests with responses
      * @param registration           service registration details
-     * @param requestResponseHandler handler for sending requests and receiving
-     *                               responses
+     * @param requestResponseHandler handler for sending requests and receiving responses
      * @param processor              callback for processing received responses
      */
     public RequestResponseChannel(String correlationId,
@@ -96,8 +95,7 @@ public void close() {
     }
 
     /**
-     * Registers success and error handlers for responses with the request-response
-     * handler.
+     * Registers success and error handlers for responses with the request-response handler.
      * This sets up the callback mechanism for processing responses.
      */
     private void registerResponseHandler() {
@@ -136,8 +134,7 @@ private ConsumerService.ResponseHandler createSuccessResponseHandler(ResponsePro
     }
 
     /**
-     * Creates a handler for error responses or exceptions that occur during
-     * processing.
+     * Creates a handler for error responses or exceptions that occur during processing.
      * This handler formats errors and forwards them to the processor.
      *
      * @param processor callback to forward the error information to

agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandler.java

@@ -6,7 +6,6 @@
 import io.confluent.pas.agent.proxy.frameworks.java.models.Key;
 import io.confluent.pas.agent.proxy.registration.kafka.ProducerService;
 import io.confluent.pas.agent.proxy.registration.kafka.ConsumerService;
-import io.micrometer.observation.ObservationRegistry;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -17,37 +16,61 @@
 import java.util.Collection;
 
 /**
- * Handle requests and responses
+ * Handles the request-response communication pattern using Kafka topics.
+ * This component manages the lifecycle of Kafka producers and consumers,
+ * handles message routing, and maintains registration of response handlers.
+ * It provides functionality to send requests and handle corresponding responses
+ * asynchronously using correlation IDs.
  */
 @Slf4j
 @Component
 public class RequestResponseHandler implements DisposableBean {
 
     private final ProducerService producerService;
     private final ConsumerService consumerService;
-    private final ObservationRegistry observationRegistry;
 
+    /**
+     * Creates a new RequestResponseHandler with auto-configured services.
+     *
+     * @param kafkaConfiguration The Kafka configuration to use for producers and consumers
+     * @param responseTimeout    The maximum time to wait for responses in milliseconds (default: 10000)
+     */
     @Autowired
     public RequestResponseHandler(KafkaConfiguration kafkaConfiguration,
-                                  ObservationRegistry observationRegistry,
                                   @Value("${kafka.response.timeout:10000}") long responseTimeout) {
         this(new ProducerService(kafkaConfiguration),
-                new ConsumerService(kafkaConfiguration, responseTimeout),
-                observationRegistry);
+                new ConsumerService(kafkaConfiguration, responseTimeout));
     }
 
+    /**
+     * Creates a new RequestResponseHandler with pre-configured services.
+     *
+     * @param producerService The service responsible for sending messages to Kafka
+     * @param consumerService The service responsible for consuming messages from Kafka
+     */
     public RequestResponseHandler(ProducerService producerService,
-                                  ConsumerService consumerService,
-                                  ObservationRegistry observationRegistry) {
+                                  ConsumerService consumerService) {
         this.producerService = producerService;
         this.consumerService = consumerService;
-        this.observationRegistry = observationRegistry;
     }
 
+    /**
+     * Adds multiple registrations to the consumer service for message handling.
+     *
+     * @param registrations Collection of registrations to be added
+     */
     public void addRegistrations(Collection<Registration> registrations) {
         consumerService.addRegistrations(registrations);
     }
 
+    /**
+     * Registers a response handler for a specific registration and correlation ID.
+     *
+     * @param registration  The registration details for the handler
+     * @param correlationId The unique identifier to correlate requests with responses
+     * @param handler       The handler for successful responses
+     * @param errorHandler  The handler for error responses
+     */
     public void registerHandler(Registration registration,
                                 String correlationId,
                                 ConsumerService.ResponseHandler handler,
@@ -59,16 +82,35 @@ public void registerHandler(Registration registration,
                 errorHandler);
     }
 
+    /**
+     * Removes a previously registered response handler.
+     *
+     * @param registration  The registration details for the handler
+     * @param correlationId The correlation ID of the handler to remove
+     */
     public void unregisterHandler(Registration registration, String correlationId) {
         consumerService.unregisterResponseHandler(registration, correlationId);
     }
 
+    /**
+     * Sends a request message to the specified Kafka topic.
+     *
+     * @param registration The registration containing the topic information
+     * @param key          The message key
+     * @param request      The request payload as JsonNode
+     * @return A Mono that completes when the message is sent
+     */
     public Mono<Void> sendRequest(Registration registration,
                                   Key key,
                                   JsonNode request) {
         return producerService.send(registration.getRequestTopicName(), key, request);
     }
 
+    /**
+     * Implements DisposableBean to properly close Kafka resources when the bean is destroyed.
+     *
+     * @throws Exception if an error occurs during resource cleanup
+     */
     @Override
     public void destroy() throws Exception {
         consumerService.close();

agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/McpResourceHandler.java

@@ -17,6 +17,7 @@
 import javax.naming.OperationNotSupportedException;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 
 /**
  * Handler for Model Context Protocol (MCP) resource registration and request
@@ -45,8 +46,7 @@ public McpResourceHandler(ResourceRegistration registration,
                               RequestResponseHandler requestResponseHandler,
                               McpAsyncServer mcpServer) {
         super(registration, schemas, mcpServer, requestResponseHandler, (payload) -> {
-            final ResourceResponse.ResponseType responseType = ResourceResponse.ResponseType.fromValue(
-                    payload.get("type").asText());
+            final String responseType = payload.get("type").asText();
 
             final McpSchema.ResourceContents content = createResourceContents(
                     JsonUtils.toMap(payload),
@@ -127,9 +127,9 @@ private McpServerFeatures.AsyncResourceSpecification createResourceRegistration(
      */
     private static McpSchema.ResourceContents createResourceContents(
             Map<String, Object> response,
-            ResourceResponse.ResponseType responseType) {
+            String responseType) {
 
-        if (responseType == ResourceResponse.ResponseType.BLOB) {
+        if (Objects.equals(responseType, ResourceResponse.BLOB_TYPE)) {
             return createBlobContents(response);
         } else {
             return createTextContents(response);