feat(mcp): preemptive task cancellation via duckdb_interrupt (#111)

src/include/mcp_task_manager.hpp

@@ -41,10 +41,19 @@ class MCPTaskManager {
         std::string error_message; // when Failed
     };
 
+    // A running task registers an interrupt closure through this callback so the
+    // manager can stop it preemptively (issue #111). The closure is invoked from
+    // another thread by cancel()/shutdown(); it must be safe to call concurrently
+    // with the work and after the work has finished (a no-op then). Typically it
+    // interrupts the DuckDB query executing on the worker thread.
+    using SetInterrupt = std::function<void(std::function<void()>)>;
+
     // The unit of work a submitted task runs: it returns the serialized MCP
     // tool result envelope (the same string a synchronous tools/call produces).
-    // It should honour `cancelled` cooperatively where possible.
-    using Work = std::function<std::string(const std::atomic<bool>& cancelled)>;
+    // It should honour `cancelled` cooperatively where possible, and may call
+    // `set_interrupt` once it has a preemptive cancellation hook to expose.
+    using Work = std::function<std::string(const std::atomic<bool>& cancelled,
+                                           const SetInterrupt& set_interrupt)>;
 
     // Optional durability: a callback that runs a SQL statement and returns its
     // rows (each a column->string map). When provided, tasks are persisted to a
@@ -80,6 +89,10 @@ class MCPTaskManager {
         Task task;
         std::atomic<bool> cancelled{false};
         Work work;
+        // Preemptive-cancel hook, set by the running work and cleared when it
+        // finishes; guarded by the manager mutex (mu_). Empty when the task is
+        // queued or already terminal.
+        std::function<void()> interrupt;
     };
 
     void workerLoop();

src/include/query_executor.hpp

@@ -5,6 +5,7 @@
 #include <stdexcept>
 #include <string>
 #include <map>
+#include <thread>
 #include <vector>
 
 #include "prepared_template_rewriter.hpp"
@@ -22,6 +23,18 @@ class BadRequestError : public std::runtime_error {
     explicit BadRequestError(const std::string& msg) : std::runtime_error(msg) {}
 };
 
+class QueryExecutor;
+
+// Active-query registry (issue #111). A QueryExecutor registers itself under the
+// thread running its query for the duration of that query; another thread can
+// then interrupt whatever query is running on a given thread by its id. This
+// lets the MCP Tasks worker make tasks/cancel and shutdown preemptive without
+// threading an executor handle through the whole tool-execution pipeline. All
+// operations are mutex-guarded, and interrupt() only touches a live executor.
+void registerActiveExecutor(std::thread::id tid, QueryExecutor* exec);
+void unregisterActiveExecutor(std::thread::id tid);
+void interruptActiveExecutor(std::thread::id tid);
+
 } // namespace flapi
 
 // Define standard vector size if not defined by DuckDB
