settled-reach/server/src/voice/worker.rs

//! Inference worker pool (D-138, Spike 2).
//!
//! Dynamic pool of worker threads, each owning a piped stdin/stdout connection
//! to its own sr-voice child process. No network ports — the model is only
//! reachable through the game server's queue (Gemma 2 T&C compliance).

use std::collections::BTreeMap;
use std::io::{BufRead, BufReader, Write as IoWrite};
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;

use crossbeam_channel::Receiver;

use crate::npc::blueprint::CultureProfile;
use crate::voice::cache::{CacheKey, VoiceCacheStore};
use crate::voice::prompt_builder;
use crate::voice::queue::VoiceRequest;

/// Minimum token count for a valid response. Below this, retry once.
const MIN_TOKENS: usize = 4;

/// Maximum time to wait for a single inference response before killing
/// the sr-voice child process. Gemma 2B at ~16 t/s should complete 64
/// tokens in ~4s; 120s covers extreme slow hardware with margin.
const INFERENCE_TIMEOUT: Duration = Duration::from_secs(120);

/// Worker pool manages inference worker threads.
pub struct WorkerPool {
    workers: Vec<WorkerHandle>,
    shutdown: Arc<AtomicBool>,
    active_count: Arc<AtomicUsize>,
}

struct WorkerHandle {
    thread: Option<JoinHandle<()>>,
    id: usize,
}

/// Shared state passed to each worker thread.
pub struct WorkerContext {
    pub receiver: Receiver<VoiceRequest>,
    pub cache: Arc<Mutex<VoiceCacheStore>>,
    pub cultures: Arc<BTreeMap<String, CultureProfile>>,
    pub shutdown: Arc<AtomicBool>,
    pub active_count: Arc<AtomicUsize>,
    pub paused: Arc<AtomicBool>,
}

/// Configuration for spawning sr-voice child processes.
#[derive(Clone)]
pub struct VoiceProcessConfig {
    pub binary_path: String,
    pub model_path: String,
    pub threads: u32,
    pub ctx_size: u32,
}

impl WorkerPool {
    /// Spawn `count` worker threads, each owning a piped sr-voice child process.
    ///
    /// `paused` should come from `VoiceQueue::paused_flag()` so workers share
    /// the same pause signal as the queue.
    pub fn spawn(
        count: usize,
        config: &VoiceProcessConfig,
        receiver: Receiver<VoiceRequest>,
        cache: Arc<Mutex<VoiceCacheStore>>,
        cultures: Arc<BTreeMap<String, CultureProfile>>,
        paused: Arc<AtomicBool>,
    ) -> Self {
        let shutdown = Arc::new(AtomicBool::new(false));
        let active_count = Arc::new(AtomicUsize::new(0));

        let mut workers = Vec::with_capacity(count);

        for id in 0..count {
            let ctx = WorkerContext {
                receiver: receiver.clone(),
                cache: Arc::clone(&cache),
                cultures: Arc::clone(&cultures),
                shutdown: Arc::clone(&shutdown),
                active_count: Arc::clone(&active_count),
                paused: Arc::clone(&paused),
            };
            let cfg = config.clone();

            match thread::Builder::new()
                .name(format!("voice-worker-{}", id))
                .spawn(move || worker_loop(id, cfg, ctx))
            {
                Ok(thread) => {
                    workers.push(WorkerHandle {
                        thread: Some(thread),
                        id,
                    });
                }
                Err(e) => {
                    tracing::error!(id, error = %e, "failed to spawn voice worker — reducing pool");
                }
            }
        }

        tracing::info!(count, "voice worker pool started");

        Self {
            workers,
            shutdown,
            active_count,
        }
    }

    /// Number of workers currently processing a request.
    pub fn active_workers(&self) -> usize {
        self.active_count.load(Ordering::Relaxed)
    }

    /// Total number of worker threads.
    pub fn worker_count(&self) -> usize {
        self.workers.len()
    }

    /// Signal all workers to shut down and join their threads.
    pub fn shutdown(&mut self) {
        self.shutdown.store(true, Ordering::SeqCst);
        for handle in &mut self.workers {
            if let Some(thread) = handle.thread.take() {
                let _ = thread.join();
                tracing::debug!(id = handle.id, "voice worker joined");
            }
        }
    }
}

impl Drop for WorkerPool {
    fn drop(&mut self) {
        self.shutdown();
    }
}

/// A piped connection to a sr-voice child process.
struct VoicePipe {
    child: Child,
    stdin: ChildStdin,
    reader: BufReader<ChildStdout>,
}

