Skip to content

Commit 99fcacf

Browse files
writePerfectCodewritePerfectCode
committed
Add cancel task
V2.20
1 parent defe5a5 commit 99fcacf

8 files changed

Lines changed: 230 additions & 176 deletions

File tree

‎examples/task_scheduler/task_scheduler_demo.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ json_session_invoke::SessionObjectHandle createCounter(
8383
auto response = scheduler.submitRequest(json_session_invoke::json{
8484
{"name", "create_counter"},
8585
{"args", {{"initial", initial}}},
86-
}).get();
86+
}).future.get();
8787
return response.at("value").get<json_session_invoke::SessionObjectHandle>();
8888
}
8989

@@ -184,11 +184,11 @@ int main()
184184
"create counter C",
185185
{{"name", "create_counter"}, {"args", {{"initial", 99}}}});
186186

187-
const auto add_response = add_a.get();
188-
const auto read_a_response = read_a.get();
189-
const auto read_b_response = read_b.get();
190-
const auto sum_response = sum.get();
191-
const auto create_c_response = create_c.get();
187+
const auto add_response = add_a.future.get();
188+
const auto read_a_response = read_a.future.get();
189+
const auto read_b_response = read_b.future.get();
190+
const auto sum_response = sum.future.get();
191+
const auto create_c_response = create_c.future.get();
192192

193193
printResponse("A add +7", add_response);
194194
printResponse("A read via string handle", read_a_response);

‎include/mcp/mcp_stdio_server.hpp‎

Lines changed: 56 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,11 @@
1313
#include <ostream>
1414
#include <sstream>
1515
#include <stdexcept>
16+
#include <stop_token>
1617
#include <string>
1718
#include <string_view>
1819
#include <thread>
20+
#include <unordered_map>
1921
#include <utility>
2022
#include <vector>
2123

@@ -33,7 +35,7 @@ struct ServerInfo {
3335
bool tools_list_changed{false};
3436
};
3537

