tests/test_engine.py

"""SDK contracts, deterministic streaming edge cases, and real native integration."""

from concurrent.futures import ThreadPoolExecutor
import threading
import time

import pytest

from cism import Engine, GenerationChunk, GenerationResult


class FakeTokenizer:
    bos_token_id = 0
    eos_token_id = None
    chat_template = "test template"

    def __init__(self, pieces=None, encoded=None, special=()):
        self.pieces = pieces or {}
        self.encoded = [1, 2] if encoded is None else encoded
        self.special = set(special)
        self.encode_calls = []
        self.chat_calls = []

    def encode(self, text, *, add_special_tokens):
        self.encode_calls.append((text, add_special_tokens))
        return self.encoded.copy()

    def decode(self, tokens, *, skip_special_tokens, clean_up_tokenization_spaces):
        assert skip_special_tokens is True
        assert clean_up_tokenization_spaces is False
        return b"".join(
            self.pieces.get(token, f"<{token}>".encode())
            for token in tokens if token not in self.special
        ).decode("utf-8", errors="replace")

    def apply_chat_template(self, messages, *, tokenize, add_generation_prompt):
        self.chat_calls.append((messages, tokenize, add_generation_prompt))
        return self.encoded.copy()


class FakeSession:
    def __init__(self, output, options):
        self.output = list(output)
        self.options = options
        self.generated_tokens = 0
        self.finish_reason = "length" if options["max_new_tokens"] == 0 else ""
        self.cancelled = threading.Event()
        self.requests = []

    def next_tokens(self, count, spec_k=0):
        self.requests.append((count, spec_k))
        batch = []
        for _ in range(count):
            if self.finish_reason:
                break
            assert self.generated_tokens < len(self.output), "Fake script exhausted"
            token = self.output[self.generated_tokens]
            self.generated_tokens += 1
            batch.append(token)
            if token in self.options["eos_token_ids"]:
                self.finish_reason = "stop"
            elif self.generated_tokens == self.options["max_new_tokens"]:
                self.finish_reason = "length"
        return batch

    def cancel(self):
        self.cancelled.set()
        if not self.finish_reason:
            self.finish_reason = "cancelled"


class FakeNative:
    info = {"precision": "fp32", "weight_bytes": 123, "threads": 1}

    def __init__(self, output):
        self.output = output
        self.sessions = []
        self.prompts = []
        self.logits_calls = []
        self.touches = 0
        self.scans = 0
        self.pages_locked = False

    def create_session(self, prompt, **options):
        self.prompts.append(prompt.copy())
        session = FakeSession(self.output, options)
        self.sessions.append(session)
        return session

    def logits(self, prompt):
        self.logits_calls.append(prompt)
        return [0.25, 0.75]

    def touch(self):
        self.touches += 1
        return 123

    def scan(self):
        self.scans += 1
        return 123

    def lock_pages(self):
        self.pages_locked = True
        return 123

    def unlock_pages(self):
        self.pages_locked = False
        return 123


def make_engine(output=(3, 4, 5, 6), *, pieces=None, tokenizer=None, **config):
    native = FakeNative(output)
    tokenizer = tokenizer or FakeTokenizer(pieces)
    engine = Engine(native, tokenizer, {
        "model_type": "llama", "vocab_size": 64, "max_position_embeddings": 256,
        **config,
    }, "fake-model", revision="pinned")
    return engine, native


def assert_released(engine, native):
    assert engine._generation_lock.acquire(blocking=False), "Generation lock leaked"
    engine._generation_lock.release()
    assert all(session.cancelled.is_set() for session in native.sessions)


