Spaces:

zlcy
/

abc

Paused

abc / test /test_image_upstream_note.py
xiaolongmr
init
85270f9 2 days ago
38.3 kB
"""超时失败时把上游 assistant 原话带回来（`upstream_note`）。
背景：生图超时的固定文案只会说「大概是被 GPT 拦了」，用户没法自己判断。
上游被拦时往往只回一段文字、既没调生图工具也没有图片资产，这段文字是唯一线索。
更关键的是：上游在 SSE 结束时就已经明说 `tool_invoked=false`（本轮没调用生图工具），
这种回合永远等不到图。再按 image gen 轮询只会把整段超时预算耗光（线上实测 82~97 秒），
所以这里要求「先做一次短探测取回原话、立刻收工」，轮询只作为探测失败时的兜底。
"""

from __future__ import annotations

import pytest

from services.image_task_service import _public_task
from services.openai_backend_api import (
    ImagePollTimeoutError,
    ImageResolutionResult,
    OpenAIBackendAPI,
)


def _backend_without_init() -> OpenAIBackendAPI:
    """只借方法，不触发网络/账号初始化。"""
    return OpenAIBackendAPI.__new__(OpenAIBackendAPI)


def _assistant(text: str, *, create_time: float, hidden: bool = False) -> dict:
    return {
        "message": {
            "author": {"role": "assistant"},
            "content": {"content_type": "text", "parts": [text]},
            "create_time": create_time,
            "metadata": {"is_visually_hidden_from_conversation": True} if hidden else {},
        }
    }


def test_latest_assistant_text_picks_newest_visible_message() -> None:
    api = _backend_without_init()
    conversation = {
        "mapping": {
            "u1": {
                "message": {
                    "author": {"role": "user"},
                    "content": {"content_type": "text", "parts": ["生成一张手机壁纸"]},
                    "create_time": 1,
                    "metadata": {},
                }
            },
            "a1": _assistant("第一次回复：正在生成。", create_time=2),
            "a2": _assistant("抱歉，我无法生成这张图片。", create_time=3),
        }
    }
    assert api._latest_assistant_text(conversation) == "抱歉，我无法生成这张图片。"


def test_latest_assistant_text_ignores_hidden_and_non_assistant() -> None:
    api = _backend_without_init()
    conversation = {
        "mapping": {
            "a1": _assistant("对用户可见的原话", create_time=2),
            "a2": _assistant("内部思考，不该给用户看", create_time=9, hidden=True),
            "t1": {
                "message": {
                    "author": {"role": "tool"},
                    "recipient": "image_gen",
                    "content": {"content_type": "text", "parts": ["工具输出"]},
                    "create_time": 10,
                    "metadata": {},
                }
            },
        }
    }
    assert api._latest_assistant_text(conversation) == "对用户可见的原话"


@pytest.mark.parametrize("payload", [None, {}, {"mapping": None}, {"mapping": {}}])
def test_latest_assistant_text_tolerates_empty_documents(payload) -> None:
    api = _backend_without_init()
    assert api._latest_assistant_text(payload) == ""


def test_poll_timeout_carries_upstream_text(monkeypatch) -> None:
    """轮询跑完预算仍没等到图片时，异常必须把最后一句话带出来。"""
    from services.config import config

    # 默认初始等待 10s（+抖动）会直接把小预算吃光，一次会话文档都拿不到；
    # 这里压掉等待和轮询间隔，让循环真的发起一轮请求。
    monkeypatch.setattr(type(config), "image_poll_initial_wait_secs", property(lambda self: 0.0))
    monkeypatch.setattr(type(config), "image_poll_interval_secs", property(lambda self: 0.05))

    api = _backend_without_init()
    conversation = {
        "mapping": {
            "a1": _assistant("我不能生成含有真实人物的图片。", create_time=5),
        }
    }

    def _fake_get_conversation(conversation_id: str) -> dict:
        return conversation

    monkeypatch.setattr(api, "_get_conversation", _fake_get_conversation, raising=False)

    with pytest.raises(ImagePollTimeoutError) as excinfo:
        api._poll_image_results("conv-upstream-text", timeout_secs=0.3)

    error = excinfo.value
    assert error.conversation_id == "conv-upstream-text"
    assert error.upstream_text == "我不能生成含有真实人物的图片。"


def test_public_task_exposes_upstream_note() -> None:
    item = _public_task(
        {
            "id": "task-1",
            "status": "error",
            "mode": "edit",
            "model": "gpt-image-2",
            "conversation_id": "cid-1",
            "error": "ChatGPT生图超时 (已等待 91 秒）。",
            "upstream_note": "抱歉，我现在无法生成这个内容。",
        }
    )
    assert item["upstream_note"] == "抱歉，我现在无法生成这个内容。"
    assert item["conversation_id"] == "cid-1"


