ModelRouter/internal/gateway/plugins_api.go

package gateway

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"llmsproxy/internal/lua"
)

var (
	errPluginName  = errors.New("plugin name must be non-empty and contain no path separator or dot")
	errNoPluginDir = errors.New("no plugin_dir configured; set plugin_dir in config.yaml to enable plugins")
)

// Plugin management API.
//
//	GET    /api/plugins                  list loaded plugins + their hooks/UI/errors
//	POST   /api/plugins                  upload/replace one plugin (.lua), hot-applied
//	DELETE /api/plugins/{name}           remove a plugin
//	GET    /api/plugins/{name}/state     the plugin's own published state
//	PUT    /api/plugins/{name}/state     replace that state (admin only)
//
// Why state is a first-class endpoint: request_end is the LAST stage, so a
// hook's return value has no downstream consumer inside the gateway. A plugin
// that accumulates numbers (the billing plugin does exactly this) therefore
// keeps them in its own Lua state and serves them here, which is what its UI
// component fetches. This keeps plugin data clearly separated from the
// gateway's own stats — see the accounting note in docs/plugins.md: the
// gateway's quota accounting stays authoritative, a plugin only reports.

// handlePluginUI serves the merged UI extensions the WebUI injects at boot.
//
// It is a single GET (not per-plugin) because the browser needs ALL extensions
// before it can build the sidebar: a page contributed by one plugin and an
// element contributed by another land in the same payload, and fetching them
// separately would mean the sidebar has to be rebuilt as each arrives.
//
// Served to any authenticated role: the WebUI is authenticated before it asks,
// and a plugin's own widgets need to render for the user whose key they show.
func (g *Gateway) handlePluginUI(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
		return
	}
	ps := g.core.Plugins()
	if ps == nil {
		writeJSON(w, http.StatusOK, map[string]interface{}{"ui": lua.UIExtension{}})
		return
	}
	// AllStages, not a hand-written list: an earlier version enumerated the
	// three stages literally here, so when chain_step was added it was silently
	// missing from this response — a plugin author reading the discovery
	// document would have believed the stage did not exist. Deriving it from
	// the one place that defines the order is the whole point of having it.
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"ui":     ps.UI(),
		"stages": pipelineStageNames(),
	})
}

// stageNames returns the pipeline stage names in firing order, for the
// discovery payload and for tests.
func pipelineStageNames() []string {
	out := make([]string, 0, len(lua.AllStages))
	for _, s := range lua.AllStages {
		out = append(out, string(s))
	}
	return out
}

func (g *Gateway) handlePluginsAPI(w http.ResponseWriter, r *http.Request) {
	ps := g.core.Plugins()
	if ps == nil {
		writeJSON(w, http.StatusOK, map[string]interface{}{"plugins": []interface{}{}})
		return
	}
	path := strings.TrimPrefix(r.URL.Path, "/api/plugins")
	path = strings.Trim(path, "/")

	// /api/plugins/{name}/state
	if strings.HasSuffix(path, "/state") {
		name := strings.TrimSuffix(path, "/state")
		if name == "" {
			writeError(w, http.StatusBadRequest, "invalid_request", "plugin name required")
			return
		}
		g.handlePluginState(w, r, name)
		return
	}

	if reqRole(r.Context()) != "admin" {
		writeError(w, http.StatusForbidden, "forbidden", "admin role required")
		return
	}

	if path == "" {
		switch r.Method {
		case http.MethodGet:
			writeJSON(w, http.StatusOK, map[string]interface{}{
				"plugins":     ps.List(),
				"on_disk":     ps.OnDisk(),
				"hook_errors": ps.HookErrors(),
				"plugin_dir":  g.pluginDir(),
			})
		case http.MethodPost:
			var body struct {
				Name string `json:"name"`
				Code string `json:"code"`
			}
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
				return
			}
			if err := g.installPlugin(body.Name, body.Code); err != nil {
				writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "name": body.Name})
		default:
			writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
		}
		return
	}

	// /api/plugins/{name} — source read (GET), enable/disable (PUT), delete.
	// Multi-segment paths (…/state) are handled earlier and never reach here.
	if !strings.Contains(path, "/") {
		switch r.Method {
		case http.MethodGet:
			// Reading the SOURCE (not the runtime state) is what an editor
			// needs; the state endpoint is /state and returns accumulated data
			// instead. Mixing the two would make a "save what I read"
			// round-trip impossible.
			code, err := g.readPluginSource(path)
			if err != nil {
				writeError(w, http.StatusNotFound, "not_found", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{
				"name": path, "code": code,
				"enabled": g.core.Plugins().Enabled(path),
			})
		case http.MethodPut:
			// Enable / disable. The intent rides in the body rather than the
			// verb, because "enable" and "disable" are the same resource and
			// PUT /{name} is the one the docs advertise.
			var body struct {
				Enabled *bool `json:"enabled"`
			}
			if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
				writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
				return
			}
			if body.Enabled == nil {
				writeError(w, http.StatusBadRequest, "invalid_request", `body must be {"enabled":true|false}`)
				return
			}
			if err := g.core.Plugins().SetEnabled(path, *body.Enabled); err != nil {
				writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{
				"ok": true, "name": path, "enabled": *body.Enabled,
			})
		case http.MethodDelete:
			if err := g.removePlugin(path); err != nil {
				writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
				return
			}
			writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
		default:
			writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
		}
		return
	}

	writeError(w, http.StatusNotFound, "not_found", "unknown plugin sub-resource: "+path)
}

