engine.h

Go to the documentation of this file.

 
#pragma once
 
#include <neograph/api.h>
#include <neograph/graph/admin.h>
#include <neograph/graph/cancel.h>
#include <neograph/graph/channel_key.h>
#include <neograph/graph/checkpoint.h>
#include <neograph/graph/compiler.h>
#include <neograph/graph/coordinator.h>
#include <neograph/graph/executor.h>
#include <neograph/graph/node.h>
#include <neograph/graph/node_cache.h>
#include <neograph/graph/provider_call_broker.h>
#include <neograph/tool_effect_broker.h>
#include <neograph/graph/registry.h>
#include <neograph/graph/run_context.h>
#include <neograph/graph/safe_point.h>
#include <neograph/graph/scheduler.h>
#include <neograph/graph/state.h>
#include <neograph/graph/store.h>
#include <neograph/graph/types.h>
#include <neograph/graph/validator.h>
#include <neograph/tool_dispatch.h>   // ToolGate (issue #89)
#include <neograph/tool_set.h>
#include <neograph/provider.h>
 
#include <asio/awaitable.hpp>
#include <asio/thread_pool.hpp>
 
#include <atomic>
#include <chrono>
#include <cstddef>
#include <map>
#include <memory>
#include <optional>
#include <set>
#include <string>
#include <vector>
 
namespace neograph {
class HookRuntime;
class RuntimeInterpositionController;
}
namespace neograph::graph {
 
class ValidatedTopology;
namespace detail {
struct SubgraphWriteJournal;
struct FailureRouting;
}  // namespace detail
 
class NEOGRAPH_API ManagedBudgetLeaseReleaseError final : public ProviderOutcomeError {
public:
    ManagedBudgetLeaseReleaseError(sp::runtime::Result outcome, std::exception_ptr execution_error,
                                  std::exception_ptr release_error)
        : ProviderOutcomeError("Managed budget lease release failed", std::move(outcome),
                               std::move(execution_error)),
          release_error_(std::move(release_error)) {}
    const std::exception_ptr& release_error() const noexcept { return release_error_; }
private:
    std::exception_ptr release_error_;
};
 
struct EngineConfig {
    NodeContext node_context;
 
    std::shared_ptr<CheckpointStore> checkpoint_store;
 
    std::shared_ptr<Store> store;
 
    std::optional<RetryPolicy> retry_policy;
 
    std::map<std::string, RetryPolicy> node_retry_policies;
 
    ToolGate tool_gate;
 
    std::shared_ptr<ToolExecutionController> tool_execution_controller;
 
    std::shared_ptr<::neograph::HookRuntime> hook_runtime;
 
    std::size_t worker_count = 1;
 
    ValidationOptions validation;
 
    bool stderr_notices = true;
 
    std::set<std::string> cached_nodes;
 
    std::size_t node_cache_max_entries = 0;
 
    std::map<std::string, CacheKeyPolicy> node_cache_policies;
 
    std::shared_ptr<::neograph::RuntimeInterpositionController> runtime_interposition;
    std::shared_ptr<sp::NativeArchive> native_history_archive;
 
    FailureTextRedactor failure_text_redactor;
};
 
struct EngineResources {
    ToolSet                              tools;
    std::shared_ptr<const GraphRegistry> registry;
};
 
struct RunResources {
    std::shared_ptr<CheckpointStore> checkpoint_store;
    std::shared_ptr<Store> store;
    std::optional<ToolGate> tool_gate;
    std::shared_ptr<ToolExecutionController> tool_execution_controller;
    std::shared_ptr<ProviderCallBroker> provider_call_broker;
    std::shared_ptr<ToolEffectBroker> tool_effect_broker;
    ToolEffectGrantIdentity tool_effect_grant;
    std::shared_ptr<::neograph::HookRuntime> hook_runtime;
};
 
struct SubgraphPathStep {
    std::string node_name;
    int parent_step = 0;
    std::string task_id;
    std::string parent_checkpoint_id;
};
 
struct NestedCheckpoint {
    std::vector<std::string> graph_path;
    std::string thread_id;
    Checkpoint checkpoint;
};
 
struct RunConfig {
    std::string thread_id;                          
    json        input;                              
    int         max_steps  = 50;                    
    StreamMode  stream_mode = StreamMode::ALL;      
    std::optional<std::vector<sp::Message>> provider_messages;
    std::shared_ptr<sp::NativeArchive> native_history_archive;
    std::function<void(const sp::Event&)> on_provider_event;
    std::shared_ptr<ProviderOutcomes> provider_outcomes;
    std::shared_ptr<ProviderLoopHistory> provider_loop_history;
 