def test_public_task_omits_upstream_note_when_absent() -> None:
    item = _public_task({"id": "task-2", "status": "error", "mode": "generate", "model": "gpt-image-2"})
    assert "upstream_note" not in item


def test_public_task_exposes_can_continue_waiting_only_when_set() -> None:
    terminal = _public_task(
        {"id": "task-3", "status": "error", "mode": "generate", "model": "gpt-image-2",
         "conversation_id": "cid-1", "can_continue_waiting": False}
    )
    assert terminal["can_continue_waiting"] is False
    # 没显式下发就不出现 —— 前台继续按旧的「有会话 id 就能继续等」规则走。
    legacy = _public_task({"id": "task-4", "status": "error", "mode": "generate", "model": "gpt-image-2"})
    assert "can_continue_waiting" not in legacy


def test_public_task_exposes_can_view_official_only_when_denied() -> None:
    """只有后端明确说「官网没有可核对的东西」时才下发 false（配额耗尽），前台据此收按钮。"""
    denied = _public_task(
        {"id": "task-5", "status": "error", "mode": "edit", "model": "gpt-image-2",
         "conversation_id": "cid-1", "error": QUOTA_TEXT, "can_view_official": False}
    )
    assert denied["can_view_official"] is False
    # 老任务 / 普通失败（超时、内容拒绝）没有这个字段 -> 前台保持原样给入口。
    legacy = _public_task(
        {"id": "task-6", "status": "error", "mode": "edit", "model": "gpt-image-2",
         "conversation_id": "cid-2", "error": "生图超时 (已等待 91 秒）。"}
    )
    assert "can_view_official" not in legacy


def test_latest_assistant_text_respects_cutoff_for_memory_turns() -> None:
    """多轮会话里不能把上一轮的回答当成本轮的结论。"""
    api = _backend_without_init()
    conversation = {
        "mapping": {
            "a_old": _assistant("上一轮的回答：图已生成。", create_time=100),
            "u_new": {
                "message": {
                    "author": {"role": "user"},
                    "content": {"content_type": "text", "parts": ["再来一张"]},
                    "create_time": 200,
                    "metadata": {},
                }
            },
        }
    }
    assert api._latest_assistant_text(conversation, min_create_time=190) == ""
    assert api._latest_assistant_text(conversation) == "上一轮的回答：图已生成。"


def test_latest_assistant_text_respects_after_message_id() -> None:
    api = _backend_without_init()
    conversation = {
        "mapping": {
            "u_prev": {
                "message": {
                    "author": {"role": "user"},
                    "content": {"content_type": "text", "parts": ["上一轮的提问"]},
                    "create_time": 90,
                    "metadata": {},
                }
            },
            "a_old": {**_assistant("上一轮的回答", create_time=100), "parent": "u_prev"},
            "u_old": {
                "message": {
                    "author": {"role": "user"},
                    "content": {"content_type": "text", "parts": ["这一次的提问"]},
                    "create_time": 200,
                    "metadata": {},
                }
            },
            "a_new": {**_assistant("这一轮的原话", create_time=300), "parent": "u_old"},
        }
    }
    assert api._latest_assistant_text(conversation, after_message_id="u_old") == "这一轮的原话"
    # 界标不在文档里（且没有时间界）时宁可不给，也别把上一轮的回答顶上来。
    assert api._latest_assistant_text(conversation, after_message_id="u_missing") == ""


def _terminal_doc(text: str = "请上传需要修改的图片。", *, create_time: float = 1791893410.45,
                  recipient: str = "all", content_type: str = "text", status: str = "finished_successfully",
                  is_complete: bool = True, children: list | None = None, hidden: bool = False) -> dict:
    """复刻真机取证 `6ace1f9f-…` 的会话形状：user → reasoning_recap → 文本叶子。"""
    return {
        "current_node": "leaf",
        "mapping": {
            "u1": {
                "parent": "root",
                "children": ["recap"],
                "message": {
                    "author": {"role": "user"},
                    "recipient": "all",
                    "content": {"content_type": "text", "parts": ["把图片上的文字改成英文"]},
                    "create_time": create_time - 8,
                    "status": "finished_successfully",
                    "metadata": {},
                },
            },
            "recap": {
                "parent": "u1",
                "children": ["leaf"],
                "message": {
                    "author": {"role": "assistant"},
                    "recipient": "all",
                    "content": {"content_type": "reasoning_recap", "parts": ["内部思考"]},
                    "create_time": create_time - 0.01,
                    "status": "finished_successfully",
                    "metadata": {"model_slug": "gpt-5-6-t-mini"},
                },
            },
            "leaf": {
                "parent": "recap",
                "children": children or [],
                "message": {
                    "author": {"role": "assistant"},
                    "recipient": recipient,
                    "content": {"content_type": content_type, "parts": [text]},
                    "create_time": create_time,
                    "status": status,
                    "metadata": {
                        "is_complete": is_complete,
                        "finish_details": {"type": "stop", "stop_tokens": [200002]},
                        **({"is_visually_hidden_from_conversation": True} if hidden else {}),
                    },
                },
            },
        },
    }


