使用 C++23 與協程打造極簡高效的 OpenAI 相容 AI Gateway

#pragma once
#include <coroutine>
#include <exception>
#include <stdexcept>
#include <utility>
#include <variant>

namespace aigw::coro {

template <typename T = void>
class [[nodiscard]] Task;

namespace detail {

struct TaskPromiseBase {
    std::coroutine_handle<> continuation_{nullptr};
    std::exception_ptr exception_{nullptr};

    struct FinalAwaiter {
        bool await_ready() const noexcept { return false; }
        
        template <typename PromiseType>
        std::coroutine_handle<> await_suspend(std::coroutine_handle<PromiseType> h) noexcept {
            // 當協程執行完畢，若有等待中的父協程，直接對稱轉移 (Symmetric Transfer) 喚醒父協程
            if (h.promise().continuation_) {
                return h.promise().continuation_;
            }
            return std::noop_coroutine();
        }
        
        void await_resume() noexcept {}
    };

    // Lazy 模式：建立協程時先暫停，直到被 co_await 時才開始執行
    std::suspend_always initial_suspend() noexcept { return {}; }
    FinalAwaiter final_suspend() noexcept { return {}; }
    void unhandled_exception() noexcept { exception_ = std::current_exception(); }
};

template <typename T>
struct TaskPromise : TaskPromiseBase {
    std::variant<std::monostate, T> result_;

    Task<T> get_return_object() noexcept;

    template <typename U>
        requires std::convertible_to<U, T>
    void return_value(U&& value) {
        result_.template emplace<T>(std::forward<U>(value));
    }

    T& result() {
        if (exception_) {
            std::rethrow_exception(exception_);
        }
        return std::get<T>(result_);
    }
};

// Task<void> 特化：無需儲存回傳值，提供 return_void()
template <>
struct TaskPromise<void> : TaskPromiseBase {
    Task<void> get_return_object() noexcept;
    void return_void() noexcept {}
    void result() {
        if (exception_) {
            std::rethrow_exception(exception_);
        }
    }
};

} // namespace detail

template <typename T>
class [[nodiscard]] Task {
public:
    using promise_type = detail::TaskPromise<T>;
    using handle_type = std::coroutine_handle<promise_type>;

    Task() noexcept : handle_(nullptr) {}
    explicit Task(handle_type handle) noexcept : handle_(handle) {}

    // Move-only 語意，嚴禁拷貝
    Task(const Task&) = delete;
    Task& operator=(const Task&) = delete;

    Task(Task&& other) noexcept : handle_(std::exchange(other.handle_, nullptr)) {}
    Task& operator=(Task&& other) noexcept {
        if (this != &other) {
            if (handle_) handle_.destroy();
            handle_ = std::exchange(other.handle_, nullptr);
        }
        return *this;
    }

    ~Task() {
        if (handle_) {
            handle_.destroy();
        }
    }

    struct Awaiter {
        handle_type handle_;

        bool await_ready() const noexcept {
            return !handle_ || handle_.done();
        }

        std::coroutine_handle<> await_suspend(std::coroutine_handle<> awaiting_coroutine) noexcept {
            // 將等待者 (caller) 的 handle 記錄進 promise，以便執行完畢時回調
            handle_.promise().continuation_ = awaiting_coroutine;
            return handle_; // 對稱轉移執行子協程
        }

        decltype(auto) await_resume() {
            if (!handle_) {
                throw std::runtime_error("Attempted to await invalid coroutine task");
            }
            return handle_.promise().result();
        }
    };

    auto operator co_await() const & = delete;
    auto operator co_await() && noexcept {
        return Awaiter{handle_};
    }

    handle_type release() noexcept { return std::exchange(handle_, nullptr); }

private:
    handle_type handle_;
};

namespace detail {
// 透過 std::coroutine_handle::from_promise(*this) 將 promise 反向轉換為型別安全的 handle
template <typename T>
Task<T> TaskPromise<T>::get_return_object() noexcept {
    return Task<T>{std::coroutine_handle<TaskPromise<T>>::from_promise(*this)};
}

inline Task<void> TaskPromise<void>::get_return_object() noexcept {
    return Task<void>{std::coroutine_handle<TaskPromise<void>>::from_promise(*this)};
}
} // namespace detail

} // namespace aigw::coro

struct AsyncReadable {
    int fd;
    EventLoop* loop{nullptr};

    bool await_ready() const noexcept { return false; }
    void await_suspend(std::coroutine_handle<> handle) {
        if (!loop) loop = EventLoop::current();
        if (loop) {
            loop->register_read(fd, handle);
        } else {
            throw std::runtime_error("AsyncReadable awaited outside EventLoop thread");
        }
    }
    void await_resume() const noexcept {}
};

ParseResult HttpParser::parse_request(
    const char* buf, size_t len, size_t last_len,
    HttpRequest& out_req, size_t& header_len) {

    const char* method;
    size_t method_len;
    const char* path;
    size_t path_len;
    int minor_version;
    struct phr_header headers[MAX_HEADERS];
    size_t num_headers = MAX_HEADERS;

    int pret = phr_parse_request(
        buf, len,
        &method, &method_len,
        &path, &path_len,
        &minor_version,
        headers, &num_headers,
        last_len
    );

    if (pret > 0) {
        header_len = static_cast<size_t>(pret);
        out_req.method = from_string_view(std::string_view(method, method_len));
        // ... 解析路徑、Query String 與 Headers
        return ParseResult::Complete;
    } else if (pret == -1) {
        return ParseResult::Error;
    }
    return ParseResult::Incomplete;
}

