Gemini + Grok

SelfCap.h

#ifndef SELF_CAP_H
#define SELF_CAP_H

#include <windows.h>
#include <io.h> // _dup2
#include <fcntl.h> // _O_TEXT
#include <queue>
#include <mutex>
#include <thread>
#include <atomic>
#include <memory>
#include <string>

class SelfOutputCapturer final {
public:
    using OutputCallback = void(*)(std::string_view line, bool is_stderr);

    explicit SelfOutputCapturer(OutputCallback cb = nullptr)
        : callback_(cb) {
        stdout_buffer_ = std::make_unique<char[]>(BUFFER_SIZE);
        stderr_buffer_ = std::make_unique<char[]>(BUFFER_SIZE);
        StartCapture();
    }

    ~SelfOutputCapturer() noexcept { Stop(); }

    void SetCallback(OutputCallback new_cb) noexcept {
        std::lock_guard<std::mutex> lock(callback_mutex_);
        callback_ = new_cb;
    }

    void Stop() noexcept;

private:
    HANDLE original_stdout_handle_{ INVALID_HANDLE_VALUE };
    HANDLE original_stderr_handle_{ INVALID_HANDLE_VALUE };
    HANDLE pipe_read_stdout_{ INVALID_HANDLE_VALUE };
    HANDLE pipe_write_stdout_{ INVALID_HANDLE_VALUE };
    HANDLE pipe_read_stderr_{ INVALID_HANDLE_VALUE };
    HANDLE pipe_write_stderr_{ INVALID_HANDLE_VALUE };

    int original_stdout_fd_{ -1 };
    int original_stderr_fd_{ -1 };

    std::atomic<bool> running_{ true };
    std::once_flag stop_once_;

    std::thread stdout_reader_;
    std::thread stderr_reader_;
    std::thread callback_dispatcher_;

    OutputCallback callback_ = nullptr;
    std::mutex callback_mutex_;

    struct LineItem {
        std::string line;
        bool is_stderr{ false };
    };
    std::queue<LineItem> line_queue_;
    std::mutex queue_mutex_;
    std::condition_variable queue_cv_;

    static constexpr size_t BUFFER_SIZE = 32768;
    static constexpr size_t MAX_QUEUE_SIZE = 10000;

    std::unique_ptr<char[]> stdout_buffer_;
    std::unique_ptr<char[]> stderr_buffer_;

    void StartCapture();
    void CleanupPipes() noexcept;
    void ReaderThread(bool is_stdout) noexcept;
    void ProcessBuffer(std::string& remainder, bool is_stderr, bool force_flush = false);
    void DispatchCallback(std::string&& line, bool is_stderr) noexcept;
    void CallbackDispatcherThread() noexcept;
};

void SelfOutputCapturer::StartCapture() {
    SECURITY_ATTRIBUTES sa{};
    sa.nLength = sizeof(sa);
    sa.bInheritHandle = FALSE;

    if (!CreatePipe(&pipe_read_stdout_, &pipe_write_stdout_, &sa, 0) ||
        !CreatePipe(&pipe_read_stderr_, &pipe_write_stderr_, &sa, 0)) {
        CleanupPipes();
        throw std::runtime_error("CreatePipe failed");
    }

    original_stdout_handle_ = GetStdHandle(STD_OUTPUT_HANDLE);
    original_stderr_handle_ = GetStdHandle(STD_ERROR_HANDLE);

    original_stdout_fd_ = _dup(_fileno(stdout));
    original_stderr_fd_ = _dup(_fileno(stderr));
    if (original_stdout_fd_ == -1 || original_stderr_fd_ == -1) {
        CleanupPipes();
        throw std::runtime_error("_dup failed");
    }

    if (!SetStdHandle(STD_OUTPUT_HANDLE, pipe_write_stdout_) ||
        !SetStdHandle(STD_ERROR_HANDLE, pipe_write_stderr_)) {
        CleanupPipes();
        throw std::runtime_error("SetStdHandle failed");
    }

    int fd_stdout = _open_osfhandle((intptr_t)pipe_write_stdout_, _O_BINARY);
    int fd_stderr = _open_osfhandle((intptr_t)pipe_write_stderr_, _O_BINARY);
    if (fd_stdout == -1 || fd_stderr == -1 ||
        _dup2(fd_stdout, _fileno(stdout)) == -1 ||
        _dup2(fd_stderr, _fileno(stderr)) == -1) {
        CleanupPipes();
        throw std::runtime_error("_dup2 failed");
    }
    _close(fd_stdout);
    _close(fd_stderr);

    stdout_reader_ = std::thread([this] { ReaderThread(true); });
    stderr_reader_ = std::thread([this] { ReaderThread(false); });
    callback_dispatcher_ = std::thread([this] { CallbackDispatcherThread(); });
}

