ServerSampling.AdvertisesSamplingWithHandler|Cancellation.CooperativeReadResourceStopsEarly|SamplingCancellationE2E

Author: rianna-parla

docs/api/client.md

@@ -76,6 +76,14 @@ This page summarizes the public client APIs with signatures and brief descriptio
     const JSONValue& includeContext)>
 - void SetSamplingHandler(SamplingHandler handler)
   - Register a handler for server-initiated sampling/createMessage.
+ - using SamplingHandlerCancelable = std::function<std::future<JSONValue>(
+    const JSONValue& messages,
+    const JSONValue& modelPreferences,
+    const JSONValue& systemPrompt,
+    const JSONValue& includeContext,
+    std::stop_token st)>
+ - void SetSamplingHandlerCancelable(SamplingHandlerCancelable handler)
+  - Register a cancelable variant that receives a `std::stop_token`. When set, the client uses this handler for `sampling/createMessage` and will request stop when it receives `notifications/cancelled` targeting the in-flight request id.
 
 ## Notifications, progress, and errors
 - using NotificationHandler = std::function<void(const std::string& method, const JSONValue& params)>;

docs/api/server.md

@@ -182,6 +182,21 @@ Tests:
   - Register a server-side handler for `sampling/createMessage`.
 - std::future<JSONValue> RequestCreateMessage(const CreateMessageParams& params)
   - Request the client to create a message (server-initiated sampling path).
+ - std::future<JSONValue> RequestCreateMessageWithId(const CreateMessageParams& params, const std::string& requestId)
+  - Same as above, but allows the server to specify the JSON-RPC request id. This is useful for end-to-end cancellation tests or workflows where the server needs to target a specific in-flight request with a `notifications/cancelled` message.
+
+Cancellation (server-initiated):
+
+- The server can request cancellation of an in-flight client request by sending:
+
+```cpp
+JSONValue::Object cancelParams; cancelParams["id"] = std::make_shared<JSONValue>(std::string("<request-id>"));
+server.SendNotification(Methods::Cancelled, JSONValue{cancelParams}).get();
+```
+
+- When the client observes this notification, it will propagate `std::stop_token` to a cancelable sampling handler (if registered) and ultimately respond with an error shaped as `{ code: -32603, message: "Cancelled" }`.
+
+Tests: see `tests/test_sampling_cancellation_e2e.cpp`.
 
  ## Keepalive / Heartbeat
  - void SetKeepaliveIntervalMs(const std::optional<int>& intervalMs)

include/mcp/Client.h

@@ -16,6 +16,7 @@
 #include <vector>
 #include <future>
 #include <functional>