@@ -115,6 +128,13 @@ class QueryExecutor {
 
     crow::json::wvalue toJson() const;
 
+    // Interrupt an in-flight query on this executor's connection from another
+    // thread. Used by the MCP Tasks extension (issue #111) to make tasks/cancel
+    // and graceful shutdown preemptive: the running duckdb_query returns an
+    // error, which execute() then surfaces as a thrown exception. Safe to call
+    // concurrently with execute(); a no-op if there is no live connection.
+    void interrupt() { if (conn) { duckdb_interrupt(conn); } }
+
     // Make these public since they're used directly in DatabaseManager
     duckdb_connection conn;
     mutable duckdb_result result;

src/mcp_route_handlers.cpp

@@ -4,7 +4,9 @@
 #include "mcp_authorization_policy.hpp"
 #include "mcp_schema_builder.hpp"
 #include "mcp_header_validation.hpp"
+#include "query_executor.hpp"
 #include <iostream>
+#include <thread>
 #include <sstream>
 #include <optional>
 
@@ -1505,7 +1507,16 @@ MCPResponse MCPRouteHandlers::handleToolsCallRequest(const MCPRequest& request,
                     // synchronous call would (content + structuredContent, or an
                     // isError text block on failure).
                     MCPToolHandler* handler = tool_handler_.get();
-                    auto work = [handler, tool_request](const std::atomic<bool>&) -> std::string {
+                    auto work = [handler, tool_request](
+                                    const std::atomic<bool>&,
+                                    const MCPTaskManager::SetInterrupt& set_interrupt) -> std::string {
+                        // Publish a preemptive-cancel hook: tasks/cancel and
+                        // shutdown interrupt the DuckDB query running on THIS
+                        // worker thread (the executor self-registers under this
+                        // id for the duration of its query — see issue #111).
+                        set_interrupt([tid = std::this_thread::get_id()]() {
+                            interruptActiveExecutor(tid);
+                        });
                         auto r = handler->executeTool(tool_request);
                         mcp::ContentResponse cr;
                         if (r.success) {

src/mcp_task_manager.cpp

@@ -227,16 +227,26 @@ void MCPTaskManager::workerLoop() {
             continue;
         }
 
+        // Let the work publish a preemptive-cancel hook (e.g. an interrupt of the
+        // DuckDB query it is about to run). Stored under mu_ so cancel()/shutdown()
+        // on another thread see a consistent value.
+        auto set_interrupt = [this, entry](std::function<void()> fn) {
+            std::lock_guard<std::mutex> lock(mu_);
+            entry->interrupt = std::move(fn);
+        };
+
         std::string result;
         std::string error;
         try {
-            result = entry->work(entry->cancelled);
+            result = entry->work(entry->cancelled, set_interrupt);
         } catch (const std::exception& e) {
             error = e.what();
         }
 
         {
             std::unique_lock<std::mutex> lock(mu_);
+            // The query has returned; a late cancel() must not fire a stale hook.
+            entry->interrupt = nullptr;
             if (entry->cancelled.load()) {
                 entry->task.status = Status::Cancelled;
             } else if (!error.empty()) {
@@ -287,11 +297,11 @@ bool MCPTaskManager::cancel(const std::string& task_id, const std::string& princ
         return false;
     }
     it->second->cancelled.store(true);
-    // If it has not started yet, mark it cancelled immediately; a running task
-    // is marked cancelled by the worker when its work returns.
-    if (it->second->task.status == Status::Working) {
-        // Leave running tasks Working until the worker observes cancellation;
-        // queued-but-not-started tasks are flipped by the worker's pre-check.
+    // Preemptively interrupt a running query if the task published a hook;
+    // otherwise the worker flips a queued task on its pre-check, and a running
+    // task without a hook is caught cooperatively when its work returns.
+    if (it->second->interrupt) {
+        it->second->interrupt();
     }
     return true;
 }
@@ -303,9 +313,14 @@ void MCPTaskManager::shutdown() {
             return;
         }
         stopping_ = true;
-        // Signal cancellation to everything so in-flight work can bail out.
+        // Signal cancellation to everything so in-flight work can bail out, and
+        // preemptively interrupt any running query so shutdown does not block on
+        // a multi-minute task finishing on its own.
         for (auto& [id, entry] : tasks_) {
             entry->cancelled.store(true);
+            if (entry->interrupt) {
+                entry->interrupt();
+            }
         }
     }
     cv_.notify_all();

src/query_executor.cpp

@@ -12,8 +12,52 @@
 // duckdb::Vector*, so dropping to Vector::GetValue() is a safe round-trip.
 #include "duckdb.hpp"
 
+#include <mutex>
+#include <unordered_map>
+
 namespace flapi {
 
+namespace {
+// thread::id -> the QueryExecutor currently running a query on that thread.
+std::mutex& activeExecMutex() {
+    static std::mutex m;
+    return m;
+}
+std::unordered_map<std::thread::id, QueryExecutor*>& activeExecMap() {
+    static std::unordered_map<std::thread::id, QueryExecutor*> m;
+    return m;
+}
+} // namespace
+
+void registerActiveExecutor(std::thread::id tid, QueryExecutor* exec) {
+    std::lock_guard<std::mutex> lock(activeExecMutex());
+    activeExecMap()[tid] = exec;
+}
+void unregisterActiveExecutor(std::thread::id tid) {
+    std::lock_guard<std::mutex> lock(activeExecMutex());
+    activeExecMap().erase(tid);
+}
+void interruptActiveExecutor(std::thread::id tid) {
+    // Hold the lock across interrupt() so the executor cannot be unregistered
+    // and destroyed between the lookup and the duckdb_interrupt call.
+    std::lock_guard<std::mutex> lock(activeExecMutex());
+    auto it = activeExecMap().find(tid);
+    if (it != activeExecMap().end() && it->second) {
+        it->second->interrupt();
+    }
+}
+
+namespace {
+// RAII: publish `exec` under the current thread for the duration of a query.
+struct ActiveExecGuard {
+    std::thread::id tid;
+    explicit ActiveExecGuard(QueryExecutor* exec) : tid(std::this_thread::get_id()) {
+        registerActiveExecutor(tid, exec);
+    }
+    ~ActiveExecGuard() { unregisterActiveExecutor(tid); }
+};
+} // namespace
+
 QueryExecutor::QueryExecutor(duckdb_database db) : has_result(false) {
     if (duckdb_connect(db, &conn) == DuckDBError) {
         throw std::runtime_error("Failed to create database connection");
@@ -34,7 +78,14 @@ void QueryExecutor::execute(const std::string& query, const std::string& context
         has_result = false;
     }
 
-    if (duckdb_query(conn, query.c_str(), &result) == DuckDBError) {
+    duckdb_state qstate;
+    {
+        // Publish this executor for the running thread so another thread can
+        // interrupt it (issue #111); unpublished the moment the query returns.
+        ActiveExecGuard guard(this);
+        qstate = duckdb_query(conn, query.c_str(), &result);
+    }
+    if (qstate == DuckDBError) {
         std::string error_message = duckdb_result_error(&result);
         std::string context_msg = context.empty() ? "" : " during " + context;
         duckdb_destroy_result(&result);
@@ -49,7 +100,12 @@ void QueryExecutor::executePrepared(duckdb_prepared_statement stmt, const std::s
         has_result = false;
     }
 
-    if (duckdb_execute_prepared(stmt, &result) == DuckDBError) {
+    duckdb_state pstate;
+    {
+        ActiveExecGuard guard(this);
+        pstate = duckdb_execute_prepared(stmt, &result);
+    }
+    if (pstate == DuckDBError) {
         std::string error_message = duckdb_result_error(&result);
         std::string context_msg = context.empty() ? "" : " during " + context;
         duckdb_destroy_result(&result);

test/cpp/mcp_task_manager_test.cpp

@@ -30,7 +30,9 @@ bool waitFor(MCPTaskManager& m, const std::string& id, const std::string& princi
 TEST_CASE("MCPTaskManager: a submitted task runs and completes", "[mcp][tasks]") {
     MCPTaskManager m(2, 16);
     auto id = m.submit("tool", "alice", 60000, 100,
-                       [](const std::atomic<bool>&) { return std::string("{\"ok\":true}"); });
+                       [](const std::atomic<bool>&, const MCPTaskManager::SetInterrupt&) {
+                           return std::string("{\"ok\":true}");
+                       });
     REQUIRE_FALSE(id.empty());
     REQUIRE(waitFor(m, id, "alice", Status::Completed));
 
@@ -43,7 +45,9 @@ TEST_CASE("MCPTaskManager: a submitted task runs and completes", "[mcp][tasks]")
 TEST_CASE("MCPTaskManager: a throwing work marks the task failed", "[mcp][tasks]") {
     MCPTaskManager m(1, 16);
     auto id = m.submit("tool", "alice", 60000, 100,
-                       [](const std::atomic<bool>&) -> std::string { throw std::runtime_error("boom"); });
+                       [](const std::atomic<bool>&, const MCPTaskManager::SetInterrupt&) -> std::string {
+                           throw std::runtime_error("boom");
+                       });
     REQUIRE(waitFor(m, id, "alice", Status::Failed));
     MCPTaskManager::Task t;
     bool found = false;
@@ -54,7 +58,9 @@ TEST_CASE("MCPTaskManager: a throwing work marks the task failed", "[mcp][tasks]
 TEST_CASE("MCPTaskManager: get re-checks principal ownership", "[mcp][tasks]") {
     MCPTaskManager m(1, 16);
     auto id = m.submit("tool", "alice", 60000, 100,
-                       [](const std::atomic<bool>&) { return std::string("{}"); });
+                       [](const std::atomic<bool>&, const MCPTaskManager::SetInterrupt&) {
+                           return std::string("{}");
+                       });
     REQUIRE(waitFor(m, id, "alice", Status::Completed));
 
     MCPTaskManager::Task t;
@@ -72,7 +78,7 @@ TEST_CASE("MCPTaskManager: cancel is honoured and cross-principal cancel is refu
     std::atomic<bool> release{false};
     // Work blocks until cancelled or released so we can cancel it mid-flight.
     auto id = m.submit("tool", "alice", 60000, 100,
-                       [&release](const std::atomic<bool>& cancelled) {
+                       [&release](const std::atomic<bool>& cancelled, const MCPTaskManager::SetInterrupt&) {
                            for (int i = 0; i < 200; ++i) {
                                if (cancelled.load() || release.load()) break;
                                std::this_thread::sleep_for(std::chrono::milliseconds(5));
@@ -85,12 +91,35 @@ TEST_CASE("MCPTaskManager: cancel is honoured and cross-principal cancel is refu
     REQUIRE(waitFor(m, id, "alice", Status::Cancelled, 3000));
 }
 
+TEST_CASE("MCPTaskManager: cancel fires the preemptive interrupt hook", "[mcp][tasks]") {
+    MCPTaskManager m(1, 16);
+    std::atomic<bool> interrupted{false};
+    // The work publishes an interrupt hook, then blocks on it — proving cancel()
+    // reaches a running task preemptively rather than waiting for cooperative
+    // polling. (In production the hook interrupts the DuckDB query.)
+    auto id = m.submit("tool", "alice", 60000, 100,
+                       [&interrupted](const std::atomic<bool>&,
+                                      const MCPTaskManager::SetInterrupt& set_interrupt) {
+                           set_interrupt([&interrupted]() { interrupted.store(true); });
+                           for (int i = 0; i < 400; ++i) {
+                               if (interrupted.load()) { break; }
+                               std::this_thread::sleep_for(std::chrono::milliseconds(5));
+                           }
+                           return std::string("{}");
+                       });
+    // Give the worker time to start and register the hook.
+    std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    REQUIRE(m.cancel(id, "alice"));
+    REQUIRE(waitFor(m, id, "alice", Status::Cancelled, 3000));
+    REQUIRE(interrupted.load());
+}
+
 TEST_CASE("MCPTaskManager: queue backpressure returns an empty id", "[mcp][tasks]") {
     // One worker, queue depth 1. Fill the worker with a blocking task, then the
     // queue with one, so a third submit is rejected.
     MCPTaskManager m(1, 1);
     std::atomic<bool> release{false};
-    auto blocker = [&release](const std::atomic<bool>& cancelled) {
+    auto blocker = [&release](const std::atomic<bool>& cancelled, const MCPTaskManager::SetInterrupt&) {
         while (!release.load() && !cancelled.load()) {
             std::this_thread::sleep_for(std::chrono::milliseconds(2));
         }

test/cpp/query_executor_test.cpp

@@ -5,6 +5,9 @@
 #include <catch2/matchers/catch_matchers_string.hpp>
 #include "query_executor.hpp"
 #include <cstdlib>
+#include <atomic>
+#include <chrono>
+#include <thread>
 
 using namespace std;
 
@@ -723,5 +726,42 @@ TEST_CASE("QueryExecutor::executeWithBindings - prepared path", "[query_executor
     duckdb_close(&database);
 }
 
+TEST_CASE("QueryExecutor::interrupt stops an in-flight query", "[query_executor][tasks]") {
+    // Foundation for preemptive task cancellation (issue #111): a long query
+    // running on one thread must be interruptible from another.
+    duckdb_database database;
+    REQUIRE(duckdb_open(NULL, &database) == DuckDBSuccess);
+
+    QueryExecutor executor(database);
+    std::atomic<bool> finished{false};
+    std::atomic<bool> threw{false};
+
+    const auto start = std::chrono::steady_clock::now();
+    std::thread runner([&]{
+        try {
+            // A count over a huge range takes many seconds unbudged; the
+            // interrupt below must cut it short.
+            executor.execute("SELECT count(*) FROM range(100000000000)");
+        } catch (const std::exception&) {
+            threw.store(true);
+        }
+        finished.store(true);
+    });
+
+    // Let the query get going, then interrupt from this thread.
+    std::this_thread::sleep_for(std::chrono::milliseconds(300));
+    executor.interrupt();
+
+    runner.join();
+    const auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
+        std::chrono::steady_clock::now() - start).count();
+
+    REQUIRE(finished.load());
+    REQUIRE(threw.load());          // interrupted query surfaces as an exception
+    REQUIRE(elapsed < 5000);        // stopped promptly, not after the full scan
+
+    duckdb_close(&database);
+}
+
 }  // namespace test
 }  // namespace flapi