ModelRouter/internal/lua/vm.go

// Package lua implements the adapter runtime: bundles/loads *.lua adapter
// scripts, exposes json/string helpers, and lets Go call the protocol
// transform functions plus a build_headers signature hook.
//
// The runtime uses LuaJIT (github.com/aarzilli/golua) so every adapter runs on
// an independent interpreter state. Because a lua.State is NOT goroutine-safe,
// each adapter owns a pool of workers (states); the pool size equals the sum
// of max_concurrent over all sources that use the adapter, and a worker is
// checked out for the duration of a single hook call.
package lua

import (
	"crypto/hmac"
	"crypto/sha256"
	"embed"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"

	// golua exposes the LuaJIT C API. Build with -tags luajit so the cgo
	// LDFLAGS resolve to -lluajit-5.1 (see golua's lua.go).
	golua "github.com/aarzilli/golua/lua"
)

//go:embed adapters/*.lua
var bundledAdapters embed.FS

// adapterGlobal is the reserved global holding the adapter table after the
// worker boots the script. A reserved name avoids clashing with user code.
const adapterGlobal = "__llmsproxy_adapter"

type APIAdapter struct {
	Name    string `json:"name"`
	Version string `json:"version"`
}

// worker wraps one LuaJIT interpreter state. Not goroutine-safe; each worker
// is used by exactly one goroutine at a time, checked out from its pool.
type worker struct {
	L *golua.State
}

func (w *worker) close() {
	if w.L != nil {
		w.L.Close()
	}
}

// staticInfo caches the immutable fields of an adapter script, extracted once
// at load time. Endpoint() and the static headers fallback read from here so
// they never churn a pooled worker.
type staticInfo struct {
	name     string
	version  string
	endpoint string
	headers  map[string]string
}

// adapterPool owns the worker states for one adapter. Idle workers are held in
// a slice guarded by a cond; up to target workers are created lazily.
type adapterPool struct {
	name   string
	script string
	static staticInfo

	mu      sync.Mutex
	cond    *sync.Cond
	idle    []*worker
	created int
	target  int
	closed  bool
}

func newAdapterPool(name, script string, static staticInfo) *adapterPool {
	p := &adapterPool{name: name, script: script, static: static, target: 1}
	p.cond = sync.NewCond(&p.mu)
	return p
}

func (p *adapterPool) setTarget(n int) {
	p.mu.Lock()
	p.target = n
	p.cond.Broadcast()
	p.mu.Unlock()
}

// boot creates a fresh LuaJIT state, registers the shared globals and executes
// the adapter script, storing the returned adapter table in the state.
func (p *adapterPool) boot() (*worker, error) {
	L := golua.NewState()
	L.OpenLibs()
	setupGlobals(L)
	if err := L.DoString(p.script); err != nil {
		L.Close()
		return nil, fmt.Errorf("compile adapter %s: %w", p.name, err)
	}
	if L.Type(-1) != golua.LUA_TTABLE {
		L.Close()
		return nil, fmt.Errorf("adapter %s must return a table", p.name)
	}
	L.SetGlobal(adapterGlobal)
	L.SetTop(0)
	return &worker{L: L}, nil
}

func (p *adapterPool) acquire() (*worker, error) {
	p.mu.Lock()
	for {
		if p.closed {
			p.mu.Unlock()
			return nil, fmt.Errorf("adapter %s pool closed", p.name)
		}
		if n := len(p.idle); n > 0 {
			w := p.idle[n-1]
			p.idle = p.idle[:n-1]
			p.mu.Unlock()
			return w, nil
		}
		if p.created < p.target {
			p.created++
			break
		}
		p.cond.Wait()
	}
	p.mu.Unlock()

	w, err := p.boot()
	if err != nil {
		p.mu.Lock()
		p.created--
		p.cond.Signal()
		p.mu.Unlock()
		return nil, err
	}
	return w, nil
}

func (p *adapterPool) release(w *worker) {
	w.L.SetTop(0)
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		w.close()
		return
	}
	p.idle = append(p.idle, w)
	p.cond.Signal()
	p.mu.Unlock()
}

func (p *adapterPool) shutdown() {
	p.mu.Lock()
	p.closed = true
	idle := p.idle
	p.idle = nil
	p.cond.Broadcast()
	p.mu.Unlock()
	for _, w := range idle {
		w.close()
	}
}

// VM aggregates per-adapter worker pools. All exported methods are safe for
// concurrent use.
type VM struct {
	mu    sync.RWMutex
	dir   string
	pools map[string]*adapterPool
}

func NewVM(dir string) *VM {
	return &VM{dir: dir, pools: map[string]*adapterPool{}}
}

