feat: 接线 TRAE callback 登录闭环

复核「写了但没接线」的模块时发现：TRAE 的登录函数全都只有定义没有调用，
README 却写了「登录 TRAE 可用」——用户点进去会发现无路可走。

- TraeProvider.start_auth：生成登录 URL（machine/device id 随 state 下发，
  保证落盘凭证与登录时用的设备标识一致）
- TraeProvider.complete_callback：回调链接 → ExchangeToken → 归一化凭证；
  回调只给 refreshToken，accessToken 必须换出来，绝不拿 refreshToken 充数
- POST /api/auth/upstream/start 支持 callback 轨道（CodeBuddy 仍走 poll）
- GET /authorize 真正完成兑换并落库（原来只存 URL 不落盘），并新增：
  · state 必须与 start 发放的一致，且消费后即失效 —— 不允许回退 URL 里的
    state，否则旧链接可以被重复兑换
  · 响应与凭证列表都不回传 token
  · 兑换失败归一为 invalid_credential，不漏 HTTP 细节到 API 层
  · 入库后立即触发额度探测

后端 541 测试 / 覆盖 100%。真实进程验证：flow=callback、state 校验、
假 token 兑换失败不入库。

Author: RobbsLuo

src/main.py

@@ -83,6 +83,8 @@ async def lifespan(app: FastAPI):
     app.state.stats_query = StatsQuery(db)
     app.state.upstream_auth = _upstream_auth(registry, config)
     app.state.pending_probes = []
+    app.state.pending_callback_state = None
+    app.state.pending_callback_user = None
 
     # ------------------------------------------------------------ 鉴权依赖
 
@@ -270,12 +272,19 @@ async def upstream_auth_start(payload: dict,
                                   principal: Principal = Depends(principal_from_request)):
         require_admin(principal)
         provider_id = str(payload.get("provider") or "")