def test_terminal_text_answer_detects_real_shape() -> None:
    """真机形状（文本叶子 + is_complete + recipient=all）必须被认成终局。"""
    api = _backend_without_init()
    assert api._terminal_text_answer(_terminal_doc()) == "请上传需要修改的图片。"


@pytest.mark.parametrize(
    "override",
    [
        {"recipient": "image_gen"},          # 工具调用中间态
        {"content_type": "reasoning_recap"},  # 内部思考，不是给用户的回答
        {"status": "in_progress"},            # 还没结束
        {"is_complete": False},
        {"children": ["tool1"]},              # 还有子节点，回合在往下走
        {"hidden": True},                     # 对用户不可见
    ],
)
def test_terminal_text_answer_rejects_non_terminal_shapes(override) -> None:
    api = _backend_without_init()
    doc = _terminal_doc(**override)
    if override.get("is_complete") is False:
        doc["mapping"]["leaf"]["message"]["metadata"]["finish_details"] = {}
    assert api._terminal_text_answer(doc) == ""


def test_terminal_text_answer_ignores_previous_turn_leaf() -> None:
    """续写场景：叶子还是上一轮的终局文本时不能当成这一轮的回答。"""
    api = _backend_without_init()
    doc = _terminal_doc(create_time=100)
    assert api._terminal_text_answer(doc, min_create_time=190) == ""
    assert api._terminal_text_answer(doc, after_message_id="leaf") == ""


@pytest.mark.parametrize("payload", [None, {}, {"mapping": None}, {"mapping": {}, "current_node": "x"}])
def test_terminal_text_answer_tolerates_empty_documents(payload) -> None:
    api = _backend_without_init()
    assert api._terminal_text_answer(payload) == ""


REFUSAL_TEXT = "非常抱歉，生成的图片可能违反了关于与第三方内容相似性的防护限制。如果你认为此判断有误，请重试或修改提示语。"


def _refusal_doc(text: str = REFUSAL_TEXT) -> dict:
    """真机取证 `6ace48ab-…`（让 GPT 生成迪士尼壁纸被内容策略拦下）的会话形状：
    user → assistant(`recipient=<工具id>`, `content_type=code`, 生图工具调用) → assistant 文本叶子。
    ⚠️ 与 `_terminal_doc` 的关键差别：文本叶子的 ``metadata`` 里**没有**
    `is_complete` / `finish_details`，只有 ``is_error: true``（上游用「错误/拒绝」收尾）。
    """
    return {
        "current_node": "leaf",
        "mapping": {
            "u1": {
                "parent": "root",
                "children": ["tool"],
                "message": {
                    "author": {"role": "user"},
                    "recipient": "all",
                    "content": {"content_type": "text", "parts": ["生成一张迪士尼手机壁纸"]},
                    "create_time": 1791903914.23,
                    "status": "finished_successfully",
                    "metadata": {"message_source": None},
                },
            },
            "tool": {
                "parent": "u1",
                "children": ["leaf"],
                "message": {
                    "author": {"role": "assistant"},
                    "recipient": "t2uay3k.sj1i4kz",
                    "content": {"content_type": "code", "parts": [""]},
                    "create_time": 1791903917.13,
                    "status": "finished_successfully",
                    "metadata": {"can_save": True, "parent_id": "u1"},
                },
            },
            "leaf": {
                "parent": "tool",
                "children": [],
                "message": {
                    "author": {"role": "assistant"},
                    "recipient": "all",
                    "content": {"content_type": "text", "parts": [text]},
                    "create_time": 1791903935.45,
                    "status": "finished_successfully",
                    "metadata": {
                        "can_save": True,
                        "is_error": True,
                        "model_switcher_deny": False,
                        "permissions": {},
                        "turn_exchange_id": "t2uay3k",
                        "working_turn_id": "t2uay3k",
                    },
                },
            },
        },
    }


