DS4Server/src/engine.rs

mod gguf;
#[cfg(any(target_os = "macos", test))]
mod kvstore;
#[cfg(target_os = "macos")]
mod metal;
mod tokenizer;
mod validation;

#[cfg(target_os = "macos")]
use crate::metrics::{KvLookup, Metrics, SsdStats};
use crate::model::{ModelChoice, validate_engine_artifacts};
#[cfg(target_os = "macos")]
use crate::settings::TurnSettings;
use crate::settings::{EngineSettings, ReasoningMode};
#[cfg(target_os = "macos")]
use base64::Engine as _;
use gguf::{
    BF16, F16, F32, Gguf, I32, IQ2_XXS, MXFP4, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value,
};
#[cfg(target_os = "macos")]
use kvstore::{KvStore, StoreReason};
use sha2::{Digest, Sha256};
#[cfg(target_os = "macos")]
use std::collections::HashMap;
use std::path::Path;
#[cfg(target_os = "macos")]
use std::path::PathBuf;
#[cfg(target_os = "macos")]
use std::sync::Arc;
#[cfg(target_os = "macos")]
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(target_os = "macos")]
use std::time::Instant;
use tokenizer::Tokenizer;
use validation::{SupportKind, validate_main, validate_support};

const VISION_DATA_START: &str = "<|ds4server_image_data|>";
const VISION_DATA_END: &str = "<|/ds4server_image_data|>";
const VISION_TOKEN_START: &str = "\u{fdd0}ds4-image:";
const VISION_TOKEN_END: &str = "\u{fdd1}";
const VISION_IMAGE_TOKEN: i32 = 154_854;
const VISION_START_TOKEN: i32 = 154_830;
const VISION_END_TOKEN: i32 = 154_831;
type VisionOverlays = Vec<(u32, metal::VisionEmbedding)>;

pub(crate) fn vision_data_marker(uri: &str) -> String {
    format!("{VISION_DATA_START}{uri}{VISION_DATA_END}")
}

#[cfg(target_os = "macos")]
unsafe extern "C" {
    fn mmap(
        address: *mut std::ffi::c_void,
        length: usize,
        protection: i32,
        flags: i32,
        fd: i32,
        offset: i64,
    ) -> *mut std::ffi::c_void;
    fn mlock(address: *const std::ffi::c_void, length: usize) -> i32;
    fn munlock(address: *const std::ffi::c_void, length: usize) -> i32;
    fn munmap(address: *mut std::ffi::c_void, length: usize) -> i32;
}

pub(crate) use validation::{validate_model_artifact, validate_vision_artifact};

#[cfg(target_os = "macos")]
pub(crate) use kvstore::sweep_unreachable as sweep_transient_cache;
#[cfg(target_os = "macos")]
pub(crate) use metal::configure_sources as configure_metal_sources;

const DENSE: &[u32] = &[BF16, Q8_0, Q4_K, Q4_0];
const ROUTED: &[u32] = &[Q8_0, IQ2_XXS, Q2_K, Q4_K, Q5_K, Q6_K, MXFP4];
const PLAIN: &[u32] = &[F16, F32];
const DSPARK_DENSE: &[u32] = &[F16, F32, Q8_0];

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ModelFamily {
    DeepSeek,
    Glm,
}

#[derive(Clone, Copy)]
struct Shape {
    model: ModelChoice,
    family: ModelFamily,
    layers: u32,
    embd: u64,
    vocab: u64,
    heads: u64,
    head_kv: u64,
    head_dim: u64,
    value_dim: u64,
    rot: u64,
    out_groups: u64,
    lora_q: u64,
    lora_o: u64,
    experts: u64,
    experts_used: u64,
    expert_shared: u64,
    ff_expert: u64,
    ff_dense: u64,
    hash_layers: u32,
    sliding_window: u64,
    indexer_heads: u64,
    indexer_head_dim: u64,
    indexer_top_k: u64,
    hc: u64,
    hc_sinkhorn: u64,
    nextn: u32,
    leading_dense: u32,
    kv_lora: u64,
    key_mla: u64,
    value_mla: u64,
    rms_epsilon: f32,
    hc_epsilon: f32,
    expert_weight_scale: f32,
    swiglu_clamp: f32,
    rope_base: f32,
    rope_scale: f32,
    rope_beta_fast: f32,
    rope_beta_slow: f32,
    compress_rope_base: f32,
    original_context: u64,
}

const DEEPSEEK_BASE: Shape = Shape {
    model: ModelChoice::DeepSeekV4Flash0731,
    family: ModelFamily::DeepSeek,
    layers: 43,
    embd: 4096,
    vocab: 129_280,
    heads: 64,
    head_kv: 1,
    head_dim: 512,
    value_dim: 512,
    rot: 64,
    out_groups: 8,
    lora_q: 1024,
    lora_o: 1024,
    experts: 256,
    experts_used: 6,
    expert_shared: 1,
    ff_expert: 2048,
    ff_dense: 0,
    hash_layers: 3,
    sliding_window: 128,
    indexer_heads: 64,
    indexer_head_dim: 128,
    indexer_top_k: 512,
    hc: 4,
    hc_sinkhorn: 20,
    nextn: 0,
    leading_dense: 0,
    kv_lora: 0,
    key_mla: 0,
    value_mla: 0,
    rms_epsilon: 1.0e-6,
    hc_epsilon: 1.0e-6,
    expert_weight_scale: 1.5,
    swiglu_clamp: 10.0,
    rope_base: 10_000.0,
    rope_scale: 16.0,
    rope_beta_fast: 32.0,
    rope_beta_slow: 1.0,
    compress_rope_base: 160_000.0,
    original_context: 65_536,
};

const FLASH_0731: Shape = Shape { ..DEEPSEEK_BASE };

const PRO: Shape = Shape {
    model: ModelChoice::DeepSeekV4Pro,
    layers: 61,
    embd: 7168,
    heads: 128,
    out_groups: 16,
    lora_q: 1536,
    experts: 384,
    ff_expert: 3072,
    indexer_top_k: 1024,
    expert_weight_scale: 2.5,
    ..DEEPSEEK_BASE
};

const GLM: Shape = Shape {
    model: ModelChoice::Glm52,
    family: ModelFamily::Glm,
    layers: 79,
    embd: 6144,
    vocab: 154_880,
    heads: 64,
    head_kv: 1,
    head_dim: 576,
    value_dim: 512,
    rot: 64,
    out_groups: 0,
    lora_q: 2048,
    lora_o: 0,
    experts: 256,
    experts_used: 8,
    expert_shared: 1,
    ff_expert: 2048,
    ff_dense: 12_288,
    hash_layers: 0,
    sliding_window: 0,
    indexer_heads: 32,
    indexer_head_dim: 128,
    indexer_top_k: 2048,
    hc: 0,
    hc_sinkhorn: 0,
    nextn: 1,
    leading_dense: 3,
    kv_lora: 512,
    key_mla: 256,
    value_mla: 256,
    rms_epsilon: 1.0e-5,
    hc_epsilon: 0.0,
    expert_weight_scale: 2.5,
    swiglu_clamp: 0.0,
    rope_base: 8_000_000.0,
    rope_scale: 1.0,
    rope_beta_fast: 0.0,
    rope_beta_slow: 0.0,
    compress_rope_base: 0.0,
    original_context: 1_048_576,
};

const GLM53_FLASH: Shape = Shape {
    model: ModelChoice::Glm53Flash,
    family: ModelFamily::Glm,
    layers: 46,
    embd: 4096,
    vocab: 154_880,
    heads: 64,
    head_kv: 1,
    head_dim: 512,
    value_dim: 256,
    rot: 0,
    out_groups: 0,
    lora_q: 1536,
    lora_o: 0,
    experts: 288,
    experts_used: 8,
    expert_shared: 1,
    ff_expert: 2048,
    ff_dense: 12_288,
    hash_layers: 0,
    sliding_window: 0,
    indexer_heads: 32,
    indexer_head_dim: 128,
    indexer_top_k: 2048,
    hc: 4,
    hc_sinkhorn: 20,
    nextn: 1,
    leading_dense: 3,
    kv_lora: 512,
    key_mla: 256,
    value_mla: 256,
    rms_epsilon: 1.0e-5,
    hc_epsilon: 1.0e-6,
    expert_weight_scale: 2.5,
    swiglu_clamp: 10.0,
    rope_base: 0.0,
    rope_scale: 0.0,
    rope_beta_fast: 0.0,
    rope_beta_slow: 0.0,
    compress_rope_base: 0.0,
    original_context: 1_048_576,
};

