tool_execution.h

Go to the documentation of this file.

 
#pragma once
 
#include <neograph/api.h>
#include <neograph/host_admission.h>
#include <neograph/json.h>
#include <asio/awaitable.hpp>
 
#include <chrono>
#include <cstdint>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
 
namespace neograph {
class Tool;
struct ToolExecutionContext;
 
namespace graph {
class CancelToken;
}
 
enum class ToolExecutionImplementation : std::uint8_t {
    NativeAsync,
    BlockingThread,
    BlockingProcess,
};
 
enum class ToolTerminalStatus : std::uint8_t {
    Succeeded,
    Failed,
    CancelledBeforeStart,
    CancellationRequested,
    Expired,
    TimedOut,
    Killed,
    InputRequired,
    ReconciliationRequired,
    Rejected,
};
 
NEOGRAPH_API std::string_view to_string(ToolTerminalStatus status) noexcept;
NEOGRAPH_API ToolTerminalStatus tool_terminal_status_from_string(std::string_view value);
 
inline bool tool_call_not_started(ToolTerminalStatus status) noexcept {
    return status == ToolTerminalStatus::CancelledBeforeStart ||
           status == ToolTerminalStatus::Expired || status == ToolTerminalStatus::Rejected;
}
 
struct NEOGRAPH_API ToolExecutionResult {
    ToolTerminalStatus status = ToolTerminalStatus::Failed;
    std::string output;
    std::string error;
    bool retryable = false;
    bool effect_uncertain = false;
    std::optional<int> exit_code;
    std::optional<int> signal_number;
    bool output_truncated = false;
    bool refused = false;
 
    [[nodiscard]] bool succeeded() const noexcept {
        return status == ToolTerminalStatus::Succeeded;
    }
};
 
class NEOGRAPH_API ToolExecutionError : public std::runtime_error {
public:
    explicit ToolExecutionError(ToolExecutionResult result);
    [[nodiscard]] const ToolExecutionResult& result() const noexcept { return result_; }
    [[nodiscard]] ToolTerminalStatus status() const noexcept { return result_.status; }
 
private:
    ToolExecutionResult result_;
};
 
class NEOGRAPH_API ToolRequestRefusal {
public:
    virtual ~ToolRequestRefusal();
 
    [[nodiscard]] virtual bool retryable() const noexcept = 0;
};
 
class NEOGRAPH_API ToolRequestRejected : public std::runtime_error, public ToolRequestRefusal {
public:
    explicit ToolRequestRejected(const std::string& message, bool retryable = false);
    ~ToolRequestRejected() override;
    [[nodiscard]] bool retryable() const noexcept override { return retryable_; }
 
private:
    bool retryable_;
};
 
enum class ToolEffectClass : std::uint8_t {
    Unknown,
    ReadOnly,
    ExternalRead,
    WorkspaceWrite,
    ExternalWrite,
};
 
enum class ToolIdempotency : std::uint8_t {
    Unknown,
    Idempotent,
    Conditional,
    NonIdempotent,
};
 
enum class ToolConcurrency : std::uint8_t {
    Reentrant,
    KeyedExclusive,
    Exclusive,
    Capacity,
    SingleFlight,
    ExternalLimited,
};
struct ToolExecutionIdentity {
    std::string owner_scope;
    std::string root_run_id;
    std::string thread_id;
    std::string request_id;
};
 
struct ToolEffectGrantIdentity {
    std::string program_version_id;
    std::string operation_id;
    std::string grant_id;
    std::uint64_t attempt = 0;
    std::string binding_fingerprint;
};
 
struct ToolExecutionPolicy {
    static constexpr std::uint32_t SCHEMA_VERSION = 2;
    static constexpr std::size_t DEFAULT_ERROR_LIMIT_BYTES = 2048;
 