def test_terminal_text_answer_detects_refusal_error_leaf() -> None:
    """内容策略拒绝的叶子没有 is_complete/finish_details，只有 is_error=true，同样要认成终局。
    真机代价：这条 21 秒就出现的拒绝文案被判成「普通超时」，白轮询 86.8s / 91.8s / 30s 三次。
    """
    api = _backend_without_init()
    assert api._terminal_text_answer(_refusal_doc()) == REFUSAL_TEXT
    # 放宽只针对上游的显式错误标记：去掉 is_error 后（既无 is_complete 也无 finish_details）
    # 不能再认成终局，否则任意中途文本都会被当成结论。
    doc = _refusal_doc()
    doc["mapping"]["leaf"]["message"]["metadata"].pop("is_error")
    assert api._terminal_text_answer(doc) == ""


def test_poll_raises_text_only_answer_for_refusal_leaf(monkeypatch) -> None:
    """拒绝文案要在**第一次轮询**就带回去，不跑满预算。"""
    from services.config import config
    from services.openai_backend_api import ImageTextOnlyAnswer

    monkeypatch.setattr(type(config), "image_poll_initial_wait_secs", property(lambda self: 0.0))
    monkeypatch.setattr(type(config), "image_poll_interval_secs", property(lambda self: 5.0))

    api = _backend_without_init()
    calls = {"count": 0}

    def _fake_get_conversation(conversation_id: str) -> dict:
        calls["count"] += 1
        return _refusal_doc()

    monkeypatch.setattr(api, "_get_conversation", _fake_get_conversation, raising=False)

    with pytest.raises(ImageTextOnlyAnswer) as excinfo:
        api._poll_image_results("conv-refusal", timeout_secs=120.0)

    assert excinfo.value.text == REFUSAL_TEXT
    assert excinfo.value.conversation_id == "conv-refusal"
    assert calls["count"] == 1


def test_poll_detects_terminal_text_at_timeout_fallback(monkeypatch) -> None:
    """预算在两次轮询之间耗尽时，收尾补取的那一次也要认出终局文本，而不是报超时。
    真机形状：拒绝文案在 21s 出现，而 30s 预算的轮询停在 20s —— 差一点点就漏判。
    """
    import time as _time

    from services.config import config
    from services.openai_backend_api import ImageTextOnlyAnswer

    monkeypatch.setattr(type(config), "image_poll_initial_wait_secs", property(lambda self: 0.0))
    monkeypatch.setattr(type(config), "image_poll_interval_secs", property(lambda self: 0.05))

    api = _backend_without_init()
    monkeypatch.setattr(api, "_get_conversation", lambda cid: _refusal_doc(), raising=False)

    real_terminal = api._terminal_text_answer
    started = _time.time()

    def _late_terminal(conversation, **kwargs):
        # 模拟「文案还没落地」：预算耗尽前的每一次轮询都判不出来，只有收尾补取那一次能判出来。
        if _time.time() - started < 0.3:
            return ""
        return real_terminal(conversation, **kwargs)

    monkeypatch.setattr(api, "_terminal_text_answer", _late_terminal, raising=False)

    with pytest.raises(ImageTextOnlyAnswer) as excinfo:
        api._poll_image_results("conv-refusal-slow", timeout_secs=0.3)

    assert excinfo.value.text == REFUSAL_TEXT
    assert excinfo.value.conversation_id == "conv-refusal-slow"