pub(crate) struct Model {
    main: Gguf,
    support: Option<Gguf>,
    vision: Option<Gguf>,
    support_kind: Option<SupportKind>,
    shape: Shape,
    tokenizer: Tokenizer,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub(crate) struct ModelSummary {
    pub(crate) model: ModelChoice,
    pub(crate) mapped_bytes: u64,
    pub(crate) tensor_count: usize,
    pub(crate) vocabulary_size: usize,
    pub(crate) support_loaded: bool,
    pub(crate) vision_loaded: bool,
}

impl Model {
    #[allow(dead_code)]
    pub(crate) fn open(settings: &EngineSettings) -> Result<Self, String> {
        validate_engine_artifacts(
            settings.model,
            settings.speculative.dspark,
            &settings.artifacts,
        )?;
        let mut model = Self::open_main(&settings.artifacts.model, settings.model)?;
        if settings.execution.warm_weights {
            model.main.warm()?;
        }
        if let Some(path) = &settings.artifacts.support {
            let support = Gguf::open(path)?;
            let kind = validate_support(&support, &model.shape)?;
            if settings.execution.warm_weights {
                support.warm()?;
            }
            model.support = Some(support);
            model.support_kind = Some(kind);
        }
        if let Some(path) = &settings.artifacts.vision {
            validate_vision_artifact(path)?;
            let vision = Gguf::open(path)?;
            if settings.execution.warm_weights && settings.speculative.keep_vision_loaded {
                vision.warm()?;
            }
            model.vision = Some(vision);
        }
        Ok(model)
    }

    fn open_main(path: &Path, expected: ModelChoice) -> Result<Self, String> {
        let main = Gguf::open(path)?;
        let shape = validate_main(&main, expected)?;
        let tokenizer = Tokenizer::load(&main, shape.family)?;
        if tokenizer.vocab_size() != shape.vocab as usize {
            return Err(format!(
                "tokenizer has {} entries, expected {}",
                tokenizer.vocab_size(),
                shape.vocab
            ));
        }
        Ok(Self {
            main,
            support: None,
            vision: None,
            support_kind: None,
            shape,
            tokenizer,
        })
    }

    pub(crate) fn summary(&self) -> ModelSummary {
        ModelSummary {
            model: self.shape.model,
            mapped_bytes: self.main.len()
                + self.support.as_ref().map_or(0, Gguf::len)
                + self.vision.as_ref().map_or(0, Gguf::len),
            tensor_count: self.main.tensors.len()
                + self
                    .support
                    .as_ref()
                    .map_or(0, |support| support.tensors.len())
                + self
                    .vision
                    .as_ref()
                    .map_or(0, |vision| vision.tensors.len()),
            vocabulary_size: self.tokenizer.vocab_size(),
            support_loaded: self.support.is_some(),
            vision_loaded: self.vision.is_some(),
        }
    }

    fn checkpoint_identity(&self) -> [u8; 32] {
        let mut hash = Sha256::new();
        hash.update(b"DS4Server model checkpoint identity v1");
        hash.update(self.main.checkpoint_identity());
        if let Some(support) = &self.support {
            hash.update(support.checkpoint_identity());
        }
        if let Some(vision) = &self.vision {
            hash.update(vision.checkpoint_identity());
        }
        hash.finalize().into()
    }

    pub(crate) fn tokenize(&self, text: &str) -> Vec<i32> {
        self.tokenizer.tokenize(text)
    }

    pub(crate) fn render_prompt(
        &self,
        system: &str,
        prompt: &str,
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.tokenizer.encode_chat(system, prompt, reasoning)
    }

    fn render_conversation(
        &self,
        system: &str,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.tokenizer
            .encode_conversation(system, messages, reasoning)
    }

    fn render_history(
        &self,
        system: &str,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
    ) -> Vec<i32> {
        self.tokenizer.encode_history(system, messages, reasoning)
    }

    fn render_continuation(
        &self,
        prompt: &str,
        reasoning: ReasoningMode,
        skip_previous_eos: bool,
    ) -> Vec<i32> {
        self.tokenizer
            .encode_continuation(prompt, reasoning, skip_previous_eos)
    }

    pub(crate) fn token_bytes(&self, token: i32) -> Option<Vec<u8>> {
        self.tokenizer.token_bytes(token)
    }

    pub(crate) fn eos_token(&self) -> i32 {
        self.tokenizer.eos()
    }

    pub(crate) fn is_stop_token(&self, token: i32) -> bool {
        self.tokenizer.is_stop(token)
    }

    pub(crate) fn is_think_start_token(&self, token: i32) -> bool {
        self.tokenizer.is_think_start(token)
    }

    pub(crate) fn is_think_end_token(&self, token: i32) -> bool {
        self.tokenizer.is_think_end(token)
    }

    pub(crate) fn is_stop_token_for_reasoning(&self, token: i32, reasoning: ReasoningMode) -> bool {
        self.is_stop_token(token)
            || (reasoning == ReasoningMode::Direct
                && (self.is_think_start_token(token) || self.is_think_end_token(token)))
    }

    pub(crate) fn tensor_data(&self, name: &str) -> Result<&[u8], String> {
        self.main.tensor_data(name)
    }
}

#[cfg(target_os = "macos")]
pub(crate) struct Generator {
    executor: metal::Executor,
    _simulated_memory: Option<SimulatedMemory>,
    checkpoint: Option<PathBuf>,
    /// Token frontier of the last transient store, so continued checkpoints are
    /// spaced like ds4's `continued_last_store_tokens`.
    last_store_tokens: u32,
    metrics: Arc<Metrics>,
    resident_sessions: HashMap<PathBuf, ResidentSlot>,
    resident_active: Option<PathBuf>,
    resident_limit: usize,
}

#[cfg(target_os = "macos")]
struct ResidentSlot {
    state: metal::ResidentState,
    last_store_tokens: u32,
}

#[cfg(target_os = "macos")]
impl Drop for Generator {
    fn drop(&mut self) {
        // Match ds4: every resident session graph must die before Metal cleanup.
        self.resident_sessions.clear();
    }
}

#[cfg(target_os = "macos")]
struct SimulatedMemory {
    address: std::ptr::NonNull<std::ffi::c_void>,
    bytes: usize,
}

#[cfg(target_os = "macos")]
impl SimulatedMemory {
    fn acquire(bytes: u64) -> Result<Option<Self>, String> {
        if bytes == 0 {
            return Ok(None);
        }
        let bytes = usize::try_from(bytes).map_err(|_| "simulated memory size is too large")?;
        let address = unsafe { mmap(std::ptr::null_mut(), bytes, 1 | 2, 2 | 0x1000, -1, 0) };
        let Some(address) = std::ptr::NonNull::new(address) else {
            return Err(format!(
                "cannot reserve simulated used memory: {}",
                std::io::Error::last_os_error()
            ));
        };
        if address.as_ptr() as isize == -1 {
            return Err(format!(
                "cannot reserve simulated used memory: {}",
                std::io::Error::last_os_error()
            ));
        }
        let chunk = 256 * 1024 * 1024;
        let mut locked = 0;
        while locked < bytes {
            let length = (bytes - locked).min(chunk);
            let start = unsafe { address.as_ptr().cast::<u8>().add(locked) };
            for page in (0..length).step_by(16 * 1024) {
                unsafe { start.add(page).write((page / (16 * 1024)) as u8) };
            }
            unsafe { start.add(length - 1).write(1) };
            if unsafe { mlock(start.cast(), length) } != 0 {
                if locked != 0 {
                    unsafe { munlock(address.as_ptr(), locked) };
                }
                unsafe { munmap(address.as_ptr(), bytes) };
                return Err(format!(
                    "cannot lock simulated used memory after {:.2} GiB: {}",
                    locked as f64 / 1_073_741_824.0,
                    std::io::Error::last_os_error()
                ));
            }
            locked += length;
        }
        Ok(Some(Self { address, bytes }))
    }
}

#[cfg(target_os = "macos")]
impl Drop for SimulatedMemory {
    fn drop(&mut self) {
        unsafe {
            munlock(self.address.as_ptr(), self.bytes);
            munmap(self.address.as_ptr(), self.bytes);
        }
    }
}

#[derive(Clone)]
pub(crate) struct ChatTurn {
    pub(crate) user: bool,
    pub(crate) tool: bool,
    pub(crate) system: bool,
    pub(crate) skip_previous_eos: bool,
    pub(crate) reasoning: Option<String>,
    pub(crate) reasoning_complete: bool,
    pub(crate) content: String,
}

pub(crate) struct GenerationOutput {
    pub(crate) message: ChatTurn,
    pub(crate) prompt_tokens: u32,
    pub(crate) cached_tokens: u32,
    pub(crate) completion_tokens: u32,
    pub(crate) finish_reason: &'static str,
    pub(crate) previous_checkpoint_bytes: Option<u64>,
    pub(crate) checkpoint_bytes: u64,
}

struct CheckpointSelection {
    found: bool,
    incompatible: bool,
}

fn checkpoint_rebuild_activity(incompatible: bool) -> &'static str {
    if incompatible {
        "Rebuilding context: the checkpoint belongs to a different model or model configuration."
    } else {
        "Rebuilding context: the saved history or generation settings changed."
    }
}