    std::shared_ptr<CancelToken> cancel_token;
 
    std::shared_ptr<UsageAccumulator> usage;
 
    std::uint64_t model_token_budget = 0;
    std::shared_ptr<std::atomic_bool> budget_exhausted;
 
    bool resume_if_exists = false;
 
};
 
class NodeRetryCeiling {
public:
    explicit NodeRetryCeiling(std::uint64_t limit) noexcept : limit_(limit) {}
    NodeRetryCeiling(const NodeRetryCeiling&) = delete;
    NodeRetryCeiling& operator=(const NodeRetryCeiling&) = delete;
 
    bool try_take() noexcept {
        auto used = used_.load(std::memory_order_relaxed);
        while (used < limit_) {
            if (used_.compare_exchange_weak(used, used + 1, std::memory_order_acq_rel,
                                            std::memory_order_relaxed))
                return true;
        }
        return false;
    }
    std::uint64_t used() const noexcept { return used_.load(std::memory_order_acquire); }
    std::uint64_t limit() const noexcept { return limit_; }
 
private:
    std::uint64_t              limit_;
    std::atomic<std::uint64_t> used_{0};
};
 
struct RunMetadata {
    std::optional<std::chrono::steady_clock::time_point> deadline;
    std::string trace_id;
    std::string run_id;
    std::string owner_scope;
    std::shared_ptr<CancelToken> budget_cancel_token;
    bool shared_scheduler = false;
    asio::any_io_executor dispatch_executor;
    std::shared_ptr<NodeRetryCeiling> node_retry_ceiling;
    std::string resume_interrupt_id;
};
 
enum class RunStatus {
    Completed,
    Interrupted,
    StepLimit,
    SafePoint,
    BarrierStalled,
    FailureExit,
};
 
struct RunResult {
    json        output;                             
    bool        interrupted       = false;          
    std::string interrupt_node;                     
    json        interrupt_value;                    
    std::string checkpoint_id;                      
    std::vector<std::string> execution_trace;
 
    sp::Usage usage;
    std::vector<sp::Message> native_messages;
    std::vector<sp::runtime::Result> provider_outcomes;
 
    std::string interrupt_id;
    std::vector<PendingInterrupt> interrupts;
 
    inline bool max_steps_exhausted() const noexcept {
        if (!output.is_object()) return false;
        const auto metadata = output["_neograph"];
        return metadata.is_object() && metadata.value("max_steps_exhausted", false);
    }
 
    inline bool safe_point_reached() const noexcept {
        if (!output.is_object()) return false;
        const auto metadata = output["_neograph"];
        return metadata.is_object() && metadata.value("safe_point", false);
    }
 
    inline bool barrier_stalled() const noexcept {
        if (!output.is_object()) return false;
        const auto metadata = output["_neograph"];
        return metadata.is_object() && metadata.contains("barrier_stalled");
    }
 
    inline bool failure_exited() const noexcept {
        if (!output.is_object()) return false;
        const auto metadata = output["_neograph"];
        return metadata.is_object() && metadata.contains("failure_exit");
    }
 
    RunStatus status() const noexcept {
        if (interrupted) return RunStatus::Interrupted;
        if (safe_point_reached()) return RunStatus::SafePoint;
        if (max_steps_exhausted()) return RunStatus::StepLimit;
        if (failure_exited()) return RunStatus::FailureExit;
        if (barrier_stalled()) return RunStatus::BarrierStalled;
        return RunStatus::Completed;
    }
 
    template <typename T>
    T channel(const std::string& name) const {
        return channel_raw(name).get<T>();
    }
 
    template <typename T>
    T channel(const ChannelKey<T>& key) const {
        return channel<T>(key.name());
    }
 
    template <typename T>
    std::optional<T> try_channel(const ChannelKey<T>& key) const {
        if (!has_channel(key.name())) return std::nullopt;
        return channel(key);
    }
 
    json channel_raw(const std::string& name) const {
        if (output.contains("channels") && output["channels"].contains(name)) {
            const auto wrapped = output["channels"][name];
            if (wrapped.contains("value")) return wrapped["value"];
        }
        if (output.contains(name)) return output[name];
        throw json::out_of_range(
            "RunResult::channel: no such channel '" + name + "' in output "
            "(checked output[\"channels\"][\"" + name + "\"][\"value\"] and output[\"" + name + "\"])");
    }
 
