games-development-ai/cheap-worker/tests/test_m3_stream_patch.py

"""m3_stream_patch 单测:合成 chunk 流复刻 amgen-93001「同 index 异 id 双完整 JSON」指纹,验证补丁分桶。

零网络/LLM:用 SimpleNamespace 假 chunk + 假 AsyncStream 直驱 agentscope 2.0.2 真聚合函数
OpenAIChatModel._parse_stream_response(补丁前原函数复刻炸弹作对照;补丁后断言两 call 正确分桶、
args 各自合法 JSON)。跑:cheap-worker/.venv/bin/python -m pytest cheap-worker/tests/test_m3_stream_patch.py -v
"""
import asyncio
import json
import logging
import sys
from datetime import datetime
from pathlib import Path
from types import SimpleNamespace

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # → cheap-worker/
import _bootstrap  # noqa: E402,F401 —— 只为 sys.path(tier2/gen-worker)兜底,不触网
import m3_stream_patch as P  # noqa: E402

from agentscope.message import ToolCallBlock  # noqa: E402
from agentscope.model import OpenAIChatModel  # noqa: E402


# ── 合成件:假 tool_call 分片 / 假 chunk / 假 openai AsyncStream(async CM + async iterator)──

def _tc(index, id=None, name=None, args=""):
    """一个流式 tool_call 分片(形状对齐 openai SDK 的 ChoiceDeltaToolCall:index/id/function.{name,arguments})。"""
    return SimpleNamespace(index=index, id=id, function=SimpleNamespace(name=name, arguments=args))


def _chunk(tool_calls=None):
    """一个流式 chunk(delta 只带 tool_calls;content/audio 置 None 走聚合函数的空分支)。"""
    delta = SimpleNamespace(content=None, tool_calls=tool_calls or [], audio=None)
    return SimpleNamespace(usage=None, id="chatcmpl-test", choices=[SimpleNamespace(delta=delta)])


class _FakeStream:
    """假 openai AsyncStream:async with 返回自身、async for 逐个吐 chunk(契约同真流,聚合函数无感)。"""

    def __init__(self, chunks):
        self._chunks = chunks

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        return False

    def __aiter__(self):
        return self._gen()

    async def _gen(self):
        for c in self._chunks:
            yield c


def _run_parse(parse_fn, chunks):
    """驱动(原始或补丁后的)_parse_stream_response 消费合成流,返回 final ChatResponse 的 ToolCallBlock 列表。"""
    model = object.__new__(OpenAIChatModel)  # 聚合函数体不读 self 属性,绕 __init__ 免建真 client

    async def _go():
        out = []
        async for resp in parse_fn(model, datetime.now(), _FakeStream(chunks)):
            out.append(resp)
        return out

    responses = asyncio.run(_go())
    assert responses, "至少应有 is_last=True 的 final ChatResponse"
    final = responses[-1]
    return [b for b in final.content if isinstance(b, ToolCallBlock)]


# 93001 指纹场景:两个并行 call **同 index=0、各带不同 id 与完整 JSON args**(M3 经 new-api 的真实形状)。
_BOMB_CHUNKS_FACTORY = lambda: [  # noqa: E731 —— 每用例新造(补丁会就地改写 index,复用会串场景)
    _chunk([_tc(0, id="call_A", name="write_file", args='{"path":"game-logic.js","content":"x"}')]),
    _chunk([_tc(0, id="call_B", name="run_check", args='{"target":"all"}')]),
]


def _ensure_patched():
    """确保补丁已装(幂等);返回补丁函数与其保留的原函数引用。"""
    P.apply_m3_stream_patch()
    patched = OpenAIChatModel.__dict__["_parse_stream_response"]
    assert getattr(patched, P._PATCH_FLAG, False), "补丁应已挂上"
    return patched, patched._m3_fix400_orig


def test_original_reproduces_bomb_fingerprint():
    # 对照组:2.0.2 原聚合函数在同 index 异 id 流上 = 拼桶(1 个 call、id 留首个、args 双 JSON 拼接非法)——
    # 坐实本测试的合成流踩的就是现场炸弹路径(93001 行 548-552 指纹),不是自造靶子。
    _, orig = _ensure_patched()
    blocks = _run_parse(orig, _BOMB_CHUNKS_FACTORY())
    assert len(blocks) == 1, "原函数应把两个 call 拼进一个桶"
    assert blocks[0].id == "call_A", "id 留首个"
    merged = blocks[0].input
    assert '{"path"' in merged and '{"target"' in merged, "args 应为双完整 JSON 拼接"
    try:
        json.loads(merged)
        raise AssertionError("拼接 args 不应是合法 JSON(否则不构成 400 炸弹)")
    except json.JSONDecodeError:
        pass


