tests/generated_cache_api_test.py

"""Dedicated-board regression for opt-in generated-state reuse; no concurrent inference."""
import argparse
import json
import time
import urllib.request
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--base", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--budget-test", action="store_true")
    args = parser.parse_args()
    client = urllib.request.build_opener(urllib.request.ProxyHandler({}))

    def open_request(path, body=None):
        return client.open(urllib.request.Request(args.base + path,
            data=None if body is None else json.dumps(body).encode(),
            headers={"Content-Type": "application/json"}), timeout=180)

    def call(path, body=None):
        with open_request(path, body) as response:
            return json.load(response)

    def infer(messages, session="a", reuse=True, count=16, **extra):
        body = {"model": "mindnano-ling3-tiny", "user": "generated-test", "messages": messages,
                "temperature": 0, "max_tokens": count, "reuse_generated_state": reuse}
        if session is not None:
            body["session_id"] = session
        return call("/v1/chat/completions", dict(body, **extra))

    def equal(a, b):
        assert a["choices"] == b["choices"], "identical numerical path changed output"

    def follow(seed, result, question="请简短说明刚才的内容。"):
        return seed + [result["choices"][0]["message"], {"role": "user", "content": question}]

    def clear():
        call("/v1/cache/clear", {})

    results = {"cases": [], "passed": False}
    assert not call("/v1/generation/status")["active"]
    assert call("/health")["capabilities"]["generated_state_reuse"]
    clear()
    try:
        seed = [{"role": "user", "content": "1" * 107}]  # 128-token stable prefix
        if args.budget_test:
            first = infer(seed, count=64)
            assert first["mindnano_metrics"]["continuation_skipped_budget"]
            assert first["mindnano_metrics"]["continuation_stored"]
            later = infer(follow(seed, first))
            assert later["mindnano_metrics"]["generated_cached_tokens"] == 63
            status = call("/v1/cache/status")
            assert status["snapshot_bytes"] <= status["budget_bytes"]
            # Fair single-session comparison: the old path also keeps its
            # original prompt cache. Compare against it, not only cold prefill.
            clear()
            first = infer(seed, session=None, count=64)
            messages = follow(seed, first)
            incremental = infer(messages, session=None)
            clear()
            reference_seed = infer(seed, session=None, reuse=False, count=64)
            equal(first, reference_seed)
            baseline = infer(messages, session=None, reuse=False)
            assert baseline["mindnano_metrics"]["cached_tokens"] == 128
            assert incremental["mindnano_metrics"]["generated_cached_tokens"] == 63
            results["cached_baseline"] = {"incremental": incremental, "original_prompt_cache": baseline,
                                          "same_reply": incremental["choices"] == baseline["choices"]}
            clear()
            long_seed = infer(seed, session=None, count=512)
            long_messages = follow(seed, long_seed)
            long_incremental = infer(long_messages, session=None)
            assert long_incremental["mindnano_metrics"]["generated_cached_tokens"] == 511
            clear()
            long_reference = infer(seed, session=None, reuse=False, count=512)
            equal(long_seed, long_reference)
            long_baseline = infer(long_messages, session=None, reuse=False)
            assert long_baseline["mindnano_metrics"]["cached_tokens"] == 128
            results["long_cached_baseline"] = {"incremental": long_incremental, "original_prompt_cache": long_baseline,
                                               "same_reply": long_incremental["choices"] == long_baseline["choices"]}
            results.update(passed=True, budget=status)
            return
        for session in (None, "a"):
            clear()
            first = infer(seed, session, count=64)
            assert first["mindnano_metrics"]["continuation_stored"]
            messages = follow(seed, first)
            if session:
                infer([{"role": "user", "content": "你好。"}], "other")
            warm = infer(messages, session)
            m = warm["mindnano_metrics"]
            assert m["generated_cached_tokens"] == 63
            assert m["cache_status"] == "generated_prefix_hit"
            assert m["cached_tokens"] + m["prompt_evaluated_tokens"] == warm["usage"]["prompt_tokens"]
            again = infer(messages, session)
            equal(warm, again)
            assert again["mindnano_metrics"]["prompt_evaluated_tokens"] == 0
            # Turning it off must also reject exact caches derived from decode state.
            strict = infer(messages, session, reuse=False)
            assert strict["mindnano_metrics"]["cached_tokens"] == 0
            cold = infer(messages, session, reuse=False, cache_prompt=False)
            equal(strict, cold)
            results["cases"].append({"session": session, "warm": warm, "cold": cold,
                                     "same_reply_as_cold": warm["choices"] == cold["choices"]})

        # Fork and cancel preserve the last committed generated prefix.
        clear()
        first = infer(seed, count=64)
        call("/v1/cache/fork", {"user": "generated-test", "source_session_id": "a", "target_session_id": "fork"})
        messages = follow(seed, first)
        branch = infer(messages, "fork")
        assert branch["mindnano_metrics"]["generated_cached_tokens"] == 63
        body = {"model": "mindnano-ling3-tiny", "user": "generated-test", "session_id": "a",
                "messages": [{"role": "user", "content": "讲个故事。"}], "stream": True,
                "flow_control": "ack", "reuse_generated_state": True, "max_tokens": 64}
        with open_request("/v1/chat/completions", body) as response:
            for line in response:
                if line.startswith(b"data:"):
                    event = json.loads(line[5:])
                    if event.get("mindnano_flow"):
                        call("/v1/cancel", {"request_id": event["id"]})
                        break
            for line in response:
                if line.strip() == b"data: [DONE]":
                    break
        deadline = time.monotonic() + 10
        while call("/v1/generation/status")["active"]:
            assert time.monotonic() < deadline
            time.sleep(.05)
        resumed = infer(messages)
        assert resumed["mindnano_metrics"]["generated_cached_tokens"] == 63
        equal(branch, resumed)

        # Changed/shorter histories and stop-filtered output cannot reuse hidden tokens.
        for variant in ("edited", "shorter", "stop"):
            clear()
            first = infer(seed, count=64, **({"stop": "11111"} if variant == "stop" else {}))
            messages = follow(seed, first)
            if variant == "edited":
                messages[1]["content"] = "修改过的回答。"
            elif variant == "shorter":
                messages = [{"role": "user", "content": "你好。"}]
            response = infer(messages)
            assert response["mindnano_metrics"]["generated_cached_tokens"] == 0
        # Sample both language modes. Different cold output is measured, not hidden.
        for think in (False, True):
            clear()
            seed = [{"role": "user", "content": "用中文介绍一次会议记录应包含什么。"}]
            first = infer(seed, count=128, enable_thinking=think)
            messages = follow(seed, first, "请简短总结。")
            warm = infer(messages, enable_thinking=think)
            cold = infer(messages, reuse=False, cache_prompt=False, enable_thinking=think)
            results["cases"].append({"thinking": think, "warm": warm, "cold": cold,
                                     "same_reply_as_cold": warm["choices"] == cold["choices"]})
        results["passed"] = True
    finally:
        clear()
        args.output.write_text(json.dumps(results, ensure_ascii=False, indent=2) + "\n")
        print(json.dumps({"generated_cache": "PASS" if results["passed"] else "FAIL",
                          "results": str(args.output)}), flush=True)


if __name__ == "__main__":
    main()