#[cfg(target_os = "macos")]
pub(crate) struct CompactionOutput {
    pub(crate) summary: String,
    pub(crate) tail_start: usize,
    pub(crate) context_tokens: u32,
    pub(crate) checkpoint: PathBuf,
}

#[cfg(target_os = "macos")]
impl Generator {
    pub(crate) fn open(settings: &EngineSettings, metrics: Arc<Metrics>) -> Result<Self, String> {
        let simulated_memory =
            SimulatedMemory::acquire(settings.diagnostics.simulated_used_memory_bytes)?;
        let model = Model::open(settings)?;
        let executor = metal::Executor::open_configured(
            model,
            settings.context_tokens.max(1) as u32,
            settings.execution.quality,
            settings.execution.prefill_chunk,
            settings.execution.power_percent,
            settings.speculative,
            settings.ssd,
            settings.steering.clone(),
            settings.diagnostics.expert_profile_path.as_deref(),
        )?;
        let generator = Self {
            executor,
            _simulated_memory: simulated_memory,
            checkpoint: None,
            last_store_tokens: 0,
            metrics,
            resident_sessions: HashMap::new(),
            resident_active: None,
            resident_limit: std::env::var("DS4_RESIDENT_SESSIONS")
                .ok()
                .and_then(|value| value.parse::<usize>().ok())
                .filter(|limit| *limit > 0)
                .unwrap_or(1)
                .saturating_sub(1),
        };
        generator.publish_execution_stats();
        Ok(generator)
    }