    std::uint32_t schema_version = SCHEMA_VERSION;
    std::string   policy_id;
    std::string   policy_source{"host-default"};
    std::string   policy_hash;
    ToolExecutionImplementation implementation = ToolExecutionImplementation::NativeAsync;
    ToolConcurrency concurrency = ToolConcurrency::KeyedExclusive;
    std::uint32_t capacity = 1;
    std::uint32_t max_pending = 128;
    std::chrono::milliseconds queue_timeout{60'000};
    std::chrono::milliseconds execution_timeout{0};
    bool cancellable = true;
    ToolEffectClass effect = ToolEffectClass::Unknown;
    ToolIdempotency idempotency = ToolIdempotency::Unknown;
    std::uint32_t retry_max_attempts = 1;
    std::chrono::milliseconds retry_backoff{0};
    std::size_t output_limit_bytes = 16U * 1024U * 1024U;
    std::size_t error_limit_bytes = DEFAULT_ERROR_LIMIT_BYTES;
    std::string resource_key_template{"tool:{tool}"};
    HostResourceVector host_resources;
    std::uint8_t host_priority = 0;
    std::uint8_t priority_ceiling = 255;
    std::uint8_t fairness_weight = 1;
 
    NEOGRAPH_API void validate() const;
};
 
struct ResourceRequest {
    std::string resource_key;
    std::string owner_scope;
    std::uint32_t capacity = 1;
    std::uint32_t max_pending = 128;
    std::chrono::milliseconds queue_timeout{60'000};
    std::uint8_t priority = 0;
    std::uint8_t fairness_weight = 1;
};
 
namespace detail {
class ResourceArbiterImpl;
 
asio::awaitable<std::string> execute_blocking_tool_async(Tool& tool, json arguments);
asio::awaitable<std::string> execute_process_tool_async(
    Tool& tool, json arguments, ToolExecutionContext context,
    std::size_t output_limit_bytes, std::chrono::milliseconds timeout);
} // namespace detail
 
class NEOGRAPH_API ResourceLease final {
public:
    ResourceLease() = default;
    ResourceLease(const ResourceLease&) = delete;
    ResourceLease& operator=(const ResourceLease&) = delete;
    ResourceLease(ResourceLease&& other) noexcept;
    ResourceLease& operator=(ResourceLease&& other) noexcept;
    ~ResourceLease() noexcept;
 
    [[nodiscard]] bool held() const noexcept;
    [[nodiscard]] std::string_view resource_key() const noexcept;
    void release() noexcept;
 
private:
    friend class ResourceArbiter;
    friend class detail::ResourceArbiterImpl;
    ResourceLease(std::shared_ptr<detail::ResourceArbiterImpl> impl,
                  std::string resource_key) noexcept;
 
    std::shared_ptr<detail::ResourceArbiterImpl> impl_;
    std::string resource_key_;
};
 
class NEOGRAPH_API ResourceArbiter final {
public:
    ResourceArbiter();
    ~ResourceArbiter();
    ResourceArbiter(const ResourceArbiter&) = delete;
    ResourceArbiter& operator=(const ResourceArbiter&) = delete;
 
    asio::awaitable<ResourceLease> acquire_async(
        ResourceRequest request,
        std::shared_ptr<graph::CancelToken> cancel_token = {});
 
private:
    std::shared_ptr<detail::ResourceArbiterImpl> impl_;
};
 
class NEOGRAPH_API ToolExecutionPolicyRegistry final {
public:
    ToolExecutionPolicyRegistry();
    ~ToolExecutionPolicyRegistry();
    ToolExecutionPolicyRegistry(const ToolExecutionPolicyRegistry&) = delete;
    ToolExecutionPolicyRegistry& operator=(const ToolExecutionPolicyRegistry&) = delete;
 
    void upsert(std::string tool_name, ToolExecutionPolicy policy);
    void erase(std::string_view tool_name);
    void set_default(ToolExecutionPolicy policy);
    [[nodiscard]] ToolExecutionPolicy resolve(std::string_view tool_name) const;
 
private:
    class Impl;
    std::unique_ptr<Impl> impl_;
};
 
struct ToolExecutionControllerConfig {
    std::shared_ptr<ToolExecutionPolicyRegistry> policies;
    std::shared_ptr<ResourceArbiter> arbiter;
    std::shared_ptr<HostAdmissionController> host_admission;
};
 
class NEOGRAPH_API ToolExecutionController final {
public:
    explicit ToolExecutionController(ToolExecutionControllerConfig config = {});
    ~ToolExecutionController();
    ToolExecutionController(const ToolExecutionController&) = delete;
    ToolExecutionController& operator=(const ToolExecutionController&) = delete;
 
    asio::awaitable<ToolExecutionResult> execute_result_async(
        Tool& tool, json arguments, ToolExecutionContext context);
 
    asio::awaitable<std::string> execute_async(Tool& tool,
                                               json arguments,
                                               ToolExecutionContext context);
 
    [[nodiscard]] std::shared_ptr<ToolExecutionPolicyRegistry> policies() const;
    [[nodiscard]] std::shared_ptr<ResourceArbiter> arbiter() const;
    [[nodiscard]] std::shared_ptr<HostAdmissionController> host_admission() const;
 
private:
    friend class NativeHookExecutionAdapter;
    asio::awaitable<ToolExecutionResult> execute_hook_target_async(
        Tool& tool, json arguments, ToolExecutionContext context);
    asio::awaitable<ToolExecutionResult> execute_result_async(
        Tool& tool, json arguments, ToolExecutionContext context, bool retry_after_timeout);
 
    class Impl;
    std::unique_ptr<Impl> impl_;
};
 
NEOGRAPH_API std::shared_ptr<ToolExecutionController> default_tool_execution_controller();
 
} // namespace neograph