Genarrative/apps/ai-game-creator-shell/src-tauri/src/runner/server.rs

use super::{dispatch::*, endpoint::*, protocol::*, state::*};
use sha2::{Digest as _, Sha256};
use std::fs;
use std::io;
use std::net::{Ipv4Addr, SocketAddrV4, TcpListener};
use std::path::{Path, PathBuf};
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

pub(super) fn refresh_external_agent_runner_heartbeat(
    state: &ExternalAgentRunnerServerState,
) -> Result<(), String> {
    let endpoint = {
        let mut endpoint = lock_unpoisoned(&state.endpoint);
        endpoint.heartbeat_at = unix_millis();
        endpoint.clone()
    };
    write_external_agent_runner_endpoint_atomic(&state.endpoint_path, &endpoint)
}

pub(super) fn bind_external_agent_runner_listener_with<T>(
    mut fallback_ports: impl FnMut() -> Vec<u16>,
    mut bind: impl FnMut(u16) -> io::Result<T>,
) -> io::Result<T> {
    let primary_error = match bind(0) {
        Ok(listener) => return Ok(listener),
        Err(error) => error,
    };
    if primary_error.kind() != io::ErrorKind::AddrInUse {
        return Err(primary_error);
    }
    for port in fallback_ports() {
        match bind(port) {
            Ok(listener) => return Ok(listener),
            Err(error) if error.kind() == io::ErrorKind::AddrInUse => {}
            Err(error) => return Err(error),
        }
    }
    Err(primary_error)
}

#[cfg(target_os = "linux")]
pub(super) fn parse_external_agent_runner_linux_ephemeral_port_range(
    content: &str,
) -> Option<(u16, u16)> {
    let mut values = content.split_whitespace();
    let start = values.next()?.parse::<u16>().ok()?;
    let end = values.next()?.parse::<u16>().ok()?;
    if values.next().is_some() || start > end {
        return None;
    }
    Some((start, end))
}

#[cfg(target_os = "linux")]
pub(super) fn parse_external_agent_runner_linux_single_port(content: &str) -> Option<u16> {
    let mut values = content.split_whitespace();
    let value = values.next()?.parse::<u16>().ok()?;
    values.next().is_none().then_some(value)
}

#[cfg(target_os = "linux")]
pub(super) fn parse_external_agent_runner_linux_reserved_ports(
    content: &str,
) -> Option<Vec<(u16, u16)>> {
    let content = content.trim();
    if content.is_empty() {
        return Some(Vec::new());
    }
    let mut ranges = Vec::new();
    for part in content.split(',') {
        let part = part.trim();
        if part.is_empty() {
            return None;
        }
        let mut bounds = part.split('-');
        let start = bounds.next()?.parse::<u16>().ok()?;
        let end = match bounds.next() {
            Some(value) => value.parse::<u16>().ok()?,
            None => start,
        };
        if bounds.next().is_some() || start > end {
            return None;
        }
        ranges.push((start, end));
    }
    Some(ranges)
}

#[cfg(target_os = "linux")]
pub(super) fn external_agent_runner_linux_fallback_ports(
    boot_id: &str,
    (ephemeral_start, ephemeral_end): (u16, u16),
    unprivileged_port_start: u16,
    reserved_ports: &[(u16, u16)],
) -> Vec<u16> {
    let start = EXTERNAL_AGENT_RUNNER_FALLBACK_PORT_START.max(unprivileged_port_start);
    let mut ports = (start..=u16::MAX)
        .filter(|port| {
            !(ephemeral_start..=ephemeral_end).contains(port)
                && !reserved_ports.iter().any(|(reserved_start, reserved_end)| {
                    (*reserved_start..=*reserved_end).contains(port)
                })
        })
        .collect::<Vec<_>>();
    if !ports.is_empty() {
        let digest = Sha256::digest(boot_id.as_bytes());
        let seed = u64::from_be_bytes([
            digest[0], digest[1], digest[2], digest[3], digest[4], digest[5], digest[6], digest[7],
        ]);
        let offset = (seed % ports.len() as u64) as usize;
        ports.rotate_left(offset);
    }
    ports
}

#[cfg(target_os = "linux")]
pub(super) fn read_external_agent_runner_linux_fallback_ports(boot_id: &str) -> Option<Vec<u16>> {
    let ephemeral_range = fs::read_to_string(EXTERNAL_AGENT_RUNNER_LINUX_EPHEMERAL_PORT_RANGE_PATH)
        .ok()
        .and_then(|content| parse_external_agent_runner_linux_ephemeral_port_range(&content))?;
    let unprivileged_port_start =
        fs::read_to_string(EXTERNAL_AGENT_RUNNER_LINUX_UNPRIVILEGED_PORT_START_PATH)
            .ok()
            .and_then(|content| parse_external_agent_runner_linux_single_port(&content))?;
    let reserved_ports = fs::read_to_string(EXTERNAL_AGENT_RUNNER_LINUX_RESERVED_PORTS_PATH)
        .ok()
        .and_then(|content| parse_external_agent_runner_linux_reserved_ports(&content))?;
    Some(external_agent_runner_linux_fallback_ports(
        boot_id,
        ephemeral_range,
        unprivileged_port_start,
        &reserved_ports,
    ))
}