+#include <stop_token>
 
 namespace mcp {
 
@@ -249,6 +250,17 @@ class IClient {
         const JSONValue& includeContext)>;
     virtual void SetSamplingHandler(SamplingHandler handler) = 0;
 
+    // Optional cancelable sampling handler variant that receives a std::stop_token.
+    // When set, this handler will be used instead of the non-cancelable variant for
+    // servicing server-initiated sampling/createMessage requests.
+    using SamplingHandlerCancelable = std::function<std::future<JSONValue>(
+        const JSONValue& messages,
+        const JSONValue& modelPreferences,
+        const JSONValue& systemPrompt,
+        const JSONValue& includeContext,
+        std::stop_token st)>;
+    virtual void SetSamplingHandlerCancelable(SamplingHandlerCancelable handler) = 0;
+
     ////////////////////////////////////////// Notification handlers ///////////////////////////////////////////
     //==========================================================================================================
     // Registers a notification handler for a specific method name.
@@ -368,6 +380,7 @@ class Client : public IClient {
                                     const JSONValue& arguments) override;
 
     void SetSamplingHandler(SamplingHandler handler) override;
+    void SetSamplingHandlerCancelable(SamplingHandlerCancelable handler) override;
 
     void SetNotificationHandler(const std::string& method, NotificationHandler handler) override;
     void RemoveNotificationHandler(const std::string& method) override;

include/mcp/Server.h

@@ -323,6 +323,18 @@ class IServer {
     //==========================================================================================================
     virtual std::future<JSONValue> RequestCreateMessage(const CreateMessageParams& params) = 0;
 
+    //==========================================================================================================
+    // Requests the client to create a message (server-initiated sampling) with a caller-provided request id.
+    // Useful for tests and cancellation E2E scenarios where the server needs to cancel an in-flight request by id.
+    // Args:
+    //   params: CreateMessage parameters.
+    //   requestId: Caller-provided id to assign to the JSON-RPC request.
+    // Returns:
+    //   Future with raw JSON result from the client's handler.
+    //==========================================================================================================
+    virtual std::future<JSONValue> RequestCreateMessageWithId(const CreateMessageParams& params,
+                                                             const std::string& requestId) = 0;
+
     /////////////////////////////////////////// Notification sending ///////////////////////////////////////////
     //==========================================================================================================
     // Sends an arbitrary JSON-RPC notification (method + params) to the client.
@@ -504,6 +516,8 @@ class Server : public IServer {
 
     // Server-initiated sampling (request client to create a message)
     std::future<JSONValue> RequestCreateMessage(const CreateMessageParams& params) override;
+    std::future<JSONValue> RequestCreateMessageWithId(const CreateMessageParams& params,
+                                                     const std::string& requestId) override;
 
     // IServer message sending
     std::future<void> SendNotification(const std::string& method, const JSONValue& params) override;

src/mcp/Client.cpp

@@ -7,6 +7,7 @@
 #include <atomic>
 #include <chrono>
 #include <mutex>
+#include <algorithm>
 #include <optional>
 #include <stdexcept>
 #include <thread>
@@ -26,7 +27,7 @@ namespace mcp {
 
 // Client implementation
 class Client::Impl {
-public:
+private:
     friend class Client; // Allow outer Client to invoke private coroutine helpers
     std::unique_ptr<ITransport> transport;
     ClientCapabilities capabilities;
@@ -37,6 +38,7 @@ class Client::Impl {
     IClient::ProgressHandler progressHandler;
     IClient::ErrorHandler errorHandler;
     IClient::SamplingHandler samplingHandler;
+    IClient::SamplingHandlerCancelable samplingHandlerCancelable;
     validation::ValidationMode validationMode{validation::ValidationMode::Off};
 
     // Listings cache (optional)
@@ -47,6 +49,85 @@ class Client::Impl {
     struct TemplatesCache { std::vector<ResourceTemplate> data; std::chrono::steady_clock::time_point ts; bool set{false}; } templatesCache;
     struct PromptsCache { std::vector<Prompt> data; std::chrono::steady_clock::time_point ts; bool set{false}; } promptsCache;
 
+    // Cancellation support for server->client requests (e.g., sampling/createMessage)
+    struct CancellationToken { std::atomic<bool> cancelled{false}; };
+    std::mutex cancelMutex;
+    std::unordered_map<std::string, std::shared_ptr<CancellationToken>> cancelMap;
+    std::unordered_map<std::string, std::vector<std::shared_ptr<std::stop_source>>> stopSources;
+
+    static std::string idToString(const JSONRPCId& id) {
+        std::string idStr;
+        std::visit([&](const auto& v){
+            using T = std::decay_t<decltype(v)>;
+            if constexpr (std::is_same_v<T, std::string>) { idStr = v; }
+            else if constexpr (std::is_same_v<T, int64_t>) { idStr = std::to_string(v); }
+            else { idStr = ""; }
+        }, id);
+        return idStr;
+    }
+
+    static std::string parseIdFromParams(const JSONValue& params) {
+        std::string idStr;
+        if (std::holds_alternative<JSONValue::Object>(params.value)) {
+            const auto& o = std::get<JSONValue::Object>(params.value);
+            auto it = o.find("id");
+            if (it != o.end() && it->second) {
+                if (std::holds_alternative<std::string>(it->second->value)) idStr = std::get<std::string>(it->second->value);
+                else if (std::holds_alternative<int64_t>(it->second->value)) idStr = std::to_string(std::get<int64_t>(it->second->value));
+            }
+        }
+        return idStr;
+    }
+
+    std::shared_ptr<CancellationToken> registerCancelToken(const std::string& idStr) {
+        if (idStr.empty()) return std::make_shared<CancellationToken>();
+        std::lock_guard<std::mutex> lk(cancelMutex);
+        auto it = cancelMap.find(idStr);
+        if (it != cancelMap.end()) return it->second;
+        auto tok = std::make_shared<CancellationToken>();
+        cancelMap[idStr] = tok;
+        return tok;
+    }
+    void unregisterCancelToken(const std::string& idStr) {
+        if (idStr.empty()) return;
+        std::lock_guard<std::mutex> lk(cancelMutex);
+        cancelMap.erase(idStr);
+        stopSources.erase(idStr);
+    }
+    void cancelById(const std::string& idStr) {
+        std::lock_guard<std::mutex> lk(cancelMutex);
+        auto it = cancelMap.find(idStr);
+        if (it == cancelMap.end() || !it->second) {
+            auto tok = std::make_shared<CancellationToken>();
+            tok->cancelled.store(true);
+            cancelMap[idStr] = tok;
+        } else {
+            it->second->cancelled.store(true);
+        }
+        auto itS = stopSources.find(idStr);
+        if (itS != stopSources.end()) {
+            for (auto& src : itS->second) { if (src) { try { src->request_stop(); } catch (...) {} } }
+        }
+    }
+    std::shared_ptr<std::stop_source> registerStopSource(const std::string& idStr) {
+        auto src = std::make_shared<std::stop_source>();
+        std::lock_guard<std::mutex> lk(cancelMutex);
+        stopSources[idStr].push_back(src);
+        auto it = cancelMap.find(idStr);
+        if (it != cancelMap.end() && it->second && it->second->cancelled.load()) {
+            try { src->request_stop(); } catch (...) {}
+        }
+        return src;
+    }
+    void unregisterStopSource(const std::string& idStr, const std::shared_ptr<std::stop_source>& src) {
+        std::lock_guard<std::mutex> lk(cancelMutex);
+        auto it = stopSources.find(idStr);
+        if (it == stopSources.end()) return;
+        auto& vec = it->second;
+        vec.erase(std::remove_if(vec.begin(), vec.end(), [&](const std::shared_ptr<std::stop_source>& p){ return p.get() == src.get(); }), vec.end());
+        if (vec.empty()) stopSources.erase(it);
+    }
+
     explicit Impl(const Implementation& info)
         : clientInfo(info) {
         // Set default client capabilities
@@ -207,6 +288,14 @@ void Client::Impl::onNotification(std::unique_ptr<JSONRPCNotification> n) {
                 const auto& o = std::get<JSONValue::Object>(n->params->value);
                 this->handleProgressNotification(o);
             }
+        } else if (n->method == Methods::Cancelled) {
+            // Server-initiated cancellation for a pending request id
+            if (n->params.has_value()) {
+                std::string idStr = parseIdFromParams(n->params.value());
+                if (!idStr.empty()) {
+                    this->cancelById(idStr);
+                }
+            }
         } else {
             this->invalidateCachesForListChanged(n->method);
             auto it = this->notificationHandlers.find(n->method);
@@ -256,7 +345,13 @@ void Client::Impl::logInvalidCreateMessageResultContext(const JSONValue& result)
 std::unique_ptr<JSONRPCResponse> Client::Impl::onRequest(const JSONRPCRequest& req) {
     try {
         if (req.method == Methods::CreateMessage) {
-            if (!this->samplingHandler) {
+            // Register cancellation and stop_source for this request id
+            const std::string idStr = Impl::idToString(req.id);
+            auto token = this->registerCancelToken(idStr);
+            struct ScopeGuard { std::function<void()> f; ~ScopeGuard(){ if (f) f(); } } guard{ [this, idStr](){ this->unregisterCancelToken(idStr); } };
+            auto src = this->registerStopSource(idStr);
+
+            if (!this->samplingHandler && !this->samplingHandlerCancelable) {
                 errors::McpError e; e.code = JSONRPCErrorCodes::MethodNotAllowed; e.message = "No sampling handler registered";
                 return errors::makeErrorResponse(req.id, e);
             }
@@ -276,8 +371,15 @@ std::unique_ptr<JSONRPCResponse> Client::Impl::onRequest(const JSONRPCRequest& r
                     return errors::makeErrorResponse(req.id, e);
                 }
             }
-            auto fut = this->samplingHandler(messages, modelPreferences, systemPrompt, includeContext);
+            std::future<JSONValue> fut = this->samplingHandler
+                ? this->samplingHandler(messages, modelPreferences, systemPrompt, includeContext)
+                : this->samplingHandlerCancelable(messages, modelPreferences, systemPrompt, includeContext, src->get_token());
             JSONValue result = fut.get();
+            // If cancelled while or after handler ran, return Cancelled
+            if (token && token->cancelled.load()) {
+                errors::McpError e; e.code = JSONRPCErrorCodes::InternalError; e.message = "Cancelled";
+                return errors::makeErrorResponse(req.id, e);
+            }
             if (this->validationMode == validation::ValidationMode::Strict) {
                 if (!validation::validateCreateMessageResultJson(result)) {
                     this->logInvalidCreateMessageResultContext(result);
@@ -1035,7 +1137,7 @@ mcp::async::Task<JSONValue> Client::Impl::coGetPrompt(const std::string& name, c
 
 
 Client::Client(const Implementation& clientInfo)
-    : pImpl(std::make_unique<Impl>(clientInfo)) {
+    : pImpl(std::unique_ptr<Impl>(new Impl(clientInfo))) {
     FUNC_SCOPE();
 }
 
@@ -1176,6 +1278,11 @@ void Client::SetSamplingHandler(SamplingHandler handler) {
     pImpl->samplingHandler = std::move(handler);
 }
 
+void Client::SetSamplingHandlerCancelable(SamplingHandlerCancelable handler) {
+    FUNC_SCOPE();
+    pImpl->samplingHandlerCancelable = std::move(handler);
+}
+
 void Client::SetErrorHandler(ErrorHandler handler) {
     FUNC_SCOPE();
     pImpl->errorHandler = std::move(handler);

src/mcp/Server.cpp

@@ -30,7 +30,7 @@ namespace mcp {
 
 // Server implementation
 class Server::Impl {
-public:
+private:
     friend class Server;
     // Coroutine helpers (declarations)
     mcp::async::Task<void> coStart(std::unique_ptr<ITransport> transport);
@@ -40,6 +40,7 @@ class Server::Impl {
     mcp::async::Task<JSONValue> coReadResource(const std::string& uri);
     mcp::async::Task<JSONValue> coGetPrompt(const std::string& name, const JSONValue& arguments);
     mcp::async::Task<JSONValue> coRequestCreateMessage(const CreateMessageParams& params);
+    mcp::async::Task<JSONValue> coRequestCreateMessageWithId(const CreateMessageParams& params, const std::string& requestId);
     mcp::async::Task<void> coSendNotification(const std::string& method, const JSONValue& params);
     mcp::async::Task<void> coNotifyResourcesListChanged();
     mcp::async::Task<void> coNotifyToolsListChanged();
@@ -1277,8 +1278,64 @@ mcp::async::Task<JSONValue> Server::Impl::coRequestCreateMessage(const CreateMes
     co_return JSONValue{};
 }
 
+mcp::async::Task<JSONValue> Server::Impl::coRequestCreateMessageWithId(const CreateMessageParams& params, const std::string& requestId) {
+    FUNC_SCOPE();
+    if (!this->transport) {
+        LOG_ERROR("RequestCreateMessageWithId called without transport");
+        co_return JSONValue{};
+    }
+    auto request = std::make_unique<JSONRPCRequest>();
+    request->method = Methods::CreateMessage;
+    request->id = requestId; // assign caller-provided id so cancellation can target this request
+    JSONValue::Object obj;
+    // messages (array)
+    JSONValue::Array msgs;
+    msgs.reserve(params.messages.size());
+    for (const auto& m : params.messages) msgs.push_back(std::make_shared<JSONValue>(m));
+    obj["messages"] = std::make_shared<JSONValue>(msgs);
+    // Optional fields
+    if (params.modelPreferences.has_value()) {
+        obj["modelPreferences"] = std::make_shared<JSONValue>(params.modelPreferences.value());
+    }
+    if (params.systemPrompt.has_value()) {
+        obj["systemPrompt"] = std::make_shared<JSONValue>(params.systemPrompt.value());
+    }
+    if (params.includeContext.has_value()) {
+        obj["includeContext"] = std::make_shared<JSONValue>(params.includeContext.value());
+    }
+    if (params.maxTokens.has_value()) {
+        obj["maxTokens"] = std::make_shared<JSONValue>(static_cast<int64_t>(params.maxTokens.value()));
+    }
+    if (params.temperature.has_value()) {
+        obj["temperature"] = std::make_shared<JSONValue>(params.temperature.value());
+    }
+    if (params.stopSequences.has_value()) {
+        JSONValue::Array arr;
+        for (const auto& s : params.stopSequences.value()) {
+            arr.push_back(std::make_shared<JSONValue>(s));
+        }
+        obj["stopSequences"] = std::make_shared<JSONValue>(arr);
+    }
+    if (params.metadata.has_value()) {
+        obj["metadata"] = std::make_shared<JSONValue>(params.metadata.value());
+    }
+    request->params = JSONValue{obj};
+
+    auto fut = this->transport->SendRequest(std::move(request));
+    try {
+        auto resp = co_await mcp::async::makeFutureAwaitable(std::move(fut));
+        if (resp) {
+            if (resp->result.has_value()) co_return resp->result.value();
+            if (resp->error.has_value()) co_return resp->error.value();
+        }
+    } catch (const std::exception& e) {
+        LOG_ERROR("RequestCreateMessageWithId exception: {}", e.what());
+    }
+    co_return JSONValue{};
+}
+
 Server::Server(const std::string& serverInfo)
-    : pImpl(std::make_unique<Impl>()) {
+    : pImpl(std::unique_ptr<Impl>(new Impl())) {
     FUNC_SCOPE();
     pImpl->serverInfo = serverInfo;
 }
@@ -1637,6 +1694,11 @@ std::future<JSONValue> Server::RequestCreateMessage(const CreateMessageParams& p
     return pImpl->coRequestCreateMessage(params).toFuture();
 }
 
+std::future<JSONValue> Server::RequestCreateMessageWithId(const CreateMessageParams& params, const std::string& requestId) {
+    FUNC_SCOPE();
+    return pImpl->coRequestCreateMessageWithId(params, requestId).toFuture();
+}
+
 void Server::RegisterResourceTemplate(const ResourceTemplate& resourceTemplate) {
     FUNC_SCOPE();
     LOG_DEBUG("Registering resource template: {}", resourceTemplate.uriTemplate);

tests/CMakeLists.txt

@@ -41,6 +41,7 @@ add_executable(mcp_tests
   test_initialize_notifications.cpp
   test_errors.cpp
   test_client_typed.cpp
+  test_sampling_cancellation_e2e.cpp
   test_validation_mode.cpp
   test_validation_lists.cpp
   test_validation_sampling.cpp