rdsh 0.6.0

//! Resident pure-compute worker for Web acceleration PoC (stage 2).
//!
//! Protocol `rdsh-compute/1` over stdio NDJSON (one JSON object per line).
//! Only pure string computation is exposed; no filesystem, no secrets,
//! no session/auth/model access. DSH stays the owner of sessions and agent
//! loop. The Node caller keeps the original path and falls back on any error.
//!
//! Request:
//!   {"v":1,"id":"<1..128 chars>","op":"tokens.estimate"|"tokens.prune",
//!    "text":"...","max_tokens":4000}
//!   {"v":1,"id":"...","op":"tokens.estimate_batch","texts":["...", ...]}
//!
//! Response (success):
//!   {"v":1,"id":"...","ok":true,"op":"...","tokens":n,"chars":n,
//!    "pruned":"...","before":n,"after":n,"budget":n,"elapsed_us":n}
//! Response (failure, fail-closed — caller must use the original path):
//!   {"v":1,"id":"...","ok":false,"error":"...","code":"..."}
//!
//! Limits: text <= 1 MiB, id 1..=128 chars, max_tokens 1..=200_000.
//! Unknown op / version / oversize / invalid UTF-8 are rejected, never
//! executed. Large histories must be chunked by the caller; one giant JSON
//! round-trip is intentionally not supported.

use std::io::{BufRead, Write};
use std::time::Instant;

pub const PROTOCOL_VERSION: u32 = 1;
pub const MAX_TEXT_BYTES: usize = 1_048_576;
// JSON can encode one text byte as six bytes (e.g. \u0000). Leave bounded
// space for batch separators, identifiers and protocol fields as well.
pub const MAX_FRAME_BYTES: usize = 6 * MAX_TEXT_BYTES + 65_536;
pub const MAX_ID_LEN: usize = 128;
pub const MAX_TOKENS: usize = 200_000;
pub const MIN_TOKENS: usize = 1;
pub const MAX_BATCH_ITEMS: usize = 64;
pub const MAX_BATCH_ITEM_BYTES: usize = 262_144;

/// Single-shot helper used by tests: parse one request line, run it.
pub fn handle_line(line: &str) -> String {
    if line.len() > MAX_FRAME_BYTES {
        return error_response("", "too_large", "request frame exceeds the byte limit");
    }
    let started = Instant::now();
    match run_request(line) {
        Ok(mut v) => {
            v["elapsed_us"] = serde_json::json!(started.elapsed().as_micros() as u64);
            v.to_string()
        }
        Err((id, code, error)) => error_response(&id, code, &error),
    }
}

fn error_response(id: &str, code: &str, error: &str) -> String {
    serde_json::json!({"v": PROTOCOL_VERSION, "id": id, "ok": false,
        "code": code, "error": error})
    .to_string()
}

fn discard_line(input: &mut impl BufRead) -> std::io::Result<()> {
    loop {
        let available = input.fill_buf()?;
        if available.is_empty() {
            return Ok(());
        }
        let end = available.iter().position(|byte| *byte == b'\n');
        let consumed = end.map_or(available.len(), |index| index + 1);
        input.consume(consumed);
        if end.is_some() {
            return Ok(());
        }
    }
}