    bool has_channel(const std::string& name) const noexcept {
        if (output.contains("channels") && output["channels"].contains(name)
            && output["channels"][name].contains("value")) {
            return true;
        }
        return output.contains(name);
    }
};
 
class NEOGRAPH_API GraphEngine {
public:
    ~GraphEngine();
 
    void set_native_history_archive(std::shared_ptr<sp::NativeArchive> archive) {
        native_history_archive_ = std::move(archive);
    }
    static std::unique_ptr<GraphEngine> link(CompiledGraph graph, EngineConfig config = {});
 
    static std::unique_ptr<GraphEngine> link(CompiledGraph   graph,
                                             EngineConfig    config,
                                             EngineResources resources);
 
    static std::unique_ptr<GraphEngine> link(
        CompiledGraph graph,
        EngineConfig config,
        EngineResources resources,
        GraphGenerationIdentity generation);
 
    static std::unique_ptr<GraphEngine> link(ValidatedTopology topology,
                                             EngineConfig config = {});
 
    static std::unique_ptr<GraphEngine> link(ValidatedTopology topology,
                                             EngineConfig config,
                                             EngineResources resources);
 
    static std::unique_ptr<GraphEngine> build(const json& definition, EngineConfig config);
 
    static std::unique_ptr<GraphEngine> build(const json&     definition,
                                               EngineConfig    config,
                                               EngineResources resources);
 
    static std::unique_ptr<GraphEngine> build_strict(const json& definition,
                                                      EngineConfig config);
 
    static std::unique_ptr<GraphEngine> build_strict(const json&     definition,
                                                      EngineConfig    config,
                                                      EngineResources resources);
 
    static std::unique_ptr<GraphEngine> compile(
        const json& definition,
        const NodeContext& default_context,
        std::shared_ptr<CheckpointStore> store = nullptr);
 
    RunResult run(const RunConfig& config);
 
    RunResult run(const RunConfig& config, const RunMetadata& metadata);
 
    asio::awaitable<RunResult> run_async(RunConfig config);
 
    asio::awaitable<RunResult> run_async(RunConfig config, RunMetadata metadata);
 
    asio::awaitable<RunResult> run_async(
        RunConfig config, RunMetadata metadata, RunResources resources);
 
    RunResult run_stream(const RunConfig& config,
                         const GraphStreamCallback& cb);
 
    RunResult run_stream(const RunConfig& config,
                         const GraphStreamCallback& cb,
                         const RunMetadata& metadata);
 
    asio::awaitable<RunResult> run_stream_async(
        RunConfig config, GraphStreamCallback cb);
 
    asio::awaitable<RunResult> run_stream_async(
        RunConfig config, GraphStreamCallback cb, RunMetadata metadata);
 
    asio::awaitable<RunResult> run_stream_async(
        RunConfig config,
        GraphStreamCallback cb,
        RunMetadata metadata,
        RunResources resources);
 
    asio::awaitable<RunResult> run_until_safe_point_async(
        RunConfig config,
        std::shared_ptr<GraphSafePointRequest> request,
        GraphStreamCallback cb = {},
        RunMetadata metadata = {},
        RunResources resources = {});
 
    RunResult resume(const std::string& thread_id,
                     const json& resume_value = json(),
                     const GraphStreamCallback& cb = nullptr);
 
    RunResult resume(const RunConfig&           config,
                     const json&                resume_value = json(),
                     const GraphStreamCallback& cb           = nullptr);
 
    RunResult resume(const RunConfig&           config,
                     const json&                resume_value,
                     const GraphStreamCallback& cb,
                     const RunMetadata&         metadata);
 
    asio::awaitable<RunResult> resume_async(
        const std::string& thread_id,
        const json& resume_value = json(),
        const GraphStreamCallback& cb = nullptr);
 
    asio::awaitable<RunResult> resume_async(RunConfig           config,
                                             json                resume_value = json(),
                                             GraphStreamCallback cb           = nullptr);
 
    asio::awaitable<RunResult> resume_async(RunConfig           config,
                                             json                resume_value,
                                             GraphStreamCallback cb,
                                             RunMetadata          metadata);
 
    asio::awaitable<RunResult> resume_async(RunConfig           config,
                                             json                resume_value,
                                             GraphStreamCallback cb,
                                             RunMetadata          metadata,
                                             RunResources        resources);
 
 
    RunResult resume_from(
        RunConfig config,
        std::string checkpoint_id,
        json resume_value = {},
        GraphStreamCallback cb = {},
        RunMetadata metadata = {},
        RunResources resources = {});
 
