ForgeLM/crates/inference-core/src/session.rs

//! GenerationSession: single-request token-by-token generation loop.
//!
//! Flow:
//!   1. Tokenize the prompt (including system + user messages).
//!   2. Prefill: feed all prompt tokens through llama_decode() in one batch.
//!   3. Generation loop: sample → decode → yield token text → check stop.
//!
//! The session borrows raw pointers from InferenceEngine and must not outlive it.
//! In Phase 7+ this will be refactored to use a shared context pool.

use std::ffi::{CStr, CString};
use std::ops::Range;
use std::sync::Arc;
use std::time::Instant;

use anyhow::{bail, Context, Result};
use tracing::{debug, info};

use llama_sys::{llama_context, llama_token, llama_vocab, LlamaLib};

use crate::sampler::SamplerConfig;

// ─── GenerationSession ─────────────────────────────────────────────────────

pub struct GenerationSession {
    lib: Arc<LlamaLib>,
    ctx: *mut llama_context,
    vocab: *const llama_vocab,
    sampler: *mut llama_sys::llama_sampler,
    stop_sequences: Vec<String>,

    // generation state
    n_decoded: i32, // number of tokens decoded so far (prefill + gen)
    prompt_tokens: u32,
    pub max_tokens: u32,
    pub tokens_generated: u32,
    pub finished: bool,

    // timing
    pub prefill_ms: Option<f64>,
    gen_start: Option<Instant>,
    pub tokens_per_sec: f64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
struct PrefillChunkPlan {
    range: Range<usize>,
    logits_index: Option<usize>,
}

// SAFETY: pointers come from InferenceEngine which is Send.
unsafe impl Send for GenerationSession {}

impl GenerationSession {
    /// Tokenize prompt and run prefill decode. Returns a session ready to generate.
    pub fn new(
        _lib: Arc<LlamaLib>,
        _ctx: *mut llama_context,
        _sampler_cfg: SamplerConfig,
    ) -> Result<Self> {
        // Get vocab from context's model (via a workaround: we pass vocab separately)
        // For now we'll get it indirectly — see Engine::new_session which passes vocab.
        // This constructor signature will be refined; use the vocab variant below.
        bail!("Use GenerationSession::with_vocab() instead")
    }

    /// Full constructor called by InferenceEngine::new_session().
    pub fn with_vocab(
        lib: Arc<LlamaLib>,
        ctx: *mut llama_context,
        vocab: *const llama_vocab,
        n_batch: u32,
        prompt: &str,
        sampler_cfg: SamplerConfig,
    ) -> Result<Self> {
        // ── 1. Tokenize prompt ─────────────────────────────────────────────
        let prompt_c = CString::new(prompt).context("Invalid UTF-8 in prompt")?;

        // First call: get token count
        let ret = unsafe {
            (lib.tokenize)(
                vocab,
                prompt_c.as_ptr(),
                prompt.len() as i32,
                std::ptr::null_mut(),
                0,
                true, // add_special (BOS)
                true, // parse_special
            )
        };

        let n_tokens_needed = if ret < 0 { -ret } else { ret };

        if n_tokens_needed == 0 {
            bail!("llama_tokenize failed to return token count");
        }

        let mut tokens: Vec<llama_token> = vec![0; n_tokens_needed as usize];
        let n_tokens = unsafe {
            (lib.tokenize)(
                vocab,
                prompt_c.as_ptr(),
                prompt.len() as i32,
                tokens.as_mut_ptr(),
                tokens.len() as i32,
                true,
                true,
            )
        };

        if n_tokens < 0 {
            bail!(
                "llama_tokenize (second call) failed (returned {})",
                n_tokens
            );
        }
        tokens.truncate(n_tokens as usize);

        let n_prompt_tokens = tokens.len() as i32;
        debug!(n_tokens = n_prompt_tokens, "Prompt tokenised");

        // ── 2. Build sampler chain ─────────────────────────────────────────
        let mut sc_params = unsafe { (lib.sampler_chain_default_params)() };
        sc_params.no_perf = true;
        let sampler_chain = unsafe { (lib.sampler_chain_init)(sc_params) };
        if sampler_chain.is_null() {
            bail!("sampler_chain_init returned NULL");
        }

        unsafe {
            if sampler_cfg.temperature <= 0.0 || sampler_cfg.temperature < 1e-6 {
                // Greedy
                let s = (lib.sampler_init_greedy)();
                (lib.sampler_chain_add)(sampler_chain, s);
            } else {
                // top-k → temp → top-p → dist (for randomness)
                let s_topk = (lib.sampler_init_top_k)(sampler_cfg.top_k);
                (lib.sampler_chain_add)(sampler_chain, s_topk);

                let s_temp = (lib.sampler_init_temp)(sampler_cfg.temperature);
                (lib.sampler_chain_add)(sampler_chain, s_temp);

                let s_topp = (lib.sampler_init_top_p)(sampler_cfg.top_p, 1);
                (lib.sampler_chain_add)(sampler_chain, s_topp);

                let s_dist = (lib.sampler_init_dist)(sampler_cfg.seed);
                (lib.sampler_chain_add)(sampler_chain, s_dist);
            }
        }

        // ── 3. Prefill decode ──────────────────────────────────────────────
        let prefill_start = Instant::now();

        if let Err(err) = decode_prefill(&lib, ctx, &tokens, n_batch) {
            unsafe { (lib.sampler_free)(sampler_chain) };
            return Err(err);
        }

        let prefill_ms = prefill_start.elapsed().as_secs_f64() * 1000.0;
        let prefill_tps = n_prompt_tokens as f64 / (prefill_ms / 1000.0);
        info!(
            "Prefill finished in {:.2} ms ({:.2} t/s)",
            prefill_ms, prefill_tps
        );

        // create session state
        Ok(Self {
            lib,
            ctx,
            vocab,
            sampler: sampler_chain, // use sampler_chain here
            stop_sequences: sampler_cfg.stop_sequences.clone(),
            n_decoded: n_prompt_tokens,
            prompt_tokens: n_prompt_tokens as u32,
            max_tokens: sampler_cfg.max_tokens,
            tokens_generated: 0,
            finished: false,
            prefill_ms: Some(prefill_ms),
            gen_start: None,
            tokens_per_sec: 0.0,
        })
    }

