ModelRouter/internal/lua/vm_test.go

package lua

import (
	"encoding/json"
	"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"} {
		if !names[want] {
			t.Errorf("missing adapter %s (got %v)", want, names)
		}
	}
}

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 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)
	}
}