    asio::awaitable<RunResult> resume_from_async(
        RunConfig config,
        std::string checkpoint_id,
        json resume_value = {},
        GraphStreamCallback cb = {},
        RunMetadata metadata = {},
        RunResources resources = {});
 
    asio::awaitable<RunResult> resume_from_until_safe_point_async(
        RunConfig config,
        std::string checkpoint_id,
        std::shared_ptr<GraphSafePointRequest> request,
        json resume_value = {},
        GraphStreamCallback cb = {},
        RunMetadata metadata = {},
        RunResources resources = {});
 
    GraphAdmin admin() noexcept;
 
    // ── State inspection & manipulation (LangGraph Checkpointer API) ──
 
    std::optional<json> get_state(const std::string& thread_id) const;
 
    std::vector<Checkpoint> get_state_history(const std::string& thread_id,
                                              int limit = 100) const;
 
    std::optional<NestedCheckpoint> inspect_nested_checkpoint(
        const std::string& root_thread_id,
        const std::vector<SubgraphPathStep>& path,
        std::shared_ptr<CheckpointStore> run_checkpoint_store = {}) const;
 
    void update_state(const std::string& thread_id,
                      const json& channel_writes,
                      const std::string& as_node = "");
 
    void update_state_writes(const std::string& thread_id,
                             const std::vector<ChannelWrite>& channel_writes,
                             const std::string& as_node = "");
 
    std::string fork(const std::string& source_thread_id,
                     const std::string& new_thread_id,
                     const std::string& checkpoint_id = "");
 
    // ── Configuration ──
 
 
    void set_checkpoint_store(std::shared_ptr<CheckpointStore> store);
 
    void set_store(std::shared_ptr<Store> store);
 
    void set_tool_gate(ToolGate gate) { tool_gate_ = std::move(gate); }
 
    void set_tool_execution_controller(
        std::shared_ptr<ToolExecutionController> controller) {
        tool_execution_controller_ = std::move(controller);
    }
 
    void set_hook_runtime(std::shared_ptr<::neograph::HookRuntime> runtime) {
        hook_runtime_ = std::move(runtime);
    }
 
    void set_runtime_interposition(
        std::shared_ptr<::neograph::RuntimeInterpositionController> controller);
 
    std::shared_ptr<Store> get_store() const { return store_; }
 
    void set_retry_policy(const RetryPolicy& policy);
 
    void set_node_retry_policy(const std::string& node_name, const RetryPolicy& policy);
 
    RetryPolicy resolved_retry_policy(const std::string& node_name) const {
        return get_retry_policy(node_name);
    }
 
    void set_node_attempt_timeout_ms(const std::string& node_name, std::optional<int> timeout_ms);
 
    std::optional<int> resolved_attempt_timeout_ms(const std::string& node_name) const;
 
    void set_worker_count(std::size_t n);
 
    void set_worker_count_auto();
 
    void set_node_cache_enabled(const std::string& node_name, bool enabled);
 
    void set_node_cache_enabled(const std::string& node_name,
                                bool enabled,
                                CacheKeyPolicy policy);
 
    void clear_node_cache();
 
    void set_node_cache_max_entries(std::size_t max_entries);
 
    const NodeCache& node_cache() const { return node_cache_; }
 
    const std::string& get_graph_name() const { return name_; }
 
private:
    friend class SubgraphNode;
 
    struct SubgraphRunResult {
        RunResult result;
        std::vector<ChannelWrite> writes;
    };
 
    struct RuntimeResources {
        std::optional<std::shared_ptr<CheckpointStore>> checkpoint_store;
        std::optional<std::shared_ptr<Store>> store;
        std::optional<ToolGate> parent_tool_gate;
        std::shared_ptr<ToolExecutionController> tool_execution_controller;
        std::shared_ptr<ProviderCallBroker> provider_call_broker;
        std::shared_ptr<ToolEffectBroker> tool_effect_broker;
        ToolEffectGrantIdentity tool_effect_grant;
        std::shared_ptr<::neograph::HookRuntime> hook_runtime;
        std::shared_ptr<::neograph::HookRuntime> parent_hook_runtime;
        std::shared_ptr<detail::SubgraphWriteJournal> subgraph_write_journal;
        std::shared_ptr<GraphSafePointRequest> safe_point_request;
    };
 
    GraphEngine() = default;
    std::shared_ptr<sp::NativeArchive> native_history_archive_;
 
