ModelRouter/internal/core/kernel_actions.go

package core

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"llmsproxy/internal/config"
	"llmsproxy/internal/lua"
)

// 内核动作实现（R4-c）
//
// 把 Core 的公开动作暴露给插件 Lua，走 lua.KernelAPI 接口注入（依赖方向是
// core → lua，所以接口定义在 lua 包、实现放在这里）。
//
// # 为什么是"动作面"而不是"给插件一个 Core 指针"
//
// 直接暴露 *Core 会让插件能碰任何未导出的细节，而且每一次内核重构都会变成
// 对插件作者的破坏性变更。动作面是一份**显式的契约**：每个动作有名字、危险级、
// 超时预算，加动作是一次有意的行为，而不是顺手多导出一个方法。
//
// # 凭据边界
//
// 这里**不实现** list_keys / find_key。Core.ListKeys() 返回的是**明文**网关
// 密钥，把它交给第三方 Lua 等于把部署的整个凭据面交出去。key 相关的读走掩码
// 形态（create_key 的返回是唯一一次明文，且是"刚创建的那一个"）。
var _ lua.KernelAPI = (*Core)(nil)

// Actions implements lua.KernelAPI.
//
// The catalog lives in the lua package because permissions must be checkable even
// when no Core is wired. This method exists so the interface is satisfied by the
// real kernel rather than by a separate adapter type.
func (c *Core) Actions() []lua.KernelAction { return lua.KernelActions() }

// Invoke implements lua.KernelAPI. It dispatches one action.
//
// Every action is expected to honour ctx. The caller (lua.callKernel) gives up on
// a deadline and returns an error to the plugin, but a kernel action that keeps
// running afterwards would be an invisible leak, so the long ones (probe_all)
// pass ctx down into the network layer.
func (c *Core) Invoke(ctx context.Context, name string, args map[string]interface{}) (interface{}, error) {
	switch name {
	// ---- 只读 ----
	case "status":
		return c.actionStatus(), nil
	case "list_sources":
		return c.actionListSources(), nil
	case "list_models":
		return c.actionListModels(), nil
	case "list_adapters":
		return c.actionListAdapters(), nil
	case "list_auto_rules":
		return c.actionListAutoRules(), nil
	case "adapter_owner":
		return c.actionAdapterOwner(args)

	// ---- 写 ----
	case "reload":
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		return map[string]interface{}{"ok": true}, c.Reload()
	case "reset_health":
		// ResetHealth already no-ops on a nil registry (it ranges over
		// Providers()), so this needs no guard beyond the nil check inside.
		c.ResetHealth()
		return map[string]interface{}{"ok": true}, nil
	case "probe_all":
		r := c.Registry()
		if r == nil {
			return nil, fmt.Errorf("the registry is not built yet")
		}
		r.ProbeAll(ctx)
		if err := ctx.Err(); err != nil {
			return nil, fmt.Errorf("probe_all was cut short: %w", err)
		}
		return map[string]interface{}{"ok": true}, nil
	case "save_auto_rules":
		return c.actionSaveAutoRules(args)
	case "save_key_auto":
		return c.actionSaveKeyAuto(args)
	case "add_source":
		return c.actionAddSource(args)
	case "create_key":
		return c.actionCreateKey(args)

	// ---- 破坏性 ----
	case "remove_source":
		return c.actionRemoveSource(args)
	case "delete_key":
		return c.actionDeleteKey(args)
	case "update_key":
		return c.actionUpdateKey(args)
	case "remove_adapter":
		return c.actionRemoveAdapter(args)
	case "set_plugin_enabled":
		return c.actionSetPluginEnabled(args)
	default:
		return nil, fmt.Errorf("unknown kernel action %q", name)
	}
}

// ---- 只读实现 ----

func (c *Core) actionStatus() interface{} {
	// A plugin can call this DURING LoadDir (its own __init or first hook), which
	// runs before rebuildRegistry builds the registry. Returning an empty status
	// rather than panicking keeps a load-time call from taking the process down;
	// the plugin sees an empty list and the next call gives real data.
	out := map[string]interface{}{
		"sources":    []interface{}{},
		"auto_slots": []interface{}{},
		"default":    c.cfg.DefaultModel,
		"ready":      false,
	}
	if r := c.Registry(); r != nil {
		out["sources"] = r.Status()
		out["ready"] = true
	}
	if slots := c.AutoSlotStates(); slots != nil {
		out["auto_slots"] = slots
	}
	return out
}

// actionListSources returns the sources with api_key MASKED.
//
// Masking is not optional: a plugin has no legitimate need for another source's
// key, and a plugin that wrote the source itself already has it in its own
// credential vault.
func (c *Core) actionListSources() interface{} {
	srcs := c.Sources()
	out := make([]map[string]interface{}, 0, len(srcs))
	for _, s := range srcs {
		models := make([]string, 0, len(s.Models))
		for _, m := range s.Models {
			models = append(models, m.ID)
		}
		out = append(out, map[string]interface{}{
			"name":           s.Name,
			"base_url":       s.BaseURL,
			"adapter":        s.Adapter,
			"models":         models,
			"has_key":        s.APIKey != "" || s.APIKeyEnv != "",
			"key_hint":       maskSecretHint(s.APIKey),
			"max_concurrent": s.MaxConcurrent,
			// rpm must be exposed too: it carries the per-model minimum gap, and
			// max_concurrent alone cannot distinguish "serial WITH a gap" from
			// "serial with no gap" — the latter still collides with the
			// upstream's per-model concurrency window.
			"rpm": s.RPM,
		})
	}
	return out
}

// maskSecretHint shows only a short suffix so an operator (and a plugin) can tell
// two keys apart without either getting the secret.
func maskSecretHint(k string) string {
	if k == "" {
		return ""
	}
	if len(k) <= 4 {
		return "***"
	}
	return "***" + k[len(k)-4:]
}

func (c *Core) actionListModels() interface{} {
	if r := c.Registry(); r != nil {
		return r.ModelList()
	}
	return []string{}
}

func (c *Core) actionListAdapters() interface{} {
	return c.ListAdapters()
}

func (c *Core) actionListAutoRules() interface{} {
	return map[string]interface{}{
		"auto":       c.AutoRules(),
		"auto_image": c.AutoImageRules(),
	}
}

func (c *Core) actionAdapterOwner(args map[string]interface{}) (interface{}, error) {
	name, _ := args["name"].(string)
	if name == "" {
		return nil, fmt.Errorf("adapter_owner requires a name")
	}
	owner := ""
	if ps := c.Plugins(); ps != nil {
		owner = ps.AdapterOwnerOf(name)
	}
	return map[string]interface{}{
		"name":  name,
		"owner": owner, // "" means operator-installed
	}, nil
}

// ---- 写实现 ----

func (c *Core) actionSaveAutoRules(args map[string]interface{}) (interface{}, error) {
	entries, err := decodeScopes(args["entries"])
	if err != nil {
		return nil, err
	}
	if err := c.SaveAutoRules(entries); err != nil {
		return nil, err
	}
	return map[string]interface{}{"ok": true, "count": len(entries)}, nil
}

func (c *Core) actionSaveKeyAuto(args map[string]interface{}) (interface{}, error) {
	keyName, _ := args["key_name"].(string)
	if keyName == "" {
		return nil, fmt.Errorf("save_key_auto requires key_name")
	}
	entries, err := decodeScopes(args["entries"])
	if err != nil {
		return nil, err
	}
	if err := c.SaveKeyAuto(keyName, entries); err != nil {
		return nil, err
	}
	return map[string]interface{}{"ok": true, "count": len(entries)}, nil
}

// actionAddSource is the capability that makes "plugin-driven upstream" real: a
// plugin that has just obtained credentials can wire its endpoint in as a source
// without the operator editing config.yaml.
func (c *Core) actionAddSource(args map[string]interface{}) (interface{}, error) {
	raw, ok := args["source"]
	if !ok {
		return nil, fmt.Errorf("add_source requires a source table")
	}
	b, err := json.Marshal(raw)
	if err != nil {
		return nil, fmt.Errorf("source is not encodable: %w", err)
	}
	var src config.Source
	if err := json.Unmarshal(b, &src); err != nil {
		return nil, fmt.Errorf("source shape is invalid: %w", err)
	}
	if strings.TrimSpace(src.Name) == "" {
		return nil, fmt.Errorf("source requires a name")
	}
	if strings.TrimSpace(src.BaseURL) == "" {
		return nil, fmt.Errorf("source %s requires a base_url", src.Name)
	}
	if err := c.AddSource(src); err != nil {
		return nil, err
	}
	return map[string]interface{}{"ok": true, "name": src.Name}, nil
}