impl VoicePipe {
    /// Spawn a sr-voice child process with piped stdin/stdout.
    fn spawn(config: &VoiceProcessConfig) -> Result<Self, String> {
        let mut child = Command::new(&config.binary_path)
            .arg("serve")
            .arg("--model")
            .arg(&config.model_path)
            .arg("--threads")
            .arg(config.threads.to_string())
            .arg("--ctx-size")
            .arg(config.ctx_size.to_string())
            .arg("--stdio")
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit())
            .spawn()
            .map_err(|e| format!("failed to spawn sr-voice: {}", e))?;

        let stdin = child.stdin.take()
            .ok_or_else(|| "failed to capture sr-voice stdin".to_string())?;
        let stdout = child.stdout.take()
            .ok_or_else(|| "failed to capture sr-voice stdout".to_string())?;

        Ok(Self {
            child,
            stdin,
            reader: BufReader::new(stdout),
        })
    }

    /// Send a prompt and read the response (JSONL: one JSON object per line).
    ///
    /// Spawns a watchdog thread that kills the child process after
    /// `INFERENCE_TIMEOUT` to prevent indefinite blocking on `read_line`.
    fn generate(&mut self, prompt: &str, seed: Option<u64>) -> Result<String, String> {
        let request = serde_json::json!({
            "prompt": prompt,
            "seed": seed,
        });

        let mut line = serde_json::to_string(&request)
            .map_err(|e| format!("failed to serialize request: {}", e))?;
        line.push('\n');

        // Check if child has already exited before writing
        if let Some(status) = self.child.try_wait().ok().flatten() {
            return Err(format!("sr-voice process exited with {}", status));
        }

        self.stdin
            .write_all(line.as_bytes())
            .map_err(|e| format!("failed to write to sr-voice stdin: {}", e))?;
        self.stdin
            .flush()
            .map_err(|e| format!("failed to flush sr-voice stdin: {}", e))?;

        // Watchdog: kill the child if it doesn't respond within the timeout.
        // This unblocks the read_line below (stdout closes → read returns empty).
        // Uses a poll loop (1s ticks) so the watchdog exits promptly on cancel.
        let child_id = self.child.id();
        let cancel = Arc::new(AtomicBool::new(false));
        let cancel_clone = Arc::clone(&cancel);
        let timeout_secs = INFERENCE_TIMEOUT.as_secs();
        let watchdog = std::thread::spawn(move || {
            for _ in 0..timeout_secs {
                std::thread::sleep(Duration::from_secs(1));
                if cancel_clone.load(Ordering::Relaxed) {
                    return;
                }
            }
            tracing::warn!(pid = child_id, "sr-voice inference timeout — killing child");
            let _ = Command::new("kill")
                .arg("-9")
                .arg(child_id.to_string())
                .output();
        });

        let mut response_line = String::new();
        let read_result = self.reader.read_line(&mut response_line);

        // Cancel the watchdog — response arrived (or EOF).
        cancel.store(true, Ordering::Relaxed);
        let _ = watchdog.join();

        read_result.map_err(|e| format!("failed to read from sr-voice stdout: {}", e))?;

        if response_line.is_empty() {
            return Err("sr-voice process closed stdout (timeout or crash)".to_string());
        }

        let body: serde_json::Value = serde_json::from_str(&response_line)
            .map_err(|e| format!("failed to parse response JSON: {}", e))?;

        if let Some(err) = body.get("error") {
            return Err(format!("sr-voice error: {}", err));
        }

        body["text"]
            .as_str()
            .map(|s| {
                // Post-processor: strip everything after the first newline.
                // Prevents prompt leakage (survey bleed, example continuation)
                // that 2B models sometimes produce after the first valid line.
                let trimmed = s.trim();
                match trimmed.find('\n') {
                    Some(pos) => trimmed[..pos].trim().to_string(),
                    None => trimmed.to_string(),
                }
            })
            .ok_or_else(|| "response missing 'text' field".to_string())
    }
}

