feat: Push based cancellation for diagnostic services (#236)

Signed-off-by: David Kornel <kornys@outlook.com>

Author: kornys

AGENTS.md

@@ -408,7 +408,7 @@ public class DiagnosticTools {
 - **Graceful degradation**: Works without Sampling/Elicitation support (gathers everything, returns raw data)
 - **Step failure resilience**: Individual step failures are recorded in `stepsFailed`, workflow continues
 - **Progress tracking**: Reports progress via `DiagnosticHelper.sendProgress()`
-- **Cancellation**: Checks `DiagnosticHelper.checkCancellation()` between steps
+- **Cancellation**: Supports both poll-based (`checkCancellation()` between steps) and push-based (`registerCancellationCallback()` for async operations)
 - **No duplication**: Calls existing domain services (KafkaService, StrimziOperatorService, etc.)
 - **Configurable token limits**: `mcp.sampling.triage-max-tokens` and `mcp.sampling.analysis-max-tokens`
 - **OpenTelemetry tracing**: Gather/triage/analysis methods are annotated with `@WithSpan` for
@@ -420,12 +420,50 @@ public class DiagnosticTools {
 
 Shared MCP framework utilities in `common/src/.../service/DiagnosticHelper.java`:
 - `sendProgress(Progress, step, totalSteps, message)` — progress update with descriptive message to MCP client
-- `checkCancellation(Cancellation)` — abort if client cancelled
+- `checkCancellation(Cancellation)` — poll-based cancellation check; throws if client cancelled (use for synchronous code paths)
+- `registerCancellationCallback(Cancellation, AtomicBoolean)` — push-based cancellation callback; immediately sets flag when cancelled (use for async operations like Sampling/Elicitation)
+- `checkAsyncCancellation(AtomicBoolean)` — check flag set by callback and throw if cancelled (use after async operations)
 - `putIfNotNull(Map, String, Object)` — conditional map insertion
 - `elicitSelection(Elicitation, message, propertyName, description, options)` — generic single-select Elicitation
 - `extractSamplingText(SamplingResponse)` — safe text extraction from Sampling response
 - `MAP_TYPE_REF` — reusable `TypeReference<Map<String, Object>>` for JSON parsing
 
+#### Cancellation Patterns
+
+Diagnostic services support two cancellation mechanisms:
+
+1. **Poll-based** (`checkCancellation()`) — For synchronous operations
+   - Checks if cancellation occurred since the last check
+   - Throws immediately if cancelled
+   - Use between discrete workflow steps (after gathering data, before next phase)
+
+2. **Push-based** (`registerCancellationCallback()`) — For async operations
+   - Registers a callback that fires immediately when cancellation occurs
+   - Sets an AtomicBoolean flag that can be checked at any point
+   - Suitable for Sampling/Elicitation calls that may take seconds to minutes
+   - Catches cancellation during the async call, not just at poll points
+
+**When to use each:**
+- Use **poll-based** for synchronous code paths (gathering data, parsing responses, building reports)
+- Use **push-based** for async operations (Sampling, Elicitation, long-running external API calls)
+- Can mix both — poll-based for sync phases, push-based for async phases
+
+**Example usage:**
+```java
+// Register callback once at start of diagnose()
+AtomicBoolean cancelled = new AtomicBoolean(false);
+DiagnosticHelper.registerCancellationCallback(cancellation, cancelled);
+
+// Check after async operations (sampling/elicitation)
+InvestigationAreas areas = performTriage(sampling, ...);
+DiagnosticHelper.checkAsyncCancellation(cancelled);  // throws if cancelled
+
+// Continue with poll-based checks for sync phases
+DiagnosticHelper.checkCancellation(cancellation);
+```
+
+See `KafkaClusterDiagnosticService` for complete reference implementation.
+
 ### BaseDiagnosticService (common module)
 
 Abstract base class in `common/src/.../service/BaseDiagnosticService.java` that all
@@ -435,8 +473,14 @@ diagnostic services must extend:
 - Binds `mcp.sampling.analysis-max-tokens` → `analysisMaxTokens`
 - Binds `mcp.log.tail-lines` → `defaultTailLines`
 - `performSampling(sampling, systemPrompt, data, maxTokens)` — sends a Sampling request and returns extracted text
+- `performSampling(sampling, systemPrompt, data, maxTokens, cancelled)` — overload with push-based cancellation support
 - `performAnalysis(sampling, systemPrompt, fullData)` — convenience wrapper using `analysisMaxTokens`
+- `performAnalysis(sampling, systemPrompt, fullData, cancelled)` — overload with push-based cancellation support
 - `performTriage(sampling, systemPrompt, phase1Summary)` — sends triage Sampling, parses JSON response to Map
+- `performTriage(sampling, systemPrompt, phase1Summary, cancelled)` — overload with push-based cancellation support
+
+All cancellation-aware overloads accept an optional `AtomicBoolean cancelled` parameter. If non-null,
+the flag is checked before and after the async `sendAndAwait()` call, returning null if cancelled.
 
 ### NamespaceElicitationHelper (common module)
 

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- **Push-based cancellation for diagnostics** — Added `DiagnosticHelper.registerCancellationCallback()` for async diagnostic operations (sampling, elicitation). The callback approach catches cancellation immediately via `Cancellation#onCancelled()` instead of polling with `skipProcessingIfCancelled()`, preventing wasted work during long-running LLM calls. Existing `checkCancellation()` remains for synchronous code paths.
 - **Elasticsearch/OpenSearch log provider** (`mcp.log.provider=streamshub-elasticsearch`) for querying logs from Elasticsearch or OpenSearch
 - **Dev environment manifests and scripts** for deploying Elasticsearch with Fluent Bit on Kind clusters and ECK Operator on OpenShift with automated API key provisioning, ClusterLogForwarder integration, and passthrough TLS routes
 - **System tests** for Elasticsearch log provider covering log collection, field mapping, time window queries, and error handling

common/src/main/java/io/streamshub/mcp/common/service/BaseDiagnosticService.java

@@ -13,6 +13,7 @@
 import org.jboss.logging.Logger;
 
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Base class for diagnostic service implementations.
@@ -53,9 +54,32 @@ protected BaseDiagnosticService() {
      */
     protected String performSampling(Sampling sampling, String systemPrompt,
                                       Map<String, Object> data, int maxTokens) {
+        return performSampling(sampling, systemPrompt, data, maxTokens, null);
+    }
+
+    /**
+     * Send a Sampling request with cancellation support.
+     *
+     * <p>Checks the cancellation flag before and after the async sampling call.
+     * Returns {@code null} if Sampling is unavailable, unsupported, or fails,
+     * or if cancelled before the call completes.</p>
+     *
+     * @param sampling     MCP Sampling interface (may be null or unsupported)
+     * @param systemPrompt the system prompt for the LLM
+     * @param data         the data map to serialize as JSON input
+     * @param maxTokens    maximum tokens for the response
+     * @param cancelled    optional flag set by push-based cancellation callback (may be null)
+     * @return the extracted text, or null on failure or cancellation
+     */
+    protected String performSampling(Sampling sampling, String systemPrompt,
+                                      Map<String, Object> data, int maxTokens,
+                                      AtomicBoolean cancelled) {
         if (sampling == null || !sampling.isSupported()) {
             return null;
         }
+        if (cancelled != null && cancelled.get()) {
+            return null;
+        }
         try {
             String dataJson = objectMapper.writeValueAsString(data);
             SamplingResponse response = sampling.requestBuilder()
@@ -64,6 +88,9 @@ protected String performSampling(Sampling sampling, String systemPrompt,
                 .setMaxTokens(maxTokens)
                 .build()
                 .sendAndAwait();
+            if (cancelled != null && cancelled.get()) {
+                return null;
+            }
             return DiagnosticHelper.extractSamplingText(response);
         } catch (Exception e) {
             getLogger().warnf("Sampling failed: %s: %s",
@@ -87,6 +114,24 @@ protected String performAnalysis(Sampling sampling, String systemPrompt,
         return performSampling(sampling, systemPrompt, fullData, analysisMaxTokens);
     }
 
+    /**
+     * Perform analysis Sampling with cancellation support.
+     *
+     * <p>Convenience wrapper that uses {@link #analysisMaxTokens} and checks
+     * the cancellation flag before and after the sampling call.</p>
+     *
+     * @param sampling     MCP Sampling interface
+     * @param systemPrompt the analysis system prompt
+     * @param fullData     the full gathered data map
+     * @param cancelled    optional flag set by push-based cancellation callback (may be null)
+     * @return the analysis text, or null on failure or cancellation
+     */
+    protected String performAnalysis(Sampling sampling, String systemPrompt,
+                                      Map<String, Object> fullData,
+                                      AtomicBoolean cancelled) {
+        return performSampling(sampling, systemPrompt, fullData, analysisMaxTokens, cancelled);
+    }
+
     /**
      * Perform triage Sampling and return the parsed JSON response as a map.
      *
@@ -101,7 +146,27 @@ protected String performAnalysis(Sampling sampling, String systemPrompt,
      */
     protected Map<String, Object> performTriage(Sampling sampling, String systemPrompt,
                                                  Map<String, Object> phase1Summary) {
-        String text = performSampling(sampling, systemPrompt, phase1Summary, triageMaxTokens);
+        return performTriage(sampling, systemPrompt, phase1Summary, null);
+    }
+
+    /**
+     * Perform triage Sampling with cancellation support.
+     *
+     * <p>Checks the cancellation flag before and after the sampling call.
+     * Returns {@code null} if Sampling is unavailable, the response cannot
+     * be parsed as JSON, or the operation was cancelled. Callers should fall
+     * back to investigating all areas when null is returned.</p>
+     *
+     * @param sampling      MCP Sampling interface
+     * @param systemPrompt  the triage system prompt
+     * @param phase1Summary the Phase 1 summary data map
+     * @param cancelled     optional flag set by push-based cancellation callback (may be null)
+     * @return the parsed triage response, or null on failure or cancellation
+     */
+    protected Map<String, Object> performTriage(Sampling sampling, String systemPrompt,
+                                                 Map<String, Object> phase1Summary,
+                                                 AtomicBoolean cancelled) {
+        String text = performSampling(sampling, systemPrompt, phase1Summary, triageMaxTokens, cancelled);
         if (text == null) {
             return null;
         }

common/src/main/java/io/streamshub/mcp/common/service/DiagnosticHelper.java

@@ -11,10 +11,12 @@
 import io.quarkiverse.mcp.server.ElicitationResponse;
 import io.quarkiverse.mcp.server.Progress;
 import io.quarkiverse.mcp.server.SamplingResponse;
+import io.quarkiverse.mcp.server.ToolCallException;
 import org.jboss.logging.Logger;
 
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
  * Shared utilities for diagnostic service implementations.
@@ -66,6 +68,51 @@ public static void checkCancellation(final Cancellation cancellation) {
         }
     }
 
+    /**
+     * Register a push-based cancellation callback for async operations.
+     *
+     * <p>This method sets up a callback that immediately fires when the MCP client
+     * cancels the operation, setting the provided AtomicBoolean to true. This is
+     * more responsive than polling with {@link #checkCancellation(Cancellation)}
+     * and is suitable for async operations like Sampling and Elicitation that may
+     * take seconds to minutes.</p>
+     *
+     * <p>The callback does not throw — it only sets the flag. Callers should check
+     * the flag before and after expensive operations using {@link #checkAsyncCancellation(AtomicBoolean)}.</p>
+     *
+     * <p>Use {@link #checkCancellation(Cancellation)} for synchronous code paths
+     * where polling is sufficient. Use this method for async operations where
+     * immediate notification is important to avoid wasted work.</p>
+     *
+     * @param cancellation  the MCP cancellation (may be null)
+     * @param cancelledFlag the flag to set when cancellation occurs (must not be null if cancellation is non-null)
+     */
+    public static void registerCancellationCallback(final Cancellation cancellation,
+                                                     final AtomicBoolean cancelledFlag) {
+        if (cancellation != null) {
+            cancellation.onCancelled(reason -> {
+                LOG.infof("Cancellation received: %s", reason.orElse("no reason provided"));
+                cancelledFlag.set(true);
+            });
+        }
+    }
+
+    /**
+     * Check if the operation was cancelled during an async operation (sampling/elicitation).
+     *
+     * <p>This method checks the flag set by {@link #registerCancellationCallback(Cancellation, AtomicBoolean)}
+     * and throws if cancellation has occurred. Use this after async operations like Sampling
+     * to detect cancellation that occurred during the operation.</p>
+     *
+     * @param cancelledFlag the flag set by the cancellation callback
+     * @throws ToolCallException if the operation was cancelled
+     */
+    public static void checkAsyncCancellation(final AtomicBoolean cancelledFlag) {
+        if (cancelledFlag != null && cancelledFlag.get()) {
+            throw new ToolCallException("Operation cancelled");
+        }
+    }
+
     /**
      * Put a value into a map only if it is not null.
      *
@@ -98,6 +145,32 @@ public static String elicitSelection(final Elicitation elicitation,
                                           final String propertyName,
                                           final String description,
                                           final List<String> options) {
+        return elicitSelection(elicitation, message, propertyName, description, options, null);
+    }
+
+    /**
+     * Ask the user to select a single value from a list via MCP Elicitation with cancellation support.
+     *
+     * <p>This is a generic Elicitation wrapper usable for any disambiguation
+     * (namespace selection, cluster selection, etc.).</p>
+     *
+     * @param elicitation  the MCP Elicitation interface
+     * @param message      the prompt message shown to the user
+     * @param propertyName the schema property name (e.g., "namespace")
+     * @param description  the property description
+     * @param options      the list of options to choose from
+     * @param cancelled    optional flag set by push-based cancellation callback (may be null)
+     * @return the selected value, or null if the user declined, elicitation failed, or operation was cancelled
+     */
+    public static String elicitSelection(final Elicitation elicitation,
+                                          final String message,
+                                          final String propertyName,
+                                          final String description,
+                                          final List<String> options,
+                                          final AtomicBoolean cancelled) {
+        if (cancelled != null && cancelled.get()) {
+            return null;
+        }
         try {
             ElicitationResponse response = elicitation.requestBuilder()
                 .setMessage(message)
@@ -109,6 +182,9 @@ public static String elicitSelection(final Elicitation elicitation,
                 .build()
                 .sendAndAwait();
 
+            if (cancelled != null && cancelled.get()) {
+                return null;
+            }
             if (response.actionAccepted()) {
                 return response.content().getString(propertyName);
             }

common/src/test/java/io/streamshub/mcp/common/service/DiagnosticHelperTest.java

@@ -0,0 +1,96 @@
+/*
+ * Copyright StreamsHub authors.
+ * License: Apache License 2.0 (see the file LICENSE or http://apache.org/licenses/LICENSE-2.0.html).
+ */
+package io.streamshub.mcp.common.service;
+
+import io.quarkiverse.mcp.server.Cancellation;
+import io.quarkiverse.mcp.server.ToolCallException;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+/**
+ * Unit tests for {@link DiagnosticHelper} cancellation callback registration.
+ */
+class DiagnosticHelperTest {
+
+    DiagnosticHelperTest() {
+    }
+
+    @Test
+    void registerCancellationCallback_withNullCancellation_doesNotThrow() {
+        AtomicBoolean cancelled = new AtomicBoolean(false);
+        assertDoesNotThrow(() -> DiagnosticHelper.registerCancellationCallback(null, cancelled));
+        assertFalse(cancelled.get(), "Flag should remain false when cancellation is null");
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    void registerCancellationCallback_withValidCancellation_setsFlag() {
+        AtomicBoolean cancelled = new AtomicBoolean(false);
+        Cancellation cancellation = mock(Cancellation.class);
+
+        ArgumentCaptor<Consumer<Optional<String>>> captor =
+            ArgumentCaptor.forClass(Consumer.class);
+
+        DiagnosticHelper.registerCancellationCallback(cancellation, cancelled);
+
+        verify(cancellation).onCancelled(captor.capture());
+        assertFalse(cancelled.get(), "Flag should be false before cancellation fires");
+
+        captor.getValue().accept(Optional.empty());
+
+        assertTrue(cancelled.get(), "Flag should be true after cancellation fires");
+    }
+
+    @Test
+    @SuppressWarnings("unchecked")
+    void registerCancellationCallback_withCancellationReason_setsFlag() {
+        AtomicBoolean cancelled = new AtomicBoolean(false);
+        Cancellation cancellation = mock(Cancellation.class);
+
+        ArgumentCaptor<Consumer<Optional<String>>> captor =
+            ArgumentCaptor.forClass(Consumer.class);
+
+        DiagnosticHelper.registerCancellationCallback(cancellation, cancelled);
+
+        verify(cancellation).onCancelled(captor.capture());
+
+        captor.getValue().accept(Optional.of("User cancelled operation"));
+
+        assertTrue(cancelled.get(), "Flag should be true after cancellation with reason");
+    }
+
+    @Test
+    void checkAsyncCancellation_withNullFlag_doesNotThrow() {
+        assertDoesNotThrow(() -> DiagnosticHelper.checkAsyncCancellation(null));
+    }
+
+    @Test
+    void checkAsyncCancellation_withFalseFlag_doesNotThrow() {
+        AtomicBoolean cancelled = new AtomicBoolean(false);
+        assertDoesNotThrow(() -> DiagnosticHelper.checkAsyncCancellation(cancelled));
+    }
+
+    @Test
+    void checkAsyncCancellation_withTrueFlag_throws() {
+        AtomicBoolean cancelled = new AtomicBoolean(true);
+        Exception exception = assertThrows(
+            ToolCallException.class,
+            () -> DiagnosticHelper.checkAsyncCancellation(cancelled)
+        );
+        assertEquals("Operation cancelled", exception.getMessage());
+    }
+}