test: verify MCP request isolation and document tracing support

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.

src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpLoggingTest.java

@@ -5,16 +5,20 @@
 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;
 import io.quarkus.test.InjectMock;
 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,215 @@ 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<Integer>();
-      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<Integer>();
+    var secondCompletion = new CompletableFuture<Integer>();
+    var completions =
+        Map.of("synthetic-token-1", firstCompletion, "synthetic-token-2", secondCompletion);
+    var started = new LinkedBlockingQueue<PendingCall>();
+    when(commandProcessor.processCommand(any(), any()))
+        .thenAnswer(
+            invocation -> {
+              CommandContext<?> context = invocation.getArgument(0);
+              var callbackMdc = new CompletableFuture<Map<String, String>>();
+              var call = new PendingCall(context, MDC.getCopyOfContextMap(), callbackMdc);
+              return Uni.createFrom()
+                  .<Integer>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}");
+
+        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);
 
-      callProviderTool();
+        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-");
+    // 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).doesNotContain("synthetic-");
+    traceMetadata(1)
+        .values()
+        .forEach(
+            value -> {
+              assertThat(response.asString()).doesNotContain(value);
+              assertThat(logs).doesNotContain(value);
+            });
+  }
+
+  private static PendingCall awaitStarted(
+      LinkedBlockingQueue<PendingCall> started, Future<Response> 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<String, String> 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<String, String> startMdc,
+      CompletableFuture<Map<String, String>> callbackMdc) {
+    String logRequestId() {
+      return startMdc.get(McpRequestLogging.REQUEST_ID);
     }
   }
 

src/test/java/io/stargate/sgv2/jsonapi/api/v1/mcp/McpResourceTest.java

@@ -6,13 +6,15 @@
 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;
 import static org.mockito.Mockito.when;
 
 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<CommandResult>();
+    var survivingResult = new CompletableFuture<CommandResult>();
+    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<ToolResponse>();
+    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);