enum class RoutingError {
    NoMatchingRoute,
    PoolNotFound,
    PoolEmpty
};

inline std::string to_string(RoutingError err) {
    switch (err) {
        case RoutingError::NoMatchingRoute: return "No route matched for requested model";
        case RoutingError::PoolNotFound: return "Target upstream pool not found in configuration";
        case RoutingError::PoolEmpty: return "Upstream pool has no configured target nodes";
    }
    std::unreachable(); // C++23: 告知編譯器此路徑不可能發生，消除無效 jump branch
}

std::expected<RouteDecision, RoutingError> ModelRouter::route(const std::string& requested_model) const {
    for (const auto& r : config_.routes) {
        if (match_pattern(r.model_pattern, requested_model)) {
            auto it = pools_.find(r.pool);
            if (it != pools_.end()) {
                if (it->second->nodes().empty()) {
                    return std::unexpected(RoutingError::PoolEmpty);
                }
                RouteDecision decision;
                decision.pool = it->second;
                decision.original_model = requested_model;
                decision.model = r.rewrite_model.empty() ? requested_model : r.rewrite_model;
                return decision;
            } else {
                return std::unexpected(RoutingError::PoolNotFound);
            }
        }
    }

    if (default_pool_) {
        if (default_pool_->nodes().empty()) {
            return std::unexpected(RoutingError::PoolEmpty);
        }
        RouteDecision decision;
        decision.pool = default_pool_;
        decision.original_model = requested_model;
        decision.model = requested_model;
        return decision;
    }

    return std::unexpected(RoutingError::NoMatchingRoute);
}

std::shared_ptr<UpstreamNode> UpstreamPool::select_node() {
    auto candidates = get_healthy_candidates();
    if (candidates.empty()) return nullptr;

    int total_weight = 0;
    for (const auto& node : candidates) {
        total_weight += std::max(1, node->weight());
    }

    size_t idx = rr_index_.fetch_add(1, std::memory_order_relaxed) % static_cast<size_t>(total_weight);
    int cumulative = 0;
    for (const auto& node : candidates) {
        cumulative += std::max(1, node->weight());
        if (idx < static_cast<size_t>(cumulative)) {
            return node;
        }
    }
    return candidates.front();
}

[package]
name = "ai-gateway"
version = "0.1.0"
cxx-standard = "c++23"

[dependencies]
nlohmann_json = { port = true, version = "^3.12.0" }
picohttpparser = { port = true, version = "^2026.4.6" }
spdlog = { port = true, version = "^1.17.0" }

[dev-dependencies]
catch2 = { port = true, version = "^3.15.1" }

// ❌ 危險：若 message 包含引號、反斜線或控制字元，將造成 JSON 語法損壞
std::string json = "{\"error\":{\"message\":\"" + message + "\",...}}";

// ✅ 安全：自動正確跳脫所有特殊字元
nlohmann::json j;
j["error"]["message"] = message;
j["error"]["type"] = type;
j["error"]["code"] = "invalid_request_error";
return make_json(status, j.dump());

volatile sig_atomic_t g_shutdown_requested = 0;

void signal_handler(int sig) {
    (void)sig;
    g_shutdown_requested = 1;
    if (g_wakeup_fd >= 0) {
        uint64_t val = 1;
        ::write(g_wakeup_fd, &val, sizeof(val)); // async-signal-safe
    }
}

std::signal(SIGPIPE, SIG_IGN);

if (s_hdr.is_chunked) {
    resp_headers.set("Transfer-Encoding", "chunked");
}
resp_headers.set("Content-Type", "text/event-stream");
resp_headers.set("Cache-Control", "no-cache");
resp_headers.set("Connection", "keep-alive");

# 編譯專案
cabin build

# 執行單元測試
cabin test

# 啟動閘道服務
cabin run -- --port 8080 --log-level debug --config config.json

{
  "host": "0.0.0.0",
  "port": 8080,
  "log_level": "info",
  "default_pool": "local_ollama",
  "pools": [
    {
      "name": "local_ollama",
      "timeout_ms": 60000,
      "max_retries": 2,
      "targets": [
        { "host": "127.0.0.1", "port": 11434, "weight": 1 }
      ]
    }
  ],
  "routes": [
    {
      "model": "gemma*",
      "pool": "local_ollama"
    },
    {
      "model": "llama*",
      "pool": "local_ollama"
    }
  ]
}

from openai import OpenAI
import sys

# 將 base_url 指向本機 C++23 AI Gateway
client = OpenAI(
    base_url="http://localhost:8080/v1",
    api_key="none"
)

# 測試即時串流對話 (使用具備 Reasoning 能力的 gemma4:26b)
response = client.chat.completions.create(
    model="gemma4:26b",
    messages=[{"role": "user", "content": "請用繁體中文以兩句話解釋 C++20 協程。"}],
    stream=True
)

in_thinking = False

for chunk in response:
    if not chunk.choices:
        continue
    delta = chunk.choices[0].delta
    
    # 處理思考/推理輸出 (Thinking / Reasoning tokens)
    reasoning = getattr(delta, "reasoning", None) or getattr(delta, "reasoning_content", None)
    if reasoning:
        if not in_thinking:
            print("\033[90m[思考中...]\033[0m ", end="", flush=True)
            in_thinking = True
        print(f"\033[90m{reasoning}\033[0m", end="", flush=True)
    elif delta.content:
        if in_thinking:
            print("\n\n\033[32m[回答]\033[0m\n", end="", flush=True)
            in_thinking = False
        print(delta.content, end="", flush=True)

print()