def test_patched_splits_buckets_by_id():
    # 主断言(工单 B):补丁后同 index 异 id 的两个 call 正确分桶,id/name 各自保留、args 各自合法 JSON。
    patched, _ = _ensure_patched()
    blocks = _run_parse(patched, _BOMB_CHUNKS_FACTORY())
    assert len(blocks) == 2, "补丁后应按 id 分成两个 call"
    by_id = {b.id: b for b in blocks}
    assert set(by_id) == {"call_A", "call_B"}
    assert by_id["call_A"].name == "write_file"
    assert by_id["call_B"].name == "run_check"
    assert json.loads(by_id["call_A"].input) == {"path": "game-logic.js", "content": "x"}
    assert json.loads(by_id["call_B"].input) == {"target": "all"}


def test_patched_normal_openai_stream_unchanged():
    # 回归护栏:标准 OpenAI 并行流(不同 call 不同 index、id 只在首帧、args 跨 chunk 分片)行为不变。
    patched, _ = _ensure_patched()
    chunks = [
        _chunk([_tc(0, id="call_X", name="write_file", args='{"pa')]),
        _chunk([_tc(1, id="call_Y", name="run_check", args='{"target"')]),
        _chunk([_tc(0, args='th":"a.js"}')]),   # 无 id 续片 → 跟 index=0 最近开的桶(call_X)
        _chunk([_tc(1, args=':"all"}')]),        # 无 id 续片 → 跟 index=1 最近开的桶(call_Y)
    ]
    blocks = _run_parse(patched, chunks)
    assert len(blocks) == 2
    by_id = {b.id: b for b in blocks}
    assert json.loads(by_id["call_X"].input) == {"path": "a.js"}
    assert json.loads(by_id["call_Y"].input) == {"target": "all"}


def test_patched_same_id_resent_per_frame_stays_one_bucket():
    # 部分网关每帧重复带同一 id:同 id 恒同桶(不能被「新帧带 id」误开新桶)。
    patched, _ = _ensure_patched()
    chunks = [
        _chunk([_tc(0, id="call_R", name="write_file", args='{"a"')]),
        _chunk([_tc(0, id="call_R", args=':1}')]),
    ]
    blocks = _run_parse(patched, chunks)
    assert len(blocks) == 1
    assert json.loads(blocks[0].input) == {"a": 1}


def test_malformed_args_logs_warning_fingerprint(caplog):
    # 聚合完 args 非空且非法 JSON → log.warning 带指纹(残余脏流可观测);合法/空 args 不告警。
    patched, _ = _ensure_patched()
    with caplog.at_level(logging.WARNING, logger="cheap.m3_stream_patch"):
        blocks = _run_parse(patched, [_chunk([_tc(0, id="call_bad", name="write_file", args='{"broken')])])
    assert len(blocks) == 1
    warns = [r for r in caplog.records if "arguments 非法 JSON" in r.getMessage()]
    assert len(warns) == 1, "坏 args 应恰好一条 warning"
    msg = warns[0].getMessage()
    assert "call_bad" in msg and "write_file" in msg, "warning 应带 id/name 指纹"

    caplog.clear()
    with caplog.at_level(logging.WARNING, logger="cheap.m3_stream_patch"):
        _run_parse(patched, _BOMB_CHUNKS_FACTORY())  # 分桶后各自合法 → 无告警
    assert not [r for r in caplog.records if "arguments 非法 JSON" in r.getMessage()]


def test_apply_is_idempotent():
    # 幂等:首次(或此前用例)已应用 → 再次 apply 返回 False、不套第二层,行为仍正确。
    P.apply_m3_stream_patch()
    assert P.apply_m3_stream_patch() is False
    patched = OpenAIChatModel.__dict__["_parse_stream_response"]
    assert getattr(patched._m3_fix400_orig, P._PATCH_FLAG, False) is False, "原函数引用不应是补丁自身(未套两层)"
    blocks = _run_parse(patched, _BOMB_CHUNKS_FACTORY())
    assert len(blocks) == 2