diff --git a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResource.java b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResource.java index 9b14d63d7f..cb66862762 100644 --- a/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResource.java +++ b/src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResource.java @@ -131,7 +131,7 @@ public Uni> buildKeyspaceContext(String keyspace, KeyspaceComm contextBuilderSupplier .getBuilder(keyspaceSchemaObject) .withCommandName(command.getClass().getSimpleName()) - .withRequestContext(createRequestContext()) + .withRequestContext(requestContext) .build()); } @@ -211,20 +211,24 @@ public Uni processCollectionCommand( } /** - * Process a general or keyspace command and return the CommandResult directly. - * - *

The Quarkus MCP Server framework automatically encodes the {@link CommandResult} to JSON via - * its built-in encoder (Jackson), so no manual serialization is needed. + * Process a command and map its result to an MCP tool response. * - * @param contextUni The command context (from buildGeneralContext or buildKeyspaceContext) + * @param contextUni The command context * @param command The command to execute - * @return Uni of CommandResult + * @return Uni of ToolResponse */ public Uni processCommand(Uni> contextUni, Command command) { return contextUni - .flatMap( - context -> { + .flatMap(context -> executeCommand(context, command)) + .map(CommandResult::toToolResponse); + } + + private Uni executeCommand(CommandContext context, Command command) { + // Defer the feature check and processor call so synchronous failures still close the context. + return Uni.createFrom() + .deferred( + () -> { if (!context.apiFeatures().isFeatureEnabled(ApiFeature.MCP)) { return Uni.createFrom() .item( @@ -234,7 +238,18 @@ public Uni processCommand(Uni> contextUni, Comma } return meteredCommandProcessor.processCommand(context, command); }) - .map(CommandResult::toToolResponse); + .onTermination() + .invoke( + () -> { + try { + context.close(); + } catch (Exception e) { + LOGGER.error( + "Error closing the command context for requestContext={}", + context.requestContext(), + e); + } + }); } /** diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResourceTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResourceTest.java new file mode 100644 index 0000000000..f4d603f45e --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResourceTest.java @@ -0,0 +1,299 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import com.fasterxml.jackson.databind.ObjectMapper; +import io.micrometer.core.instrument.MeterRegistry; +import io.quarkus.security.identity.SecurityIdentity; +import io.smallrye.mutiny.Uni; +import io.stargate.sgv2.jsonapi.TestConstants; +import io.stargate.sgv2.jsonapi.api.model.command.CommandContext; +import io.stargate.sgv2.jsonapi.api.model.command.CommandResult; +import io.stargate.sgv2.jsonapi.api.model.command.KeyspaceCommand; +import io.stargate.sgv2.jsonapi.api.model.command.tracing.RequestTracing; +import io.stargate.sgv2.jsonapi.api.request.EmbeddingCredentials; +import io.stargate.sgv2.jsonapi.api.request.EmbeddingCredentialsResolver; +import io.stargate.sgv2.jsonapi.api.request.RequestContext; +import io.stargate.sgv2.jsonapi.api.request.tenant.RequestTenantResolver; +import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; +import io.stargate.sgv2.jsonapi.config.feature.ApiFeatures; +import io.stargate.sgv2.jsonapi.metrics.JsonProcessingMetricsReporter; +import io.stargate.sgv2.jsonapi.service.cqldriver.CQLSessionCache; +import io.stargate.sgv2.jsonapi.service.cqldriver.CqlSessionCacheSupplier; +import io.stargate.sgv2.jsonapi.service.embedding.operation.EmbeddingProviderFactory; +import io.stargate.sgv2.jsonapi.service.processor.MeteredCommandProcessor; +import io.stargate.sgv2.jsonapi.service.reranking.operation.RerankingProviderFactory; +import io.stargate.sgv2.jsonapi.service.schema.SchemaObjectCache; +import io.stargate.sgv2.jsonapi.service.schema.SchemaObjectCacheSupplier; +import io.stargate.sgv2.jsonapi.service.schema.collections.CollectionSchemaObject; +import io.vertx.core.MultiMap; +import io.vertx.core.http.HttpServerRequest; +import io.vertx.ext.web.RoutingContext; +import jakarta.enterprise.inject.Instance; +import jakarta.inject.Provider; +import java.security.Principal; +import java.time.Duration; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +/** + * Unit tests for the request-context lifetime in {@link McpResource}: one context per tool call, + * closed exactly once when the call ends. + */ +class McpResourceTest { + + private static final TestConstants TEST_CONSTANTS = new TestConstants(); + + private final MeteredCommandProcessor commandProcessor = mock(MeteredCommandProcessor.class); + private final SchemaObjectCache schemaObjectCache = mock(SchemaObjectCache.class); + + @Test + void closesContextOnSuccess() throws Exception { + var context = commandContext(true); + var command = mock(KeyspaceCommand.class); + var result = CommandResult.statusOnlyBuilder(RequestTracing.NO_OP).build(); + when(commandProcessor.processCommand(context, command)) + .thenReturn(Uni.createFrom().item(result)); + + var response = + mcpResource() + .processCommand(Uni.createFrom().item(context), command) + .await() + .atMost(Duration.ofSeconds(5)); + + assertThat(response.isError()).isFalse(); + verify(context, times(1)).close(); + } + + @Test + void closesContextOnFailure() throws Exception { + var context = commandContext(true); + var command = mock(KeyspaceCommand.class); + when(commandProcessor.processCommand(context, command)) + .thenReturn(Uni.createFrom().failure(new RuntimeException("processor failed"))); + + assertThatThrownBy( + () -> + mcpResource() + .processCommand(Uni.createFrom().item(context), command) + .await() + .atMost(Duration.ofSeconds(5))) + .hasMessage("processor failed"); + verify(context, times(1)).close(); + } + + @Test + void closesContextWhenProcessorThrowsSynchronously() throws Exception { + var context = commandContext(true); + var command = mock(KeyspaceCommand.class); + var failure = new IllegalStateException("processor failed before returning a Uni"); + when(commandProcessor.processCommand(context, command)).thenThrow(failure); + + assertThatThrownBy( + () -> + mcpResource() + .processCommand(Uni.createFrom().item(context), command) + .await() + .atMost(Duration.ofSeconds(5))) + .isSameAs(failure); + verify(context).close(); + } + + @Test + void closesContextWhenFeatureConfigurationIsInvalid() throws Exception { + var context = commandContext(true); + when(context.apiFeatures()) + .thenReturn(ApiFeatures.fromConfigAndRequest(Map.of(ApiFeature.MCP, "invalid"), null)); + + assertThatThrownBy( + () -> + mcpResource() + .processCommand(Uni.createFrom().item(context), mock(KeyspaceCommand.class)) + .await() + .atMost(Duration.ofSeconds(5))) + .isInstanceOf(IllegalArgumentException.class); + verifyNoInteractions(commandProcessor); + verify(context).close(); + } + + @Test + void closesContextWhenFeatureDisabled() throws Exception { + var context = commandContext(false); + var command = mock(KeyspaceCommand.class); + + var response = + mcpResource() + .processCommand(Uni.createFrom().item(context), command) + .await() + .atMost(Duration.ofSeconds(5)); + + assertThat(response.isError()).isTrue(); + verifyNoInteractions(commandProcessor); + verify(context, times(1)).close(); + } + + @Test + void closesContextOnCancellation() throws Exception { + var context = commandContext(true); + var command = mock(KeyspaceCommand.class); + when(commandProcessor.processCommand(context, command)) + .thenReturn(Uni.createFrom().emitter(emitter -> {})); + + var cancellable = + mcpResource() + .processCommand(Uni.createFrom().item(context), command) + .subscribe() + .with(response -> {}, failure -> {}); + cancellable.cancel(); + + verify(context, times(1)).close(); + } + + @Test + void closeFailureDoesNotReplaceResult() throws Exception { + var context = commandContext(true); + doThrow(new Exception("close failed")).when(context).close(); + var command = mock(KeyspaceCommand.class); + var result = CommandResult.statusOnlyBuilder(RequestTracing.NO_OP).build(); + when(commandProcessor.processCommand(context, command)) + .thenReturn(Uni.createFrom().item(result)); + + var response = + mcpResource() + .processCommand(Uni.createFrom().item(context), command) + .await() + .atMost(Duration.ofSeconds(5)); + + assertThat(response.isError()).isFalse(); + verify(context, times(1)).close(); + } + + @Test + void closeFailureDoesNotReplaceError() throws Exception { + var context = commandContext(true); + doThrow(new Exception("close failed")).when(context).close(); + var command = mock(KeyspaceCommand.class); + when(commandProcessor.processCommand(context, command)) + .thenReturn(Uni.createFrom().failure(new RuntimeException("processor failed"))); + + assertThatThrownBy( + () -> + mcpResource() + .processCommand(Uni.createFrom().item(context), command) + .await() + .atMost(Duration.ofSeconds(5))) + .hasMessage("processor failed"); + verify(context, times(1)).close(); + } + + @Test + void keyspaceContextUsesOneRequestContext() { + var requestContextCaptor = ArgumentCaptor.forClass(RequestContext.class); + when(schemaObjectCache.getKeyspace(requestContextCaptor.capture(), any(), any(), eq(true))) + .thenReturn(Uni.createFrom().item(TEST_CONSTANTS.KEYSPACE_SCHEMA_OBJECT)); + + var context = + mcpResourceWithRequest() + .buildKeyspaceContext(TEST_CONSTANTS.KEYSPACE_NAME, mock(KeyspaceCommand.class)) + .await() + .atMost(Duration.ofSeconds(5)); + + // The schema lookup and the built command context must see the same request, so the + // call has one request ID from start to finish. + assertThat(context.requestContext()).isSameAs(requestContextCaptor.getValue()); + } + + @SuppressWarnings("unchecked") + private CommandContext commandContext(boolean mcpEnabled) { + CommandContext context = mock(CommandContext.class); + var apiFeatures = + mcpEnabled + ? ApiFeatures.fromConfigAndRequest(Map.of(ApiFeature.MCP, "true"), null) + : ApiFeatures.empty(); + when(context.apiFeatures()).thenReturn(apiFeatures); + when(context.requestTracing()).thenReturn(RequestTracing.NO_OP); + return context; + } + + private McpResource mcpResource() { + return mcpResource( + () -> mock(RoutingContext.class), + () -> mock(SecurityIdentity.class), + instanceOf(mock(RequestTenantResolver.class)), + instanceOf(mock(EmbeddingCredentialsResolver.class))); + } + + private McpResource mcpResource( + Provider routingContextProvider, + Provider securityIdentityProvider, + Instance tenantResolver, + Instance embeddingCredentialsResolver) { + return new McpResource( + mock(ObjectMapper.class), + commandProcessor, + mock(MeterRegistry.class), + mock(JsonProcessingMetricsReporter.class), + sessionCacheSupplier(), + mock(EmbeddingProviderFactory.class), + mock(RerankingProviderFactory.class), + schemaCacheSupplier(), + routingContextProvider, + securityIdentityProvider, + tenantResolver, + embeddingCredentialsResolver); + } + + /** An {@link McpResource} whose request can build a real {@link RequestContext}. */ + private McpResource mcpResourceWithRequest() { + var httpRequest = mock(HttpServerRequest.class); + when(httpRequest.headers()).thenReturn(mock(MultiMap.class)); + var routingContext = mock(RoutingContext.class); + when(routingContext.request()).thenReturn(httpRequest); + + var principal = mock(Principal.class); + when(principal.getName()).thenReturn("test-token"); + var securityIdentity = mock(SecurityIdentity.class); + when(securityIdentity.getPrincipal()).thenReturn(principal); + + var tenantResolver = mock(RequestTenantResolver.class); + when(tenantResolver.resolve(any(), any())).thenReturn(TEST_CONSTANTS.TENANT); + var embeddingCredentialsResolver = mock(EmbeddingCredentialsResolver.class); + when(embeddingCredentialsResolver.resolveEmbeddingCredentials(any(), any())) + .thenReturn(mock(EmbeddingCredentials.class)); + + return mcpResource( + () -> routingContext, + () -> securityIdentity, + instanceOf(tenantResolver), + instanceOf(embeddingCredentialsResolver)); + } + + private CqlSessionCacheSupplier sessionCacheSupplier() { + var supplier = mock(CqlSessionCacheSupplier.class); + when(supplier.get()).thenReturn(mock(CQLSessionCache.class)); + return supplier; + } + + private SchemaObjectCacheSupplier schemaCacheSupplier() { + var supplier = mock(SchemaObjectCacheSupplier.class); + when(supplier.get()).thenReturn(schemaObjectCache); + return supplier; + } + + @SuppressWarnings("unchecked") + private static Instance instanceOf(T value) { + Instance instance = mock(Instance.class); + when(instance.get()).thenReturn(value); + return instance; + } +}