ModelRouter/internal/lua/vm_test.go

package lua

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// freshAdapterDir returns a path under a temp dir that does not exist yet, so
// VM.Start() treats it as first-run and seeds the bundled adapters.
func freshAdapterDir(t *testing.T) string {
	t.Helper()
	return filepath.Join(t.TempDir(), "adapters")
}

func TestLoadBundledAdapters(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatalf("start: %v", err)
	}
	defer vm.Stop()
	adapters := vm.ListAdapters()
	if len(adapters) == 0 {
		t.Fatal("no adapters loaded")
	}
	names := map[string]bool{}
	for _, a := range adapters {
		names[a.Name] = true
	}
	for _, want := range []string{"openai", "deepseek", "anthropic", "gemini", "ollama", "kimicode", "opencode"} {
		if !names[want] {
			t.Errorf("missing adapter %s (got %v)", want, names)
		}
	}
}

func TestFirstRunSeedsEveryBundledAdapter(t *testing.T) {
	dir := filepath.Join(t.TempDir(), "adapters")
	vm := NewVM(dir)
	if err := vm.Start(); err != nil {
		t.Fatalf("start: %v", err)
	}
	defer vm.Stop()
	// every embedded .lua must have been copied out (and no extra, non-.lua files)
	embedded, err := bundledAdapters.ReadDir("adapters")
	if err != nil {
		t.Fatal(err)
	}
	seeded := 0
	for _, e := range embedded {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".lua") {
			continue
		}
		seeded++
		if _, err := os.Stat(filepath.Join(dir, e.Name())); err != nil {
			t.Errorf("embedded adapter %s not seeded", e.Name())
		}
	}
	if seeded == 0 {
		t.Fatal("no embedded adapters found")
	}
	written := 0
	entries, _ := os.ReadDir(dir)
	for _, e := range entries {
		if !e.IsDir() && strings.HasSuffix(e.Name(), ".lua") {
			written++
		}
	}
	if written != seeded {
		t.Fatalf("seeded %d files but %d embedded adapters exist", written, seeded)
	}
	// re-start with the existing dir: authoritative, never rewritten
	vm2 := NewVM(dir)
	if err := vm2.Start(); err != nil {
		t.Fatalf("second start: %v", err)
	}
	defer vm2.Stop()
}

func TestTransformRequest(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	out, err := vm.Transform("openai", "transform_request", `{"model":"x","disable_thinking":true,"messages":[]}`)
	if err != nil {
		t.Fatalf("transform: %v", err)
	}
	if strings.Contains(out, "disable_thinking") {
		t.Fatalf("disable_thinking not stripped: %s", out)
	}
}

func TestBuildHeadersFallbackStatic(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	hdrs, err := vm.BuildHeaders("anthropic", nil)
	if err != nil {
		t.Fatalf("build headers: %v", err)
	}
	if hdrs["anthropic-version"] != "2023-06-01" {
		t.Fatalf("static header missing: %v", hdrs)
	}
}

func TestOpenCodeAdapterFingerprint(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	hdrs, err := vm.BuildHeaders("opencode", map[string]interface{}{"api_key": "public"})
	if err != nil {
		t.Fatalf("build headers: %v", err)
	}
	if ua := hdrs["User-Agent"]; ua != "opencode/0.1.0" {
		t.Fatalf("opencode adapter must send the opencode client UA, got %q", ua)
	}
	if hdrs["Authorization"] != "" {
		t.Fatalf("opencode adapter must not hardcode Authorization (config api_key supplies it), got %q", hdrs["Authorization"])
	}
	// passthrough behaves like the openai adapter
	out, err := vm.Transform("opencode", "transform_request", `{"model":"x","disable_thinking":true,"extra_body":{},"messages":[{"role":"user","content":"hi"}]}`)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(out, "disable_thinking") || strings.Contains(out, "extra_body") {
		t.Fatalf("opencode transform_request must strip provider fields: %s", out)
	}
}

