#11 抓包默认流式 + 单一聚合预算 + store 收敛

capture/ai_capture_addon.py

@@ -543,7 +543,12 @@ def request(self, flow):
         req = extract_request(variant, body)
         bounded_messages, request_truncated = _truncate_value(req["messages"])
         bounded_system, system_truncated = _truncate_value(req["system"])
+        acc = CaptureAccumulator(budget=MAX_CAPTURE_FLOW_BYTES)
+        acc.reserve_request_snapshot(
+            json.dumps({"system": bounded_system, "messages": bounded_messages},
+                       ensure_ascii=False).encode("utf-8"))
         flow.metadata["ai_capture"] = {
+            "_acc": acc,
             "ts": _utc_now_iso(), "flow_id": flow.id, "method": flow.request.method,
             "url": flow.request.url, "host": flow.request.pretty_host,
             "provider": provider, "variant": variant, "model": req["model"],
@@ -553,37 +558,30 @@ def request(self, flow):
         }
 
     def responseheaders(self, flow):
-        # Opt-in: tee streamed chunks (pass through unchanged) to preserve live
-        # browser streaming while still capturing. Default is buffered (below).
-        if not self.preserve_streaming:
-            return
+        # issue #11：AI 流式响应默认 tee/pass-through——首块延迟不再等待
+        # 完整响应（此前需 MAGIC_PROXY_PRESERVE_STREAMING 选择性开启）。
+        # 捕获异常或预算耗尽都不影响下发：tee 原样返回 chunk。
         meta = flow.metadata.get("ai_capture")
-        if not (meta and meta["stream"]):
+        if not meta:
             return
-        buffer = bytearray()
-        meta["_response_total_bytes"] = 0
-
-        def tee(chunk):
-            meta["_response_total_bytes"] += len(chunk)
-            remaining = MAX_CAPTURE_FLOW_BYTES - len(buffer)
-            if remaining > 0:
-                buffer.extend(chunk[:remaining])
-            if len(chunk) > remaining:
-                meta["response_truncated"] = True
-            return chunk  # read-only: never modify the proxied bytes
-
-        flow.response.stream = tee
-        meta["_tee"] = buffer
+        if not meta["stream"] and not self.capture_raw_sse:
+            return  # 非流式走 buffered 路径（response 钩子读 content）
+        acc = meta.get("_acc")
+        if acc is None:
+            acc = CaptureAccumulator(budget=MAX_CAPTURE_FLOW_BYTES)
+            meta["_acc"] = acc
+        flow.response.stream = acc.tee
 
     def response(self, flow):
         meta = flow.metadata.get("ai_capture")
         if not meta or meta.get("_written"):
             return
         try:
-            if "_tee" in meta:
-                raw = bytes(meta["_tee"])
-                resp_text = raw.decode("utf-8", "replace")
-                bytes_down = meta.get("_response_total_bytes", len(raw))
+            acc = meta.get("_acc")
+            if acc is not None and acc.total_seen:
+                resp_text = acc.captured().decode("utf-8", "replace")
+                bytes_down = acc.total_seen
+                meta["response_truncated"] = acc.truncated
             else:
                 content = flow.response.content or b""
                 bytes_down = len(flow.response.raw_content or b"")
@@ -606,7 +604,8 @@ def error(self, flow):
             return
         try:
             meta["duration_ms"] = int((time.monotonic() - meta["t0"]) * 1000)
-            raw = bytes(meta["_tee"]).decode("utf-8", "replace") if "_tee" in meta else ""
+            acc0 = meta.get("_acc")
+            raw = acc0.captured().decode("utf-8", "replace") if acc0 is not None else ""
             record = build_record(meta, None, raw, len(raw), capture_raw_sse=self.capture_raw_sse)
             record["capture_error"] = record["capture_error"] or "flow error / aborted before completion"
             record.setdefault("response", {})["raw"] = raw
@@ -617,3 +616,61 @@ def error(self, flow):
 
 
 addons = [AICaptureAddon()]