def test_encode_chat_and_metadata():
    engine, native = make_engine()
    prompt = [4, 5]
    assert engine.encode(prompt) == prompt
    assert engine.encode(prompt) is not prompt
    assert engine.encode("hello") == [1, 2]
    assert engine.tokenizer.encode_calls == [("hello", True)]
    messages = [{"role": "user", "content": "hello"}]
    assert engine.format_chat(messages) == [1, 2]
    assert engine.tokenizer.chat_calls == [(messages, True, True)]
    assert engine.logits(prompt) == [0.25, 0.75]
    assert native.logits_calls == [prompt]
    assert engine.info == {
        **native.info, "model": "fake-model", "revision": "pinned",
        "context_length": 256, "architecture": "llama",
        "capabilities": ["text", "streaming", "sampling"],
        "keep_warm": False,
    }
    engine.info["precision"] = "changed"
    assert native.info["precision"] == "fp32"
    engine.tokenizer.chat_template = None
    with pytest.raises(ValueError, match="chat template"):
        engine.format_chat(messages)


@pytest.mark.parametrize("prompt", [[], ""])
def test_empty_prompt_bos_and_usage(prompt):
    engine, native = make_engine(tokenizer=FakeTokenizer(encoded=[]), bos_token_id=7)
    result = engine.generate(prompt, max_new_tokens=1)
    assert native.prompts == [[7]]
    assert result.prompt_tokens == 1
    del engine.config["bos_token_id"]
    assert engine.encode(prompt) == [0]
    engine.tokenizer.bos_token_id = None
    with pytest.raises(ValueError, match="BOS"):
        engine.encode(prompt)


@pytest.mark.parametrize("prompt", [[-1], [64], [True], [1.0], ["1"], [None]])
def test_invalid_prompt_tokens(prompt):
    engine, native = make_engine()
    with pytest.raises(ValueError, match="Token IDs"):
        engine.generate(prompt, max_new_tokens=1)
    assert not native.sessions
    assert_released(engine, native)


@pytest.mark.parametrize("prompt", [(1, 2), None, 3, b"hello"])
def test_invalid_prompt_type(prompt):
    engine, _ = make_engine()
    with pytest.raises(TypeError, match="prompt"):
        engine.encode(prompt)


@pytest.mark.parametrize("eos", [-1, 64, True, [1, 64], [False]])
def test_invalid_eos_ids(eos):
    with pytest.raises(ValueError, match="Token IDs"):
        make_engine(eos_token_id=eos)


@pytest.mark.parametrize("options", [
    {"chunk_size": 0}, {"chunk_size": 1025}, {"chunk_size": True}, {"chunk_size": 1.5},
    {"max_new_tokens": -1}, {"max_new_tokens": True}, {"max_new_tokens": 1.5},
    {"seed": -1}, {"seed": 2**64}, {"seed": True}, {"seed": 1.5},
    {"top_k": -1}, {"top_k": True}, {"top_k": 1.5},
    {"temperature": -1}, {"temperature": float("nan")},
    {"temperature": float("inf")}, {"temperature": True}, {"temperature": "1"},
    {"top_p": 0}, {"top_p": 1.01}, {"top_p": float("nan")},
    {"top_p": float("inf")}, {"top_p": False}, {"top_p": None},
    {"stop": "end"}, {"stop": [""]}, {"stop": [1]}, {"stop": ("end",)},
])
def test_invalid_options_before_native_call(options):
    engine, native = make_engine()
    with pytest.raises(ValueError):
        list(engine.stream([1], **{"max_new_tokens": 1, **options}))
    assert not native.sessions
    assert_released(engine, native)


@pytest.mark.parametrize("chunk_size", [1, 2, 16, 1024])
def test_stream_generate_consistency_and_sampling_forwarding(chunk_size):
    engine, native = make_engine(pieces={3: b"a", 4: b" b", 5: b"c", 6: b"!"})
    options = dict(max_new_tokens=4, temperature=0.7, top_p=0.8, top_k=9,
                   seed=2**64 - 1, chunk_size=chunk_size)
    chunks = list(engine.stream("prompt", **options))
    assert all(isinstance(chunk, GenerationChunk) for chunk in chunks)
    assert chunks[0].token_ids == [3]
    assert chunks[-1].finish_reason == "length"
    assert all(chunk.finish_reason is None for chunk in chunks[:-1])
    total = 0
    for chunk in chunks:
        total += len(chunk.token_ids)
        assert chunk.completion_tokens == total
        assert chunk.prompt_tokens == 2
    result = engine.generate("prompt", **options)
    assert result == GenerationResult("a bc!", [3, 4, 5, 6], "length", 2, 4)
    assert result.text == "".join(chunk.text for chunk in chunks)
    assert result.token_ids == [token for chunk in chunks for token in chunk.token_ids]
    assert native.sessions[0].options == {
        "max_new_tokens": 4, "temperature": 0.7, "top_p": 0.8,
        "top_k": 9, "seed": 2**64 - 1, "eos_token_ids": [],
    }
    assert native.sessions[0].requests == [(1, 0)] + [(chunk_size, 0)] * (len(chunks) - 1)
    assert_released(engine, native)


@pytest.mark.parametrize("pieces,stops,text,count", [
    ([b"hello ", b"<", b"EN", b"D>tail", b"never"], ["<END>"], "hello ", 4),
    ([b"hello END ignored", b"never"], ["END"], "hello ", 1),
    ([b"x", b"ab", b"ab!", b"never"], ["abab"], "x", 3),
    ([b"x", b"ab", b"c", b"!"], ["abcd"], "xabc!", 4),
    ([b"x", b"ab"], ["abcd"], "xab", 2),
    ([b"x", b"later END trailing"], ["END", "later"], "x", 2),
    ([b"prefix END"], ["END", "prefix"], "", 1),
    ([b"x", b"ENDtail"], ["END", "ENDtail"], "x", 2),
])
@pytest.mark.parametrize("chunk_size", [1, 16])
def test_stop_sequences(pieces, stops, text, count, chunk_size):
    ids = list(range(3, 3 + len(pieces)))
    engine, native = make_engine(ids, pieces=dict(zip(ids, pieces)))
    chunks = list(engine.stream([1, 2], max_new_tokens=len(ids), stop=stops,
                                chunk_size=chunk_size))
    expected_reason = "stop" if any(stop in b"".join(pieces).decode() for stop in stops) else "length"
    assert "".join(chunk.text for chunk in chunks) == text
    assert all(text.startswith("".join(chunk.text for chunk in chunks[:i + 1]))
               for i in range(len(chunks)))
    assert chunks[-1].finish_reason == expected_reason
    assert chunks[-1].completion_tokens == count
    assert [token for chunk in chunks for token in chunk.token_ids] == ids[:count]
    assert native.sessions[0].generated_tokens == count
    assert native.sessions[0].requests == [(1, 0)] * count
    assert engine.generate([1, 2], max_new_tokens=len(ids), stop=stops,
                           chunk_size=chunk_size) == GenerationResult(
        text, ids[:count], expected_reason, 2, count)
    assert_released(engine, native)


@pytest.mark.parametrize("chunk_size", [1, 2, 16])
def test_partial_utf8_is_held_until_complete(chunk_size):
    engine, native = make_engine(pieces={3: b"x\xe2", 4: b"\x82", 5: b"\xac", 6: b"!"})
    chunks = list(engine.stream([1], max_new_tokens=4, chunk_size=chunk_size))
    assert chunks[0].text == "x"
    assert all("\ufffd" not in chunk.text for chunk in chunks)
    assert "".join(chunk.text for chunk in chunks) == "x\u20ac!"
    assert_released(engine, native)


def test_partial_utf8_at_length_is_flushed():
    engine, _ = make_engine((3,), pieces={3: b"x\xe2"})
    assert engine.generate([1], max_new_tokens=1).text == "x\ufffd"


@pytest.mark.parametrize("pieces", [b"x\xe2", b"x\xef\xbf\xbd"])
def test_replacement_stop_is_honored_at_terminal_decode(pieces):
    engine, _ = make_engine((3,), pieces={3: pieces})
    assert engine.generate([1], max_new_tokens=1, stop=["\ufffd"]) == GenerationResult(
        "x", [3], "stop", 1, 1)


def test_unicode_stop_spans_byte_tokens():
    engine, native = make_engine(pieces={3: b"x\xe2", 4: b"\x82", 5: b"\xac", 6: b"!"})
    result = engine.generate([1], max_new_tokens=4, stop=["\u20ac"])
    assert result == GenerationResult("x", [3, 4, 5], "stop", 1, 3)
    assert native.sessions[0].generated_tokens == 3


@pytest.mark.parametrize("stop", ["\ufffd", "x\ufffd"])
def test_stop_does_not_match_temporary_utf8_replacement(stop):
    engine, _ = make_engine(pieces={3: b"x\xe2", 4: b"\x82", 5: b"\xac", 6: b"!"})
    result = engine.generate([1], max_new_tokens=4, stop=[stop])
    assert result == GenerationResult("x\u20ac!", [3, 4, 5, 6], "length", 1, 4)


@pytest.mark.parametrize("eos", [4, [4, 5]])
def test_eos_counted_but_not_rendered(eos):
    tokenizer = FakeTokenizer({3: b"a", 4: b"<eos>"}, special=[4, 5])
    engine, native = make_engine(tokenizer=tokenizer, eos_token_id=eos)
    result = engine.generate([4, 1], max_new_tokens=4)
    assert result == GenerationResult("a", [3, 4], "stop", 2, 2)
    assert native.sessions[0].options["eos_token_ids"] == ([eos] if isinstance(eos, int) else eos)
    assert_released(engine, native)


def test_eos_tokenizer_fallback_and_eos_at_budget():
    tokenizer = FakeTokenizer(special=[3])
    tokenizer.eos_token_id = 3
    engine, _ = make_engine(tokenizer=tokenizer)
    assert engine.generate([1], max_new_tokens=1) == GenerationResult("", [3], "stop", 1, 1)


def test_over_budget_is_clamped_and_long_prompt_rejected():
    engine, native = make_engine(max_position_embeddings=2)
    assert list(engine.stream([1, 2], max_new_tokens=0)) == [GenerationChunk("", [], "length", 2, 0)]
    assert engine.generate([1, 2], max_new_tokens=0) == GenerationResult("", [], "length", 2, 0)
    assert engine.generate([1], max_new_tokens=1).completion_tokens == 1
    # Over-budget requests clamp to the remaining window instead of failing.
    assert engine.generate([1], max_new_tokens=9) == GenerationResult("<3>", [3], "length", 1, 1)
    assert engine.generate([1, 2], max_new_tokens=5) == GenerationResult("", [], "length", 2, 0)
    with pytest.raises(ValueError, match="Prompt length exceeds"):
        engine.generate([1, 2, 3], max_new_tokens=0)
    assert len(native.sessions) == 5
    assert_released(engine, native)


def test_terminal_chunk_releases_engine_before_iterator_resumed():
    engine, native = make_engine((3,))
    stream = engine.stream([1], max_new_tokens=1)
    chunk = next(stream)
    assert chunk.finish_reason == "length"
    assert_released(engine, native)
    assert engine.generate([1], max_new_tokens=1).finish_reason == "length"
    stream.close()


def test_precancelled_stream_and_generate():
    engine, native = make_engine()
    cancel = threading.Event()
    cancel.set()
    assert list(engine.stream([1], cancel=cancel)) == []
    with pytest.raises(RuntimeError, match="cancelled"):
        engine.generate([1], cancel=cancel)
    assert not native.sessions
    assert_released(engine, native)


@pytest.mark.parametrize("action", ["close", "cancel", "throw"])
def test_suspended_generator_cleanup(action):
    engine, native = make_engine()
    cancel = threading.Event()
    before = set(threading.enumerate())
    stream = engine.stream([1], max_new_tokens=4, cancel=cancel)
    assert next(stream).token_ids == [3]
    monitors = {thread for thread in threading.enumerate() if thread not in before and thread.name == "cism-cancel"}
    assert monitors
    assert engine._generation_lock.locked()
    if action == "cancel":
        cancel.set()
        assert native.sessions[0].cancelled.wait(2)
        assert list(stream) == []
    elif action == "throw":
        with pytest.raises(LookupError, match="consumer failed"):
            stream.throw(LookupError("consumer failed"))
    else:
        stream.close()
    assert all(not thread.is_alive() for thread in monitors)
    assert_released(engine, native)
    assert engine.generate([1], max_new_tokens=1).token_ids == [3]


def test_close_before_start_does_not_allocate():
    engine, native = make_engine()
    stream = engine.stream([1])
    stream.close()
    assert not native.sessions
    assert_released(engine, native)


@pytest.mark.parametrize("operation", ["stream", "generate"])
def test_cancel_interrupts_blocked_native_decode(monkeypatch, operation):
    engine, native = make_engine()
    entered = threading.Event()

    def blocked_decode(session, count, spec_k=0):
        entered.set()
        assert session.cancelled.wait(2), "Cancellation monitor failed to interrupt native work"
        return []

    monkeypatch.setattr(FakeSession, "next_tokens", blocked_decode)
    cancel = threading.Event()

    def request():
        if operation == "generate":
            return engine.generate([1], cancel=cancel)
        return list(engine.stream([1], cancel=cancel))

    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(request)
        try:
            assert entered.wait(2)
            cancel.set()
            if operation == "generate":
                with pytest.raises(RuntimeError, match="cancelled"):
                    future.result(timeout=3)
            else:
                assert future.result(timeout=3) == []
        finally:
            cancel.set()
    assert_released(engine, native)


def test_generate_cancellation_after_partial_output_raises(monkeypatch):
    engine, native = make_engine()
    cancel = threading.Event()
    original_decode = engine.tokenizer.decode

    def decode(tokens, **kwargs):
        cancel.set()
        return original_decode(tokens, **kwargs)

    monkeypatch.setattr(engine.tokenizer, "decode", decode)
    with pytest.raises(RuntimeError, match="cancelled"):
        engine.generate([1], max_new_tokens=4, cancel=cancel)
    assert native.sessions[0].generated_tokens == 1
    assert_released(engine, native)


def test_cancel_while_waiting_for_generation_lock():
    engine, native = make_engine()
    first = engine.stream([1], max_new_tokens=4)
    next(first)
    cancel = threading.Event()
    started = threading.Event()

    def second_request():
        started.set()
        return list(engine.stream([2], max_new_tokens=1, cancel=cancel))

    try:
        with ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(second_request)
            try:
                assert started.wait(2)
                assert not future.done()
                cancel.set()
                assert future.result(timeout=2) == []
                assert len(native.sessions) == 1
                assert engine._generation_lock.locked()
            finally:
                cancel.set()
                first.close()
    finally:
        first.close()
    assert_released(engine, native)


@pytest.mark.parametrize("operation", ["generate", "logits"])
def test_generation_and_logits_share_lock(operation):
    engine, native = make_engine()
    first = engine.stream([1], max_new_tokens=4)
    next(first)
    entered = threading.Event()

    def second_request():
        entered.set()
        if operation == "generate":
            return engine.generate([2], max_new_tokens=1)
        return engine.logits([2])

    try:
        with ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(second_request)
            try:
                assert entered.wait(2)
                with pytest.raises(TimeoutError):
                    future.result(timeout=0.1)
                assert len(native.sessions) == 1
                assert not native.logits_calls
            finally:
                first.close()
            assert future.result(timeout=2) is not None
    finally:
        first.close()
    assert_released(engine, native)


@pytest.mark.parametrize("failure", ["create", "decode", "tokenizer", "no_progress", "rewrite"])
def test_failures_release_session_and_lock(monkeypatch, failure):
    engine, native = make_engine()

    def fail(*args, **kwargs):
        raise RuntimeError("injected failure")

    error, message = RuntimeError, "injected failure"
    if failure == "create":
        monkeypatch.setattr(native, "create_session", fail)
    elif failure == "decode":
        monkeypatch.setattr(FakeSession, "next_tokens", fail)
    elif failure == "tokenizer":
        monkeypatch.setattr(engine.tokenizer, "decode", fail)
    elif failure == "no_progress":
        monkeypatch.setattr(FakeSession, "next_tokens", lambda *args, **kwargs: [])
        message = "no progress"
    else:
        monkeypatch.setattr(engine.tokenizer, "decode", lambda tokens, **kwargs: "a" if len(tokens) == 1 else "b")
        error, message = ValueError, "rewrote emitted text"
    with pytest.raises(error, match=message):
        list(engine.stream([1], max_new_tokens=4))
    assert_released(engine, native)


def test_monitor_start_failure_releases_lock(monkeypatch):
    engine, native = make_engine()

    def fail_start(self):
        raise RuntimeError("cannot start monitor")

    monkeypatch.setattr(threading.Thread, "start", fail_start)
    with pytest.raises(RuntimeError):
        list(engine.stream([1], max_new_tokens=1))
    assert_released(engine, native)


def test_monitor_start_failure_preserves_original_error(monkeypatch):
    engine, _ = make_engine()

    def fail_start(self):
        raise RuntimeError("cannot start monitor")

    monkeypatch.setattr(threading.Thread, "start", fail_start)
    with pytest.raises(RuntimeError, match="^cannot start monitor$"):
        list(engine.stream([1], max_new_tokens=1))


def test_from_pretrained_forwards_import_options_and_resolved_revision(monkeypatch):
    from cism import loader

    calls = []
    loaded = loader.LoadedModel({}, {}, FakeTokenizer(), "resolved-local-path", "commit-sha")

    def import_model(model, **options):
        calls.append((model, options))
        return loaded

    def from_weights(cls, config, weights, tokenizer, **options):
        assert cls is Engine
        assert config is loaded.config
        assert weights is loaded.weights
        assert tokenizer is loaded.tokenizer
        assert options == {"model_id": "org/model", "precision": "int8", "revision": "commit-sha",
                           "threads": 1, "act_precision": "fp32"}
        return "constructed-engine"

    monkeypatch.setattr(loader, "import_model", import_model)
    monkeypatch.setattr(Engine, "from_weights", classmethod(from_weights))
    assert Engine.from_pretrained("org/model", precision="int8", revision="main",
                                  trust_remote_code=True, local_files_only=True,
                                  cache_dir="cache", threads=1) == "constructed-engine"
    assert calls == [("org/model", {
        "revision": "main", "trust_remote_code": True,
        "local_files_only": True, "cache_dir": "cache",
    })]


@pytest.fixture(params=["llama", "qwen3"])
def real_engine(request):
    # Lazy import keeps fake-only SDK tests independent of the native extension.
    from test_native import reference_logits, tiny_model

    config, weights = tiny_model(request.param)
    engine = Engine.from_weights(config, weights, FakeTokenizer(),
                                 model_id="tiny-" + request.param, revision="local-test")
    return engine, config, weights, reference_logits


def test_from_weights_logits_and_greedy_oracle(real_engine):
    import numpy as np

    engine, config, weights, reference_logits = real_engine
    prompt = [2, 7, 1, 19, 4]
    np.testing.assert_allclose(engine.logits(prompt), reference_logits(config, weights, prompt),
                               atol=8e-6, rtol=8e-5)
    expected = []
    for _ in range(9):
        expected.append(int(np.argmax(reference_logits(config, weights, prompt + expected))))
    result = engine.generate(prompt, max_new_tokens=9)
    assert result.token_ids == expected
    assert result.finish_reason == "length"
    assert result.prompt_tokens == len(prompt)
    assert result.completion_tokens == 9
    assert result.text == "".join(f"<{token}>" for token in expected)
    assert engine.info["architecture"] == config["model_type"]
    assert engine.info["revision"] == "local-test"
    assert engine.info["precision"] == "fp32"


def test_nll_windowing_and_validation(real_engine):
    import math

    import numpy as np

    engine, config, weights, reference_logits = real_engine
    tokens = [2, 7, 1, 19, 4, 11, 3, 5]
    scores = engine.nll(tokens)
    assert set(scores) == {"nll", "perplexity", "scored_tokens"}
    assert scores["scored_tokens"] == len(tokens) - 1
    assert abs(scores["perplexity"] - math.exp(scores["nll"])) < 1e-9 * max(1.0, scores["perplexity"])
    # Windowed scoring equals single-window scoring when everything fits.
    whole = engine.nll(tokens, window=64)
    assert whole["nll"] == scores["nll"]
    # Deterministic across repeated calls and identical via the text path.
    assert engine.nll(tokens)["nll"] == scores["nll"]
    # Split windows are scored with truncated context, so they can differ;
    # both must stay deterministic and aggregate over all scored tokens.
    split = engine.nll(tokens, window=4)
    assert split["scored_tokens"] == len(tokens) - 1
    assert split["nll"] > 0
    assert engine.nll(tokens[:2])["scored_tokens"] == 1
    for bad in (0, 1, -3, "x"):
        with pytest.raises(ValueError, match="window"):
            engine.nll(tokens, window=bad)
    with pytest.raises(ValueError, match="at least 2 tokens"):
        engine.nll([5])


@pytest.mark.parametrize("chunk_size", [1, 2, 16])
def test_from_weights_sampled_stream_consistency(real_engine, chunk_size):
    engine, _, _, _ = real_engine
    options = dict(max_new_tokens=13, temperature=0.9, top_p=0.8, top_k=11, seed=876)
    result = engine.generate([3, 1], **options, chunk_size=1)
    chunks = list(engine.stream([3, 1], **options, chunk_size=chunk_size))
    assert [token for chunk in chunks for token in chunk.token_ids] == result.token_ids
    assert "".join(chunk.text for chunk in chunks) == result.text
    assert chunks[-1].completion_tokens == result.completion_tokens == 13
    assert chunks[-1].finish_reason == result.finish_reason == "length"
    assert chunks[-1].prompt_tokens == result.prompt_tokens == 2


@pytest.mark.parametrize("precision", ["int8", "hybrid-int4"])
def test_from_weights_forwards_precision(real_engine, precision):
    import numpy as np
    from test_native import reconstructed_weights

    _, config, weights, reference_logits = real_engine
    engine = Engine.from_weights(config, weights, FakeTokenizer(), precision=precision)
    assert engine.info["precision"] == precision
    np.testing.assert_allclose(engine.logits([1, 2, 3]),
                               reference_logits(config, reconstructed_weights(weights, precision), [1, 2, 3]),
                               atol=8e-6, rtol=8e-5)


def test_from_weights_eos_stop_and_zero_budget(real_engine):
    import numpy as np

    engine, config, weights, _ = real_engine
    eos = int(np.argmax(engine.logits([1, 2])))
    stopped = Engine.from_weights({**config, "eos_token_id": [eos]}, weights,
                                  FakeTokenizer(special=[eos]))
    assert stopped.generate([1, 2], max_new_tokens=1) == GenerationResult("", [eos], "stop", 2, 1)
    assert engine.generate([1, 2], max_new_tokens=4, stop=[f"<{eos}>"]) == GenerationResult(
        "", [eos], "stop", 2, 1)
    assert engine.generate([1] * 80, max_new_tokens=0) == GenerationResult("", [], "length", 80, 0)
    assert engine.generate([1] * 80, max_new_tokens=1) == GenerationResult("", [], "length", 80, 0)
    with pytest.raises(ValueError):
        engine.logits([1] * 81)
    with pytest.raises(ValueError, match="top_k"):
        engine.generate([1], max_new_tokens=1, top_k=config["vocab_size"] + 1)
    assert not engine._generation_lock.locked()


