yao/llmprovider/registry_test.go

package llmprovider_test

import (
	"fmt"
	"os"
	"sync"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/yaoapp/gou/connector"
	"github.com/yaoapp/gou/store"
	"github.com/yaoapp/yao/config"
	"github.com/yaoapp/yao/llmprovider"
	"github.com/yaoapp/yao/test"
)

func TestMain(m *testing.M) {
	test.Prepare(nil, config.Conf)
	defer test.Clean()
	os.Exit(m.Run())
}

func setupRegistry(t *testing.T) *llmprovider.Registry {
	t.Helper()
	test.Prepare(t, config.Conf)

	err := llmprovider.Init()
	require.NoError(t, err)

	t.Cleanup(func() {
		s, _ := store.Get("__yao.store")
		if s != nil {
			s.Del("llmprovider:*")
		}
		c, _ := store.Get("__yao.cache")
		if c != nil {
			c.Del("llmprovider:*")
		}
		test.Clean()
	})

	return llmprovider.Global
}

var testProvider = llmprovider.Provider{
	Key:        "test-openai",
	Name:       "Test OpenAI",
	Type:       "openai",
	APIURL:     "https://api.openai.com",
	APIKey:     "sk-test-xxxxx",
	Models:     []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"vision", "tool_calls", "streaming"}, Enabled: true}},
	Enabled:    true,
	RequireKey: true,
	Owner:      llmprovider.ProviderOwner{Type: "system"},
}

func TestCreate(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	created, err := r.Create(&p)
	require.NoError(t, err)
	assert.Equal(t, "test-openai", created.Key)
	assert.Equal(t, llmprovider.ProviderSourceDynamic, created.Source)
	assert.NotEmpty(t, created.ConnectorID)

	// Verify store persistence
	s, _ := store.Get("__yao.store")
	assert.True(t, s.Has("llmprovider:p:test-openai"))

	// Verify connector registered
	_, err = connector.Select(created.ConnectorID)
	assert.NoError(t, err)
}

func TestCreateDuplicate(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	dup := testProvider
	_, err = r.Create(&dup)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "already exists")
}

func TestGet(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	got, err := r.Get("test-openai")
	require.NoError(t, err)
	assert.Equal(t, "Test OpenAI", got.Name)
	assert.Equal(t, "openai", got.Type)
	assert.Equal(t, "https://api.openai.com", got.APIURL)
	assert.Len(t, got.Models, 1)
	assert.Equal(t, "gpt-4o", got.Models[0].ID)
}

func TestGetNotFound(t *testing.T) {
	r := setupRegistry(t)

	_, err := r.Get("nonexistent")
	assert.Error(t, err)
}

func TestGetMasked(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	got, err := r.GetMasked("test-openai")
	require.NoError(t, err)
	assert.NotEqual(t, "sk-test-xxxxx", got.APIKey)
	assert.True(t, len(got.APIKey) > 0)
	// Last 4 chars should be visible
	assert.Contains(t, got.APIKey, "xxxx")
}

// ---------------------------------------------------------------------------
// withKey behavior
// ---------------------------------------------------------------------------

func TestGet_DefaultMasked(t *testing.T) {
	r := setupRegistry(t)
	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	got, err := r.Get("test-openai")
	require.NoError(t, err)
	assert.NotEqual(t, "sk-test-xxxxx", got.APIKey, "Get() default should mask APIKey")
	assert.Contains(t, got.APIKey, "*")
}

func TestGet_WithKeyTrue(t *testing.T) {
	r := setupRegistry(t)
	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	got, err := r.Get("test-openai", true)
	require.NoError(t, err)
	assert.Equal(t, "sk-test-xxxxx", got.APIKey, "Get(key, true) should return plain text APIKey")
}

func TestGetByConnectorID_DefaultMasked(t *testing.T) {
	r := setupRegistry(t)
	p := testProvider
	created, err := r.Create(&p)
	require.NoError(t, err)

	got, err := r.GetByConnectorID(created.ConnectorID)
	require.NoError(t, err)
	assert.NotEqual(t, "sk-test-xxxxx", got.APIKey, "GetByConnectorID() default should mask")
	assert.Contains(t, got.APIKey, "*")
}

func TestGetByConnectorID_WithKeyTrue(t *testing.T) {
	r := setupRegistry(t)
	p := testProvider
	created, err := r.Create(&p)
	require.NoError(t, err)

	got, err := r.GetByConnectorID(created.ConnectorID, true)
	require.NoError(t, err)
	assert.Equal(t, "sk-test-xxxxx", got.APIKey, "GetByConnectorID(cid, true) should return plain text")
}

