KiloStar/kilostar/api/plugin.py

from __future__ import annotations

import json
from typing import Optional

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from kilostar.utils.access import Accessor, TokenData
from kilostar.utils.ray_hook import ray_actor_hook
from kilostar.utils.settings import get_plugin_dir

plugin_router = APIRouter(prefix="/api/v1/plugin", tags=["plugin"])


class SubmitRequest(BaseModel):
    org_name: str
    task_description: str
    context: Optional[dict] = None


@plugin_router.post("/submit")
async def submit_task(
    req: SubmitRequest,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    plugins = await pm.list_plugins.remote()
    if req.org_name not in plugins:
        raise HTTPException(404, f"Plugin '{req.org_name}' not found")

    org = ray_actor_hook(f"org_{req.org_name}").get(f"org_{req.org_name}")
    ctx = req.context or {}
    ctx["user"] = token_data.username
    task_id = await org.submit.remote(req.task_description, ctx)
    return {"task_id": task_id}


@plugin_router.get("/task/{task_id}")
async def get_task_status(
    task_id: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    db = ray_actor_hook("postgres_database").postgres_database
    task = await db.get_org_task.remote(task_id)
    if not task:
        raise HTTPException(404, "Task not found")
    return task


@plugin_router.get("/task/{task_id}/events")
async def get_task_events(
    task_id: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    db = ray_actor_hook("postgres_database").postgres_database
    events = await db.query_org_events.remote(task_id)
    return {"events": events}


@plugin_router.get("/task/{task_id}/stream")
async def stream_task(
    task_id: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    import asyncio

    org_name = None
    db = ray_actor_hook("postgres_database").postgres_database
    task = await db.get_org_task.remote(task_id)
    if not task:
        raise HTTPException(404, "Task not found")
    org_name = task["org_name"]

    org = ray_actor_hook(f"org_{org_name}").get(f"org_{org_name}")

    async def _generate():
        async for event in await org.stream.remote(task_id):
            yield f"data: {event}\n\n"

    return StreamingResponse(_generate(), media_type="text/event-stream")


@plugin_router.get("/list")
async def list_plugins(
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    plugins = await pm.list_plugins.remote()
    return {"plugins": plugins}


@plugin_router.post("/install")
async def install_plugin(
    name: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    await pm.install.remote(name)
    return {"status": "ok", "name": name}


@plugin_router.post("/reload/{name}")
async def reload_plugin(
    name: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    pm = ray_actor_hook("global_plugin_manager").global_plugin_manager
    await pm.reload.remote(name)
    return {"status": "ok", "name": name}


@plugin_router.get("/{name}/ui-manifest")
async def ui_manifest(
    name: str,
    token_data: TokenData = Depends(Accessor.get_current_user),
):
    """返回 Web Component 注入所需的元数据。

    读插件 build 产物 ``<plugin>/frontend/dist/wc-manifest.json``，把 ``js`` / ``css`` 路径
    转成绝对静态路径（与 ``/plugin-ui/<name>/`` 静态挂载对齐）。dist 不存在 → 404。
    """
    plugin_dir = get_plugin_dir() / name
    if not plugin_dir.is_dir():
        raise HTTPException(404, f"plugin {name!r} not found")
    wc_path = plugin_dir / "frontend" / "dist" / "wc-manifest.json"
    if not wc_path.is_file():
        raise HTTPException(404, f"plugin {name!r} has no built UI")
    try:
        wc = json.loads(wc_path.read_text(encoding="utf-8"))
    except Exception as e:
        raise HTTPException(500, f"invalid wc-manifest.json: {e}")

    js_rel = wc.get("js")
    if not js_rel or not isinstance(js_rel, str):
        raise HTTPException(500, "wc-manifest.json missing 'js'")
    tag = wc.get("tag")
    if not tag or not isinstance(tag, str):
        raise HTTPException(500, "wc-manifest.json missing 'tag'")

    base_url = f"/plugin-ui/{name}"
    css_rel = wc.get("css") or []
    if isinstance(css_rel, str):
        css_rel = [css_rel]

    # 顺便把 manifest.json 的展示元数据带回前端，少一次往返
    display_name = name
    icon = None
    try:
        manifest_data = json.loads((plugin_dir / "manifest.json").read_text(encoding="utf-8"))
        display_name = manifest_data.get("display_name") or name
        icon = (manifest_data.get("ui") or {}).get("icon")
    except Exception:
        pass

    return {
        "name": name,
        "tag": tag,
        "js": f"{base_url}/{js_rel.lstrip('/')}",
        "css": [f"{base_url}/{c.lstrip('/')}" for c in css_rel],
        "display_name": display_name,
        "icon": icon,
    }