worldquant-alpha-system/backend/app/mcp_api/token_routes.py

"""Cookie-authenticated PAT administration; MCP bearer tokens grant no access here."""

from datetime import timezone

from fastapi import APIRouter, Depends, HTTPException, Query, Request
from pydantic import BaseModel, ConfigDict, Field
from sqlalchemy import func, select

from ..models import Account, MCPToken, now
from ..security import require_auth
from .auth import SCOPES, create_token

router = APIRouter(prefix="/api/v1/mcp-tokens", tags=["mcp-tokens"], dependencies=[Depends(require_auth)])


class TokenInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    name: str = Field(min_length=1, max_length=100)
    days: int = Field(default=90, ge=1, le=365, strict=True)
    scopes: list[str] = Field(default_factory=lambda: ["research:read"], max_length=len(SCOPES))


class TokenPermissionsInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    scopes: list[str] = Field(max_length=len(SCOPES))


def token_output(row, account):
    """Return public metadata only, including whether the current binding is usable."""
    def timestamp(value):
        return value.replace(tzinfo=value.tzinfo or timezone.utc) if value else None

    expires = timestamp(row.expires_at)
    status = (
        "revoked" if row.revoked_at else
        "expired" if expires <= now() else
        "invalid_binding" if not account or account.wq_user_id != row.wq_user_id else
        "active"
    )
    return {
        "id": row.id, "name": row.name, "scopes": row.scopes,
        "created_at": timestamp(row.created_at), "expires_at": expires,
        "revoked_at": timestamp(row.revoked_at), "status": status,
    }


@router.get("")
async def list_tokens(request: Request, limit: int = Query(25, ge=1, le=100), offset: int = Query(0, ge=0)):
    async with request.app.state.sessions() as db:
        account = await db.get(Account, 1)
        owned = (MCPToken.admin_id == 1, MCPToken.account_id == 1)
        total = await db.scalar(select(func.count()).select_from(MCPToken).where(*owned))
        rows = await db.scalars(select(MCPToken).where(*owned).order_by(
            MCPToken.created_at.desc(), MCPToken.id.desc()).offset(offset).limit(limit))
        return {
            "items": [token_output(row, account) for row in rows], "total": total,
            "limit": limit, "offset": offset, "has_more": offset + limit < total,
            "enabled": request.app.state.settings.mcp_enabled,
            "endpoint": request.app.state.settings.public_origin.rstrip("/") + "/api/v1/mcp/",
            "can_create": bool(account and account.wq_user_id),
        }


@router.post("", status_code=201)
async def issue_token(body: TokenInput, request: Request):
    # The existing single-admin browser session is the authority, never request-supplied IDs.
    async with request.app.state.sessions.begin() as db:
        try:
            row, secret = await create_token(db, body.name, body.scopes, body.days)
        except ValueError as exc:
            raise HTTPException(422, str(exc)) from exc
        result = token_output(row, await db.get(Account, 1))
    # Do not expose the secret until the transaction successfully commits.
    return {**result, "token": secret}


@router.patch("/{token_id}")
async def update_token_permissions(token_id: str, body: TokenPermissionsInput, request: Request):
    """Update owned active-token scopes without rotating or revealing its secret."""
    scopes = set(body.scopes)
    if not scopes <= SCOPES or "research:read" not in scopes:
        raise HTTPException(422, "权限无效；所有令牌必须包含 research:read")
    async with request.app.state.sessions.begin() as db:
        row = await db.scalar(select(MCPToken).where(
            MCPToken.id == token_id, MCPToken.admin_id == 1, MCPToken.account_id == 1,
        ).with_for_update())
        if not row:
            raise HTTPException(404, "MCP Key 不存在")
        account = await db.get(Account, 1)
        if token_output(row, account)["status"] != "active":
            raise HTTPException(409, "仅有效的 MCP Key 可以编辑权限")
        row.scopes = sorted(scopes)
        result = token_output(row, account)
    return result


@router.post("/{token_id}/revoke")
async def revoke_token(token_id: str, request: Request):
    async with request.app.state.sessions.begin() as db:
        row = await db.scalar(select(MCPToken).where(
            MCPToken.id == token_id, MCPToken.admin_id == 1, MCPToken.account_id == 1,
        ).with_for_update())
        if not row:
            raise HTTPException(404, "MCP Key 不存在")
        row.revoked_at = row.revoked_at or now()
        result = token_output(row, await db.get(Account, 1))
    return result