impl Drop for VoicePipe {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

/// Main worker loop: spawn sr-voice child, receive requests, process them.
fn worker_loop(id: usize, config: VoiceProcessConfig, ctx: WorkerContext) {
    tracing::debug!(id, "voice worker starting sr-voice child process");

    let mut pipe = match VoicePipe::spawn(&config) {
        Ok(p) => {
            tracing::info!(id, "voice worker connected to sr-voice via stdio");
            p
        }
        Err(e) => {
            tracing::error!(id, error = %e, "voice worker failed to spawn sr-voice — exiting");
            return;
        }
    };

    loop {
        if ctx.shutdown.load(Ordering::SeqCst) {
            break;
        }

        // Wait for a request (with timeout so we can check shutdown)
        let request = match ctx.receiver.recv_timeout(Duration::from_secs(1)) {
            Ok(req) => req,
            Err(crossbeam_channel::RecvTimeoutError::Timeout) => continue,
            Err(crossbeam_channel::RecvTimeoutError::Disconnected) => break,
        };

        // During zone transitions the queue is paused — wait rather than
        // process (priorities may be stale). Sleep briefly and re-check.
        if ctx.paused.load(Ordering::SeqCst) {
            // Don't drop the request — sleep and retry the pause check.
            // The request stays in our local variable until we can process it.
            while ctx.paused.load(Ordering::SeqCst) {
                if ctx.shutdown.load(Ordering::SeqCst) {
                    break;
                }
                std::thread::sleep(Duration::from_millis(50));
            }
            if ctx.shutdown.load(Ordering::SeqCst) {
                break;
            }
        }

        ctx.active_count.fetch_add(1, Ordering::Relaxed);
        process_request(id, &mut pipe, &request, &ctx);
        ctx.active_count.fetch_sub(1, Ordering::Relaxed);
    }

    tracing::debug!(id, "voice worker stopped");
    // VoicePipe::drop kills the child process
}

/// Process a single voice request: build prompt → infer → validate → cache.
fn process_request(
    worker_id: usize,
    pipe: &mut VoicePipe,
    request: &VoiceRequest,
    ctx: &WorkerContext,
) {
    let culture = match ctx.cultures.get(&request.culture_id) {
        Some(c) => c,
        None => {
            tracing::warn!(
                culture_id = %request.culture_id,
                "unknown culture — serving base text"
            );
            cache_base_text(request, ctx);
            return;
        }
    };

    // Build prompt
    let built = prompt_builder::build_prompt(
        culture,
        &request.base_text,
        request.content_type,
        request.tell_state,
        request.seed,
    );

    // Call sr-voice via stdio pipe
    let result = pipe.generate(&built.prompt, Some(request.seed));

    match result {
        Ok(text) if text.split_whitespace().count() >= MIN_TOKENS => {
            cache_result(request, &text, ctx);
        }
        Ok(_short_text) => {
            // Empty output guard: retry once with different seed
            tracing::debug!(
                worker_id,
                npc = request.npc_stable_id,
                "short output — retrying with different seed"
            );
            let retry_built = prompt_builder::build_prompt(
                culture,
                &request.base_text,
                request.content_type,
                request.tell_state,
                request.seed.wrapping_add(1),
            );
            match pipe.generate(&retry_built.prompt, Some(request.seed.wrapping_add(1))) {
                Ok(text) if text.split_whitespace().count() >= MIN_TOKENS => {
                    cache_result(request, &text, ctx);
                }
                _ => {
                    tracing::debug!(
                        worker_id,
                        npc = request.npc_stable_id,
                        "retry also short — caching base text"
                    );
                    cache_base_text(request, ctx);
                }
            }
        }
        Err(e) => {
            tracing::warn!(
                worker_id,
                error = %e,
                "sr-voice request failed — serving base text"
            );
            cache_base_text(request, ctx);
        }
    }
}

/// Cache the inference result.
fn cache_result(request: &VoiceRequest, text: &str, ctx: &WorkerContext) {
    let key = cache_key_from_request(request);
    if let Ok(mut cache) = ctx.cache.lock() {
        cache.store(request.zone_id, key, text.to_string());
    }
}

/// Cache the base text as fallback (graceful degradation).
fn cache_base_text(request: &VoiceRequest, ctx: &WorkerContext) {
    let key = cache_key_from_request(request);
    if let Ok(mut cache) = ctx.cache.lock() {
        cache.store(request.zone_id, key, request.base_text.clone());
    }
}

fn cache_key_from_request(request: &VoiceRequest) -> CacheKey {
    CacheKey {
        culture_id: request.culture_id.clone(),
        npc_stable_id: request.npc_stable_id,
        content_type: request.content_type,
        content_index: request.content_index,
        tell_state: request.tell_state,
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cache_key_from_request_maps_fields() {
        let request = VoiceRequest {
            priority: crate::voice::queue::Priority::High,
            npc_stable_id: 42,
            zone_id: 100,
            culture_id: "krenn".into(),
            base_text: "Test.".into(),
            content_type: crate::voice::prompt_builder::ContentType::Dialogue,
            content_index: 5,
            tell_state: Some(crate::npc::tell_state::TellCategory::Angry),
            seed: 99,
        };
        let key = cache_key_from_request(&request);
        assert_eq!(key.culture_id, "krenn");
        assert_eq!(key.npc_stable_id, 42);
        assert_eq!(key.content_index, 5);
    }
}