Support task timeout

V2.19

Author: writePerfectCode

include/runtime/runtime_facade.hpp

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <chrono>
 #include <cstddef>
 #include <future>
 #include <optional>
@@ -40,6 +41,10 @@ struct InvokeResult {
     json raw_response = json::object();
 };
 
+struct InvokeOptions {
+    std::optional<std::chrono::milliseconds> timeout;
+};
+
 template<bool EnableThreadSafety = false>
 class BasicRuntimeFacade {
 public:
@@ -72,32 +77,101 @@ class BasicRuntimeFacade {
         return describeSchema(adapter_.getToolSchemaJson(tool_name));
     }
 
-    InvokeResult invoke(InvokeRequest request)
-    {
-        return submitInvoke(std::move(request)).get();
-    }
-
-    std::future<InvokeResult> submitInvoke(InvokeRequest request)
+    InvokeResult invoke(InvokeRequest request, InvokeOptions options = {})
     {
         const std::string tool_name = request.tool_name;
 
+        std::optional<task_scheduler::TaskDeadline> deadline;
+        if (options.timeout.has_value())
+        {
+            deadline = std::chrono::steady_clock::now() + *options.timeout;
+        }
+
         try
         {
-            auto response_future = scheduler_.submitRequest(makeRequestJson(std::move(request)));
-            return std::async(std::launch::deferred, [tool_name, response_future = std::move(response_future)]() mutable {
-                try
-                {
-                    return normalizeResponse(tool_name, response_future.get());
-                }
-                catch (const json_invoke::JsonInvokeError& e)
+            auto response_future = scheduler_.submitRequest(makeRequestJson(std::move(request)),
+                std::move(deadline));
+
+            if (options.timeout.has_value())
+            {
+                if (response_future.wait_for(*options.timeout) == std::future_status::timeout)
                 {
-                    return makeErrorResult(tool_name, e.code(), e.what());
+                    return makeErrorResult(tool_name, "timeout",
+                        "tool invocation exceeded timeout of " +
+                        std::to_string(options.timeout->count()) + "ms");
                 }
-                catch (const std::exception& e)
+            }
+
+            try
+            {
+                return normalizeResponse(tool_name, response_future.get());
+            }
+            catch (const std::runtime_error& e)
+            {
+                if (std::string(e.what()) == "task_timeout")
                 {
-                    return makeErrorResult(tool_name, "unknown_error", e.what());
+                    return makeErrorResult(tool_name, "timeout",
+                        "tool invocation timed out in queue");
                 }
-            });
+                return makeErrorResult(tool_name, "unknown_error", e.what());
+            }
+            catch (const json_invoke::JsonInvokeError& e)
+            {
+                return makeErrorResult(tool_name, e.code(), e.what());
+            }
+            catch (const std::exception& e)
+            {
+                return makeErrorResult(tool_name, "unknown_error", e.what());
+            }
+        }
+        catch (const json_invoke::JsonInvokeError& e)
+        {
+            return makeErrorResult(tool_name, e.code(), e.what());
+        }
+        catch (const std::exception& e)
+        {
+            return makeErrorResult(tool_name, "unknown_error", e.what());
+        }
+    }
+
+    std::future<InvokeResult> submitInvoke(InvokeRequest request, InvokeOptions options = {})
+    {
+        const std::string tool_name = request.tool_name;
+
+        std::optional<task_scheduler::TaskDeadline> deadline;
+        if (options.timeout.has_value())
+        {
+            deadline = std::chrono::steady_clock::now() + *options.timeout;
+        }
+
+        try
+        {
+            auto response_future = scheduler_.submitRequest(makeRequestJson(std::move(request)),
+                std::move(deadline));
+            return std::async(std::launch::deferred,
+                [tool_name, response_future = std::move(response_future)]() mutable {
+                    try
+                    {
+                        return normalizeResponse(tool_name, response_future.get());
+                    }
+                    catch (const std::runtime_error& e)
+                    {
+                        if (std::string(e.what()) == "task_timeout")
+                        {
+                            return makeErrorResult(tool_name, "timeout",
+                                "tool invocation timed out in queue");
+                        }
+                        return makeErrorResult(tool_name, "unknown_error", e.what());
+                    }
+                    catch (const json_invoke::JsonInvokeError& e)
+                    {
+                        return makeErrorResult(tool_name, e.code(), e.what());
+                    }
+                    catch (const std::exception& e)
+                    {
+                        return makeErrorResult(tool_name, "unknown_error", e.what());
+                    }
+                });
         }
         catch (const json_invoke::JsonInvokeError& e)
         {

include/task_scheduler/task_scheduler.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <algorithm>
+#include <chrono>
 #include <condition_variable>
 #include <cstddef>
 #include <deque>
@@ -9,6 +10,7 @@
 #include <future>
 #include <memory>
 #include <mutex>
+#include <optional>
 #include <stdexcept>
 #include <thread>
 #include <unordered_map>
@@ -21,9 +23,12 @@ namespace task_scheduler {
 
 using ScheduledTaskFunction = std::function<json()>;
 
+using TaskDeadline = std::chrono::steady_clock::time_point;
+
 struct ScheduledTask {
     TaskExecutionPlan plan;
     ScheduledTaskFunction run;
+    std::optional<TaskDeadline> deadline;
 };
 
 class ITaskScheduler {
@@ -123,13 +128,22 @@ class KeyedTaskScheduler final : public ITaskScheduler {
                 pending_tasks_.erase(ready_it);
             }
 
-            try
+            if (item.task.deadline.has_value() &&
+                std::chrono::steady_clock::now() > *item.task.deadline)
             {
-                item.promise.set_value(item.task.run());
+                item.promise.set_exception(std::make_exception_ptr(
+                    std::runtime_error("task_timeout")));
             }
-            catch (...)
+            else
             {
-                item.promise.set_exception(std::current_exception());
+                try
+                {
+                    item.promise.set_value(item.task.run());
+                }
+                catch (...)
+                {
+                    item.promise.set_exception(std::current_exception());
+                }
             }
 
             {
@@ -312,11 +326,13 @@ class KeyedTaskScheduler final : public ITaskScheduler {
 };
 
 template<typename Fn>
-ScheduledTask makeScheduledTask(TaskExecutionPlan plan, Fn&& fn)
+ScheduledTask makeScheduledTask(TaskExecutionPlan plan, Fn&& fn,
+    std::optional<TaskDeadline> deadline = std::nullopt)
 {
     return ScheduledTask{
         std::move(plan),
         ScheduledTaskFunction(std::forward<Fn>(fn)),
+        std::move(deadline),
     };
 }
 

include/task_scheduler/task_scheduler_facade.hpp

@@ -1,8 +1,10 @@
 #pragma once
 
+#include <chrono>
 #include <exception>
 #include <functional>
 #include <future>
+#include <optional>
 #include <utility>
 
 #include <task_scheduler/task_scheduler.hpp>
@@ -38,10 +40,21 @@ class BasicTaskSchedulerFacade {
 
     std::future<json> submitRequest(json request)
     {
-        return submitRequest(std::move(request), RequestExecutionObserver{});
+        return submitRequest(std::move(request), RequestExecutionObserver{}, std::nullopt);
     }
 
     std::future<json> submitRequest(json request, RequestExecutionObserver observer)
+    {
+        return submitRequest(std::move(request), std::move(observer), std::nullopt);
+    }
+
+    std::future<json> submitRequest(json request, std::optional<TaskDeadline> deadline)
+    {
+        return submitRequest(std::move(request), RequestExecutionObserver{}, std::move(deadline));
+    }
+
+    std::future<json> submitRequest(json request, RequestExecutionObserver observer,
+        std::optional<TaskDeadline> deadline)
     {
         auto plan = classifier_.classify(request);
         if (observer.on_classified)
@@ -74,7 +87,8 @@ class BasicTaskSchedulerFacade {
                     }
                     throw;
                 }
-            }));
+            },
+            std::move(deadline)));
     }
 
 private: