|
| 1 | +#!/usr/bin/env python3 |
| 2 | +"""临时诊断：TRAE 实际上下文上限 + max_mode 字段是否有效（E5 验证）。 |
| 3 | +
|
| 4 | +背景（2026-10-18 实测，结论见 docs/trae-endpoint-verification-2026-10-18.md）： |
| 5 | +- 模型表声明 `context_window_tokens.dev` 为 200K/232K，但实测可到 ~1M token； |
| 6 | +- Trae2api-cn 的 `max_mode` 字段组加与不加，行为完全一致。 |
| 7 | +
|
| 8 | +做法：针在开头、问题在结尾、中间填随机 ASCII。 |
| 9 | +
|
| 10 | + [暗号 <随机8位>] + <N 字符填充> + 「开头提到的暗号是什么？只回答暗号本身」 |
| 11 | +
|
| 12 | +答对暗号 ⇒ 开头（针）与结尾（问题）必须同时在窗口内 ⇒ 窗口横跨整段填充。 |
| 13 | +这比只看「是否报错」更严格：静默截断也能被这一步抓出来。 |
| 14 | +
|
| 15 | +用法： |
| 16 | + # 默认测几个定点，确认声明值与实际值的差距 |
| 17 | + uv run python scripts/probe_trae_context_limit.py |
| 18 | + # 只测一个模型 / 一个长度 |
| 19 | + uv run python scripts/probe_trae_context_limit.py --model glm-5.2 --chars 1450000 |
| 20 | +
|
| 21 | +成本：单次长输入约 60 万–100 万 token。按 0.77 credits/1M 估算， |
| 22 | +默认组合约 5–8 credits。**跑之前先确认凭证额度充足。** |
| 23 | +
|
| 24 | +只读：不改库、不改 src/。 |
| 25 | +""" |
| 26 | + |
| 27 | +from __future__ import annotations |
| 28 | + |
| 29 | +import argparse |
| 30 | +import asyncio |
| 31 | +import contextlib |
| 32 | +import json |
| 33 | +import random |
| 34 | +import string |
| 35 | +import sys |
| 36 | +import uuid |
| 37 | +from pathlib import Path |
| 38 | +from typing import Any |
| 39 | + |
| 40 | +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) |
| 41 | + |
| 42 | +from src.config import load_settings # noqa: E402 |
| 43 | +from src.db.conn import Database # noqa: E402 |
| 44 | +from src.db.crypto import CredentialCipher # noqa: E402 |
| 45 | +from src.engine.sse import iter_frames # noqa: E402 |
| 46 | +from src.provider.trae.client import TraeClient, prepare_body, solo_headers # noqa: E402 |
| 47 | +from src.provider.trae.credential import TraeCredential # noqa: E402 |
| 48 | + |
| 49 | +# Trae2api-cn 用的那套「1M Max」字段。实测证明全部无效，保留只为对照。 |
| 50 | +MAX_MODE_FIELDS: dict[str, Any] = { |
| 51 | + "model_selection_strategy": "max", |
| 52 | + "model_auto_selection": {"strategy": "max"}, |
| 53 | + "context_window_size": 1_000_000, |
| 54 | + "prompt_max_tokens": 936_000, |
| 55 | +} |
| 56 | + |
| 57 | +# 默认定点：声明 232768 的模型测到 1.45M 字符（≈1.03M token）， |
| 58 | +# 再加一个已知超限点确认错误码。 |
| 59 | +DEFAULT_CASES: tuple[tuple[str, int], ...] = ( |
| 60 | + ("seed-code-pro-0430", 900_000), |
| 61 | + ("seed-code-pro-0430", 1_450_000), |
| 62 | + ("seed-code-pro-0430", 1_470_000), |
| 63 | + ("glm-5.2", 1_450_000), |
| 64 | +) |
| 65 | + |
| 66 | + |
| 67 | +def filler(chars: int) -> str: |
| 68 | + """确定性的随机 ASCII 填充（同 seed 保证可复现）。""" |
| 69 | + alphabet = string.ascii_letters + string.digits + " " |
| 70 | + rng = random.Random(20261020) |
| 71 | + chunk = "".join(rng.choices(alphabet, k=4096)) |
| 72 | + return (chunk * ((chars // len(chunk)) + 1))[:chars] |
| 73 | + |
| 74 | + |
| 75 | +def build_prompt(chars: int, needle: str) -> str: |
| 76 | + return ( |
| 77 | + f"暗号是 {needle}，请记住它，稍后会问你。\n\n" |
| 78 | + f"{filler(chars)}\n\n" |
| 79 | + "上面这段长文本开头提到的暗号是什么？只回答暗号本身，不要解释。" |
| 80 | + ) |
| 81 | + |
| 82 | + |
| 83 | +def load_credential() -> TraeCredential: |
| 84 | + settings = load_settings() |
| 85 | + db = Database(str(Path(settings.data_dir) / "coding2api.sqlite3")) |
| 86 | + cipher = CredentialCipher(settings.app_secret) |
| 87 | + row = db.connect().execute( |
| 88 | + "SELECT data_enc FROM credentials WHERE provider='trae' AND enabled=1 " |
| 89 | + "ORDER BY health DESC LIMIT 1", |
| 90 | + ).fetchone() |
| 91 | + return TraeCredential.from_dict(json.loads(cipher.decrypt(row["data_enc"]))) |
| 92 | + |
| 93 | + |
| 94 | +def build_client() -> TraeClient: |
| 95 | + """长输入体积大：默认 write=10s 会在上传阶段 WriteTimeout。 |
| 96 | +
|
| 97 | + 本脚本单独放宽写/连超时（read 仍 None，长流必须自己兜超时）。 |
| 98 | + """ |
| 99 | + import httpx |
| 100 | + |
| 101 | + return TraeClient(stream_client=httpx.AsyncClient( |
| 102 | + timeout=httpx.Timeout(connect=30.0, read=None, write=300.0, pool=30.0), |
| 103 | + follow_redirects=True, |
| 104 | + )) |
| 105 | + |
| 106 | + |
| 107 | +async def ask(client: TraeClient, credential: TraeCredential, model: str, |
| 108 | + chars: int, extra: dict[str, Any] | None = None) -> tuple[bool, dict]: |
| 109 | + """发一次带暗号的长输入，返回 (是否答对暗号, 观测信息)。""" |
| 110 | + needle = uuid.uuid4().hex[:8] |
| 111 | + body = prepare_body({ |
| 112 | + "messages": [{"role": "user", "content": build_prompt(chars, needle)}], |
| 113 | + "max_tokens": 32, |
| 114 | + }, model) |
| 115 | + if extra: |
| 116 | + body.update(extra) |
| 117 | + text, errors, usage = "", [], {} |
| 118 | + async with client._stream().stream( |
| 119 | + "POST", f"{client.agent_host}/api/agent/v3/llm_utils_chat", |
| 120 | + json=body, headers=solo_headers(credential), |
| 121 | + ) as response: |
| 122 | + if response.status_code >= 400: |
| 123 | + raw = await response.aread() |
| 124 | + return False, {"http": response.status_code, |
| 125 | + "body": raw.decode("utf-8", "replace")[:200]} |
| 126 | + async for frame in iter_frames(response.aiter_bytes()): |
| 127 | + name = frame.event.strip() |
| 128 | + if name == "error": |
| 129 | + errors.append(frame.data[:160]) |
| 130 | + elif name == "token_usage": |
| 131 | + usage = json.loads(frame.data) |
| 132 | + elif name == "output": |
| 133 | + with contextlib.suppress(ValueError): |
| 134 | + text += json.loads(frame.data).get("response") or "" |
| 135 | + return needle in text, { |
| 136 | + "needle": needle, |
| 137 | + "text": text.strip()[:40], |
| 138 | + "prompt_tokens": usage.get("prompt_tokens"), |
| 139 | + "errors": errors, |
| 140 | + } |
| 141 | + |
| 142 | + |
| 143 | +async def main() -> None: |
| 144 | + parser = argparse.ArgumentParser() |
| 145 | + parser.add_argument("--model", default=None, help="只测该模型（配 --chars）") |
| 146 | + parser.add_argument("--chars", type=int, default=None, help="只测该字符数") |
| 147 | + parser.add_argument("--compare-max-mode", action="store_true", |
| 148 | + help="对每个点额外发一次带 max_mode 字段组的请求做对照") |
| 149 | + args = parser.parse_args() |
| 150 | + |
| 151 | + cases = DEFAULT_CASES |
| 152 | + if args.model and args.chars: |
| 153 | + cases = ((args.model, args.chars),) |
| 154 | + |
| 155 | + credential = load_credential() |
| 156 | + client = build_client() |
| 157 | + try: |
| 158 | + for model, chars in cases: |
| 159 | + ok, info = await ask(client, credential, model, chars) |
| 160 | + tag = "✅ 读到" if ok else "❌ 未读到" |
| 161 | + print(f"{model:22s} {chars:>9} chars -> {tag} " |
| 162 | + f"prompt_tokens={info.get('prompt_tokens')} " |
| 163 | + f"{info.get('errors') or ''}") |
| 164 | + if not args.compare_max_mode: |
| 165 | + continue |
| 166 | + ok_max, info_max = await ask(client, credential, model, chars, |
| 167 | + MAX_MODE_FIELDS) |
| 168 | + tag_max = "✅ 读到" if ok_max else "❌ 未读到" |
| 169 | + print(f"{' └ max_mode 字段组':22s} {chars:>9} chars -> {tag_max} " |
| 170 | + f"prompt_tokens={info_max.get('prompt_tokens')} " |
| 171 | + f"{info_max.get('errors') or ''}") |
| 172 | + finally: |
| 173 | + await client.aclose() |
| 174 | + |
| 175 | + |
| 176 | +if __name__ == "__main__": |
| 177 | + asyncio.run(main()) |
0 commit comments