    /// Sample and decode the next token. Returns `None` when generation ends.
    ///
    /// Each call advances the context by one token.
    pub fn next_token(&mut self) -> Option<String> {
        if self.finished {
            return None;
        }
        if self.tokens_generated >= self.max_tokens {
            self.finished = true;
            self.finalize_timing();
            return None;
        }

        // Start timing on first generated token
        if self.gen_start.is_none() {
            self.gen_start = Some(Instant::now());
        }

        unsafe {
            // Sample the next token
            let new_token = (self.lib.sampler_sample)(self.sampler, self.ctx, -1);

            // Check end-of-generation
            if (self.lib.vocab_is_eog)(self.vocab, new_token) {
                self.finished = true;
                self.finalize_timing();
                return None;
            }

            // Decode (advance context with the new token)
            let mut token_arr = [new_token];
            let batch = (self.lib.batch_get_one)(token_arr.as_mut_ptr(), 1);
            let ret = (self.lib.decode)(self.ctx, batch);
            if ret != 0 {
                debug!(ret, "llama_decode (gen step) error — stopping");
                self.finished = true;
                self.finalize_timing();
                return None;
            }

            self.n_decoded += 1;
            self.tokens_generated += 1;

            // Detokenize
            let token_text = self.detokenize_token(new_token);

            // Check stop sequences
            // (simple check: if generated text so far ends with a stop seq)
            // For streaming we check per-token — good enough for Phase 2.
            for stop in &self.stop_sequences {
                if token_text.trim_end() == stop.as_str() {
                    self.finished = true;
                    self.finalize_timing();
                    return None;
                }
            }

            Some(token_text)
        }
    }

    fn finalize_timing(&mut self) {
        if let Some(start) = self.gen_start.take() {
            let elapsed = start.elapsed().as_secs_f64();
            if elapsed > 0.0 {
                self.tokens_per_sec = self.tokens_generated as f64 / elapsed;
            }
        }
        info!(
            tokens_generated = self.tokens_generated,
            tps = format!("{:.1}", self.tokens_per_sec),
            prefill_ms = self
                .prefill_ms
                .map(|v| format!("{:.1}", v))
                .unwrap_or_default(),
            "Generation complete"
        );
    }

    pub fn tokens_generated(&self) -> u32 {
        self.tokens_generated
    }

    pub fn prompt_tokens(&self) -> u32 {
        self.prompt_tokens
    }

    pub fn tps(&self) -> f64 {
        self.tokens_per_sec
    }

    pub fn prefill_ms(&self) -> Option<f64> {
        self.prefill_ms
    }

