mirror/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2026-02-18 08:09:22 -06:00
Code Issues Packages Projects Releases Wiki Activity

feat: respect context and add request cancellation (#7187)

Browse Source 
* feat: respect context

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* workaround fasthttp

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(ui): allow to abort call

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Refactor

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* chore: improving error

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Respect context also with MCP

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Tie to both contexts

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Make detection more robust

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto
2025-11-09 18:19:19 +01:00
committed by GitHub
parent 4730b52461
commit 679d43c2f5
8 changed files with 240 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
backend/cpp/llama-cpp/grpc-server.cpp
View File

@@ -822,6 +822,12 @@ public:
}
ctx_server.receive_cmpl_results_stream(task_ids, [&](server_task_result_ptr & result) -> bool {
// Check if context is cancelled before processing result
if (context->IsCancelled()) {
ctx_server.cancel_tasks(task_ids);
return false;
}
json res_json = result->to_json();
if (res_json.is_array()) {
for (const auto & res : res_json) {
@@ -875,13 +881,18 @@ public:
reply.set_message(error_data.value("content", ""));
writer->Write(reply);
return true;
}, [&]() {
// NOTE: we should try to check when the writer is closed here
return false;
}, [&context]() {
// Check if the gRPC context is cancelled
return context->IsCancelled();
});
ctx_server.queue_results.remove_waiting_task_ids(task_ids);
// Check if context was cancelled during processing
if (context->IsCancelled()) {
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
return grpc::Status::OK;
}
@@ -1145,6 +1156,14 @@ public:
std::cout << "[DEBUG] Waiting for results..." << std::endl;
// Check cancellation before waiting for results
if (context->IsCancelled()) {
ctx_server.cancel_tasks(task_ids);
ctx_server.queue_results.remove_waiting_task_ids(task_ids);
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
ctx_server.receive_multi_results(task_ids, [&](std::vector<server_task_result_ptr> & results) {
std::cout << "[DEBUG] Received " << results.size() << " results" << std::endl;
if (results.size() == 1) {
@@ -1176,13 +1195,20 @@ public:
}, [&](const json & error_data) {
std::cout << "[DEBUG] Error in results: " << error_data.value("content", "") << std::endl;
reply->set_message(error_data.value("content", ""));
}, [&]() {
return false;
}, [&context]() {
// Check if the gRPC context is cancelled
// This is checked every HTTP_POLLING_SECONDS (1 second) during receive_multi_results
return context->IsCancelled();
});
ctx_server.queue_results.remove_waiting_task_ids(task_ids);
std::cout << "[DEBUG] Predict request completed successfully" << std::endl;
// Check if context was cancelled during processing
if (context->IsCancelled()) {
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
return grpc::Status::OK;
}
@@ -1234,6 +1260,13 @@ public:
ctx_server.queue_tasks.post(std::move(tasks));
}
// Check cancellation before waiting for results
if (context->IsCancelled()) {
ctx_server.cancel_tasks(task_ids);
ctx_server.queue_results.remove_waiting_task_ids(task_ids);
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
// get the result
ctx_server.receive_multi_results(task_ids, [&](std::vector<server_task_result_ptr> & results) {
for (auto & res : results) {
@@ -1242,12 +1275,18 @@ public:
}
}, [&](const json & error_data) {
error = true;
}, [&]() {
return false;
}, [&context]() {
// Check if the gRPC context is cancelled
return context->IsCancelled();
});
ctx_server.queue_results.remove_waiting_task_ids(task_ids);
// Check if context was cancelled during processing
if (context->IsCancelled()) {
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
if (error) {
return grpc::Status(grpc::StatusCode::INTERNAL, "Error in receiving results");
}
@@ -1325,6 +1364,13 @@ public:
ctx_server.queue_tasks.post(std::move(tasks));
}
// Check cancellation before waiting for results
if (context->IsCancelled()) {
ctx_server.cancel_tasks(task_ids);
ctx_server.queue_results.remove_waiting_task_ids(task_ids);
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
// Get the results
ctx_server.receive_multi_results(task_ids, [&](std::vector<server_task_result_ptr> & results) {
for (auto & res : results) {
@@ -1333,12 +1379,18 @@ public:
}
}, [&](const json & error_data) {
error = true;
}, [&]() {
return false;
}, [&context]() {
// Check if the gRPC context is cancelled
return context->IsCancelled();
});
ctx_server.queue_results.remove_waiting_task_ids(task_ids);
// Check if context was cancelled during processing
if (context->IsCancelled()) {
return grpc::Status(grpc::StatusCode::CANCELLED, "Request cancelled by client");
}
if (error) {
return grpc::Status(grpc::StatusCode::INTERNAL, "Error in receiving results");
}