From 20640bb8dc54278fc2512de7a2728a6a0310a33e Mon Sep 17 00:00:00 2001 From: Joachim Rosskopf Date: Mon, 31 Aug 2026 18:57:51 +0200 Subject: [PATCH 1/2] =?UTF-8?q?feat(mcp):=20add=20QueryExecutor::interrupt?= =?UTF-8?q?()=20=E2=80=94=20foundation=20for=20preemptive=20task=20cancel?= =?UTF-8?q?=20(#111)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First increment toward interruptible tasks: QueryExecutor already owns its duckdb_connection for its lifetime, so expose interrupt() which calls duckdb_interrupt(conn) from another thread. Additive and REST-safe (the REST path never calls it); a running query returns an error, surfaced as an exception by execute(). Verified by a unit test that starts a huge range-count on one thread and interrupts it from another — it stops in ~1.4s (vs many seconds unbudged) and throws. Next: thread this through a cancellable async-tool execution path and wire tasks/cancel + shutdown to it (see the plan on #111). --- src/include/query_executor.hpp | 7 ++++++ test/cpp/query_executor_test.cpp | 40 ++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/include/query_executor.hpp b/src/include/query_executor.hpp index b18df01..540e01f 100644 --- a/src/include/query_executor.hpp +++ b/src/include/query_executor.hpp @@ -115,6 +115,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; diff --git a/test/cpp/query_executor_test.cpp b/test/cpp/query_executor_test.cpp index 345ce08..666bc3c 100644 --- a/test/cpp/query_executor_test.cpp +++ b/test/cpp/query_executor_test.cpp @@ -5,6 +5,9 @@ #include #include "query_executor.hpp" #include +#include +#include +#include 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 finished{false}; + std::atomic 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::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 From 49633b5e1e41502e71b531078a214e848c58eab8 Mon Sep 17 00:00:00 2001 From: Joachim Rosskopf Date: Tue, 1 Sep 2026 03:31:47 +0200 Subject: [PATCH 2/2] feat: preemptive MCP task cancellation via query interrupt (#111) - Add a thread-id registry so a QueryExecutor self-registers for the duration of its DuckDB query; another thread can then interrupt it by thread id without threading an executor handle through the whole executeTool -> DatabaseManager -> QueryExecutor pipeline. - MCPTaskManager gains a per-task interrupt hook: the worker publishes an interrupt closure while running, and tasks/cancel and shutdown fire it to stop the in-flight query preemptively instead of waiting for cooperative polling. Cancelled-flag still wins over the resulting error, so an interrupted task lands in 'cancelled', not 'failed'. - shutdown() now interrupts running tasks before joining, so a graceful stop no longer blocks on a multi-minute query finishing on its own. - Tests: unit test that cancel fires the interrupt hook; integration test that a long (100B-row) async query is cancelled in ~1s instead of ~32s. --- src/include/mcp_task_manager.hpp | 17 ++++++++- src/include/query_executor.hpp | 13 +++++++ src/mcp_route_handlers.cpp | 13 ++++++- src/mcp_task_manager.cpp | 29 +++++++++++---- src/query_executor.cpp | 60 +++++++++++++++++++++++++++++- test/cpp/mcp_task_manager_test.cpp | 39 ++++++++++++++++--- test/integration/test_mcp_tasks.py | 27 ++++++++++++++ 7 files changed, 181 insertions(+), 17 deletions(-) diff --git a/src/include/mcp_task_manager.hpp b/src/include/mcp_task_manager.hpp index a26b123..0a69856 100644 --- a/src/include/mcp_task_manager.hpp +++ b/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)>; + // 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& 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& 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 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 interrupt; }; void workerLoop(); diff --git a/src/include/query_executor.hpp b/src/include/query_executor.hpp index 540e01f..3560093 100644 --- a/src/include/query_executor.hpp +++ b/src/include/query_executor.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #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 diff --git a/src/mcp_route_handlers.cpp b/src/mcp_route_handlers.cpp index 30bbc7f..ee7ebfd 100644 --- a/src/mcp_route_handlers.cpp +++ b/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 +#include #include #include @@ -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&) -> std::string { + auto work = [handler, tool_request]( + const std::atomic&, + 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) { diff --git a/src/mcp_task_manager.cpp b/src/mcp_task_manager.cpp index 0c0bbd9..36139bd 100644 --- a/src/mcp_task_manager.cpp +++ b/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 fn) { + std::lock_guard 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 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(); diff --git a/src/query_executor.cpp b/src/query_executor.cpp index 0436c6d..f9d40b6 100644 --- a/src/query_executor.cpp +++ b/src/query_executor.cpp @@ -12,8 +12,52 @@ // duckdb::Vector*, so dropping to Vector::GetValue() is a safe round-trip. #include "duckdb.hpp" +#include +#include + 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& activeExecMap() { + static std::unordered_map m; + return m; +} +} // namespace + +void registerActiveExecutor(std::thread::id tid, QueryExecutor* exec) { + std::lock_guard lock(activeExecMutex()); + activeExecMap()[tid] = exec; +} +void unregisterActiveExecutor(std::thread::id tid) { + std::lock_guard 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 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); diff --git a/test/cpp/mcp_task_manager_test.cpp b/test/cpp/mcp_task_manager_test.cpp index 4538583..ff4b04f 100644 --- a/test/cpp/mcp_task_manager_test.cpp +++ b/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&) { return std::string("{\"ok\":true}"); }); + [](const std::atomic&, 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&) -> std::string { throw std::runtime_error("boom"); }); + [](const std::atomic&, 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&) { return std::string("{}"); }); + [](const std::atomic&, 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 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& cancelled) { + [&release](const std::atomic& 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 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&, + 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 release{false}; - auto blocker = [&release](const std::atomic& cancelled) { + auto blocker = [&release](const std::atomic& cancelled, const MCPTaskManager::SetInterrupt&) { while (!release.load() && !cancelled.load()) { std::this_thread::sleep_for(std::chrono::milliseconds(2)); } diff --git a/test/integration/test_mcp_tasks.py b/test/integration/test_mcp_tasks.py index d640c33..757c7be 100644 --- a/test/integration/test_mcp_tasks.py +++ b/test/integration/test_mcp_tasks.py @@ -69,6 +69,12 @@ def server() -> Iterator[str]: f.write("SELECT 7 AS answer\n") with open(os.path.join(sqls, "rep.yaml"), "w") as f: f.write("template-source: rep.sql\nconnection: [inmem]\nmcp-tool: {name: report, description: r, async: true}\n") + # A deliberately long-running query, used to prove tasks/cancel is preemptive + # (interrupts the in-flight DuckDB query) rather than merely cooperative. + with open(os.path.join(sqls, "slow.sql"), "w") as f: + f.write("SELECT count(*) AS n FROM range(100000000000)\n") + with open(os.path.join(sqls, "slow.yaml"), "w") as f: + f.write("template-source: slow.sql\nconnection: [inmem]\nmcp-tool: {name: slow, description: s, async: true}\n") log = open(os.path.join(tmp, "server.log"), "w") proc = subprocess.Popen([binary, "-c", os.path.join(tmp, "flapi.yaml"), "--no-telemetry"], @@ -166,6 +172,27 @@ def test_tasks_cancel(self, server): assert "error" not in body, body assert body["result"]["taskId"] == task_id + def test_tasks_cancel_preempts_long_query(self, server): + # Submit a query that would run for many minutes, then cancel it. With + # preemptive cancellation (issue #111) the in-flight DuckDB query is + # interrupted, so the task reaches `cancelled` in seconds rather than + # hanging until the query finishes on its own. + r = _call(server, META_TASKS, name="slow") + result = r.json()["result"] + assert result["resultType"] == "task", result + task_id = result["task"]["taskId"] + + # Let the worker actually start the query before cancelling. + time.sleep(0.5) + c = requests.post(f"{server}/mcp/jsonrpc", headers=_headers("tasks/cancel"), + json={"jsonrpc": "2.0", "id": 3, "method": "tasks/cancel", + "params": {"taskId": task_id, "_meta": META_TASKS}}, timeout=10) + assert "error" not in c.json(), c.text + + # Must transition to cancelled quickly — proof the query was interrupted. + cancelled = _wait_status(server, task_id, "cancelled", tries=20) + assert cancelled is not None, "long query was not cancelled preemptively" + def _wait_status(base, task_id, want, tries=40): for _ in range(tries):