pub(crate) fn bind_loopback_listener_with_linux_fallback(seed: &str) -> io::Result<TcpListener> {
    #[cfg(target_os = "linux")]
    {
        return bind_external_agent_runner_listener_with(
            || read_external_agent_runner_linux_fallback_ports(seed).unwrap_or_default(),
            |port| TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)),
        );
    }

    #[cfg(not(target_os = "linux"))]
    {
        let _ = seed;
        TcpListener::bind(SocketAddrV4::new(Ipv4Addr::LOCALHOST, 0))
    }
}

#[cfg(test)]
pub(super) fn external_agent_runner_shutdown_if_gui_owner_lost(
    state: &ExternalAgentRunnerServerState,
) -> Result<bool, String> {
    if !state.gui_owner_attached.load(Ordering::Acquire) {
        return Ok(false);
    }
    if external_agent_runner_gui_owner_is_locked(&state.gui_owner_lock_path)? {
        return Ok(false);
    }
    request_external_agent_runner_forced_shutdown(state);
    Ok(true)
}

pub(super) fn spawn_external_agent_runner_gui_owner_watchdog(
    state: Arc<ExternalAgentRunnerServerState>,
    endpoint_path: PathBuf,
    boot_id: String,
) -> Result<(), String> {
    thread::Builder::new()
        .name("agent-runner-gui-owner-watchdog".to_string())
        .spawn(move || loop {
            if !state.gui_owner_attached.load(Ordering::Acquire) {
                thread::sleep(EXTERNAL_AGENT_RUNNER_GUI_OWNER_CHECK_INTERVAL);
                continue;
            }
            let owner_lost =
                match external_agent_runner_gui_owner_is_locked(&state.gui_owner_lock_path) {
                    Ok(locked) => !locked,
                    Err(_) => true,
                };
            if !owner_lost {
                thread::sleep(EXTERNAL_AGENT_RUNNER_GUI_OWNER_CHECK_INTERVAL);
                continue;
            }

            state.draining.store(true, Ordering::Release);
            state
                .force_shutdown_requested
                .store(true, Ordering::Release);
            state.shutdown_requested.store(true, Ordering::Release);
            thread::sleep(EXTERNAL_AGENT_RUNNER_GUI_OWNER_WATCHDOG_HARD_EXIT_TIMEOUT);
            let _ = crate::agent::shutdown_game_creator_codex_app_servers();
            remove_external_agent_runner_endpoint_if_boot_matches(&endpoint_path, &boot_id);
            std::process::exit(1);
        })
        .map(|_| ())
        .map_err(|error| format!("启动 Agent Runner GUI owner watchdog 失败：{error}"))
}

pub(super) fn resolve_external_agent_runner_initial_gui_owner(
    gui_owner_required: bool,
    gui_owner_present: bool,
) -> Result<bool, String> {
    if gui_owner_required && !gui_owner_present {
        return Err("GUI owner 在 Agent Runner 启动完成前已释放".to_string());
    }
    Ok(gui_owner_present)
}

