From 81ed051aef6dbda8a502ce4242f54e4fe924666d Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Tue, 27 Oct 2026 16:03:39 -0700 Subject: [PATCH] fix: handle MCP tool setup failures Defer MCP context setup and handle request, schema, and provider setup failures at one shared boundary. Known non-security Data API errors keep their codes and messages. Unexpected setup failures return a generic error. Authentication, security, protocol, cancellation, and processor failures keep their existing handling. Rebased onto fix/mcp-request-context-cleanup without the readable-results change. Tool results keep empty text content, so the tests check structuredContent and _meta instead. --- .../sgv2/jsonapi/api/v1/mcp/McpResource.java | 232 ++++++---- .../v1/mcp/McpSetupErrorsIntegrationTest.java | 162 +++++++ .../api/v1/mcp/McpSetupErrorsTest.java | 417 ++++++++++++++++++ 3 files changed, 717 insertions(+), 94 deletions(-) create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsIntegrationTest.java create mode 100644 src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsTest.java 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 cb66862762..82e56a039b 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 @@ -5,7 +5,11 @@ import com.fasterxml.jackson.databind.ObjectMapper; import io.micrometer.core.instrument.MeterRegistry; +import io.quarkiverse.mcp.server.Cancellation; +import io.quarkiverse.mcp.server.InputRequiredException; +import io.quarkiverse.mcp.server.McpException; import io.quarkiverse.mcp.server.ToolResponse; +import io.quarkus.security.AuthenticationException; import io.quarkus.security.identity.SecurityIdentity; import io.smallrye.mutiny.Uni; import io.stargate.sgv2.jsonapi.ConfigPreLoader; @@ -15,7 +19,10 @@ 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.exception.APIException; +import io.stargate.sgv2.jsonapi.exception.APISecurityException; import io.stargate.sgv2.jsonapi.exception.SchemaException; +import io.stargate.sgv2.jsonapi.exception.ServerException; import io.stargate.sgv2.jsonapi.metrics.JsonProcessingMetricsReporter; import io.stargate.sgv2.jsonapi.service.cqldriver.CqlSessionCacheSupplier; import io.stargate.sgv2.jsonapi.service.cqldriver.executor.*; @@ -29,6 +36,9 @@ import jakarta.enterprise.inject.Instance; import jakarta.inject.Inject; import jakarta.inject.Provider; +import jakarta.ws.rs.WebApplicationException; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletionException; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -95,20 +105,23 @@ public McpResource( * io.stargate.sgv2.jsonapi.api.v1.GeneralResource} */ public Uni> buildDatabaseContext(GeneralCommand command) { + return Uni.createFrom() + .deferred( + () -> { + var requestContext = createRequestContext(); + var dbIdentifier = SchemaObjectIdentifier.forDatabase(requestContext.tenant()); - var requestContext = createRequestContext(); - var dbIdentifier = SchemaObjectIdentifier.forDatabase(requestContext.tenant()); - - return schemaObjectCacheSupplier - .get() - .getDatabase(requestContext, dbIdentifier, requestContext.userAgent()) - .map( - databaseSchemaObject -> - contextBuilderSupplier - .getBuilder(databaseSchemaObject) - .withCommandName(command.getClass().getSimpleName()) - .withRequestContext(requestContext) - .build()); + return schemaObjectCacheSupplier + .get() + .getDatabase(requestContext, dbIdentifier, requestContext.userAgent()) + .map( + databaseSchemaObject -> + contextBuilderSupplier + .getBuilder(databaseSchemaObject) + .withCommandName(command.getClass().getSimpleName()) + .withRequestContext(requestContext) + .build()); + }); } /** @@ -116,23 +129,26 @@ public Uni> buildDatabaseContext(GeneralCommand command) { * io.stargate.sgv2.jsonapi.api.v1.KeyspaceResource} */ public Uni> buildKeyspaceContext(String keyspace, KeyspaceCommand command) { + return Uni.createFrom() + .deferred( + () -> { + var requestContext = createRequestContext(); + var keyspaceIdentifier = + SchemaObjectIdentifier.forKeyspace( + requestContext.tenant(), cqlIdentifierFromUserInput(keyspace)); - var requestContext = createRequestContext(); - var keyspaceIdentifier = - SchemaObjectIdentifier.forKeyspace( - requestContext.tenant(), cqlIdentifierFromUserInput(keyspace)); - - // Force refresh on all keyspace commands because they are all DDL commands - return schemaObjectCacheSupplier - .get() - .getKeyspace(requestContext, keyspaceIdentifier, requestContext.userAgent(), true) - .map( - keyspaceSchemaObject -> - contextBuilderSupplier - .getBuilder(keyspaceSchemaObject) - .withCommandName(command.getClass().getSimpleName()) - .withRequestContext(requestContext) - .build()); + // Force refresh on all keyspace commands because they are all DDL commands + return schemaObjectCacheSupplier + .get() + .getKeyspace(requestContext, keyspaceIdentifier, requestContext.userAgent(), true) + .map( + keyspaceSchemaObject -> + contextBuilderSupplier + .getBuilder(keyspaceSchemaObject) + .withCommandName(command.getClass().getSimpleName()) + .withRequestContext(requestContext) + .build()); + }); } /** @@ -142,76 +158,69 @@ public Uni> buildKeyspaceContext(String keyspace, KeyspaceComm */ public Uni processCollectionCommand( String keyspace, String collection, CollectionCommand command) { + Uni> contextUni = + Uni.createFrom() + .deferred( + () -> { + var requestContext = createRequestContext(); + var unscopedSchemaIdentifier = + new UnscopedSchemaObjectIdentifier.DefaultKeyspaceScopedName( + cqlIdentifierFromUserInput(keyspace), + cqlIdentifierFromUserInput(collection)); - var requestContext = createRequestContext(); + return schemaObjectCacheSupplier + .get() + .getTableBased( + requestContext, + unscopedSchemaIdentifier, + requestContext.userAgent(), + command.isForceSchemaRefresh()) + .map( + schemaObject -> { + VectorColumnDefinition vectorColDef = null; + if (schemaObject.type() == SchemaObjectType.COLLECTION) { + vectorColDef = + schemaObject + .vectorConfig() + .getColumnDefinition(VECTOR_EMBEDDING_TEXT_FIELD) + .orElse(null); + } else if (schemaObject.type() == SchemaObjectType.TABLE) { + vectorColDef = + schemaObject + .vectorConfig() + .getFirstVectorColumnWithVectorizeDefinition() + .orElse(null); + } - var unscopedSchemaIdentifier = - new UnscopedSchemaObjectIdentifier.DefaultKeyspaceScopedName( - cqlIdentifierFromUserInput(keyspace), cqlIdentifierFromUserInput(collection)); - - return schemaObjectCacheSupplier - .get() - .getTableBased( - requestContext, - unscopedSchemaIdentifier, - requestContext.userAgent(), - command.isForceSchemaRefresh()) - .onItemOrFailure() - .transformToUni( - (schemaObject, throwable) -> { - if (throwable != null) { - // If schema resolution or authorization fails, return an error ToolResponse - CommandResult errorResult = - CommandResult.statusOnlyBuilder(RequestTracing.NO_OP) - .addThrowable(throwable) - .build(); - return Uni.createFrom().item(errorResult.toToolResponse()); - } else { - VectorColumnDefinition vectorColDef = null; - if (schemaObject.type() == SchemaObjectType.COLLECTION) { - vectorColDef = - schemaObject - .vectorConfig() - .getColumnDefinition(VECTOR_EMBEDDING_TEXT_FIELD) - .orElse(null); - } else if (schemaObject.type() == SchemaObjectType.TABLE) { - vectorColDef = - schemaObject - .vectorConfig() - .getFirstVectorColumnWithVectorizeDefinition() - .orElse(null); - } + EmbeddingProvider embeddingProvider = null; + if (vectorColDef != null + && vectorColDef.vectorizeDefinition() != null) { + embeddingProvider = + embeddingProviderFactory.create( + requestContext.tenant(), + requestContext.authToken(), + vectorColDef.vectorizeDefinition().provider(), + vectorColDef.vectorizeDefinition().modelName(), + vectorColDef.vectorSize(), + vectorColDef.vectorizeDefinition().parameters(), + vectorColDef.vectorizeDefinition().authentication(), + command.getClass().getSimpleName()); + } - EmbeddingProvider embeddingProvider = null; - - if (vectorColDef != null && vectorColDef.vectorizeDefinition() != null) { - embeddingProvider = - embeddingProviderFactory.create( - requestContext.tenant(), - requestContext.authToken(), - vectorColDef.vectorizeDefinition().provider(), - vectorColDef.vectorizeDefinition().modelName(), - vectorColDef.vectorSize(), - vectorColDef.vectorizeDefinition().parameters(), - vectorColDef.vectorizeDefinition().authentication(), - command.getClass().getSimpleName()); - } - - var commandContext = - contextBuilderSupplier - .getBuilder(schemaObject) - .withEmbeddingProvider(embeddingProvider) - .withCommandName(command.getClass().getSimpleName()) - .withRequestContext(requestContext) - .build(); - - return processCommand(Uni.createFrom().item(commandContext), command); - } - }); + return contextBuilderSupplier + .getBuilder(schemaObject) + .withEmbeddingProvider(embeddingProvider) + .withCommandName(command.getClass().getSimpleName()) + .withRequestContext(requestContext) + .build(); + }); + }); + return processCommand(contextUni, command); } /** - * Process a command and map its result to an MCP tool response. + * Normalize context setup failures, then process a command and map its result to an MCP tool + * response. * * @param contextUni The command context * @param command The command to execute @@ -220,7 +229,10 @@ public Uni processCollectionCommand( public Uni processCommand(Uni> contextUni, Command command) { return contextUni - .flatMap(context -> executeCommand(context, command)) + .onItemOrFailure() + .transformToUni( + (context, failure) -> + failure == null ? executeCommand(context, command) : handleSetupFailure(failure)) .map(CommandResult::toToolResponse); } @@ -252,6 +264,38 @@ private Uni executeCommand(CommandContext context, Command com }); } + private Uni handleSetupFailure(Throwable failure) { + // Synchronous cache reads use CompletableFuture.join(), unlike the asynchronous cache path. + while (failure instanceof CompletionException && failure.getCause() != null) { + failure = failure.getCause(); + } + + // Authentication, protocol errors, and framework control signals retain their own handling. + if (failure instanceof Error + || failure instanceof SecurityException + || failure instanceof APISecurityException + || failure instanceof AuthenticationException + || failure instanceof McpException + || failure instanceof InputRequiredException + || failure instanceof Cancellation.OperationCancellationException + || failure instanceof CancellationException + || (failure instanceof WebApplicationException webFailure + && (webFailure.getResponse().getStatus() == 401 + || webFailure.getResponse().getStatus() == 403))) { + return Uni.createFrom().failure(failure); + } + + // Only APIException messages are intended for clients. Arbitrary setup failures may include + // credentials or provider payloads, so do not pass them to the generic throwable mapper. + var error = + failure instanceof APIException apiException + ? apiException + : ServerException.Code.UNEXPECTED_SERVER_ERROR.withPreformattedMessage( + "Unable to prepare the tool request. Retry the request or contact support."); + return Uni.createFrom() + .item(CommandResult.statusOnlyBuilder(RequestTracing.NO_OP).addThrowable(error).build()); + } + /** * Create a {@link RequestContext} for the current MCP request. Uses the MCP-specific constructor * that resolves the auth token from {@link SecurityIdentity} instead of JAX-RS {@link diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsIntegrationTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsIntegrationTest.java new file mode 100644 index 0000000000..a4f0494133 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsIntegrationTest.java @@ -0,0 +1,162 @@ +package io.stargate.sgv2.jsonapi.api.v1.mcp; + +import static io.stargate.sgv2.jsonapi.api.v1.mcp.McpHttpTestSupport.assertResponseComplete; +import static io.stargate.sgv2.jsonapi.api.v1.mcp.McpHttpTestSupport.rpcPayload; +import static io.stargate.sgv2.jsonapi.api.v1.mcp.McpHttpTestSupport.toolPayload; +import static org.assertj.core.api.Assertions.assertThat; + +import io.quarkiverse.mcp.server.JsonRpcErrorCodes; +import io.quarkus.test.common.WithTestResource; +import io.quarkus.test.junit.QuarkusIntegrationTest; +import io.restassured.response.Response; +import io.restassured.specification.RequestSpecification; +import io.stargate.sgv2.jsonapi.api.model.command.CommandName; +import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; +import io.stargate.sgv2.jsonapi.exception.APISecurityException; +import io.stargate.sgv2.jsonapi.exception.SchemaException; +import io.stargate.sgv2.jsonapi.testresource.DseTestResource; +import io.vertx.core.json.JsonObject; +import java.util.Map; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +/** + * Verifies application setup errors without changing authentication or framework error handling. + */ +@QuarkusIntegrationTest +@WithTestResource(value = DseTestResource.class) +class McpSetupErrorsIntegrationTest extends McpIntegrationTestBase { + @BeforeAll + void createTestKeyspace() { + createKeyspace(keyspaceName); + } + + @AfterAll + void dropTestKeyspace() { + dropKeyspace(keyspaceName); + } + + @Test + void unknownKeyspaceReturnsAnApplicationErrorThenAValidCallSucceeds() { + var response = + callTool( + 1, CommandName.Names.FIND_COLLECTIONS, Map.of("keyspace", keyspaceName + "_missing")); + + assertApplicationError(response, 1, SchemaException.Code.UNKNOWN_KEYSPACE); + assertProviderDiscoverySucceeds(2); + } + + @Test + void unknownCollectionReturnsAnApplicationErrorThenAValidCallSucceeds() { + var response = + callTool( + 3, + CommandName.Names.COUNT_DOCUMENTS, + Map.of("keyspace", keyspaceName, "collection", collectionName)); + + assertApplicationError(response, 3, SchemaException.Code.UNKNOWN_COLLECTION_OR_TABLE); + assertProviderDiscoverySucceeds(4); + } + + @Test + void missingTokenRemainsAnHttpDenialThenAValidCallSucceeds() { + var response = + request("tools/call", CommandName.Names.FIND_EMBEDDING_PROVIDERS, false) + .body(toolPayload(5, CommandName.Names.FIND_EMBEDDING_PROVIDERS, Map.of())) + .post(MCP_PATH); + + assertThat(response.statusCode()).as(response.asString()).isEqualTo(401); + var body = new JsonObject(response.asString()); + assertThat(body.containsKey("result")).isFalse(); + assertThat(body.getJsonArray("errors").getJsonObject(0).getString("errorCode")) + .isEqualTo(APISecurityException.Code.MISSING_AUTHENTICATION_TOKEN.name()); + assertProviderDiscoverySucceeds(6); + } + + @Test + void malformedJsonRemainsAProtocolError() { + var response = + request("tools/call", CommandName.Names.FIND_EMBEDDING_PROVIDERS, true) + .body("{invalid-json") + .post(MCP_PATH); + + // Quarkiverse 2.0.3 reports parse errors in a JSON-RPC envelope with HTTP 200. + assertProtocolError(response, null, JsonRpcErrorCodes.PARSE_ERROR); + assertProviderDiscoverySucceeds(7); + } + + @Test + void unknownMethodAndToolRemainProtocolErrors() { + String method = "tools/unknownMethod"; + var response = request(method, null, true).body(rpcPayload(8, method, Map.of())).post(MCP_PATH); + + assertProtocolError(response, 8, JsonRpcErrorCodes.METHOD_NOT_FOUND); + assertProtocolError(callTool(9, "unknownTool", Map.of()), 9, JsonRpcErrorCodes.INVALID_PARAMS); + assertProviderDiscoverySucceeds(10); + } + + @Test + void missingArgumentRemainsAFrameworkToolError() { + var result = + assertResponseComplete(callTool(11, CommandName.Names.FIND_COLLECTIONS, Map.of()), 11); + + // The framework's default input-validation strategy returns a plain tool error. + assertThat(result.getBoolean("isError")).isTrue(); + assertThat(result.getValue("structuredContent")).isNull(); + var content = result.getJsonArray("content"); + assertThat(content).hasSize(1); + assertThat(content.getJsonObject(0).getString("type")).isEqualTo("text"); + assertThat(content.getJsonObject(0).getString("text")) + .isEqualTo("Missing required argument: keyspace"); + assertProviderDiscoverySucceeds(12); + } + + private void assertProviderDiscoverySucceeds(int id) { + var result = + assertResponseComplete( + callTool(id, CommandName.Names.FIND_EMBEDDING_PROVIDERS, Map.of()), id); + assertThat(result.getBoolean("isError")).isFalse(); + var status = result.getJsonObject("_meta").getJsonObject("status"); + assertThat(status.getJsonObject("embeddingProviders")).isNotNull(); + assertThat(result.getValue("structuredContent")).isNull(); + assertThat(result.getJsonArray("content")).isEmpty(); + } + + private static void assertApplicationError(Response response, int id, SchemaException.Code code) { + var result = assertResponseComplete(response, id); + assertThat(result.getBoolean("isError")).isTrue(); + var structured = result.getJsonObject("structuredContent"); + var errors = structured.getJsonArray("errors"); + assertThat(errors).hasSize(1); + assertThat(errors.getJsonObject(0).getString("errorCode")).isEqualTo(code.name()); + assertThat(errors.getJsonObject(0).getString("message")).isNotBlank(); + assertThat(result.getJsonArray("content")).isEmpty(); + } + + private static void assertProtocolError(Response response, Integer id, int errorCode) { + assertThat(response.statusCode()).as(response.asString()).isEqualTo(200); + var body = new JsonObject(response.asString()); + assertThat(body.getInteger("id")).isEqualTo(id); + assertThat(body.containsKey("result")).isFalse(); + assertThat(body.getJsonObject("error").getInteger("code")).isEqualTo(errorCode); + } + + private Response callTool(int id, String tool, Map arguments) { + return request("tools/call", tool, true).body(toolPayload(id, tool, arguments)).post(MCP_PATH); + } + + private RequestSpecification request(String method, String tool, boolean authenticated) { + var request = + McpHttpTestSupport.request(method).header(ApiFeature.MCP.httpHeaderName(), "true"); + if (tool != null) { + request.header("Mcp-Name", tool); + } + if (authenticated) { + for (var header : authHeaders()) { + request.header(header.getKey(), header.getValue()); + } + } + return request; + } +} diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsTest.java new file mode 100644 index 0000000000..572fb2fcb6 --- /dev/null +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpSetupErrorsTest.java @@ -0,0 +1,417 @@ +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.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doCallRealMethod; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.micrometer.core.instrument.MeterRegistry; +import io.quarkiverse.mcp.server.Cancellation; +import io.quarkiverse.mcp.server.InputRequiredException; +import io.quarkiverse.mcp.server.JsonRpcErrorCodes; +import io.quarkiverse.mcp.server.McpException; +import io.quarkiverse.mcp.server.MetaKey; +import io.quarkiverse.mcp.server.ToolResponse; +import io.quarkus.security.AuthenticationException; +import io.quarkus.security.AuthenticationFailedException; +import io.quarkus.security.ForbiddenException; +import io.quarkus.security.UnauthorizedException; +import io.quarkus.security.identity.SecurityIdentity; +import io.smallrye.mutiny.Uni; +import io.stargate.sgv2.jsonapi.ConfigPreLoader; +import io.stargate.sgv2.jsonapi.TestConstants; +import io.stargate.sgv2.jsonapi.api.model.command.CollectionCommand; +import io.stargate.sgv2.jsonapi.api.model.command.CommandConfig; +import io.stargate.sgv2.jsonapi.api.model.command.CommandResult; +import io.stargate.sgv2.jsonapi.api.model.command.CommandStatus; +import io.stargate.sgv2.jsonapi.api.model.command.GeneralCommand; +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.EmbeddingCredentialsResolver; +import io.stargate.sgv2.jsonapi.api.request.tenant.RequestTenantResolver; +import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; +import io.stargate.sgv2.jsonapi.config.feature.FeaturesConfig; +import io.stargate.sgv2.jsonapi.exception.APIException; +import io.stargate.sgv2.jsonapi.exception.APISecurityException; +import io.stargate.sgv2.jsonapi.exception.EmbeddingProviderException; +import io.stargate.sgv2.jsonapi.exception.SchemaException; +import io.stargate.sgv2.jsonapi.exception.ServerException; +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.vertx.core.MultiMap; +import io.vertx.core.http.HttpServerRequest; +import io.vertx.ext.web.RoutingContext; +import jakarta.enterprise.inject.Instance; +import jakarta.ws.rs.WebApplicationException; +import java.security.Principal; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletionException; +import java.util.stream.Stream; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.MockedStatic; +import org.slf4j.MDC; + +/** Setup failures become tool results. Transport and processor failures keep their owners. */ +class McpSetupErrorsTest { + + private static final TestConstants TEST_CONSTANTS = new TestConstants(); + private static final String SECRET = "secret-token-from-internal-setup"; + private static final Duration TIMEOUT = Duration.ofSeconds(5); + + private final ObjectMapper objectMapper = new ObjectMapper(); + private final MeteredCommandProcessor commandProcessor = mock(MeteredCommandProcessor.class); + private final SchemaObjectCache schemaObjectCache = mock(SchemaObjectCache.class); + private final EmbeddingProviderFactory embeddingProviderFactory = + mock(EmbeddingProviderFactory.class); + private final RequestTenantResolver tenantResolver = mock(RequestTenantResolver.class); + private final GeneralCommand generalCommand = mock(GeneralCommand.class); + private final KeyspaceCommand keyspaceCommand = mock(KeyspaceCommand.class); + private final CollectionCommand collectionCommand = mock(CollectionCommand.class); + private MockedStatic configPreLoader; + private McpResource resource; + + @BeforeEach + void setUp() { + // Supply only application configuration. MCP enablement still comes from real request headers. + var commandConfig = mock(CommandConfig.class); + when(commandConfig.get(FeaturesConfig.class)).thenReturn(() -> Map.of()); + configPreLoader = mockStatic(ConfigPreLoader.class); + configPreLoader.when(ConfigPreLoader::getPreLoadOrEmpty).thenReturn(commandConfig); + + var headers = MultiMap.caseInsensitiveMultiMap().add(ApiFeature.MCP.httpHeaderName(), "true"); + var httpRequest = mock(HttpServerRequest.class); + when(httpRequest.headers()).thenReturn(headers); + 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); + when(tenantResolver.resolve(any(), any())).thenReturn(TEST_CONSTANTS.TENANT); + var credentialsResolver = mock(EmbeddingCredentialsResolver.class); + when(credentialsResolver.resolveEmbeddingCredentials(any(), any())) + .thenReturn(TEST_CONSTANTS.EMBEDDING_CREDENTIALS); + var sessionCacheSupplier = mock(CqlSessionCacheSupplier.class); + when(sessionCacheSupplier.get()).thenReturn(mock(CQLSessionCache.class)); + var schemaCacheSupplier = mock(SchemaObjectCacheSupplier.class); + when(schemaCacheSupplier.get()).thenReturn(schemaObjectCache); + + resource = + new McpResource( + objectMapper, + commandProcessor, + mock(MeterRegistry.class), + mock(JsonProcessingMetricsReporter.class), + sessionCacheSupplier, + embeddingProviderFactory, + mock(RerankingProviderFactory.class), + schemaCacheSupplier, + () -> routingContext, + () -> securityIdentity, + instanceOf(tenantResolver), + instanceOf(credentialsResolver)); + } + + @AfterEach + void tearDown() { + configPreLoader.close(); + MDC.clear(); + } + + @ParameterizedTest + @EnumSource(CommandScope.class) + void requestContextFailureIsDeferredAndReturnedAsAToolError(CommandScope scope) throws Exception { + var failure = SchemaException.Code.UNKNOWN_KEYSPACE.get("keyspace", "missing-keyspace"); + when(tenantResolver.resolve(any(), any())).thenThrow(failure); + + var response = assertDoesNotThrow(() -> call(scope)); + + verifyNoInteractions(tenantResolver, schemaObjectCache, commandProcessor); + assertSetupError(response, failure); + verifyNoInteractions(schemaObjectCache); + } + + @Test + void nullKeyspaceIdentifierBecomesADeferredToolError() throws Exception { + var response = + assertDoesNotThrow( + () -> + resource.processCommand( + resource.buildKeyspaceContext(null, keyspaceCommand), keyspaceCommand)); + + verifyNoInteractions(tenantResolver, schemaObjectCache, commandProcessor); + assertSetupError(response, ServerException.Code.UNEXPECTED_SERVER_ERROR.name()); + verifyNoInteractions(schemaObjectCache); + } + + @Test + void nullCollectionIdentifierBecomesADeferredToolError() throws Exception { + // Collection names are validated inside the schema lookup before any cache or database access. + doCallRealMethod().when(schemaObjectCache).getTableBased(any(), any(), any(), anyBoolean()); + var response = + assertDoesNotThrow( + () -> + resource.processCollectionCommand( + TEST_CONSTANTS.KEYSPACE_NAME, null, collectionCommand)); + + verifyNoInteractions(tenantResolver, schemaObjectCache, commandProcessor); + assertSetupError(response, ServerException.Code.UNEXPECTED_SERVER_ERROR.name()); + verify(schemaObjectCache).getTableBased(any(), any(), any(), anyBoolean()); + } + + @Test + void databaseSchemaSecurityFailureRemainsATransportFailure() { + var failure = APISecurityException.Code.UNAUTHENTICATED_REQUEST.get(); + when(schemaObjectCache.getDatabase(any(), any(), any())) + .thenReturn(Uni.createFrom().failure(failure)); + + var response = assertDoesNotThrow(() -> call(CommandScope.DATABASE)); + + assertThatThrownBy(() -> response.await().atMost(TIMEOUT)).isSameAs(failure); + verifyNoInteractions(commandProcessor); + } + + @Test + void keyspaceSchemaFailureReturnsItsActionableError() throws Exception { + var failure = SchemaException.Code.UNKNOWN_KEYSPACE.get("keyspace", "missing-keyspace"); + when(schemaObjectCache.getKeyspace(any(), any(), any(), eq(true))) + .thenReturn(Uni.createFrom().failure(failure)); + + assertSetupError(call(CommandScope.KEYSPACE), failure); + } + + @Test + void collectionSchemaFailureReturnsItsActionableError() throws Exception { + var failure = + SchemaException.Code.UNKNOWN_COLLECTION_OR_TABLE.get( + "keyspace", + TEST_CONSTANTS.KEYSPACE_NAME, + "table", + "missing-collection", + "allTables", + "[]"); + when(schemaObjectCache.getTableBased(any(), any(), any(), anyBoolean())) + .thenReturn(Uni.createFrom().failure(failure)); + + assertSetupError(call(CommandScope.COLLECTION), failure); + } + + @Test + void providerConstructionFailureReturnsItsActionableError() throws Exception { + doReturn(Uni.createFrom().item(TEST_CONSTANTS.VECTORIZE_RERANK_COLLECTION_SCHEMA_OBJECT)) + .when(schemaObjectCache) + .getTableBased(any(), any(), any(), anyBoolean()); + var failure = + EmbeddingProviderException.Code.EMBEDDING_PROVIDER_UNAVAILABLE.get( + "errorMessage", "The selected provider is unavailable."); + when(embeddingProviderFactory.create( + any(), any(), any(), any(), anyInt(), any(), any(), anyString())) + .thenThrow(failure); + + var response = assertDoesNotThrow(() -> call(CommandScope.COLLECTION)); + + assertSetupError(response, failure); + verify(embeddingProviderFactory) + .create(any(), any(), any(), any(), anyInt(), any(), any(), anyString()); + } + + @Test + void unknownSetupFailureDoesNotExposeExceptionDetails() throws Exception { + when(schemaObjectCache.getTableBased(any(), any(), any(), anyBoolean())) + .thenReturn(Uni.createFrom().failure(new IllegalStateException(SECRET))); + + var error = + assertSetupError( + call(CommandScope.COLLECTION), ServerException.Code.UNEXPECTED_SERVER_ERROR.name()); + + assertThat(error.toString()).doesNotContain(SECRET, IllegalStateException.class.getName()); + } + + @Test + void completionWrappedSetupErrorRetainsItsActionableMessage() throws Exception { + var failure = SchemaException.Code.UNKNOWN_KEYSPACE.get("keyspace", "missing-keyspace"); + // Cache reads can synchronously surface a failed future through CompletionException. + doThrow(new CompletionException(failure)) + .when(schemaObjectCache) + .getTableBased(any(), any(), any(), anyBoolean()); + + var response = assertDoesNotThrow(() -> call(CommandScope.COLLECTION)); + + assertSetupError(response, failure); + } + + @Test + void completionWrappedUnauthorizedFailureRemainsATransportFailure() { + var failure = new UnauthorizedException("authentication required"); + doThrow(new CompletionException(failure)) + .when(schemaObjectCache) + .getTableBased(any(), any(), any(), anyBoolean()); + + var response = assertDoesNotThrow(() -> call(CommandScope.COLLECTION)); + + assertThatThrownBy(() -> response.await().atMost(TIMEOUT)).isSameAs(failure); + verifyNoInteractions(commandProcessor); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("transportFailures") + void collectionSetupPreservesTransportAndProtocolFailures(Throwable failure) { + when(schemaObjectCache.getTableBased(any(), any(), any(), anyBoolean())) + .thenReturn(Uni.createFrom().failure(failure)); + + var response = assertDoesNotThrow(() -> call(CommandScope.COLLECTION)); + + assertThatThrownBy(() -> response.await().atMost(TIMEOUT)).isSameAs(failure); + verifyNoInteractions(commandProcessor); + } + + private static Stream transportFailures() { + return Stream.of( + new SecurityException("resource access denied"), + new UnauthorizedException("authentication required"), + new ForbiddenException("access denied"), + new AuthenticationFailedException("authentication failed"), + new AuthenticationMarkerFailure(), + new WebApplicationException("authentication required", 401), + new WebApplicationException("access denied", 403), + new McpException("invalid protocol request", JsonRpcErrorCodes.INVALID_REQUEST), + InputRequiredException.builder().addRootsRequest("roots").build(), + new Cancellation.OperationCancellationException(), + new CancellationException("request canceled")); + } + + @Test + void processorFailureIsNotConvertedByTheSetupErrorHandler() { + when(schemaObjectCache.getKeyspace(any(), any(), any(), eq(true))) + .thenReturn(Uni.createFrom().item(TEST_CONSTANTS.KEYSPACE_SCHEMA_OBJECT)); + var failure = new IllegalStateException("processor owns this failure"); + when(commandProcessor.processCommand(any(), eq(keyspaceCommand))) + .thenReturn(Uni.createFrom().failure(failure)); + + var response = call(CommandScope.KEYSPACE); + + assertThatThrownBy(() -> response.await().atMost(TIMEOUT)).isSameAs(failure); + verify(commandProcessor).processCommand(any(), eq(keyspaceCommand)); + } + + @Test + void validCallSucceedsAfterASetupErrorOnTheSameResource() throws Exception { + var failure = SchemaException.Code.UNKNOWN_KEYSPACE.get("keyspace", "missing-keyspace"); + when(schemaObjectCache.getKeyspace(any(), any(), any(), eq(true))) + .thenReturn(Uni.createFrom().failure(failure)) + .thenReturn(Uni.createFrom().item(TEST_CONSTANTS.KEYSPACE_SCHEMA_OBJECT)); + + assertSetupError(call(CommandScope.KEYSPACE), failure); + + assertNextKeyspaceCallSucceeds(); + } + + @Test + void validCallSucceedsAfterATransportDenialOnTheSameResource() throws Exception { + var failure = new UnauthorizedException("authentication required"); + when(schemaObjectCache.getKeyspace(any(), any(), any(), eq(true))) + .thenReturn(Uni.createFrom().failure(failure)) + .thenReturn(Uni.createFrom().item(TEST_CONSTANTS.KEYSPACE_SCHEMA_OBJECT)); + + var response = call(CommandScope.KEYSPACE); + assertThatThrownBy(() -> response.await().atMost(TIMEOUT)).isSameAs(failure); + verifyNoInteractions(commandProcessor); + + assertNextKeyspaceCallSucceeds(); + } + + private void assertNextKeyspaceCallSucceeds() { + var result = + CommandResult.statusOnlyBuilder(RequestTracing.NO_OP) + .addStatus(CommandStatus.OK, 1) + .build(); + when(commandProcessor.processCommand(any(), eq(keyspaceCommand))) + .thenReturn(Uni.createFrom().item(result)); + + var response = call(CommandScope.KEYSPACE).await().atMost(TIMEOUT); + + assertThat(response.isError()).isFalse(); + assertThat(response.content()).isEmpty(); + assertThat(response.structuredContent()).isNull(); + assertThat(response._meta()) + .containsExactlyEntriesOf(Map.of(MetaKey.of("status"), result.status())); + verify(commandProcessor).processCommand(any(), eq(keyspaceCommand)); + } + + private Uni call(CommandScope scope) { + return switch (scope) { + case DATABASE -> + resource.processCommand(resource.buildDatabaseContext(generalCommand), generalCommand); + case KEYSPACE -> + resource.processCommand( + resource.buildKeyspaceContext(TEST_CONSTANTS.KEYSPACE_NAME, keyspaceCommand), + keyspaceCommand); + case COLLECTION -> + resource.processCollectionCommand( + TEST_CONSTANTS.KEYSPACE_NAME, TEST_CONSTANTS.COLLECTION_NAME, collectionCommand); + }; + } + + private void assertSetupError(Uni responseUni, APIException failure) { + var error = assertSetupError(responseUni, failure.code); + assertThat(error.path("message").asText()).isEqualTo(failure.body); + } + + private JsonNode assertSetupError(Uni responseUni, String code) { + var response = responseUni.await().atMost(TIMEOUT); + assertThat(response.isError()).isTrue(); + assertThat(response.content()).isEmpty(); + JsonNode structured = objectMapper.valueToTree(response.structuredContent()); + assertThat(structured.path("errors")).hasSize(1); + var error = structured.path("errors").get(0); + assertThat(error.path("errorCode").asText()).isEqualTo(code); + assertThat(error.path("message").asText()).isNotBlank(); + assertThat(response._meta()).isEmpty(); + verifyNoInteractions(commandProcessor); + return error; + } + + @SuppressWarnings("unchecked") + private static Instance instanceOf(T value) { + Instance instance = mock(Instance.class); + when(instance.get()).thenReturn(value); + return instance; + } + + private enum CommandScope { + DATABASE, + KEYSPACE, + COLLECTION + } + + /** Quarkus treats this marker as authentication even without a SecurityException superclass. */ + private static final class AuthenticationMarkerFailure extends RuntimeException + implements AuthenticationException {} +}