From cca7071dbdf48371852a0603214b88f923aad1fa Mon Sep 17 00:00:00 2001
From: iomgaa <iomgaaycz@gmail.com>
Date: Tue, 21 Jul 2026 01:17:46 -0400
Subject: [PATCH] feat: add soak corpus loaders and scenario generators

---
 tests/unit/test_soak_corpus.py | 131 +++++++++++++++++++++++
 tools/soak/__init__.py         |   4 +
 tools/soak/corpus.py           | 124 +++++++++++++++++++++
 tools/soak/scenarios.py        | 190 +++++++++++++++++++++++++++++++++
 4 files changed, 449 insertions(+)
 create mode 100644 tests/unit/test_soak_corpus.py
 create mode 100644 tools/soak/__init__.py
 create mode 100644 tools/soak/corpus.py
 create mode 100644 tools/soak/scenarios.py

diff --git a/tests/unit/test_soak_corpus.py b/tests/unit/test_soak_corpus.py
new file mode 100644
index 0000000..64a8c1a
--- /dev/null
+++ b/tests/unit/test_soak_corpus.py
@@ -0,0 +1,131 @@
+"""soak harness 纯函数层测试(M2 设计 §8;T10/T11)。
+
+语料还原器与不变量断言函数是纯函数,用真实 schema 的二次构造微型库测试
+(真实场景优先: schema 与 data/soak 实库逐字一致)。
+"""
+
+import json
+import sqlite3
+
+import pytest
+
+from tools.soak.corpus import (
+    assemble_frame_messages,
+    chs_image_messages,
+    load_replay_payloads,
+    load_trace_chains,
+)
+from tools.soak.scenarios import weighted_mix
+
+
+def _mini_harness_db(path):
+    """predictions 表 schema 与 data/soak/harness.db 逐字一致的微型库。"""
+    conn = sqlite3.connect(path)
+    conn.execute(
+        "CREATE TABLE predictions (run_id TEXT, timestamp TEXT, video_id TEXT,"
+        " question_id TEXT, task_type TEXT, prediction TEXT, answer TEXT, evidence TEXT,"
+        " reasoning TEXT, steps_used INTEGER, prompt_tokens INTEGER,"
+        " completion_tokens INTEGER, stop_reason TEXT, steps_json TEXT)"
+    )
+    steps = [
+        {"thought": f"思考{i}", "tool_call": {"tool": "view_node", "args": {"i": i}},
+         "tool_output": f"输出{i}" * 50}
+        for i in range(3)
+    ]
+    conn.execute(
+        "INSERT INTO predictions VALUES ('r1','t','vid1','q1','synopsis','A','A','','',3,1,1,'stop',?)",
+        (json.dumps(steps, ensure_ascii=False),),
+    )
+    conn.execute(
+        "INSERT INTO predictions VALUES ('r1','t','vid2','q2','synopsis','B','B','','',0,1,1,'stop','[]')"
+    )
+    conn.commit()
+    conn.close()
+
+
+def _mini_telemetry_db(path):
+    conn = sqlite3.connect(path)
+    conn.execute(
+        "CREATE TABLE llm_calls (call_id TEXT PRIMARY KEY, parent_call_id TEXT,"
+        " session_id TEXT, model_name TEXT, provider TEXT, messages TEXT, response TEXT,"
+        " thinking TEXT, prompt_tokens INTEGER, completion_tokens INTEGER, latency_ms INTEGER,"
+        " ttft_ms REAL, max_inter_token_ms REAL, cache_hit INTEGER, error TEXT, created_at TEXT)"
+    )
+    msgs = json.dumps([{"role": "user", "content": "回放负载"}])
+    conn.execute(
+        "INSERT INTO llm_calls VALUES ('c1',NULL,'s1','m','p',?,'r','',1,1,1,NULL,NULL,0,NULL,'t')",
+        (msgs,),
+    )
+    conn.execute(
+        "INSERT INTO llm_calls VALUES ('c2',NULL,'s1','m','p','not-json','r','',1,1,1,NULL,NULL,0,NULL,'t')"
+    )
+    conn.commit()
+    conn.close()
+
+
+class TestTraceChains:
+    def test_cumulative_chain_shapes(self, tmp_path):
+        db = tmp_path / "h.db"
+        _mini_harness_db(db)
+        chains = load_trace_chains(db)
+        assert len(chains) == 1  # steps_used=0 的行剔除
+        chain = chains[0]
+        assert len(chain) == 3  # 每步一个累积快照
+        # 快照单调增长: system+user 起步,每步追加 assistant+user
+        assert [len(snap) for snap in chain] == [4, 6, 8]
+        first = chain[0]
+        assert first[0]["role"] == "system" and first[1]["role"] == "user"
+        assert first[2]["role"] == "assistant" and "思考0" in first[2]["content"]
+        assert first[3]["role"] == "user" and "输出0" in first[3]["content"]
+        # 后续快照是前缀扩展(累积语义)
+        assert chain[1][:4] == chain[0]
+
+    def test_max_steps_cap(self, tmp_path):
+        db = tmp_path / "h.db"
+        _mini_harness_db(db)
+        assert [len(s) for s in load_trace_chains(db, max_steps=2)[0]] == [4, 6]
+
+
+class TestReplayPayloads:
+    def test_loads_and_skips_unparseable(self, tmp_path):
+        db = tmp_path / "t.db"
+        _mini_telemetry_db(db)
+        payloads = load_replay_payloads(db)
+        assert len(payloads) == 1  # 坏 JSON 行剔除
+        assert payloads[0][0]["content"] == "回放负载"
+
+
+class TestImageAssembly:
+    def test_frame_messages_two_tiers(self, tmp_path):
+        vid = tmp_path / "vid1" / "frames"
+        vid.mkdir(parents=True)
+        for i in range(6):
+            (vid / f"f{i}.jpg").write_bytes(b"\xff\xd8fakejpeg" + bytes([i]))
+        msgs = assemble_frame_messages(tmp_path, n_frames=4, rng=lambda: 0.0)
+        parts = msgs[0]["content"]
+        images = [p for p in parts if p["type"] == "image_url"]
+        assert len(images) == 4
+        assert images[0]["image_url"]["url"].startswith("data:image/jpeg;base64,")
+
+    def test_chs_single_image_with_instruction(self, tmp_path):
+        (tmp_path / "chs_0001.jpg").write_bytes(b"\xff\xd8fake")
+        msgs = chs_image_messages(tmp_path / "chs_0001.jpg")
+        parts = msgs[0]["content"]
+        assert any(p["type"] == "image_url" for p in parts)
+        text = next(p["text"] for p in parts if p["type"] == "text")
+        assert "JSON" in text  # 固定结构化指令
+
+
+class TestWeightedMix:
+    def test_ratio_approximation(self):
+        import random
+
+        rng = random.Random(42).random
+        weights = {"P1": 0.1, "P2": 0.2, "P3": 0.5, "P5": 0.2}
+        picks = [weighted_mix(weights, rng) for _ in range(10_000)]
+        for key, w in weights.items():
+            assert picks.count(key) / 10_000 == pytest.approx(w, abs=0.05)
+
+    def test_invalid_weights_rejected(self):
+        with pytest.raises(ValueError):
+            weighted_mix({"P1": 0.0}, lambda: 0.5)
diff --git a/tools/soak/__init__.py b/tools/soak/__init__.py
new file mode 100644
index 0000000..af5a637
--- /dev/null
+++ b/tools/soak/__init__.py
@@ -0,0 +1,4 @@
+"""真实数据压测 harness(M2 设计 §8;findings/2026-07-20-m2-soak-workload.md)。
+
+独立工具,不被 src/ import、不入 pytest 门(纯函数部分有 unit 测试)。
+"""
diff --git a/tools/soak/corpus.py b/tools/soak/corpus.py
new file mode 100644
index 0000000..26d008e
--- /dev/null
+++ b/tools/soak/corpus.py
@@ -0,0 +1,124 @@
+"""压测语料装载(纯函数): 真实库还原请求负载,不做任何网络调用。
+
+数据源(data/soak/,不入 git):
+- harness.db predictions.steps_json: VT agent 完整轨迹 → P1 长上下文累积链
+- generate_questions_telemetry.db llm_calls.messages: 376 条真实重载 → P2 回放
+- vt_frames/<video>/frames/*.jpg: 多模态组装;chs_images/chs_NNNN.jpg: 单图短指令
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import sqlite3
+from pathlib import Path
+from typing import Any
+
+Messages = list[dict[str, Any]]
+
+_TRACE_SYSTEM = (
+    "You are a video-understanding agent replayed for gateway soak testing. "
+    "Answer the tool-use conversation faithfully and concisely."
+)
+
+# CHS 形态的固定结构化短指令(业务 schema 只存在于 tools/,不入 src——零业务假设铁律)
+CHS_INSTRUCTION = (
+    "Extract the following 12 fields from this ultrasound image and reply in strict JSON: "
+    "modality, orientation, depth_cm, gain, focus_zone, vessel_visible, lesion_present, "
+    "lesion_location, echo_pattern, doppler_signal, image_quality, notes. "
+    "Reply with a single JSON object only."
+)
+
+
+def load_trace_chains(harness_db: Path | str, *, max_steps: int = 40) -> list[list[Messages]]:
+    """还原 P1 长上下文链: 每条 prediction → 逐步累积的 messages 快照序列。
+
+    快照 k = system + user(任务行)+ 前 k 步的 (assistant: thought+tool_call,
+    user: tool_output);steps_used=0 或 steps_json 坏行剔除。
+    """
+    conn = sqlite3.connect(harness_db)
+    try:
+        rows = conn.execute(
+            "SELECT video_id, question_id, task_type, steps_json FROM predictions"
+        ).fetchall()
+    finally:
+        conn.close()
+    chains: list[list[Messages]] = []
+    for video_id, question_id, task_type, steps_json in rows:
+        try:
+            steps = json.loads(steps_json or "[]")
+        except json.JSONDecodeError:
+            continue
+        if not steps:
+            continue
+        base: Messages = [
+            {"role": "system", "content": _TRACE_SYSTEM},
+            {
+                "role": "user",
+                "content": f"[soak replay] video={video_id} question={question_id} task={task_type}",
+            },
+        ]
+        snapshots: list[Messages] = []
+        for step in steps[:max_steps]:
+            base = base + [
+                {
+                    "role": "assistant",
+                    "content": f"{step.get('thought', '')}\n{json.dumps(step.get('tool_call', {}), ensure_ascii=False)}",
+                },
+                {"role": "user", "content": str(step.get("tool_output", ""))},
+            ]
+            snapshots.append(base)
+        chains.append(snapshots)
+    return chains
+
+
+def load_replay_payloads(telemetry_db: Path | str) -> list[Messages]:
+    """P2 回放负载: 原样取真实 messages(内嵌帧的多模态重载),坏 JSON 行剔除。"""
+    conn = sqlite3.connect(telemetry_db)
+    try:
+        rows = conn.execute("SELECT messages FROM llm_calls").fetchall()
+    finally:
+        conn.close()
+    payloads: list[Messages] = []
+    for (raw,) in rows:
+        try:
+            messages = json.loads(raw)
+        except (json.JSONDecodeError, TypeError):
+            continue
+        if isinstance(messages, list) and messages:
+            payloads.append(messages)
+    return payloads
+
+
+def _data_uri(path: Path) -> str:
+    return "data:image/jpeg;base64," + base64.b64encode(path.read_bytes()).decode("ascii")
+
+
+def assemble_frame_messages(frames_root: Path | str, *, n_frames: int, rng) -> Messages:
+    """VT 帧组装(P2 两档: 1-4 帧 / 5-6 帧): 随机选视频取前 n 帧 + 简短指令。"""
+    root = Path(frames_root)
+    videos = sorted(d for d in root.iterdir() if (d / "frames").is_dir())
+    if not videos:
+        raise ValueError(f"{root} 下没有 <video>/frames/ 目录")
+    video = videos[int(rng() * len(videos)) % len(videos)]
+    frames = sorted((video / "frames").glob("*.jpg"))[:n_frames]
+    if len(frames) < n_frames:
+        raise ValueError(f"{video} 帧数不足 {n_frames}")
+    parts: list[dict[str, Any]] = [
+        {"type": "text", "text": f"Describe the key visual content of these {n_frames} frames."}
+    ]
+    parts.extend({"type": "image_url", "image_url": {"url": _data_uri(f)}} for f in frames)
+    return [{"role": "user", "content": parts}]
+
+
+def chs_image_messages(image_path: Path | str) -> Messages:
+    """CHS 形态(P3): 单图 + 固定 12 字段结构化短指令。"""
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": CHS_INSTRUCTION},
+                {"type": "image_url", "image_url": {"url": _data_uri(Path(image_path))}},
+            ],
+        }
+    ]
diff --git a/tools/soak/scenarios.py b/tools/soak/scenarios.py
new file mode 100644
index 0000000..1642b51
--- /dev/null
+++ b/tools/soak/scenarios.py
@@ -0,0 +1,190 @@
+"""P1-P6 场景请求生成器(findings §2 矩阵;async 生成器,产出调用参数)。
+
+每项产出 `(kind, kwargs)`: kind ∈ {"chat"};kwargs 直接喂
+`GatewayClient.chat(**kwargs)`。回放/组装场景一律掺 `cache_salt=run_id`
+破缓存(缓存行为归 P4);P4 子流量特意重复 messages 且不掺 salt。
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, AsyncIterator
+
+from pydantic import BaseModel
+
+from tools.soak.corpus import (
+    assemble_frame_messages,
+    chs_image_messages,
+    load_replay_payloads,
+    load_trace_chains,
+)
+
+Item = tuple[str, dict[str, Any]]
+
+
+class ChsExtraction(BaseModel):
+    """P3 结构化档的真实形态 schema(12 字段,仅存在于 tools/)。"""
+
+    modality: str
+    orientation: str
+    depth_cm: float | None = None
+    gain: str | None = None
+    focus_zone: str | None = None
+    vessel_visible: bool | None = None
+    lesion_present: bool | None = None
+    lesion_location: str | None = None
+    echo_pattern: str | None = None
+    doppler_signal: str | None = None
+    image_quality: str
+    notes: str | None = None
+
+
+@dataclass
+class SoakCorpus:
+    """一次 run 的语料句柄(装载一次,场景间共享)。"""
+
+    harness_db: Path
+    telemetry_db: Path
+    frames_root: Path
+    images_root: Path
+
+    def __post_init__(self) -> None:
+        self.chains = load_trace_chains(self.harness_db)
+        self.replays = load_replay_payloads(self.telemetry_db)
+        self.images = sorted(Path(self.images_root).glob("chs_*.jpg"))
+        if not (self.chains and self.replays and self.images):
+            raise ValueError("语料不完整: 请确认 data/soak/ 已按 findings §6 拉取")
+
+
+def weighted_mix(weights: dict[str, float], rng) -> str:
+    """按权重抽一个场景 key;权重和必须 > 0。"""
+    total = sum(weights.values())
+    if total <= 0:
+        raise ValueError("场景权重之和必须 > 0")
+    point = rng() * total
+    acc = 0.0
+    for key, w in weights.items():
+        acc += w
+        if point < acc:
+            return key
+    return next(reversed(weights))
+
+
+async def p1_trace_chains(corpus: SoakCorpus, run_id: str, rng=random.random) -> AsyncIterator[Item]:
+    """P1 文本长上下文回放: 单链串行,session/parent 链路照原样语义。"""
+    for chain_idx, chain in enumerate(corpus.chains):
+        session_id = f"{run_id}-p1-{chain_idx}"
+        parent: str | None = None
+        for snapshot in chain:
+            yield (
+                "chat",
+                {
+                    "messages": snapshot,
+                    "session_id": session_id,
+                    "parent_call_id": parent,
+                    "cache_salt": run_id,
+                },
+            )
+            parent = session_id  # 链内父子: 以 session 为锚(真实 call_id 由库生成)
+
+
+async def p2_multimodal_replay(
+    corpus: SoakCorpus, run_id: str, rng=random.random
+) -> AsyncIterator[Item]:
+    """P2 多模态重载: 真实 376 条原样回放 + 帧组装两档交错。"""
+    for i, messages in enumerate(corpus.replays):
+        yield ("chat", {"messages": messages, "cache_salt": run_id, "session_id": f"{run_id}-p2"})
+        if i % 3 == 0:
+            n = 4 if rng() < 0.5 else 6
+            yield (
+                "chat",
+                {
+                    "messages": assemble_frame_messages(corpus.frames_root, n_frames=n, rng=rng),
+                    "cache_salt": run_id,
+                    "session_id": f"{run_id}-p2",
+                },
+            )
+
+
+async def p3_single_image(
+    corpus: SoakCorpus, run_id: str, rng=random.random, *, structured_ratio: float = 0.5
+) -> AsyncIterator[Item]:
+    """P3 单图短指令高频主场景: 无状态,50% 走 structured 真实 schema。"""
+    while True:
+        image = corpus.images[int(rng() * len(corpus.images)) % len(corpus.images)]
+        kwargs: dict[str, Any] = {
+            "messages": chs_image_messages(image),
+            "cache_salt": run_id,
+            "session_id": f"{run_id}-p3",
+        }
+        if rng() < structured_ratio:
+            kwargs["structured"] = ChsExtraction
+        yield ("chat", kwargs)
+
+
+async def p4_cache_bidirectional(
+    corpus: SoakCorpus, run_id: str, rng=random.random
+) -> AsyncIterator[Item]:
+    """P4 缓存双向: 固定小图池重复 messages(命中侧,不掺 salt)+ salt 对照(强制 miss)。"""
+    pool = corpus.images[:5]
+    while True:
+        image = pool[int(rng() * len(pool)) % len(pool)]
+        messages = chs_image_messages(image)
+        if rng() < 0.5:
+            yield ("chat", {"messages": messages, "session_id": f"{run_id}-p4"})  # 可命中
+        else:
+            yield (
+                "chat",
+                {
+                    "messages": messages,
+                    "cache_salt": f"{run_id}-{rng()}",  # 强制 miss 对照
+                    "session_id": f"{run_id}-p4",
+                },
+            )
+
+
+async def p5_fault_mixed(corpus: SoakCorpus, run_id: str, rng=random.random) -> AsyncIterator[Item]:
+    """P5 故障源混编: 请求形态同 P3;故障性来自 SOAK scope 源池配置(findings §3),
+    生成器本身不造故障——真实故障由坏 key/黑洞/紧闸源在协议层自然发生。"""
+    async for item in p3_single_image(corpus, run_id, rng, structured_ratio=0.3):
+        yield item
+
+
+async def p6_mixed_soak(
+    corpus: SoakCorpus,
+    run_id: str,
+    rng=random.random,
+    *,
+    weights: dict[str, float] | None = None,
+) -> AsyncIterator[Item]:
+    """P6 混合浸泡: 按签字比例(设计 §8.1)加权混合;P4 已并入 P3 权重的 3/10。"""
+    weights = weights or {"P1": 0.10, "P2": 0.20, "P3": 0.35, "P4": 0.15, "P5": 0.20}
+    gens = {
+        "P1": p1_trace_chains(corpus, run_id, rng),
+        "P2": p2_multimodal_replay(corpus, run_id, rng),
+        "P3": p3_single_image(corpus, run_id, rng),
+        "P4": p4_cache_bidirectional(corpus, run_id, rng),
+        "P5": p5_fault_mixed(corpus, run_id, rng),
+    }
+    while True:
+        key = weighted_mix(weights, rng)
+        try:
+            yield await gens[key].__anext__()
+        except StopAsyncIteration:
+            # 有限语料场景(P1/P2)耗尽后重启一轮(浸泡语义: 语料循环使用)
+            gens[key] = {
+                "P1": p1_trace_chains,
+                "P2": p2_multimodal_replay,
+            }[key](corpus, f"{run_id}-r{rng()}", rng)
+
+
+SCENARIOS = {
+    "P1": p1_trace_chains,
+    "P2": p2_multimodal_replay,
+    "P3": p3_single_image,
+    "P4": p4_cache_bidirectional,
+    "P5": p5_fault_mixed,
+    "P6": p6_mixed_soak,
+}