func TestList_DefaultMasked(t *testing.T) {
	r := setupRegistry(t)
	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	list, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceDynamic})
	require.NoError(t, err)
	require.True(t, len(list) > 0)

	for _, item := range list {
		assert.NotEqual(t, "sk-test-xxxxx", item.APIKey, "List() default should mask all APIKeys")
	}
}

func TestList_WithKeyTrue(t *testing.T) {
	r := setupRegistry(t)
	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	list, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceDynamic}, true)
	require.NoError(t, err)

	found := false
	for _, item := range list {
		if item.Key == "test-openai" {
			found = true
			assert.Equal(t, "sk-test-xxxxx", item.APIKey, "List(filter, true) should return plain text")
		}
	}
	assert.True(t, found)
}

func TestGetMasked_EqualsGetDefault(t *testing.T) {
	r := setupRegistry(t)
	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	fromGet, err := r.Get("test-openai")
	require.NoError(t, err)

	fromGetMasked, err := r.GetMasked("test-openai")
	require.NoError(t, err)

	assert.Equal(t, fromGet.APIKey, fromGetMasked.APIKey, "GetMasked should equal Get (both masked by default)")
}

func TestListModels_ConnectorHasRealKey(t *testing.T) {
	r := setupRegistry(t)

	owner := llmprovider.ProviderOwner{Type: "user", UserID: "rk-user"}
	p := llmprovider.Provider{
		Key:    llmprovider.ScopedKey(&owner, "realkey-prov"),
		Name:   "RealKey Test",
		Type:   "openai",
		APIURL: "https://api.openai.com",
		APIKey: "sk-real-secret-key-12345",
		Models: []llmprovider.ModelInfo{
			{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true},
		},
		Enabled: true,
		Owner:   owner,
	}
	_, err := r.Create(&p)
	require.NoError(t, err)

	opts := r.ListModelsByUser("rk-user")
	require.True(t, len(opts) > 0, "should have at least one model option")

	var modelCID string
	for _, o := range opts {
		if o.Label == "RealKey Test / GPT-4o" {
			modelCID = o.Value
			break
		}
	}
	require.NotEmpty(t, modelCID, "should find the model option")

	conn, err := connector.Select(modelCID)
	require.NoError(t, err, "model connector should be registered")

	s := conn.Setting()
	key, _ := s["key"].(string)
	assert.Equal(t, "sk-real-secret-key-12345", key, "connector should have the real API key, not masked")
}

func TestGetLazy(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	created, err := r.Create(&p)
	require.NoError(t, err)

	// Manually unregister the connector
	err = connector.Unregister(created.ConnectorID)
	require.NoError(t, err)

	// Verify it's gone
	_, err = connector.Select(created.ConnectorID)
	assert.Error(t, err)

	// Get should lazily re-register
	got, err := r.Get("test-openai")
	require.NoError(t, err)
	assert.Equal(t, "test-openai", got.Key)

	// Connector should be back
	_, err = connector.Select(got.ConnectorID)
	assert.NoError(t, err)
}

