Skip to content

Handle MCP tool setup failures #2653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
erichare wants to merge 1 commit into from
Open

Handle MCP tool setup failures #2653

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
232 changes: 138 additions & 94 deletions src/main/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResource.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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.*;
Expand All @@ -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;

Expand Down Expand Up @@ -95,44 +105,50 @@ public McpResource(
* io.stargate.sgv2.jsonapi.api.v1.GeneralResource}
*/
public Uni<CommandContext<?>> 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());
});
}

/**
* Build a CommandContext for keyspace-level (keyspace) commands. similar to {@link
* io.stargate.sgv2.jsonapi.api.v1.KeyspaceResource}
*/
public Uni<CommandContext<?>> 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());
});
}

/**
Expand All @@ -142,76 +158,69 @@ public Uni<CommandContext<?>> buildKeyspaceContext(String keyspace, KeyspaceComm
*/
public Uni<ToolResponse> processCollectionCommand(
String keyspace, String collection, CollectionCommand command) {
Uni<CommandContext<?>> 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
Expand All @@ -220,7 +229,10 @@ public Uni<ToolResponse> processCollectionCommand(
public Uni<ToolResponse> processCommand(Uni<CommandContext<?>> contextUni, Command command) {

return contextUni
.flatMap(context -> executeCommand(context, command))
.onItemOrFailure()
.transformToUni(
(context, failure) ->
failure == null ? executeCommand(context, command) : handleSetupFailure(failure))
.map(CommandResult::toToolResponse);
}

Expand Down Expand Up @@ -252,6 +264,38 @@ private Uni<CommandResult> executeCommand(CommandContext<?> context, Command com
});
}

private Uni<CommandResult> 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
Expand Down
Loading
Loading