// actionCreateKey creates a gateway key.
//
// The plaintext key is returned ONCE, in this response — the same posture as the
// admin API. It is not retrievable later, so a plugin that needs it must store it
// in its own credential vault immediately.
func (c *Core) actionCreateKey(args map[string]interface{}) (interface{}, error) {
	name, _ := args["name"].(string)
	role, _ := args["role"].(string)
	if role == "" {
		role = "user"
	}
	models, err := decodeScopes(args["models"])
	if err != nil {
		return nil, err
	}
	k, err := c.CreateKey(name, role, models, "created by plugin")
	if err != nil {
		return nil, err
	}
	return map[string]interface{}{
		"key":    k.Key,
		"name":   k.Name,
		"role":   k.Role,
		"models": k.Models,
	}, nil
}

// ---- 破坏性实现 ----

func (c *Core) actionRemoveSource(args map[string]interface{}) (interface{}, error) {
	name, _ := args["name"].(string)
	if name == "" {
		return nil, fmt.Errorf("remove_source requires a name")
	}
	// The last source is not removable: an empty registry turns every request
	// into a 503, and the operator would have to fix it from the filesystem.
	if len(c.Sources()) <= 1 {
		return nil, fmt.Errorf("refusing to remove the last source: the gateway would have nothing to route to")
	}
	if err := c.RemoveSource(name); err != nil {
		return nil, err
	}
	return map[string]interface{}{"ok": true, "name": name}, nil
}

func (c *Core) actionDeleteKey(args map[string]interface{}) (interface{}, error) {
	key, _ := args["key"].(string)
	if key == "" {
		return nil, fmt.Errorf("delete_key requires the key to delete")
	}
	ok, err := c.DeleteKey(key)
	if err != nil {
		return nil, err
	}
	if !ok {
		return nil, fmt.Errorf("no such key")
	}
	return map[string]interface{}{"ok": true}, nil
}

func (c *Core) actionUpdateKey(args map[string]interface{}) (interface{}, error) {
	key, _ := args["key"].(string)
	if key == "" {
		return nil, fmt.Errorf("update_key requires the key")
	}
	name, _ := args["name"].(string)
	role, _ := args["role"].(string)
	models, err := decodeScopes(args["models"])
	if err != nil {
		return nil, err
	}
	k, err := c.UpdateKey(key, name, role, models, "")
	if err != nil {
		return nil, err
	}
	return map[string]interface{}{"ok": true, "name": k.Name, "role": k.Role}, nil
}

// actionRemoveAdapter deletes an adapter, but only one the CALLER plugin owns —
// the ownership check lives in Plugins.RemoveAdapter (see the adapter bridge).
// A plugin can always remove its own; removing an operator's or a built-in is
// refused there, and that refusal is the security property.
func (c *Core) actionRemoveAdapter(args map[string]interface{}) (interface{}, error) {
	name, _ := args["name"].(string)
	if name == "" {
		return nil, fmt.Errorf("remove_adapter requires a name")
	}
	ps := c.Plugins()
	if ps == nil {
		return nil, fmt.Errorf("no plugins loaded")
	}
	plugin, _ := args["plugin"].(string)
	if plugin == "" {
		return nil, fmt.Errorf("remove_adapter needs the calling plugin's name")
	}
	if err := ps.RemoveAdapter(plugin, name); err != nil {
		return nil, err
	}
	return map[string]interface{}{"ok": true, "name": name}, nil
}

func (c *Core) actionSetPluginEnabled(args map[string]interface{}) (interface{}, error) {
	name, _ := args["name"].(string)
	if name == "" {
		return nil, fmt.Errorf("set_plugin_enabled requires a name")
	}
	enabled, ok := args["enabled"].(bool)
	if !ok {
		return nil, fmt.Errorf("set_plugin_enabled requires enabled (boolean)")
	}
	ps := c.Plugins()
	if ps == nil {
		return nil, fmt.Errorf("no plugins loaded")
	}
	if err := ps.SetEnabled(name, enabled); err != nil {
		return nil, err
	}
	return map[string]interface{}{"ok": true, "name": name, "enabled": enabled}, nil
}

// decodeScopes converts a Lua-decoded array of model-scope tables into
// []config.ModelScope. It accepts both an explicit array and a single table so a
// one-entry chain does not need a wrapper.
func decodeScopes(v interface{}) ([]config.ModelScope, error) {
	if v == nil {
		return nil, nil
	}
	b, err := json.Marshal(v)
	if err != nil {
		return nil, fmt.Errorf("entries are not encodable: %w", err)
	}
	var list []config.ModelScope
	if err := json.Unmarshal(b, &list); err == nil {
		return list, nil
	}
	var one config.ModelScope
	if err := json.Unmarshal(b, &one); err != nil {
		return nil, fmt.Errorf("entries shape is invalid: %w", err)
	}
	return []config.ModelScope{one}, nil
}