fn run_request(line: &str) -> Result<serde_json::Value, (String, &'static str, String)> {
    let v: serde_json::Value = serde_json::from_str(line)
        .map_err(|e| (String::new(), "bad_json", format!("invalid JSON: {e}")))?;
    let id = v
        .get("id")
        .and_then(|x| x.as_str())
        .unwrap_or_default()
        .to_string();
    let fail = |code: &'static str, msg: String| -> (String, &'static str, String) {
        (id.clone(), code, msg)
    };
    if v.get("v").and_then(|x| x.as_u64()) != Some(PROTOCOL_VERSION as u64) {
        return Err(fail(
            "bad_version",
            format!("unsupported protocol, want v{PROTOCOL_VERSION}"),
        ));
    }
    if id.is_empty() || id.chars().count() > MAX_ID_LEN {
        return Err(fail("bad_id", "id must be 1..=128 chars".to_string()));
    }
    let op = v.get("op").and_then(|x| x.as_str()).unwrap_or_default();
    let text = v.get("text").and_then(|x| x.as_str()).unwrap_or_default();
    if text.len() > MAX_TEXT_BYTES {
        return Err(fail(
            "too_large",
            format!("text exceeds {MAX_TEXT_BYTES} bytes"),
        ));
    }
    match op {
        "tokens.estimate" => {
            let tokens = crate::tokens::estimate_tokens(text);
            Ok(serde_json::json!({
                "v": PROTOCOL_VERSION,
                "id": id,
                "ok": true,
                "op": op,
                "tokens": tokens,
                "chars": text.chars().count(),
            }))
        }
        "tokens.prune" => {
            let max = v.get("max_tokens").and_then(|x| x.as_u64()).unwrap_or(4000) as usize;
            if !(MIN_TOKENS..=MAX_TOKENS).contains(&max) {
                return Err(fail(
                    "bad_budget",
                    format!("max_tokens must be {MIN_TOKENS}..={MAX_TOKENS}"),
                ));
            }
            let before = crate::tokens::estimate_tokens(text);
            let pruned = crate::tokens::prune_to_budget(text, max);
            let after = crate::tokens::estimate_tokens(&pruned);
            Ok(serde_json::json!({
                "v": PROTOCOL_VERSION,
                "id": id,
                "ok": true,
                "op": op,
                "pruned": pruned,
                "before": before,
                "after": after,
                "budget": max,
            }))
        }
        "tokens.estimate_batch" => {
            let arr = v
                .get("texts")
                .and_then(|x| x.as_array())
                .cloned()
                .unwrap_or_default();
            if arr.is_empty() || arr.len() > MAX_BATCH_ITEMS {
                return Err(fail(
                    "bad_batch",
                    format!("texts must hold 1..={MAX_BATCH_ITEMS} items"),
                ));
            }
            let mut texts: Vec<&str> = Vec::with_capacity(arr.len());
            let mut total = 0usize;
            for item in &arr {
                let s = item
                    .as_str()
                    .ok_or_else(|| fail("bad_batch", "every item must be a string".to_string()))?;
                if s.len() > MAX_BATCH_ITEM_BYTES {
                    return Err(fail(
                        "too_large",
                        format!("one item exceeds {MAX_BATCH_ITEM_BYTES} bytes"),
                    ));
                }
                total += s.len();
                if total > MAX_TEXT_BYTES {
                    return Err(fail(
                        "too_large",
                        format!("batch total exceeds {MAX_TEXT_BYTES} bytes"),
                    ));
                }
                texts.push(s);
            }
            // One pass over all items: a single IPC round-trip replaces N.
            let tokens: Vec<usize> = texts
                .iter()
                .map(|s| crate::tokens::estimate_tokens(s))
                .collect();
            let chars: Vec<usize> = texts.iter().map(|s| s.chars().count()).collect();
            let total_tokens: usize = tokens.iter().sum();
            Ok(serde_json::json!({
                "v": PROTOCOL_VERSION,
                "id": id,
                "ok": true,
                "op": op,
                "tokens": tokens,
                "chars": chars,
                "total": total_tokens,
            }))
        }
        _ => Err(fail("bad_op", format!("unknown op: {op}"))),
    }
}

/// Resident loop: read NDJSON from stdin, write NDJSON to stdout.
/// EOF ends the loop (exit 0). A malformed line yields an error object,
/// never a crash. `--once` handles a single line (for IPC cost comparison).
pub fn cmd_compute(once: bool, _locale: crate::locale::Locale) -> anyhow::Result<()> {
    let stdin = std::io::stdin();
    let mut out = std::io::BufWriter::with_capacity(256 * 1024, std::io::stdout().lock());
    run_lines(stdin.lock(), &mut out, once)
}