void SelfOutputCapturer::Stop() noexcept {
    std::call_once(stop_once_, [this] {
        running_.store(false, std::memory_order_release);

        if (original_stdout_fd_ != -1) {
            _dup2(original_stdout_fd_, _fileno(stdout));
            _close(original_stdout_fd_);
            original_stdout_fd_ = -1;
        }
        if (original_stderr_fd_ != -1) {
            _dup2(original_stderr_fd_, _fileno(stderr));
            _close(original_stderr_fd_);
            original_stderr_fd_ = -1;
        }

        if (original_stdout_handle_ != INVALID_HANDLE_VALUE)
            SetStdHandle(STD_OUTPUT_HANDLE, original_stdout_handle_);
        if (original_stderr_handle_ != INVALID_HANDLE_VALUE)
            SetStdHandle(STD_ERROR_HANDLE, original_stderr_handle_);

        if (pipe_write_stdout_ != INVALID_HANDLE_VALUE) CloseHandle(pipe_write_stdout_);
        if (pipe_write_stderr_ != INVALID_HANDLE_VALUE) CloseHandle(pipe_write_stderr_);

        if (stdout_reader_.joinable()) stdout_reader_.join();
        if (stderr_reader_.joinable()) stderr_reader_.join();

        queue_cv_.notify_all();

        if (callback_dispatcher_.joinable()) callback_dispatcher_.join();

        CleanupPipes();
    });
}

void SelfOutputCapturer::CleanupPipes() noexcept {
    if (pipe_read_stdout_ != INVALID_HANDLE_VALUE) CloseHandle(pipe_read_stdout_);
    if (pipe_write_stdout_ != INVALID_HANDLE_VALUE) CloseHandle(pipe_write_stdout_);
    if (pipe_read_stderr_ != INVALID_HANDLE_VALUE) CloseHandle(pipe_read_stderr_);
    if (pipe_write_stderr_ != INVALID_HANDLE_VALUE) CloseHandle(pipe_write_stderr_);

    pipe_read_stdout_ = pipe_write_stdout_ = pipe_read_stderr_ = pipe_write_stderr_ = INVALID_HANDLE_VALUE;
}

void SelfOutputCapturer::ReaderThread(bool is_stdout) noexcept {
    HANDLE hRead = is_stdout ? pipe_read_stdout_ : pipe_read_stderr_;
    char* buffer = is_stdout ? stdout_buffer_.get() : stderr_buffer_.get();
    std::string remainder;

    while (running_.load(std::memory_order_relaxed)) {
        DWORD bytes_read = 0;
        if (!ReadFile(hRead, buffer, BUFFER_SIZE - 1, &bytes_read, nullptr)) {
            DWORD err = GetLastError();
            if (err == ERROR_BROKEN_PIPE || err == ERROR_HANDLE_EOF) break;
            break;
        }
        if (bytes_read == 0) continue;

        remainder.append(buffer, bytes_read);
        ProcessBuffer(remainder, !is_stdout);
    }

    if (!remainder.empty()) {
        ProcessBuffer(remainder, !is_stdout, true);
    }
}

void SelfOutputCapturer::ProcessBuffer(std::string& remainder, bool is_stderr, bool force_flush) {
    size_t pos = 0;
    const size_t size = remainder.size();

    while (pos < size) {
        size_t newline_pos = remainder.find('\n', pos);

        if (newline_pos == std::string::npos) {
            if (force_flush && pos < size) {
                std::string line = remainder.substr(pos);
                if (!line.empty()) {
                    DispatchCallback(std::move(line), is_stderr);
                }
                remainder.clear();
            }
            break;
        }

        bool has_cr = (newline_pos > pos && remainder[newline_pos - 1] == '\r');
        size_t line_end = has_cr ? newline_pos - 1 : newline_pos;

        std::string line = remainder.substr(pos, line_end - pos);
        if (!line.empty()) {
            DispatchCallback(std::move(line), is_stderr);
        }

        pos = newline_pos + 1;
        if (has_cr && pos < size && remainder[pos] == '\n') {
            ++pos;
        }
    }

    if (pos > 0) {
        if (pos < remainder.size()) {
            std::memmove(&remainder[0], &remainder[pos], remainder.size() - pos);
            remainder.resize(remainder.size() - pos);
        } else {
            remainder.clear();
        }
    }
}

void SelfOutputCapturer::DispatchCallback(std::string&& line, bool is_stderr) noexcept {
    if (line.empty() || !callback_) return;

    {
        std::unique_lock<std::mutex> lock(queue_mutex_);

        queue_cv_.wait(lock, [this] {
            return line_queue_.size() < MAX_QUEUE_SIZE || !running_.load(std::memory_order_acquire);
        });

        if (!running_.load()) return;

        line_queue_.push({ std::move(line), is_stderr });
    }
    queue_cv_.notify_one();
}

void SelfOutputCapturer::CallbackDispatcherThread() noexcept {
    while (true) {
        LineItem item;

        {
            std::unique_lock<std::mutex> lock(queue_mutex_);
            queue_cv_.wait(lock, [this] {
                return !line_queue_.empty() || !running_.load(std::memory_order_acquire);
            });

            if (line_queue_.empty()) {
                break;
            }

            item = std::move(line_queue_.front());
            line_queue_.pop();
        }

        {
            std::lock_guard<std::mutex> cb_lock(callback_mutex_);
            if (callback_) {
                callback_(item.line, item.is_stderr);
            }
        }
    }

    while (true) {
        LineItem item;
        {
            std::lock_guard<std::mutex> lock(queue_mutex_);
            if (line_queue_.empty()) break;
            item = std::move(line_queue_.front());
            line_queue_.pop();
        }

        {
            std::lock_guard<std::mutex> cb_lock(callback_mutex_);
            if (callback_) {
                callback_(item.line, item.is_stderr);
            }
        }
    }
}

#endif // SELF_CAP_H