    pub(crate) fn summary(&self) -> ModelSummary {
        self.executor.model().summary()
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn generate(
        &mut self,
        checkpoint: &Path,
        bootstrap: Option<&Path>,
        messages: &[ChatTurn],
        settings: &TurnSettings,
        cancelled: &AtomicBool,
        mut emit: impl FnMut(bool, String),
        mut progress: impl FnMut(u32, u32, Option<f32>),
        mut phase: impl FnMut(&'static str),
    ) -> Result<GenerationOutput, String> {
        let checkpoint_present = checkpoint.is_file();
        let selected = self.select_checkpoint(checkpoint, |tag| {
            checkpoint_matches_prefix(
                tag,
                &settings.system_prompt,
                settings.reasoning_mode,
                messages,
            )
        })?;
        if checkpoint_present && !selected.found {
            phase(checkpoint_rebuild_activity(selected.incompatible));
        }
        if !selected.found
            && let Some(directory) = bootstrap
        {
            self.prepare_bootstrap(directory, settings, cancelled, &mut progress, &mut phase)?;
        }
        let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
        self.publish_execution_stats();
        let (mut output, prompt_complete) = result?;
        let mut completed = messages.to_vec();
        completed.push(output.message.clone());
        output.previous_checkpoint_bytes =
            std::fs::metadata(checkpoint).ok().map(|item| item.len());
        self.save_checkpoint(
            checkpoint,
            if prompt_complete {
                conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed)
            } else {
                [0; 32]
            },
        )?;
        output.checkpoint_bytes = std::fs::metadata(checkpoint)
            .map(|item| item.len())
            .unwrap_or(0);
        self.checkpoint = Some(checkpoint.to_owned());
        self.resident_active = Some(checkpoint.to_owned());
        Ok(output)
    }

    fn publish_execution_stats(&self) {
        let stats = self.executor.execution_stats();
        self.metrics.speculative_stats(
            stats.speculative_mode,
            stats.speculative_cycles,
            stats.drafted_tokens,
            stats.accepted_draft_tokens,
            stats.verifier_passes,
            stats.verifier_ms,
        );
        self.metrics.ssd_stats(SsdStats {
            enabled: stats.ssd_enabled,
            resident_bytes: stats.ssd_resident_bytes,
            cache_bytes: stats.ssd_cache_bytes,
            cache_experts: stats.ssd_cache_experts,
            cache_entries: stats.ssd_cache_entries,
            preloaded_experts: stats.ssd_preloaded_experts,
            cache_hits: stats.ssd_cache_hits,
            cache_misses: stats.ssd_cache_misses,
            cache_evictions: stats.ssd_cache_evictions,
            cache_wraps: stats.ssd_cache_wraps,
            buffer_allocs: stats.ssd_buffer_allocs,
            buffer_reuses: stats.ssd_buffer_reuses,
            pread_bytes: stats.ssd_pread_bytes,
            pread_ms: stats.ssd_pread_ms,
            evict_advise_bytes: stats.ssd_evict_advise_bytes,
            willneed_advise_bytes: stats.ssd_willneed_advise_bytes,
            selected_requests: stats.ssd_selected_requests,
            requested_bytes: stats.ssd_requested_bytes,
            wait_ms: stats.ssd_wait_ms,
        });
    }

    pub(crate) fn generate_transient(
        &mut self,
        directory: &Path,
        messages: &[ChatTurn],
        settings: &TurnSettings,
        cancelled: &AtomicBool,
        mut emit: impl FnMut(bool, String),
        mut progress: impl FnMut(u32, u32, Option<f32>),
    ) -> Result<GenerationOutput, String> {
        let history = messages
            .split_last()
            .map_or(messages, |(_, history)| history);
        let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
        let history_key =
            conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
        let history_tag: [u8; 32] = Sha256::digest(&history_key).into();
        let previous_checkpoint =
            self.restore_cached_prefix(directory, &store, &history_key, history_tag)?;

        let result = self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress);
        self.publish_execution_stats();
        let (mut output, prompt_complete) = result?;
        let mut completed = messages.to_vec();
        completed.push(output.message.clone());
        output.previous_checkpoint_bytes = previous_checkpoint
            .as_deref()
            .and_then(|path| std::fs::metadata(path).ok())
            .map(|item| item.len());
        if prompt_complete {
            let completed_key =
                conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed);
            let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into();
            if !settings.kv_cache.stores(
                self.executor.position(),
                history.is_empty(),
                self.last_store_tokens,
            ) {
                // A gated store still leaves the finished conversation in the
                // live KV, so mark it and drop the stale file association.
                self.executor.note_checkpoint_tag(completed_tag);
                self.checkpoint = None;
                self.resident_active = Some(resident_key(directory, completed_tag));
                return Ok(output);
            }
            let completed_checkpoint = store.checkpoint_path(&completed_key);
            self.save_checkpoint(&completed_checkpoint, completed_tag)?;
            let retained = store.record(
                &completed_checkpoint,
                &completed_key,
                completed_tag,
                self.executor.position(),
                self.executor.context(),
                if history.is_empty() {
                    StoreReason::Cold
                } else {
                    StoreReason::Continued
                },
            )?;
            output.checkpoint_bytes = retained
                .then(|| std::fs::metadata(&completed_checkpoint).ok())
                .flatten()
                .map_or(0, |item| item.len());
            if retained {
                self.last_store_tokens = self.executor.position();
            }
            self.checkpoint = retained.then_some(completed_checkpoint);
            if let Some(checkpoint) = &self.checkpoint {
                self.resident_active = Some(checkpoint.clone());
            } else {
                self.resident_active = Some(resident_key(directory, completed_tag));
            }
        }
        Ok(output)
    }

    fn prepare_bootstrap(
        &mut self,
        directory: &Path,
        settings: &TurnSettings,
        cancelled: &AtomicBool,
        progress: &mut impl FnMut(u32, u32, Option<f32>),
        phase: &mut impl FnMut(&'static str),
    ) -> Result<(), String> {
        let key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &[]);
        let tag: [u8; 32] = Sha256::digest(&key).into();
        let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
        self.restore_cached_prefix(directory, &store, &key, tag)?;

        let tokens = self.executor.model().render_history(
            &settings.system_prompt,
            &[],
            settings.reasoning_mode,
        );
        if tokens.len() >= self.executor.context() as usize {
            return Err(format!(
                "System prompt has {} tokens, but the configured context size is {} tokens",
                tokens.len(),
                self.executor.context()
            ));
        }
        let reused = self.executor.align_prompt(&tokens)?;
        if reused == tokens.len() {
            return Ok(());
        }

        phase("Updating system prompt cache…");
        let completed = self.prefill_suffix(&tokens, reused, false, cancelled, progress)?;
        if completed != tokens.len() - reused {
            return Err("generation cancelled while updating the system prompt cache".into());
        }
        self.executor.note_checkpoint_tag(tag);
        if !settings.kv_cache.stores(self.executor.position(), true, 0) {
            self.checkpoint = None;
            self.resident_active = Some(resident_key(directory, tag));
            return Ok(());
        }

        let checkpoint = store.checkpoint_path(&key);
        self.save_checkpoint(&checkpoint, tag)?;
        let retained = store.record(
            &checkpoint,
            &key,
            tag,
            self.executor.position(),
            self.executor.context(),
            StoreReason::Cold,
        )?;
        self.checkpoint = retained.then_some(checkpoint.clone());
        self.resident_active = Some(if retained {
            checkpoint
        } else {
            resident_key(directory, tag)
        });
        Ok(())
    }

    fn restore_cached_prefix(
        &mut self,
        directory: &Path,
        store: &KvStore,
        key: &[u8],
        tag: [u8; 32],
    ) -> Result<Option<PathBuf>, String> {
        if self.executor.checkpoint_tag() == tag {
            self.metrics.kv_lookup(KvLookup::MemoryHit);
            return Ok(self.checkpoint.clone());
        }
        if let Some(entry) = store.find(key, self.executor.context()) {
            if self
                .select_checkpoint(&entry.checkpoint, |tag| tag == entry.tag)?
                .found
            {
                store.touch(&entry)?;
                self.last_store_tokens = entry.tokens;
                return Ok(Some(entry.checkpoint));
            }
            store.discard(&entry);
        } else {
            let key = resident_key(directory, tag);
            let restored = self.activate_resident(key)?;
            if !restored || self.executor.checkpoint_tag() != tag {
                self.executor.reset()?;
            }
            self.metrics.kv_lookup(KvLookup::Miss);
        }
        self.checkpoint = None;
        self.last_store_tokens = 0;
        Ok(None)
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn compact(
        &mut self,
        messages: &[ChatTurn],
        settings: &TurnSettings,
        rebuild_system_prompt: &str,
        reason: &str,
        checkpoint: &Path,
        cancelled: &AtomicBool,
        mut progress: impl FnMut(u32, u32, Option<f32>),
        mut phase: impl FnMut(&'static str),
    ) -> Result<CompactionOutput, String> {
        self.activate_resident(checkpoint.to_owned())?;
        let _ = std::fs::remove_file(checkpoint);
        self.executor.reset()?;
        self.checkpoint = None;
        let result = (|| {
            phase("Compacting durable task state…");
            let mut private_messages = messages.to_vec();
            private_messages.push(ChatTurn {
                user: true,
                tool: false,
                system: false,
                skip_previous_eos: false,
                reasoning: None,
                reasoning_complete: true,
                content: crate::compaction::summary_prompt(reason),
            });
            let mut private_settings = settings.clone();
            private_settings.reasoning_mode = ReasoningMode::Direct;
            let private_prompt = self.executor.model().render_conversation(
                &private_settings.system_prompt,
                &private_messages,
                private_settings.reasoning_mode,
            );
            private_settings.max_generated_tokens = crate::compaction::summary_budget(
                private_prompt.len().min(u32::MAX as usize) as u32,
                self.executor.context(),
            )
            .ok_or_else(|| "not enough context left to request compaction summary".to_owned())?;
            private_settings.temperature = 0.0;
            private_settings.stops = vec![
                "<｜DSML｜".into(),
                "<DSML｜".into(),
                "<tool_call>".into(),
                "<think>".into(),
                "</think>".into(),
            ];
            let (output, prompt_complete) = self.generate_inner(
                &private_messages,
                &private_settings,
                cancelled,
                &mut |_, _| {},
                &mut progress,
            )?;
            if cancelled.load(Ordering::Relaxed) || !prompt_complete {
                return Err("context compaction interrupted".into());
            }
            let summary = crate::compaction::sanitize_summary(&output.message.content);
            if summary.is_empty() {
                return Err("context compaction produced an empty summary".into());
            }

            // The private request must never become the rebuilt session prefix.
            phase("Rebuilding compacted context…");
            self.executor.reset()?;
            self.checkpoint = None;

            let full = self.executor.model().render_conversation(
                &settings.system_prompt,
                messages,
                settings.reasoning_mode,
            );
            let mut starts = Vec::with_capacity(messages.len());
            for index in 0..messages.len() {
                starts.push(
                    self.executor
                        .model()
                        .render_conversation(
                            &settings.system_prompt,
                            &messages[..index],
                            settings.reasoning_mode,
                        )
                        .len() as u32,
                );
            }
            let start = crate::compaction::tail_start(
                messages,
                &starts,
                full.len() as u32,
                crate::compaction::tail_budget(self.executor.context()),
            );
            let tail = messages[start..].to_vec();
            let rebuilt_system =
                crate::compaction::summary_system_prompt(rebuild_system_prompt, Some(&summary));
            let history_tokens = self.executor.model().render_history(
                &rebuilt_system,
                &tail,
                settings.reasoning_mode,
            );
            if history_tokens.len() >= self.executor.context() as usize {
                return Err("compacted context does not fit the configured context".into());
            }
            let context = self.executor.context();
            let completed = self.executor.prefill(&history_tokens, |used| {
                progress(used, context, None);
                !cancelled.load(Ordering::Relaxed)
            })?;
            if completed != history_tokens.len() || cancelled.load(Ordering::Relaxed) {
                return Err("context compaction interrupted during rebuild".into());
            }
            let tag = conversation_tag(&rebuilt_system, settings.reasoning_mode, &tail);
            phase("Saving compacted context…");
            self.save_checkpoint(checkpoint, tag)?;
            Ok(CompactionOutput {
                summary,
                tail_start: start,
                context_tokens: history_tokens.len() as u32,
                checkpoint: checkpoint.to_owned(),
            })
        })();
        if result.is_err() {
            let _ = std::fs::remove_file(checkpoint);
            self.executor.reset()?;
            self.checkpoint = None;
        } else {
            self.checkpoint = Some(checkpoint.to_owned());
        }
        result
    }

    pub(crate) fn rendered_history_tokens(
        &self,
        messages: &[ChatTurn],
        settings: &TurnSettings,
    ) -> Result<u32, String> {
        u32::try_from(
            self.executor
                .model()
                .render_history(&settings.system_prompt, messages, settings.reasoning_mode)
                .len(),
        )
        .map_err(|_| "rendered conversation is too large".to_owned())
    }

    fn select_checkpoint(
        &mut self,
        checkpoint: &Path,
        matches: impl Fn([u8; 32]) -> bool,
    ) -> Result<CheckpointSelection, String> {
        let resident_hit = self.activate_resident(checkpoint.to_owned())?;
        if !checkpoint.is_file() {
            self.executor.reset()?;
            self.checkpoint = None;
            self.metrics.kv_lookup(KvLookup::Miss);
            return Ok(CheckpointSelection {
                found: false,
                incompatible: false,
            });
        }
        if resident_hit && matches(self.executor.checkpoint_tag()) {
            self.checkpoint = Some(checkpoint.to_owned());
            self.metrics.kv_lookup(KvLookup::MemoryHit);
            return Ok(CheckpointSelection {
                found: true,
                incompatible: false,
            });
        }
        if self.checkpoint.as_deref() == Some(checkpoint) {
            let found = matches(self.executor.checkpoint_tag());
            self.metrics.kv_lookup(if found {
                KvLookup::MemoryHit
            } else {
                KvLookup::Miss
            });
            return Ok(CheckpointSelection {
                found,
                incompatible: false,
            });
        }
        self.executor.reset()?;
        self.metrics.kv_read_started();
        let started = Instant::now();
        let loaded = self
            .executor
            .load_checkpoint(checkpoint, &mut |bytes| self.metrics.kv_read_bytes(bytes));
        self.metrics
            .kv_read_finished(started.elapsed(), loaded.is_err());
        let found = matches!(loaded, Ok(true)) && matches(self.executor.checkpoint_tag());
        let incompatible = loaded.is_err();
        let lookup = match loaded {
            Ok(true) if found => KvLookup::DiskHit,
            Ok(true) | Ok(false) => KvLookup::Miss,
            Err(_) => {
                self.executor.reset()?;
                let _ = std::fs::remove_file(checkpoint);
                KvLookup::Invalid
            }
        };
        self.metrics.kv_lookup(lookup);
        self.checkpoint = Some(checkpoint.to_owned());
        Ok(CheckpointSelection {
            found,
            incompatible,
        })
    }

    fn activate_resident(&mut self, key: PathBuf) -> Result<bool, String> {
        if self.resident_active.as_ref() == Some(&key) {
            return Ok(true);
        }
        let restored = self.resident_sessions.contains_key(&key);
        let (mut incoming_state, incoming_last_store) = self
            .resident_sessions
            .remove(&key)
            .map_or((None, 0), |slot| (Some(slot.state), slot.last_store_tokens));
        self.executor.swap_resident_state(&mut incoming_state)?;
        if let (Some(previous), Some(outgoing)) = (self.resident_active.take(), incoming_state) {
            // DS4 defaults to one active session; opt into its batched-server
            // behavior with DS4_RESIDENT_SESSIONS when memory permits.
            if self.resident_limit != 0 {
                if self.resident_sessions.len() >= self.resident_limit {
                    let evicted = self.resident_sessions.keys().next().cloned();
                    if let Some(evicted) = evicted {
                        self.resident_sessions.remove(&evicted);
                    }
                }
                self.resident_sessions.insert(
                    previous,
                    ResidentSlot {
                        state: outgoing,
                        last_store_tokens: self.last_store_tokens,
                    },
                );
            }
        }
        self.resident_active = Some(key);
        self.checkpoint = None;
        self.last_store_tokens = incoming_last_store;
        Ok(restored)
    }

    fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> {
        self.metrics.kv_write_started();
        let started = Instant::now();
        let result = self
            .executor
            .save_checkpoint(checkpoint, tag, &mut |bytes| {
                self.metrics.kv_write_bytes(bytes);
            });
        self.metrics
            .kv_write_finished(started.elapsed(), result.is_err());
        result
    }

    fn prefill_suffix(
        &mut self,
        tokens: &[i32],
        reused: usize,
        has_vision: bool,
        cancelled: &AtomicBool,
        progress: &mut impl FnMut(u32, u32, Option<f32>),
    ) -> Result<usize, String> {
        let suffix = &tokens[reused..];
        if has_vision || (reused == 0 && tokens.len() > 1) || suffix.len() >= 4 {
            let context = self.executor.context();
            self.executor.prefill(suffix, |used| {
                progress(used, context, None);
                !cancelled.load(Ordering::Relaxed)
            })
        } else {
            let mut completed = 0;
            for &token in suffix {
                if cancelled.load(Ordering::Relaxed) {
                    break;
                }
                self.executor.eval(token)?;
                completed += 1;
                progress(self.executor.position(), self.executor.context(), None);
            }
            Ok(completed)
        }
    }

    fn render_multimodal_conversation(
        &mut self,
        system: &str,
        messages: &[ChatTurn],
        reasoning: ReasoningMode,
    ) -> Result<(Vec<i32>, VisionOverlays), String> {
        let mut rendered = messages.to_vec();
        let mut images = Vec::new();
        let mut total_images = 0_usize;
        let mut total_bytes = 0_usize;
        for message in &rendered {
            if !message.content.contains(VISION_DATA_START) {
                continue;
            }
            if !message.user && !message.tool {
                return Err("vision input is allowed only in user or tool messages".into());
            }
            let mut rest = message.content.as_str();
            while let Some(start) = rest.find(VISION_DATA_START) {
                let encoded = &rest[start + VISION_DATA_START.len()..];
                let end = encoded
                    .find(VISION_DATA_END)
                    .ok_or("unterminated image input")?;
                let uri = &encoded[..end];
                let payload = uri
                    .strip_prefix("data:image/png;base64,")
                    .or_else(|| uri.strip_prefix("data:image/jpeg;base64,"))
                    .ok_or("image input must be an inline PNG or JPEG data URI")?;
                total_images += 1;
                if total_images > 16 {
                    return Err("a request may contain at most 16 images".into());
                }
                let bytes = base64::engine::general_purpose::STANDARD
                    .decode(payload)
                    .map_err(|_| "image data URI contains invalid base64")?;
                total_bytes = total_bytes
                    .checked_add(bytes.len())
                    .ok_or("image input size overflow")?;
                if total_bytes > 64 * 1024 * 1024 {
                    return Err("image inputs exceed the 64 MiB request limit".into());
                }
                images.push(bytes);
                rest = &encoded[end + VISION_DATA_END.len()..];
            }
        }
        let mut encoded_images = self.executor.encode_visions(&images)?.into_iter();
        let mut embeddings = Vec::with_capacity(images.len());
        for message in &mut rendered {
            if !message.content.contains(VISION_DATA_START) {
                continue;
            }
            let mut content = String::with_capacity(message.content.len());
            let mut rest = message.content.as_str();
            while let Some(start) = rest.find(VISION_DATA_START) {
                content.push_str(&rest[..start]);
                let encoded = &rest[start + VISION_DATA_START.len()..];
                let end = encoded
                    .find(VISION_DATA_END)
                    .ok_or("unterminated image input")?;
                let embedding = encoded_images
                    .next()
                    .ok_or("vision encoder returned too few embeddings")?;
                content.push_str(VISION_TOKEN_START);
                content.push_str(&embedding.tokens.to_string());
                content.push_str(VISION_TOKEN_END);
                embeddings.push(embedding);
                rest = &encoded[end + VISION_DATA_END.len()..];
            }
            content.push_str(rest);
            message.content = content;
        }
        if encoded_images.next().is_some() {
            return Err("vision encoder returned too many embeddings".into());
        }
        let tokens = self
            .executor
            .model()
            .render_conversation(system, &rendered, reasoning);
        let mut overlays = Vec::with_capacity(embeddings.len());
        let mut cursor = 0_usize;
        for embedding in embeddings {
            let count = embedding.tokens as usize;
            let relative = tokens[cursor..]
                .windows(count + 2)
                .position(|window| {
                    window[0] == VISION_START_TOKEN
                        && window[count + 1] == VISION_END_TOKEN
                        && window[1..count + 1]
                            .iter()
                            .all(|token| *token == VISION_IMAGE_TOKEN)
                })
                .ok_or("rendered prompt lost an image placeholder")?;
            let start = cursor + relative + 1;
            overlays.push((
                u32::try_from(start).map_err(|_| "image prompt position overflow")?,
                embedding,
            ));
            cursor = start + count + 1;
        }
        Ok((tokens, overlays))
    }

    fn generate_inner(
        &mut self,
        messages: &[ChatTurn],
        settings: &TurnSettings,
        cancelled: &AtomicBool,
        emit: &mut impl FnMut(bool, String),
        progress: &mut impl FnMut(u32, u32, Option<f32>),
    ) -> Result<(GenerationOutput, bool), String> {
        let has_vision = messages
            .iter()
            .any(|message| message.content.contains(VISION_DATA_START));
        let (tokens, overlays) = if has_vision {
            self.render_multimodal_conversation(
                &settings.system_prompt,
                messages,
                settings.reasoning_mode,
            )?
        } else {
            let tokens = match messages.split_last() {
                Some((latest, history))
                    if latest.user
                        && self.executor.checkpoint_tag()
                            == conversation_tag(
                                &settings.system_prompt,
                                settings.reasoning_mode,
                                history,
                            ) =>
                {
                    let mut tokens = self.executor.tokens().to_vec();
                    tokens.extend(self.executor.model().render_continuation(
                        &latest.content,
                        settings.reasoning_mode,
                        latest.skip_previous_eos,
                    ));
                    tokens
                }
                _ => self.executor.model().render_conversation(
                    &settings.system_prompt,
                    messages,
                    settings.reasoning_mode,
                ),
            };
            (tokens, Vec::new())
        };
        if tokens.is_empty() {
            return Err("the rendered prompt is empty".into());
        }
        let max_context = self.executor.context() as usize;
        if tokens.len() >= max_context {
            return Err(format!(
                "Prompt has {} tokens, but the configured context size is {max_context} tokens",
                tokens.len()
            ));
        }
        let reused = self.executor.align_prompt(&tokens)?;
        self.executor.set_vision_overlays(overlays)?;
        self.metrics.kv_prefix_reused(reused);
        progress(self.executor.position(), self.executor.context(), None);
        let mut rng = Rng::new(settings.seed.unwrap_or(0x4453_3453_4552_5645));
        let mut reasoning = settings.reasoning_mode != ReasoningMode::Direct;
        let mut generated = ChatTurn {
            user: false,
            tool: false,
            system: false,
            skip_previous_eos: false,
            reasoning: reasoning.then(String::new),
            reasoning_complete: !reasoning,
            content: String::new(),
        };
        let mut emitted_reasoning = 0;
        let mut emitted_content = 0;
        let mut pending_utf8 = Vec::new();
        let prompt_tokens = tokens.len();
        let suffix = &tokens[reused..];
        let completed = self.prefill_suffix(&tokens, reused, has_vision, cancelled, progress)?;
        self.publish_execution_stats();
        if completed != suffix.len() {
            return Ok((
                GenerationOutput {
                    message: generated,
                    prompt_tokens: prompt_tokens as u32,
                    cached_tokens: reused as u32,
                    completion_tokens: 0,
                    finish_reason: "stop",
                    previous_checkpoint_bytes: None,
                    checkpoint_bytes: 0,
                },
                false,
            ));
        }
        progress(self.executor.position(), self.executor.context(), Some(0.0));
        let generation_started = Instant::now();
        let mut generated_tokens = 0_u32;
        let generation_limit = settings
            .max_generated_tokens
            .max(0)
            .min((max_context - self.executor.position() as usize) as i32)
            as u32;
        while generated_tokens < generation_limit {
            if cancelled.load(Ordering::Relaxed) {
                append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                flush_generated(
                    &mut generated,
                    &mut emitted_reasoning,
                    &mut emitted_content,
                    &settings.stops,
                    emit,
                );
                return Ok((
                    GenerationOutput {
                        message: generated,
                        prompt_tokens: prompt_tokens as u32,
                        cached_tokens: reused as u32,
                        completion_tokens: generated_tokens,
                        finish_reason: "stop",
                        previous_checkpoint_bytes: None,
                        checkpoint_bytes: 0,
                    },
                    true,
                ));
            }
            let token = sample(
                self.executor.logits(),
                settings.temperature,
                settings.top_p,
                settings.min_p,
                settings.top_k,
                &mut rng,
            );
            if self
                .executor
                .model()
                .is_stop_token_for_reasoning(token, settings.reasoning_mode)
            {
                append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                flush_generated(
                    &mut generated,
                    &mut emitted_reasoning,
                    &mut emitted_content,
                    &settings.stops,
                    emit,
                );
                return Ok((
                    GenerationOutput {
                        message: generated,
                        prompt_tokens: prompt_tokens as u32,
                        cached_tokens: reused as u32,
                        completion_tokens: generated_tokens,
                        finish_reason: "stop",
                        previous_checkpoint_bytes: None,
                        checkpoint_bytes: 0,
                    },
                    true,
                ));
            }
            let cycle = if settings.temperature <= 0.0 {
                self.executor.eval_speculative_greedy(
                    token,
                    generation_limit - generated_tokens,
                    settings.reasoning_mode,
                    cancelled,
                )?
            } else {
                self.executor.eval_speculative_sampled(
                    token,
                    generation_limit - generated_tokens,
                    settings.reasoning_mode,
                    settings.temperature,
                    settings.top_p,
                    settings.min_p,
                    settings.top_k,
                    &mut rng,
                    cancelled,
                )?
            };
            self.publish_execution_stats();
            for token in cycle {
                if generated_tokens >= generation_limit
                    || self
                        .executor
                        .model()
                        .is_stop_token_for_reasoning(token, settings.reasoning_mode)
                {
                    append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                    flush_generated(
                        &mut generated,
                        &mut emitted_reasoning,
                        &mut emitted_content,
                        &settings.stops,
                        emit,
                    );
                    return Ok((
                        GenerationOutput {
                            message: generated,
                            prompt_tokens: prompt_tokens as u32,
                            cached_tokens: reused as u32,
                            completion_tokens: generated_tokens,
                            finish_reason: "stop",
                            previous_checkpoint_bytes: None,
                            checkpoint_bytes: 0,
                        },
                        true,
                    ));
                }
                if self.executor.model().is_think_start_token(token) {
                    append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                    reasoning = true;
                    generated.reasoning.get_or_insert_default();
                } else if self.executor.model().is_think_end_token(token) {
                    append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
                    if let Some(reasoning_text) = &mut generated.reasoning
                        && emit_safe_text(
                            reasoning_text,
                            &mut emitted_reasoning,
                            &settings.stops,
                            true,
                            true,
                            emit,
                        )
                    {
                        return Ok((
                            GenerationOutput {
                                message: generated,
                                prompt_tokens: prompt_tokens as u32,
                                cached_tokens: reused as u32,
                                completion_tokens: generated_tokens + 1,
                                finish_reason: "stop",
                                previous_checkpoint_bytes: None,
                                checkpoint_bytes: 0,
                            },
                            false,
                        ));
                    }
                    reasoning = false;
                    generated.reasoning_complete = true;
                    emit(false, String::new());
                } else if let Some(bytes) = self.executor.model().token_bytes(token) {
                    append_generated_bytes(
                        &mut generated,
                        reasoning,
                        &mut pending_utf8,
                        &bytes,
                        false,
                    );
                    let stopped = if reasoning {
                        let text = generated.reasoning.get_or_insert_default();
                        emit_safe_text(
                            text,
                            &mut emitted_reasoning,
                            &settings.stops,
                            false,
                            true,
                            emit,
                        )
                    } else {
                        generated.reasoning_complete = true;
                        emit_safe_text(
                            &mut generated.content,
                            &mut emitted_content,
                            &settings.stops,
                            false,
                            false,
                            emit,
                        )
                    };
                    if stopped {
                        return Ok((
                            GenerationOutput {
                                message: generated,
                                prompt_tokens: prompt_tokens as u32,
                                cached_tokens: reused as u32,
                                completion_tokens: generated_tokens + 1,
                                finish_reason: "stop",
                                previous_checkpoint_bytes: None,
                                checkpoint_bytes: 0,
                            },
                            false,
                        ));
                    }
                }
                generated_tokens += 1;
                progress(
                    self.executor.position(),
                    self.executor.context(),
                    Some(
                        generated_tokens as f32
                            / generation_started.elapsed().as_secs_f32().max(1.0e-6),
                    ),
                );
            }
        }
        append_generated_bytes(&mut generated, reasoning, &mut pending_utf8, &[], true);
        flush_generated(
            &mut generated,
            &mut emitted_reasoning,
            &mut emitted_content,
            &settings.stops,
            emit,
        );
        Ok((
            GenerationOutput {
                message: generated,
                prompt_tokens: prompt_tokens as u32,
                cached_tokens: reused as u32,
                completion_tokens: generated_tokens,
                finish_reason: "length",
                previous_checkpoint_bytes: None,
                checkpoint_bytes: 0,
            },
            true,
        ))
    }
}

#[cfg(any(target_os = "macos", test))]
fn append_generated_bytes(
    generated: &mut ChatTurn,
    reasoning: bool,
    pending: &mut Vec<u8>,
    bytes: &[u8],
    final_flush: bool,
) {
    let text = if reasoning {
        generated.reasoning.get_or_insert_default()
    } else {
        &mut generated.content
    };
    pending.extend_from_slice(bytes);
    loop {
        match std::str::from_utf8(pending) {
            Ok(valid) => {
                text.push_str(valid);
                pending.clear();
                return;
            }
            Err(error) => {
                let valid = error.valid_up_to();
                text.push_str(std::str::from_utf8(&pending[..valid]).unwrap());
                pending.drain(..valid);
                match error.error_len() {
                    Some(length) => {
                        text.push('\u{fffd}');
                        pending.drain(..length);
                    }
                    None if final_flush => {
                        text.push_str(&String::from_utf8_lossy(pending));
                        pending.clear();
                        return;
                    }
                    None => return,
                }
            }
        }
    }
}

#[cfg(target_os = "macos")]
fn flush_generated(
    generated: &mut ChatTurn,
    emitted_reasoning: &mut usize,
    emitted_content: &mut usize,
    stops: &[String],
    emit: &mut impl FnMut(bool, String),
) {
    if let Some(reasoning) = &mut generated.reasoning {
        let _ = emit_safe_text(reasoning, emitted_reasoning, stops, true, true, emit);
    }
    let _ = emit_safe_text(
        &mut generated.content,
        emitted_content,
        stops,
        true,
        false,
        emit,
    );
}

#[cfg(any(target_os = "macos", test))]
fn emit_safe_text(
    text: &mut String,
    emitted: &mut usize,
    stops: &[String],
    final_flush: bool,
    reasoning: bool,
    emit: &mut impl FnMut(bool, String),
) -> bool {
    let stop = stops
        .iter()
        .filter_map(|stop| {
            text[*emitted..]
                .find(stop)
                .map(|position| *emitted + position)
        })
        .min();
    if let Some(stop) = stop {
        if stop > *emitted {
            emit(reasoning, text[*emitted..stop].to_owned());
        }
        text.truncate(stop);
        *emitted = stop;
        return true;
    }
    let hold = if final_flush {
        0
    } else {
        stops
            .iter()
            .map(|stop| stop.len().saturating_sub(1))
            .max()
            .unwrap_or(0)
    };
    let mut safe = text.len().saturating_sub(hold);
    while safe > *emitted && !text.is_char_boundary(safe) {
        safe -= 1;
    }
    if safe > *emitted {
        emit(reasoning, text[*emitted..safe].to_owned());
        *emitted = safe;
    }
    false
}

#[cfg(any(target_os = "macos", test))]
fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> Vec<u8> {
    fn text(output: &mut Vec<u8>, value: &str) {
        output.extend_from_slice(&(value.len() as u64).to_le_bytes());
        output.extend_from_slice(value.as_bytes());
    }

    let mut output = b"DS4Server chat checkpoint v5".to_vec();
    text(&mut output, system);
    output.push(match reasoning {
        ReasoningMode::Direct => 0,
        ReasoningMode::Low => 1,
        ReasoningMode::High => 2,
        ReasoningMode::Max => 3,
    });
    for message in messages {
        output.push(u8::from(message.user));
        output.push(u8::from(message.tool));
        output.push(u8::from(message.system));
        output.push(u8::from(message.skip_previous_eos));
        match &message.reasoning {
            Some(reasoning) => {
                output.push(1);
                text(&mut output, reasoning);
            }
            None => output.push(0),
        }
        output.push(u8::from(message.reasoning_complete));
        text(&mut output, &message.content);
    }
    output
}

#[cfg(any(target_os = "macos", test))]
fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] {
    Sha256::digest(conversation_key(system, reasoning, messages)).into()
}