func (v *VM) Start() error {
	if v.dir == "" {
		return nil
	}
	if err := os.MkdirAll(v.dir, 0755); err != nil {
		return fmt.Errorf("mkdir adapter dir: %w", err)
	}
	if err := v.writeBundledAdapters(); err != nil {
		return err
	}
	entries, err := os.ReadDir(v.dir)
	if err != nil {
		return err
	}
	for _, e := range entries {
		if filepath.Ext(e.Name()) != ".lua" {
			continue
		}
		if err := v.LoadAdapter(filepath.Join(v.dir, e.Name())); err != nil {
			fmt.Printf("[lua] preload %s: %v\n", e.Name(), err)
		}
	}
	return nil
}

func (v *VM) Stop() {
	v.mu.Lock()
	pools := v.pools
	v.pools = map[string]*adapterPool{}
	v.mu.Unlock()
	for _, p := range pools {
		p.shutdown()
	}
}

// ConfigureConcurrency sets each adapter's worker target to the sum of
// max_concurrent of every source using it. Values below one are clamped to 1.
func (v *VM) ConfigureConcurrency(adapterConcurrency map[string]int) {
	v.mu.RLock()
	defer v.mu.RUnlock()
	for name, n := range adapterConcurrency {
		if p, ok := v.pools[name]; ok {
			if n < 1 {
				n = 1
			}
			p.setTarget(n)
		}
	}
}

func (v *VM) ListAdapters() []APIAdapter {
	v.mu.RLock()
	defer v.mu.RUnlock()
	list := make([]APIAdapter, 0, len(v.pools))
	for name, p := range v.pools {
		list = append(list, APIAdapter{Name: name, Version: p.static.version})
	}
	sort.Slice(list, func(i, j int) bool { return list[i].Name < list[j].Name })
	return list
}

// LoadAdapter compiles and registers an adapter script from a file.
func (v *VM) LoadAdapter(path string) error {
	data, err := os.ReadFile(path)
	if err != nil {
		return fmt.Errorf("read adapter: %w", err)
	}
	name := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path))
	return v.LoadAdapterSource(name, string(data))
}

// LoadAdapterSource registers an adapter from source code. Replacing an
// existing adapter shuts the old pool down and re-boots lazily.
func (v *VM) LoadAdapterSource(name, code string) error {
	static, err := inspectScript(name, code)
	if err != nil {
		return err
	}
	target := 1
	v.mu.RLock()
	if p, ok := v.pools[static.name]; ok {
		target = p.target
	}
	v.mu.RUnlock()

	v.mu.Lock()
	if p, ok := v.pools[static.name]; ok {
		p.shutdown()
	}
	p := newAdapterPool(static.name, code, *static)
	p.setTarget(target)
	v.pools[static.name] = p
	v.mu.Unlock()
	return nil
}

// RemoveAdapter evicts an adapter and shuts its workers down.
func (v *VM) RemoveAdapter(name string) {
	v.mu.Lock()
	p := v.pools[name]
	delete(v.pools, name)
	v.mu.Unlock()
	if p != nil {
		p.shutdown()
	}
}

func (v *VM) pool(name string) *adapterPool {
	v.mu.RLock()
	defer v.mu.RUnlock()
	return v.pools[name]
}

func (v *VM) writeBundledAdapters() error {
	known := []string{"openai", "anthropic", "deepseek", "gemini", "github", "groq", "mistral", "ollama", "kimicode"}
	for _, name := range known {
		dst := filepath.Join(v.dir, name+".lua")
		if _, err := os.Stat(dst); err == nil {
			continue
		}
		data, err := bundledAdapters.ReadFile("adapters/" + name + ".lua")
		if err != nil {
			continue
		}
		if err := os.WriteFile(dst, data, 0644); err != nil {
			return err
		}
	}
	return nil
}

// Transform runs adapter.Fn(raw) and returns the resulting string.
func (v *VM) Transform(name, fn, raw string) (string, error) {
	p := v.pool(name)
	if p == nil {
		return "", fmt.Errorf("adapter %s not loaded", name)
	}
	w, err := p.acquire()
	if err != nil {
		return "", err
	}
	defer p.release(w)
	L := w.L

	L.SetTop(0)
	defer L.SetTop(0)

	L.GetGlobal(adapterGlobal)
	if L.IsNil(-1) {
		return "", fmt.Errorf("adapter %s has no table", name)
	}
	L.GetField(-1, fn)
	if L.IsNil(-1) || !L.IsFunction(-1) {
		return "", fmt.Errorf("adapter %s missing %s", name, fn)
	}
	L.PushString(raw)
	if err := L.Call(1, 1); err != nil {
		return "", fmt.Errorf("%s: %w", fn, err)
	}
	return L.ToString(-1), nil
}