func TestOpenCodeStripsMultiModalParts(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	body := `{"model":"x","messages":[
		{"role":"user","content":[
			{"type":"text","text":"describe"},
			{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}},
			{"type":"input_audio","input_audio":{"data":"QQ","format":"wav"}}
		]},
		{"role":"user","content":[{"type":"image_url","image_url":{"url":"https://ex.com/a.png"}}]},
		{"role":"assistant","content":"plain"}
	]}`
	out, err := vm.Transform("opencode", "transform_request", body)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(out, "image_url") || strings.Contains(out, "input_audio") {
		t.Fatalf("opencode must strip multimodal parts: %s", out)
	}
	if !strings.Contains(out, "describe") {
		t.Fatalf("text parts must survive: %s", out)
	}
	// image-only message dropped: only the text message and plain assistant message remain
	if got := strings.Count(out, `"role"`); got != 2 {
		t.Fatalf("image-only message must be dropped, got %d messages: %s", got, out)
	}
}

func TestBuildHeadersCustomHook(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	hdrs, err := vm.BuildHeaders("kimicode", map[string]interface{}{
		"timestamp": int64(12345),
		"api_key":   "k",
		"body":      "{}",
		"method":    "POST",
		"url":       "http://x/chat",
		"source": map[string]interface{}{"meta": map[string]interface{}{
			"app_id": "app-9", "app_secret": "s", "api_key": "k",
		}},
	})
	if err != nil {
		t.Fatalf("build headers: %v", err)
	}
	if hdrs["X-App-Id"] != "app-9" {
		t.Fatalf("x-app-id = %q", hdrs["X-App-Id"])
	}
	if hdrs["X-App-Sign"] == "" {
		t.Fatal("expected signature header")
	}
	if hdrs["X-Timestamp"] != "12345" {
		t.Fatalf("timestamp = %q", hdrs["X-Timestamp"])
	}
}

func TestDisableThinkingPassthrough(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	body := `{"model":"x","disable_thinking":true,"messages":[{"role":"user","content":"hi"}]}`
	out, err := vm.Transform("deepseek", "transform_request", body)
	if err != nil {
		t.Fatalf("deepseek transform: %v", err)
	}
	var req struct {
		ExtraBody map[string]interface{} `json:"extra_body"`
	}
	if err := json.Unmarshal([]byte(out), &req); err != nil {
		t.Fatalf("unmarshal: %v (%s)", err, out)
	}
	thinking, ok := req.ExtraBody["thinking"].(map[string]interface{})
	if !ok {
		t.Fatalf("deepseek should emit extra_body.thinking on disable_thinking: %s", out)
	}
	if thinking["type"] != "disabled" {
		t.Fatalf("thinking.type = %v", thinking["type"])
	}

	// anthropic: disable_thinking removes the thinking block
	out2, err := vm.Transform("anthropic", "transform_request", body)
	if err != nil {
		t.Fatalf("anthropic transform: %v", err)
	}
	if strings.Contains(out2, "thinking") {
		t.Fatalf("anthropic should drop thinking when disable_thinking: %s", out2)
	}
	out3, err := vm.Transform("anthropic", "transform_request", `{"model":"x","messages":[{"role":"user","content":"hi"}]}`)
	if err != nil {
		t.Fatalf("anthropic transform: %v", err)
	}
	if !strings.Contains(out3, "enabled") {
		t.Fatalf("anthropic should enable thinking by default: %s", out3)
	}
}