// Streaming NDJSON loop: one request line in, one response line out.
// Invalid UTF-8 yields a bad_utf8 error object, never a worker crash.
fn run_lines<R: BufRead, W: Write>(mut input: R, output: &mut W, once: bool) -> anyhow::Result<()> {
    let mut buf: Vec<u8> = Vec::with_capacity(64 * 1024);
    loop {
        buf.clear();
        let n = std::io::Read::take(&mut input, (MAX_FRAME_BYTES + 1) as u64)
            .read_until(b'\n', &mut buf)?;
        if n == 0 {
            break; // EOF
        }
        if n > MAX_FRAME_BYTES {
            if buf.last() != Some(&b'\n') {
                discard_line(&mut input)?;
            }
            writeln!(
                output,
                "{}",
                error_response("", "too_large", "request frame exceeds the byte limit")
            )?;
            output.flush()?;
            if once {
                break;
            }
            continue;
        }
        // Strip one trailing newline / CR for tolerance on Windows pipes.
        if buf.last() == Some(&b'\n') {
            buf.pop();
        }
        if buf.last() == Some(&b'\r') {
            buf.pop();
        }
        if buf.iter().all(|b| b.is_ascii_whitespace()) {
            if once {
                break;
            }
            continue;
        }
        let resp = match std::str::from_utf8(&buf) {
            Ok(s) => handle_line(s),
            Err(_) => error_response("", "bad_utf8", "request is not valid UTF-8"),
        };
        writeln!(output, "{resp}")?;
        output.flush()?;
        if once {
            break;
        }
    }
    output.flush()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    #[test]
    fn maximally_escaped_text_still_fits_the_frame_limit() {
        let text = "\0".repeat(super::MAX_TEXT_BYTES);
        let line =
            serde_json::json!({"v": 1, "id": "escaped", "op": "tokens.estimate", "text": text})
                .to_string();
        let response: serde_json::Value = serde_json::from_str(&super::handle_line(&line)).unwrap();
        assert_eq!(response["ok"], true);
        assert_eq!(response["chars"], super::MAX_TEXT_BYTES);
    }
    #[test]
    fn oversized_frames_are_rejected_and_the_next_request_survives() {
        let large = "x".repeat(7 * super::MAX_TEXT_BYTES);
        let first = serde_json::json!({"v": 1, "id": "oversized", "op": "tokens.estimate",
            "text": "small", "padding": large})
        .to_string();
        let input = format!(
            "{first}\n{{\"v\":1,\"id\":\"after\",\"op\":\"tokens.estimate\",\"text\":\"ok\"}}\n"
        );
        let mut output = Vec::new();
        super::run_lines(std::io::Cursor::new(input), &mut output, false).unwrap();
        let responses: Vec<serde_json::Value> = std::str::from_utf8(&output)
            .unwrap()
            .lines()
            .map(|line| serde_json::from_str(line).unwrap())
            .collect();
        assert_eq!(responses.len(), 2);
        assert_eq!(responses[0]["ok"], false);
        assert_eq!(responses[0]["code"], "too_large");
        assert_eq!(responses[1]["id"], "after");
        assert_eq!(responses[1]["ok"], true);
    }
    use super::*;

    fn req(id: &str, op: &str, text: &str, max: Option<usize>) -> String {
        let mut v = serde_json::json!({"v": 1, "id": id, "op": op, "text": text});
        if let Some(m) = max {
            v["max_tokens"] = serde_json::json!(m);
        }
        v.to_string()
    }

    #[test]
    fn estimate_matches_direct() {
        let text = "hello world, token test";
        let resp: serde_json::Value =
            serde_json::from_str(&handle_line(&req("a1", "tokens.estimate", text, None))).unwrap();
        assert_eq!(resp["ok"], true);
        assert_eq!(resp["tokens"], crate::tokens::estimate_tokens(text));
        assert_eq!(resp["chars"], text.chars().count());
        assert!(resp["elapsed_us"].as_u64().is_some());
    }

    // #389: chars は文字数（Unicode scalar 数）であってバイト数ではない。
    #[test]
    fn estimate_reports_char_count_not_byte_count() {
        for (text, want) in [("\u{65e5}\u{672c}\u{8a9e}", 3), ("\u{e9}", 1), ("abc", 3)] {
            let resp: serde_json::Value =
                serde_json::from_str(&handle_line(&req("c1", "tokens.estimate", text, None)))
                    .unwrap();
            assert_eq!(resp["chars"], want, "text {text:?}");
        }
        let line = serde_json::json!({"v": 1, "id": "c2", "op": "tokens.estimate_batch",
            "texts": ["\u{65e5}\u{672c}\u{8a9e}", "\u{e9}", "abc"]})
        .to_string();
        let resp: serde_json::Value = serde_json::from_str(&handle_line(&line)).unwrap();
        assert_eq!(resp["chars"], serde_json::json!([3, 1, 3]));
    }

    #[test]
    fn prune_matches_direct_and_respects_budget() {
        let text = "x".repeat(20000);
        let resp: serde_json::Value =
            serde_json::from_str(&handle_line(&req("p1", "tokens.prune", &text, Some(100))))
                .unwrap();
        assert_eq!(resp["ok"], true);
        let pruned = resp["pruned"].as_str().unwrap();
        assert_eq!(pruned, crate::tokens::prune_to_budget(&text, 100));
        assert!(resp["after"].as_u64().unwrap() <= 100);
    }

    #[test]
    fn batch_matches_sum_of_direct() {
        let texts = vec!["abcd1234", "hello", "あいう"];
        let line =
            serde_json::json!({"v": 1, "id": "m1", "op": "tokens.estimate_batch", "texts": texts})
                .to_string();
        let resp: serde_json::Value = serde_json::from_str(&handle_line(&line)).unwrap();
        assert_eq!(resp["ok"], true);
        let want: Vec<usize> = texts
            .iter()
            .map(|s| crate::tokens::estimate_tokens(s))
            .collect();
        let got: Vec<u64> = serde_json::from_value(resp["tokens"].clone()).unwrap();
        assert_eq!(got, want.iter().map(|&n| n as u64).collect::<Vec<_>>());
        assert_eq!(resp["total"], want.iter().sum::<usize>() as u64);

        let empty =
            serde_json::json!({"v": 1, "id": "m2", "op": "tokens.estimate_batch", "texts": []})
                .to_string();
        let r: serde_json::Value = serde_json::from_str(&handle_line(&empty)).unwrap();
        assert_eq!(r["code"], "bad_batch");
    }

    #[test]
    fn rejects_unknown_op_version_and_size() {
        let bad_op: serde_json::Value =
            serde_json::from_str(&handle_line(&req("x", "nope", "hi", None))).unwrap();
        assert_eq!(bad_op["ok"], false);
        assert_eq!(bad_op["code"], "bad_op");

        let bad_v = serde_json::json!({"v": 999, "id": "x", "op": "tokens.estimate", "text": "hi"})
            .to_string();
        let r: serde_json::Value = serde_json::from_str(&handle_line(&bad_v)).unwrap();
        assert_eq!(r["code"], "bad_version");

        let big = "y".repeat(MAX_TEXT_BYTES + 1);
        let r: serde_json::Value =
            serde_json::from_str(&handle_line(&req("big", "tokens.estimate", &big, None))).unwrap();
        assert_eq!(r["code"], "too_large");

        let bad_budget: serde_json::Value = serde_json::from_str(&handle_line(&req(
            "b",
            "tokens.prune",
            "hello",
            Some(MAX_TOKENS + 1),
        )))
        .unwrap();
        assert_eq!(bad_budget["code"], "bad_budget");
    }

    #[test]
    fn rejects_bad_json_and_id() {
        let r: serde_json::Value = serde_json::from_str(&handle_line("{oops")).unwrap();
        assert_eq!(r["ok"], false);
        let r: serde_json::Value =
            serde_json::from_str(&handle_line(&req("", "tokens.estimate", "hi", None))).unwrap();
        assert_eq!(r["code"], "bad_id");
    }

    // run_lines は NDJSON のフェイルセーフ境界: 壊れた行でも落ちずに応答を返し、
    // 空行は飛ばし、once では最初の1件で止まる。
    #[test]
    fn run_lines_answers_each_request_in_order() {
        let input = format!(
            "{}\r\n\r\n{}\n",
            req("a1", "tokens.estimate", "one", None),
            req("a2", "tokens.estimate", "two", None)
        );
        let mut out: Vec<u8> = Vec::new();
        run_lines(input.as_bytes(), &mut out, false).unwrap();
        let text = String::from_utf8(out).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert_eq!(lines.len(), 2, "a blank line must not produce a response");
        let first: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
        let second: serde_json::Value = serde_json::from_str(lines[1]).unwrap();
        assert_eq!(first["id"], "a1");
        assert_eq!(second["id"], "a2");
    }

    #[test]
    fn run_lines_reports_bad_utf8_and_keeps_going() {
        let mut input: Vec<u8> = vec![0xff, 0xfe, b'\n'];
        input.extend_from_slice(req("a2", "tokens.estimate", "ok", None).as_bytes());
        input.push(b'\n');
        let mut out: Vec<u8> = Vec::new();
        run_lines(&input[..], &mut out, false).unwrap();
        let text = String::from_utf8(out).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert_eq!(lines.len(), 2);
        let bad: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
        assert_eq!(bad["ok"], false);
        assert_eq!(bad["code"], "bad_utf8");
        let good: serde_json::Value = serde_json::from_str(lines[1]).unwrap();
        assert_eq!(good["id"], "a2");
    }

    #[test]
    fn run_lines_once_stops_after_the_first_request() {
        let input = format!(
            "{}\n{}\n",
            req("a1", "tokens.estimate", "one", None),
            req("a2", "tokens.estimate", "two", None)
        );
        let mut out: Vec<u8> = Vec::new();
        run_lines(input.as_bytes(), &mut out, true).unwrap();
        let text = String::from_utf8(out).unwrap();
        assert_eq!(text.lines().count(), 1);
    }

    #[test]
    fn run_lines_answers_without_a_trailing_newline() {
        let input = req("a1", "tokens.estimate", "tail", None);
        let mut out: Vec<u8> = Vec::new();
        run_lines(input.as_bytes(), &mut out, false).unwrap();
        let text = String::from_utf8(out).unwrap();
        let lines: Vec<&str> = text.lines().collect();
        assert_eq!(lines.len(), 1);
        let resp: serde_json::Value = serde_json::from_str(lines[0]).unwrap();
        assert_eq!(resp["id"], "a1");
    }

    // id の上限は仕様・エラー文どおり「文字数」。バイト数で判定すると
    // 非ASCII の id が 128 文字でも拒否される。
    #[test]
    fn accepts_a_128_char_non_ascii_id() {
        let id = "\u{3042}".repeat(MAX_ID_LEN);
        let resp: serde_json::Value =
            serde_json::from_str(&handle_line(&req(&id, "tokens.estimate", "hi", None))).unwrap();
        assert_eq!(resp["ok"], true);
        assert_eq!(resp["id"], id);
    }

    #[test]
    fn rejects_an_id_over_the_char_limit() {
        let id = "\u{3042}".repeat(MAX_ID_LEN + 1);
        let resp: serde_json::Value =
            serde_json::from_str(&handle_line(&req(&id, "tokens.estimate", "hi", None))).unwrap();
        assert_eq!(resp["ok"], false);
        assert_eq!(resp["code"], "bad_id");
    }
}