def test_pool_turns_refusal_into_terminal_error(monkeypatch) -> None:
    """拒绝文案经号池层要变成「终局错误」：带 conversation_id + terminal_no_image。
    这样前台同时拿到两件事：卡片上有「查看官网」入口，且不会出现没用的「继续等待」按钮。
    """
    from services.protocol import conversation as conv

    class _StubBackend:
        def __init__(self, **kwargs) -> None:
            self.progress_callback = None

        def close(self) -> None:
            return None

    monkeypatch.setattr(conv.account_service, "get_available_access_token",
                        lambda *args, **kwargs: "tok-refusal", raising=False)
    monkeypatch.setattr(conv.account_service, "get_account",
                        lambda token: {"email": "refusal@duck.com"}, raising=False)
    monkeypatch.setattr(conv, "record_image_failure", lambda *args, **kwargs: None, raising=False)
    monkeypatch.setattr(conv, "OpenAIBackendAPI", _StubBackend, raising=False)
    # ⚠️ finally 里的会话隐藏/取证登记走的是**真实台账**（本机 config 开着取证）：
    # 不拦住 register_conversation，假 conversation_id="cid-refusal" 会写进
    # data/share_evidence.json，worker 之后反复尝试给这条假会话创建分享链接。
    monkeypatch.setattr(
        "services.share_evidence_service.share_evidence_service.register_conversation",
        lambda **kwargs: None,
    )

    def _fake_stream(backend, request, index, total, **kwargs):
        yield conv.ImageOutput(
            kind="message",
            model=request.model,
            index=index,
            total=total,
            text=REFUSAL_TEXT,
            conversation_id="cid-refusal",
        )

    monkeypatch.setattr(conv, "stream_image_outputs", _fake_stream)

    with pytest.raises(conv.ImageGenerationError) as excinfo:
        list(conv._stream_image_outputs_with_pool(_request()))

    exc = excinfo.value
    assert exc.terminal_no_image is True
    assert exc.conversation_id == "cid-refusal"
    assert exc.code == "content_policy_violation"
    assert exc.can_view_official is True
    assert str(exc) == REFUSAL_TEXT


def test_quota_failure_task_hides_official_entry(tmp_path) -> None:
    """端到端：号池层抛出的「配额用完」错误要落进任务记录，并下发 can_view_official=False。
    只测单个函数会漏掉 `_update_task(**{...})` 那一层写入的拼写错误 —— 那会让整个功能静默失效
    （前端永远拿到 undefined = 照旧显示按钮）。
    """
    import time

    from services.image_task_service import ImageTaskService
    from services.protocol.conversation import ImageGenerationError

    def handler(_payload):
        raise ImageGenerationError(
            QUOTA_TEXT,
            status_code=429,
            error_type="rate_limit_error",
            code="image_quota_exhausted",
            conversation_id="cid-quota",
            terminal_no_image=True,
            can_view_official=False,
        )

    service = ImageTaskService(
        tmp_path / "image_tasks.json",
        generation_handler=handler,
        edit_handler=handler,
        retention_days_getter=lambda: 30,
    )
    owner = {"id": "owner-1", "name": "Owner", "role": "admin"}
    service.submit_generation(
        owner,
        client_task_id="quota-task",
        prompt="cat",
        model="gpt-image-2",
        size=None,
        base_url="http://local.test",
    )

    deadline = time.time() + 3
    item = None
    while time.time() < deadline:
        item = (service.list_tasks(owner, ["quota-task"]).get("items") or [None])[0]
        if item and item.get("status") == "error":
            break
        time.sleep(0.02)

    assert item and item.get("status") == "error", item
    assert item["can_view_official"] is False
    # 会话 id 仍然落库（排查/日志要用），只是前台不给入口。
    assert item["conversation_id"] == "cid-quota"


def test_pool_switches_account_when_quota_exhausted(monkeypatch) -> None:
    """配额用完的账号要当场移出号池 + 本轮换号重试：别把上游账号额度问题变成用户的失败。
    只把错误标成 no-official 是不够的：本地 quota 对免费号常常是「未知」（不计数），
    不主动移出池子的话这个账号会被后续每个请求反复挑中，每个用户都吃一次同样的报错。
    """
    from unittest import mock

    from services.protocol import conversation as conv

    switched_tokens: list[str] = []

    class _StubBackend:
        def __init__(self, access_token: str = "", **kwargs) -> None:
            self.access_token = access_token
            self.progress_callback = None

        def close(self) -> None:
            return None

    def _fake_get_token(*args, excluded_tokens=None, **kwargs):
        return "tok-alive" if "tok-dead" in set(excluded_tokens or set()) else "tok-dead"

    monkeypatch.setattr(conv.account_service, "get_available_access_token", _fake_get_token, raising=False)
    monkeypatch.setattr(
        conv.account_service,
        "get_account",
        lambda token: {"email": f"{token}@duck.com"},
        raising=False,
    )
    monkeypatch.setattr(conv, "record_image_failure", lambda *args, **kwargs: None, raising=False)
    monkeypatch.setattr(conv, "OpenAIBackendAPI", _StubBackend, raising=False)
    monkeypatch.setattr(
        "services.share_evidence_service.share_evidence_service.register_conversation",
        lambda **kwargs: None,
    )
    mark_quota_exhausted = mock.MagicMock()
    monkeypatch.setattr(conv.account_service, "mark_image_quota_exhausted", mark_quota_exhausted, raising=False)

    def _fake_stream(backend, request, index, total, **kwargs):
        if backend.access_token == "tok-dead":
            # 上游只回文字 = 终局结论：本账号配额用完。
            yield conv.ImageOutput(
                kind="message",
                model=request.model,
                index=index,
                total=total,
                text=QUOTA_TEXT,
                conversation_id="cid-dead",
            )
            return
        switched_tokens.append(backend.access_token)
        yield conv.ImageOutput(
            kind="result",
            model=request.model,
            index=index,
            total=total,
            data=[{"b64_json": "x"}],
            conversation_id="cid-alive",
        )

    monkeypatch.setattr(conv, "stream_image_outputs", _fake_stream)

    outputs = list(conv._stream_image_outputs_with_pool(_request()))

    # 换号成功：用户拿到图，而不是「Free 套餐额度用完」的报错。
    assert switched_tokens == ["tok-alive"]
    assert [item.kind for item in outputs] == ["result"]
    # 额度用完的账号被当场移出可用池（否则它会被后续每个请求反复挑中）。
    mark_quota_exhausted.assert_called_once()
    assert mark_quota_exhausted.call_args.args[0] == "tok-dead"