@pytest.mark.parametrize("change", [
    {"max_position_embeddings": 0}, {"max_position_embeddings": 1048577},
    {"hidden_size": True}, {"num_hidden_layers": 257}, {"head_dim": 7},
    {"num_key_value_heads": 3}, {"model_type": "qwen3_moe"},
    {"rope_scaling": {"rope_type": "linear", "factor": 2}},
])
def test_from_weights_rejects_invalid_config(real_engine, change):
    _, config, weights, _ = real_engine
    with pytest.raises(ValueError):
        Engine.from_weights({**config, **change}, weights, FakeTokenizer())


def test_lock_and_unlock_pages_forwarding():
    engine, native = make_engine()
    assert engine.lock_pages() == 123
    assert native.pages_locked
    assert engine.unlock_pages() == 123
    assert not native.pages_locked


def test_scan_wrapper_is_serialized_and_forwarded():
    engine, native = make_engine()
    assert engine.scan() == 123
    assert native.scans == 1
    assert not engine._generation_lock.locked()


@pytest.mark.parametrize("interval", [0, -1, "2", float("nan"), 3601])
def test_keep_warm_rejects_invalid_interval(interval):
    engine, _ = make_engine()
    with pytest.raises(ValueError, match="interval"):
        engine.start_keep_warm(interval)
    assert not engine.keep_warm_active


def test_keep_warm_sweeps_until_stopped():
    engine, native = make_engine()
    engine.start_keep_warm(0.005)
    deadline = time.monotonic() + 2
    while native.touches < 2 and time.monotonic() < deadline:
        time.sleep(0.005)
    assert native.touches >= 2
    engine.stop_keep_warm()
    final = native.touches
    time.sleep(0.03)
    assert native.touches == final
    assert not engine.keep_warm_active
    assert engine.keep_warm_error is None


def test_keep_warm_skips_generation_and_resumes():
    engine, native = make_engine()
    engine._generation_lock.acquire()
    engine.start_keep_warm(0.005)
    try:
        time.sleep(0.06)
        assert native.touches == 0
    finally:
        engine._generation_lock.release()
    try:
        deadline = time.monotonic() + 2
        while native.touches == 0 and time.monotonic() < deadline:
            time.sleep(0.005)
        assert native.touches >= 1
    finally:
        engine.stop_keep_warm()
    final = native.touches
    time.sleep(0.03)
    assert native.touches == final


def test_keep_warm_double_start_and_touch_failure():
    engine, native = make_engine()
    engine.start_keep_warm(0.05)
    try:
        with pytest.raises(RuntimeError, match="already running"):
            engine.start_keep_warm(0.05)
    finally:
        engine.stop_keep_warm()

    def broken():
        raise RuntimeError("touch failed")
    monkeypatch = pytest.MonkeyPatch()
    monkeypatch.setattr(native, "touch", broken)
    try:
        engine.start_keep_warm(0.005)
        deadline = time.monotonic() + 2
        while engine.keep_warm_active and time.monotonic() < deadline:
            time.sleep(0.005)
        assert isinstance(engine.keep_warm_error, RuntimeError)
        assert not engine.keep_warm_active
    finally:
        monkeypatch.undo()
        engine.stop_keep_warm()
    # A failed sweep leaves the engine usable.
    assert engine.generate([1], max_new_tokens=1).finish_reason == "length"


def test_stop_keep_warm_without_start_is_noop():
    engine, _ = make_engine()
    engine.stop_keep_warm()
    assert not engine.keep_warm_active