// BuildHeaders executes adapter.build_headers(meta). If the adapter defines no
// build_headers hook, it falls back to the static adapter.headers table.
func (v *VM) BuildHeaders(name string, meta map[string]interface{}) (map[string]string, error) {
	p := v.pool(name)
	if p == nil {
		return nil, fmt.Errorf("adapter %s not loaded", name)
	}

	w, err := p.acquire()
	if err != nil {
		return nil, err
	}
	defer p.release(w)
	L := w.L

	L.SetTop(0)
	defer L.SetTop(0)

	L.GetGlobal(adapterGlobal)
	if L.IsNil(-1) {
		return nil, fmt.Errorf("adapter %s has no table", name)
	}
	L.GetField(-1, "build_headers")
	if !L.IsFunction(-1) {
		// no hook: fall back to the static headers captured at load time
		p.mu.Lock()
		hdr := make(map[string]string, len(p.static.headers))
		for k, h := range p.static.headers {
			hdr[k] = h
		}
		p.mu.Unlock()
		return hdr, nil
	}

	L.SetTop(0)
	L.GetGlobal(adapterGlobal)
	L.GetField(-1, "build_headers")
	pushGoValue(L, meta)
	if err := L.Call(1, 1); err != nil {
		return nil, fmt.Errorf("build_headers: %w", err)
	}
	if L.Type(-1) != golua.LUA_TTABLE {
		return nil, fmt.Errorf("build_headers returned non-table")
	}
	headers := map[string]string{}
	luaTableToMap(L, -1, headers)
	return headers, nil
}