    fn detokenize_token(&self, token: llama_token) -> String {
        if let Some(token_to_piece) = self.lib.token_to_piece {
            return decode_token_piece_with(|buf, len| unsafe {
                token_to_piece(self.vocab, token, buf, len, 0, true)
            });
        }

        let text_ptr = unsafe { (self.lib.vocab_get_text)(self.vocab, token) };
        if text_ptr.is_null() {
            String::new()
        } else {
            unsafe { CStr::from_ptr(text_ptr) }
                .to_string_lossy()
                .into_owned()
        }
    }
}

fn build_prefill_plan(token_count: usize, max_batch_tokens: usize) -> Vec<PrefillChunkPlan> {
    let chunk_size = max_batch_tokens.max(1);
    let mut plans = Vec::new();
    let mut start = 0usize;

    while start < token_count {
        let end = (start + chunk_size).min(token_count);
        let logits_index = if end == token_count {
            Some(end - start - 1)
        } else {
            None
        };
        plans.push(PrefillChunkPlan {
            range: start..end,
            logits_index,
        });
        start = end;
    }

    plans
}

fn decode_prefill(
    lib: &LlamaLib,
    ctx: *mut llama_context,
    tokens: &[llama_token],
    max_batch_tokens: u32,
) -> Result<()> {
    for plan in build_prefill_plan(tokens.len(), max_batch_tokens as usize) {
        let chunk = &tokens[plan.range.clone()];
        let mut batch = unsafe { (lib.batch_init)(chunk.len() as i32, 0, 1) };
        batch.n_tokens = chunk.len() as i32;

        let ret = unsafe {
            for (i, token) in chunk.iter().copied().enumerate() {
                *batch.token.add(i) = token;
                *batch.pos.add(i) = (plan.range.start + i) as i32;
                *batch.n_seq_id.add(i) = 1;

                let seq_ids = *batch.seq_id.add(i);
                *seq_ids.add(0) = 0;

                *batch.logits.add(i) = if plan.logits_index == Some(i) { 1 } else { 0 };
            }

            let decode_batch = std::ptr::read(&batch);
            (lib.decode)(ctx, decode_batch)
        };

        unsafe { (lib.batch_free)(batch) };

        if ret != 0 {
            bail!(
                "llama_decode (prefill) returned {} while decoding prompt chunk {:?}",
                ret,
                plan.range
            );
        }
    }

    Ok(())
}

fn decode_token_piece_with<F>(mut call: F) -> String
where
    F: FnMut(*mut std::ffi::c_char, i32) -> i32,
{
    let mut capacity = 64usize;

    loop {
        let mut buffer = vec![0u8; capacity];
        let written = call(
            buffer.as_mut_ptr() as *mut std::ffi::c_char,
            capacity as i32,
        );

        if written < 0 {
            capacity = (-written) as usize;
            continue;
        }

        let written = written as usize;
        return String::from_utf8_lossy(&buffer[..written]).into_owned();
    }
}

#[cfg(test)]
mod tests {
    use super::{build_prefill_plan, decode_token_piece_with, PrefillChunkPlan};

    #[test]
    fn retries_when_piece_buffer_is_too_small() {
        let mut calls = 0;

        let decoded = decode_token_piece_with(|buf, len| {
            calls += 1;

            if len < 17 {
                return -17;
            }

            let bytes = b" Quantum computer";
            unsafe {
                std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
            }
            bytes.len() as i32
        });

        assert_eq!(decoded, " Quantum computer");
        assert_eq!(calls, 1);
    }

    #[test]
    fn decodes_piece_after_resizing_buffer() {
        let mut calls = 0;

        let decoded = decode_token_piece_with(|buf, len| {
            calls += 1;

            if calls == 1 {
                return -18;
            }

            assert!(len >= 18);
            let bytes = b" Quantum computer";
            unsafe {
                std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
            }
            bytes.len() as i32
        });

        assert_eq!(decoded, " Quantum computer");
        assert_eq!(calls, 2);
    }

    #[test]
    fn prefill_plan_requests_logits_only_for_last_prompt_token_at_batch_limit() {
        let plan = build_prefill_plan(128, 128);

        assert_eq!(
            plan,
            vec![PrefillChunkPlan {
                range: 0..128,
                logits_index: Some(127),
            }]
        );
    }

    #[test]
    fn prefill_plan_splits_large_prompt_and_only_requests_final_logits() {
        let plan = build_prefill_plan(300, 128);

        assert_eq!(
            plan,
            vec![
                PrefillChunkPlan {
                    range: 0..128,
                    logits_index: None,
                },
                PrefillChunkPlan {
                    range: 128..256,
                    logits_index: None,
                },
                PrefillChunkPlan {
                    range: 256..300,
                    logits_index: Some(43),
                },
            ]
        );
    }
}

impl Drop for GenerationSession {
    fn drop(&mut self) {
        if !self.sampler.is_null() {
            unsafe { (self.lib.sampler_free)(self.sampler) };
            self.sampler = std::ptr::null_mut();
        }
    }
}