func TestList(t *testing.T) {
	r := setupRegistry(t)

	p2Owner := llmprovider.ProviderOwner{Type: "user", UserID: "123"}
	providers := []llmprovider.Provider{
		{Key: "p1", Name: "Provider 1", Type: "openai", Enabled: true,
			Models: []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"vision", "tool_calls"}, Enabled: true}},
			Owner:  llmprovider.ProviderOwner{Type: "system"}},
		{Key: llmprovider.ScopedKey(&p2Owner, "p2"), Name: "Provider 2", Type: "anthropic", Enabled: false,
			Models: []llmprovider.ModelInfo{{ID: "claude-3", Name: "Claude 3", Capabilities: []string{"tool_calls"}, Enabled: true}},
			Owner:  p2Owner},
		{Key: "p3", Name: "Provider 3", Type: "openai", Enabled: true,
			Models: []llmprovider.ModelInfo{{ID: "gpt-4o-mini", Name: "GPT-4o Mini", Capabilities: []string{"streaming"}, Enabled: true}},
			Owner:  llmprovider.ProviderOwner{Type: "system"}},
	}
	for i := range providers {
		_, err := r.Create(&providers[i])
		require.NoError(t, err)
	}

	t.Run("AllDynamic", func(t *testing.T) {
		list, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceDynamic})
		require.NoError(t, err)
		assert.GreaterOrEqual(t, len(list), 3)
	})

	t.Run("FilterByType", func(t *testing.T) {
		typ := "openai"
		list, err := r.List(&llmprovider.ProviderFilter{
			Source: llmprovider.ProviderSourceDynamic,
			Type:   &typ,
		})
		require.NoError(t, err)
		for _, p := range list {
			assert.Equal(t, "openai", p.Type)
		}
	})

	t.Run("FilterByEnabled", func(t *testing.T) {
		enabled := true
		list, err := r.List(&llmprovider.ProviderFilter{
			Source:  llmprovider.ProviderSourceDynamic,
			Enabled: &enabled,
		})
		require.NoError(t, err)
		for _, p := range list {
			assert.True(t, p.Enabled)
		}
	})

	t.Run("FilterByOwner", func(t *testing.T) {
		list, err := r.List(&llmprovider.ProviderFilter{
			Source: llmprovider.ProviderSourceDynamic,
			Owner:  &llmprovider.ProviderOwner{Type: "user", UserID: "123"},
		})
		require.NoError(t, err)
		for _, p := range list {
			assert.Equal(t, "user", p.Owner.Type)
			assert.Equal(t, "123", p.Owner.UserID)
		}
	})

	t.Run("FilterByCapabilities", func(t *testing.T) {
		list, err := r.List(&llmprovider.ProviderFilter{
			Source:       llmprovider.ProviderSourceDynamic,
			Capabilities: []string{"vision", "tool_calls"},
		})
		require.NoError(t, err)
		for _, p := range list {
			found := false
			for _, m := range p.Models {
				capSet := map[string]bool{}
				for _, c := range m.Capabilities {
					capSet[c] = true
				}
				if capSet["vision"] && capSet["tool_calls"] {
					found = true
					break
				}
			}
			assert.True(t, found, "provider %s should have model matching vision+tool_calls", p.Key)
		}
	})

	t.Run("FilterByKeyword", func(t *testing.T) {
		list, err := r.List(&llmprovider.ProviderFilter{
			Source:  llmprovider.ProviderSourceDynamic,
			Keyword: "Provider 2",
		})
		require.NoError(t, err)
		found := false
		for _, p := range list {
			if p.Name == "Provider 2" {
				found = true
			}
		}
		assert.True(t, found)
	})
}

func TestUpdate(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	created, err := r.Create(&p)
	require.NoError(t, err)

	updated := *created
	updated.APIURL = "https://custom.openai.com"
	updated.APIKey = "sk-new-key"

	result, err := r.Update("test-openai", &updated)
	require.NoError(t, err)
	assert.Equal(t, "https://custom.openai.com", result.APIURL)

	// Verify store updated
	got, err := r.Get("test-openai")
	require.NoError(t, err)
	assert.Equal(t, "https://custom.openai.com", got.APIURL)
}

func TestDelete(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	created, err := r.Create(&p)
	require.NoError(t, err)
	cid := created.ConnectorID

	err = r.Delete("test-openai")
	require.NoError(t, err)

	// Verify removed from store
	_, err = r.Get("test-openai")
	assert.Error(t, err)

	// Verify connector unregistered
	_, err = connector.Select(cid)
	assert.Error(t, err)
}

func TestReload(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	_, err := r.Create(&p)
	require.NoError(t, err)

	// Clear cache to simulate stale state
	c, _ := store.Get("__yao.cache")
	if c != nil {
		c.Del("llmprovider:*")
	}

	err = r.Reload()
	require.NoError(t, err)

	// Should still be able to get the provider
	got, err := r.Get("test-openai")
	require.NoError(t, err)
	assert.Equal(t, "Test OpenAI", got.Name)
}

func TestImportFromConnectors(t *testing.T) {
	r := setupRegistry(t)

	// After Init, builtin connectors should be imported
	list, err := r.List(&llmprovider.ProviderFilter{
		Source: llmprovider.ProviderSourceAll,
	})
	require.NoError(t, err)

	builtinCount := 0
	for _, p := range list {
		if p.Source == llmprovider.ProviderSourceBuiltIn {
			builtinCount++
		}
	}

	// Should have imported some from connector.AIConnectors (if test app has connectors)
	t.Logf("Imported %d builtin providers from connector.AIConnectors (total AIConnectors: %d)", builtinCount, len(connector.AIConnectors))
}

func TestGetPresets(t *testing.T) {
	presets := llmprovider.GetPresets()
	assert.Greater(t, len(presets), 0, "should have at least one preset")

	// Verify openai preset exists
	var openai *llmprovider.ProviderPreset
	for i := range presets {
		if presets[i].Key == "openai" {
			openai = &presets[i]
			break
		}
	}
	require.NotNil(t, openai, "openai preset should exist")
	assert.Equal(t, "OpenAI", openai.Name)
	assert.Equal(t, "openai", openai.Type)
	assert.True(t, openai.RequireKey)
	assert.Greater(t, len(openai.DefaultModels), 0)
}