+
+
+class CaptureAccumulator:
+    """每 flow 单一聚合内存预算（issue #11）.
+
+    tee() 原样立即下发（对代理字节流零影响），同时把 chunk 收进预算内
+    缓冲；request 快照先预留同一预算池。截断按字节累计——大量小
+    message 也无法绕过。captured() 保证 UTF-8 安全（截点回退到字符
+    边界由消费方 replace 兜底，此处按字节预算硬停）。
+    """
+
+    def __init__(self, budget=MAX_CAPTURE_FLOW_BYTES):
+        self.budget = budget
+        self._buf = bytearray()       # 仅响应字节（resp_text 来源）
+        self._reserved = 0            # 请求快照占去的预算（不进 _buf）
+        self.total_seen = 0
+        self.truncated = False
+
+    def reserve_request_snapshot(self, data: bytes):
+        """请求快照占用预算池（与响应累计共享同一总额，但不进响应缓冲）。"""
+        take = min(len(data), self.budget)
+        self._reserved = take
+        if len(data) > take:
+            self.truncated = True
+
+    def tee(self, chunk: bytes) -> bytes:
+        """流量字节原样下发；预算内收集。绝不修改 chunk。"""
+        self.total_seen += len(chunk)
+        remaining = self.budget - self._reserved - len(self._buf)
+        if remaining <= 0:
+            self.truncated = True
+            return chunk
+        take = min(len(chunk), remaining)
+        self._buf += chunk[:take]
+        if take < len(chunk):
+            self.truncated = True
+        return chunk
+
+    def total_budgeted(self) -> int:
+        """聚合占用 = 请求快照预留 + 响应缓冲。"""
+        return self._reserved + len(self._buf)
+
+    def captured(self) -> bytes:
+        """预算内响应缓冲；截点回退到 UTF-8 字符边界（不产生半个字符）。"""
+        buf = bytes(self._buf)
+        if not self.truncated:
+            return buf
+        # 回退到字符边界：先跳过尾部的后续字节（10xxxxxx），再检查
+        # 其起始字节（>=0xC0）——被截断的序列按其前缀位数判断是否完整
+        end = len(buf)
+        while end > 0 and (buf[end - 1] & 0xC0) == 0x80:
+            end -= 1
+        if end > 0 and buf[end - 1] >= 0xC0:
+            start = buf[end - 1]
+            need = (2 if start < 0xE0 else 3 if start < 0xF0 else 4)
+            if end - 1 + need > len(buf):
+                end -= 1  # 序列不完整——丢弃整个残字符
+        return buf[:end]

capture/capture_store.py

@@ -10,6 +10,7 @@
 MARKER = ".magic-proxy-capture-store"
 MAX_FILE_BYTES = 50 * 1024 * 1024
 MAX_STORE_BYTES = 200 * 1024 * 1024
+MAX_RECORD_BYTES = 8 * 1024 * 1024  # 单条 record 上限：append 前拒收
 
 
 def _home_dir():
@@ -90,6 +91,19 @@ def _trim_store(directory):
         total -= size
 
 
+def _converge_after_append(directory):
+    """append 后即时预算收敛（_trim_store：留新删旧；绝非全清）。
+
+    best-effort——失败只记日志，写入结果不受影响。
+    """
+    try:
+        _trim_store(directory)
+    except OSError:
+        import logging
+        logging.getLogger("magic-proxy.capture_store").debug(
+            "post-append converge skipped", exc_info=True)
+
+
 def append_json(record, directory):
     directory = prepare(directory)
     _trim_store(directory)
@@ -101,6 +115,10 @@ def append_json(record, directory):
     name = datetime.now().strftime("%Y-%m-%d") + ".jsonl"
     flags = os.O_WRONLY | os.O_APPEND | os.O_CREAT | getattr(os, "O_NOFOLLOW", 0)
     try:
+        payload = json.dumps(record, ensure_ascii=False)
+        if len(payload.encode("utf-8")) > MAX_RECORD_BYTES:
+            raise OSError(
+                f"单条抓包记录超过 {MAX_RECORD_BYTES} 字节上限，已拒收")
         try:
             info = os.stat(name, dir_fd=dir_fd, follow_symlinks=False)
             if not stat.S_ISREG(info.st_mode):
@@ -120,7 +138,9 @@ def append_json(record, directory):
             raise OSError("抓包文件所有者或类型不安全")
         os.fchmod(fd, 0o600)
         with os.fdopen(fd, "a", encoding="utf-8") as fh:
-            fh.write(json.dumps(record, ensure_ascii=False) + "\n")
+            fh.write(payload + "\n")
+        # append 后容量收敛：单条超大也不让总量长期超限
+        _converge_after_append(directory)
         return os.path.join(directory, name)
     finally:
         os.close(dir_fd)

tests/test_ai_capture_addon.py

@@ -549,28 +549,35 @@ def test_load_defaults_when_env_absent(self, monkeypatch):
         assert a.capture_raw_sse is False
         assert a.preserve_streaming is False
 
-    def test_preserve_streaming_off_by_default_no_tee(self, tmp_path):
-        a = self._addon(tmp_path)  # preserve_streaming defaults False
-        body = {"model": "gpt-4o", "stream": True, "messages": [{"role": "user", "content": "x"}]}
-        flow = _flow("api.openai.com", "/v1/chat/completions", req_body=body, resp_text="")
+    def test_streaming_tee_on_by_default(self, tmp_path):
+        """issue #11：AI 流式默认 tee——不再需要 MAGIC_PROXY_PRESERVE_STREAMING。"""
+        a = self._addon(tmp_path)
+        body = {"model": "gpt-4o", "stream": True,
+                "messages": [{"role": "user", "content": "x"}]}
+        flow = _flow("api.openai.com", "/v1/chat/completions",
+                     req_body=body, resp_text="")
         a.request(flow)
         a.responseheaders(flow)
-        assert getattr(flow.response, "stream", None) is None  # no tee installed
+        assert callable(getattr(flow.response, "stream", None))
+
 
-    def test_preserve_streaming_tees_chunks_unchanged_and_captures(self, tmp_path):
+    def test_streaming_tees_chunks_unchanged_and_captures(self, tmp_path):
         a = self._addon(tmp_path)
-        a.preserve_streaming = True
-        body = {"model": "gpt-4o", "stream": True, "messages": [{"role": "user", "content": "x"}]}
-        flow = _flow("api.openai.com", "/v1/chat/completions", req_body=body, resp_text="")
+        body = {"model": "gpt-4o", "stream": True,
+                "messages": [{"role": "user", "content": "x"}]}
+        flow = _flow("api.openai.com", "/v1/chat/completions",
+                     req_body=body, resp_text="")
         a.request(flow)
         a.responseheaders(flow)
-        assert callable(flow.response.stream)
-        chunk = b'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n'
-        assert flow.response.stream(chunk) == chunk  # read-only: passes through unchanged
-        flow.response.stream(b"data: [DONE]\n\n")
+        out1 = flow.response.stream(b"chunk-1")
+        out2 = flow.response.stream(b"chunk-2")
+        assert (out1, out2) == (b"chunk-1", b"chunk-2")  # 原样下发
         a.response(flow)
         rec = _only_jsonl_record(tmp_path)
-        assert rec["response"]["reassembled"] == "hi"
+        # 默认不落 raw（opt-in）；断言捕获量与未截断
+        assert rec["response"]["truncated"] is False
+        assert rec["response"]["captured_bytes"] == len(b"chunk-1chunk-2")
+
 
     def test_error_hook_writes_partial_record_with_capture_error(self, tmp_path):
         a = self._addon(tmp_path)
@@ -584,7 +591,7 @@ def test_error_hook_writes_partial_record_with_capture_error(self, tmp_path):
     def test_stream_capture_is_bounded_and_marked_truncated(self, tmp_path, monkeypatch):
         monkeypatch.setattr(addon, "MAX_CAPTURE_FLOW_BYTES", 16)
         a = self._addon(tmp_path)