    static std::unique_ptr<GraphEngine> link_impl(CompiledGraph graph,
                                                   EngineConfig config,
                                                   EngineResources resources,
                                                   bool validate);
 
    ToolGate tool_gate_;
 
    std::shared_ptr<ToolExecutionController> tool_execution_controller_;
 
    std::shared_ptr<::neograph::HookRuntime> hook_runtime_;
 
    std::shared_ptr<::neograph::RuntimeInterpositionController> runtime_interposition_;
    // Admission state: -1 = exclusive administration; >=0 = active executions.
    // No heap allocation or lock on the execution fast path.
    void enter_execution();
    void leave_execution() noexcept;
    void enter_administration() const;
    void leave_administration() const noexcept;
    struct ExecutionGuard {
        explicit ExecutionGuard(GraphEngine& engine);
        ExecutionGuard(const ExecutionGuard&) = delete;
        ExecutionGuard& operator=(const ExecutionGuard&) = delete;
        ~ExecutionGuard();
        GraphEngine& engine;
    };
    struct AdministrationGuard {
        explicit AdministrationGuard(const GraphEngine& engine);
        AdministrationGuard(const AdministrationGuard&) = delete;
        AdministrationGuard& operator=(const AdministrationGuard&) = delete;
        ~AdministrationGuard();
        const GraphEngine& engine;
    };
 
    void init_state(GraphState& state) const;
    void apply_input(GraphState& state, const json& input) const;
 
    asio::awaitable<RunResult> resume_execute_async(
        RunConfig config,
        json resume_value,
        GraphStreamCallback cb,
        RunMetadata metadata,
        RuntimeResources resources,
        std::optional<std::string> checkpoint_id);
 
    asio::awaitable<RunResult> run_async_with_runtime(
        RunConfig config,
        GraphStreamCallback cb,
        RunMetadata metadata,
        RuntimeResources resources);
 
    asio::awaitable<RunResult> resume_async_with_runtime(
        RunConfig config,
        json resume_value,
        GraphStreamCallback cb,
        RunMetadata metadata,
        RuntimeResources resources,
        std::optional<std::string> checkpoint_id = std::nullopt);
 
    asio::awaitable<SubgraphRunResult> run_subgraph_async(
        RunConfig config,
        const RunContext& parent,
        GraphStreamCallback cb,
        SubgraphPersistence persistence);
 
    asio::awaitable<RunResult> execute_graph_async(
        const RunConfig& config,
        const GraphStreamCallback& cb,
        std::optional<ResumeContext> resume_context = std::nullopt,
        const json* resume_value = nullptr,
        const RunMetadata& metadata = {},
        const RuntimeResources* resources = nullptr);
 
    RetryPolicy get_retry_policy(const std::string& node_name) const;
    const GraphGenerationIdentity* bound_generation_identity() const noexcept;
    std::string managed_budget_graph_identity() const;
 
    // --- Graph definition ---
    std::string name_;
    std::string budget_graph_identity_;
 
    std::vector<ChannelDef> channel_defs_;
 
    ToolSet                             tools_;
    std::vector<std::unique_ptr<Tool>> owned_tools_;
    std::map<std::string, std::unique_ptr<GraphNode>> nodes_;
    bool has_stateful_subgraph_ = false;
    std::vector<Edge>            edges_;
    std::vector<ConditionalEdge> conditional_edges_;
 
    std::unique_ptr<Scheduler> scheduler_;
 
    std::unique_ptr<NodeExecutor> executor_;
 
    std::set<std::string> interrupt_before_;
    std::set<std::string> interrupt_after_;
 
    std::shared_ptr<CheckpointStore> checkpoint_store_;
    std::shared_ptr<Store>           store_;
 
    // Retry policies
    RetryPolicy default_retry_policy_;
    std::map<std::string, RetryPolicy> node_retry_policies_;
    std::map<std::string, NodeRetryPolicySpec> sealed_retry_policies_;
    std::shared_ptr<const detail::FailureRouting> failure_routing_;
    std::map<std::string, int> node_attempt_timeouts_;
 
    std::unique_ptr<asio::thread_pool> pool_;
 
    bool stderr_notices_ = true;
 
    NodeCache node_cache_;
    std::atomic<std::uint64_t> next_cache_execution_id_{1};
 
    mutable std::atomic<int> active_runs_{0};
 
    std::shared_ptr<detail::AttemptQuiescence> attempt_quiescence_;
};
 
} // namespace neograph::graph