func TestGetPreset(t *testing.T) {
	p := llmprovider.GetPreset("anthropic")
	require.NotNil(t, p)
	assert.Equal(t, "Anthropic", p.Name)

	none := llmprovider.GetPreset("nonexistent")
	assert.Nil(t, none)
}

func TestEncryptionRoundTrip(t *testing.T) {
	r := setupRegistry(t)
	r.SetEncryptionKey("my-super-secret-key-for-tests")

	p := testProvider
	p.Key = "test-encrypted"
	_, err := r.Create(&p)
	require.NoError(t, err)

	got, err := r.Get("test-encrypted", true)
	require.NoError(t, err)
	assert.Equal(t, "sk-test-xxxxx", got.APIKey, "APIKey should be decrypted on read with withKey=true")

	masked, err := r.Get("test-encrypted")
	require.NoError(t, err)
	assert.NotEqual(t, "sk-test-xxxxx", masked.APIKey, "Get without withKey should mask")
	assert.Contains(t, masked.APIKey, "xxxx")

	// Verify raw store value is encrypted
	s, _ := store.Get("__yao.store")
	raw, ok := s.Get("llmprovider:p:test-encrypted")
	require.True(t, ok)
	m := raw.(map[string]interface{})
	storedKey, _ := m["api_key"].(string)
	assert.True(t, len(storedKey) > 0)
	assert.NotEqual(t, "sk-test-xxxxx", storedKey, "raw stored value should be encrypted")
}

func TestGetConnector(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	p.Key = "test-getconn"
	_, err := r.Create(&p)
	require.NoError(t, err)

	conn, err := r.GetConnector("test-getconn")
	require.NoError(t, err)
	assert.NotNil(t, conn)

	setting := conn.Setting()
	assert.NotNil(t, setting)
	host, _ := setting["host"].(string)
	assert.Equal(t, "https://api.openai.com", host)
}

func TestGetSetting(t *testing.T) {
	r := setupRegistry(t)

	p := testProvider
	p.Key = "test-getsetting"
	_, err := r.Create(&p)
	require.NoError(t, err)

	setting, err := r.GetSetting("test-getsetting")
	require.NoError(t, err)
	assert.NotNil(t, setting)
	host, _ := setting["host"].(string)
	assert.Equal(t, "https://api.openai.com", host)
}

func TestGetConnectorNotFound(t *testing.T) {
	r := setupRegistry(t)
	_, err := r.GetConnector("not-exist")
	assert.Error(t, err)
}

func TestGetSettingNotFound(t *testing.T) {
	r := setupRegistry(t)
	_, err := r.GetSetting("not-exist")
	assert.Error(t, err)
}

func TestCreateEmptyKey(t *testing.T) {
	r := setupRegistry(t)
	p := llmprovider.Provider{Name: "No Key"}
	_, err := r.Create(&p)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key is required")
}

func TestCreateDisabled(t *testing.T) {
	r := setupRegistry(t)
	p := llmprovider.Provider{
		Key:     "test-disabled",
		Name:    "Disabled Provider",
		Type:    "openai",
		APIURL:  "https://api.openai.com",
		Enabled: false,
		Models:  []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true}},
		Owner:   llmprovider.ProviderOwner{Type: "system"},
	}
	created, err := r.Create(&p)
	require.NoError(t, err)
	assert.Equal(t, "unconfigured", created.Status)

	// Disabled provider should not have its connector registered
	_, err = connector.Select(created.ConnectorID)
	assert.Error(t, err, "disabled provider should not register connector")
}

func TestOwnerPrefixedIDs(t *testing.T) {
	r := setupRegistry(t)

	cases := []struct {
		key    string
		owner  llmprovider.ProviderOwner
		prefix string
	}{
		{"owner-sys", llmprovider.ProviderOwner{Type: "system"}, "s."},
		{"owner-user", llmprovider.ProviderOwner{Type: "user", UserID: "42"}, "u42."},
		{"owner-team", llmprovider.ProviderOwner{Type: "team", TeamID: "99"}, "t99."},
	}

	for _, tc := range cases {
		t.Run(tc.key, func(t *testing.T) {
			p := llmprovider.Provider{
				Key:     tc.key,
				Name:    tc.key,
				Type:    "openai",
				APIURL:  "https://api.openai.com",
				Enabled: true,
				Models:  []llmprovider.ModelInfo{{ID: "m1", Name: "M1", Capabilities: []string{"streaming"}, Enabled: true}},
				Owner:   tc.owner,
			}
			created, err := r.Create(&p)
			require.NoError(t, err)
			assert.Contains(t, created.ConnectorID, tc.prefix,
				"ConnectorID for %s owner should contain prefix %s", tc.owner.Type, tc.prefix)

			// Verify connector is registered with the prefixed ID
			_, err = connector.Select(created.ConnectorID)
			assert.NoError(t, err)
		})
	}
}