def test_pool_reports_upstream_quota_text_when_switch_budget_runs_out(monkeypatch) -> None:
    """连着撞到额度用尽的账号、换号预算用完后，出站的是**上游原话**。
    站长按这句原话配了「报错文案映射」（→ 大白话），所以出站文案不能被内部改写。
    """
    from unittest import mock

    from services.protocol import conversation as conv

    class _StubBackend:
        def __init__(self, access_token: str = "", **kwargs) -> None:
            self.access_token = access_token
            self.progress_callback = None

        def close(self) -> None:
            return None

    def _fake_get_token(*args, excluded_tokens=None, **kwargs):
        excluded = set(excluded_tokens or set())
        if len(excluded) >= 4:
            raise RuntimeError("no available image quota")
        return f"tok-{len(excluded) + 1}"

    monkeypatch.setattr(conv.account_service, "get_available_access_token", _fake_get_token, raising=False)
    monkeypatch.setattr(conv.account_service, "get_account", lambda token: {"email": "dead@duck.com"}, raising=False)
    monkeypatch.setattr(conv, "record_image_failure", lambda *args, **kwargs: None, raising=False)
    monkeypatch.setattr(conv, "OpenAIBackendAPI", _StubBackend, raising=False)
    monkeypatch.setattr(
        "services.share_evidence_service.share_evidence_service.register_conversation",
        lambda **kwargs: None,
    )
    monkeypatch.setattr(conv.account_service, "mark_image_quota_exhausted", mock.MagicMock(), raising=False)

    def _fake_stream(backend, request, index, total, **kwargs):
        yield conv.ImageOutput(
            kind="message",
            model=request.model,
            index=index,
            total=total,
            text=QUOTA_TEXT,
            conversation_id="cid-dead",
        )

    monkeypatch.setattr(conv, "stream_image_outputs", _fake_stream)

    with pytest.raises(conv.ImageGenerationError) as excinfo:
        list(conv._stream_image_outputs_with_pool(_request()))

    exc = excinfo.value
    assert exc.code == "image_quota_exhausted"
    assert exc.can_view_official is False
    # 换号次数有预算（至少 2 次尝试 = 至少换 1 次），不能把整池挨个 20s 试一遍把用户拖住。
    assert exc.account_email == "dead@duck.com"
    assert str(exc) == QUOTA_TEXT


def test_pool_reports_internal_quota_error_when_pool_drains(monkeypatch) -> None:
    """号池被掏空时要冒出内部英文 `no available image quota`——站长的报错文案规则按它配。
    ⚠️ 这串英文是「对外的内部哨兵」：api/support.py 按它回 429、log_service 按它标
    insufficient_quota、站长的「报错文案映射」按它换大白话。所以池空时**不能**用
    上游原话（那是另一条规则的事）把它替掉，否则规则永远匹配不上。
    """
    from unittest import mock

    from services.protocol import conversation as conv

    class _StubBackend:
        def __init__(self, access_token: str = "", **kwargs) -> None:
            self.access_token = access_token
            self.progress_callback = None

        def close(self) -> None:
            return None

    def _fake_get_token(*args, excluded_tokens=None, **kwargs):
        # 只有 1 个账号：它被移出池子后，号池就空了。
        if set(excluded_tokens or set()):
            raise RuntimeError("no available image quota")
        return "tok-only"

    monkeypatch.setattr(conv.account_service, "get_available_access_token", _fake_get_token, raising=False)
    monkeypatch.setattr(conv.account_service, "get_account", lambda token: {"email": "only@duck.com"}, raising=False)
    monkeypatch.setattr(conv, "record_image_failure", lambda *args, **kwargs: None, raising=False)
    monkeypatch.setattr(conv, "OpenAIBackendAPI", _StubBackend, raising=False)
    monkeypatch.setattr(
        "services.share_evidence_service.share_evidence_service.register_conversation",
        lambda **kwargs: None,
    )
    mark_quota_exhausted = mock.MagicMock()
    monkeypatch.setattr(conv.account_service, "mark_image_quota_exhausted", mark_quota_exhausted, raising=False)

    def _fake_stream(backend, request, index, total, **kwargs):
        yield conv.ImageOutput(
            kind="message",
            model=request.model,
            index=index,
            total=total,
            text=QUOTA_TEXT,
            conversation_id="cid-only",
        )

    monkeypatch.setattr(conv, "stream_image_outputs", _fake_stream)

    with pytest.raises(conv.ImageGenerationError) as excinfo:
        list(conv._stream_image_outputs_with_pool(_request()))

    assert "no available image quota" in str(excinfo.value)
    mark_quota_exhausted.assert_called_once()