-        a.preserve_streaming = True
+        a.preserve_streaming = True  # legacy 属性仍可设（不再门控）
         flow = _flow(
             "api.openai.com", "/v1/chat/completions",
             req_body={"model": "gpt", "stream": True, "messages": []},

tests/test_capture_accumulator.py

@@ -0,0 +1,90 @@
+"""CaptureAccumulator（issue #11）：每 flow 单一聚合内存预算.
+
+流量字节原样立即下发；tee 收集受总预算约束（request 快照 + response
+累计共享同一预算）；大量小 messages/blocks 也无法绕过（按累计字节计，
+不按条目数）。truncated/captured_bytes 语义对 UTF-8 安全（不切多字节）。
+"""
+import pathlib
+import unittest
+
+from capture.ai_capture_addon import CaptureAccumulator
+
+
+class TestBudget(unittest.TestCase):
+    def test_default_tee_passes_bytes_through_unchanged(self):
+        acc = CaptureAccumulator()
+        chunk = b"data: hello\n\n"
+        out = acc.tee(chunk)
+        self.assertEqual(out, chunk, "流量字节原样立即下发")
+
+    def test_single_aggregate_budget_request_plus_response(self):
+        acc = CaptureAccumulator(budget=100)
+        acc.reserve_request_snapshot(b"x" * 60)   # 请求快照占 60
+        out = acc.tee(b"y" * 60)                  # 响应只剩 40 预算
+        self.assertEqual(out, b"y" * 60, "下发不受预算影响")
+        self.assertEqual(len(acc.captured()), 40, "响应缓冲拿剩余预算")
+        self.assertEqual(acc.total_budgeted(), 100, "聚合（快照+响应）恰好停在预算")
+        self.assertTrue(acc.truncated)
+
+    def test_many_small_messages_cannot_bypass_budget(self):
+        acc = CaptureAccumulator(budget=1000)
+        acc.reserve_request_snapshot(b"m" * 800)
+        for _ in range(500):
+            acc.tee(b"chunk")   # 500×5=2500 字节涌入
+        self.assertLessEqual(len(acc.captured()), 1000)
+
+    def test_utf8_safe_truncation(self):
+        acc = CaptureAccumulator(budget=7)
+        acc.tee("中文中文".encode("utf-8"))  # 每字 3 字节；7 预算切 2 字留 6
+        text = acc.captured().decode("utf-8", "replace")
+        self.assertNotIn("�", text, "预算切点不产生半个字符")
+
+    def test_exception_state_never_affects_passthrough(self):
+        acc = CaptureAccumulator(budget=1)
+        for i in range(10):
+            self.assertEqual(acc.tee(b"abc"), b"abc")
+
+    def test_captured_bytes_counts_budgeted_not_total(self):
+        acc = CaptureAccumulator(budget=10)
+        acc.tee(b"a" * 50)
+        self.assertEqual(acc.total_seen, 50)
+        self.assertEqual(len(acc.captured()), 10)
+
+
+class TestStoreConvergence(unittest.TestCase):
+    """issue #11：单条超大拒收 + append 后总量收敛。"""
+
+    def test_oversized_record_rejected(self):
+        import json as _json
+        from capture import capture_store as cs
+        import tempfile, pathlib, os
+        with tempfile.TemporaryDirectory(dir=os.path.expanduser("~")) as d:
+            cs.append_json({"seed": 1}, os.path.join(d, "cap"))  # 首条建目录+marker
+            d = os.path.join(d, "cap")
+            with self.assertRaises(OSError):
+                cs.append_json({"pad": "x" * (cs.MAX_RECORD_BYTES + 1)}, d)
+
+    def test_post_append_trims_to_store_budget(self):
+        from capture import capture_store as cs
+        import tempfile, os
+        with tempfile.TemporaryDirectory(dir=os.path.expanduser("~")) as d:
+            cs.append_json({"seed": 1}, os.path.join(d, "cap"))  # 建目录
+            d = os.path.join(d, "cap")
+            small = cs.MAX_STORE_BYTES
+            old_total = cs.MAX_STORE_BYTES
+            try:
+                # 两文件各 ~3/4 上限 → append 第二条后总量超限触发收敛
+                cs.MAX_STORE_BYTES = 100
+                cs.MAX_FILE_BYTES = 1000
+                cs.append_json({"d": "y" * 60}, d)
+                # 人造第二个更旧文件使总量超限
+                import time
+                p = pathlib.Path(d) / "2026-01-01.jsonl"
+                p.write_text("z" * 80)
+                st = os.stat(p); os.utime(p, (st.st_atime - 100, st.st_mtime - 100))
+                cs.append_json({"d": "y" * 5}, d)
+                total = sum(f.stat().st_size for f in pathlib.Path(d).glob("*.jsonl*"))
+                self.assertLessEqual(total, cs.MAX_STORE_BYTES + 20,
+                                     "append 后总量收敛（旧文件被 trim）")
+            finally:
+                cs.MAX_STORE_BYTES = old_total