diff --git a/capture/ai_capture_addon.py b/capture/ai_capture_addon.py index 4a6e41c..0c397fb 100644 --- a/capture/ai_capture_addon.py +++ b/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] diff --git a/capture/capture_store.py b/capture/capture_store.py index 050bbd8..61b0a79 100644 --- a/capture/capture_store.py +++ b/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) diff --git a/tests/test_ai_capture_addon.py b/tests/test_ai_capture_addon.py index b08c607..9fabbc1 100644 --- a/tests/test_ai_capture_addon.py +++ b/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": []}, diff --git a/tests/test_capture_accumulator.py b/tests/test_capture_accumulator.py new file mode 100644 index 0000000..0c0b5f7 --- /dev/null +++ b/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