def test_pool_marks_quota_notice_as_no_official_evidence(monkeypatch) -> None:
    """账号生图配额用完时不能再挂「查看官网对话」入口。
    真机报错（HF 日志，账号 defy-scruffy-armed@duck.com，20s 失败）：
    「你已达到 Free 套餐的图像生成请求上限。上限将在 23小时 后重置，届时可创建更多图像。」
    上游这一轮**什么都没生成**，点「查看官网」进去只有同一句提示 → 纯噪声。
    同时它也不该被谎称成 content_policy_violation（那会让这个账号逃过配额校验）。
    """
    from services.protocol import conversation as conv

    class _StubBackend:
        def __init__(self, **kwargs) -> None:
            self.progress_callback = None

        def close(self) -> None:
            return None

    monkeypatch.setattr(conv.account_service, "get_available_access_token",
                        lambda *args, **kwargs: "tok-quota", raising=False)
    monkeypatch.setattr(conv.account_service, "get_account",
                        lambda token: {"email": "quota@duck.com"}, raising=False)
    monkeypatch.setattr(conv, "record_image_failure", lambda *args, **kwargs: None, raising=False)
    monkeypatch.setattr(conv, "OpenAIBackendAPI", _StubBackend, raising=False)
    monkeypatch.setattr(
        "services.share_evidence_service.share_evidence_service.register_conversation",
        lambda **kwargs: None,
    )

    def _fake_stream(backend, request, index, total, **kwargs):
        yield conv.ImageOutput(
            kind="message",
            model=request.model,
            index=index,
            total=total,
            text=QUOTA_TEXT,
            conversation_id="cid-quota",
        )

    monkeypatch.setattr(conv, "stream_image_outputs", _fake_stream)

    with pytest.raises(conv.ImageGenerationError) as excinfo:
        list(conv._stream_image_outputs_with_pool(_request()))

    exc = excinfo.value
    assert exc.terminal_no_image is True
    assert exc.can_view_official is False
    assert exc.code == "image_quota_exhausted"
    assert exc.status_code == 429
    assert exc.conversation_id == "cid-quota"
    assert str(exc) == QUOTA_TEXT


QUOTA_TEXT = "你已达到 Free 套餐的图像生成请求上限。上限将在 23小时 后重置，届时可创建更多图像。"


@pytest.mark.parametrize(
    "message",
    [
        QUOTA_TEXT,
        # 套餐名/时长变化、只剩后半句时同样要认出来。
        "你已达到 Plus 套餐的图像生成请求上限。",
        "上限将在 1小时 后重置，届时可创建更多图像。",
        # 英文原话（账号 locale 为英文时）。
        "You've reached the Free plan's limit for image generation requests. "
        "The limit will reset in 23 hours.",
        "image generation request limit reached",
    ],
)
def test_is_image_quota_exhausted_message_matches_quota_notices(message: str) -> None:
    from services.protocol.conversation import is_image_quota_exhausted_message

    assert is_image_quota_exhausted_message(message) is True


@pytest.mark.parametrize(
    "message",
    [
        # 内容策略拒绝：官网会话里**有原话可核对**，「查看官网对话」必须保留。
        REFUSAL_TEXT,
        # 上传限流是另一类（换号即可），不是生图配额用完。
        "你已达到文件上传上限。请17小时 内重试。",
        # 空值 / 普通超时文案。
        "",
        "生图超时 (已等待 91 秒）。",
    ],
)
def test_is_image_quota_exhausted_message_ignores_other_failures(message: str) -> None:
    from services.protocol.conversation import is_image_quota_exhausted_message

    assert is_image_quota_exhausted_message(message) is False


