From 211b2341ab733210d7c0a9d02128f1307b173849 Mon Sep 17 00:00:00 2001 From: Pascal Vantrepote Date: Sun, 20 Apr 2025 21:52:09 -0400 Subject: [PATCH 1/3] Move schemas to their own package --- .../registration/RegistrationCoordinator.java | 30 +- .../registration/RegistrationHandler.java | 4 +- .../registration/RequestResponseHandler.java | 13 +- .../events/DeletedRegistrationEvent.java | 6 +- .../events/NewRegistrationEvent.java | 6 +- .../handlers/{ => mcp}/ResourceHandler.java | 28 +- .../handlers/{ => mcp}/ToolHandler.java | 10 +- .../registration/kafka/ConsumerService.java | 10 +- .../schemas/RegistrationSchemas.java | 4 +- .../proxy/rest/ControlAPIController.java | 8 +- .../proxy/rest/OpenAPIConfiguration.java | 17 +- .../proxy/rest/ToolRestConfiguration.java | 7 +- .../agent/proxy/rest/ToolRestController.java | 24 +- .../pas/agent/proxy/TestProxyIT.java | 4 +- .../RegistrationCoordinatorTest.java | 33 +- .../RequestResponseHandlerTest.java | 13 +- .../handlers/ResourceHandlerTest.java | 11 +- .../handlers/ToolHandlerTest.java | 5 +- .../kafka/ConsumerServiceTest.java | 8 +- .../common/services/RegistrationService.java | 6 +- .../services/RegistrationServiceHandler.java | 8 +- .../pas/agent/common/services/Schemas.java | 537 ------------------ .../schemas/BlobResourceResponse.java | 73 +++ .../common/services/schemas/Registration.java | 116 ++++ .../services/schemas/RegistrationKey.java | 56 ++ .../schemas/ResourceRegistration.java | 126 ++++ .../services/schemas/ResourceRequest.java | 33 ++ .../services/schemas/ResourceResponse.java | 109 ++++ .../schemas/TextResourceResponse.java | 73 +++ .../RegistrationServiceHandlerTest.java | 18 +- .../services/RegistrationServiceTest.java | 22 +- .../pas/mcp/exemple/ResourceAgent.java | 9 +- .../frameworks/java/SubscriptionHandler.java | 31 +- .../spring/annotation/AgentRegistrar.java | 14 +- .../java/spring/annotation/Resource.java | 4 +- .../McpRegistrationAutoConfiguration.java | 10 +- .../mcp/AsyncMcpToolCallbackProvider.java | 6 +- .../java/spring/mcp/McpToolFilters.java | 4 +- .../mcp/SyncMcpToolCallbackProvider.java | 6 +- .../java/SubscriptionHandlerTest.java | 17 +- .../spring/annotation/AgentRegistrarTest.java | 9 +- .../java/spring/mcp/McpToolFiltersTest.java | 4 +- .../mcp/proxy/frameworks/client/Agent.java | 4 +- 43 files changed, 806 insertions(+), 730 deletions(-) rename agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/{ => mcp}/ResourceHandler.java (82%) rename agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/{ => mcp}/ToolHandler.java (92%) delete mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/Schemas.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/RegistrationKey.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRequest.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java 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 39a2c49..4983773 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 @@ -4,11 +4,13 @@ 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.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.RegistrationKey; +import io.confluent.pas.agent.common.services.schemas.ResourceRegistration; import io.confluent.pas.agent.proxy.registration.events.DeletedRegistrationEvent; import io.confluent.pas.agent.proxy.registration.events.NewRegistrationEvent; -import io.confluent.pas.agent.proxy.registration.handlers.ResourceHandler; -import io.confluent.pas.agent.proxy.registration.handlers.ToolHandler; +import io.confluent.pas.agent.proxy.registration.handlers.mcp.ResourceHandler; +import io.confluent.pas.agent.proxy.registration.handlers.mcp.ToolHandler; import io.modelcontextprotocol.server.McpAsyncServer; import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.DisposableBean; @@ -32,7 +34,7 @@ public class RegistrationCoordinator implements DisposableBean { private final McpAsyncServer mcpServer; private final Map> handlers = new ConcurrentHashMap<>(); private final SchemaRegistryClient schemaRegistryClient; - private final RegistrationService registrationService; + private final RegistrationService registrationService; private final ApplicationEventPublisher applicationEventPublisher; @Autowired @@ -58,15 +60,15 @@ public RegistrationCoordinator(KafkaConfiguration kafkaConfiguration, this.applicationEventPublisher = applicationEventPublisher; this.registrationService = new RegistrationService<>( kafkaConfiguration, - Schemas.RegistrationKey.class, - Schemas.Registration.class, + RegistrationKey.class, + Registration.class, this::onRegistration); } public RegistrationCoordinator(RequestResponseHandler requestResponseHandler, McpAsyncServer mcpServer, SchemaRegistryClient schemaRegistryClient, - RegistrationService registrationService, + RegistrationService registrationService, ApplicationEventPublisher applicationEventPublisher) { this.requestResponseHandler = requestResponseHandler; this.mcpServer = mcpServer; @@ -109,7 +111,7 @@ public boolean isRegistered(String name) { * * @return The registrations */ - public List getAllRegistrations() { + public List getAllRegistrations() { return registrationService.getAllRegistrations(); } @@ -118,8 +120,8 @@ public List getAllRegistrations() { * * @param registration The registration */ - public void register(Schemas.Registration registration) { - registrationService.register(new Schemas.RegistrationKey(registration.getName()), registration); + public void register(Registration registration) { + registrationService.register(new RegistrationKey(registration.getName()), registration); } /** @@ -128,7 +130,7 @@ public void register(Schemas.Registration registration) { * @param name The registration name to delete */ public void unregister(String name) { - registrationService.unregister(new Schemas.RegistrationKey(name)); + registrationService.unregister(new RegistrationKey(name)); } /** @@ -136,13 +138,13 @@ public void unregister(String name) { * * @param registrations The registrations */ - void onRegistration(Map registrations) { + void onRegistration(Map registrations) { requestResponseHandler.addRegistrations(registrations.values()); registrations.forEach(this::onRegistration); } - private void onRegistration(Schemas.RegistrationKey key, Schemas.Registration registration) { + private void onRegistration(RegistrationKey key, Registration registration) { final String registrationName = key.getName(); // Unregister? @@ -164,7 +166,7 @@ private void onRegistration(Schemas.RegistrationKey key, Schemas.Registration re } try { - final RegistrationHandler handler = (registration instanceof Schemas.ResourceRegistration rcsRegistration) + final RegistrationHandler handler = (registration instanceof ResourceRegistration rcsRegistration) ? new ResourceHandler(rcsRegistration, schemaRegistryClient, requestResponseHandler) : new ToolHandler(registration, schemaRegistryClient, requestResponseHandler); diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationHandler.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationHandler.java index 538a50c..62bbe76 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationHandler.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationHandler.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.registration; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.proxy.registration.schemas.RegistrationSchemas; import io.modelcontextprotocol.server.McpAsyncServer; import reactor.core.publisher.Mono; @@ -17,7 +17,7 @@ public interface RegistrationHandler { * * @return the registration */ - Schemas.Registration getRegistration(); + Registration getRegistration(); /** * Gets the registration schemas for the tool or resource. 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 b883950..d216cd8 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 @@ -2,8 +2,8 @@ import com.fasterxml.jackson.databind.JsonNode; import io.confluent.pas.agent.common.services.KafkaConfiguration; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.proxy.registration.kafka.ProducerService; -import io.confluent.pas.agent.common.services.Schemas; import io.confluent.pas.agent.proxy.registration.kafka.ConsumerService; import io.confluent.pas.agent.proxy.registration.schemas.RegistrationSchemas; import io.micrometer.observation.Observation; @@ -18,6 +18,7 @@ import java.util.Collection; import java.util.Map; +import java.util.Objects; import java.util.concurrent.ExecutionException; /** @@ -48,7 +49,7 @@ public RequestResponseHandler(ProducerService producerService, this.observationRegistry = observationRegistry; } - public void addRegistrations(Collection registrations) { + public void addRegistrations(Collection registrations) { consumerService.addRegistrations(registrations); } @@ -63,7 +64,7 @@ public void addRegistrations(Collection registrations) { * @throws ExecutionException if the request fails * @throws InterruptedException if the request is interrupted */ - public Mono sendRequestResponse(Schemas.Registration registration, + public Mono sendRequestResponse(Registration registration, RegistrationSchemas schemas, String correlationId, Map request) @@ -75,11 +76,11 @@ public Mono sendRequestResponse(Schemas.Registration registration, .lowCardinalityKeyValue("correlationId", correlationId) .highCardinalityKeyValue("name", registration.getName()); - return observation.observe(() -> sendRequestResponse( + return Objects.requireNonNull(observation.observe(() -> sendRequestResponse( registration, correlationId, schemas.getRequestKeySchema().envelope(key), - schemas.getRequestSchema().envelope(request))) + schemas.getRequestSchema().envelope(request)))) .doOnError(observation::error) .doFinally(signalType -> observation.stop()); } @@ -93,7 +94,7 @@ public Mono sendRequestResponse(Schemas.Registration registration, * @param request the request * @return the response */ - public Mono sendRequestResponse(Schemas.Registration registration, + public Mono sendRequestResponse(Registration registration, String correlationId, JsonNode key, JsonNode request) { diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/DeletedRegistrationEvent.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/DeletedRegistrationEvent.java index 90a42a8..614b33c 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/DeletedRegistrationEvent.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/DeletedRegistrationEvent.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.registration.events; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import lombok.Getter; import org.springframework.context.ApplicationEvent; @@ -10,9 +10,9 @@ @Getter public class DeletedRegistrationEvent extends ApplicationEvent { - private final Schemas.Registration registration; + private final Registration registration; - public DeletedRegistrationEvent(Object source, Schemas.Registration registration) { + public DeletedRegistrationEvent(Object source, Registration registration) { super(source); this.registration = registration; } diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/NewRegistrationEvent.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/NewRegistrationEvent.java index 047a089..3f4877c 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/NewRegistrationEvent.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/events/NewRegistrationEvent.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.registration.events; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import lombok.Getter; import org.springframework.context.ApplicationEvent; @@ -10,9 +10,9 @@ @Getter public class NewRegistrationEvent extends ApplicationEvent { - private final Schemas.Registration registration; + private final Registration registration; - public NewRegistrationEvent(Object source, Schemas.Registration registration) { + public NewRegistrationEvent(Object source, Registration registration) { super(source); this.registration = registration; } diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/ResourceHandler.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/ResourceHandler.java similarity index 82% rename from agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/ResourceHandler.java rename to agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/ResourceHandler.java index 579423b..e3cf335 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/ResourceHandler.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/ResourceHandler.java @@ -1,13 +1,17 @@ -package io.confluent.pas.agent.proxy.registration.handlers; +package io.confluent.pas.agent.proxy.registration.handlers.mcp; import com.fasterxml.jackson.databind.JsonNode; import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.ResourceRegistration; +import io.confluent.pas.agent.common.services.schemas.ResourceRequest; +import io.confluent.pas.agent.common.services.schemas.ResourceResponse; import io.confluent.pas.agent.common.utils.JsonUtils; import io.confluent.pas.agent.proxy.registration.RegistrationHandler; import io.confluent.pas.agent.proxy.registration.RequestResponseHandler; import io.confluent.pas.agent.proxy.registration.schemas.RegistrationSchemas; +import io.confluent.pas.mcp.common.schemas.BlobResourceResponse; +import io.confluent.pas.mcp.common.schemas.TextResourceResponse; import io.modelcontextprotocol.server.McpAsyncServer; import io.modelcontextprotocol.server.McpServerFeatures; import io.modelcontextprotocol.spec.McpSchema; @@ -26,15 +30,15 @@ @Slf4j @AllArgsConstructor -public class ResourceHandler implements RegistrationHandler { +public class ResourceHandler implements RegistrationHandler { @Getter - private final Schemas.ResourceRegistration registration; + private final ResourceRegistration registration; @Getter private final RegistrationSchemas schemas; private final RequestResponseHandler requestResponseHandler; - public ResourceHandler(Schemas.ResourceRegistration registration, + public ResourceHandler(ResourceRegistration registration, SchemaRegistryClient schemaRegistryClient, RequestResponseHandler requestResponseHandler) throws RestClientException, IOException { this.requestResponseHandler = requestResponseHandler; @@ -115,11 +119,11 @@ public Mono unregister(McpAsyncServer mcpServer) { } @Override - public Mono sendRequest(Schemas.ResourceRequest request) { + public Mono sendRequest(ResourceRequest request) { final Map arguments = JsonUtils.toMap(request); return sendRequest(arguments) - .map(response -> JsonUtils.toObject(response, Schemas.ResourceResponse.class)); + .map(response -> JsonUtils.toObject(response, ResourceResponse.class)); } /** @@ -149,21 +153,21 @@ protected Mono sendRequest(Map arguments) { * @param request the read resource request * @param sink the sink to send the response to */ - protected void sendRequest(McpSchema.ReadResourceRequest request, MonoSink sink) { + public void sendRequest(McpSchema.ReadResourceRequest request, MonoSink sink) { final Map arguments = JsonUtils.toMap(request); sendRequest(arguments).subscribe(response -> { - final Schemas.ResourceResponse.ResponseType responseType = Schemas.ResourceResponse.ResponseType.fromValue(response.get("type").asText()); + final ResourceResponse.ResponseType responseType = ResourceResponse.ResponseType.fromValue(response.get("type").asText()); final McpSchema.ResourceContents content; - if (responseType == Schemas.ResourceResponse.ResponseType.BLOB) { - Schemas.BlobResourceResponse resource = JsonUtils.toObject(response, Schemas.BlobResourceResponse.class); + if (responseType == ResourceResponse.ResponseType.BLOB) { + BlobResourceResponse resource = JsonUtils.toObject(response, BlobResourceResponse.class); content = new McpSchema.BlobResourceContents( resource.getUri(), resource.getMimeType(), resource.getBlob()); } else { - Schemas.TextResourceResponse resource = JsonUtils.toObject(response, Schemas.TextResourceResponse.class); + TextResourceResponse resource = JsonUtils.toObject(response, TextResourceResponse.class); content = new McpSchema.TextResourceContents( resource.getUri(), resource.getMimeType(), diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/ToolHandler.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/ToolHandler.java similarity index 92% rename from agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/ToolHandler.java rename to agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/ToolHandler.java index 4a23abd..92d592b 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/ToolHandler.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/handlers/mcp/ToolHandler.java @@ -1,10 +1,10 @@ -package io.confluent.pas.agent.proxy.registration.handlers; +package io.confluent.pas.agent.proxy.registration.handlers.mcp; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.JsonNode; import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.common.utils.JsonUtils; import io.confluent.pas.agent.proxy.registration.RegistrationHandler; import io.confluent.pas.agent.proxy.registration.RequestResponseHandler; @@ -31,12 +31,12 @@ @AllArgsConstructor public class ToolHandler implements RegistrationHandler, JsonNode> { @Getter - private final Schemas.Registration registration; + private final Registration registration; @Getter private final RegistrationSchemas schemas; private final RequestResponseHandler requestResponseHandler; - public ToolHandler(Schemas.Registration registration, + public ToolHandler(Registration registration, SchemaRegistryClient schemaRegistryClient, RequestResponseHandler requestResponseHandler) throws RestClientException, IOException { this.requestResponseHandler = requestResponseHandler; @@ -101,7 +101,7 @@ public Mono sendRequest(Map arguments) { * @param arguments the arguments to send * @param sink the sink to send the response to */ - protected void sendToolRequest(Map arguments, MonoSink sink) { + public void sendToolRequest(Map arguments, MonoSink sink) { sendRequest(arguments).subscribe(response -> { // Serialize the response try { 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 7c71992..7d58eeb 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 @@ -2,7 +2,7 @@ import com.fasterxml.jackson.databind.JsonNode; import io.confluent.pas.agent.common.services.KafkaConfiguration; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -76,7 +76,7 @@ public record RegistrationHandler( * @param registrationHandlers Map of correlation IDs to their handlers */ public record RegistrationItem( - Schemas.Registration registration, + Registration registration, Map registrationHandlers) { } @@ -129,14 +129,14 @@ public ConsumerService(Consumer consumer, long responseTimeo * * @param registrations The collection of registrations to subscribe to */ - public void addRegistrations(Collection registrations) { + public void addRegistrations(Collection registrations) { if (registrations == null || registrations.isEmpty()) { log.warn("No registrations provided to add"); return; } List topics = registrations.stream() - .map(Schemas.Registration::getResponseTopicName) + .map(Registration::getResponseTopicName) .collect(Collectors.toList()); log.info("Subscribing to response topics: {}", topics); @@ -154,7 +154,7 @@ public void addRegistrations(Collection registrations) { * @throws NullPointerException if any parameter is null */ public void registerResponseHandler( - Schemas.Registration registration, + Registration registration, String correlationId, ResponseHandler handler, ErrorHandler errorHandler) { diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/schemas/RegistrationSchemas.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/schemas/RegistrationSchemas.java index 0d41aff..b910a87 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/schemas/RegistrationSchemas.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/schemas/RegistrationSchemas.java @@ -2,7 +2,7 @@ import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.common.utils.Lazy; import lombok.extern.slf4j.Slf4j; @@ -33,7 +33,7 @@ public RegistrationSchema getResponseSchema() { } public RegistrationSchemas(SchemaRegistryClient client, - Schemas.Registration registration) throws IOException, RestClientException { + Registration registration) throws IOException, RestClientException { this.requestKeySchema = new Lazy<>(() -> { try { return getSchema(registration.getRequestTopicName(), true, client); diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ControlAPIController.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ControlAPIController.java index e239249..a007acf 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ControlAPIController.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ControlAPIController.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.rest; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.proxy.registration.RegistrationCoordinator; import io.confluent.pas.agent.proxy.registration.RegistrationHandler; import lombok.extern.slf4j.Slf4j; @@ -21,7 +21,7 @@ public ControlAPIController(RegistrationCoordinator coordinator) { } @GetMapping("/control/registrations") - public List getRegistrations() { + public List getRegistrations() { return coordinator .getAllRegistrationHandlers() .stream() @@ -30,7 +30,7 @@ public List getRegistrations() { } @PostMapping("/control/registration") - public void register(Schemas.Registration registration) { + public void register(Registration registration) { if (coordinator.isRegistered(registration.getName())) { throw new ResponseStatusException( HttpStatus.CONFLICT, @@ -42,7 +42,7 @@ public void register(Schemas.Registration registration) { } @PatchMapping("/control/registration") - public void update(Schemas.Registration registration) { + public void update(Registration registration) { if (coordinator.isRegistered(registration.getName())) { coordinator.register(registration); } else { diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/OpenAPIConfiguration.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/OpenAPIConfiguration.java index ed74d2c..16d1b3a 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/OpenAPIConfiguration.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/OpenAPIConfiguration.java @@ -4,7 +4,8 @@ import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.annotation.JsonProperty; import com.fasterxml.jackson.core.JsonProcessingException; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.ResourceRegistration; import io.confluent.pas.agent.common.utils.JsonUtils; import io.confluent.pas.agent.common.utils.UriTemplate; import io.confluent.pas.agent.proxy.registration.RegistrationCoordinator; @@ -107,9 +108,9 @@ private Map buildPathsFromRegistrations() { registrationHandlers.forEach(handler -> { try { - final Schemas.Registration registration = handler.getRegistration(); + final Registration registration = handler.getRegistration(); - if (registration instanceof Schemas.ResourceRegistration resourceRegistration) { + if (registration instanceof ResourceRegistration resourceRegistration) { addResourcePathItem(resourceRegistration, registration, pathItems); } else { addStandardPathItem(registration, handler.getSchemas(), pathItems); @@ -129,7 +130,7 @@ private Map buildPathsFromRegistrations() { * @param schemas the registration schemas * @param pathItems the path items map to update */ - private void addStandardPathItem(Schemas.Registration registration, + private void addStandardPathItem(Registration registration, RegistrationSchemas schemas, Map pathItems) { final String path = registration.getName(); @@ -160,9 +161,9 @@ private void addStandardPathItem(Schemas.Registration registration, * @param registration the general registration * @param pathItems the path items map to update */ - private void addResourcePathItem(Schemas.ResourceRegistration resourceRegistration, - Schemas.Registration registration, - Map pathItems) { + private void addResourcePathItem(ResourceRegistration resourceRegistration, + Registration registration, + Map pathItems) { final PathItem pathItem = new PathItem(); final String urlPath = resourceRegistration.getUrl(); @@ -244,7 +245,7 @@ private Content createRequestBody(String requestSchema) throws JsonProcessingExc * @param registration the resource registration * @return the API response */ - private ApiResponse createApiResponse(Schemas.ResourceRegistration registration) { + private ApiResponse createApiResponse(ResourceRegistration registration) { final Content content = new Content(); content.addMediaType(registration.getMimeType(), new MediaType()); diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestConfiguration.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestConfiguration.java index 8c7616b..50e34ac 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestConfiguration.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestConfiguration.java @@ -1,6 +1,7 @@ package io.confluent.pas.agent.proxy.rest; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.ResourceRegistration; import io.confluent.pas.agent.proxy.registration.RegistrationCoordinator; import io.confluent.pas.agent.proxy.registration.RegistrationHandler; import org.springframework.context.annotation.Bean; @@ -35,8 +36,8 @@ public RouterFunction createRoute(RegistrationCoordinator regist final List> registrationHandlers = registrationCoordinator.getAllRegistrationHandlers(); registrationHandlers.stream() .map(RegistrationHandler::getRegistration) - .filter(Schemas.Registration::isResource) - .map(r -> (Schemas.ResourceRegistration) r) + .filter(Registration::isResource) + .map(r -> (ResourceRegistration) r) .forEach(registration -> { final String url = registration.getUrl(); route.GET("/rcs/" + url, accept(APPLICATION_JSON), toolRestController::processResourceRequest); diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestController.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestController.java index 3f14d24..0f851c1 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestController.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/ToolRestController.java @@ -1,10 +1,14 @@ package io.confluent.pas.agent.proxy.rest; import com.fasterxml.jackson.databind.JsonNode; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.ResourceRegistration; +import io.confluent.pas.agent.common.services.schemas.ResourceRequest; +import io.confluent.pas.agent.common.services.schemas.ResourceResponse; import io.confluent.pas.agent.proxy.registration.RegistrationCoordinator; import io.confluent.pas.agent.proxy.registration.RegistrationHandler; -import io.confluent.pas.agent.proxy.registration.handlers.ResourceHandler; +import io.confluent.pas.agent.proxy.registration.handlers.mcp.ResourceHandler; +import io.confluent.pas.agent.common.services.schemas.BlobResourceResponse; +import io.confluent.pas.agent.common.services.schemas.TextResourceResponse; import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; import org.springframework.core.ParameterizedTypeReference; @@ -57,7 +61,7 @@ private void initializeResourceHandlers() { .filter(handler -> handler instanceof ResourceHandler) .map(handler -> (ResourceHandler) handler) .forEach(handler -> { - final Schemas.ResourceRegistration registration = handler.getRegistration(); + final ResourceRegistration registration = handler.getRegistration(); final String urlPattern = registration.getUrl(); resourceHandlersByUrlPattern.put(urlPattern, handler); log.debug("Registered resource handler for pattern: {}", urlPattern); @@ -125,7 +129,7 @@ public Mono processResourceRequest(ServerRequest request) { // Get the appropriate handler for the resource request return findResourceHandler(urlParts) .map(handler -> { - final Schemas.ResourceRequest resourceRequest = new Schemas.ResourceRequest( + final ResourceRequest resourceRequest = new ResourceRequest( StringUtils.join(urlParts, "/")); // Send the resource request to the handler @@ -164,7 +168,7 @@ private List extractUrlParts(ServerRequest request) { * @param response the resource response * @return the server response Mono */ - private Mono createResourceResponse(Schemas.ResourceResponse response) { + private Mono createResourceResponse(ResourceResponse response) { final MediaType mediaType = MediaType.parseMediaType(response.getMimeType()); final String responseContent = extractResponseContent(response); @@ -179,10 +183,10 @@ private Mono createResourceResponse(Schemas.ResourceResponse res * @param response the resource response * @return the response content as a string */ - private String extractResponseContent(Schemas.ResourceResponse response) { - if (response instanceof Schemas.BlobResourceResponse blobResponse) { + private String extractResponseContent(ResourceResponse response) { + if (response instanceof BlobResourceResponse blobResponse) { return blobResponse.getBlob(); - } else if (response instanceof Schemas.TextResourceResponse textResponse) { + } else if (response instanceof TextResourceResponse textResponse) { return textResponse.getText(); } else { throw new IllegalArgumentException("Unsupported resource response type: " + response.getClass().getName()); @@ -208,11 +212,11 @@ private Mono createErrorResponse(HttpStatus status, String messa * @param urlParts the URL parts * @return an Optional containing the handler if found, empty otherwise */ - private Optional> findResourceHandler( + private Optional> findResourceHandler( List urlParts) { return resourceHandlersByUrlPattern.keySet().stream() .filter(pattern -> isUrlPatternMatch(pattern, urlParts)) - .map(pattern -> (RegistrationHandler) resourceHandlersByUrlPattern + .map(pattern -> (RegistrationHandler) resourceHandlersByUrlPattern .get(pattern)) .findFirst(); } diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/TestProxyIT.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/TestProxyIT.java index 4fc6cd1..3fa80a5 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/TestProxyIT.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/TestProxyIT.java @@ -1,7 +1,7 @@ package io.confluent.pas.agent.proxy; import io.confluent.pas.agent.common.services.KafkaConfiguration; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.proxy.frameworks.java.SubscriptionHandler; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; import io.confluent.pas.agent.proxy.registration.RegistrationCoordinator; @@ -69,7 +69,7 @@ public void testToolRegistration() throws InterruptedException { String.class, String.class); - Schemas.Registration registration = new Schemas.Registration( + Registration registration = new Registration( "test", "Sample registration", "sample_req", diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinatorTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinatorTest.java index e6e67ee..4cd1772 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinatorTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinatorTest.java @@ -5,7 +5,8 @@ import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; import io.confluent.pas.agent.common.services.KafkaConfiguration; import io.confluent.pas.agent.common.services.RegistrationService; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.RegistrationKey; import io.modelcontextprotocol.server.McpAsyncServer; import io.modelcontextprotocol.server.McpServerFeatures; import org.apache.commons.lang3.RandomStringUtils; @@ -38,7 +39,7 @@ class RegistrationCoordinatorTest { private SchemaRegistryClient schemaRegistryClient; @Mock - private RegistrationService registrationService; + private RegistrationService registrationService; @InjectMocks private RegistrationCoordinator registrationCoordinator; @@ -82,9 +83,9 @@ void testIsRegistered() { String toolName = "testTool"; assertFalse(registrationCoordinator.isRegistered(toolName)); - registrationCoordinator.register(new Schemas.Registration(toolName, "description", "request", "response")); + registrationCoordinator.register(new Registration(toolName, "description", "request", "response")); - verify(registrationService, times(1)).register(any(Schemas.RegistrationKey.class), any(Schemas.Registration.class)); + verify(registrationService, times(1)).register(any(RegistrationKey.class), any(Registration.class)); } @Test @@ -92,9 +93,9 @@ void testGetRegistrationHandler() { String toolName = "testTool"; assertNull(registrationCoordinator.getRegistrationHandler(toolName)); - var registration = new Schemas.Registration(toolName, "description", "request", "response"); + var registration = new Registration(toolName, "description", "request", "response"); - registrationCoordinator.onRegistration(Map.of(new Schemas.RegistrationKey(toolName), registration)); + registrationCoordinator.onRegistration(Map.of(new RegistrationKey(toolName), registration)); assertNotNull(registrationCoordinator.getRegistrationHandler(toolName)); } @@ -102,34 +103,34 @@ void testGetRegistrationHandler() { void testGetAllRegistrationHandlers() { assertTrue(registrationCoordinator.getAllRegistrationHandlers().isEmpty()); - registrationCoordinator.register(new Schemas.Registration("testTool1", "description", "request", "response")); - registrationCoordinator.register(new Schemas.Registration("testTool2", "description", "request", "response")); + registrationCoordinator.register(new Registration("testTool1", "description", "request", "response")); + registrationCoordinator.register(new Registration("testTool2", "description", "request", "response")); - verify(registrationService, times(2)).register(any(Schemas.RegistrationKey.class), any(Schemas.Registration.class)); + verify(registrationService, times(2)).register(any(RegistrationKey.class), any(Registration.class)); } @Test void testRegister() { - Schemas.Registration registration = new Schemas.Registration("testTool", "description", "request", "response"); + Registration registration = new Registration("testTool", "description", "request", "response"); registrationCoordinator.register(registration); - verify(registrationService, times(1)).register(any(Schemas.RegistrationKey.class), eq(registration)); + verify(registrationService, times(1)).register(any(RegistrationKey.class), eq(registration)); } @Test void testUnregister() { String toolName = "testTool"; - registrationCoordinator.register(new Schemas.Registration(toolName, "description", "request", "response")); + registrationCoordinator.register(new Registration(toolName, "description", "request", "response")); registrationCoordinator.unregister(toolName); - verify(registrationService, times(1)).unregister(any(Schemas.RegistrationKey.class)); + verify(registrationService, times(1)).unregister(any(RegistrationKey.class)); } @Test void testOnRegistration() { - Schemas.RegistrationKey key = new Schemas.RegistrationKey("testTool"); - Schemas.Registration registration = new Schemas.Registration("testTool", "description", "request", "response"); - Map registrations = Collections.singletonMap(key, registration); + RegistrationKey key = new RegistrationKey("testTool"); + Registration registration = new Registration("testTool", "description", "request", "response"); + Map registrations = Collections.singletonMap(key, registration); registrationCoordinator.onRegistration(registrations); diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java index 15f8c0b..181b521 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java @@ -1,24 +1,17 @@ package io.confluent.pas.agent.proxy.registration; -import com.fasterxml.jackson.databind.JsonNode; import io.confluent.pas.agent.common.services.KafkaConfiguration; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.proxy.registration.kafka.ConsumerService; import io.confluent.pas.agent.proxy.registration.kafka.ProducerService; -import io.confluent.pas.agent.proxy.registration.schemas.RegistrationSchema; -import io.confluent.pas.agent.proxy.registration.schemas.RegistrationSchemas; import io.micrometer.observation.ObservationRegistry; import org.apache.commons.lang3.RandomStringUtils; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.*; -import reactor.core.publisher.Mono; import java.util.Collection; -import java.util.Map; -import java.util.concurrent.ExecutionException; -import static org.junit.jupiter.api.Assertions.*; import static org.mockito.Mockito.*; class RequestResponseHandlerTest { @@ -56,14 +49,14 @@ void setUp() { @Test void testAddRegistrations() { - Collection registrations = mock(Collection.class); + Collection registrations = mock(Collection.class); requestResponseHandler.addRegistrations(registrations); verify(consumerService, times(1)).addRegistrations(registrations); } // @Test // void testSendRequestResponse() throws ExecutionException, InterruptedException { -// Schemas.Registration registration = new Schemas.Registration("testTool", "description", "requestTopic", "responseTopic"); +// Registration registration = new Registration("testTool", "description", "requestTopic", "responseTopic"); // RegistrationSchemas schemas = mock(RegistrationSchemas.class); // String correlationId = "testCorrelationId"; // Map request = Map.of("key", "value"); diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ResourceHandlerTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ResourceHandlerTest.java index 69776af..d8c28d2 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ResourceHandlerTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ResourceHandlerTest.java @@ -3,8 +3,11 @@ import com.fasterxml.jackson.databind.JsonNode; import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.ResourceRegistration; +import io.confluent.pas.agent.common.services.schemas.ResourceRequest; +import io.confluent.pas.agent.common.services.schemas.ResourceResponse; import io.confluent.pas.agent.proxy.registration.RequestResponseHandler; +import io.confluent.pas.agent.proxy.registration.handlers.mcp.ResourceHandler; import io.modelcontextprotocol.server.McpAsyncServer; import io.modelcontextprotocol.spec.McpSchema; import org.junit.jupiter.api.BeforeEach; @@ -24,7 +27,7 @@ class ResourceHandlerTest { @Mock - private Schemas.ResourceRegistration registration; + private ResourceRegistration registration; @Mock private SchemaRegistryClient schemaRegistryClient; @@ -77,13 +80,13 @@ void testUnregister() { @Test void testSendRequest() throws ExecutionException, InterruptedException { - Schemas.ResourceRequest request = new Schemas.ResourceRequest(); + ResourceRequest request = new ResourceRequest(); Map arguments = Map.of("key", "value"); when(requestResponseHandler.sendRequestResponse(any(), any(), anyString(), anyMap())) .thenReturn(Mono.just(mock(JsonNode.class))); - Mono result = resourceHandler.sendRequest(request); + Mono result = resourceHandler.sendRequest(request); assertNotNull(result); verify(requestResponseHandler).sendRequestResponse(any(), any(), anyString(), anyMap()); } diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ToolHandlerTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ToolHandlerTest.java index 692a44a..7ef8012 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ToolHandlerTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/handlers/ToolHandlerTest.java @@ -4,8 +4,9 @@ import io.confluent.kafka.schemaregistry.client.SchemaMetadata; import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.proxy.registration.RequestResponseHandler; +import io.confluent.pas.agent.proxy.registration.handlers.mcp.ToolHandler; import io.modelcontextprotocol.server.McpAsyncServer; import io.modelcontextprotocol.server.McpServerFeatures; import io.modelcontextprotocol.spec.McpSchema; @@ -26,7 +27,7 @@ class ToolHandlerTest { @Mock - private Schemas.Registration registration; + private Registration registration; @Mock private SchemaRegistryClient schemaRegistryClient; diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java index b72d00c..ce2c927 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java @@ -2,7 +2,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -31,7 +31,7 @@ void setUp() { @Test void testAddRegistrations() { - Schemas.Registration registration = new Schemas.Registration("testTool", "testDescription", "requestTopic", "responseTopic"); + Registration registration = new Registration("testTool", "testDescription", "requestTopic", "responseTopic"); consumerService.addRegistrations(Collections.singletonList(registration)); verify(consumer).subscribe(Collections.singletonList("responseTopic")); @@ -39,7 +39,7 @@ void testAddRegistrations() { @Test void testRegisterResponseHandler() { - Schemas.Registration registration = new Schemas.Registration("testTool", "testDescription", "requestTopic", "responseTopic"); + Registration registration = new Registration("testTool", "testDescription", "requestTopic", "responseTopic"); ConsumerService.ResponseHandler handler = mock(ConsumerService.ResponseHandler.class); ConsumerService.ErrorHandler errorHandler = mock(ConsumerService.ErrorHandler.class); @@ -56,7 +56,7 @@ void testRegisterResponseHandler() { @Test void testHandleResponse() throws IOException { - Schemas.Registration registration = new Schemas.Registration("testTool", "testDescription", "requestTopic", "responseTopic"); + Registration registration = new Registration("testTool", "testDescription", "requestTopic", "responseTopic"); ConsumerService.ResponseHandler handler = mock(ConsumerService.ResponseHandler.class); ConsumerService.ErrorHandler errorHandler = mock(ConsumerService.ErrorHandler.class); diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java b/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java index c3cf001..2d2149b 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java @@ -5,6 +5,8 @@ import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializerConfig; import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializer; import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializerConfig; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.RegistrationKey; import io.confluent.pas.agent.common.utils.SchemaUtils; import io.kcache.KafkaCache; import lombok.extern.slf4j.Slf4j; @@ -23,7 +25,7 @@ */ @Slf4j -public class RegistrationService implements Closeable { +public class RegistrationService implements Closeable { private final Map registrationCache; @@ -176,7 +178,7 @@ public void unregister(K key) { * (optional, can be null) * @return the initialized Kafka cache */ - private static KafkaCache initialize( + private static KafkaCache initialize( KafkaConfiguration kafkaConfiguration, Class registrationKeyClass, Class registrationClass, diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationServiceHandler.java b/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationServiceHandler.java index cc8661d..632e20c 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationServiceHandler.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationServiceHandler.java @@ -1,5 +1,7 @@ package 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.kcache.CacheUpdateHandler; import lombok.Getter; import lombok.extern.slf4j.Slf4j; @@ -17,7 +19,7 @@ * @param the type of registration */ @Slf4j -public class RegistrationServiceHandler implements CacheUpdateHandler { +public class RegistrationServiceHandler implements CacheUpdateHandler { /** * Interface for handling registration updates. @@ -25,7 +27,7 @@ public class RegistrationServiceHandler the type of registration key * @param the type of registration */ - public interface Handler { + public interface Handler { void handleRegistrations(Map registrations); } @@ -58,7 +60,7 @@ public void cacheInitialized(int count, Map checkpoints) { if (count == 0) { // No event, we might need to register the schemas - log.info("No registration found in the cache, registering schemas."); + log.info("No registration found in the cache, registering "); empty = true; } else if (!accumulator.isEmpty()) { handler.handleRegistrations(accumulator); diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/Schemas.java b/common/src/main/java/io/confluent/pas/agent/common/services/Schemas.java deleted file mode 100644 index 2934e26..0000000 --- a/common/src/main/java/io/confluent/pas/agent/common/services/Schemas.java +++ /dev/null @@ -1,537 +0,0 @@ -package io.confluent.pas.agent.common.services; - -import com.fasterxml.jackson.annotation.*; -import io.confluent.kafka.schemaregistry.annotations.Schema; -import io.confluent.pas.agent.common.utils.UriUtils; -import lombok.*; -import org.apache.commons.lang3.StringUtils; -import org.jetbrains.annotations.NotNull; - -public class Schemas { - - @Schema(value = """ - { - "properties":{ - "registrationType":{ - "connect.index":5, - "type":"string" - }, - "correlationIdFieldName":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "description":{ - "connect.index":1, - "type":"string" - }, - "name":{ - "connect.index":0, - "type":"string" - }, - "requestTopicName":{ - "connect.index":2, - "type":"string" - }, - "responseTopicName":{ - "connect.index":3, - "type":"string" - }, - "mimeType":{ - "connect.index":6, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "url":{ - "connect.index":7, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "name", - "description", - "registrationType", - "requestTopicName", - "responseTopicName" - ], - "additionalProperties":false, - "title":"Record", - "type":"object" - }""", refs = {}) - @Getter - @Setter - @NoArgsConstructor - @AllArgsConstructor - @JsonIgnoreProperties(ignoreUnknown = true) - @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "registrationType", defaultImpl = Registration.class) - @JsonSubTypes({ - @JsonSubTypes.Type(value = ResourceRegistration.class, name = Registration.RESOURCE), - @JsonSubTypes.Type(value = Registration.class, name = Registration.TOOL)}) - @Builder - public static class Registration { - public final static String TOOL = "tool"; - public final static String RESOURCE = "resource"; - public final static String CORRELATION_ID_FIELD_NAME = "correlationId"; - - @JsonProperty(value = "registrationType", required = true, defaultValue = TOOL) - private String registrationType; - @JsonProperty(value = "name", required = true) - private String name; - @JsonProperty(value = "description", required = true) - private String description; - @JsonProperty(value = "requestTopicName", required = true) - private String requestTopicName; - @JsonProperty(value = "responseTopicName", required = true) - private String responseTopicName; - @JsonProperty(value = "correlationIdFieldName", defaultValue = CORRELATION_ID_FIELD_NAME) - private String correlationIdFieldName; - - public Registration(String name, String description, String requestTopicName, String responseTopicName) { - this(TOOL, name, description, requestTopicName, responseTopicName, CORRELATION_ID_FIELD_NAME); - } - - public Registration(String name, String description, String requestTopicName, String responseTopicName, String correlationIdFieldName) { - this(TOOL, name, description, requestTopicName, responseTopicName, correlationIdFieldName); - } - - @JsonIgnore - public boolean isResource() { - return StringUtils.equals(registrationType, RESOURCE); - } - } - - @Setter - @Getter - @Schema(value = """ - { - "properties":{ - "registrationType":{ - "connect.index":5, - "type":"string" - }, - "correlationIdFieldName":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "description":{ - "connect.index":1, - "type":"string" - }, - "name":{ - "connect.index":0, - "type":"string" - }, - "requestTopicName":{ - "connect.index":2, - "type":"string" - }, - "responseTopicName":{ - "connect.index":3, - "type":"string" - }, - "mimeType":{ - "connect.index":6, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "url":{ - "connect.index":7, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "name", - "description", - "registrationType", - "requestTopicName", - "responseTopicName" - ], - "additionalProperties":false, - "title":"Record", - "type":"object" - }""", refs = {}) - @AllArgsConstructor - @NoArgsConstructor - @JsonIgnoreProperties(ignoreUnknown = true) - public static class ResourceRegistration extends Registration { - private String mimeType; - private String url; - - public void setUrl(String url) { - if (StringUtils.isBlank(url)) { - throw new IllegalArgumentException("url cannot be blank"); - } - - this.url = url.startsWith("/") ? url.substring(1) : url; - } - - @JsonIgnore - public boolean isTemplate() { - return UriUtils.isTemplate(url); - } - - public ResourceRegistration(String name, - String description, - String requestTopicName, - String responseTopicName, - String correlationIdFieldName, - String mimeType, - String url) { - super(RESOURCE, name, description, requestTopicName, responseTopicName, correlationIdFieldName); - this.mimeType = mimeType; - setUrl(url); - } - - public ResourceRegistration(String name, - String description, - String requestTopicName, - String responseTopicName, - String mimeType, - String url) { - super(RESOURCE, name, description, requestTopicName, responseTopicName, CORRELATION_ID_FIELD_NAME); - this.mimeType = mimeType; - setUrl(url); - } - } - - - /** - * A key for a registration. - */ - @Schema(value = """ - { - "properties": { - "name": { - "connect.index": 0, - "type": "string" - } - }, - "required": [ - "name" - ], - "title": "Record", - "type": "object" - } - """, - refs = {}) - @JsonIgnoreProperties(ignoreUnknown = true) - @Getter - @Setter - @AllArgsConstructor() - @NoArgsConstructor() - public static class RegistrationKey implements Comparable { - - @JsonProperty(value = "name", required = true) - private String name; - - @Override - public int hashCode() { - return name.hashCode(); - } - - @Override - public boolean equals(Object obj) { - return obj == this || obj instanceof RegistrationKey key && StringUtils.equals(name, key.getName()); - } - - @Override - public int compareTo(@NotNull RegistrationKey o) { - return StringUtils.compare(name, o.getName()); - } - } - - @Schema(value = """ - { - "properties":{ - "type":{ - "connect.index":0, - "enum":[ - "text", - "blob" - ], - "default": "text", - "type":"string" - }, - "uri":{ - "connect.index":1, - "type":"string" - }, - "mimeType":{ - "connect.index":2, - "type":"string" - }, - "text":{ - "connect.index":3, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "blob":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "type", - "uri", - "mimeType" - ], - "title":"Record", - "type":"object" - }""", refs = {}) - @Getter - @Setter - @AllArgsConstructor - @NoArgsConstructor - @JsonIgnoreProperties(ignoreUnknown = true) - @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type", defaultImpl = TextResourceResponse.class) - @JsonSubTypes({ - @JsonSubTypes.Type(value = TextResourceResponse.class, name = "text"), - @JsonSubTypes.Type(value = BlobResourceResponse.class, name = "blob")}) - public static class ResourceResponse { - public enum ResponseType { - TEXT("text"), - BLOB("blob"); - - private final String value; - - ResponseType(String value) { - this.value = value; - } - - @JsonValue - public String getValue() { - return value; - } - - @Override - public String toString() { - return value; - } - - @JsonCreator - public static ResourceResponse.ResponseType fromValue(String value) { - for (ResourceResponse.ResponseType e : ResourceResponse.ResponseType.values()) { - if (e.value.equals(value)) { - return e; - } - } - - throw new IllegalArgumentException("Unexpected value '" + value + "'"); - } - } - - @JsonProperty(value = "type", required = true) - private ResponseType type; - @JsonProperty(value = "uri", required = true) - private String uri; - @JsonProperty(value = "mimeType", required = true) - private String mimeType; - } - - @Schema(value = """ - { - "properties":{ - "type":{ - "connect.index":0, - "enum":[ - "text", - "blob" - ], - "default": "text", - "type":"string" - }, - "uri":{ - "connect.index":1, - "type":"string" - }, - "mimeType":{ - "connect.index":2, - "type":"string" - }, - "text":{ - "connect.index":3, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "blob":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "type", - "uri", - "mimeType" - ], - "title":"Record", - "type":"object" - }""", refs = {}) - @Getter - @Setter - @NoArgsConstructor - @JsonIgnoreProperties(ignoreUnknown = true) - public static class BlobResourceResponse extends ResourceResponse { - @JsonProperty(value = "blob", required = true) - private String blob; - - public BlobResourceResponse(String uri, String mimeType, String blob) { - super(ResponseType.BLOB, uri, mimeType); - this.blob = blob; - } - } - - @Schema(value = """ - { - "properties":{ - "type":{ - "connect.index":0, - "enum":[ - "text", - "blob" - ], - "default": "text", - "type":"string" - }, - "uri":{ - "connect.index":1, - "type":"string" - }, - "mimeType":{ - "connect.index":2, - "type":"string" - }, - "text":{ - "connect.index":3, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "blob":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "type", - "uri", - "mimeType" - ], - "title":"Record", - "type":"object" - }""", refs = {}) - @Getter - @Setter - @NoArgsConstructor - @JsonIgnoreProperties(ignoreUnknown = true) - public static class TextResourceResponse extends ResourceResponse { - @JsonProperty(value = "text", required = true) - private String text; - - public TextResourceResponse(String uri, String mimeType, String text) { - super(ResponseType.TEXT, uri, mimeType); - this.text = text; - } - } - - @Schema(value = """ - { - "properties":{ - "uri":{ - "connect.index":0, - "type":"string" - } - }, - "required":[ - "uri" - ], - "title":"Record", - "type":"object" - }""", refs = {}) - @Getter - @Setter - @AllArgsConstructor - @NoArgsConstructor - @JsonIgnoreProperties(ignoreUnknown = true) - public static class ResourceRequest { - @JsonProperty(value = "uri", required = true) - private String uri; - } - -} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java new file mode 100644 index 0000000..eb890ed --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java @@ -0,0 +1,73 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.confluent.kafka.schemaregistry.annotations.Schema; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +@Schema(value = """ + { + "properties":{ + "type":{ + "connect.index":0, + "enum":[ + "text", + "blob" + ], + "default": "text", + "type":"string" + }, + "uri":{ + "connect.index":1, + "type":"string" + }, + "mimeType":{ + "connect.index":2, + "type":"string" + }, + "text":{ + "connect.index":3, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + }, + "blob":{ + "connect.index":4, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + } + }, + "required":[ + "type", + "uri", + "mimeType" + ], + "title":"Record", + "type":"object" + }""", refs = {}) +@Getter +@Setter +@NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +public class BlobResourceResponse extends ResourceResponse { + @JsonProperty(value = "blob", required = true) + private String blob; + + public BlobResourceResponse(String uri, String mimeType, String blob) { + super(ResponseType.BLOB, uri, mimeType); + this.blob = blob; + } +} \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java new file mode 100644 index 0000000..2bed9fd --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java @@ -0,0 +1,116 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.*; +import io.confluent.kafka.schemaregistry.annotations.Schema; +import lombok.*; +import org.apache.commons.lang3.StringUtils; + +@Schema(value = """ + { + "properties":{ + "registrationType":{ + "connect.index":5, + "type":"string" + }, + "correlationIdFieldName":{ + "connect.index":4, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + }, + "description":{ + "connect.index":1, + "type":"string" + }, + "name":{ + "connect.index":0, + "type":"string" + }, + "requestTopicName":{ + "connect.index":2, + "type":"string" + }, + "responseTopicName":{ + "connect.index":3, + "type":"string" + }, + "mimeType":{ + "connect.index":6, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + }, + "url":{ + "connect.index":7, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + } + }, + "required":[ + "name", + "description", + "registrationType", + "requestTopicName", + "responseTopicName" + ], + "additionalProperties":false, + "title":"Record", + "type":"object" + }""", refs = {}) +@Getter +@Setter +@NoArgsConstructor +@AllArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "registrationType", defaultImpl = Registration.class) +@JsonSubTypes({ + @JsonSubTypes.Type(value = ResourceRegistration.class, name = Registration.RESOURCE), + @JsonSubTypes.Type(value = Registration.class, name = Registration.TOOL)}) +@Builder +public class Registration { + public final static String TOOL = "tool"; + public final static String RESOURCE = "resource"; + public final static String CORRELATION_ID_FIELD_NAME = "correlationId"; + + @JsonProperty(value = "registrationType", required = true, defaultValue = TOOL) + private String registrationType; + @JsonProperty(value = "name", required = true) + private String name; + @JsonProperty(value = "description", required = true) + private String description; + @JsonProperty(value = "requestTopicName", required = true) + private String requestTopicName; + @JsonProperty(value = "responseTopicName", required = true) + private String responseTopicName; + @JsonProperty(value = "correlationIdFieldName", defaultValue = CORRELATION_ID_FIELD_NAME) + private String correlationIdFieldName; + + public Registration(String name, String description, String requestTopicName, String responseTopicName) { + this(TOOL, name, description, requestTopicName, responseTopicName, CORRELATION_ID_FIELD_NAME); + } + + public Registration(String name, String description, String requestTopicName, String responseTopicName, String correlationIdFieldName) { + this(TOOL, name, description, requestTopicName, responseTopicName, correlationIdFieldName); + } + + @JsonIgnore + public boolean isResource() { + return StringUtils.equals(registrationType, RESOURCE); + } +} \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/RegistrationKey.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/RegistrationKey.java new file mode 100644 index 0000000..7eeead2 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/RegistrationKey.java @@ -0,0 +1,56 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.confluent.kafka.schemaregistry.annotations.Schema; +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; +import org.apache.commons.lang3.StringUtils; +import org.jetbrains.annotations.NotNull; + +/** + * A key for a registration. + */ +@Schema(value = """ + { + "properties": { + "name": { + "connect.index": 0, + "type": "string" + } + }, + "required": [ + "name" + ], + "title": "Record", + "type": "object" + } + """, + refs = {}) +@JsonIgnoreProperties(ignoreUnknown = true) +@Getter +@Setter +@AllArgsConstructor() +@NoArgsConstructor() +public class RegistrationKey implements Comparable { + + @JsonProperty(value = "name", required = true) + private String name; + + @Override + public int hashCode() { + return name.hashCode(); + } + + @Override + public boolean equals(Object obj) { + return obj == this || obj instanceof RegistrationKey key && StringUtils.equals(name, key.getName()); + } + + @Override + public int compareTo(@NotNull RegistrationKey o) { + return StringUtils.compare(name, o.getName()); + } +} \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java new file mode 100644 index 0000000..ad013b1 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java @@ -0,0 +1,126 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import io.confluent.kafka.schemaregistry.annotations.Schema; +import io.confluent.pas.agent.common.utils.UriUtils; +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; +import org.apache.commons.lang3.StringUtils; + +@Setter +@Getter +@Schema(value = """ + { + "properties":{ + "registrationType":{ + "connect.index":5, + "type":"string" + }, + "correlationIdFieldName":{ + "connect.index":4, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + }, + "description":{ + "connect.index":1, + "type":"string" + }, + "name":{ + "connect.index":0, + "type":"string" + }, + "requestTopicName":{ + "connect.index":2, + "type":"string" + }, + "responseTopicName":{ + "connect.index":3, + "type":"string" + }, + "mimeType":{ + "connect.index":6, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + }, + "url":{ + "connect.index":7, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + } + }, + "required":[ + "name", + "description", + "registrationType", + "requestTopicName", + "responseTopicName" + ], + "additionalProperties":false, + "title":"Record", + "type":"object" + }""", refs = {}) +@AllArgsConstructor +@NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +public class ResourceRegistration extends Registration { + private String mimeType; + private String url; + + public void setUrl(String url) { + if (StringUtils.isBlank(url)) { + throw new IllegalArgumentException("url cannot be blank"); + } + + this.url = url.startsWith("/") ? url.substring(1) : url; + } + + @JsonIgnore + public boolean isTemplate() { + return UriUtils.isTemplate(url); + } + + public ResourceRegistration(String name, + String description, + String requestTopicName, + String responseTopicName, + String correlationIdFieldName, + String mimeType, + String url) { + super(RESOURCE, name, description, requestTopicName, responseTopicName, correlationIdFieldName); + this.mimeType = mimeType; + setUrl(url); + } + + public ResourceRegistration(String name, + String description, + String requestTopicName, + String responseTopicName, + String mimeType, + String url) { + super(RESOURCE, name, description, requestTopicName, responseTopicName, CORRELATION_ID_FIELD_NAME); + this.mimeType = mimeType; + setUrl(url); + } +} + diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRequest.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRequest.java new file mode 100644 index 0000000..794055b --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRequest.java @@ -0,0 +1,33 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.confluent.kafka.schemaregistry.annotations.Schema; +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +@Schema(value = """ + { + "properties":{ + "uri":{ + "connect.index":0, + "type":"string" + } + }, + "required":[ + "uri" + ], + "title":"Record", + "type":"object" + }""", refs = {}) +@Getter +@Setter +@AllArgsConstructor +@NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +public class ResourceRequest { + @JsonProperty(value = "uri", required = true) + private String uri; +} \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java new file mode 100644 index 0000000..fb65191 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java @@ -0,0 +1,109 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.*; +import io.confluent.kafka.schemaregistry.annotations.Schema; +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +@Schema(value = """ + { + "properties":{ + "type":{ + "connect.index":0, + "enum":[ + "text", + "blob" + ], + "default": "text", + "type":"string" + }, + "uri":{ + "connect.index":1, + "type":"string" + }, + "mimeType":{ + "connect.index":2, + "type":"string" + }, + "text":{ + "connect.index":3, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + }, + "blob":{ + "connect.index":4, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + } + }, + "required":[ + "type", + "uri", + "mimeType" + ], + "title":"Record", + "type":"object" + }""", refs = {}) +@Getter +@Setter +@AllArgsConstructor +@NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type", defaultImpl = TextResourceResponse.class) +@JsonSubTypes({ + @JsonSubTypes.Type(value = TextResourceResponse.class, name = "text"), + @JsonSubTypes.Type(value = BlobResourceResponse.class, name = "blob")}) +public class ResourceResponse { + public enum ResponseType { + TEXT("text"), + BLOB("blob"); + + private final String value; + + ResponseType(String value) { + this.value = value; + } + + @JsonValue + public String getValue() { + return value; + } + + @Override + public String toString() { + return value; + } + + @JsonCreator + public static ResourceResponse.ResponseType fromValue(String value) { + for (ResourceResponse.ResponseType e : ResourceResponse.ResponseType.values()) { + if (e.value.equals(value)) { + return e; + } + } + + throw new IllegalArgumentException("Unexpected value '" + value + "'"); + } + } + + @JsonProperty(value = "type", required = true) + private ResponseType type; + @JsonProperty(value = "uri", required = true) + private String uri; + @JsonProperty(value = "mimeType", required = true) + private String mimeType; +} \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java new file mode 100644 index 0000000..dca54a7 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java @@ -0,0 +1,73 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; +import com.fasterxml.jackson.annotation.JsonProperty; +import io.confluent.kafka.schemaregistry.annotations.Schema; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +@Schema(value = """ + { + "properties":{ + "type":{ + "connect.index":0, + "enum":[ + "text", + "blob" + ], + "default": "text", + "type":"string" + }, + "uri":{ + "connect.index":1, + "type":"string" + }, + "mimeType":{ + "connect.index":2, + "type":"string" + }, + "text":{ + "connect.index":3, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + }, + "blob":{ + "connect.index":4, + "oneOf":[ + { + "type":"null" + }, + { + "type":"string" + } + ] + } + }, + "required":[ + "type", + "uri", + "mimeType" + ], + "title":"Record", + "type":"object" + }""", refs = {}) +@Getter +@Setter +@NoArgsConstructor +@JsonIgnoreProperties(ignoreUnknown = true) +public class TextResourceResponse extends ResourceResponse { + @JsonProperty(value = "text", required = true) + private String text; + + public TextResourceResponse(String uri, String mimeType, String text) { + super(ResponseType.TEXT, uri, mimeType); + this.text = text; + } +} \ No newline at end of file diff --git a/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceHandlerTest.java b/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceHandlerTest.java index 23ca02a..558a01d 100644 --- a/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceHandlerTest.java +++ b/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceHandlerTest.java @@ -1,5 +1,7 @@ package 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 org.apache.kafka.common.TopicPartition; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -15,9 +17,9 @@ public class RegistrationServiceHandlerTest { @Mock - private RegistrationServiceHandler.Handler handler; + private RegistrationServiceHandler.Handler handler; - private RegistrationServiceHandler registrationServiceHandler; + private RegistrationServiceHandler registrationServiceHandler; @BeforeEach public void setUp() { @@ -45,8 +47,8 @@ public void testCacheInitializedWithEntries() { @Test public void testHandleUpdateBeforeInitialization() { - Schemas.RegistrationKey key = new Schemas.RegistrationKey("key"); - Schemas.Registration value = new Schemas.Registration(); + RegistrationKey key = new RegistrationKey("key"); + Registration value = new Registration(); TopicPartition tp = new TopicPartition("topic", 0); registrationServiceHandler.handleUpdate(key, value, null, tp, 0L, 0L); @@ -56,8 +58,8 @@ public void testHandleUpdateBeforeInitialization() { @Test public void testHandleUpdateAfterInitialization() { - Schemas.RegistrationKey key = new Schemas.RegistrationKey("Key"); - Schemas.Registration value = new Schemas.Registration(); + RegistrationKey key = new RegistrationKey("Key"); + Registration value = new Registration(); TopicPartition tp = new TopicPartition("topic", 0); registrationServiceHandler.cacheInitialized(1, new HashMap<>()); @@ -68,11 +70,11 @@ public void testHandleUpdateAfterInitialization() { @Test public void testHandleUpdateWithNullValue() { - Schemas.RegistrationKey key = new Schemas.RegistrationKey("Name"); + RegistrationKey key = new RegistrationKey("Name"); TopicPartition tp = new TopicPartition("topic", 0); registrationServiceHandler.cacheInitialized(1, new HashMap<>()); - registrationServiceHandler.handleUpdate(key, null, new Schemas.Registration(), tp, 0L, 0L); + registrationServiceHandler.handleUpdate(key, null, new Registration(), tp, 0L, 0L); verify(handler, times(1)).handleRegistrations(new HashMap<>() {{ put(key, null); diff --git a/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceTest.java b/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceTest.java index f042a60..15baddf 100644 --- a/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceTest.java +++ b/common/src/test/java/io/confluent/pas/agent/common/services/RegistrationServiceTest.java @@ -1,5 +1,7 @@ package 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.kcache.KafkaCache; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -18,12 +20,12 @@ public class RegistrationServiceTest { private KafkaConfiguration kafkaConfiguration; @Mock - private RegistrationServiceHandler.Handler handler; + private RegistrationServiceHandler.Handler handler; @Mock - private KafkaCache registrationCache; + private KafkaCache registrationCache; - private RegistrationService registrationService; + private RegistrationService registrationService; @BeforeEach public void setUp() { @@ -34,18 +36,18 @@ public void setUp() { @Test public void testGetAllRegistrations() { - List registrations = List.of(new Schemas.Registration()); + List registrations = List.of(new Registration()); when(registrationCache.values()).thenReturn(registrations); - List result = registrationService.getAllRegistrations(); + List result = registrationService.getAllRegistrations(); assertEquals(registrations, result); } @Test public void testIsRegistered() { - Schemas.RegistrationKey key = new Schemas.RegistrationKey(); - when(registrationCache.get(key)).thenReturn(new Schemas.Registration()); + RegistrationKey key = new RegistrationKey(); + when(registrationCache.get(key)).thenReturn(new Registration()); boolean result = registrationService.isRegistered(key); @@ -54,8 +56,8 @@ public void testIsRegistered() { @Test public void testRegister() { - Schemas.RegistrationKey key = new Schemas.RegistrationKey(); - Schemas.Registration registration = new Schemas.Registration(); + RegistrationKey key = new RegistrationKey(); + Registration registration = new Registration(); registrationService.register(key, registration); @@ -64,7 +66,7 @@ public void testRegister() { @Test public void testUnregister() { - Schemas.RegistrationKey key = new Schemas.RegistrationKey(); + RegistrationKey key = new RegistrationKey(); registrationService.unregister(key); diff --git a/examples/JavaRCSProvider/src/main/java/io/confluent/pas/mcp/exemple/ResourceAgent.java b/examples/JavaRCSProvider/src/main/java/io/confluent/pas/mcp/exemple/ResourceAgent.java index b07e958..7575d11 100644 --- a/examples/JavaRCSProvider/src/main/java/io/confluent/pas/mcp/exemple/ResourceAgent.java +++ b/examples/JavaRCSProvider/src/main/java/io/confluent/pas/mcp/exemple/ResourceAgent.java @@ -1,6 +1,7 @@ package io.confluent.pas.mcp.exemple; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.ResourceRequest; +import io.confluent.pas.agent.common.services.schemas.TextResourceResponse; import io.confluent.pas.agent.common.utils.UriTemplate; import io.confluent.pas.agent.proxy.frameworks.java.Request; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; @@ -40,16 +41,16 @@ public ResourceAgent() { response_topic = "resource-response", contentType = MIME_TYPE, path = URI, - responseClass = Schemas.TextResourceResponse.class + responseClass = TextResourceResponse.class ) - public void onRequest(Request request) { + public void onRequest(Request request) { log.info("Received request: {}", request.getRequest().getUri()); // Extract values from the URI using the template final Map values = this.template.match(request.getRequest().getUri()); // Respond to the request with a message containing the client_id - request.respond(new Schemas.TextResourceResponse( + request.respond(new TextResourceResponse( request.getRequest().getUri(), MIME_TYPE, "{ \"message\": \"Hello, " + values.get("client_id") + "!\" }" diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java index aed9397..b497418 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java @@ -6,7 +6,8 @@ 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.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.RegistrationKey; import io.confluent.pas.agent.proxy.frameworks.java.kafka.TopicManagement; import io.confluent.pas.agent.proxy.frameworks.java.kafka.impl.TopicManagementImpl; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; @@ -57,7 +58,7 @@ public interface RequestHandler { } private final KafkaConfiguration kafkaConfiguration; - private final RegistrationService registrationService; + private final RegistrationService registrationService; private final Class keyClass; private final Class requestClass; private final Class responseClass; @@ -86,8 +87,8 @@ public SubscriptionHandler(KafkaConfiguration kafkaConfiguration, responseClass, new RegistrationService<>( kafkaConfiguration, - Schemas.RegistrationKey.class, - Schemas.Registration.class), + RegistrationKey.class, + Registration.class), () -> new TopicManagementImpl(kafkaConfiguration), KafkaStreams::new ); @@ -106,7 +107,7 @@ public SubscriptionHandler(KafkaConfiguration kafkaConfiguration, Class keyClass, Class requestClass, Class responseClass, - RegistrationService registrationService, + RegistrationService registrationService, Supplier topicManagementSupplier, KStreamsSupplier kafkaStreamsSupplier) { this.kafkaConfiguration = kafkaConfiguration; @@ -128,7 +129,7 @@ public SubscriptionHandler(KafkaConfiguration kafkaConfiguration, * @param handler Handler to process incoming requests * @throws SubscriptionException if subscription setup fails */ - public void subscribeWith(Schemas.Registration registration, + public void subscribeWith(Registration registration, RequestHandler handler) throws SubscriptionException { log.info("Subscribing for registration: {}", registration.getName()); @@ -141,7 +142,7 @@ public void subscribeWith(Schemas.Registration registration, } /** - * Subscribes to a registration using explicit JSON schemas. + * Subscribes to a registration using explicit JSON * * @param registration Registration containing topic and name information * @param requestSchema Schema for request validation @@ -149,7 +150,7 @@ public void subscribeWith(Schemas.Registration registration, * @param handler Handler to process incoming requests * @throws SubscriptionException if subscription setup fails */ - public void subscribeWith(Schemas.Registration registration, + public void subscribeWith(Registration registration, JsonSchema requestSchema, JsonSchema responseSchema, RequestHandler handler) throws SubscriptionException { @@ -209,7 +210,7 @@ private Serdes.WrapperSerde createSerde(Class valueClass, boolean isKe /** * Creates topics using class types. */ - private void createTopics(Schemas.Registration registration, + private void createTopics(Registration registration, Class keyClass, Class requestClass, Class responseClass) throws Exception { @@ -221,9 +222,9 @@ private void createTopics(Schemas.Registration registration, } /** - * Creates topics using explicit schemas. + * Creates topics using explicit */ - private void createTopicsWithSchemas(Schemas.Registration registration, + private void createTopicsWithSchemas(Registration registration, JsonSchema requestSchema, JsonSchema responseSchema) throws Exception { try (TopicManagement topicManagement = topicManagementSupplier.get()) { @@ -236,7 +237,7 @@ private void createTopicsWithSchemas(Schemas.Registration registration, /** * Registers capability and starts Kafka Streams processing. */ - private void startSubscription(Schemas.Registration registration, + private void startSubscription(Registration registration, RequestHandler handler) { registerCapability(registration); setupAndStartKafkaStreams(registration, handler); @@ -245,8 +246,8 @@ private void startSubscription(Schemas.Registration registration, /** * Registers the capability in the registration service. */ - private void registerCapability(Schemas.Registration registration) { - final Schemas.RegistrationKey registrationKey = new Schemas.RegistrationKey(registration.getName()); + private void registerCapability(Registration registration) { + final RegistrationKey registrationKey = new RegistrationKey(registration.getName()); if (!registrationService.isRegistered(registrationKey)) { log.info("Registering capability: {}", registration.getName()); @@ -259,7 +260,7 @@ private void registerCapability(Schemas.Registration registration) { /** * Sets up and starts the Kafka Streams topology. */ - private void setupAndStartKafkaStreams(Schemas.Registration registration, + private void setupAndStartKafkaStreams(Registration registration, RequestHandler handler) { StreamsBuilder builder = new StreamsBuilder(); diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrar.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrar.java index 8cba53e..c812ee4 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrar.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrar.java @@ -1,7 +1,9 @@ package io.confluent.pas.agent.proxy.frameworks.java.spring.annotation; import io.confluent.pas.agent.common.services.KafkaConfiguration; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.ResourceRegistration; +import io.confluent.pas.agent.common.services.schemas.ResourceRequest; import io.confluent.pas.agent.proxy.frameworks.java.SubscriptionHandler; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; import lombok.Builder; @@ -57,7 +59,7 @@ record InvocationHandler(MethodHandle method, * @param registration The registration information for the subscription. * @return The current InvocationHandler instance. */ - public InvocationHandler subscribe(Schemas.Registration registration) { + public InvocationHandler subscribe(Registration registration) { subscriptionHandler.subscribeWith( registration, (request) -> { @@ -178,7 +180,7 @@ private InvocationHandler getSubscriptionHandler(Method method, Resource resourc log.info("Found resource {} on method {}", resource.name(), method.getName()); // Create registration info for the resource - final Schemas.ResourceRegistration registration = new Schemas.ResourceRegistration( + final ResourceRegistration registration = new ResourceRegistration( resource.name(), resource.description(), resource.request_topic(), @@ -189,7 +191,7 @@ private InvocationHandler getSubscriptionHandler(Method method, Resource resourc // Create and start a subscription handler for the resource var subscriptionHandler = subscriptionHandlerSupplier.get( resource.keyClass(), - Schemas.ResourceRequest.class, + ResourceRequest.class, resource.responseClass()); return getInvocationHandler(method, bean, registration, subscriptionHandler); @@ -208,7 +210,7 @@ private InvocationHandler getSubscriptionHandler(Method method, Agent agent, Obj log.info("Found agent {} on method {}", agent.name(), method.getName()); // Create registration info for the agent - final Schemas.Registration registration = new Schemas.Registration( + final Registration registration = new Registration( agent.name(), agent.description(), agent.request_topic(), @@ -233,7 +235,7 @@ private InvocationHandler getSubscriptionHandler(Method method, Agent agent, Obj * @return Invocation handler for the method */ @NotNull - private InvocationHandler getInvocationHandler(Method method, Object bean, Schemas.Registration registration, SubscriptionHandler subscriptionHandler) { + private InvocationHandler getInvocationHandler(Method method, Object bean, Registration registration, SubscriptionHandler subscriptionHandler) { try { // Create a MethodHandle for the method to allow dynamic invocation MethodHandles.Lookup lookup = MethodHandles.lookup(); diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/Resource.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/Resource.java index 5112af5..f45b234 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/Resource.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/Resource.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.frameworks.java.spring.annotation; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.ResourceResponse; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; import java.lang.annotation.ElementType; @@ -57,5 +57,5 @@ * The class type for the response message payload. * Must extend ResourceResponse. */ - Class responseClass(); + Class responseClass(); } \ No newline at end of file diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java index 5274edf..11b3a9c 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java @@ -1,6 +1,5 @@ package io.confluent.pas.agent.proxy.frameworks.java.spring.autoconfig; -import io.confluent.pas.agent.common.services.Schemas; import lombok.extern.slf4j.Slf4j; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.autoconfigure.AutoConfiguration; @@ -9,6 +8,9 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.context.annotation.Bean; +import io.confluent.pas.agent.common.services.schemas.Registration; +import static io.confluent.pas.agent.common.services.schemas.Registration.CORRELATION_ID_FIELD_NAME; + /** * Auto-configuration class for MCP (Model Control Protocol) agent registration. * Provides configuration for registering agents in the MCP ecosystem. @@ -47,7 +49,7 @@ public class McpRegistrationAutoConfiguration { /** * Field name for request-response correlation, defaults to standard name */ - @Value("${agent.correlation-id:" + Schemas.Registration.CORRELATION_ID_FIELD_NAME + "}") + @Value("${agent.correlation-id:" + CORRELATION_ID_FIELD_NAME + "}") private String correlationIdFieldName; /** @@ -58,8 +60,8 @@ public class McpRegistrationAutoConfiguration { */ @Bean @ConditionalOnMissingBean - public Schemas.Registration getRegistration() { - return Schemas.Registration.builder() + public Registration getRegistration() { + return Registration.builder() .name(name) .description(agentDescription) .requestTopicName(requestTopic) diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/AsyncMcpToolCallbackProvider.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/AsyncMcpToolCallbackProvider.java index 988b848..a542061 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/AsyncMcpToolCallbackProvider.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/AsyncMcpToolCallbackProvider.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.frameworks.java.spring.mcp; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.modelcontextprotocol.client.McpAsyncClient; import lombok.extern.slf4j.Slf4j; import org.jetbrains.annotations.NotNull; @@ -30,7 +30,7 @@ public class AsyncMcpToolCallbackProvider extends McpToolFilters clients) { + public AsyncMcpToolCallbackProvider(Registration registration, List clients) { super(registration); this.clients = clients; } @@ -41,7 +41,7 @@ public AsyncMcpToolCallbackProvider(Schemas.Registration registration, List> { * * @param registration Tool registration information. */ - public McpToolFilters(Schemas.Registration registration) { + public McpToolFilters(Registration registration) { deny(registration.getName()); } diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/SyncMcpToolCallbackProvider.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/SyncMcpToolCallbackProvider.java index 8484827..2566e6e 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/SyncMcpToolCallbackProvider.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/SyncMcpToolCallbackProvider.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.frameworks.java.spring.mcp; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.modelcontextprotocol.client.McpSyncClient; import lombok.extern.slf4j.Slf4j; import org.jetbrains.annotations.NotNull; @@ -29,7 +29,7 @@ public class SyncMcpToolCallbackProvider extends McpToolFilters clients) { + public SyncMcpToolCallbackProvider(Registration registration, List clients) { super(registration); this.clients = clients; } @@ -40,7 +40,7 @@ public SyncMcpToolCallbackProvider(Schemas.Registration registration, List registrationService; + private RegistrationService registrationService; @Mock private TopicManagement topicManagement; @@ -66,7 +67,7 @@ public void setUp() { public void testSubscribeWith() throws Exception { subscriptionHandler.subscribeWith( - new Schemas.Registration( + new Registration( "Name", "Description", "requestTopic", @@ -82,8 +83,8 @@ public void testSubscribeWith() throws Exception { verify(topicManagement, times(1)).createTopic(eq("requestTopic"), any(Class.class), any(Class.class)); verify(topicManagement, times(1)).createTopic(eq("responseTopic"), any(Class.class), any(Class.class)); verify(topicManagement, times(2)).close(); - verify(registrationService, times(1)).isRegistered(any(Schemas.RegistrationKey.class)); - verify(registrationService, times(1)).register(any(Schemas.RegistrationKey.class), any(Schemas.Registration.class)); + verify(registrationService, times(1)).isRegistered(any(RegistrationKey.class)); + verify(registrationService, times(1)).register(any(RegistrationKey.class), any(Registration.class)); } @Test @@ -92,7 +93,7 @@ public void testSubscribeWithSchema() throws Exception { final String resSchema = "{\"type\":\"record\",\"name\":\"Response\",\"fields\":[{\"name\":\"result\",\"type\":\"int\"}]}"; subscriptionHandler.subscribeWith( - new Schemas.Registration( + new Registration( "Name", "Description", "requestTopic", @@ -110,7 +111,7 @@ public void testSubscribeWithSchema() throws Exception { verify(topicManagement, times(1)).createTopic(eq("requestTopic"), any(Class.class), any(JsonSchema.class)); verify(topicManagement, times(1)).createTopic(eq("responseTopic"), any(Class.class), any(JsonSchema.class)); verify(topicManagement, times(2)).close(); - verify(registrationService, times(1)).isRegistered(any(Schemas.RegistrationKey.class)); - verify(registrationService, times(1)).register(any(Schemas.RegistrationKey.class), any(Schemas.Registration.class)); + verify(registrationService, times(1)).isRegistered(any(RegistrationKey.class)); + verify(registrationService, times(1)).register(any(RegistrationKey.class), any(Registration.class)); } } diff --git a/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrarTest.java b/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrarTest.java index 352d05d..afff76c 100644 --- a/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrarTest.java +++ b/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/annotation/AgentRegistrarTest.java @@ -1,7 +1,8 @@ package io.confluent.pas.agent.proxy.frameworks.java.spring.annotation; import io.confluent.pas.agent.common.services.KafkaConfiguration; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; +import io.confluent.pas.agent.common.services.schemas.ResourceResponse; import io.confluent.pas.agent.proxy.frameworks.java.SubscriptionHandler; import lombok.Getter; import lombok.Setter; @@ -26,7 +27,7 @@ record Request(int a, int b) { @Getter @Setter - static class Response extends Schemas.ResourceResponse { + static class Response extends ResourceResponse { public int result; } @@ -70,10 +71,10 @@ public void testAfterPropertiesSet() throws Exception { agentRegistrar.afterPropertiesSet(); // Verify the subscription handler is created and subscribed - ArgumentCaptor registrationCaptor = ArgumentCaptor.forClass(Schemas.Registration.class); + ArgumentCaptor registrationCaptor = ArgumentCaptor.forClass(Registration.class); verify(subscriptionHandler, times(2)).subscribeWith(registrationCaptor.capture(), any()); - List registrations = registrationCaptor.getAllValues(); + List registrations = registrationCaptor.getAllValues(); assertEquals(2, registrations.size()); assertEquals("testAgent", registrations.get(0).getName()); } diff --git a/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/McpToolFiltersTest.java b/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/McpToolFiltersTest.java index 749f069..2807caa 100644 --- a/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/McpToolFiltersTest.java +++ b/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/spring/mcp/McpToolFiltersTest.java @@ -1,6 +1,6 @@ package io.confluent.pas.agent.proxy.frameworks.java.spring.mcp; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.modelcontextprotocol.spec.McpSchema; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -15,7 +15,7 @@ public class McpToolFiltersTest { @BeforeEach public void setUp() { - Schemas.Registration registration = new Schemas.Registration("testTool", "description", "requestTopic", "responseTopic"); + Registration registration = new Registration("testTool", "description", "requestTopic", "responseTopic"); mcpToolFilters = new McpToolFilters<>(registration); } diff --git a/mpc-client-proxy/src/main/java/io/confluent/pas/mcp/proxy/frameworks/client/Agent.java b/mpc-client-proxy/src/main/java/io/confluent/pas/mcp/proxy/frameworks/client/Agent.java index 1dc91eb..fc526b3 100644 --- a/mpc-client-proxy/src/main/java/io/confluent/pas/mcp/proxy/frameworks/client/Agent.java +++ b/mpc-client-proxy/src/main/java/io/confluent/pas/mcp/proxy/frameworks/client/Agent.java @@ -16,7 +16,7 @@ import com.fasterxml.jackson.databind.JsonNode; import io.confluent.kafka.schemaregistry.json.JsonSchema; import io.confluent.pas.agent.common.services.KafkaConfiguration; -import io.confluent.pas.agent.common.services.Schemas; +import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.common.utils.JsonUtils; import io.confluent.pas.mcp.proxy.frameworks.client.internal.*; import io.confluent.pas.agent.proxy.frameworks.java.SubscriptionHandler; @@ -102,7 +102,7 @@ public void destroy() { */ private void setupSubscriptionHandler(List agentToolHandlers) { agentToolHandlers.forEach(handler -> { - final Schemas.Registration registration = new Schemas.Registration( + final Registration registration = new Registration( handler.mcpTool().name(), handler.mcpTool().description(), handler.tool().getRequest_topic(), From 39cc27602d843ba731270c9d70923c279b6c844a Mon Sep 17 00:00:00 2001 From: Pascal Vantrepote Date: Fri, 16 May 2025 15:28:57 -0400 Subject: [PATCH 2/3] Code cleanup --- agent-proxy/pom.xml | 8 +- .../registration/RegistrationCoordinator.java | 12 +- .../registration/RequestResponseChannel.java | 9 +- .../registration/RequestResponseHandler.java | 10 +- .../proxy/registration/kafka/Consumer.java | 219 +++-------------- .../registration/kafka/ConsumerService.java | 9 +- .../registration/kafka/impl/ConsumerImpl.java | 227 ++++++++++++++++++ .../kafka/impl/DistributedConsumer.java | 28 +++ .../proxy/registration/models/WorkItem.java | 19 ++ .../proxy/registration/models/WorkItems.java | 20 ++ .../agent/proxy/rest/a2a/A2AController.java | 2 +- .../src/main/resources/application.yaml | 12 +- .../RequestResponseHandlerTest.java | 2 +- .../kafka/ConsumerServiceTest.java | 1 + .../registration/kafka/ConsumerTest.java | 5 +- .../src/test/resources/application.yaml | 4 +- common/pom.xml | 15 ++ .../pas/agent/common/services/Cache.java | 187 +++++++++++++++ .../agent/common/services/CacheHandler.java | 115 +++++++++ .../services/KafkaPropertiesFactory.java | 25 +- .../common/services/RegistrationService.java | 111 +-------- .../pas/agent/common/services/Streaming.java | 32 +++ .../common/services/cache/LocalCache.java | 170 +++++++++++++ .../services/kstream/StreamingHandler.java | 48 ++++ .../services/kstream/StreamingImpl.java | 132 ++++++++++ .../services/kstream/StreamingProcessor.java | 56 +++++ .../services/kstream/StreamingSupplier.java | 34 +++ .../common/services/schemas/Registration.java | 2 +- examples/JavaAgent/README.md | 4 +- examples/client_info.md | 6 +- frameworks/agent-proxy-framework/pom.xml | 1 - .../java/RequestResponseHandler.java | 52 ++++ .../frameworks/java/SubscriptionHandler.java | 97 ++------ .../java/SubscriptionHandlerProcessor.java | 54 ----- .../java/SubscriptionHandlerSupplier.java | 20 -- .../subscription/SubscriptionRequest.java | 1 - .../java/SubscriptionHandlerTest.java | 22 +- pom.xml | 19 ++ shell/src/main/resources/application.yaml | 2 +- 39 files changed, 1288 insertions(+), 504 deletions(-) create mode 100644 agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/ConsumerImpl.java create mode 100644 agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/DistributedConsumer.java create mode 100644 agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItem.java create mode 100644 agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItems.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/Cache.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/CacheHandler.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/Streaming.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/cache/LocalCache.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingHandler.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingImpl.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingProcessor.java create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingSupplier.java create mode 100644 frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/RequestResponseHandler.java delete mode 100644 frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerProcessor.java delete mode 100644 frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerSupplier.java 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/registration/RegistrationCoordinator.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RegistrationCoordinator.java index 3f3e67c..cd6e53b 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; @@ -29,7 +26,7 @@ *

* It coordinates the following processes: * - Listening for new registrations on the registration topic - * - Processing incoming registrations and unregistrations + * - Processing incoming registrations and un-registrations * - Creating and managing handlers for each registered tool * - Maintaining the lifecycle of registrations * - Broadcasting registration events to other components @@ -50,9 +47,6 @@ public class RegistrationCoordinator implements DisposableBean { @Getter private final A2AAsyncServer a2AAsyncServer; - /** - * REST server for handling HTTP/REST communications - */ @Getter private final AgentAsyncServer restServer; @@ -105,7 +99,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); } 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..b28a083 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; @@ -25,23 +24,18 @@ public class RequestResponseHandler implements DisposableBean { private final ProducerService producerService; private final ConsumerService consumerService; - private final ObservationRegistry observationRegistry; @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)); } public RequestResponseHandler(ProducerService producerService, - ConsumerService consumerService, - ObservationRegistry observationRegistry) { + ConsumerService consumerService) { this.producerService = producerService; this.consumerService = consumerService; - this.observationRegistry = observationRegistry; } public void addRegistrations(Collection registrations) { 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: *

    *
  • Subscribes to response topics for registered services
    • *
  • Routes messages to the correct handler using correlation IDs
    • @@ -82,8 +82,7 @@ public record RegistrationItem( } /** - * Map of topic names to their registration items, which include handlers - * indexed by correlation ID. + * Map of topic names to their registration items, which include handlers indexed by correlation ID. */ @Getter private final Map responseHandlers = new ConcurrentHashMap<>(); @@ -103,7 +102,7 @@ public record RegistrationItem( * timing out */ public ConsumerService(KafkaConfiguration kafkaConfiguration, long responseTimeout) { - this.consumer = new Consumer<>( + this.consumer = new ConsumerImpl<>( kafkaConfiguration, JsonNode.class, this::handleResponse, diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/ConsumerImpl.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/ConsumerImpl.java new file mode 100644 index 0000000..8ffb2ff --- /dev/null +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/ConsumerImpl.java @@ -0,0 +1,227 @@ +package io.confluent.pas.agent.proxy.registration.kafka.impl; + +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 io.confluent.pas.agent.proxy.registration.kafka.Consumer; +import io.confluent.pas.agent.proxy.registration.kafka.ConsumerHandler; +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; + +/** + * Consumer class for managing Kafka topic subscriptions and message processing. + * + * @param Key type for Kafka messages + * @param Value type for Kafka messages + */ +@Slf4j +public class ConsumerImpl implements Closeable, Consumer { + + 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; + + /** + * 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 + */ + public ConsumerImpl(KafkaConfiguration kafkaConfiguration, + Class requestClass, + ConsumerHandler consumerHandler, + TimeoutChecker timeoutChecker) { + this(consumerHandler, + new KafkaConsumer<>(KafkaPropertiesFactory.getConsumerProperties( + kafkaConfiguration, + false, + Key.class, + requestClass)), + timeoutChecker); + } + + /** + * Constructs a Consumer instance with the specified Kafka consumer and handler. + * + * @param consumerHandler The handler to process messages + * @param kafkaConsumer The Kafka consumer to use + * @param timeoutChecker The timeout checker to use + */ + public ConsumerImpl(ConsumerHandler consumerHandler, + KafkaConsumer kafkaConsumer, + TimeoutChecker timeoutChecker) { + this.kafkaConsumer = kafkaConsumer; + this.consumerHandler = consumerHandler; + this.timeoutChecker = timeoutChecker; + + executorSvc.submit(this::runLoop); + } + + @Override + public boolean isSubscribed(String topic) { + return topics.contains(topic); + } + + /** + * Subscribes to a Kafka topic with the specified handler. + * + * @param topic The topic to subscribe to + */ + @Override + public void subscribe(String topic) { + if (topics.contains(topic)) { + throw new IllegalArgumentException("Subscription already exists for topic: " + topic); + } + + topics.add(topic); + subscriptionUpdated = true; + } + + /** + * Subscribes to a Kafka topic with the specified handler. + * + * @param topicsToAdd The topics to subscribe to + */ + @Override + public void subscribe(Collection topicsToAdd) { + topics.addAll(topicsToAdd); + subscriptionUpdated = true; + } + + /** + * Unsubscribes from a Kafka topic. + * + * @param topic The topic to unsubscribe from + */ + @Override + public void unsubscribe(String topic) { + topics.remove(topic); + subscriptionUpdated = true; + } + + /** + * Closes the consumer, stopping the polling loop and releasing resources. + * + * @throws IOException if an error occurs while closing the consumer + */ + @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. + */ + public 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 + */ + @Override + public 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 diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/DistributedConsumer.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/DistributedConsumer.java new file mode 100644 index 0000000..edd6f88 --- /dev/null +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/kafka/impl/DistributedConsumer.java @@ -0,0 +1,28 @@ +package io.confluent.pas.agent.proxy.registration.kafka.impl; + +import com.github.benmanes.caffeine.cache.Cache; +import com.github.benmanes.caffeine.cache.Caffeine; +import io.confluent.pas.agent.common.services.KafkaConfiguration; +import io.confluent.pas.agent.proxy.registration.kafka.ConsumerHandler; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.kafka.clients.consumer.KafkaConsumer; + +import java.io.Closeable; + +public class DistributedConsumer extends ConsumerImpl implements Closeable { + + private final Cache> cache = Caffeine.newBuilder() + .maximumSize(10000) + .expireAfterWrite(5, java.util.concurrent.TimeUnit.MINUTES) + .build(); + + public DistributedConsumer(KafkaConfiguration kafkaConfiguration, Class requestClass, ConsumerHandler consumerHandler, TimeoutChecker timeoutChecker) { + super(kafkaConfiguration, requestClass, consumerHandler, timeoutChecker); + } + + public DistributedConsumer(ConsumerHandler consumerHandler, KafkaConsumer kafkaConsumer, TimeoutChecker timeoutChecker) { + super(consumerHandler, kafkaConsumer, timeoutChecker); + } + + +} diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItem.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItem.java new file mode 100644 index 0000000..e766995 --- /dev/null +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItem.java @@ -0,0 +1,19 @@ +package io.confluent.pas.agent.proxy.registration.models; + +import io.confluent.pas.agent.proxy.frameworks.java.models.Request; +import io.confluent.pas.agent.proxy.frameworks.java.models.Response; +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +@Getter +@Setter +@NoArgsConstructor +@AllArgsConstructor +public class WorkItem { + + private Request request; + private Response response; + +} diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItems.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItems.java new file mode 100644 index 0000000..1e3bf78 --- /dev/null +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/models/WorkItems.java @@ -0,0 +1,20 @@ +package io.confluent.pas.agent.proxy.registration.models; + +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +import java.util.ArrayList; +import java.util.List; + +@Getter +@Setter +@AllArgsConstructor +@NoArgsConstructor +public class WorkItems { + + private String correlationId; + private List workItems = new ArrayList<>(); + +} diff --git a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/a2a/A2AController.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/a2a/A2AController.java index 77d4ce8..30fc504 100644 --- a/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/a2a/A2AController.java +++ b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/rest/a2a/A2AController.java @@ -9,7 +9,7 @@ import java.util.List; -//@RestController +@RestController @Tag(name = "A2A", description = "Google A2A compliant API") public class A2AController { private final static String AGENT_PATH = "/a2a/"; diff --git a/agent-proxy/src/main/resources/application.yaml b/agent-proxy/src/main/resources/application.yaml index bf2c57a..8831955 100644 --- a/agent-proxy/src/main/resources/application.yaml +++ b/agent-proxy/src/main/resources/application.yaml @@ -3,7 +3,7 @@ spring: activate: on-profile: "production" application: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" mvc: async: request-timeout: 60000 @@ -11,7 +11,7 @@ spring: web-application-type: reactive mcp: server: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" version: 1.0.0 mode: "sse" kafka: @@ -48,7 +48,7 @@ spring: activate: on-profile: "default" application: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" mvc: async: request-timeout: 60000 @@ -56,7 +56,7 @@ spring: web-application-type: reactive mcp: server: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" version: 1.0.0 mode: "sse" kafka: @@ -87,13 +87,13 @@ spring: activate: on-profile: "stdio" application: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" main: banner-mode: off web-application-type: none mcp: server: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" version: 1.0.0 mode: "stdio" kafka: diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java index 181b521..df4e617 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandlerTest.java @@ -44,7 +44,7 @@ void setUp() { when(kafkaConfiguration.topicConfiguration()).thenReturn(new KafkaConfiguration.DefaultTopicConfiguration()); when(kafkaConfiguration.saslMechanism()).thenReturn(KafkaConfiguration.DEFAULT_SASL_MECHANISM); - requestResponseHandler = new RequestResponseHandler(producerService, consumerService, observationRegistry); + requestResponseHandler = new RequestResponseHandler(producerService, consumerService); } @Test diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java index eff04a9..0ceb139 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerServiceTest.java @@ -4,6 +4,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; 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 org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; diff --git a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerTest.java b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerTest.java index 5de4470..c1105a6 100644 --- a/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerTest.java +++ b/agent-proxy/src/test/java/io/confluent/pas/agent/proxy/registration/kafka/ConsumerTest.java @@ -1,6 +1,7 @@ package io.confluent.pas.agent.proxy.registration.kafka; import io.confluent.pas.agent.common.services.KafkaConfiguration; +import io.confluent.pas.agent.proxy.registration.kafka.impl.ConsumerImpl; import org.apache.commons.lang3.RandomStringUtils; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.KafkaConsumer; @@ -28,7 +29,7 @@ class ConsumerTest { @Mock private KafkaConfiguration kafkaConfiguration; - private Consumer consumer; + private ConsumerImpl consumer; @BeforeEach void setUp() { @@ -44,7 +45,7 @@ void setUp() { when(kafkaConfiguration.saslMechanism()).thenReturn(KafkaConfiguration.DEFAULT_SASL_MECHANISM); - consumer = new Consumer<>(consumerHandler, kafkaConsumer, mock(Consumer.TimeoutChecker.class)); + consumer = new ConsumerImpl<>(consumerHandler, kafkaConsumer, mock(ConsumerImpl.TimeoutChecker.class)); } @Test diff --git a/agent-proxy/src/test/resources/application.yaml b/agent-proxy/src/test/resources/application.yaml index fb98d7a..77bbd99 100644 --- a/agent-proxy/src/test/resources/application.yaml +++ b/agent-proxy/src/test/resources/application.yaml @@ -1,6 +1,6 @@ spring: application: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" mvc: async: request-timeout: 60000 @@ -8,7 +8,7 @@ spring: web-application-type: reactive mcp: server: - name: "Confluent MCP Proxy" + name: "Confluent Agent Proxy" version: 1.0.0 mode: "sse" authentication: diff --git a/common/pom.xml b/common/pom.xml index bc5211f..b42539c 100644 --- a/common/pom.xml +++ b/common/pom.xml @@ -22,6 +22,17 @@ + + + com.github.ben-manes.caffeine + caffeine + + + + org.apache.kafka + kafka-streams + + com.github.victools jsonschema-generator @@ -46,6 +57,10 @@ io.kcache kcache + + io.kcache + kcache-caffeine + org.mockito diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/Cache.java b/common/src/main/java/io/confluent/pas/agent/common/services/Cache.java new file mode 100644 index 0000000..9e3d787 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/Cache.java @@ -0,0 +1,187 @@ +package io.confluent.pas.agent.common.services; + +import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializer; +import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializerConfig; +import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializer; +import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializerConfig; +import io.confluent.pas.agent.common.services.cache.LocalCache; +import io.kcache.CacheUpdateHandler; +import io.kcache.KafkaCache; +import org.apache.kafka.common.serialization.Serde; +import org.apache.kafka.common.serialization.Serdes; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +import java.time.Duration; +import java.time.temporal.ChronoUnit; +import java.util.Collection; +import java.util.Map; +import java.util.Set; + +/** + * A generic cache implementation that wraps KafkaCache to provide distributed caching capabilities. + * This cache uses Kafka as a backing store and supports JSON schema serialization/deserialization. + * + * @param The type of keys in the cache + * @param The type of values in the cache + */ +public class Cache implements Map { + + private final KafkaCache cache; + + /** + * Creates a new Cache instance with the specified configuration. + * + * @param kafkaConfiguration The Kafka configuration settings + * @param cacheName The name of the cache + * @param kClass The class type of the keys + * @param vClass The class type of the values + * @param handler The handler for cache updates + * @param readOnly Whether the cache is read-only + * @param topicName The name of the Kafka topic to use + */ + public Cache(KafkaConfiguration kafkaConfiguration, + String cacheName, + Class kClass, + Class vClass, + CacheHandler.Handler handler, + boolean readOnly, + String topicName) { + this.cache = initialize( + kafkaConfiguration, + cacheName, + kClass, + vClass, + handler, + readOnly, + topicName); + } + + public Cache(KafkaCache cache) { + this.cache = cache; + } + + @Override + public int size() { + return cache.size(); + } + + @Override + public boolean isEmpty() { + return cache.isEmpty(); + } + + @Override + public boolean containsKey(Object key) { + return cache.containsKey(key); + } + + @Override + public boolean containsValue(Object value) { + return cache.containsValue(value); + } + + @Override + public V get(Object key) { + return cache.get(key); + } + + @Override + public @Nullable V put(K key, V value) { + return cache.put(key, value); + } + + @Override + public V remove(Object key) { + return cache.remove(key); + } + + @Override + public void putAll(@NotNull Map m) { + cache.putAll(m); + } + + @Override + public void clear() { + cache.clear(); + } + + @Override + public @NotNull Set keySet() { + return cache.keySet(); + } + + @Override + public @NotNull Collection values() { + return cache.values(); + } + + @Override + public @NotNull Set> entrySet() { + return cache.entrySet(); + } + + /** + * Initializes the KafkaCache with the specified configuration. + * Sets up JSON schema serialization/deserialization and configures the cache storage. + * + * @param kafkaConfiguration The Kafka configuration settings + * @param cacheName The name of the cache + * @param kClass The class type of the keys + * @param vClass The class type of the values + * @param handler The handler for cache updates + * @param readOnly Whether the cache is read-only + * @param topicName The name of the Kafka topic to use + * @return The initialized KafkaCache instance + */ + private KafkaCache initialize(KafkaConfiguration kafkaConfiguration, + String cacheName, + Class kClass, + Class vClass, + CacheHandler.Handler handler, + boolean readOnly, + String topicName) { + // Retrieve the Schema Registry configuration settings from KafkaPropertiesFactory. + final Map srConfig = KafkaPropertiesFactory.getSchemaRegistryConfig(kafkaConfiguration); + srConfig.put(KafkaJsonSchemaDeserializerConfig.JSON_KEY_TYPE, kClass); + srConfig.put(KafkaJsonSchemaDeserializerConfig.JSON_VALUE_TYPE, vClass); + srConfig.put(KafkaJsonSchemaSerializerConfig.AUTO_REGISTER_SCHEMAS, true); + + // Create and configure the Serde for the registration key using Kafka JSON schema serializer and deserializer. + final Serde keySerdes = new Serdes.WrapperSerde<>( + new KafkaJsonSchemaSerializer<>(), + new KafkaJsonSchemaDeserializer<>() + ); + keySerdes.configure(srConfig, true); + + // Create and configure the Serde for the registration value using Kafka JSON schema serializer and deserializer. + final Serde valueSerdes = new Serdes.WrapperSerde<>( + new KafkaJsonSchemaSerializer<>(), + new KafkaJsonSchemaDeserializer<>() + ); + valueSerdes.configure(srConfig, false); + + final CacheUpdateHandler serviceHandler = handler != null + ? new CacheHandler<>(handler) + : null; + + final KafkaCache cache = new KafkaCache<>( + KafkaPropertiesFactory.getCacheConfig( + kafkaConfiguration, + readOnly, + topicName, + cacheName), + keySerdes, + valueSerdes, + serviceHandler, + new LocalCache<>( + 1000, + Duration.of(5, ChronoUnit.MINUTES), + null) + ); + + cache.init(); + + return cache; + } +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/CacheHandler.java b/common/src/main/java/io/confluent/pas/agent/common/services/CacheHandler.java new file mode 100644 index 0000000..5f658a1 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/CacheHandler.java @@ -0,0 +1,115 @@ +package io.confluent.pas.agent.common.services; + +import io.kcache.CacheUpdateHandler; +import lombok.Getter; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.common.TopicPartition; + +import java.util.HashMap; +import java.util.Map; + + +/** + * Handles cache updates and initialization for a Kafka-based caching system. + * This class implements CacheUpdateHandler to manage cache events and updates, + * providing a way to accumulate and process cache entries during initialization. + * + * @param The type of keys in the cache + * @param The type of values in the cache + */ +@Slf4j +public class CacheHandler implements CacheUpdateHandler { + + /** + * Interface defining the contract for handling cache values. + * Implementations process batches of key-value pairs from the cache. + * + * @param The type of keys in the cache + * @param The type of values in the cache + */ + public interface Handler { + /** + * Processes a batch of cache entries. + * + * @param registrations Map of key-value pairs to process + */ + void onValues(Map registrations); + } + + /** + * Handler for processing cache values + */ + private final Handler handler; + /** + * Temporary storage for cache entries during initialization + */ + private final Map accumulator = new HashMap<>(); + /** + * Flag indicating if the cache has been initialized + */ + private boolean initialized = false; + + /** + * Flag indicating if the cache is empty after initialization + */ + @Getter + private boolean empty; + + public CacheHandler(Handler handler) { + this.handler = handler; + } + + /** + * Called when the cache is initialized. + * Processes accumulated entries if any exist, or marks the cache as empty if no entries were found. + * + * @param count Number of entries in the cache + * @param checkpoints Map of topic partitions and their offsets + */ + @Override + public void cacheInitialized(int count, Map checkpoints) { + initialized = true; + + if (count == 0) { + // No event, we might need to register the schemas + empty = true; + } else if (!accumulator.isEmpty()) { + handler.onValues(accumulator); + accumulator.clear(); + } + + log.info("Registration cache initialized."); + } + + /** + * Handles updates to the cache entries. + * If the cache is initialized, processes updates immediately. + * Otherwise, accumulates updates for later processing. + * + * @param key Key being updated + * @param value New value + * @param oldValue Previous value + * @param tp Topic partition + * @param offset Kafka offset + * @param ts Timestamp of the update + */ + @Override + public void handleUpdate(K key, + V value, + V oldValue, + TopicPartition tp, + long offset, + long ts) { + // If the cache is not initialized, store the value in the accumulator + // This will help to keep the last event for each key + if (initialized) { + handler.onValues(Map.of(key, value)); + } else { + if (value == null) { + accumulator.remove(key); + } else { + accumulator.put(key, value); + } + } + } +} \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/KafkaPropertiesFactory.java b/common/src/main/java/io/confluent/pas/agent/common/services/KafkaPropertiesFactory.java index bcc9bcc..a33afec 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/KafkaPropertiesFactory.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/KafkaPropertiesFactory.java @@ -146,11 +146,28 @@ public static Properties getConsumerProperties(KafkaConfiguration configration, * @param readOnly Whether the cache should be read-only * @return KafkaCacheConfig configured for the cache */ - public static KafkaCacheConfig getCacheConfig(KafkaConfiguration configration, boolean readOnly) { + public static KafkaCacheConfig getRegistrationCacheConfig(KafkaConfiguration configration, boolean readOnly) { + return getCacheConfig(configration, readOnly, configration.registrationTopicName(), "registration"); + } + + /** + * Creates configuration for a Kafka cache. + * Sets up topic, client ID, and group ID with appropriate suffixes. + * + * @param configration The Kafka configuration containing connection and auth details + * @param readOnly Whether the cache should be read-only + * @param topicName The name of the topic + * @param clientIdSuffix The suffix to add to the client ID to create a unique ID for the cache + * @return KafkaCacheConfig configured for the cache + */ + public static KafkaCacheConfig getCacheConfig(KafkaConfiguration configration, + boolean readOnly, + String topicName, + String clientIdSuffix) { Properties properties = getDefaultProperties(configration, "kafkacache."); - properties.put(KafkaCacheConfig.KAFKACACHE_TOPIC_CONFIG, configration.registrationTopicName()); - properties.put(KafkaCacheConfig.KAFKACACHE_CLIENT_ID_CONFIG, configration.applicationId() + "-registration" + "-" + configration.clientId()); - properties.put(KafkaCacheConfig.KAFKACACHE_GROUP_ID_CONFIG, configration.applicationId() + "-registration" + "-group"); + properties.put(KafkaCacheConfig.KAFKACACHE_TOPIC_CONFIG, topicName); + properties.put(KafkaCacheConfig.KAFKACACHE_CLIENT_ID_CONFIG, configration.applicationId() + "-" + clientIdSuffix + "-" + configration.clientId()); + properties.put(KafkaCacheConfig.KAFKACACHE_GROUP_ID_CONFIG, configration.applicationId() + "-" + clientIdSuffix + "-group"); properties.put(KafkaCacheConfig.KAFKACACHE_TOPIC_READ_ONLY_CONFIG, readOnly); return new KafkaCacheConfig(properties); } diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java b/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java index 2d2149b..7b89509 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/RegistrationService.java @@ -1,17 +1,8 @@ package io.confluent.pas.agent.common.services; -import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient; -import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializer; -import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializerConfig; -import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializer; -import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializerConfig; import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.common.services.schemas.RegistrationKey; -import io.confluent.pas.agent.common.utils.SchemaUtils; -import io.kcache.KafkaCache; import lombok.extern.slf4j.Slf4j; -import org.apache.kafka.common.serialization.Serde; -import org.apache.kafka.common.serialization.Serdes; import java.io.Closeable; import java.io.IOException; @@ -42,12 +33,15 @@ public RegistrationService(KafkaConfiguration kafkaConfiguration, Class registrationKeyClass, Class registrationClass, boolean readOnly, - RegistrationServiceHandler.Handler handler) { - this(initialize(kafkaConfiguration, + CacheHandler.Handler handler) { + registrationCache = new Cache<>( + kafkaConfiguration, + "registration-cache", registrationKeyClass, registrationClass, + handler, readOnly, - handler)); + kafkaConfiguration.registrationTopicName()); } /** @@ -70,24 +64,12 @@ public RegistrationService(Map registrationCache) { public RegistrationService(KafkaConfiguration kafkaConfiguration, Class registrationKeyClass, Class registrationClass, - RegistrationServiceHandler.Handler handler) { - this(kafkaConfiguration, registrationKeyClass, registrationClass, false, handler); - } - - - /** - * Constructor for RegistrationService without a handler. - * - * @param kafkaConfiguration the Kafka configuration - * @param registrationKeyClass the class type of the registration key - * @param registrationClass the class type of the registration - * @param readOnly whether the service is read-only - */ - public RegistrationService(KafkaConfiguration kafkaConfiguration, - Class registrationKeyClass, - Class registrationClass, - boolean readOnly) { - this(kafkaConfiguration, registrationKeyClass, registrationClass, readOnly, null); + CacheHandler.Handler handler) { + this(kafkaConfiguration, + registrationKeyClass, + registrationClass, + false, + handler); } /** @@ -164,73 +146,4 @@ public void unregister(K key) { // Remove the registration associated with the given key from the Kafka cache. registrationCache.remove(key); } - - - /** - * Initialize the registration service. - * Configures the Kafka serializers and deserializers, and initializes the Kafka cache. - * - * @param kafkaConfiguration the Kafka configuration - * @param registrationKeyClass the class type of the registration key - * @param registrationClass the class type of the registration - * @param readOnly whether the service is read-only - * @param handler the handler for processing registration updates - * (optional, can be null) - * @return the initialized Kafka cache - */ - private static KafkaCache initialize( - KafkaConfiguration kafkaConfiguration, - Class registrationKeyClass, - Class registrationClass, - boolean readOnly, - RegistrationServiceHandler.Handler handler) { - // Retrieve the Schema Registry configuration settings from KafkaPropertiesFactory. - final Map srConfig = KafkaPropertiesFactory.getSchemaRegistryConfig(kafkaConfiguration); - srConfig.put(KafkaJsonSchemaDeserializerConfig.JSON_KEY_TYPE, registrationKeyClass); - srConfig.put(KafkaJsonSchemaDeserializerConfig.JSON_VALUE_TYPE, registrationClass); - srConfig.put(KafkaJsonSchemaSerializerConfig.AUTO_REGISTER_SCHEMAS, true); - - // Create and configure the Serde for the registration key using Kafka JSON schema serializer and deserializer. - final Serde keySerdes = new Serdes.WrapperSerde<>( - new KafkaJsonSchemaSerializer<>(), - new KafkaJsonSchemaDeserializer<>() - ); - keySerdes.configure(srConfig, true); - - // Create and configure the Serde for the registration value using Kafka JSON schema serializer and deserializer. - final Serde valueSerdes = new Serdes.WrapperSerde<>( - new KafkaJsonSchemaSerializer<>(), - new KafkaJsonSchemaDeserializer<>() - ); - valueSerdes.configure(srConfig, false); - - final RegistrationServiceHandler serviceHandler = handler != null - ? new RegistrationServiceHandler<>(handler) - : null; - - final KafkaCache cache = new KafkaCache<>( - KafkaPropertiesFactory.getCacheConfig(kafkaConfiguration, readOnly), - keySerdes, - valueSerdes, - serviceHandler, - null - ); - - cache.init(); - - if (serviceHandler != null && serviceHandler.isEmpty()) { - final String registrationTopic = kafkaConfiguration.registrationTopicName(); - - // If a service handler is provided and is empty, register the necessary schemas in the Schema Registry. - try (SchemaRegistryClient schemaRegistryClient = KafkaPropertiesFactory.getSchemRegistryClient(kafkaConfiguration)) { - SchemaUtils.registerSchemaIfMissing(registrationTopic, registrationKeyClass, true, schemaRegistryClient); - SchemaUtils.registerSchemaIfMissing(registrationTopic, registrationClass, false, schemaRegistryClient); - } catch (Throwable e) { - log.error("Error registering schemas", e); - } - - } - - return cache; - } } \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/Streaming.java b/common/src/main/java/io/confluent/pas/agent/common/services/Streaming.java new file mode 100644 index 0000000..f445826 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/Streaming.java @@ -0,0 +1,32 @@ +package io.confluent.pas.agent.common.services; + +import io.confluent.pas.agent.common.services.kstream.StreamingHandler; + +/** + * Interface for managing streaming operations with Kafka. + * + * @param The type of the message key + * @param The type of the input message value + * @param The type of the output message value + */ +public interface Streaming extends AutoCloseable { + + /** + * Initializes the streaming process with the given configuration and topics. + * + * @param kafkaConfiguration Configuration for Kafka connection + * @param inTopicName Name of the input topic + * @param outTopicName Name of the output topic + * @param streamingHandler Handler for processing the stream + */ + void init(KafkaConfiguration kafkaConfiguration, + String inTopicName, + String outTopicName, + StreamingHandler streamingHandler); + + /** + * Starts the streaming process. + * This method initializes and begins the Kafka streams processing. + */ + void start(); +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/cache/LocalCache.java b/common/src/main/java/io/confluent/pas/agent/common/services/cache/LocalCache.java new file mode 100644 index 0000000..a3673ef --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/cache/LocalCache.java @@ -0,0 +1,170 @@ +package io.confluent.pas.agent.common.services.cache; + +import com.github.benmanes.caffeine.cache.*; +import io.kcache.CacheLoader; +import io.kcache.utils.InMemoryCache; +import org.jetbrains.annotations.NotNull; + +import java.time.Duration; +import java.util.Comparator; +import java.util.Map; + +/** + * A local caching implementation that extends InMemoryCache with Caffeine caching capabilities. + * This cache supports maximum size limits, expiration policies, and optional loading functionality. + * + * @param The type of keys maintained by this cache + * @param The type of mapped values + */ +public class LocalCache extends InMemoryCache { + private final Cache cache; + /** + * The loader used to load values when they are not present in the cache. + * Can be null if no loading functionality is required. + */ + private final CacheLoader loader; + + /** + * Constructs a LocalCache with a specified comparator for key ordering. + * + * @param comparator the comparator to determine the ordering of keys + */ + public LocalCache(Comparator comparator) { + this(null, null, null, comparator); + } + + /** + * Constructs a LocalCache with size limits, expiration policy, and loading functionality. + * + * @param maximumSize the maximum number of entries the cache may contain + * @param expireAfterWrite the time after which entries should be automatically removed + * @param loader the cache loader to use when entries are not found + */ + public LocalCache(Integer maximumSize, Duration expireAfterWrite, CacheLoader loader) { + this.loader = loader; + this.cache = this.createCache(maximumSize, expireAfterWrite); + } + + /** + * Constructs a LocalCache with size limits, expiration policy, loading functionality, and key ordering. + * + * @param maximumSize the maximum number of entries the cache may contain + * @param expireAfterWrite the time after which entries should be automatically removed + * @param loader the cache loader to use when entries are not found + * @param comparator the comparator to determine the ordering of keys + */ + public LocalCache(Integer maximumSize, Duration expireAfterWrite, CacheLoader loader, Comparator comparator) { + super(comparator); + this.loader = loader; + this.cache = this.createCache(maximumSize, expireAfterWrite); + } + + /** + * Returns true if this cache contains a mapping for the specified key. + * + * @param key key whose presence in this cache is to be tested + * @return true if this cache contains a mapping for the specified key + */ + public boolean containsKey(Object key) { + return this.get(key) != null; + } + + /** + * Returns the value associated with the specified key, or null if there is no cached value. + * If a loader is configured and the key isn't present, it will attempt to load the value. + * + * @param key the key whose associated value is to be returned + * @return the value associated with the specified key, or null if none + */ + @SuppressWarnings("unchecked") + public V get(Object key) { + return (cache instanceof LoadingCache loadingCache) + ? loadingCache.get((K) key) + : super.get(key); + } + + /** + * Associates the specified value with the specified key in this cache. + * + * @param key key with which the specified value is to be associated + * @param value value to be associated with the specified key + * @return the previous value associated with key, or null if there was no mapping + */ + public V put(K key, V value) { + V originalValue = this.get(key); + + this.cache.put(key, value); + delegate().put(key, value); + + return originalValue; + } + + /** + * Copies all entries from the specified map to this cache. + * + * @param entries mappings to be stored in this cache + */ + public void putAll(@NotNull Map entries) { + this.cache.putAll(entries); + for (Map.Entry entry : entries.entrySet()) { + delegate().put(entry.getKey(), entry.getValue()); + } + } + + /** + * Removes the entry for the specified key if present. + * + * @param key key whose mapping is to be removed from the cache + * @return the previous value associated with key, or null if there was no mapping + */ + @SuppressWarnings("unchecked") + public V remove(Object key) { + V originalValue = this.get(key); + if (key != null) { + this.cache.invalidate((K) key); + } + + return originalValue; + } + + /** + * Removes all entries from this cache. + */ + public void clear() { + this.cache.invalidateAll(); + } + + /** + * Creates a new Caffeine cache instance with the specified configuration. + * + * @param maximumSize the maximum number of entries the cache may contain + * @param expireAfterWrite the time after which entries should be automatically removed + * @return a new configured Cache instance + */ + private Cache createCache(Integer maximumSize, Duration expireAfterWrite) { + Caffeine caffeine = Caffeine + .newBuilder() + .evictionListener((key, value, cause) -> this.delegate().remove(key, value)); + + if (maximumSize != null && maximumSize >= 0) { + caffeine.maximumSize((long) maximumSize); + } + + if (expireAfterWrite != null && !expireAfterWrite.isNegative()) { + caffeine.scheduler(Scheduler.systemScheduler()).expireAfterWrite(expireAfterWrite); + } + + if (this.loader != null) { + return caffeine.build((key) -> { + V value = this.loader.load(key); + if (value != null) { + this.delegate().put(key, value); + } + + return value; + }); + } + + return caffeine.build(); + } +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingHandler.java b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingHandler.java new file mode 100644 index 0000000..5a2f9c7 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingHandler.java @@ -0,0 +1,48 @@ +package io.confluent.pas.agent.common.services.kstream; + +import reactor.core.publisher.Mono; + +import java.util.function.Consumer; + +/** + * Interface defining handler for processing streaming data. + * Implementations of this interface handle the processing logic for Kafka Streams. + * + * @param The type of the message key + * @param The type of the input message value + * @param The type of the output message value + */ +public interface StreamingHandler { + + /** + * Gets the Class object representing the key type. + * + * @return The Class object for the key type + */ + Class getKeyClass(); + + /** + * Gets the Class object representing the input type. + * + * @return The Class object for the input type + */ + Class getInClass(); + + /** + * Gets the Class object representing the output type. + * + * @return The Class object for the output type + */ + Class getOutClass(); + + + /** + * Handles the processing of a single message in the stream. + * + * @param key The key of the message to be processed + * @param in The input value to be processed + * @param sink Consumer function to emit the processed output value + */ + void handle(Key key, In in, Consumer sink); + +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingImpl.java b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingImpl.java new file mode 100644 index 0000000..c829b36 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingImpl.java @@ -0,0 +1,132 @@ +package io.confluent.pas.agent.common.services.kstream; + +import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializer; +import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializer; +import io.confluent.pas.agent.common.services.KafkaConfiguration; +import io.confluent.pas.agent.common.services.KafkaPropertiesFactory; +import io.confluent.pas.agent.common.services.Streaming; +import lombok.Builder; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.common.serialization.Serdes; +import org.apache.kafka.streams.KafkaStreams; +import org.apache.kafka.streams.StreamsBuilder; +import org.apache.kafka.streams.Topology; +import org.apache.kafka.streams.kstream.Consumed; +import org.apache.kafka.streams.kstream.Produced; + +import java.util.Map; +import java.util.Properties; + + +/** + * Implementation of the Streaming interface for Kafka Streams processing. + * + * @param The type of the message key + * @param The type of the input message value + * @param The type of the output message value + */ +@Slf4j +public class StreamingImpl implements Streaming { + + private KafkaStreams kafkaStreams; + + /** + * Initializes the streaming process with the given configuration and topics. + * + * @param kafkaConfiguration Configuration for Kafka connection + * @param inTopicName Name of the input topic + * @param outTopicName Name of the output topic + * @param streamingHandler Handler for processing the stream + */ + @Override + public void init(KafkaConfiguration kafkaConfiguration, + String inTopicName, + String outTopicName, + StreamingHandler streamingHandler) { + this.kafkaStreams = setupAndStartKafkaStreams( + kafkaConfiguration, + inTopicName, + outTopicName, + streamingHandler); + } + + /** + * Starts the Kafka Streams processing if the streams instance is initialized. + */ + @Override + public void start() { + if (kafkaStreams != null) { + kafkaStreams.start(); + log.info("Kafka Streams started"); + } + } + + /** + * Closes the Kafka Streams instance and releases all resources. + */ + public void close() { + if (kafkaStreams != null) { + kafkaStreams.close(); + log.info("Kafka Streams closed"); + } + } + + /** + * Creates a Serde (Serializer/Deserializer) for a specific class type. + * + * @param kafkaConfiguration Configuration for Kafka connection + * @param valueClass Class type for which to create the Serde + * @param isKey Whether this Serde is for a key or value + * @return A configured WrapperSerde instance + */ + private Serdes.WrapperSerde createSerde(KafkaConfiguration kafkaConfiguration, + Class valueClass, + boolean isKey) { + final Map configuration = + KafkaPropertiesFactory.getSchemaRegistryConfig(kafkaConfiguration, valueClass, isKey); + + final Serdes.WrapperSerde serde = new Serdes.WrapperSerde<>( + new KafkaJsonSchemaSerializer<>(), + new KafkaJsonSchemaDeserializer<>()); + + serde.configure(configuration, isKey); + return serde; + } + + /** + * Sets up and initializes a KafkaStreams instance with the specified configuration and topology. + * + * @param kafkaConfiguration Configuration for Kafka connection + * @param inTopicName Name of the input topic + * @param outTopicName Name of the output topic + * @param streamingHandler Handler for processing the stream + * @return Configured KafkaStreams instance + */ + private KafkaStreams setupAndStartKafkaStreams(KafkaConfiguration kafkaConfiguration, + String inTopicName, + String outTopicName, + StreamingHandler streamingHandler) { + final Serdes.WrapperSerde keySerde = createSerde(kafkaConfiguration, + streamingHandler.getKeyClass(), + true); + final Serdes.WrapperSerde inSerde = createSerde( + kafkaConfiguration, + streamingHandler.getInClass(), + false); + final Serdes.WrapperSerde outSerde = createSerde(kafkaConfiguration, + streamingHandler.getOutClass(), + false); + + StreamsBuilder builder = new StreamsBuilder(); + + builder.stream(inTopicName, Consumed.with(keySerde, inSerde)) + .process(new StreamingSupplier<>(streamingHandler)) + .to(outTopicName, Produced.with(keySerde, outSerde)); + + final Properties configuration = KafkaPropertiesFactory.getKStreamsProperties(kafkaConfiguration); + final Topology topology = builder.build(); + kafkaStreams = new KafkaStreams(topology, configuration); + + return kafkaStreams; + } +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingProcessor.java b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingProcessor.java new file mode 100644 index 0000000..f9c6556 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingProcessor.java @@ -0,0 +1,56 @@ +package io.confluent.pas.agent.common.services.kstream; + +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorContext; +import org.apache.kafka.streams.processor.api.Record; + +/** + * A generic Kafka Streams processor that handles stream processing operations. + * + * @param The type of the record key + * @param The type of the input record value + * @param The type of the output record value + */ +@Slf4j +@RequiredArgsConstructor +public class StreamingProcessor implements Processor { + + /** + * Handler responsible for processing the streaming records + */ + private final StreamingHandler handler; + /** + * Processor context for forwarding processed records + */ + private ProcessorContext context; + + /** + * Initializes the processor with the given context. + * + * @param context The processor context used for forwarding records + */ + @Override + public void init(ProcessorContext context) { + this.context = context; + } + + /** + * Processes the input record by delegating to the handler and forwarding the result. + * + * @param record The input record to process + */ + @Override + public void process(Record record) { + try { + handler.handle(record.key(), record.value(), out -> { + log.debug("Forwarding record: {}", record); + context.forward(record.withValue(out)); + }); + } catch (Exception e) { + log.error("Error processing record: {}", record, e); + // TODO: Handle error appropriately + } + } +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingSupplier.java b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingSupplier.java new file mode 100644 index 0000000..e61d438 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/kstream/StreamingSupplier.java @@ -0,0 +1,34 @@ +package io.confluent.pas.agent.common.services.kstream; + +import lombok.AllArgsConstructor; +import org.apache.kafka.streams.processor.api.Processor; +import org.apache.kafka.streams.processor.api.ProcessorSupplier; + + +/** + * A supplier class for streaming processors that handles the creation of processor instances + * based on a given streaming handler. + * + * @param the type of the message key + * @param the type of the input value + * @param the type of the output value + */ +@AllArgsConstructor +public class StreamingSupplier implements ProcessorSupplier { + + /** + * The streaming handler that will be used by the created processors + * to process the input messages. + */ + private final StreamingHandler streamingHandler; + + /** + * Creates a new processor instance using the configured streaming handler. + * + * @return a new processor instance that will use the configured streaming handler + */ + @Override + public Processor get() { + return new StreamingProcessor<>(streamingHandler); + } +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java index 5b94f05..b1bc609 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java @@ -88,7 +88,7 @@ @Getter @Setter @NoArgsConstructor -@AllArgsConstructor +@AllArgsConstructor(access = AccessLevel.PROTECTED) @JsonIgnoreProperties(ignoreUnknown = true) @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "registrationType", defaultImpl = Registration.class) @JsonSubTypes({ diff --git a/examples/JavaAgent/README.md b/examples/JavaAgent/README.md index 0f8dce4..195c931 100644 --- a/examples/JavaAgent/README.md +++ b/examples/JavaAgent/README.md @@ -114,9 +114,9 @@ llm gemini ### Example Interaction in MCP Shell ```sh -shell:> mcp list tools "Confluent MCP Proxy" +shell:> mcp list tools "Confluent Agent Proxy" -------------------- -Tools for server: Confluent MCP Proxy +Tools for server: Confluent Agent Proxy .................... Name: JavaSentimentAgent Description: This agent analyzes sentiment in human queries. diff --git a/examples/client_info.md b/examples/client_info.md index 77ad24e..4439c23 100644 --- a/examples/client_info.md +++ b/examples/client_info.md @@ -212,10 +212,10 @@ llm gemini Below is an example of how the shell interacts with the MCP proxy and LLM: ``` -Connected to server: Confluent MCP Proxy (1.0.0) -shell:>mcp list tools "Confluent MCP Proxy" +Connected to server: Confluent Agent Proxy (1.0.0) +shell:>mcp list tools "Confluent Agent Proxy" -------------------- -Tools for server: Confluent MCP Proxy +Tools for server: Confluent Agent Proxy .................... Name: findLastName Description: Find the last name of a given user diff --git a/frameworks/agent-proxy-framework/pom.xml b/frameworks/agent-proxy-framework/pom.xml index d54a35f..fc2a78f 100644 --- a/frameworks/agent-proxy-framework/pom.xml +++ b/frameworks/agent-proxy-framework/pom.xml @@ -16,7 +16,6 @@ 4.38.0 1.0.0-M7 - 7.9.0-ce diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/RequestResponseHandler.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/RequestResponseHandler.java new file mode 100644 index 0000000..ff747c1 --- /dev/null +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/RequestResponseHandler.java @@ -0,0 +1,52 @@ +package io.confluent.pas.agent.proxy.frameworks.java; + +import io.confluent.pas.agent.common.services.kstream.StreamingHandler; +import io.confluent.pas.agent.common.utils.JsonUtils; +import io.confluent.pas.agent.proxy.frameworks.java.models.Key; +import io.confluent.pas.agent.proxy.frameworks.java.models.Request; +import io.confluent.pas.agent.proxy.frameworks.java.models.Response; +import io.confluent.pas.agent.proxy.frameworks.java.models.ResponseStatus; +import io.confluent.pas.agent.proxy.frameworks.java.subscription.SubscriptionRequest; +import io.confluent.pas.agent.proxy.frameworks.java.subscription.SubscriptionResponse; + +import java.util.function.Consumer; + +public record RequestResponseHandler(Class requestClass, Class responseClass, + SubscriptionHandler.RequestHandler subscriptionHandler) + implements StreamingHandler { + + @Override + public Class getKeyClass() { + return Key.class; + } + + @Override + public Class getInClass() { + return Request.class; + } + + @Override + public Class getOutClass() { + return Response.class; + } + + @Override + public void handle(Key key, Request request, Consumer sink) { + final REQ req = JsonUtils.toObject(request.getPayload(), requestClass); + + final SubscriptionRequest subscriptionRequest = new SubscriptionRequest<>( + key, + req, + subscriptionResponse -> sendResponse(subscriptionResponse, sink)); + + subscriptionHandler.onRequest(subscriptionRequest); + } + + private void sendResponse(SubscriptionResponse subscriptionResponse, Consumer sink) { + final Response response = new Response(); + response.setStatus(ResponseStatus.COMPLETED); + response.setPayload(JsonUtils.toMap(subscriptionResponse.response())); + + sink.accept(response); + } +} diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java index add71b4..f5105b4 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandler.java @@ -1,11 +1,8 @@ package io.confluent.pas.agent.proxy.frameworks.java; import io.confluent.kafka.schemaregistry.json.JsonSchema; -import io.confluent.kafka.serializers.json.KafkaJsonSchemaDeserializer; -import io.confluent.kafka.serializers.json.KafkaJsonSchemaSerializer; -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.*; +import io.confluent.pas.agent.common.services.kstream.StreamingImpl; import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.common.services.schemas.RegistrationKey; import io.confluent.pas.agent.proxy.frameworks.java.kafka.TopicManagement; @@ -15,16 +12,8 @@ import io.confluent.pas.agent.proxy.frameworks.java.models.Response; import io.confluent.pas.agent.proxy.frameworks.java.subscription.SubscriptionRequest; import lombok.extern.slf4j.Slf4j; -import org.apache.kafka.common.serialization.Serdes; -import org.apache.kafka.streams.KafkaStreams; -import org.apache.kafka.streams.StreamsBuilder; -import org.apache.kafka.streams.Topology; -import org.apache.kafka.streams.kstream.Consumed; -import org.apache.kafka.streams.kstream.Produced; import java.io.Closeable; -import java.util.Map; -import java.util.Properties; import java.util.function.Supplier; /** @@ -41,13 +30,6 @@ @Slf4j public class SubscriptionHandler implements Closeable { - /** - * Supplier for creating Kafka Streams instances. - */ - public interface KStreamsSupplier { - KafkaStreams get(Topology topology, Properties kStreamsProperties); - } - /** * Interface for handling incoming requests from Kafka topics. * @@ -62,12 +44,8 @@ public interface RequestHandler { private final RegistrationService registrationService; private final Class requestClass; private final Class responseClass; - private final Serdes.WrapperSerde keySerde; - private final Serdes.WrapperSerde requestSerde; - private final Serdes.WrapperSerde responseSerde; private final Supplier topicManagementSupplier; - private final KStreamsSupplier kafkaStreamsSupplier; - private KafkaStreams kafkaStreams; + private final Streaming streaming; /** * Creates a new subscription handler with the specified message types. @@ -87,33 +65,32 @@ public SubscriptionHandler(KafkaConfiguration kafkaConfiguration, RegistrationKey.class, Registration.class), () -> new TopicManagementImpl(kafkaConfiguration), - KafkaStreams::new + new StreamingImpl<>() ); } /** - * Creates a new subscription handler with the specified message types. + * Creates a new subscription handler with the specified message types and registration service. * - * @param kafkaConfiguration Kafka cluster configuration - * @param requestClass Class type for request payloads - * @param responseClass Class type for response payloads - * @param registrationService Registration service for storing capabilities + * @param kafkaConfiguration Kafka cluster configuration + * @param requestClass Class type for request payloads + * @param responseClass Class type for response payloads + * @param registrationService Registration service for managing registrations + * @param topicManagementSupplier Supplier for topic management + * @param streaming Streaming service for processing requests */ public SubscriptionHandler(KafkaConfiguration kafkaConfiguration, Class requestClass, Class responseClass, RegistrationService registrationService, Supplier topicManagementSupplier, - KStreamsSupplier kafkaStreamsSupplier) { + Streaming streaming) { this.kafkaConfiguration = kafkaConfiguration; this.requestClass = requestClass; this.responseClass = responseClass; this.registrationService = registrationService; - this.keySerde = createSerde(Key.class, true); - this.requestSerde = createSerde(Request.class, false); - this.responseSerde = createSerde(Response.class, false); this.topicManagementSupplier = topicManagementSupplier; - this.kafkaStreamsSupplier = kafkaStreamsSupplier; + this.streaming = streaming; } /** @@ -165,9 +142,9 @@ public void subscribeWith(Registration registration, public void close() { log.info("Closing subscription handler resources"); - if (kafkaStreams != null) { + if (streaming != null) { try { - kafkaStreams.close(); + streaming.close(); log.debug("Kafka Streams closed successfully"); } catch (Exception e) { log.warn("Error closing Kafka Streams", e); @@ -182,32 +159,12 @@ public void close() { } } - /** - * Creates a Kafka Serde for serialization/deserialization. - * - * @param valueClass Class to create serde for - * @param isKey Whether this is for a key (true) or value (false) - * @return Configured Serde instance - */ - private Serdes.WrapperSerde createSerde(Class valueClass, boolean isKey) { - final Map configuration = - KafkaPropertiesFactory.getSchemaRegistryConfig(kafkaConfiguration, valueClass, isKey); - - final Serdes.WrapperSerde serde = new Serdes.WrapperSerde<>( - new KafkaJsonSchemaSerializer<>(), - new KafkaJsonSchemaDeserializer<>()); - - serde.configure(configuration, isKey); - return serde; - } - /** * Creates topics using class types. */ private void createTopics(Registration registration, Class requestClass, Class responseClass) throws Exception { - final JsonSchema reqSchema = Request.getSchema(requestClass); final JsonSchema resSchema = Response.getSchema(responseClass); @@ -240,7 +197,11 @@ private void createTopicsWithSchemas(Registration registration, private void startSubscription(Registration registration, RequestHandler handler) { registerCapability(registration); - setupAndStartKafkaStreams(registration, handler); + streaming.init(kafkaConfiguration, + registration.getRequestTopicName(), + registration.getResponseTopicName(), + new RequestResponseHandler<>(requestClass, responseClass, handler)); + streaming.start(); } /** @@ -257,22 +218,4 @@ private void registerCapability(Registration registration) { } } - /** - * Sets up and starts the Kafka Streams topology. - */ - private void setupAndStartKafkaStreams(Registration registration, - RequestHandler handler) { - StreamsBuilder builder = new StreamsBuilder(); - - builder.stream(registration.getRequestTopicName(), Consumed.with(keySerde, requestSerde)) - .process(new SubscriptionHandlerSupplier<>(handler, requestClass)) - .to(registration.getResponseTopicName(), Produced.with(keySerde, responseSerde)); - - final Topology topology = builder.build(); - kafkaStreams = kafkaStreamsSupplier.get(topology, - KafkaPropertiesFactory.getKStreamsProperties(kafkaConfiguration)); - - log.info("Starting Kafka Streams for registration: {}", registration.getName()); - kafkaStreams.start(); - } } \ No newline at end of file diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerProcessor.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerProcessor.java deleted file mode 100644 index e98d42f..0000000 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerProcessor.java +++ /dev/null @@ -1,54 +0,0 @@ -package io.confluent.pas.agent.proxy.frameworks.java; - -import io.confluent.pas.agent.common.utils.JsonUtils; -import io.confluent.pas.agent.proxy.frameworks.java.models.Key; -import io.confluent.pas.agent.proxy.frameworks.java.models.Request; -import io.confluent.pas.agent.proxy.frameworks.java.models.Response; -import io.confluent.pas.agent.proxy.frameworks.java.models.ResponseStatus; -import io.confluent.pas.agent.proxy.frameworks.java.subscription.SubscriptionRequest; -import io.confluent.pas.agent.proxy.frameworks.java.subscription.SubscriptionResponse; -import org.apache.kafka.streams.processor.api.Processor; -import org.apache.kafka.streams.processor.api.ProcessorContext; -import org.apache.kafka.streams.processor.api.Record; - -public class SubscriptionHandlerProcessor implements Processor { - - private final SubscriptionHandler.RequestHandler subscriptionHandler; - private final Class requestClass; - - private ProcessorContext context; - - public SubscriptionHandlerProcessor(SubscriptionHandler.RequestHandler subscriptionHandler, - Class requestClass) { - this.subscriptionHandler = subscriptionHandler; - this.requestClass = requestClass; - } - - @Override - public void init(ProcessorContext context) { - this.context = context; - } - - @Override - public void process(Record record) { - final REQ request = JsonUtils.toObject(record.value().getPayload(), requestClass); - - final SubscriptionRequest subscriptionRequest = new SubscriptionRequest<>( - record.key(), - request, - this::sendResponse); - - subscriptionHandler.onRequest(subscriptionRequest); - } - - void sendResponse(SubscriptionResponse subscriptionResponse) { - final Response response = new Response(); - response.setStatus(ResponseStatus.COMPLETED); - response.setPayload(JsonUtils.toMap(subscriptionResponse.response())); - - context.forward(new Record<>( - subscriptionResponse.key(), - response, - System.currentTimeMillis())); - } -} diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerSupplier.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerSupplier.java deleted file mode 100644 index 366bbde..0000000 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerSupplier.java +++ /dev/null @@ -1,20 +0,0 @@ -package io.confluent.pas.agent.proxy.frameworks.java; - -import io.confluent.pas.agent.proxy.frameworks.java.models.Key; -import io.confluent.pas.agent.proxy.frameworks.java.models.Request; -import io.confluent.pas.agent.proxy.frameworks.java.models.Response; -import lombok.AllArgsConstructor; -import org.apache.kafka.streams.processor.api.Processor; -import org.apache.kafka.streams.processor.api.ProcessorSupplier; - -@AllArgsConstructor -public class SubscriptionHandlerSupplier implements ProcessorSupplier { - - private final SubscriptionHandler.RequestHandler subscriptionHandler; - private final Class requestClass; - - @Override - public Processor get() { - return new SubscriptionHandlerProcessor<>(subscriptionHandler, requestClass); - } -} diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/subscription/SubscriptionRequest.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/subscription/SubscriptionRequest.java index 13cf159..40a3456 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/subscription/SubscriptionRequest.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/subscription/SubscriptionRequest.java @@ -1,7 +1,6 @@ package io.confluent.pas.agent.proxy.frameworks.java.subscription; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; -import io.confluent.pas.agent.proxy.frameworks.java.models.Response; import lombok.AllArgsConstructor; import lombok.Getter; import lombok.extern.slf4j.Slf4j; diff --git a/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerTest.java b/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerTest.java index 557901f..c69a588 100644 --- a/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerTest.java +++ b/frameworks/agent-proxy-framework/src/test/java/io/confluent/pas/agent/proxy/frameworks/java/SubscriptionHandlerTest.java @@ -3,12 +3,14 @@ import io.confluent.kafka.schemaregistry.json.JsonSchema; import io.confluent.pas.agent.common.services.KafkaConfiguration; import io.confluent.pas.agent.common.services.RegistrationService; +import io.confluent.pas.agent.common.services.Streaming; import io.confluent.pas.agent.common.services.schemas.Registration; import io.confluent.pas.agent.common.services.schemas.RegistrationKey; import io.confluent.pas.agent.proxy.frameworks.java.kafka.TopicManagement; import io.confluent.pas.agent.proxy.frameworks.java.models.Key; +import io.confluent.pas.agent.proxy.frameworks.java.models.Request; +import io.confluent.pas.agent.proxy.frameworks.java.models.Response; import org.apache.commons.lang3.RandomStringUtils; -import org.apache.kafka.streams.KafkaStreams; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.Mock; @@ -18,10 +20,10 @@ public class SubscriptionHandlerTest { - private record Request(int a, int b) { + private record HandlingRequest(int a, int b) { } - private record Response(int result) { + private record HandlingResponse(int result) { } @@ -35,9 +37,9 @@ private record Response(int result) { private TopicManagement topicManagement; @Mock - private KafkaStreams kStreams; + private Streaming kStreams; - private SubscriptionHandler subscriptionHandler; + private SubscriptionHandler subscriptionHandler; @BeforeEach public void setUp() { @@ -55,11 +57,11 @@ public void setUp() { subscriptionHandler = new SubscriptionHandler<>( kafkaConfiguration, - Request.class, - Response.class, + HandlingRequest.class, + HandlingResponse.class, registrationService, () -> topicManagement, - (topology, properties) -> kStreams); + kStreams); } @Test @@ -73,7 +75,7 @@ public void testSubscribeWith() throws Exception { "responseTopic" ), (request) -> { - request.respond(new Response(request.getRequest().a() + request.getRequest().b())) + request.respond(new HandlingResponse(request.getRequest().a() + request.getRequest().b())) .block(); } ); @@ -100,7 +102,7 @@ public void testSubscribeWithSchema() throws Exception { new JsonSchema(reqSchema), new JsonSchema(resSchema), (request) -> { - request.respond(new Response(request.getRequest().a() + request.getRequest().b())) + request.respond(new HandlingResponse(request.getRequest().a() + request.getRequest().b())) .block(); } ); diff --git a/pom.xml b/pom.xml index cb175ce..6e9ac54 100644 --- a/pom.xml +++ b/pom.xml @@ -48,6 +48,8 @@ 3.7.5 1.21.0 0.9.0 + 7.9.0-ce + 3.2.0 @@ -69,6 +71,18 @@ + + com.github.ben-manes.caffeine + caffeine + ${caffeine.version} + + + + org.apache.kafka + kafka-streams + ${kstream.version} + + io.modelcontextprotocol.sdk mcp @@ -112,6 +126,11 @@ kcache ${kcache.version} + + io.kcache + kcache-caffeine + ${kcache.version} + org.projectlombok lombok diff --git a/shell/src/main/resources/application.yaml b/shell/src/main/resources/application.yaml index ed6ef7a..a9c4668 100644 --- a/shell/src/main/resources/application.yaml +++ b/shell/src/main/resources/application.yaml @@ -1,6 +1,6 @@ spring: application: - name: "Confluent MCP Proxy Client" + name: "Confluent Agent Proxy Client" shell: interactive: enabled: true From 00286d5d57b335fc481b0c14aa794b020a7e2097 Mon Sep 17 00:00:00 2001 From: Pascal Vantrepote Date: Tue, 16 Sep 2025 09:38:35 -0400 Subject: [PATCH 3/3] TMP push --- ...cation.java => AgentProxyApplication.java} | 4 +- .../registration/RegistrationCoordinator.java | 40 +++-- .../registration/RequestResponseHandler.java | 50 +++++- .../handlers/mcp/McpResourceHandler.java | 8 +- .../services/schemas/AbstractSchema.java | 49 ++++++ .../schemas/BlobResourceResponse.java | 69 ++------ .../common/services/schemas/Registration.java | 93 +++++----- .../schemas/ResourceRegistration.java | 166 +++++++++--------- .../services/schemas/ResourceResponse.java | 149 +++++++--------- .../schemas/TextResourceResponse.java | 71 ++------ .../McpRegistrationAutoConfiguration.java | 1 - 11 files changed, 353 insertions(+), 347 deletions(-) rename agent-proxy/src/main/java/io/confluent/pas/agent/proxy/{ConfluentMcpProxyApplication.java => AgentProxyApplication.java} (78%) create mode 100644 common/src/main/java/io/confluent/pas/agent/common/services/schemas/AbstractSchema.java 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 cd6e53b..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 @@ -16,37 +16,47 @@ 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 un-registrations + * - 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 synchronous HTTP communications with agents + */ @Getter private final AgentAsyncServer restServer; @@ -57,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<>(); @@ -234,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(); } /** @@ -253,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( @@ -263,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/RequestResponseHandler.java b/agent-proxy/src/main/java/io/confluent/pas/agent/proxy/registration/RequestResponseHandler.java index b28a083..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 @@ -16,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,6 +29,12 @@ public class RequestResponseHandler implements DisposableBean { private final ProducerService producerService; private final ConsumerService consumerService; + /** + * 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, @Value("${kafka.response.timeout:10000}") long responseTimeout) { @@ -32,16 +42,35 @@ public RequestResponseHandler(KafkaConfiguration kafkaConfiguration, 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) { this.producerService = producerService; this.consumerService = consumerService; } + /** + * 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, @@ -53,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/common/src/main/java/io/confluent/pas/agent/common/services/schemas/AbstractSchema.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/AbstractSchema.java new file mode 100644 index 0000000..9e179b1 --- /dev/null +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/AbstractSchema.java @@ -0,0 +1,49 @@ +package io.confluent.pas.agent.common.services.schemas; + +import com.fasterxml.jackson.annotation.JsonProperty; +import lombok.AllArgsConstructor; +import lombok.NoArgsConstructor; + +import java.util.HashMap; +import java.util.Map; + +@AllArgsConstructor(access = lombok.AccessLevel.PROTECTED) +@NoArgsConstructor(access = lombok.AccessLevel.PROTECTED) +public abstract class AbstractSchema { + private Map metaData; + + @JsonProperty(value = "metaData") + public Map getMetaData() { + return metaData; + } + + @JsonProperty(value = "metaData") + public void setMetaData(Map metaData) { + this.metaData = metaData; + } + + protected T getMetaDataValue(String key, Class type) { + if (metaData == null || !metaData.containsKey(key)) { + return null; + } + + Object value = metaData.get(key); + if (value == null) { + return null; + } + + if (type.isInstance(value)) { + return type.cast(value); + } + + throw new IllegalArgumentException("Meta data value for key '" + key + "' is not of type " + type.getName()); + } + + protected void setMetaDataValue(String key, Object value) { + if (metaData == null) { + metaData = new HashMap<>(); + } + + metaData.put(key, value); + } +} diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java index eb890ed..57b9386 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/BlobResourceResponse.java @@ -1,73 +1,28 @@ package io.confluent.pas.agent.common.services.schemas; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; -import com.fasterxml.jackson.annotation.JsonProperty; -import io.confluent.kafka.schemaregistry.annotations.Schema; import lombok.Getter; import lombok.NoArgsConstructor; import lombok.Setter; -@Schema(value = """ - { - "properties":{ - "type":{ - "connect.index":0, - "enum":[ - "text", - "blob" - ], - "default": "text", - "type":"string" - }, - "uri":{ - "connect.index":1, - "type":"string" - }, - "mimeType":{ - "connect.index":2, - "type":"string" - }, - "text":{ - "connect.index":3, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "blob":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "type", - "uri", - "mimeType" - ], - "title":"Record", - "type":"object" - }""", refs = {}) @Getter @Setter @NoArgsConstructor @JsonIgnoreProperties(ignoreUnknown = true) public class BlobResourceResponse extends ResourceResponse { - @JsonProperty(value = "blob", required = true) - private String blob; + private final static String BLOB = "blob"; + + public String getBlob() { + return getMetaDataValue(BLOB, String.class); + } + + public void setBlob(String blob) { + setMetaDataValue(BLOB, blob); + } public BlobResourceResponse(String uri, String mimeType, String blob) { - super(ResponseType.BLOB, uri, mimeType); - this.blob = blob; + super(ResourceResponse.BLOB_TYPE, uri, mimeType); + + setBlob(blob); } } \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java index b1bc609..c4f80da 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/Registration.java @@ -5,6 +5,9 @@ import lombok.*; import org.apache.commons.lang3.StringUtils; +import java.util.HashMap; +import java.util.Map; + @Schema(value = """ { "properties":{ @@ -12,17 +15,6 @@ "connect.index":5, "type":"string" }, - "correlationIdFieldName":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, "description":{ "connect.index":1, "type":"string" @@ -33,25 +25,19 @@ }, "requestTopicName":{ "connect.index":2, - "type":"string" + "type":"string" }, "responseTopicName":{ "connect.index":3, "type":"string" }, - "mimeType":{ - "connect.index":6, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] + "metaData":{ + "connect.index":4, + "type":"object", + "additionalProperties":true }, - "url":{ - "connect.index":7, + "version":{ + "connect.index":6, "oneOf":[ { "type":"null" @@ -59,20 +45,9 @@ { "type":"string" } - ] - }, - "version":{ - "connect.index":8, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ], - "default":"N/A" - } + ], + "default":"N/A" + } }, "required":[ "name", @@ -88,17 +63,16 @@ @Getter @Setter @NoArgsConstructor -@AllArgsConstructor(access = AccessLevel.PROTECTED) @JsonIgnoreProperties(ignoreUnknown = true) @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "registrationType", defaultImpl = Registration.class) @JsonSubTypes({ @JsonSubTypes.Type(value = ResourceRegistration.class, name = Registration.RESOURCE), @JsonSubTypes.Type(value = Registration.class, name = Registration.TOOL)}) @Builder -public class Registration { +public class Registration extends AbstractSchema { public final static String TOOL = "tool"; public final static String RESOURCE = "resource"; - public final static String CORRELATION_ID_FIELD_NAME = "correlationId"; + public final static String AGENT = "agent"; @JsonProperty(value = "registrationType", required = true, defaultValue = TOOL) private String registrationType; @@ -113,12 +87,45 @@ public class Registration { @JsonProperty(value = "version", defaultValue = "N/A") private String version; - public Registration(String name, String description, String requestTopicName, String responseTopicName) { - this(TOOL, name, description, requestTopicName, responseTopicName, "N/A"); + public Registration(String name, + String description, + String requestTopicName, + String responseTopicName) { + this(TOOL, name, description, requestTopicName, responseTopicName, "N/A", new HashMap<>()); + } + + public Registration(String registrationType, String name, String description, String requestTopicName, String responseTopicName, String version) { + super(new HashMap<>()); + + this.registrationType = registrationType; + this.name = name; + this.description = description; + this.requestTopicName = requestTopicName; + this.responseTopicName = responseTopicName; + this.version = version; + } + + protected Registration(String registrationType, + String name, + String description, + String requestTopicName, + String responseTopicName, + String version, + Map metaData) { + super(metaData); + + this.registrationType = registrationType; + this.name = name; + this.description = description; + this.requestTopicName = requestTopicName; + this.responseTopicName = responseTopicName; + this.version = version; } @JsonIgnore public boolean isResource() { return StringUtils.equals(registrationType, RESOURCE); } + + } \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java index 08dd065..5807001 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceRegistration.java @@ -2,113 +2,115 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; -import io.confluent.kafka.schemaregistry.annotations.Schema; import io.confluent.pas.agent.common.utils.UriUtils; -import lombok.AllArgsConstructor; -import lombok.Getter; import lombok.NoArgsConstructor; -import lombok.Setter; import org.apache.commons.lang3.StringUtils; -@Setter -@Getter -@Schema(value = """ - { - "properties":{ - "registrationType":{ - "connect.index":5, - "type":"string" - }, - "correlationIdFieldName":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "description":{ - "connect.index":1, - "type":"string" - }, - "name":{ - "connect.index":0, - "type":"string" - }, - "requestTopicName":{ - "connect.index":2, - "type":"string" - }, - "responseTopicName":{ - "connect.index":3, - "type":"string" - }, - "mimeType":{ - "connect.index":6, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "url":{ - "connect.index":7, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "name", - "description", - "registrationType", - "requestTopicName", - "responseTopicName" - ], - "additionalProperties":false, - "title":"Record", - "type":"object" - }""", refs = {}) -@AllArgsConstructor +import java.util.HashMap; + +/** + * Represents a resource registration in the system, extending the base Registration class. + * This class manages resource-specific attributes such as URLs and MIME types. + *

    + * Each resource registration contains: + * - A URL identifying the resource location + * - A MIME type specifying the resource format + * - Metadata storage for additional resource properties + *

    + * The class provides methods to manage these attributes while ensuring data validity + * through input validation. It supports template URLs and follows RESTful principles + * for resource identification. + */ @NoArgsConstructor @JsonIgnoreProperties(ignoreUnknown = true) public class ResourceRegistration extends Registration { - private String mimeType; - private String url; + /** + * Key used to store and retrieve the MIME type in the metadata map. + * This constant ensures consistent access to MIME type information. + */ + private final static String MIME_TYPE = "mimeType"; + + /** + * Key used to store and retrieve the URL in the metadata map. + * This constant ensures consistent access to URL information. + */ + private final static String URL = "url"; + + /** + * Retrieves the URL associated with this resource. + * + * @return The URL string stored in metadata + */ + public String getUrl() { + return getMetaDataValue(URL, String.class); + } - // TODO: Should not require to remove the leading slash + /** + * Sets the URL for this resource after validation. + * + * @param url The URL to set + * @throws IllegalArgumentException if the URL is blank + */ public void setUrl(String url) { if (StringUtils.isBlank(url)) { throw new IllegalArgumentException("url cannot be blank"); } - this.url = url.startsWith("/") ? url.substring(1) : url; + setMetaDataValue(URL, url.startsWith("/") ? url.substring(1) : url); + } + + /** + * Retrieves the MIME type of this resource. + * + * @return The MIME type string stored in metadata + */ + public String getMimeType() { + return getMetaDataValue(MIME_TYPE, String.class); } + /** + * Sets the MIME type for this resource after validation. + * + * @param mimeType The MIME type to set + * @throws IllegalArgumentException if the MIME type is blank + */ + public void setMimeType(String mimeType) { + if (StringUtils.isBlank(mimeType)) { + throw new IllegalArgumentException("mimeType cannot be blank"); + } + + setMetaDataValue(MIME_TYPE, mimeType); + } + + /** + * Checks if the resource URL is a template. + * + * @return true if the URL contains template parameters, false otherwise + */ @JsonIgnore public boolean isTemplate() { - return UriUtils.isTemplate(url); + return UriUtils.isTemplate(getUrl()); } + /** + * Creates a new ResourceRegistration with the specified parameters. + * + * @param name The name of the resource + * @param description The description of the resource + * @param requestTopicName The name of the request topic + * @param responseTopicName The name of the response topic + * @param mimeType The MIME type of the resource + * @param url The URL of the resource + */ public ResourceRegistration(String name, String description, String requestTopicName, String responseTopicName, String mimeType, String url) { - super(RESOURCE, name, description, requestTopicName, responseTopicName, "N/A"); - this.mimeType = mimeType; + super(RESOURCE, name, description, requestTopicName, responseTopicName, "N/A", new HashMap<>()); + + setMimeType(mimeType); setUrl(url); } } diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java index fb65191..896e4cc 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/ResourceResponse.java @@ -7,103 +7,80 @@ import lombok.NoArgsConstructor; import lombok.Setter; +import java.util.HashMap; +import java.util.Map; + @Schema(value = """ - { - "properties":{ - "type":{ - "connect.index":0, - "enum":[ - "text", - "blob" - ], - "default": "text", - "type":"string" - }, - "uri":{ - "connect.index":1, - "type":"string" - }, - "mimeType":{ - "connect.index":2, - "type":"string" - }, - "text":{ - "connect.index":3, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "blob":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "type", - "uri", - "mimeType" - ], - "title":"Record", - "type":"object" - }""", refs = {}) + { + "properties":{ + "type":{ + "connect.index":0, + "enum":[ + "text", + "blob" + ], + "default": "text", + "type":"string" + }, + "uri":{ + "connect.index":1, + "type":"string" + }, + "mimeType":{ + "connect.index":2, + "type":"string" + }, + "metaData":{ + "connect.index":3, + "type":"object", + "additionalProperties":true + } + }, + "required":[ + "type", + "uri", + "mimeType" + ], + "title":"Record", + "type":"object" + }""", refs = {}) @Getter @Setter -@AllArgsConstructor @NoArgsConstructor @JsonIgnoreProperties(ignoreUnknown = true) @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type", defaultImpl = TextResourceResponse.class) @JsonSubTypes({ - @JsonSubTypes.Type(value = TextResourceResponse.class, name = "text"), - @JsonSubTypes.Type(value = BlobResourceResponse.class, name = "blob")}) -public class ResourceResponse { - public enum ResponseType { - TEXT("text"), - BLOB("blob"); - - private final String value; - - ResponseType(String value) { - this.value = value; - } - - @JsonValue - public String getValue() { - return value; - } - - @Override - public String toString() { - return value; - } - - @JsonCreator - public static ResourceResponse.ResponseType fromValue(String value) { - for (ResourceResponse.ResponseType e : ResourceResponse.ResponseType.values()) { - if (e.value.equals(value)) { - return e; - } - } - - throw new IllegalArgumentException("Unexpected value '" + value + "'"); - } - } + @JsonSubTypes.Type(value = TextResourceResponse.class, name = ResourceResponse.TEXT_TYPE), + @JsonSubTypes.Type(value = BlobResourceResponse.class, name = ResourceResponse.BLOB_TYPE)}) +public class ResourceResponse extends AbstractSchema { + public static final String TEXT_TYPE = "text"; + public static final String BLOB_TYPE = "blob"; @JsonProperty(value = "type", required = true) - private ResponseType type; + private String type; @JsonProperty(value = "uri", required = true) private String uri; @JsonProperty(value = "mimeType", required = true) private String mimeType; + + protected ResourceResponse(String type, + String uri, + String mimeType, + Map metaData) { + super(metaData); + + this.type = type; + this.uri = uri; + this.mimeType = mimeType; + } + + protected ResourceResponse(String type, + String uri, + String mimeType) { + super(new HashMap<>()); + + this.type = type; + this.uri = uri; + this.mimeType = mimeType; + } } \ No newline at end of file diff --git a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java index dca54a7..1f3be33 100644 --- a/common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java +++ b/common/src/main/java/io/confluent/pas/agent/common/services/schemas/TextResourceResponse.java @@ -1,73 +1,30 @@ package io.confluent.pas.agent.common.services.schemas; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; -import com.fasterxml.jackson.annotation.JsonProperty; -import io.confluent.kafka.schemaregistry.annotations.Schema; import lombok.Getter; import lombok.NoArgsConstructor; import lombok.Setter; -@Schema(value = """ - { - "properties":{ - "type":{ - "connect.index":0, - "enum":[ - "text", - "blob" - ], - "default": "text", - "type":"string" - }, - "uri":{ - "connect.index":1, - "type":"string" - }, - "mimeType":{ - "connect.index":2, - "type":"string" - }, - "text":{ - "connect.index":3, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - }, - "blob":{ - "connect.index":4, - "oneOf":[ - { - "type":"null" - }, - { - "type":"string" - } - ] - } - }, - "required":[ - "type", - "uri", - "mimeType" - ], - "title":"Record", - "type":"object" - }""", refs = {}) +import java.util.HashMap; + @Getter @Setter @NoArgsConstructor @JsonIgnoreProperties(ignoreUnknown = true) public class TextResourceResponse extends ResourceResponse { - @JsonProperty(value = "text", required = true) - private String text; + private final static String TEXT = "text"; + + public String getText() { + return getMetaDataValue(TEXT, String.class); + } + + public void setText(String text) { + setMetaDataValue(TEXT, text); + } public TextResourceResponse(String uri, String mimeType, String text) { - super(ResponseType.TEXT, uri, mimeType); - this.text = text; + super(ResourceResponse.TEXT_TYPE, uri, mimeType, new HashMap<>()); + + setText(text); } } \ No newline at end of file diff --git a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java index 38c44a6..4174451 100644 --- a/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java +++ b/frameworks/agent-proxy-framework/src/main/java/io/confluent/pas/agent/proxy/frameworks/java/spring/autoconfig/McpRegistrationAutoConfiguration.java @@ -9,7 +9,6 @@ import org.springframework.context.annotation.Bean; import io.confluent.pas.agent.common.services.schemas.Registration; -import static io.confluent.pas.agent.common.services.schemas.Registration.CORRELATION_ID_FIELD_NAME; /** * Auto-configuration class for MCP (Model Control Protocol) agent registration.