From b07ca8d685353e9a10835d3157585342069b95a5 Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Thu, 5 Nov 2026 11:05:19 -0800 Subject: [PATCH 1/2] test: verify MCP request isolation and document tracing support --- docs/mcp.md | 8 + .../jsonapi/api/v1/mcp/McpLoggingTest.java | 243 +++++++++++++++--- .../jsonapi/api/v1/mcp/McpResourceTest.java | 48 ++++ 3 files changed, 264 insertions(+), 35 deletions(-) diff --git a/docs/mcp.md b/docs/mcp.md index 2b283f9d7f..a8aae2cc05 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -26,6 +26,14 @@ The [endpoint comparison tests](../src/test/java/io/stargate/sgv2/jsonapi/api/v1 Use the schemas returned by `tools/list` for accepted arguments. The [query specification](dataapi-spec.md) describes command semantics, and the [configuration guide](../CONFIGURATION.md#operations-configuration) records configurable operation limits. +## Request isolation and tracing + +Send `Feature-Flag-request-tracing` with value `true` to include diagnostic events for a command. Each call has its own request context, generated request ID, and diagnostic trace. The [overlapping-call tests](../src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java) check that callbacks and returned events stay with their original call. The [subscription cancellation test](../src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResourceTest.java) checks cleanup while another command remains active. + +These diagnostic events are separate from distributed tracing. The current build does not include the OpenTelemetry extension. It does not import `traceparent`, `tracestate`, or `baggage` from request `_meta` or HTTP headers into a distributed trace. Those fields do not supply authentication or the generated request IDs. + +Enabling distributed tracing and defining allowed baggage keys and size limits remain tracked in [#2629](https://github.com/stargate/data-api/issues/2629). + ## Diagnostics Follow [MCP logging](../CONFIGURATION.md#mcp-logging) for request correlation, redaction, and temporary diagnostics. That section also covers custom endpoint paths and log handlers. diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java index 958b349720..38c4cbdc26 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java @@ -5,6 +5,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.awaitility.Awaitility.await; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.when; import com.fasterxml.jackson.databind.node.JsonNodeFactory; @@ -12,9 +13,12 @@ import io.quarkus.test.junit.QuarkusTest; import io.quarkus.test.junit.TestProfile; import io.restassured.http.ContentType; +import io.restassured.response.Response; +import io.restassured.specification.RequestSpecification; import io.smallrye.mutiny.Uni; 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.CommandStatus; import io.stargate.sgv2.jsonapi.api.model.command.tracing.RequestTracing; import io.stargate.sgv2.jsonapi.config.feature.ApiFeature; import io.stargate.sgv2.jsonapi.service.processor.CommandProcessor; @@ -26,11 +30,13 @@ import java.nio.file.Path; import java.time.Duration; import java.util.Arrays; +import java.util.List; import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; -import java.util.regex.Pattern; import org.eclipse.microprofile.config.ConfigProvider; import org.junit.jupiter.api.Test; import org.slf4j.LoggerFactory; @@ -217,43 +223,210 @@ void asynchronousCommandFailuresStayOutOfLogs() throws IOException { } @Test - void driverThreadCallbacksKeepRequestIdAndRestoreWorkerContext() throws Exception { - try (var worker = Executors.newSingleThreadExecutor()) { - // Start the worker before a request can contribute inherited context. - assertThat(worker.submit(() -> MDC.get(McpRequestLogging.REQUEST_ID)).get()).isNull(); - var completed = new CompletableFuture(); - when(commandProcessor.processCommand(any(), any())) - .thenAnswer( - ignored -> - Uni.createFrom() - .completionStage(completed) - .map( - value -> { - LoggerFactory.getLogger(McpLoggingTest.class) - .info("synthetic-driver-payload-secret"); - return CommandResult.statusOnlyBuilder(RequestTracing.NO_OP).build(); - })); - var completion = - worker.submit( - () -> { - await() - .atMost(Duration.ofSeconds(10)) - .until(() -> completed.getNumberOfDependents() > 0); - completed.complete(1); - return MDC.get(McpRequestLogging.REQUEST_ID); - }); + void interleavedHttpCallsKeepTheirRequestTracingAndLogContext() throws Exception { + var firstCompletion = new CompletableFuture(); + var secondCompletion = new CompletableFuture(); + var completions = + Map.of("synthetic-token-1", firstCompletion, "synthetic-token-2", secondCompletion); + var started = new LinkedBlockingQueue(); + when(commandProcessor.processCommand(any(), any())) + .thenAnswer( + invocation -> { + CommandContext context = invocation.getArgument(0); + var callbackMdc = new CompletableFuture>(); + var call = new PendingCall(context, MDC.getCopyOfContextMap(), callbackMdc); + return Uni.createFrom() + .emitter( + emitter -> { + completions + .get(context.requestContext().authToken()) + .thenAccept(emitter::complete); + // Both requests remain open until the test completes their futures. + started.add(call); + }) + .map( + value -> { + callbackMdc.complete(MDC.getCopyOfContextMap()); + context.requestTracing().maybeTrace("completed-call-" + value); + LoggerFactory.getLogger(McpLoggingTest.class) + .info("synthetic-driver-payload-secret"); + var result = + CommandResult.statusOnlyBuilder(context.requestTracing()) + .addStatus(CommandStatus.OK, value); + if (value == 2) { + result.addThrowable( + new IllegalArgumentException("synthetic-call-2-error")); + } + return result.build(); + }); + }); + + try (var clients = Executors.newFixedThreadPool(2); + var worker = Executors.newSingleThreadExecutor()) { + // Start the worker before either request can contribute inherited context. + worker.submit(() -> MDC.put(McpRequestLogging.REQUEST_ID, "worker-context")).get(); long start = Files.size(LOG_FILE); + try { + var firstResponse = + clients.submit( + () -> + tracedProviderRequest(1).header("Token", "synthetic-token-1").post(mcpPath())); + var first = awaitStarted(started, firstResponse); + var secondResponse = + clients.submit( + () -> + tracedProviderRequest(2).header("Token", "synthetic-token-2").post(mcpPath())); + var second = awaitStarted(started, secondResponse); + + assertThat(first.context().requestContext()).isNotSameAs(second.context().requestContext()); + assertThat(first.context().requestContext().requestId()) + .isNotEqualTo(second.context().requestContext().requestId()); + assertThat(first.context().requestTracing().enabled()).isTrue(); + assertThat(second.context().requestTracing().enabled()).isTrue(); + assertThat(first.context().requestTracing().getSession().orElseThrow()) + .isNotSameAs(second.context().requestTracing().getSession().orElseThrow()); + assertThat(first.logRequestId()) + .matches("[a-f0-9-]{36}") + .isNotEqualTo(second.logRequestId()); + assertThat(second.logRequestId()).matches("[a-f0-9-]{36}"); - callProviderTool(); + assertThat( + worker + .submit( + () -> { + secondCompletion.complete(2); + return MDC.get(McpRequestLogging.REQUEST_ID); + }) + .get(10, TimeUnit.SECONDS)) + .isEqualTo("worker-context"); + var secondResult = assertResponseComplete(secondResponse.get(10, TimeUnit.SECONDS), 2); + assertThat(firstResponse.isDone()).isFalse(); + assertThat(firstCompletion.isDone()).isFalse(); + assertThat( + worker + .submit( + () -> { + firstCompletion.complete(1); + return MDC.get(McpRequestLogging.REQUEST_ID); + }) + .get(10, TimeUnit.SECONDS)) + .isEqualTo("worker-context"); + var firstResult = assertResponseComplete(firstResponse.get(10, TimeUnit.SECONDS), 1); + + assertIsolatedResult(first, firstResult, 1); + assertIsolatedResult(second, secondResult, 2); + var logs = logsSince(start); + for (var call : List.of(first, second)) { + assertThat(logs).contains("MCP request completed requestId=" + call.logRequestId()); + traceMetadata(call == first ? 1 : 2) + .values() + .forEach(value -> assertThat(logs).doesNotContain(value)); + } + assertThat(logs) + .contains( + "MCP command completed requestId=" + + first.logRequestId() + + " command=findEmbeddingProviders isError=false", + "MCP command completed requestId=" + + second.logRequestId() + + " command=findEmbeddingProviders isError=true", + "MCP event details redacted") + .doesNotContain("synthetic-", "requestId=null"); + } finally { + firstCompletion.complete(1); + secondCompletion.complete(2); + } + } + } + + @Test + void traceMetadataCannotAuthenticateAToolCall() throws IOException { + long start = Files.size(LOG_FILE); + var response = tracedProviderRequest(1).post(mcpPath()); + + assertThat(response.statusCode()).isEqualTo(401); + verifyNoInteractions(commandProcessor); + assertThat(response.asString()).doesNotContain("synthetic-"); + var logs = logsSince(start); + assertThat(logs).contains("status=401").doesNotContain("synthetic-"); + traceMetadata(1) + .values() + .forEach( + value -> { + assertThat(response.asString()).doesNotContain(value); + assertThat(logs).doesNotContain(value); + }); + } + + private static PendingCall awaitStarted( + LinkedBlockingQueue started, Future response) throws Exception { + var call = started.poll(10, TimeUnit.SECONDS); + if (call == null && response.isDone()) { + throw new AssertionError( + "The request ended before reaching the command processor: " + response.get().asString()); + } + assertThat(call).as("The command processor received the request").isNotNull(); + return call; + } + + private void assertIsolatedResult(PendingCall call, JsonObject result, int id) throws Exception { + assertThat(result.getBoolean("isError")).isEqualTo(id == 2); + var status = result.getJsonObject("_meta").getJsonObject("status"); + assertThat(status.getInteger("ok")).isEqualTo(id); + var trace = status.getJsonObject("trace").getJsonObject("TraceSession"); + assertThat(trace.getString("requestId")) + .isEqualTo(call.context().requestContext().requestId()) + .isEqualTo(call.logRequestId()); + assertThat(trace.getJsonArray("events")).hasSize(1); + assertThat( + trace + .getJsonArray("events") + .getJsonObject(0) + .getJsonObject("TraceEvent") + .getString("message")) + .isEqualTo("completed-call-" + id); + var callbackMdc = call.callbackMdc().get(10, TimeUnit.SECONDS); + assertThat(callbackMdc).containsEntry(McpRequestLogging.REQUEST_ID, call.logRequestId()); + for (var mdc : List.of(call.startMdc(), callbackMdc)) { + assertThat(mdc) + .doesNotContainKeys("traceId", "spanId", "traceparent", "tracestate", "baggage"); + } + for (int otherId : List.of(1, 2)) { + for (String value : traceMetadata(otherId).values()) { + assertThat(result.encode()).doesNotContain(value); + assertThat(call.startMdc()).doesNotContainValue(value); + assertThat(callbackMdc).doesNotContainValue(value); + } + } + } + + private RequestSpecification tracedProviderRequest(int id) { + var payload = + new JsonObject(McpHttpTestSupport.toolPayload(id, "findEmbeddingProviders", Map.of())) + .copy(); + var meta = payload.getJsonObject("params").getJsonObject("_meta"); + traceMetadata(id).forEach(meta::put); + return McpHttpTestSupport.request("tools/call") + .header("Mcp-Name", "findEmbeddingProviders") + .header(ApiFeature.REQUEST_TRACING.httpHeaderName(), "true") + .headers(traceMetadata(id)) + .body(payload.encode()); + } + + private static Map traceMetadata(int id) { + return Map.of( + "traceparent", + "00-" + String.valueOf(id).repeat(32) + "-" + String.valueOf(id).repeat(16) + "-01", + "tracestate", "vendor=synthetic-state-" + id, + "baggage", "token=synthetic-baggage-token-" + id); + } - assertThat(completion.get(10, TimeUnit.SECONDS)).isNull(); - String logs = logsSince(start); - var requestId = Pattern.compile("MCP request completed requestId=([a-f0-9-]+)").matcher(logs); - assertThat(requestId.find()).isTrue(); - assertThat(logs.contains("MCP event details redacted")).isTrue(); - assertThat(logs.contains("MCP command completed requestId=" + requestId.group(1))).isTrue(); - assertThat(logs.contains("synthetic-")).as("Payloads stay out of logs").isFalse(); - assertThat(logs.contains("requestId=null")).isFalse(); + private record PendingCall( + CommandContext context, + Map startMdc, + CompletableFuture> callbackMdc) { + String logRequestId() { + return startMdc.get(McpRequestLogging.REQUEST_ID); } } 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 index efc3bb2a41..e45043aa6a 100644 --- 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 @@ -6,6 +6,7 @@ import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyNoInteractions; @@ -13,6 +14,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import io.micrometer.core.instrument.MeterRegistry; +import io.quarkiverse.mcp.server.ToolResponse; import io.quarkus.security.identity.SecurityIdentity; import io.smallrye.mutiny.Uni; import io.stargate.sgv2.jsonapi.TestConstants; @@ -45,6 +47,8 @@ import java.security.Principal; import java.time.Duration; import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -161,6 +165,50 @@ void closesContextOnCancellation() throws Exception { verify(context, times(1)).close(); } + @Test + void cancelingOneOfTwoSubscriptionsDoesNotCloseOrReplaceTheOther() throws Exception { + var canceledContext = commandContext(true); + var survivingContext = commandContext(true); + var command = keyspaceCommand(); + var canceledResult = new CompletableFuture(); + var survivingResult = new CompletableFuture(); + when(commandProcessor.processCommand(canceledContext, command)) + .thenReturn(Uni.createFrom().completionStage(canceledResult)); + when(commandProcessor.processCommand(survivingContext, command)) + .thenReturn(Uni.createFrom().completionStage(survivingResult)); + var resource = mcpResource(); + var canceledResponse = new CompletableFuture(); + var canceledSubscription = + resource + .processCommand(() -> Uni.createFrom().item(canceledContext), command) + .subscribe() + .with(canceledResponse::complete, canceledResponse::completeExceptionally); + var survivingResponse = + resource + .processCommand(() -> Uni.createFrom().item(survivingContext), command) + .subscribeAsCompletionStage(); + + verify(commandProcessor).processCommand(canceledContext, command); + verify(commandProcessor).processCommand(survivingContext, command); + verify(canceledContext, never()).close(); + verify(survivingContext, never()).close(); + + canceledSubscription.cancel(); + + verify(canceledContext).close(); + verify(survivingContext, never()).close(); + assertThat(canceledResponse).isNotDone(); + assertThat(survivingResponse).isNotDone(); + + var result = CommandResult.statusOnlyBuilder(RequestTracing.NO_OP).build(); + survivingResult.complete(result); + + assertThat(survivingResponse.get(5, TimeUnit.SECONDS).isError()).isFalse(); + assertThat(canceledResponse).isNotDone(); + verify(canceledContext, times(1)).close(); + verify(survivingContext, times(1)).close(); + } + @Test void closeFailureDoesNotReplaceResult() throws Exception { var context = commandContext(true); From d07404c8598c382b426de772e14c44f4b0bd7d21 Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Thu, 5 Nov 2026 21:43:01 +0000 Subject: [PATCH 2/2] test: wait for the MCP 401 completion log line instead of reading it right away --- .../stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java index 38c4cbdc26..a7d953fc92 100644 --- a/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java +++ b/src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java @@ -347,8 +347,13 @@ void traceMetadataCannotAuthenticateAToolCall() throws IOException { assertThat(response.statusCode()).isEqualTo(401); verifyNoInteractions(commandProcessor); assertThat(response.asString()).doesNotContain("synthetic-"); + // The completion line is written by a response end handler, which can run after the client + // has already read the 401, so wait for it rather than reading the log right away. + await() + .atMost(Duration.ofSeconds(5)) + .untilAsserted(() -> assertThat(logsSince(start)).contains("status=401")); var logs = logsSince(start); - assertThat(logs).contains("status=401").doesNotContain("synthetic-"); + assertThat(logs).doesNotContain("synthetic-"); traceMetadata(1) .values() .forEach(