def test_poll_raises_text_only_answer_instead_of_waiting(monkeypatch) -> None:
    """终局纯文字要在**第一次轮询**就抛出来，不能把 120s 预算跑满。"""
    from services.config import config
    from services.openai_backend_api import ImageTextOnlyAnswer

    monkeypatch.setattr(type(config), "image_poll_initial_wait_secs", property(lambda self: 0.0))
    monkeypatch.setattr(type(config), "image_poll_interval_secs", property(lambda self: 5.0))

    api = _backend_without_init()
    calls = {"count": 0}

    def _fake_get_conversation(conversation_id: str) -> dict:
        calls["count"] += 1
        return _terminal_doc("请上传需要修改的图片。")

    monkeypatch.setattr(api, "_get_conversation", _fake_get_conversation, raising=False)

    with pytest.raises(ImageTextOnlyAnswer) as excinfo:
        api._poll_image_results("conv-text-only", timeout_secs=120.0)

    assert excinfo.value.text == "请上传需要修改的图片。"
    assert excinfo.value.conversation_id == "conv-text-only"
    assert calls["count"] == 1


def test_poll_keeps_waiting_when_turn_not_terminal(monkeypatch) -> None:
    """回合还没结束（叶子有子节点）时必须继续轮询，等到真超时。"""
    from services.config import config

    monkeypatch.setattr(type(config), "image_poll_initial_wait_secs", property(lambda self: 0.0))
    monkeypatch.setattr(type(config), "image_poll_interval_secs", property(lambda self: 0.05))

    api = _backend_without_init()
    doc = _terminal_doc(children=["tool1"])
    monkeypatch.setattr(api, "_get_conversation", lambda cid: doc, raising=False)

    with pytest.raises(ImagePollTimeoutError) as excinfo:
        api._poll_image_results("conv-still-running", timeout_secs=0.3)

    assert excinfo.value.conversation_id == "conv-still-running"


def _request(**overrides):
    from services.protocol.conversation import ConversationRequest

    return ConversationRequest(
        prompt="生成一张手机壁纸",
        model="gpt-image-2",
        message_as_error=True,
        **overrides,
    )


def _stub_stream_events(monkeypatch, *, turn_use_case: str = "image gen") -> None:
    from services.protocol import conversation as conv

    def _fake_events(*args, **kwargs):
        yield {
            "type": "conversation.done",
            "text": "",
            "conversation_id": "cid-9",
            "parent_message_id": "",
            "file_ids": [],
            "sediment_ids": [],
            "blocked": False,
            "tool_invoked": False,
            "turn_use_case": turn_use_case,
            "done": True,
        }

    monkeypatch.setattr(conv, "conversation_events", _fake_events)


def test_stream_returns_upstream_text_when_turn_ended_with_text(monkeypatch) -> None:
    """轮询识别出「终局纯文字」-> 直接把它当结果返回（前台显示上游原话）。"""
    from services.openai_backend_api import ImageTextOnlyAnswer
    from services.protocol import conversation as conv

    _stub_stream_events(monkeypatch)
    api = _backend_without_init()

    def _fake_resolve(*args, **kwargs):
        raise ImageTextOnlyAnswer("请上传需要修改的图片。", conversation_id="cid-9")

    monkeypatch.setattr(api, "resolve_conversation_image_result", _fake_resolve, raising=False)

    outputs = list(conv.stream_image_outputs(api, _request()))

    assert [output.kind for output in outputs] == ["message"]
    assert outputs[0].text == "请上传需要修改的图片。"
    # 会话 id 必须带出去，前台才有「查看官网」入口。
    assert outputs[0].conversation_id == "cid-9"


def test_stream_still_polls_and_returns_nothing_when_no_text(monkeypatch) -> None:
    """没有终局文字时行为不变：照常轮询，轮询返回空就什么都不产出。"""
    from services.protocol import conversation as conv

    _stub_stream_events(monkeypatch)
    api = _backend_without_init()
    resolved: list[str] = []

    def _fake_resolve(conversation_id, file_ids, sediment_ids, **kwargs):
        resolved.append(conversation_id)
        return ImageResolutionResult(urls=[], message_id="")

    monkeypatch.setattr(api, "resolve_conversation_image_result", _fake_resolve, raising=False)

    outputs = list(conv.stream_image_outputs(api, _request()))

    assert resolved == ["cid-9"]
    assert outputs == []