// readPluginSource returns a plugin's file contents for the editor.
func (g *Gateway) readPluginSource(name string) (string, error) {
	if err := validPluginName(name); err != nil {
		return "", err
	}
	if g.pluginDir() == "" {
		return "", errNoPluginDir
	}
	b, err := os.ReadFile(filepath.Join(g.pluginDir(), name+".lua"))
	if err != nil {
		return "", fmt.Errorf("no plugin source named %q", name)
	}
	return string(b), nil
}

// validPluginName rejects anything that could escape the plugin directory.
// Shared by install / remove / read so the check cannot drift between them.
func validPluginName(name string) error {
	if name == "" || strings.ContainsAny(name, `/\.`) {
		return errPluginName
	}
	return nil
}

// handlePluginState serves GET (read state) and PUT (replace state).
func (g *Gateway) handlePluginState(w http.ResponseWriter, r *http.Request, name string) {
	ps := g.core.Plugins()
	switch r.Method {
	case http.MethodGet:
		// Any role may read: plugin state is reporting data (cost, counts),
		// and the caller has already been authenticated. Admin-only would stop
		// a user key's own billing widget from rendering.
		// prices rides along: the billing rules editor must show the table
		// that is actually in effect, not a reconstruction from config. It is
		// absent for plugins with no prices (every plugin but billing).
		payload := map[string]interface{}{
			"plugin": name,
			"state":  ps.State(name),
		}
		if pr := ps.Prices(name); pr != nil {
			payload["prices"] = pr
		}
		writeJSON(w, http.StatusOK, payload)
	case http.MethodPut:
		if reqRole(r.Context()) != "admin" {
			writeError(w, http.StatusForbidden, "forbidden", "admin role required")
			return
		}
		var state interface{}
		if err := json.NewDecoder(r.Body).Decode(&state); err != nil {
			writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
			return
		}
		if err := ps.SetState(name, state); err != nil {
			writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
			return
		}
		writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
	default:
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
	}
}

func (g *Gateway) pluginDir() string {
	if c := g.core.Config(); c != nil {
		return c.PluginDir
	}
	return ""
}

// installPlugin writes a plugin to disk and hot-loads it. A syntax error is
// returned to the caller AND the file is left on disk so the operator can fix
// it, matching how adapters behave (the file is authoritative once present).
func (g *Gateway) installPlugin(name, code string) error {
	if err := validPluginName(name); err != nil {
		return err
	}
	dir := g.pluginDir()
	if dir == "" {
		return errNoPluginDir
	}
	if err := os.MkdirAll(dir, 0755); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, name+".lua"), []byte(code), 0644); err != nil {
		return err
	}
	return g.core.Plugins().LoadSource(name, code)
}

func (g *Gateway) removePlugin(name string) error {
	if err := validPluginName(name); err != nil {
		return err
	}
	if g.pluginDir() == "" {
		return errNoPluginDir
	}
	_ = os.Remove(filepath.Join(g.pluginDir(), name+".lua"))
	return g.core.Plugins().Unload(name)
}