#[cfg(any(target_os = "macos", test))]
fn checkpoint_matches_prefix(
    checkpoint: [u8; 32],
    system: &str,
    reasoning: ReasoningMode,
    messages: &[ChatTurn],
) -> bool {
    // ponytail: appended control messages are few; carry incremental hashes if
    // scanning a genuinely changed, very long history becomes measurable.
    (0..messages.len())
        .rev()
        .any(|end| checkpoint == conversation_tag(system, reasoning, &messages[..end]))
}

#[cfg(target_os = "macos")]
fn resident_key(directory: &Path, tag: [u8; 32]) -> PathBuf {
    let mut name = String::with_capacity(64);
    for byte in tag {
        use std::fmt::Write;
        let _ = write!(name, "{byte:02x}");
    }
    directory.join("resident").join(name)
}

#[cfg(any(target_os = "macos", test))]
fn sample(
    logits: &[f32],
    temperature: f32,
    top_p: f32,
    min_p: f32,
    top_k: i32,
    rng: &mut Rng,
) -> i32 {
    let probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k);
    sample_probabilities(&probabilities, rng, None)
}

#[cfg(any(target_os = "macos", test))]
fn sampling_probabilities(
    logits: &[f32],
    temperature: f32,
    top_p: f32,
    min_p: f32,
    top_k: i32,
) -> Vec<(usize, f32)> {
    let greedy = || {
        vec![(
            logits
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.total_cmp(b.1))
                .map_or(0, |(index, _)| index),
            1.0,
        )]
    };
    if temperature <= 0.0 {
        return greedy();
    }
    let maximum = logits
        .iter()
        .copied()
        .filter(|value| value.is_finite())
        .fold(f32::NEG_INFINITY, f32::max);
    if !maximum.is_finite() {
        return greedy();
    }
    let top_p = if top_p <= 0.0 || top_p > 1.0 {
        1.0
    } else {
        top_p
    };
    let min_p = min_p.max(0.0);
    let mut probabilities: Vec<(usize, f32)> = logits
        .iter()
        .enumerate()
        .filter(|(_, logit)| logit.is_finite())
        .map(|(index, logit)| (index, ((*logit - maximum) / temperature).exp()))
        .collect();
    if probabilities.is_empty() {
        return greedy();
    }
    if top_p < 1.0 || top_k > 0 || min_p > 0.0 {
        probabilities.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
        if top_k > 0 {
            probabilities.truncate(probabilities.len().min((top_k as usize).min(1024)));
        }
    }
    let total: f32 = probabilities
        .iter()
        .map(|(_, probability)| probability)
        .sum();
    let mut kept = 0.0;
    let mut count = 0;
    for (_, probability) in &probabilities {
        if count > 0 && *probability < min_p {
            break;
        }
        kept += *probability;
        count += 1;
        if kept / total >= top_p {
            break;
        }
    }
    probabilities.truncate(count);
    if probabilities.is_empty() || !kept.is_finite() || kept <= 0.0 {
        return greedy();
    }
    for (_, probability) in &mut probabilities {
        *probability /= kept;
    }
    if top_p >= 1.0 && top_k <= 0 {
        probabilities.sort_unstable_by_key(|(token, _)| *token);
    }
    probabilities
}