36-
template<bool EnableThreadSafety = false>
38+
template<bool EnableThreadSafety = true>
3739
class BasicMcpStdioServer {
3840
public:
3941
using RuntimeType = runtime::BasicRuntimeFacade<EnableThreadSafety>;
@@ -49,28 +51,8 @@ class BasicMcpStdioServer {
4951
{
5052
}
5153

52-
std::optional<json> handleMessage(const json& message)
53-
{
54-
validateBaseMessage(message);
55-
56-
const std::string method = message.at("method").get<std::string>();
57-
const auto params_it = message.find("params");
58-
const json& params = params_it == message.end() ? empty_object() : *params_it;
59-
60-
const auto id_it = message.find("id");
61-
if (id_it == message.end())
62-
{
63-
handleNotification(method, params);
64-
return std::nullopt;
65-
}
66-
67-
return handleRequest(*id_it, method, params);
68-
}
69-
7054
void serveConcurrent(std::istream& input, std::ostream& output)
7155
{
72-
static_assert(EnableThreadSafety, "serveConcurrent requires EnableThreadSafety=true");
73-
7456
ResponseQueue response_queue;
7557
auto enqueue_response = [&response_queue](json response) {
7658
response_queue.push(std::move(response));
@@ -231,6 +213,16 @@ class BasicMcpStdioServer {
231213
return;
232214
}
233215

216+
if (method == "notifications/cancelled")
217+
{
218+
const auto id_it = params.find("requestId");
219+
if (id_it != params.end() && (id_it->is_string() || id_it->is_number()))
220+
{
221+
cancelInFlightRequest(id_it->dump());
222+
}
223+
return;
224+
}
225+
234226
if (!params.is_object())
235227
{
236228
throw std::runtime_error("notification params must be a JSON object");
@@ -352,12 +344,12 @@ class BasicMcpStdioServer {
352344
return error_response;
353345
}
354346

355-
const auto invoke_result = runtime_.invoke(runtime::InvokeRequest{
347+
auto handle = runtime_.submitInvoke(runtime::InvokeRequest{
356348
std::move(request->name),
357349
std::move(request->arguments),
358350
});
359351

360-
return makeToolsCallResponse(id, invoke_result);
352+
return makeToolsCallResponse(id, handle.future.get());
361353
}
362354

363355
struct ToolCallRequest {
@@ -467,18 +459,30 @@ class BasicMcpStdioServer {
467459
}
468460

469461
const json id = *id_it;
470-
auto invoke_future = runtime_.submitInvoke(runtime::InvokeRequest{
462+
const std::string id_key = id.dump();
463+
auto handle = runtime_.submitInvoke(runtime::InvokeRequest{
471464
std::move(request->name),
472465
std::move(request->arguments),
473466
});
474467

468+
registerInFlightHandle(id_key, handle.stop_source);
469+
475470
try
476471
{
477472
completion_tasks.push_back(std::async(std::launch::async,
478-
[this, id, invoke_future = std::move(invoke_future), enqueue_response]() mutable {
473+
[this, id, id_key, handle = std::move(handle), enqueue_response]() mutable {
474+
auto result = handle.future.get();
475+
unregisterInFlightHandle(id_key);
476+
477+
if (result.error.has_value() &&
478+
(result.error->code == "cancelled" || result.error->code == "result_discarded"))
479+
{
480+
return;
481+
}
482+
479483
try
480484
{
481-
enqueue_response(makeToolsCallResponse(id, invoke_future.get()));
485+
enqueue_response(makeToolsCallResponse(id, result));
482486
}
483487
catch (const std::exception& e)
484488
{
@@ -492,6 +496,7 @@ class BasicMcpStdioServer {
492496
}
493497
catch (const std::exception& e)
494498
{
499+
unregisterInFlightHandle(id_key);
495500
enqueue_response(makeErrorResponse(id, -32603, e.what()));
496501
}
497502
}
@@ -626,14 +631,38 @@ class BasicMcpStdioServer {
626631
value = std::string(begin, end);
627632
}
628633

634+
void registerInFlightHandle(const std::string& id_key, std::stop_source stop_source)
635+
{
636+
std::lock_guard<std::mutex> lock(in_flight_mutex_);
637+
in_flight_handles_.emplace(id_key, std::move(stop_source));
638+
}
639+
640+
void unregisterInFlightHandle(const std::string& id_key)
641+
{
642+
std::lock_guard<std::mutex> lock(in_flight_mutex_);
643+
in_flight_handles_.erase(id_key);
644+
}
645+
646+
void cancelInFlightRequest(const std::string& id_key)
647+
{
648+
std::lock_guard<std::mutex> lock(in_flight_mutex_);
649+
const auto it = in_flight_handles_.find(id_key);
650+
if (it != in_flight_handles_.end())
651+
{
652+
it->second.request_stop();
653+
}
654+
}
655+
629656
RuntimeType& runtime_;
630657
ServerInfo server_info_{};
631658
mutable std::mutex state_mutex_;
632659
bool initialize_seen_{false};
633660
bool initialized_notification_seen_{false};
661+
mutable std::mutex in_flight_mutex_;
662+
std::unordered_map<std::string, std::stop_source> in_flight_handles_;
634663
};
635664

665+
using McpStdioServer = BasicMcpStdioServer<true>;
636666
using McpStdioServerThreadSafe = BasicMcpStdioServer<true>;
637-
using McpStdioServerUnsafe = BasicMcpStdioServer<false>;
638667

639668
} // namespace mcp

‎include/runtime/runtime_facade.hpp‎

Lines changed: 30 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include <cstddef>
55
#include <future>
66
#include <optional>
7+
#include <stop_token>
78
#include <string>
89
#include <utility>
910
#include <vector>
@@ -45,6 +46,11 @@ struct InvokeOptions {
4546
std::optional<std::chrono::milliseconds> timeout;
4647
};
4748

49+
struct InvokeHandle {
50+
std::future<InvokeResult> future;
51+
std::stop_source stop_source;
52+
};
53+
4854
template<bool EnableThreadSafety = false>
4955
class BasicRuntimeFacade {
5056
public:
@@ -77,7 +83,7 @@ class BasicRuntimeFacade {
7783
return describeSchema(adapter_.getToolSchemaJson(tool_name));
7884
}
7985

80-
InvokeResult invoke(InvokeRequest request, InvokeOptions options = {})
86+
InvokeHandle submitInvoke(InvokeRequest request, InvokeOptions options = {})
8187
{
8288
const std::string tool_name = request.tool_name;
8389

@@ -89,79 +95,34 @@ class BasicRuntimeFacade {
8995

9096
try
9197
{
92-
auto response_future = scheduler_.submitRequest(makeRequestJson(std::move(request)),
98+
auto scheduler_result = scheduler_.submitRequest(makeRequestJson(std::move(request)),
9399
std::move(deadline));
94100

95-
if (options.timeout.has_value())
96-
{
97-
if (response_future.wait_for(*options.timeout) == std::future_status::timeout)
98-
{
99-
return makeErrorResult(tool_name, "timeout",
100-
"tool invocation exceeded timeout of " +
101-
std::to_string(options.timeout->count()) + "ms");
102-
}
103-
}
104-
105-
try
106-
{
107-
return normalizeResponse(tool_name, response_future.get());
108-
}
109-
catch (const std::runtime_error& e)
110-
{
111-
if (std::string(e.what()) == "task_timeout")
112-
{
113-
return makeErrorResult(tool_name, "timeout",
114-
"tool invocation timed out in queue");
115-
}
116-
return makeErrorResult(tool_name, "unknown_error", e.what());
117-
}
118-
catch (const json_invoke::JsonInvokeError& e)
119-
{
120-
return makeErrorResult(tool_name, e.code(), e.what());
121-
}
122-
catch (const std::exception& e)
123-
{
124-
return makeErrorResult(tool_name, "unknown_error", e.what());
125-
}
126-
}
127-
catch (const json_invoke::JsonInvokeError& e)
128-
{
129-
return makeErrorResult(tool_name, e.code(), e.what());
130-
}
131-
catch (const std::exception& e)
132-
{
133-
return makeErrorResult(tool_name, "unknown_error", e.what());
134-
}
135-
}
136-
137-
std::future<InvokeResult> submitInvoke(InvokeRequest request, InvokeOptions options = {})
138-
{
139-
const std::string tool_name = request.tool_name;
140-
141-
std::optional<task_scheduler::TaskDeadline> deadline;
142-
if (options.timeout.has_value())
143-
{
144-
deadline = std::chrono::steady_clock::now() + *options.timeout;
145-
}
146-
147-
try
148-
{
149-
auto response_future = scheduler_.submitRequest(makeRequestJson(std::move(request)),
150-
std::move(deadline));
151-
return std::async(std::launch::deferred,
152-
[tool_name, response_future = std::move(response_future)]() mutable {
101+
auto future = std::async(std::launch::deferred,
102+
[tool_name, response_future = std::move(scheduler_result.future)]() mutable {
153103
try
154104
{
155105
return normalizeResponse(tool_name, response_future.get());
156106
}
157107
catch (const std::runtime_error& e)
158108
{
159-
if (std::string(e.what()) == "task_timeout")
109+
const std::string what = e.what();
110+
if (what == "task_cancelled")
111+
{
112+
return makeErrorResult(tool_name, "cancelled",
113+
"tool invocation was cancelled before execution");
114+
}
115+
if (what == "task_result_discarded")
116+
{
117+
return makeErrorResult(tool_name, "result_discarded",
118+
"tool invocation completed but result was discarded after cancellation");
119+
}
120+
if (what == "task_timeout")
160121
{
161122
return makeErrorResult(tool_name, "timeout",
162123
"tool invocation timed out in queue");
163124
}
164-
return makeErrorResult(tool_name, "unknown_error", e.what());
125+
return makeErrorResult(tool_name, "unknown_error", what);
165126
}
166127
catch (const json_invoke::JsonInvokeError& e)
167128
{
@@ -172,14 +133,19 @@ class BasicRuntimeFacade {
172133
return makeErrorResult(tool_name, "unknown_error", e.what());
173134
}
174135
});
136+
137+
return InvokeHandle{
138+
std::move(future),
139+
std::move(scheduler_result.stop_source),
140+
};
175141
}
176142
catch (const json_invoke::JsonInvokeError& e)
177143
{
178-
return makeReadyFuture(makeErrorResult(tool_name, e.code(), e.what()));
144+
return InvokeHandle{makeReadyFuture(makeErrorResult(tool_name, e.code(), e.what())), {}};
179145
}
180146
catch (const std::exception& e)
181147
{
182-
return makeReadyFuture(makeErrorResult(tool_name, "unknown_error", e.what()));
148+
return InvokeHandle{makeReadyFuture(makeErrorResult(tool_name, "unknown_error", e.what())), {}};
183149
}
184150
}
185151

0 commit comments

Comments
 (0)