func TestListBuiltInFilter(t *testing.T) {
	r := setupRegistry(t)

	builtinList, err := r.List(&llmprovider.ProviderFilter{Source: llmprovider.ProviderSourceBuiltIn})
	require.NoError(t, err)
	for _, p := range builtinList {
		assert.Equal(t, llmprovider.ProviderSourceBuiltIn, p.Source)
	}
}

func TestListPresetKeyFilter(t *testing.T) {
	r := setupRegistry(t)

	p := llmprovider.Provider{
		Key:       "from-preset",
		Name:      "From Preset",
		Type:      "openai",
		PresetKey: "openai",
		Enabled:   true,
		Models:    []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true}},
		Owner:     llmprovider.ProviderOwner{Type: "system"},
	}
	_, err := r.Create(&p)
	require.NoError(t, err)

	pk := "openai"
	list, err := r.List(&llmprovider.ProviderFilter{
		Source:    llmprovider.ProviderSourceDynamic,
		PresetKey: &pk,
	})
	require.NoError(t, err)
	found := false
	for _, item := range list {
		if item.Key == "from-preset" {
			found = true
			assert.Equal(t, "openai", item.PresetKey)
		}
	}
	assert.True(t, found)
}

func TestDefaultModelFallback(t *testing.T) {
	r := setupRegistry(t)

	// Provider with no enabled models — should use first model ID as default
	p := llmprovider.Provider{
		Key:     "test-fallback",
		Name:    "Fallback",
		Type:    "openai",
		APIURL:  "https://api.openai.com",
		Enabled: true,
		Models:  []llmprovider.ModelInfo{{ID: "only-model", Name: "Only", Capabilities: []string{"streaming"}, Enabled: false}},
		Owner:   llmprovider.ProviderOwner{Type: "system"},
	}
	created, err := r.Create(&p)
	require.NoError(t, err)

	// Connector should still be registered using the fallback model
	conn, cerr := connector.Select(created.ConnectorID)
	require.NoError(t, cerr)
	setting := conn.Setting()
	model, _ := setting["model"].(string)
	assert.Equal(t, "only-model", model)
}

func TestMaskShortKey(t *testing.T) {
	r := setupRegistry(t)

	p := llmprovider.Provider{
		Key:     "test-shortkey",
		Name:    "Short",
		Type:    "openai",
		APIKey:  "ab",
		Enabled: true,
		Models:  []llmprovider.ModelInfo{{ID: "m", Name: "M", Capabilities: []string{"streaming"}, Enabled: true}},
		Owner:   llmprovider.ProviderOwner{Type: "system"},
	}
	_, err := r.Create(&p)
	require.NoError(t, err)

	masked, err := r.GetMasked("test-shortkey")
	require.NoError(t, err)
	// Short keys should be fully masked
	assert.Equal(t, "**", masked.APIKey)
}

func TestConcurrency(t *testing.T) {
	r := setupRegistry(t)

	var wg sync.WaitGroup
	errCh := make(chan error, 30)

	// Concurrent creates with unique keys
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			p := llmprovider.Provider{
				Key:     fmt.Sprintf("conc-%d", idx),
				Name:    fmt.Sprintf("Concurrent %d", idx),
				Type:    "openai",
				APIURL:  "https://api.openai.com",
				Enabled: true,
				Models:  []llmprovider.ModelInfo{{ID: "gpt-4o", Name: "GPT-4o", Capabilities: []string{"streaming"}, Enabled: true}},
				Owner:   llmprovider.ProviderOwner{Type: "system"},
			}
			if _, err := r.Create(&p); err != nil {
				errCh <- err
			}
		}(i)
	}

	wg.Wait()

	// Concurrent reads
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			_, err := r.Get(fmt.Sprintf("conc-%d", idx))
			if err != nil {
				errCh <- err
			}
		}(i)
	}

	wg.Wait()

	// Concurrent deletes
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func(idx int) {
			defer wg.Done()
			if err := r.Delete(fmt.Sprintf("conc-%d", idx)); err != nil {
				errCh <- err
			}
		}(i)
	}

	wg.Wait()
	close(errCh)

	for err := range errCh {
		t.Errorf("concurrent operation error: %v", err)
	}
}