tests/test_upstream_errors.py

"""上游错误码建议与上下文裁剪的验证。

背景：上游对超长上下文会返回模糊错误（例如错误码 1096）。gemini-webapi 只映射了
1013/1037/1050/1052/1060，其余走兜底分支并附一句"可能是临时 Google 服务问题"——那是库
自己的猜测，对 1096 并不准确（1096 属于后端/会话层面的临时故障，常见诱因是上下文过长）。
所以本服务自己补了一张"怎么办"的表，并在超长时裁剪历史。
"""

import os
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

os.environ["SECURE_1PSID"] = "PSID-ERR"
os.environ["SECURE_1PSIDTS"] = "TS-ERR"
os.environ["GEMINI_COOKIE_PATH"] = tempfile.mkdtemp(prefix="gs-err-")

import main as app_main  # noqa: E402

PASS, FAIL = [], []


def check(name, cond, extra=""):
	(PASS if cond else FAIL).append(name)
	print(("[OK]   " if cond else "[FAIL] ") + name + (f"  {extra}" if extra and not cond else ""))


def main():
	# ---------------------------------------------------------------- 错误码建议
	code_1096 = (
		"APIError: Failed to generate contents (stream). Unknown API error code: 1096. "
		"This might be a temporary Google service issue."
	)
	hint = app_main.describe_upstream_error(code_1096)
	check("1 认得 1096", "错误码 1096" in hint, hint)
	check("2 并给出可执行建议（上下文/新对话）", "上下文过长" in hint and "新对话" in hint, hint)
	check("3 明确提示不要连续重试", "不要连续重试" in hint, hint)

	check("4 认得 1037（配额）", "配额" in app_main.describe_upstream_error("error code: 1037"))
	check("5 认得 1060（IP 被限）", "IP" in app_main.describe_upstream_error("error code: 1060"))
	check("6 认得 1050（换模型）", "模型" in app_main.describe_upstream_error("error code: 1050"))

	unknown = app_main.describe_upstream_error("APIError: ... error code: 9999 ...")
	check("7 未知错误码给出兜底建议", "9999" in unknown and "issue" in unknown, unknown)
	check("8 没有错误码时不瞎猜", app_main.describe_upstream_error("UsageLimitExceededError: 配额用完") == "")
	check("9 空输入不报错", app_main.describe_upstream_error("") == "")

	# ---------------------------------------------------------------- 上下文裁剪
	check("10 默认上限为正数", app_main.MAX_CONTEXT_CHARS > 0, str(app_main.MAX_CONTEXT_CHARS))

	system = app_main.Message(role="system", content="S" * 100)
	history = [app_main.Message(role="user", content="X" * 1000) for _ in range(100)]
	kept = app_main.trim_messages([system, *history])
	check("11 超长时确实裁剪了", len(kept) < 101, str(len(kept)))
	check("12 system 消息始终保留", kept[0].role == "system", kept[0].role)
	check("13 最后一条消息始终保留", kept[-1] is history[-1])
	check("14 裁剪后总量不超上限（最后一条例外）", sum(len(str(m.content)) for m in kept) <= app_main.MAX_CONTEXT_CHARS + 1000)

	short = [app_main.Message(role="user", content="hi")]
	check("15 短对话原样返回", app_main.trim_messages(short) == short)
	check("16 空列表原样返回", app_main.trim_messages([]) == [])

	# 单条就超限时也要保留（否则请求没意义）
	huge = [app_main.Message(role="user", content="Y" * (app_main.MAX_CONTEXT_CHARS * 2))]
	check("17 单条超限也保留", len(app_main.trim_messages(huge)) == 1)

	# 关掉上限时不裁剪
	saved = app_main.MAX_CONTEXT_CHARS
	try:
		app_main.MAX_CONTEXT_CHARS = 0
		check("18 上限为 0 时不裁剪", len(app_main.trim_messages([system, *history])) == 101)
	finally:
		app_main.MAX_CONTEXT_CHARS = saved

	# ---------------------------------------------------------------- 接线：prepare_conversation 会裁剪
	saved = app_main.MAX_CONTEXT_CHARS
	try:
		app_main.MAX_CONTEXT_CHARS = 2000
		conversation, _ = app_main.prepare_conversation([system, *history])
		check("19 prepare_conversation 应用了裁剪", len(conversation) < 100000, str(len(conversation)))
		check("20 系统指令仍在最前面", conversation.startswith("System: "), repr(conversation[:30]))
	finally:
		app_main.MAX_CONTEXT_CHARS = saved

	print(f"\n通过 {len(PASS)} / {len(PASS) + len(FAIL)}")
	if FAIL:
		print("失败项：")
		for name in FAIL:
			print("  -", name)
	return 1 if FAIL else 0


if __name__ == "__main__":
	sys.exit(main())