#[cfg(any(target_os = "macos", test))]
fn sample_probabilities(
    probabilities: &[(usize, f32)],
    rng: &mut Rng,
    excluded: Option<usize>,
) -> i32 {
    let total: f32 = probabilities
        .iter()
        .filter(|(token, _)| Some(*token) != excluded)
        .map(|(_, probability)| probability)
        .sum();
    let mut choice = rng.unit() * total;
    for (token, probability) in probabilities {
        if Some(*token) == excluded {
            continue;
        }
        choice -= probability;
        if choice <= 0.0 {
            return *token as i32;
        }
    }
    probabilities
        .iter()
        .rev()
        .find(|(token, _)| Some(*token) != excluded)
        .map_or(0, |(token, _)| *token as i32)
}

#[cfg(any(target_os = "macos", test))]
fn exact_delta_sample(
    logits: &[f32],
    draft: i32,
    temperature: f32,
    top_p: f32,
    min_p: f32,
    top_k: i32,
    rng: &mut Rng,
) -> (i32, bool) {
    let mut probabilities = sampling_probabilities(logits, temperature, top_p, min_p, top_k);
    let draft_probability = probabilities
        .iter()
        .find(|(token, _)| *token == draft as usize)
        .map_or(0.0, |(_, probability)| *probability);
    if rng.unit() <= draft_probability {
        return (draft, true);
    }
    probabilities.sort_unstable_by_key(|(token, _)| *token);
    (
        sample_probabilities(&probabilities, rng, Some(draft as usize)),
        false,
    )
}

