diff --git a/agent-proxy/pom.xml b/agent-proxy/pom.xml index f4d090b..cb1d8d2 100644 --- a/agent-proxy/pom.xml +++ b/agent-proxy/pom.xml @@ -40,10 +40,10 @@ 21 7.9.1-22 7.9.0-ce - 3.2.0 3.7.5 2.8.6 5.2.1 + 3.2.0 0.9.0 3.17.0 @@ -109,6 +109,12 @@ ${kcache.version} + + io.kcache + kcache-caffeine + ${kcache.version} + + io.modelcontextprotocol.sdk mcp diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/ConfluentMcpProxyApplication.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/AgentProxyApplication.java similarity index 78% rename from agent-proxy/src/main/java/io/confluent/pas/agent/proxy/ConfluentMcpProxyApplication.java rename to agent-proxy/src/main/java/io/confluent/pas/agent/proxy/AgentProxyApplication.java index 5531522..7ecde16 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/ConfluentMcpProxyApplication.java +++ b/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); } } \ No newline at end of file diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinator.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinator.java index 3f3e67c..f6c1283 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinator.java +++ b/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. *

* 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 *

* 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 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 registrationHandler = registrations -> { + CacheHandler.Handler 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 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); } } diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseChannel.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseChannel.java index ca1bcf6..198f456 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseChannel.java +++ b/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 diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandler.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandler.java index ae5416d..f0f89b2 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandler.java +++ b/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,7 +16,11 @@ 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 @@ -25,29 +28,49 @@ 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 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 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(); diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/McpResourceHandler.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/McpResourceHandler.java index 08cbb89..f809487 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/McpResourceHandler.java +++ b/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 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); diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/Consumer.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/Consumer.java index 6c5376a..8699f3f 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/Consumer.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/Consumer.java @@ -1,225 +1,64 @@ package io.confluent.pas.agent.proxy.registration.kafka; -import io.confluent.pas.agent.common.services.KafkaConfiguration; -import io.confluent.pas.agent.common.services.KafkaPropertiesFactory; -import io.confluent.pas.agent.proxy.frameworks.java.models.Key; -import lombok.extern.slf4j.Slf4j; -import org.apache.kafka.clients.consumer.KafkaConsumer; import org.apache.kafka.clients.consumer.ConsumerRecord; -import org.apache.kafka.common.errors.WakeupException; import java.io.Closeable; -import java.io.IOException; -import java.time.Duration; -import java.util.*; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; +import java.util.Collection; + /** - * Consumer class for managing Kafka topic subscriptions and message processing. + * Interface for managing Kafka topic subscriptions and message consumption. * * @param Key type for Kafka messages * @param Value type for Kafka messages */ -@Slf4j -public class Consumer implements Closeable { - - @FunctionalInterface - public interface TimeoutChecker { - void checkTimeouts(long lastCheck); - } - - private final ExecutorService executorSvc = Executors.newSingleThreadExecutor(); - private final KafkaConsumer kafkaConsumer; - - private final ConsumerHandler consumerHandler; - private final List topics = Collections.synchronizedList(new ArrayList<>()); - private final TimeoutChecker timeoutChecker; - private volatile boolean subscriptionUpdated = false; - private volatile boolean stopRequested = false; - +public interface Consumer extends Closeable { /** - * Constructs a Consumer instance with the specified Kafka configuration and - * message types. - * - * @param kafkaConfiguration Kafka configuration containing connection and auth details - * @param requestClass Class type for message values - * @param timeoutChecker The timeout checker to use + * Functional interface for checking timeouts in message processing. */ - public Consumer(KafkaConfiguration kafkaConfiguration, - Class requestClass, - ConsumerHandler consumerHandler, - TimeoutChecker timeoutChecker) { - this(consumerHandler, - new KafkaConsumer<>(KafkaPropertiesFactory.getConsumerProperties( - kafkaConfiguration, - false, - Key.class, - requestClass)), - timeoutChecker); + @FunctionalInterface + interface TimeoutChecker { + /** + * Check for any timeouts since the last check. + * + * @param lastCheck timestamp of the last timeout check + */ + void checkTimeouts(long lastCheck); } /** - * Constructs a Consumer instance with the specified Kafka consumer and handler. + * Checks if there is an active subscription for the specified topic. * - * @param consumerHandler The handler to process messages - * @param kafkaConsumer The Kafka consumer to use - * @param timeoutChecker The timeout checker to use + * @param topic the topic to check + * @return true if subscribed to the topic, false otherwise */ - public Consumer(ConsumerHandler consumerHandler, - KafkaConsumer kafkaConsumer, - TimeoutChecker timeoutChecker) { - this.kafkaConsumer = kafkaConsumer; - this.consumerHandler = consumerHandler; - this.timeoutChecker = timeoutChecker; - - executorSvc.submit(this::runLoop); - } - - public boolean isSubscribed(String topic) { - return topics.contains(topic); - } + boolean isSubscribed(String topic); /** - * Subscribes to a Kafka topic with the specified handler. + * Subscribes to a single Kafka topic. * - * @param topic The topic to subscribe to + * @param topic the topic to subscribe to */ - public void subscribe(String topic) { - if (topics.contains(topic)) { - throw new IllegalArgumentException("Subscription already exists for topic: " + topic); - } - - topics.add(topic); - subscriptionUpdated = true; - } + void subscribe(String topic); /** - * Subscribes to a Kafka topic with the specified handler. + * Subscribes to multiple Kafka topics. * - * @param topicsToAdd The topics to subscribe to + * @param topicsToAdd collection of topics to subscribe to */ - public void subscribe(Collection topicsToAdd) { - topics.addAll(topicsToAdd); - subscriptionUpdated = true; - } + void subscribe(Collection topicsToAdd); /** * Unsubscribes from a Kafka topic. * - * @param topic The topic to unsubscribe from + * @param topic the topic to unsubscribe from */ - public void unsubscribe(String topic) { - topics.remove(topic); - subscriptionUpdated = true; - } + void unsubscribe(String topic); /** - * Closes the consumer, stopping the polling loop and releasing resources. + * Processes a single Kafka record. * - * @throws IOException if an error occurs while closing the consumer + * @param record the Kafka record to process */ - @Override - public void close() throws IOException { - stopRequested = true; - - kafkaConsumer.wakeup(); - - try { - executorSvc.shutdown(); - - boolean done = executorSvc.awaitTermination(10, TimeUnit.SECONDS); - if (!done) { - log.error("Executor did not terminate"); - } - } catch (InterruptedException e) { - log.error("Error waiting for executor to terminate", e); - } - } - - /** - * Main loop for polling Kafka messages and processing them. - */ - void runLoop() { - log.info("Starting Consumer Service"); - - long lastCheck = System.currentTimeMillis(); - - while (!stopRequested) { - updateSubscriptions(); - if (topics.isEmpty()) { - try { - Thread.sleep(100); - } catch (InterruptedException e) { - log.warn("Interrupted while sleeping", e); - Thread.currentThread().interrupt(); - } - - continue; - } - - try { - var records = kafkaConsumer.poll(Duration.ofMillis(100)); - - // Process the records before checking for timeouts - records.forEach(this::processRecord); - - try { - // Check for timeouts in the registered handlers - timeoutChecker.checkTimeouts(lastCheck); - - // Update the last check time - lastCheck = System.currentTimeMillis(); - } catch (Exception e) { - log.error("Error checking timeouts", e); - } - - } catch (WakeupException e) { - if (!stopRequested) { - log.error("Unexpected WakeupException", e); - throw e; - } - } - } - - kafkaConsumer.close(); - - log.info("Consumer Service stopped"); - } - - /** - * Updates the Kafka topic subscriptions based on the current handlers. - */ - private void updateSubscriptions() { - if (subscriptionUpdated) { - log.info("Updating subscriptions"); - kafkaConsumer.unsubscribe(); - if (!topics.isEmpty()) { - kafkaConsumer.subscribe(topics); - } else { - log.info("No topics to subscribe to"); - } - subscriptionUpdated = false; - } - } - - /** - * Processes a single Kafka record by invoking the appropriate handler. - * - * @param record The Kafka record to process - */ - void processRecord(ConsumerRecord record) { - if (!topics.contains(record.topic())) { - log.warn("Received message from unregistered topic: {}", record.topic()); - return; - } - - try { - log.info("Processing message from topic: {}", record.topic()); - consumerHandler.onMessage(record.topic(), record.key(), record.value()); - } catch (Exception e) { - log.error("Failed to process message", e); - } - } -} \ No newline at end of file + void processRecord(ConsumerRecord record); +} diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerService.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerService.java index db1725f..752dc15 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerService.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerService.java @@ -4,6 +4,7 @@ import io.confluent.pas.agent.common.services.KafkaConfiguration; import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; +import io.confluent.pas.agent.proxy.registration.kafka.impl.ConsumerImpl; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -17,8 +18,7 @@ /** * Service that handles responses from Kafka topics by routing messages to - * appropriate handlers - * based on correlation IDs. This service: + * appropriate handlers based on correlation IDs. This service: *