-        oauth = app.state.upstream_auth.get(provider_id)
-        if oauth is None:
-            raise InvalidRequest(f"provider {provider_id!r} does not support polling login")
-        session = await oauth.start(principal.username)
+        if provider_id in app.state.upstream_auth:
+            session = await app.state.upstream_auth[provider_id].start(principal.username)
+        else:
+            provider = registry.get(provider_id)
+            builder = getattr(provider, "start_auth", None)
+            if not callable(builder):
+                raise InvalidRequest(f"provider {provider_id!r} does not support login")
+            session = builder(resolve_public_callback_url(config))
+            app.state.pending_callback_state = session.state
+            app.state.pending_callback_user = principal.username
+            # 回调轨道没有本地轮询：登录结果由 /authorize 落库后由前端查凭证列表
         return {"flow": session.flow, "state": session.state, "auth_url": session.auth_url,
-                "interval": session.interval}
+                "interval": session.interval, "callback_url": session.callback_url}
 
     @app.post("/api/auth/upstream/poll")
     async def upstream_auth_poll(payload: dict,
@@ -412,9 +421,36 @@ async def stats_by_provider(principal: Principal = Depends(principal_from_reques
 
     @app.get("/authorize")
     async def authorize(request: Request):
-        """TRAE 浏览器 302 落点：只捕获 query，不落盘。"""
+        """TRAE 浏览器 302 落点。
+
+        回调不需要 API Key（浏览器不会带），因此这里不做鉴权，但：
+        - 只接受带 refreshToken 的链接，其他一律拒绝
+        - 换到的凭证直接落库，响应里绝不回传 token
+        - state 必须与 start_auth 发放的一致，防止任意回调被塞进池子
+        """
         raw = str(request.url)
         app.state.last_callback_url = raw
+        state = request.query_params.get("state")
+        provider = registry.get("trae")
+        pending = app.state.pending_callback_state
+        if provider is None or pending is None or state != pending:
+            # 没有进行中的登录，或 state 不匹配（含已消费后的重放）→ 拒绝。
+            # 不允许回退到 URL 里的 state，否则旧链接可以被重复兑换。
+            return JSONResponse(status_code=400, content=error_payload(
+                "no pending TRAE login in progress", "invalid_request", 400))
+        if not request.query_params.get("refreshToken") and not request.query_params.get("userJwt"):
+            return JSONResponse(status_code=400, content=error_payload(
+                "callback missing refreshToken", "invalid_request", 400))
+        try:
+            credential_data = await provider.complete_callback(raw, state)
+        except UpstreamProtocolViolation as error:
+            return JSONResponse(status_code=400, content=error_payload(
+                str(error), "invalid_credential", 400))
+        app.state.pending_callback_state = None
+        credential_id = credentials.add(provider="trae", credential_data=credential_data,
+                                        nickname=str(credential_data.get("nickname") or ""),
+                                        added_by=app.state.pending_callback_user or "")
+        schedule_probe(credential_id)
         return {"ok": True, "captured": True, "at": int(time.time())}
 
     # 静态资源必须最后注册：catch-all 会匹配所有未命中的路径

src/provider/trae/client.py

@@ -13,9 +13,15 @@
 import httpx
 
 from ...engine.sse import iter_frames
-from ...provider.base import ErrKind, Event, Model, Quota
+from ...provider.base import AuthSession, ErrKind, Event, Model, Quota
 from . import events as trae_events
-from .callback import build_login_url, parse_callback_url
+from .callback import (
+    CallbackInfo,
+    build_login_url,
+    credential_from_callback,
+    new_machine_identity,
+    parse_callback_url,
+)
 from .credential import TraeCredential, parse_credential
 from .events import (
     AGENT_HOST,
@@ -323,10 +329,55 @@ async def aclose(self) -> None:
         """释放内部 HTTP 连接池。"""
         await self.client.aclose()
 
-    def parse_login_url(self, raw_url: str) -> dict:
-        """解析 TRAE 登录回调链接（Q17=C callback 轨道）。"""
+    # ------------------------------------------------- callback 轨道（Q17=C）
+
+    def parse_login_url(self, raw_url: str) -> CallbackInfo:
+        """解析 TRAE 登录回调链接。"""
         return parse_callback_url(raw_url)
 
     def build_login_url(self, callback_url: str, *, machine_id: str,
                         device_id: str) -> str:
         return build_login_url(callback_url, machine_id=machine_id, device_id=device_id)
+
+    def start_auth(self, callback_url: str) -> AuthSession:
+        """生成登录 URL。machine/device id 由调用方保管，落盘凭证必须复用同一对。"""
+        machine_id, device_id = new_machine_identity()
+        return AuthSession(
+            flow="callback", state=f"{machine_id}:{device_id}",
+            callback_url=callback_url,
+            auth_url=build_login_url(callback_url, machine_id=machine_id,
+                                     device_id=device_id),
+        )
+
+    async def complete_callback(self, raw_url: str, state: str) -> dict:
+        """回调链接 → ExchangeToken → 归一化凭证。
+
+        state 形如 ``machine_id:device_id``，用于保证落盘凭证与登录时用的
+        设备标识一致（原实现里这两者不一致会导致登录态与凭证不匹配）。
+        """
+        info = parse_callback_url(raw_url)
+        machine_id, _, device_id = state.partition(":")
+        if not machine_id or not device_id:
+            raise UpstreamProtocolViolation("auth state missing machine/device id")
+
+        # 回调只给 refreshToken，accessToken 必须由 ExchangeToken 换出来。
+        # 换失败就没有可用凭证，绝不能用 refreshToken 充当 accessToken 混进池子。
+        credential = credential_from_callback(
+            info, "", machine_id=machine_id, device_id=device_id)
+        try:
+            credential = await self.client.refresh_token(credential)
+        except UpstreamHTTPError as error:
+            # 上游拒绝兑换（refreshToken 过期/失效）→ 归一为协议违规，
+            # 让调用方按「凭证无效」处理，而不是把 HTTP 细节漏到 API 层
+            raise UpstreamProtocolViolation(
+                f"token exchange rejected: {error.status}") from error
+        if not credential.uid:
+            uid, nickname = await self.client.get_user_info(credential)
+            credential = TraeCredential(
+                uid=uid or credential.uid, access_token=credential.access_token,
+                refresh_token=credential.refresh_token, expires_at=credential.expires_at,
+                domain=credential.domain, api_host=credential.api_host,
+                machine_id=credential.machine_id, device_id=credential.device_id,
+                enterprise_id=credential.enterprise_id,
+                nickname=nickname or credential.nickname)
+        return credential.to_dict()

tests/test_m15_operations.py

@@ -1012,11 +1012,20 @@ def admin_client(tmp_path):
         yield app, client
 
 
-def test_upstream_auth_start_requires_admin_and_known_provider(admin_client):
-    _app, client = admin_client
-    started = client.post("/api/auth/upstream/start", json={"provider": "codebuddy"})
-    assert started.status_code in (200, 400, 502)
-    assert client.post("/api/auth/upstream/start", json={"provider": "trae"}).status_code == 400
+def test_upstream_auth_start_supports_both_providers(admin_client):
+    """CodeBuddy 走 poll 轨道，TRAE 走 callback 轨道，unknown provider 报 400。"""
+    app, client = admin_client
+    # TRAE 的 callback 轨道不需要出网，必定成功并带入参回调地址
+    trae = client.post("/api/auth/upstream/start", json={"provider": "trae"})
+    assert trae.status_code == 200
+    body = trae.json()
+    assert body["flow"] == "callback"
+    assert body["callback_url"].endswith("/authorize")
+    assert "auth_callback_url" in body["auth_url"]
+    assert app.state.pending_callback_state == body["state"]
+
+    assert client.post("/api/auth/upstream/start",
+                       json={"provider": "unknown"}).status_code == 400
 
 
 def test_upstream_auth_poll_unknown_state(admin_client):
@@ -2052,3 +2061,195 @@ def test_schedule_probe_returns_early_when_credential_unreadable(admin_client):
     assert created.status_code == 200
     time.sleep(0.05)
     assert len(app.state.pending_probes) == before
+
+
+# ------------------------------------------- TRAE callback 登录闭环
+
+def test_trae_start_auth_builds_login_url_with_public_callback():
+    from src.main import resolve_public_callback_url
+    from src.provider.trae.client import TraeProvider
+
+    settings = Settings(_env_file=None, APP_SECRET="s", PUBLIC_BASE_URL="https://gw.example")
+    provider = TraeProvider()
+    session = provider.start_auth(resolve_public_callback_url(settings))
+
+    assert session.flow == "callback"
+    assert session.callback_url == "https://gw.example/authorize"
+    assert "auth_callback_url=https%3A%2F%2Fgw.example%2Fauthorize" in session.auth_url
+    machine_id, _, device_id = session.state.partition(":")
+    assert len(machine_id) == 32 and len(device_id) == 32
+
+
+async def test_trae_complete_callback_exchanges_token():
+    """回调链接必须真的换 token，而不是只存 refreshToken。"""
+    import httpx as _httpx
+
+    from src.provider.trae.client import TraeClient, TraeProvider
+
+    def handler(request: _httpx.Request) -> _httpx.Response:
+        if request.url.path.endswith("ExchangeToken"):
+            return _httpx.Response(200, json={"Result": {
+                "Token": "ACCESS", "RefreshToken": "RT2", "TokenExpireAt": 1_800_000_000_000}})
+        return _httpx.Response(200, json={"Result": {"UserID": "uid-9", "ScreenName": "昵称"}})
+
+    transport = _httpx.MockTransport(handler)
+    provider = TraeProvider(client=TraeClient(
+        stream_client=_httpx.AsyncClient(transport=transport, timeout=None),
+        short_client=_httpx.AsyncClient(transport=transport, timeout=None)))
+
+    session = provider.start_auth("https://gw.example/authorize")
+    url = ("https://gw.example/authorize?refreshToken=RT1&userInfo="
+           "%7B%22uid%22%3A%22%22%7D")
+    data = await provider.complete_callback(url, session.state)
+
+    assert data["accessToken"] == "ACCESS"
+    assert data["refreshToken"] == "RT2"
+    assert data["expiresAt"] == 1_800_000_000
+    assert data["uid"] == "uid-9"          # 回调没给 uid 时回退 GetUserInfo
+    assert data["nickname"] == "昵称"
+    assert data["machineId"] == session.state.partition(":")[0]
+
+
+async def test_trae_complete_callback_rejects_bad_state():
+    from src.provider.trae.client import TraeProvider
+    from src.provider.trae.events import UpstreamProtocolViolation
+
+    provider = TraeProvider()
+    with pytest.raises(UpstreamProtocolViolation):
+        await provider.complete_callback("https://x/authorize?refreshToken=RT", "no-colon")
+
+
+def test_authorize_completes_trae_login_end_to_end(tmp_path):
+    """完整闭环：start → 浏览器回调 → 凭证入库 → 立即探测。"""
+    import httpx as _httpx
+
+    from src.provider.trae.client import TraeClient, TraeProvider
+
+    def handler(request: _httpx.Request) -> _httpx.Response:
+        if request.url.path.endswith("ExchangeToken"):
+            return _httpx.Response(200, json={"Result": {"Token": "ACCESS",
+                                                         "RefreshToken": "RT2"}})
+        if request.url.endswith("ide_user_ent_usage"):
+            return _httpx.Response(200, json={"user_entitlement_pack_list": [
+                {"entitlement_base_info": {"quota": {"credits_limit": 100}},
+                 "usage": {"credits_amount": 25}}]})
+        return _httpx.Response(200, json={"Result": {"UserID": "u1", "ScreenName": "n"}})
+
+    transport = _httpx.MockTransport(handler)
+    trae = TraeProvider(client=TraeClient(
+        stream_client=_httpx.AsyncClient(transport=transport, timeout=None),
+        short_client=_httpx.AsyncClient(transport=transport, timeout=None)))
+
+    settings = Settings(_env_file=None, APP_SECRET="s", DATA_DIR=str(tmp_path),
+                        ADMIN_USERNAMES="root")
+    app = build_app(settings, providers={"trae": trae})
+
+    with TestClient(app) as client:
+        client.post("/api/auth/login", json={"username": "root", "password": "rootpw"})
+        started = client.post("/api/auth/upstream/start", json={"provider": "trae"}).json()
+
+        callback = ("/authorize?refreshToken=RT1&userInfo=%7B%22uid%22%3A%22u1%22%7D"
+                    f"&state={started['state']}")
+        response = client.get(callback)
+        assert response.status_code == 200 and response.json()["captured"] is True
+
+        listed = client.get("/api/credentials").json()["credentials"]
+        assert len(listed) == 1 and listed[0]["provider"] == "trae"
+        # 响应与列表都不得泄漏 token
+        assert "ACCESS" not in response.text
+        assert "data_enc" not in listed[0]
+
+    # state 已消费，同一回调不可重放
+    with TestClient(app) as replay:
+        assert replay.get(callback).status_code == 400
+
+
+def test_authorize_rejects_callback_without_token(tmp_path):
+    """带 state 但没有 refreshToken/userJwt 的回调必须拒绝。"""
+    settings = Settings(_env_file=None, APP_SECRET="s", DATA_DIR=str(tmp_path),
+                        ADMIN_USERNAMES="root")
+    app = build_app(settings)
+    with TestClient(app) as client:
+        client.post("/api/auth/login", json={"username": "root", "password": "rootpw"})
+        started = client.post("/api/auth/upstream/start", json={"provider": "trae"}).json()
+        response = client.get(f"/authorize?state={started['state']}")
+    assert response.status_code == 400
+    assert response.json()["error"]["code"] == "invalid_request"
+
+
+def test_upstream_start_uses_poll_track_for_codebuddy(admin_client):
+    """provider 在 upstream_auth 里时走 poll 轨道（main 275-276）。"""
+    app, client = admin_client
+
+    class FakeOAuth:
+        def __init__(self) -> None:
+            self.store = type("S", (), {"cancel": staticmethod(lambda *_: True)})()
+
+        async def start(self, username):
+            from src.provider.base import AuthSession
+
+            return AuthSession(flow="poll", state="local-reservation",
+                               auth_url="https://auth.example/x", interval=5)
+
+    app.state.upstream_auth["codebuddy"] = FakeOAuth()
+    body = client.post("/api/auth/upstream/start", json={"provider": "codebuddy"}).json()
+    assert body["flow"] == "poll"
+    assert body["state"] == "local-reservation"
+    assert body["callback_url"] is None
+
+
+async def test_complete_callback_does_not_use_refresh_token_as_access_token():
+    """ExchangeToken 失败时不得把 refreshToken 当 accessToken 塞进池子。"""
+    import httpx as _httpx
+
+    from src.provider.trae.client import TraeClient, TraeProvider
+    from src.provider.trae.events import UpstreamProtocolViolation
+
+    def handler(request: _httpx.Request) -> _httpx.Response:
+        if request.url.path.endswith("ExchangeToken"):
+            return _httpx.Response(400, content=b"bad refresh token")
+        return _httpx.Response(200, json={"Result": {"UserID": "u", "ScreenName": "n"}})
+
+    transport = _httpx.MockTransport(handler)
+    provider = TraeProvider(client=TraeClient(
+        stream_client=_httpx.AsyncClient(transport=transport, timeout=None),
+        short_client=_httpx.AsyncClient(transport=transport, timeout=None)))
+
+    session = provider.start_auth("https://gw.example/authorize")
+    with pytest.raises(UpstreamProtocolViolation):
+        await provider.complete_callback(
+            "https://gw.example/authorize?refreshToken=RT", session.state)
+
+
+async def test_complete_callback_raises_when_no_token_available():
+    import httpx as _httpx
+
+    from src.provider.trae.client import TraeClient, TraeProvider
+    from src.provider.trae.events import UpstreamProtocolViolation
+
+    def handler(_request: _httpx.Request) -> _httpx.Response:
+        return _httpx.Response(400, content=b"rejected")
+
+    transport = _httpx.MockTransport(handler)
+    provider = TraeProvider(client=TraeClient(
+        stream_client=_httpx.AsyncClient(transport=transport, timeout=None),
+        short_client=_httpx.AsyncClient(transport=transport, timeout=None)))
+    session = provider.start_auth("https://gw.example/authorize")
+    with pytest.raises(UpstreamProtocolViolation):
+        await provider.complete_callback(
+            "https://gw.example/authorize?refreshToken=RT", session.state)
+
+
+def test_authorize_reports_invalid_credential(tmp_path):
+    """回调结构损坏 → 400 invalid_credential（main 446-448）。"""
+    settings = Settings(_env_file=None, APP_SECRET="s", DATA_DIR=str(tmp_path),
+                        ADMIN_USERNAMES="root")
+    app = build_app(settings)
+    with TestClient(app) as client:
+        client.post("/api/auth/login", json={"username": "root", "password": "rootpw"})
+        client.post("/api/auth/upstream/start", json={"provider": "trae"})
+        # 有 refreshToken 但状态串损坏 → complete_callback 抛协议违规
+        app.state.pending_callback_state = "broken"
+        response = client.get("/authorize?refreshToken=RT&state=broken")
+    assert response.status_code == 400
+    assert response.json()["error"]["code"] == "invalid_credential"

tests/test_m1a_trae.py

@@ -876,7 +876,9 @@ def test_non_admin_cannot_write(client):
                                                  "credential": {}}).status_code == 403
 
 
-def test_authorize_callback_capture(client):
+def test_authorize_rejects_callback_without_pending_login(client):
+    """没有进行中的登录时，任意回调不得被塞进凭证池。"""
     response = client.get("/authorize?refreshToken=RT")
-    assert response.status_code == 200 and response.json()["captured"] is True
+    assert response.status_code == 400
+    assert response.json()["error"]["code"] == "invalid_request"
     assert "refreshToken=RT" in client.app.state.last_callback_url