memra-server 0.81.2

//! Durable per-request usage and cost receipts.
//!
//! The worker remains the sole source of prompt/cache/completion counts.  The HTTP
//! task appends and syncs the corresponding cost row before publishing a terminal
//! response, so receipt I/O never runs on the CUDA-owner thread.

use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::Write;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};

use serde_json::json;

use super::OpenRouterModelMetadata;

const FORMAT: &str = "memra.request-cost.v1";
const MAX_DECIMAL_SCALE: u32 = 18;

#[derive(Debug, Clone, PartialEq, Eq)]
struct Decimal {
    coefficient: u128,
    scale: u32,
}

impl Decimal {
    fn parse(value: &str) -> Result<Self, String> {
        let (whole, fraction) = value.split_once('.').unwrap_or((value, ""));
        if whole.is_empty()
            || !whole.bytes().all(|b| b.is_ascii_digit())
            || !fraction.bytes().all(|b| b.is_ascii_digit())
        {
            return Err(format!("invalid non-negative decimal {value:?}"));
        }
        let scale = u32::try_from(fraction.len())
            .map_err(|_| format!("decimal scale is too large in {value:?}"))?;
        if scale > MAX_DECIMAL_SCALE {
            return Err(format!(
                "decimal {value:?} has scale {scale}; maximum supported scale is {MAX_DECIMAL_SCALE}"
            ));
        }
        let digits = format!("{whole}{fraction}");
        let coefficient = digits
            .parse::<u128>()
            .map_err(|_| format!("decimal coefficient overflows u128 in {value:?}"))?;
        Ok(Self { coefficient, scale })
    }

    fn checked_mul_u64(&self, count: u64) -> Result<Self, String> {
        let coefficient = self
            .coefficient
            .checked_mul(count as u128)
            .ok_or_else(|| "request cost multiplication overflowed u128".to_string())?;
        Ok(Self { coefficient, scale: self.scale })
    }

    fn checked_add(&self, other: &Self) -> Result<Self, String> {
        let scale = self.scale.max(other.scale);
        let left = scale_up(self.coefficient, scale - self.scale)?;
        let right = scale_up(other.coefficient, scale - other.scale)?;
        let coefficient = left
            .checked_add(right)
            .ok_or_else(|| "request cost addition overflowed u128".to_string())?;
        Ok(Self { coefficient, scale })
    }

    fn checked_cmp(&self, other: &Self) -> Result<std::cmp::Ordering, String> {
        let scale = self.scale.max(other.scale);
        Ok(scale_up(self.coefficient, scale - self.scale)?
            .cmp(&scale_up(other.coefficient, scale - other.scale)?))
    }

    fn render(&self) -> String {
        if self.scale == 0 {
            return self.coefficient.to_string();
        }
        let mut digits = self.coefficient.to_string();
        let scale = self.scale as usize;
        if digits.len() <= scale {
            digits.insert_str(0, &"0".repeat(scale + 1 - digits.len()));
        }
        let split = digits.len() - scale;
        digits.insert(split, '.');
        digits
    }
}

fn scale_up(value: u128, places: u32) -> Result<u128, String> {
    let mut factor = 1u128;
    for _ in 0..places {
        factor = factor
            .checked_mul(10)
            .ok_or_else(|| "request cost decimal scale overflowed u128".to_string())?;
    }
    value
        .checked_mul(factor)
        .ok_or_else(|| "request cost decimal alignment overflowed u128".to_string())
}

#[derive(Debug, Clone)]
struct PriceSchedule {
    prompt_text: String,
    cached_prompt_text: String,
    completion_text: String,
    request_text: String,
    prompt: Decimal,
    cached_prompt: Decimal,
    completion: Decimal,
    request: Decimal,
}