pub(crate) fn run_external_agent_runner_server(
    config_dir: impl AsRef<Path>,
    gui_owner_required: bool,
) -> Result<(), String> {
    let config_dir = normalize_external_agent_runner_config_dir(config_dir.as_ref())?;
    EXTERNAL_AGENT_RUNNER_SERVER_PROCESS.store(true, Ordering::Release);
    crate::set_game_creator_runtime_config_dir(config_dir.clone());
    set_external_agent_runner_config_dir(config_dir.clone());

    let boot_id = random_identifier(b"genarrative-agent-runner-boot-id")?;
    crate::initialize_process_session_boot_id(&boot_id)?;
    let token = random_identifier(b"genarrative-agent-runner-token")?;
    let _instance_lock = acquire_external_agent_runner_instance_lock(
        &external_agent_runner_lock_path(&config_dir),
        &boot_id,
    )?;
    let gui_owner_present_at_start = resolve_external_agent_runner_initial_gui_owner(
        gui_owner_required,
        external_agent_runner_gui_owner_is_locked(&external_agent_runner_gui_owner_lock_path(
            &config_dir,
        ))?,
    )?;
    let executable_fingerprint = current_external_agent_runner_executable_fingerprint()?;
    let listener = bind_loopback_listener_with_linux_fallback(&boot_id)
        .map_err(|error| format!("绑定 Agent Runner loopback 端口失败：{error}"))?;
    listener
        .set_nonblocking(true)
        .map_err(|error| format!("配置 Agent Runner listener 失败：{error}"))?;
    let port = listener
        .local_addr()
        .map_err(|error| format!("读取 Agent Runner loopback 地址失败：{error}"))?
        .port();
    let endpoint = ExternalAgentRunnerEndpoint {
        protocol_version: EXTERNAL_AGENT_RUNNER_PROTOCOL_VERSION,
        pid: std::process::id(),
        boot_id: boot_id.clone(),
        port,
        token,
        heartbeat_at: unix_millis(),
        executable_fingerprint: Some(executable_fingerprint),
        process_start_identity: external_agent_runner_process_start_identity(std::process::id())?,
    };
    let endpoint_path = external_agent_runner_endpoint_path(&config_dir);
    write_external_agent_runner_endpoint_atomic(&endpoint_path, &endpoint)?;
    let _endpoint_guard = ExternalAgentRunnerEndpointGuard {
        path: endpoint_path.clone(),
        boot_id,
    };
    let state = Arc::new(ExternalAgentRunnerServerState::new(endpoint_path, endpoint));
    state
        .gui_owner_attached
        .store(gui_owner_present_at_start, Ordering::Release);
    spawn_external_agent_runner_gui_owner_watchdog(
        Arc::clone(&state),
        state.endpoint_path.clone(),
        state.endpoint_snapshot().boot_id,
    )?;
    let mut last_heartbeat = Instant::now();
    let mut server_error = None;

    loop {
        if state.shutdown_requested.load(Ordering::Acquire) {
            if state.force_shutdown_requested.load(Ordering::Acquire)
                || state.active_connections.load(Ordering::Acquire) == 0
            {
                break;
            }
            thread::sleep(EXTERNAL_AGENT_RUNNER_LOOP_INTERVAL);
            continue;
        }

        match listener.accept() {
            Ok((stream, _)) => {
                let previous = state.active_connections.fetch_add(1, Ordering::AcqRel);
                if previous >= EXTERNAL_AGENT_RUNNER_MAX_CONNECTIONS {
                    state.active_connections.fetch_sub(1, Ordering::AcqRel);
                    drop(stream);
                    continue;
                }
                let worker_state = Arc::clone(&state);
                if thread::Builder::new()
                    .name("agent-runner-connection".to_string())
                    .spawn(move || {
                        let _ = handle_external_agent_runner_connection(stream, worker_state);
                    })
                    .is_err()
                {
                    state.active_connections.fetch_sub(1, Ordering::AcqRel);
                }
            }
            Err(error) if error.kind() == io::ErrorKind::WouldBlock => {}
            Err(error) => {
                server_error = Some(format!("接受 Agent Runner 连接失败：{error}"));
                break;
            }
        }

        if last_heartbeat.elapsed() >= EXTERNAL_AGENT_RUNNER_HEARTBEAT_INTERVAL {
            if let Err(error) = refresh_external_agent_runner_heartbeat(&state) {
                server_error = Some(error);
                break;
            }
            last_heartbeat = Instant::now();
        }
        thread::sleep(EXTERNAL_AGENT_RUNNER_LOOP_INTERVAL);
    }

    state.shutdown_requested.store(true, Ordering::Release);
    let forced = state.force_shutdown_requested.load(Ordering::Acquire);
    let forced_deadline =
        forced.then(|| Instant::now() + EXTERNAL_AGENT_RUNNER_FORCED_TOTAL_DRAIN_TIMEOUT);
    let worker_deadline = Instant::now()
        + if forced {
            EXTERNAL_AGENT_RUNNER_FORCED_WORKER_DRAIN_TIMEOUT
        } else {
            EXTERNAL_AGENT_RUNNER_IO_TIMEOUT
        };
    while state.active_connections.load(Ordering::Acquire) > 0 && Instant::now() < worker_deadline {
        thread::sleep(EXTERNAL_AGENT_RUNNER_LOOP_INTERVAL);
    }
    let forced_roots = forced.then(|| {
        let roots = state.known_roots_snapshot();
        crate::interrupt_game_creator_agent_runtime_provider_requests_for_roots(&roots);
        roots
    });
    let process_timeout = forced_deadline
        .map(|deadline| deadline.saturating_duration_since(Instant::now()))
        .unwrap_or(Duration::from_secs(3));
    let mut process_shutdown = crate::shutdown_all_process_sessions_and_wait(process_timeout);
    if forced {
        let roots = forced_roots.as_deref().unwrap_or_default();
        let provider_deadline = forced_deadline.expect("forced shutdown has a deadline");
        while crate::game_creator_agent_runtime_provider_request_count_for_roots(roots) > 0
            && Instant::now() < provider_deadline
        {
            thread::sleep(EXTERNAL_AGENT_RUNNER_LOOP_INTERVAL);
        }
        if crate::game_creator_agent_runtime_provider_request_count_for_roots(roots) > 0 {
            let provider_error = "Runner 退出前未能中断全部 Provider 请求".to_string();
            process_shutdown = Err(match process_shutdown {
                Ok(()) => provider_error,
                Err(process_error) => format!("{process_error}；{provider_error}"),
            });
        }
    }
    if let Err(app_server_error) = crate::agent::shutdown_game_creator_codex_app_servers() {
        process_shutdown = Err(match process_shutdown {
            Ok(()) => app_server_error,
            Err(process_error) => format!("{process_error}；{app_server_error}"),
        });
    }
    if let Some(error) = server_error {
        Err(match process_shutdown {
            Ok(()) => error,
            Err(process_error) => format!("{error}；{process_error}"),
        })
    } else {
        process_shutdown
    }
}