func (v *VM) Endpoint(name string) string {
	p := v.pool(name)
	if p == nil {
		return ""
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.static.endpoint
}

// ---- static inspection (compile-once at load time) ----

func inspectScript(name, code string) (*staticInfo, error) {
	L := golua.NewState()
	L.OpenLibs()
	setupGlobals(L)
	if err := L.DoString(code); err != nil {
		L.Close()
		return nil, fmt.Errorf("compile adapter: %w", err)
	}
	if L.Type(-1) != golua.LUA_TTABLE {
		L.Close()
		return nil, fmt.Errorf("adapter script must return a table")
	}
	info := &staticInfo{name: name, headers: map[string]string{}}
	if n := readString(L, -1, "name"); n != "" {
		info.name = n
	}
	info.version = readString(L, -1, "version")
	info.endpoint = readString(L, -1, "endpoint")
	if getFieldBatch(L, "headers") {
		luaTableToMap(L, -1, info.headers)
		L.Pop(1)
	} else {
		L.Pop(1)
	}
	L.Close()
	return info, nil
}

// getFieldBatch fetches field of the table at the top of the stack and returns
// true when the fetched value (now at the top) is a table.
func getFieldBatch(L *golua.State, field string) bool {
	L.GetField(-1, field)
	return L.Type(-1) == golua.LUA_TTABLE
}

func readString(L *golua.State, tableIdx int, field string) string {
	L.GetField(tableIdx, field)
	defer L.Pop(1)
	if L.Type(-1) == golua.LUA_TNIL || L.Type(-1) == golua.LUA_TNONE {
		return ""
	}
	return normalizeString(L, -1)
}

func normalizeString(L *golua.State, idx int) string {
	if L.Type(idx) == golua.LUA_TNUMBER {
		n := L.ToNumber(idx)
		return strconv.FormatFloat(n, 'f', -1, 64)
	}
	return L.ToString(idx)
}

// ---- shared globals installed into every worker state ----

// restorePcall puts the real pcall/xpcall back in place. golua's OpenLibs
// replaces them with unsafe_ variants (it captures errors via its own panic
// handler); the built-in adapters rely on pcall for soft error handling.
func restorePcall(L *golua.State) {
	for _, name := range []string{"pcall", "xpcall"} {
		L.GetGlobal("unsafe_" + name)
		if !L.IsNil(-1) {
			L.SetGlobal(name)
		} else {
			L.Pop(1)
		}
	}
}

func setupGlobals(L *golua.State) {
	restorePcall(L)
	buildJSONTable(L)

	registerFn(L, "hmac_sha256_hex", func(L *golua.State) int {
		key := L.ToString(1)
		data := L.ToString(2)
		m := hmac.New(sha256.New, []byte(key))
		m.Write([]byte(data))
		L.PushString(hex.EncodeToString(m.Sum(nil)))
		return 1
	})
	registerFn(L, "sha256_hex", func(L *golua.State) int {
		h := sha256.Sum256([]byte(L.ToString(1)))
		L.PushString(hex.EncodeToString(h[:]))
		return 1
	})
	registerFn(L, "base64_encode", func(L *golua.State) int {
		L.PushString(base64.StdEncoding.EncodeToString([]byte(L.ToString(1))))
		return 1
	})
	registerFn(L, "tohex", func(L *golua.State) int {
		L.PushString(hex.EncodeToString([]byte(L.ToString(1))))
		return 1
	})
	registerFn(L, "log", func(L *golua.State) int {
		fmt.Printf("[adapter/%s] %s\n", L.ToString(1), L.ToString(2))
		return 0
	})
}

func buildJSONTable(L *golua.State) {
	L.NewTable()
	for _, name := range []string{"encode", "decode"} {
		var fn golua.LuaGoFunction
		switch name {
		case "encode":
			fn = func(L *golua.State) int {
				b, err := jsonEncode(luaValueToGo(L, 1))
				if err != nil {
					L.PushString("null")
					return 1
				}
				L.PushString(string(b))
				return 1
			}
		case "decode":
			fn = func(L *golua.State) int {
				v, err := jsonDecode(L.ToString(1))
				if err != nil {
					L.PushNil()
					return 1
				}
				pushGoValue(L, v)
				return 1
			}
		}
		L.PushGoClosure(fn)
		L.SetField(-2, name)
	}
	L.SetGlobal("json")
}

func registerFn(L *golua.State, name string, fn golua.LuaGoFunction) {
	L.PushGoClosure(fn)
	L.SetGlobal(name)
}

// ---- value conversions ----

func luaValueToGo(L *golua.State, idx int) interface{} {
	switch L.Type(idx) {
	case golua.LUA_TSTRING:
		return L.ToString(idx)
	case golua.LUA_TNUMBER:
		return L.ToNumber(idx)
	case golua.LUA_TBOOLEAN:
		return L.ToBoolean(idx)
	case golua.LUA_TTABLE:
		return luaTableToGoValue(L, idx)
	default:
		return nil
	}
}

// luaTableToGo value converts a Lua table to either []interface{} (array 1..n)
// or map[string]interface{}.
func luaTableToGoValue(L *golua.State, idx int) interface{} {
	abs := idx
	if idx < 0 {
		abs = L.GetTop() + idx + 1
	}
	m := map[string]interface{}{}
	allInt, n, maxK := true, 0, 0
	L.PushNil()
	for L.Next(abs) != 0 {
		kt := L.Type(-2)
		if kt == golua.LUA_TNUMBER {
			k := int(L.ToNumber(-2))
			if k > maxK {
				maxK = k
			}
		} else {
			allInt = false
		}
		m[keyToString(L, -2)] = luaValueToGo(L, -1)
		L.Pop(1)
		n++
	}
	if allInt && maxK > 0 && n == maxK {
		arr := make([]interface{}, maxK)
		for i := 1; i <= maxK; i++ {
			arr[i-1] = m[strconv.Itoa(i)]
		}
		return arr
	}
	return m
}

func luaTableToMap(L *golua.State, idx int, dst map[string]string) {
	abs := idx
	if idx < 0 {
		abs = L.GetTop() + idx + 1
	}
	L.PushNil()
	for L.Next(abs) != 0 {
		dst[keyToString(L, -2)] = normalizeString(L, -1)
		L.Pop(1)
	}
}

// keyToString reads the stack value at idx as a string without modifying the
// original (golua's ToString mutates numeric stack values in place, which
// would break the next lua_next call), then pops the temporary copy.
func keyToString(L *golua.State, idx int) string {
	L.PushValue(idx)
	defer L.Pop(1)
	return normalizeString(L, -1)
}

func pushGoValue(L *golua.State, v interface{}) {
	switch x := v.(type) {
	case nil:
		L.PushNil()
	case string:
		L.PushString(x)
	case bool:
		L.PushBoolean(x)
	case int:
		L.PushNumber(float64(x))
	case int64:
		L.PushNumber(float64(x))
	case float64:
		L.PushNumber(x)
	case []interface{}:
		L.NewTable()
		for i, it := range x {
			pushGoValue(L, it)
			L.RawSeti(-2, i+1)
		}
	case map[string]interface{}:
		L.NewTable()
		for k, it := range x {
			pushGoValue(L, it)
			L.SetField(-2, k)
		}
	default:
		if b, err := json.Marshal(x); err == nil {
			var iv interface{}
			if json.Unmarshal(b, &iv) == nil {
				pushGoValue(L, iv)
				return
			}
		}
		L.PushNil()
	}
}

func jsonEncode(v interface{}) ([]byte, error) { return json.Marshal(v) }
func jsonDecode(s string) (interface{}, error) {
	var v interface{}
	if err := json.Unmarshal([]byte(s), &v); err != nil {
		return nil, err
	}
	return v, nil
}