impl PriceSchedule {
    fn from_metadata(alias: &str, metadata: &OpenRouterModelMetadata) -> Result<Self, String> {
        let required = |field: &str, value: &Option<String>| {
            value.clone().ok_or_else(|| {
                format!(
                    "model {alias:?}: {field} is required when MEMRA_REQUEST_LEDGER is enabled"
                )
            })
        };
        let prompt_text = required("pricing.prompt", &metadata.pricing.prompt)?;
        let cached_prompt_text =
            required("pricing.cached_prompt", &metadata.pricing.cached_prompt)?;
        let completion_text = required("pricing.completion", &metadata.pricing.completion)?;
        let request_text = metadata
            .pricing
            .request
            .clone()
            .unwrap_or_else(|| "0".into());
        let schedule = Self {
            prompt: Decimal::parse(&prompt_text)
                .map_err(|e| format!("model {alias:?}: pricing.prompt: {e}"))?,
            cached_prompt: Decimal::parse(&cached_prompt_text)
                .map_err(|e| format!("model {alias:?}: pricing.cached_prompt: {e}"))?,
            completion: Decimal::parse(&completion_text)
                .map_err(|e| format!("model {alias:?}: pricing.completion: {e}"))?,
            request: Decimal::parse(&request_text)
                .map_err(|e| format!("model {alias:?}: pricing.request: {e}"))?,
            prompt_text,
            cached_prompt_text,
            completion_text,
            request_text,
        };
        if schedule.cached_prompt.checked_cmp(&schedule.prompt)?
            != std::cmp::Ordering::Less
        {
            return Err(format!(
                "model {alias:?}: pricing.cached_prompt must be lower than pricing.prompt"
            ));
        }
        Ok(schedule)
    }

    fn cost(&self, usage: Usage) -> Result<serde_json::Value, String> {
        let ordinary_prompt = usage
            .prompt_tokens
            .checked_sub(usage.cached_prompt_tokens)
            .ok_or_else(|| {
                format!(
                    "cached prompt tokens {} exceed total prompt tokens {}",
                    usage.cached_prompt_tokens, usage.prompt_tokens
                )
            })?;
        let ordinary_cost = self.prompt.checked_mul_u64(ordinary_prompt)?;
        let cached_cost = self.cached_prompt.checked_mul_u64(usage.cached_prompt_tokens)?;
        let completion_cost = self.completion.checked_mul_u64(usage.completion_tokens)?;
        let request_cost = self.request.clone();
        let total = ordinary_cost
            .checked_add(&cached_cost)?
            .checked_add(&completion_cost)?
            .checked_add(&request_cost)?;
        Ok(json!({
            "usage": {
                "prompt_tokens": usage.prompt_tokens,
                "cached_prompt_tokens": usage.cached_prompt_tokens,
                "ordinary_prompt_tokens": ordinary_prompt,
                "completion_tokens": usage.completion_tokens,
                "total_tokens": usage.prompt_tokens.checked_add(usage.completion_tokens)
                    .ok_or_else(|| "request token total overflowed u64".to_string())?,
            },
            "unit_prices_usd": {
                "prompt": self.prompt_text,
                "cached_prompt": self.cached_prompt_text,
                "completion": self.completion_text,
                "request": self.request_text,
            },
            "cost_usd": {
                "ordinary_prompt": ordinary_cost.render(),
                "cached_prompt": cached_cost.render(),
                "completion": completion_cost.render(),
                "request": request_cost.render(),
                "total": total.render(),
            },
        }))
    }
}

#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct Usage {
    pub(crate) prompt_tokens: u64,
    pub(crate) cached_prompt_tokens: u64,
    pub(crate) completion_tokens: u64,
}

#[derive(Clone)]
pub(crate) struct Ledger {
    inner: Arc<LedgerInner>,
}

struct LedgerInner {
    path: PathBuf,
    file: Mutex<File>,
    prices: HashMap<String, PriceSchedule>,
    failed: AtomicBool,
}

