holy-python/app.py

import secrets
from pathlib import Path
from urllib.parse import quote

from flask import Flask, g, jsonify, redirect, request, send_file

from auth_service import (
    build_dingtalk_login_url,
    exchange_dingtalk_user,
    get_user_by_session_token,
    revoke_session,
    upsert_user_and_session,
)
from config import Config
from database import init_database
from rule_service import (
    delete_rule_document,
    get_rule_document,
    get_rule_html_path,
    save_rule_document,
)
from task_service import (
    TaskError,
    cancel_task,
    create_task,
    get_file,
    get_task,
    get_task_extracted,
    list_tasks,
    recollect_task,
    resume_task,
    resume_unfinished_tasks,
)


def create_app() -> Flask:
    app = Flask(__name__)
    init_database()

    @app.before_request
    def assign_request_id():
        g.request_id = request.headers.get("X-Request-Id") or f"req_{secrets.token_hex(10)}"
        if request.method == "OPTIONS":
            return None
        open_paths = ("/health", "/api/v1/auth/", "/api/v1/skill/")
        if request.path == "/health" or any(request.path.startswith(path) for path in open_paths):
            return None
        user = get_user_by_session_token(request.cookies.get(Config.SESSION_COOKIE_NAME))
        if user is None and request.path.startswith("/api/"):
            return jsonify({
                "error": {"code": "UNAUTHORIZED", "message": "请先通过钉钉授权登录"},
                "request_id": g.request_id,
            }), 401
        g.current_user = user

    @app.after_request
    def add_headers(response):
        response.headers["X-Request-Id"] = g.get("request_id", "")
        response.headers["Access-Control-Allow-Origin"] = Config.CORS_ORIGIN
        response.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization,Idempotency-Key,X-Request-Id"
        response.headers["Access-Control-Allow-Methods"] = "GET,POST,DELETE,OPTIONS"
        response.headers["Access-Control-Allow-Credentials"] = "true"
        return response

    @app.errorhandler(TaskError)
    def handle_task_error(error: TaskError):
        return jsonify({
            "error": {"code": error.code, "message": error.message, "details": error.details},
            "request_id": g.request_id,
        }), error.status_code

    @app.errorhandler(Exception)
    def handle_unexpected(error: Exception):
        app.logger.exception("未处理异常")
        return jsonify({
            "error": {"code": "INTERNAL_ERROR", "message": "服务器内部错误"},
            "request_id": g.get("request_id"),
        }), 500

    @app.route("/health", methods=["GET"])
    def health():
        return jsonify({"status": "ok", "request_id": g.request_id})

    @app.route("/api/v1/auth/dingtalk/login", methods=["GET"])
    def dingtalk_login():
        return jsonify({
            "data": build_dingtalk_login_url(request.args.get("redirect", "/tasks")),
            "request_id": g.request_id,
        })

    @app.route("/api/v1/auth/dingtalk/callback", methods=["GET"])
    def dingtalk_callback():
        code = request.args.get("authCode") or request.args.get("code")
        state = request.args.get("state", "")
        frontend_redirect = "/tasks"
        if ":" in state:
            frontend_redirect = state.split(":", 1)[1] or frontend_redirect
        if not code:
            return redirect(f"{Config.FRONTEND_BASE_URL}/login?error={quote('钉钉授权缺少 code')}")
        user_payload = exchange_dingtalk_user(code)
        user, token, expires_at = upsert_user_and_session(user_payload)
        response = redirect(f"{Config.FRONTEND_BASE_URL}{frontend_redirect}")
        response.set_cookie(
            Config.SESSION_COOKIE_NAME,
            token,
            expires=expires_at,
            httponly=True,
            secure=Config.SESSION_COOKIE_SECURE,
            samesite=Config.SESSION_COOKIE_SAMESITE,
            path="/",
        )
        # 给前端“Link AI”复制提示词使用。后端鉴权仍读取 hc_session；
        # 前端只读取这个辅助 Cookie 的 value，并在提示词中拼成 hc_session=<value>。
        response.set_cookie(
            "hc_session_agent",
            token,
            expires=expires_at,
            httponly=False,
            secure=Config.SESSION_COOKIE_SECURE,
            samesite=Config.SESSION_COOKIE_SAMESITE,
            path="/",
        )
        app.logger.info("钉钉用户登录成功 user_id=%s name=%s", user["id"], user["name"])
        return response

    @app.route("/api/v1/auth/direct-token", methods=["GET"])
    def auth_direct_token():
        """直接用钉钉 authCode 换取 session token，AI 无需浏览器 Cookie。

        调用方式：
            GET /api/v1/auth/direct-token?authCode=钉钉返回的code

        返回：
            {"data": {"token": "session_token", "user": {...}, "expires_at": "..."}}

        AI 拿到 token 后，后续请求带上：
            Cookie: hc_session=<token>
        """
        code = request.args.get("authCode") or request.args.get("code")
        if not code:
            return jsonify({"code": "MISSING_AUTH_CODE", "msg": "缺少 authCode 参数"}), 400
        try:
            dingtalk_user = exchange_dingtalk_user(code)
            user, token, expires_at_dt = upsert_user_and_session(dingtalk_user)
            return jsonify({
                "data": {
                    "token": token,
                    "user": user,
                    "expires_at": expires_at_dt.isoformat(timespec="seconds"),
                },
                "request_id": g.request_id,
            })
        except TaskError as e:
            return jsonify({"code": e.code, "msg": e.message, "data": e.data}), e.http_status

    # ── Skill 版本管理 ────────────────────────────────────────────
    SKILL_BASE = Path(__file__).parent / "skills"

    @app.route("/api/v1/skill/holy-crab/version", methods=["GET"])
    def skill_holy_crab_version():
        """返回当前 Holy Crab Skill 的版本信息。"""
        import json
        skill_dir = SKILL_BASE / "holy-crab"
        version_file = skill_dir / "VERSION"
        if not version_file.exists():
            return jsonify({"code": "SKILL_NOT_FOUND", "msg": "Skill 文件不存在"}), 404
        info = json.loads(version_file.read_text(encoding="utf-8"))
        return jsonify({"data": {"version": info}, "request_id": g.request_id})

    @app.route("/api/v1/skill/holy-crab/download", methods=["GET"])
    def skill_holy_crab_download():
        """下载 Holy Crab Skill 完整 zip 包（含所有文件）。"""
        import hashlib, io, zipfile
        skill_dir = SKILL_BASE / "holy-crab"
        if not skill_dir.exists():
            return jsonify({"code": "SKILL_NOT_FOUND", "msg": "Skill 文件不存在"}), 404

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
            for fpath in skill_dir.rglob("*"):
                if fpath.is_file():
                    arcname = str(fpath.relative_to(skill_dir))
                    zf.writestr(arcname, fpath.read_bytes())
        buf.seek(0)
        return buf.getvalue(), 200, {
            "Content-Type": "application/zip",
            "Content-Disposition": "attachment; filename=holy-crab-skill.zip",
            "X-Skill-Version": json.loads((skill_dir / "VERSION").read_text(encoding="utf-8")).get("version", "unknown"),
        }

    @app.route("/api/v1/auth/me", methods=["GET"])
    def auth_me():
        user = get_user_by_session_token(request.cookies.get(Config.SESSION_COOKIE_NAME))
        if not user:
            raise TaskError("UNAUTHORIZED", "请先通过钉钉授权登录", 401)
        return jsonify({"data": {"user": user}, "request_id": g.request_id})

    @app.route("/api/v1/auth/logout", methods=["POST"])
    def auth_logout():
        revoke_session(request.cookies.get(Config.SESSION_COOKIE_NAME))
        response = jsonify({"data": {"ok": True}, "request_id": g.request_id})
        response.delete_cookie(Config.SESSION_COOKIE_NAME, path="/")
        response.delete_cookie("hc_session_agent", path="/")
        return response

    @app.route("/api/v1/tasks", methods=["POST"])
    def create_research_task():
        payload = request.get_json(silent=True) or {}
        current_user = getattr(g, "current_user", None) or {}
        payload["created_by_id"] = current_user.get("id") or payload.get("created_by_id")
        payload["created_by_name"] = current_user.get("name") or payload.get("created_by_name")
        task = create_task(payload)
        return jsonify({"data": task, "request_id": g.request_id}), 201

    @app.route("/api/v1/tasks", methods=["GET"])
    def query_research_tasks():
        result = list_tasks(
            q=request.args.get("q", "").strip(),
            status=request.args.get("status", "").strip(),
            page=int(request.args.get("page", "1")),
            page_size=int(request.args.get("page_size", "10")),
        )
        return jsonify({"data": result, "request_id": g.request_id})

    @app.route("/api/v1/tasks/<task_id>", methods=["GET"])
    def get_research_task(task_id: str):
        if request.args.get("view") in {"extracted", "cleaned", "parsed"}:
            return jsonify({"data": get_task_extracted(task_id), "request_id": g.request_id})
        return jsonify({"data": get_task(task_id), "request_id": g.request_id})

    @app.route("/api/v1/tasks/<task_id>/extracted", methods=["GET"])
    def get_research_task_extracted(task_id: str):
        return jsonify({"data": get_task_extracted(task_id), "request_id": g.request_id})

    @app.route("/api/v1/tasks/<task_id>/cancel", methods=["POST"])
    def cancel_research_task(task_id: str):
        body = request.get_json(silent=True) or {}
        return jsonify({
            "data": cancel_task(task_id, str(body.get("reason") or "")),
            "request_id": g.request_id,
        })

    @app.route("/api/v1/tasks/<task_id>/retry", methods=["POST"])
    @app.route("/api/v1/tasks/<task_id>/resume", methods=["POST"])
    def resume_research_task(task_id: str):
        return jsonify({"data": resume_task(task_id), "request_id": g.request_id})

    @app.route("/api/v1/tasks/<task_id>/recollect", methods=["POST"])
    def recollect_research_task(task_id: str):
        return jsonify({
            "data": recollect_task(task_id),
            "request_id": g.request_id,
        }), 201

    @app.route("/api/v1/tasks/<task_id>/files", methods=["GET"])
    def list_task_files(task_id: str):
        task = get_task(task_id)
        return jsonify({"data": {"items": task["files"]}, "request_id": g.request_id})

    @app.route("/api/v1/tasks/<task_id>/files/<file_id>/preview", methods=["GET"])
    def preview_task_file(task_id: str, file_id: str):
        metadata, path = get_file(task_id, file_id)
        if not metadata["previewable"]:
            raise TaskError("FILE_TOO_LARGE", "文件超过预览大小限制", 413)
        return send_file(path, as_attachment=False)

    @app.route("/api/v1/tasks/<task_id>/files/<file_id>/download", methods=["GET"])
    def download_task_file(task_id: str, file_id: str):
        _, path = get_file(task_id, file_id)
        return send_file(path, as_attachment=True, download_name=Path(path).name)

    @app.route("/api/v1/rules", methods=["GET"])
    def get_current_rule():
        return jsonify({"data": get_rule_document(), "request_id": g.request_id})

    @app.route("/api/v1/rules/current.html", methods=["GET"])
    def preview_current_rule_html():
        return send_file(get_rule_html_path(), mimetype="text/html; charset=utf-8")

    @app.route("/api/v1/rules/html", methods=["POST"])
    def upload_rule_html():
        return jsonify({"data": save_rule_document(request), "request_id": g.request_id})

    @app.route("/api/v1/rules/html", methods=["DELETE"])
    def remove_rule_html():
        return jsonify({"data": delete_rule_document(), "request_id": g.request_id})

    resume_unfinished_tasks()
    return app


app = create_app()


if __name__ == "__main__":
    app.run(host=Config.HOST, port=Config.PORT, threaded=True)