func TestMultimodalTransform(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatal(err)
	}
	defer vm.Stop()
	body := `{"model":"x","messages":[{"role":"user","content":[
		{"type":"text","text":"what is this?"},
		{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}},
		{"type":"image_url","image_url":{"url":"https://ex.com/a.png"}}
	]}]}`

	// anthropic: image_url -> image block (base64/url), text preserved
	out, err := vm.Transform("anthropic", "transform_request", body)
	if err != nil {
		t.Fatalf("anthropic: %v", err)
	}
	for _, want := range []string{`"media_type":"image/png"`, `"data":"QUJD"`, `"type":"url","url":"https://ex.com/a.png"`, `"what is this?"`} {
		if !strings.Contains(out, want) {
			t.Fatalf("anthropic multimodal missing %s: %s", want, out)
		}
	}

	// gemini: image_url -> inline_data
	gout, err := vm.Transform("gemini", "transform_request", body)
	if err != nil {
		t.Fatalf("gemini: %v", err)
	}
	var g struct {
		Contents []struct {
			Parts []map[string]interface{} `json:"parts"`
		} `json:"contents"`
	}
	if err := json.Unmarshal([]byte(gout), &g); err != nil {
		t.Fatalf("gemini unmarshal: %v", err)
	}
	if len(g.Contents) == 0 {
		t.Fatalf("gemini no contents")
	}
	var found bool
	for _, p := range g.Contents[0].Parts {
		if v, ok := p["inline_data"].(map[string]interface{}); ok && v["data"] == "QUJD" && v["mime_type"] == "image/png" {
			found = true
		}
	}
	if !found {
		t.Fatalf("gemini missing inline_data image: %s", gout)
	}

	// ollama: image_url -> images base64 array
	out, err = vm.Transform("ollama", "transform_request", body)
	if err != nil {
		t.Fatalf("ollama: %v", err)
	}
	if !strings.Contains(out, `"images":["QUJD"]`) {
		t.Fatalf("ollama multimodal missing images: %s", out)
	}

	// openai passthrough keeps the content array intact
	po, _ := vm.Transform("openai", "transform_request", body)
	if !strings.Contains(po, `"image_url"`) || !strings.Contains(po, `,QUJD"`) {
		t.Fatalf("openai passthrough lost multimodal content: %s", po)
	}
}

func TestOpenCodeAdapterNormalizesDeveloperRole(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatalf("start: %v", err)
	}
	defer vm.Stop()

	raw := `{"model":"zen","messages":[` +
		`{"role":"developer","content":"be concise"},` +
		`{"role":"user","content":"hi"},` +
		`{"role":"assistant","content":"hello"},` +
		`{"role":"function","content":"{\"a\":1}"},` +
		`{"role":"latest_reminder","content":"remind"}]}`
	out, err := vm.Transform("opencode", "transform_request", raw)
	if err != nil {
		t.Fatalf("transform_request: %v", err)
	}
	var req struct {
		Messages []struct {
			Role string `json:"role"`
		} `json:"messages"`
	}
	if err := json.Unmarshal([]byte(out), &req); err != nil {
		t.Fatalf("output not JSON: %v\n%s", err, out)
	}
	want := []string{"system", "user", "assistant", "system", "latest_reminder"}
	if len(req.Messages) != len(want) {
		t.Fatalf("got %d messages, want %d: %s", len(req.Messages), len(want), out)
	}
	for i, w := range want {
		if req.Messages[i].Role != w {
			t.Errorf("message %d role = %q, want %q", i, req.Messages[i].Role, w)
		}
	}
}

func TestOpenCodeAdapterKeepsWhitelistedRolesAndDropsMultimodal(t *testing.T) {
	vm := NewVM(freshAdapterDir(t))
	if err := vm.Start(); err != nil {
		t.Fatalf("start: %v", err)
	}
	defer vm.Stop()

	raw := `{"messages":[` +
		`{"role":"system","content":"sys"},` +
		`{"role":"user","content":[{"type":"text","text":"keep"},{"type":"image_url","url":"x"}]},` +
		`{"role":"user","content":[{"type":"image_url","url":"x"}]}]}`
	out, err := vm.Transform("opencode", "transform_request", raw)
	if err != nil {
		t.Fatalf("transform_request: %v", err)
	}
	var req struct {
		Messages []struct {
			Role    string `json:"role"`
			Content any    `json:"content"`
		} `json:"messages"`
	}
	if err := json.Unmarshal([]byte(out), &req); err != nil {
		t.Fatalf("output not JSON: %v\n%s", err, out)
	}
	if len(req.Messages) != 2 {
		t.Fatalf("got %d messages, want 2 (multimodal-only dropped): %s", len(req.Messages), out)
	}
	if req.Messages[0].Role != "system" || req.Messages[1].Role != "user" {
		t.Fatalf("unexpected roles: %s", out)
	}
}