impl Ledger {
    pub(crate) fn from_env(
        models: &[(String, String, Option<String>)],
        metadata: &HashMap<String, OpenRouterModelMetadata>,
    ) -> Result<Option<Self>, String> {
        let Some(path) = std::env::var_os("MEMRA_REQUEST_LEDGER") else {
            return Ok(None);
        };
        if path.is_empty() {
            return Err("MEMRA_REQUEST_LEDGER must not be empty".into());
        }
        let mut prices = HashMap::new();
        for (alias, _, _) in models {
            let model_metadata = metadata.get(alias).ok_or_else(|| {
                format!(
                    "model {alias:?}: MEMRA_MODEL_METADATA entry is required when \
                     MEMRA_REQUEST_LEDGER is enabled"
                )
            })?;
            prices.insert(alias.clone(), PriceSchedule::from_metadata(alias, model_metadata)?);
        }
        let ledger = Self::open(Path::new(&path), prices)?;
        eprintln!(
            "[ledger] durable request-cost ledger enabled: {}",
            ledger.inner.path.display()
        );
        Ok(Some(ledger))
    }

    fn open(path: &Path, prices: HashMap<String, PriceSchedule>) -> Result<Self, String> {
        #[cfg(unix)]
        use std::os::unix::fs::OpenOptionsExt as _;

        let mut options = OpenOptions::new();
        options.create(true).append(true);
        #[cfg(unix)]
        options.mode(0o640).custom_flags(libc::O_NOFOLLOW);
        let file = options
            .open(path)
            .map_err(|e| format!("open request ledger {}: {e}", path.display()))?;
        let metadata = file
            .metadata()
            .map_err(|e| format!("stat request ledger {}: {e}", path.display()))?;
        if !metadata.is_file() {
            return Err(format!("request ledger {} is not a regular file", path.display()));
        }
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt as _;
            let mode = metadata.permissions().mode() & 0o777;
            if mode & 0o137 != 0 {
                return Err(format!(
                    "request ledger {} must have 0600 or 0640-class permissions; found {mode:04o}",
                    path.display(),
                ));
            }
        }
        Ok(Self {
            inner: Arc::new(LedgerInner {
                path: path.to_path_buf(),
                file: Mutex::new(file),
                prices,
                failed: AtomicBool::new(false),
            }),
        })
    }

    pub(crate) fn start(
        &self,
        request_id: &str,
        tenant: &str,
        model: &str,
        route: &'static str,
        lane: &'static str,
        stream: bool,
    ) -> PendingReceipt {
        PendingReceipt {
            ledger: self.clone(),
            request_id: request_id.into(),
            tenant: tenant.into(),
            model: model.into(),
            route,
            lane,
            stream,
            started_unix_ms: unix_ms(),
            partial_usage: Usage::default(),
            finalized: false,
        }
    }

    fn append(&self, row: &serde_json::Value) -> Result<(), String> {
        if self.inner.failed.load(Ordering::Acquire) {
            return Err(format!(
                "request ledger {} is latched unavailable after an earlier write failure",
                self.inner.path.display()
            ));
        }
        let mut encoded = serde_json::to_vec(row)
            .map_err(|e| format!("serialize request ledger row: {e}"))?;
        encoded.push(b'\n');
        let mut file = match self.inner.file.lock() {
            Ok(file) => file,
            Err(_) => {
                self.inner.failed.store(true, Ordering::Release);
                return Err("request ledger writer mutex is poisoned".into());
            }
        };
        if let Err(err) = file.write_all(&encoded) {
            self.inner.failed.store(true, Ordering::Release);
            return Err(format!(
                "append request ledger {}: {err}",
                self.inner.path.display()
            ));
        }
        if let Err(err) = file.sync_data() {
            self.inner.failed.store(true, Ordering::Release);
            return Err(format!(
                "sync request ledger {}: {err}",
                self.inner.path.display()
            ));
        }
        Ok(())
    }

    #[cfg(test)]
    pub(crate) fn for_test(
        path: &Path,
        alias: &str,
        metadata: &OpenRouterModelMetadata,
    ) -> Self {
        let schedule = PriceSchedule::from_metadata(alias, metadata).unwrap();
        Self::open(path, HashMap::from([(alias.to_string(), schedule)])).unwrap()
    }
}