#[cfg(any(target_os = "macos", test))]
struct Rng(u64);

#[cfg(any(target_os = "macos", test))]
impl Rng {
    fn new(seed: u64) -> Self {
        Self(seed.max(1))
    }

    fn unit(&mut self) -> f32 {
        let mut value = self.0;
        if value == 0 {
            value = 0x9e37_79b9_7f4a_7c15;
        }
        value ^= value >> 12;
        value ^= value << 25;
        value ^= value >> 27;
        self.0 = value;
        let value = value.wrapping_mul(0x2545_f491_4f6c_dd1d);
        ((value >> 40) & 0xff_ffff) as f32 / 16_777_216.0
    }
}

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

    #[test]
    fn zero_temperature_is_greedy() {
        let mut rng = Rng::new(1);
        assert_eq!(sample(&[1.0, 4.0, 2.0], 0.0, 1.0, 0.0, 0, &mut rng), 1);
    }

    #[test]
    fn top_k_and_stream_stops_are_applied_before_output() {
        let mut rng = Rng::new(1);
        assert_eq!(sample(&[1.0, 4.0, 2.0], 1.0, 1.0, 0.0, 1, &mut rng), 1);

        let mut text = "hello STOP hidden".to_owned();
        let mut emitted = 0;
        let mut chunks = Vec::new();
        assert!(emit_safe_text(
            &mut text,
            &mut emitted,
            &["STOP".into()],
            false,
            false,
            &mut |_, chunk| chunks.push(chunk),
        ));
        assert_eq!(text, "hello ");
        assert_eq!(chunks, ["hello "]);
    }

    #[test]
    fn exact_delta_sampling_accepts_or_corrects_the_draft() {
        let mut accept_rng = Rng::new(2);
        let (accepted, was_draft) =
            exact_delta_sample(&[10.0, 0.0], 0, 1.0, 1.0, 0.0, 0, &mut accept_rng);
        assert_eq!((accepted, was_draft), (0, true));

        let mut reject_rng = Rng::new(1);
        let (replacement, was_draft) =
            exact_delta_sample(&[0.0, 10.0], 0, 1.0, 1.0, 0.0, 0, &mut reject_rng);
        assert_eq!((replacement, was_draft), (1, false));
    }

    #[test]
    fn sampling_probabilities_match_ds4_filter_order() {
        let probabilities = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 0.8, 0.2, 0);
        assert_eq!(probabilities.len(), 2);
        assert_eq!(probabilities[0].0, 1);
        assert_eq!(probabilities[1].0, 2);
        assert!((probabilities.iter().map(|(_, value)| value).sum::<f32>() - 1.0).abs() < 1e-6);

        let min_p_only = sampling_probabilities(&[0.0, 2.0, 1.0], 1.0, 1.0, 0.2, 0);
        assert_eq!(
            min_p_only
                .iter()
                .map(|(token, _)| *token)
                .collect::<Vec<_>>(),
            [1, 2]
        );
    }

    #[test]
    fn split_utf8_token_bytes_are_joined_before_decoding() {
        let mut generated = ChatTurn {
            user: false,
            tool: false,
            system: false,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: String::new(),
        };
        let mut pending = Vec::new();

        append_generated_bytes(&mut generated, false, &mut pending, &[0xf0, 0x9f], false);
        assert!(generated.content.is_empty());
        append_generated_bytes(&mut generated, false, &mut pending, &[0x98, 0x8a], false);

        assert_eq!(generated.content, "😊");
        assert!(pending.is_empty());
    }

    #[test]
    fn checkpoint_tag_covers_the_canonical_chat_state() {
        let mut messages = vec![ChatTurn {
            user: true,
            tool: false,
            system: false,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: "Hello".into(),
        }];
        let tag = conversation_tag("System", ReasoningMode::High, &messages);
        assert_eq!(
            tag,
            conversation_tag("System", ReasoningMode::High, &messages)
        );
        assert_ne!(
            tag,
            conversation_tag("Changed", ReasoningMode::High, &messages)
        );
        assert_ne!(
            tag,
            conversation_tag("System", ReasoningMode::Direct, &messages)
        );
        messages[0].system = true;
        assert_ne!(
            tag,
            conversation_tag("System", ReasoningMode::High, &messages)
        );
        messages[0].system = false;
        messages[0].skip_previous_eos = true;
        assert_ne!(
            tag,
            conversation_tag("System", ReasoningMode::High, &messages)
        );
        let prefix = conversation_key("System", ReasoningMode::High, &messages);
        messages.push(ChatTurn {
            user: false,
            tool: false,
            system: false,
            skip_previous_eos: false,
            reasoning: Some("because".into()),
            reasoning_complete: true,
            content: "Hi".into(),
        });
        assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix));
    }

    #[test]
    fn checkpoint_tag_accepts_an_unchanged_prefix_before_reminders() {
        let mut messages = vec![ChatTurn {
            user: true,
            tool: false,
            system: false,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: "Question".into(),
        }];
        let checkpoint = conversation_tag("System", ReasoningMode::High, &messages);
        messages.extend([
            ChatTurn {
                user: false,
                tool: true,
                system: false,
                skip_previous_eos: false,
                reasoning: None,
                reasoning_complete: true,
                content: "Tool result".into(),
            },
            ChatTurn {
                user: false,
                tool: false,
                system: true,
                skip_previous_eos: false,
                reasoning: None,
                reasoning_complete: true,
                content: "System prompt reminder".into(),
            },
        ]);

        assert!(checkpoint_matches_prefix(
            checkpoint,
            "System",
            ReasoningMode::High,
            &messages,
        ));
        assert!(!checkpoint_matches_prefix(
            checkpoint,
            "Changed",
            ReasoningMode::High,
            &messages,
        ));
    }

    #[test]
    fn workspace_instruction_identity_participates_in_append_only_kv_keys() {
        let system = "System";
        let bootstrap = conversation_key(system, ReasoningMode::High, &[]);
        let mut messages = vec![
            ChatTurn {
                user: false,
                tool: false,
                system: true,
                skip_previous_eos: false,
                reasoning: None,
                reasoning_complete: true,
                content: "<system-reminder>\nWorkspace instruction identity: abc\nInstructions from: /project/AGENTS.md\nkeep this\n</system-reminder>".into(),
            },
            ChatTurn {
                user: true,
                tool: false,
                system: false,
                skip_previous_eos: false,
                reasoning: None,
                reasoning_complete: true,
                content: "hello".into(),
            },
        ];

        assert!(conversation_key(system, ReasoningMode::High, &messages).starts_with(&bootstrap));
        let mut changed_opening = messages[0].clone();
        changed_opening.content = changed_opening
            .content
            .replace("identity: abc", "identity: def");
        assert_ne!(
            conversation_tag(system, ReasoningMode::High, &messages[..1]),
            conversation_tag(system, ReasoningMode::High, &[changed_opening])
        );
        let baseline = conversation_key(system, ReasoningMode::High, &messages);
        messages.push(ChatTurn {
            user: false,
            tool: false,
            system: true,
            skip_previous_eos: false,
            reasoning: None,
            reasoning_complete: true,
            content: "<system-reminder>\nWorkspace instruction identity: def\nReplacement instructions from: /project/AGENTS.md\nkeep that\n</system-reminder>".into(),
        });
        let changed = conversation_key(system, ReasoningMode::High, &messages);
        assert!(changed.starts_with(&baseline));
    }

    #[test]
    fn checkpoint_rebuilds_explain_compatibility_and_history_misses() {
        assert!(checkpoint_rebuild_activity(true).contains("different model"));
        assert!(checkpoint_rebuild_activity(false).contains("history"));
    }

    #[test]
    #[cfg(target_os = "macos")]
    #[ignore = "requires the 80 GiB Flash checkpoint and Apple Metal"]
    fn metal_executes_real_flash_token() {
        configure_metal_sources().unwrap();
        let path = crate::model::engine_artifacts(
            ModelChoice::DeepSeekV4Flash0731,
            false,
            &crate::app::models_path(),
        )
        .model;
        let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash0731).unwrap();
        let tokens = model.render_prompt(
            "You are a helpful assistant",
            "Hello",
            ReasoningMode::Direct,
        );
        assert_eq!(tokens.len(), 10);
        let mut executor = metal::Executor::open(model, 32_768, false, 0).unwrap();
        assert_eq!(executor.context(), 32_768);
        for &token in &tokens {
            executor.eval(token).unwrap();
        }
        assert!(executor.logits().iter().all(|logit| logit.is_finite()));
        let argmax = executor
            .logits()
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap();
        eprintln!(
            "Rust logits: argmax={} value={} logit0={}",
            argmax.0,
            argmax.1,
            executor.logits()[0]
        );
        assert_eq!(argmax.0, 19_923);
        assert!((executor.logits()[0] - -7.675_424).abs() < 0.1);
        let next = argmax.0 as i32;
        let checkpoint =
            std::env::temp_dir().join(format!("ds4-rust-kv-{}.bin", std::process::id()));
        executor
            .save_checkpoint(&checkpoint, [7; 32], &mut |_| {})
            .unwrap();
        executor.eval(next).unwrap();
        let continued_logit = executor.logits()[0];
        let continued_argmax = executor
            .logits()
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0;
        drop(executor);

        let model = Model::open_main(&path, ModelChoice::DeepSeekV4Flash0731).unwrap();
        let mut restored = metal::Executor::open(model, 32_768, false, 0).unwrap();
        assert!(restored.load_checkpoint(&checkpoint, &mut |_| {}).unwrap());
        assert_eq!(restored.position(), tokens.len() as u32);
        assert_eq!(restored.tokens(), tokens);
        assert_eq!(restored.checkpoint_tag(), [7; 32]);
        restored.eval(next).unwrap();
        assert_eq!(
            restored
                .logits()
                .iter()
                .enumerate()
                .max_by(|a, b| a.1.total_cmp(b.1))
                .unwrap()
                .0,
            continued_argmax
        );
        assert!((restored.logits()[0] - continued_logit).abs() < 1.0e-5);
        std::fs::remove_file(checkpoint).unwrap();
    }
}