pub(crate) struct PendingReceipt {
    ledger: Ledger,
    request_id: String,
    tenant: String,
    model: String,
    route: &'static str,
    lane: &'static str,
    stream: bool,
    started_unix_ms: u64,
    partial_usage: Usage,
    finalized: bool,
}

impl PendingReceipt {
    pub(crate) fn record_prompt_usage(
        &mut self,
        prompt_tokens: u64,
        cached_prompt_tokens: u64,
    ) -> Result<(), String> {
        if cached_prompt_tokens > prompt_tokens {
            return Err(format!(
                "cached prompt tokens {cached_prompt_tokens} exceed total prompt tokens {prompt_tokens}"
            ));
        }
        self.partial_usage.prompt_tokens = prompt_tokens;
        self.partial_usage.cached_prompt_tokens = cached_prompt_tokens;
        Ok(())
    }

    pub(crate) fn record_completion_token(&mut self) -> Result<(), String> {
        self.partial_usage.completion_tokens = self
            .partial_usage
            .completion_tokens
            .checked_add(1)
            .ok_or_else(|| "partial completion token count overflowed u64".to_string())?;
        Ok(())
    }

    pub(crate) fn complete(&mut self, usage: Usage, worker_elapsed_s: f64) -> Result<(), String> {
        let mut row = self.base_row("completed", 200, None);
        self.add_accounting(&mut row, usage)?;
        row["worker_elapsed_s"] = json!(worker_elapsed_s);
        self.finalize(row)
    }

    pub(crate) fn reject(&mut self, status: u16, error_code: &str) -> Result<(), String> {
        let row = self.base_row("rejected", status, Some(error_code));
        self.finalize(row)
    }

    fn base_row(
        &self,
        outcome: &str,
        http_status: u16,
        error_code: Option<&str>,
    ) -> serde_json::Value {
        json!({
            "format": FORMAT,
            "request_id": self.request_id,
            "started_unix_ms": self.started_unix_ms,
            "finished_unix_ms": unix_ms(),
            "tenant": self.tenant,
            "model": self.model,
            "route": self.route,
            "lane": self.lane,
            "stream": self.stream,
            "outcome": outcome,
            "http_status": http_status,
            "error_code": error_code,
            "usage": null,
            "unit_prices_usd": null,
            "cost_usd": null,
        })
    }

    fn add_accounting(
        &self,
        row: &mut serde_json::Value,
        usage: Usage,
    ) -> Result<(), String> {
        let schedule = self
            .ledger
            .inner
            .prices
            .get(&self.model)
            .ok_or_else(|| format!("request ledger has no price schedule for {:?}", self.model))?;
        let accounting = schedule.cost(usage)?;
        row["usage"] = accounting["usage"].clone();
        row["unit_prices_usd"] = accounting["unit_prices_usd"].clone();
        row["cost_usd"] = accounting["cost_usd"].clone();
        Ok(())
    }

    fn finalize(&mut self, row: serde_json::Value) -> Result<(), String> {
        if self.finalized {
            return Err(format!(
                "request {} already has a terminal ledger row",
                self.request_id
            ));
        }
        // An append whose sync fails has an indeterminate durability state. Never retry the
        // same request id from Drop and risk a duplicate bill; fail the HTTP completion loud.
        self.finalized = true;
        self.ledger.append(&row)
    }
}

impl Drop for PendingReceipt {
    fn drop(&mut self) {
        if self.finalized {
            return;
        }
        let mut row = self.base_row(
            "abandoned",
            499,
            Some("client_disconnected_or_handler_dropped"),
        );
        if let Err(err) = self.add_accounting(&mut row, self.partial_usage) {
            eprintln!(
                "[ledger] ERROR: could not price abandoned request {}: {err}",
                self.request_id
            );
        }
        self.finalized = true;
        if let Err(err) = self.ledger.append(&row) {
            eprintln!(
                "[ledger] ERROR: could not persist abandoned request {}: {err}",
                self.request_id
            );
        }
    }
}

fn unix_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX))
        .unwrap_or(0)
}

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

    fn schedule() -> PriceSchedule {
        let metadata = OpenRouterModelMetadata {
            pricing: super::super::OpenRouterPricing {
                prompt: Some("0.000000289".into()),
                cached_prompt: Some("0.0000000289".into()),
                completion: Some("0.0000024".into()),
                ..Default::default()
            },
            ..Default::default()
        };
        PriceSchedule::from_metadata("q27", &metadata).unwrap()
    }

    #[test]
    fn decimal_math_is_exact_across_different_scales() {
        let accounting = schedule()
            .cost(Usage {
                prompt_tokens: 100,
                cached_prompt_tokens: 90,
                completion_tokens: 5,
            })
            .unwrap();
        assert_eq!(accounting["usage"]["ordinary_prompt_tokens"], 10);
        assert_eq!(accounting["cost_usd"]["ordinary_prompt"], "0.000002890");
        assert_eq!(accounting["cost_usd"]["cached_prompt"], "0.0000026010");
        assert_eq!(accounting["cost_usd"]["completion"], "0.0000120");
        assert_eq!(accounting["cost_usd"]["total"], "0.0000174910");
    }

    #[test]
    fn cached_tokens_cannot_exceed_prompt_tokens() {
        let err = schedule()
            .cost(Usage {
                prompt_tokens: 9,
                cached_prompt_tokens: 10,
                completion_tokens: 1,
            })
            .unwrap_err();
        assert!(err.contains("exceed total prompt tokens"));
    }

    #[test]
    fn completed_and_rejected_rows_are_durable_jsonl() {
        let dir = std::env::temp_dir().join(format!("memra-ledger-{}", super::unix_ms()));
        std::fs::create_dir(&dir).unwrap();
        let path = dir.join("requests.jsonl");
        let ledger = Ledger::open(
            &path,
            HashMap::from([("q27".into(), schedule())]),
        )
        .unwrap();
        let mut completed = ledger.start(
            "chatcmpl-a",
            "tenant-a",
            "q27",
            "/v1/chat/completions",
            "interactive",
            true,
        );
        completed
            .complete(
                Usage {
                    prompt_tokens: 100,
                    cached_prompt_tokens: 90,
                    completion_tokens: 5,
                },
                0.25,
            )
            .unwrap();
        assert!(completed
            .complete(
                Usage {
                    prompt_tokens: 100,
                    cached_prompt_tokens: 90,
                    completion_tokens: 5,
                },
                0.25,
            )
            .unwrap_err()
            .contains("already has a terminal ledger row"));
        let mut rejected = ledger.start(
            "chatcmpl-b",
            "tenant-a",
            "q27",
            "/v1/chat/completions",
            "interactive",
            false,
        );
        rejected.reject(429, "rate_limit_exceeded").unwrap();
        drop((completed, rejected, ledger));

        let rows: Vec<serde_json::Value> = std::fs::read_to_string(&path)
            .unwrap()
            .lines()
            .map(|line| serde_json::from_str(line).unwrap())
            .collect();
        assert_eq!(rows.len(), 2);
        assert_eq!(rows[0]["format"], FORMAT);
        assert_eq!(rows[0]["cost_usd"]["total"], "0.0000174910");
        assert_eq!(rows[1]["http_status"], 429);
        assert!(rows[1]["cost_usd"].is_null());
        std::fs::remove_dir_all(dir).unwrap();
    }

    #[test]
    fn abandoned_row_carries_partial_usage_and_cost() {
        let dir = std::env::temp_dir().join(format!(
            "memra-ledger-abandoned-{}-{}",
            std::process::id(),
            super::unix_ms(),
        ));
        std::fs::create_dir(&dir).unwrap();
        let path = dir.join("requests.jsonl");
        let ledger = Ledger::open(
            &path,
            HashMap::from([("q27".into(), schedule())]),
        )
        .unwrap();
        {
            let mut receipt = ledger.start(
                "chatcmpl-partial",
                "tenant-a",
                "q27",
                "/v1/chat/completions",
                "interactive",
                true,
            );
            receipt.record_prompt_usage(100, 90).unwrap();
            receipt.record_completion_token().unwrap();
            receipt.record_completion_token().unwrap();
        }
        drop(ledger);

        let row: serde_json::Value = serde_json::from_str(
            std::fs::read_to_string(&path).unwrap().trim(),
        )
        .unwrap();
        assert_eq!(row["outcome"], "abandoned");
        assert_eq!(row["http_status"], 499);
        assert_eq!(row["usage"]["prompt_tokens"], 100);
        assert_eq!(row["usage"]["cached_prompt_tokens"], 90);
        assert_eq!(row["usage"]["ordinary_prompt_tokens"], 10);
        assert_eq!(row["usage"]["completion_tokens"], 2);
        assert_eq!(row["usage"]["total_tokens"], 102);
        assert_eq!(row["cost_usd"]["ordinary_prompt"], "0.000002890");
        assert_eq!(row["cost_usd"]["cached_prompt"], "0.0000026010");
        assert_eq!(row["cost_usd"]["completion"], "0.0000048");
        assert_eq!(row["cost_usd"]["total"], "0.0000102910");
        std::fs::remove_dir_all(dir).unwrap();
    }

    #[cfg(unix)]
    #[test]
    fn ledger_open_accepts_0640_class_and_refuses_unsafe_modes() {
        use std::os::unix::fs::PermissionsExt as _;

        let dir = std::env::temp_dir().join(format!(
            "memra-ledger-modes-{}-{}",
            std::process::id(),
            super::unix_ms(),
        ));
        std::fs::create_dir(&dir).unwrap();
        for mode in [0o600, 0o640] {
            let path = dir.join(format!("accepted-{mode:04o}.jsonl"));
            std::fs::write(&path, b"").unwrap();
            std::fs::set_permissions(&path, std::fs::Permissions::from_mode(mode)).unwrap();
            drop(Ledger::open(
                &path,
                HashMap::from([("q27".into(), schedule())]),
            )
            .unwrap());
        }
        for mode in [0o660, 0o644, 0o610] {
            let path = dir.join(format!("refused-{mode:04o}.jsonl"));
            std::fs::write(&path, b"").unwrap();
            std::fs::set_permissions(&path, std::fs::Permissions::from_mode(mode)).unwrap();
            let err = Ledger::open(
                &path,
                HashMap::from([("q27".into(), schedule())]),
            )
            .err()
            .expect("unsafe ledger mode must be refused");
            assert!(err.contains(&format!("found {mode:04o}")), "{err}");
        }
        std::fs::remove_dir_all(dir).unwrap();
    }

    #[cfg(unix)]
    #[test]
    fn ledger_open_does_not_follow_final_symlink() {
        use std::os::unix::fs::{PermissionsExt as _, symlink};

        let dir = std::env::temp_dir().join(format!(
            "memra-ledger-symlink-{}-{}",
            std::process::id(),
            super::unix_ms(),
        ));
        std::fs::create_dir(&dir).unwrap();
        let target = dir.join("target.jsonl");
        let link = dir.join("requests.jsonl");
        std::fs::write(&target, b"").unwrap();
        std::fs::set_permissions(&target, std::fs::Permissions::from_mode(0o600)).unwrap();
        symlink(&target, &link).unwrap();
        assert!(Ledger::open(
            &link,
            HashMap::from([("q27".into(), schedule())]),
        )
        .is_err());
        std::fs::remove_dir_all(dir).unwrap();
    }

    #[test]
    fn cached_price_must_be_strictly_lower_than_prompt_price() {
        let metadata = OpenRouterModelMetadata {
            pricing: super::super::OpenRouterPricing {
                prompt: Some("0.1".into()),
                cached_prompt: Some("0.10".into()),
                completion: Some("0.2".into()),
                ..Default::default()
            },
            ..Default::default()
        };
        let err = PriceSchedule::from_metadata("m", &metadata).unwrap_err();
        assert!(err.contains("must be lower"));
    }
}