load_test.go 33 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250
  1. package config
  2. import (
  3. "io"
  4. "log/slog"
  5. "os"
  6. "path/filepath"
  7. "strings"
  8. "testing"
  9. "github.com/charmbracelet/catwalk/pkg/catwalk"
  10. "github.com/charmbracelet/crush/internal/csync"
  11. "github.com/charmbracelet/crush/internal/env"
  12. "github.com/stretchr/testify/assert"
  13. "github.com/stretchr/testify/require"
  14. )
  15. func TestMain(m *testing.M) {
  16. slog.SetDefault(slog.New(slog.NewTextHandler(io.Discard, nil)))
  17. exitVal := m.Run()
  18. os.Exit(exitVal)
  19. }
  20. func TestConfig_LoadFromReaders(t *testing.T) {
  21. data1 := strings.NewReader(`{"providers": {"openai": {"api_key": "key1", "base_url": "https://api.openai.com/v1"}}}`)
  22. data2 := strings.NewReader(`{"providers": {"openai": {"api_key": "key2", "base_url": "https://api.openai.com/v2"}}}`)
  23. data3 := strings.NewReader(`{"providers": {"openai": {}}}`)
  24. loadedConfig, err := loadFromReaders([]io.Reader{data1, data2, data3})
  25. require.NoError(t, err)
  26. require.NotNil(t, loadedConfig)
  27. require.Equal(t, 1, loadedConfig.Providers.Len())
  28. pc, _ := loadedConfig.Providers.Get("openai")
  29. require.Equal(t, "key2", pc.APIKey)
  30. require.Equal(t, "https://api.openai.com/v2", pc.BaseURL)
  31. }
  32. func TestConfig_setDefaults(t *testing.T) {
  33. cfg := &Config{}
  34. cfg.setDefaults("/tmp", "")
  35. require.NotNil(t, cfg.Options)
  36. require.NotNil(t, cfg.Options.TUI)
  37. require.NotNil(t, cfg.Options.ContextPaths)
  38. require.NotNil(t, cfg.Providers)
  39. require.NotNil(t, cfg.Models)
  40. require.NotNil(t, cfg.LSP)
  41. require.NotNil(t, cfg.MCP)
  42. require.Equal(t, filepath.Join("/tmp", ".crush"), cfg.Options.DataDirectory)
  43. require.Equal(t, "AGENTS.md", cfg.Options.InitializeAs)
  44. for _, path := range defaultContextPaths {
  45. require.Contains(t, cfg.Options.ContextPaths, path)
  46. }
  47. require.Equal(t, "/tmp", cfg.workingDir)
  48. }
  49. func TestConfig_configureProviders(t *testing.T) {
  50. knownProviders := []catwalk.Provider{
  51. {
  52. ID: "openai",
  53. APIKey: "$OPENAI_API_KEY",
  54. APIEndpoint: "https://api.openai.com/v1",
  55. Models: []catwalk.Model{{
  56. ID: "test-model",
  57. }},
  58. },
  59. }
  60. cfg := &Config{}
  61. cfg.setDefaults("/tmp", "")
  62. env := env.NewFromMap(map[string]string{
  63. "OPENAI_API_KEY": "test-key",
  64. })
  65. resolver := NewEnvironmentVariableResolver(env)
  66. err := cfg.configureProviders(env, resolver, knownProviders)
  67. require.NoError(t, err)
  68. require.Equal(t, 1, cfg.Providers.Len())
  69. // We want to make sure that we keep the configured API key as a placeholder
  70. pc, _ := cfg.Providers.Get("openai")
  71. require.Equal(t, "$OPENAI_API_KEY", pc.APIKey)
  72. }
  73. func TestConfig_configureProvidersWithOverride(t *testing.T) {
  74. knownProviders := []catwalk.Provider{
  75. {
  76. ID: "openai",
  77. APIKey: "$OPENAI_API_KEY",
  78. APIEndpoint: "https://api.openai.com/v1",
  79. Models: []catwalk.Model{{
  80. ID: "test-model",
  81. }},
  82. },
  83. }
  84. cfg := &Config{
  85. Providers: csync.NewMap[string, ProviderConfig](),
  86. }
  87. cfg.Providers.Set("openai", ProviderConfig{
  88. APIKey: "xyz",
  89. BaseURL: "https://api.openai.com/v2",
  90. Models: []catwalk.Model{
  91. {
  92. ID: "test-model",
  93. Name: "Updated",
  94. },
  95. {
  96. ID: "another-model",
  97. },
  98. },
  99. })
  100. cfg.setDefaults("/tmp", "")
  101. env := env.NewFromMap(map[string]string{
  102. "OPENAI_API_KEY": "test-key",
  103. })
  104. resolver := NewEnvironmentVariableResolver(env)
  105. err := cfg.configureProviders(env, resolver, knownProviders)
  106. require.NoError(t, err)
  107. require.Equal(t, 1, cfg.Providers.Len())
  108. // We want to make sure that we keep the configured API key as a placeholder
  109. pc, _ := cfg.Providers.Get("openai")
  110. require.Equal(t, "xyz", pc.APIKey)
  111. require.Equal(t, "https://api.openai.com/v2", pc.BaseURL)
  112. require.Len(t, pc.Models, 2)
  113. require.Equal(t, "Updated", pc.Models[0].Name)
  114. }
  115. func TestConfig_configureProvidersWithNewProvider(t *testing.T) {
  116. knownProviders := []catwalk.Provider{
  117. {
  118. ID: "openai",
  119. APIKey: "$OPENAI_API_KEY",
  120. APIEndpoint: "https://api.openai.com/v1",
  121. Models: []catwalk.Model{{
  122. ID: "test-model",
  123. }},
  124. },
  125. }
  126. cfg := &Config{
  127. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  128. "custom": {
  129. APIKey: "xyz",
  130. BaseURL: "https://api.someendpoint.com/v2",
  131. Models: []catwalk.Model{
  132. {
  133. ID: "test-model",
  134. },
  135. },
  136. },
  137. }),
  138. }
  139. cfg.setDefaults("/tmp", "")
  140. env := env.NewFromMap(map[string]string{
  141. "OPENAI_API_KEY": "test-key",
  142. })
  143. resolver := NewEnvironmentVariableResolver(env)
  144. err := cfg.configureProviders(env, resolver, knownProviders)
  145. require.NoError(t, err)
  146. // Should be to because of the env variable
  147. require.Equal(t, cfg.Providers.Len(), 2)
  148. // We want to make sure that we keep the configured API key as a placeholder
  149. pc, _ := cfg.Providers.Get("custom")
  150. require.Equal(t, "xyz", pc.APIKey)
  151. // Make sure we set the ID correctly
  152. require.Equal(t, "custom", pc.ID)
  153. require.Equal(t, "https://api.someendpoint.com/v2", pc.BaseURL)
  154. require.Len(t, pc.Models, 1)
  155. _, ok := cfg.Providers.Get("openai")
  156. require.True(t, ok, "OpenAI provider should still be present")
  157. }
  158. func TestConfig_configureProvidersBedrockWithCredentials(t *testing.T) {
  159. knownProviders := []catwalk.Provider{
  160. {
  161. ID: catwalk.InferenceProviderBedrock,
  162. APIKey: "",
  163. APIEndpoint: "",
  164. Models: []catwalk.Model{{
  165. ID: "anthropic.claude-sonnet-4-20250514-v1:0",
  166. }},
  167. },
  168. }
  169. cfg := &Config{}
  170. cfg.setDefaults("/tmp", "")
  171. env := env.NewFromMap(map[string]string{
  172. "AWS_ACCESS_KEY_ID": "test-key-id",
  173. "AWS_SECRET_ACCESS_KEY": "test-secret-key",
  174. })
  175. resolver := NewEnvironmentVariableResolver(env)
  176. err := cfg.configureProviders(env, resolver, knownProviders)
  177. require.NoError(t, err)
  178. require.Equal(t, cfg.Providers.Len(), 1)
  179. bedrockProvider, ok := cfg.Providers.Get("bedrock")
  180. require.True(t, ok, "Bedrock provider should be present")
  181. require.Len(t, bedrockProvider.Models, 1)
  182. require.Equal(t, "anthropic.claude-sonnet-4-20250514-v1:0", bedrockProvider.Models[0].ID)
  183. }
  184. func TestConfig_configureProvidersBedrockWithoutCredentials(t *testing.T) {
  185. knownProviders := []catwalk.Provider{
  186. {
  187. ID: catwalk.InferenceProviderBedrock,
  188. APIKey: "",
  189. APIEndpoint: "",
  190. Models: []catwalk.Model{{
  191. ID: "anthropic.claude-sonnet-4-20250514-v1:0",
  192. }},
  193. },
  194. }
  195. cfg := &Config{}
  196. cfg.setDefaults("/tmp", "")
  197. env := env.NewFromMap(map[string]string{})
  198. resolver := NewEnvironmentVariableResolver(env)
  199. err := cfg.configureProviders(env, resolver, knownProviders)
  200. require.NoError(t, err)
  201. // Provider should not be configured without credentials
  202. require.Equal(t, cfg.Providers.Len(), 0)
  203. }
  204. func TestConfig_configureProvidersBedrockWithoutUnsupportedModel(t *testing.T) {
  205. knownProviders := []catwalk.Provider{
  206. {
  207. ID: catwalk.InferenceProviderBedrock,
  208. APIKey: "",
  209. APIEndpoint: "",
  210. Models: []catwalk.Model{{
  211. ID: "some-random-model",
  212. }},
  213. },
  214. }
  215. cfg := &Config{}
  216. cfg.setDefaults("/tmp", "")
  217. env := env.NewFromMap(map[string]string{
  218. "AWS_ACCESS_KEY_ID": "test-key-id",
  219. "AWS_SECRET_ACCESS_KEY": "test-secret-key",
  220. })
  221. resolver := NewEnvironmentVariableResolver(env)
  222. err := cfg.configureProviders(env, resolver, knownProviders)
  223. require.Error(t, err)
  224. }
  225. func TestConfig_configureProvidersVertexAIWithCredentials(t *testing.T) {
  226. knownProviders := []catwalk.Provider{
  227. {
  228. ID: catwalk.InferenceProviderVertexAI,
  229. APIKey: "",
  230. APIEndpoint: "",
  231. Models: []catwalk.Model{{
  232. ID: "gemini-pro",
  233. }},
  234. },
  235. }
  236. cfg := &Config{}
  237. cfg.setDefaults("/tmp", "")
  238. env := env.NewFromMap(map[string]string{
  239. "VERTEXAI_PROJECT": "test-project",
  240. "VERTEXAI_LOCATION": "us-central1",
  241. })
  242. resolver := NewEnvironmentVariableResolver(env)
  243. err := cfg.configureProviders(env, resolver, knownProviders)
  244. require.NoError(t, err)
  245. require.Equal(t, cfg.Providers.Len(), 1)
  246. vertexProvider, ok := cfg.Providers.Get("vertexai")
  247. require.True(t, ok, "VertexAI provider should be present")
  248. require.Len(t, vertexProvider.Models, 1)
  249. require.Equal(t, "gemini-pro", vertexProvider.Models[0].ID)
  250. require.Equal(t, "test-project", vertexProvider.ExtraParams["project"])
  251. require.Equal(t, "us-central1", vertexProvider.ExtraParams["location"])
  252. }
  253. func TestConfig_configureProvidersVertexAIWithoutCredentials(t *testing.T) {
  254. knownProviders := []catwalk.Provider{
  255. {
  256. ID: catwalk.InferenceProviderVertexAI,
  257. APIKey: "",
  258. APIEndpoint: "",
  259. Models: []catwalk.Model{{
  260. ID: "gemini-pro",
  261. }},
  262. },
  263. }
  264. cfg := &Config{}
  265. cfg.setDefaults("/tmp", "")
  266. env := env.NewFromMap(map[string]string{
  267. "GOOGLE_GENAI_USE_VERTEXAI": "false",
  268. "GOOGLE_CLOUD_PROJECT": "test-project",
  269. "GOOGLE_CLOUD_LOCATION": "us-central1",
  270. })
  271. resolver := NewEnvironmentVariableResolver(env)
  272. err := cfg.configureProviders(env, resolver, knownProviders)
  273. require.NoError(t, err)
  274. // Provider should not be configured without proper credentials
  275. require.Equal(t, cfg.Providers.Len(), 0)
  276. }
  277. func TestConfig_configureProvidersVertexAIMissingProject(t *testing.T) {
  278. knownProviders := []catwalk.Provider{
  279. {
  280. ID: catwalk.InferenceProviderVertexAI,
  281. APIKey: "",
  282. APIEndpoint: "",
  283. Models: []catwalk.Model{{
  284. ID: "gemini-pro",
  285. }},
  286. },
  287. }
  288. cfg := &Config{}
  289. cfg.setDefaults("/tmp", "")
  290. env := env.NewFromMap(map[string]string{
  291. "GOOGLE_GENAI_USE_VERTEXAI": "true",
  292. "GOOGLE_CLOUD_LOCATION": "us-central1",
  293. })
  294. resolver := NewEnvironmentVariableResolver(env)
  295. err := cfg.configureProviders(env, resolver, knownProviders)
  296. require.NoError(t, err)
  297. // Provider should not be configured without project
  298. require.Equal(t, cfg.Providers.Len(), 0)
  299. }
  300. func TestConfig_configureProvidersSetProviderID(t *testing.T) {
  301. knownProviders := []catwalk.Provider{
  302. {
  303. ID: "openai",
  304. APIKey: "$OPENAI_API_KEY",
  305. APIEndpoint: "https://api.openai.com/v1",
  306. Models: []catwalk.Model{{
  307. ID: "test-model",
  308. }},
  309. },
  310. }
  311. cfg := &Config{}
  312. cfg.setDefaults("/tmp", "")
  313. env := env.NewFromMap(map[string]string{
  314. "OPENAI_API_KEY": "test-key",
  315. })
  316. resolver := NewEnvironmentVariableResolver(env)
  317. err := cfg.configureProviders(env, resolver, knownProviders)
  318. require.NoError(t, err)
  319. require.Equal(t, cfg.Providers.Len(), 1)
  320. // Provider ID should be set
  321. pc, _ := cfg.Providers.Get("openai")
  322. require.Equal(t, "openai", pc.ID)
  323. }
  324. func TestConfig_EnabledProviders(t *testing.T) {
  325. t.Run("all providers enabled", func(t *testing.T) {
  326. cfg := &Config{
  327. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  328. "openai": {
  329. ID: "openai",
  330. APIKey: "key1",
  331. Disable: false,
  332. },
  333. "anthropic": {
  334. ID: "anthropic",
  335. APIKey: "key2",
  336. Disable: false,
  337. },
  338. }),
  339. }
  340. enabled := cfg.EnabledProviders()
  341. require.Len(t, enabled, 2)
  342. })
  343. t.Run("some providers disabled", func(t *testing.T) {
  344. cfg := &Config{
  345. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  346. "openai": {
  347. ID: "openai",
  348. APIKey: "key1",
  349. Disable: false,
  350. },
  351. "anthropic": {
  352. ID: "anthropic",
  353. APIKey: "key2",
  354. Disable: true,
  355. },
  356. }),
  357. }
  358. enabled := cfg.EnabledProviders()
  359. require.Len(t, enabled, 1)
  360. require.Equal(t, "openai", enabled[0].ID)
  361. })
  362. t.Run("empty providers map", func(t *testing.T) {
  363. cfg := &Config{
  364. Providers: csync.NewMap[string, ProviderConfig](),
  365. }
  366. enabled := cfg.EnabledProviders()
  367. require.Len(t, enabled, 0)
  368. })
  369. }
  370. func TestConfig_IsConfigured(t *testing.T) {
  371. t.Run("returns true when at least one provider is enabled", func(t *testing.T) {
  372. cfg := &Config{
  373. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  374. "openai": {
  375. ID: "openai",
  376. APIKey: "key1",
  377. Disable: false,
  378. },
  379. }),
  380. }
  381. require.True(t, cfg.IsConfigured())
  382. })
  383. t.Run("returns false when no providers are configured", func(t *testing.T) {
  384. cfg := &Config{
  385. Providers: csync.NewMap[string, ProviderConfig](),
  386. }
  387. require.False(t, cfg.IsConfigured())
  388. })
  389. t.Run("returns false when all providers are disabled", func(t *testing.T) {
  390. cfg := &Config{
  391. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  392. "openai": {
  393. ID: "openai",
  394. APIKey: "key1",
  395. Disable: true,
  396. },
  397. "anthropic": {
  398. ID: "anthropic",
  399. APIKey: "key2",
  400. Disable: true,
  401. },
  402. }),
  403. }
  404. require.False(t, cfg.IsConfigured())
  405. })
  406. }
  407. func TestConfig_setupAgentsWithNoDisabledTools(t *testing.T) {
  408. cfg := &Config{
  409. Options: &Options{
  410. DisabledTools: []string{},
  411. },
  412. }
  413. cfg.SetupAgents()
  414. coderAgent, ok := cfg.Agents[AgentCoder]
  415. require.True(t, ok)
  416. assert.Equal(t, allToolNames(), coderAgent.AllowedTools)
  417. taskAgent, ok := cfg.Agents[AgentTask]
  418. require.True(t, ok)
  419. assert.Equal(t, []string{"glob", "grep", "ls", "sourcegraph", "view"}, taskAgent.AllowedTools)
  420. }
  421. func TestConfig_setupAgentsWithDisabledTools(t *testing.T) {
  422. cfg := &Config{
  423. Options: &Options{
  424. DisabledTools: []string{
  425. "edit",
  426. "download",
  427. "grep",
  428. },
  429. },
  430. }
  431. cfg.SetupAgents()
  432. coderAgent, ok := cfg.Agents[AgentCoder]
  433. require.True(t, ok)
  434. assert.Equal(t, []string{"agent", "bash", "job_output", "job_kill", "multiedit", "lsp_diagnostics", "lsp_references", "fetch", "agentic_fetch", "glob", "ls", "sourcegraph", "todos", "view", "write"}, coderAgent.AllowedTools)
  435. taskAgent, ok := cfg.Agents[AgentTask]
  436. require.True(t, ok)
  437. assert.Equal(t, []string{"glob", "ls", "sourcegraph", "view"}, taskAgent.AllowedTools)
  438. }
  439. func TestConfig_setupAgentsWithEveryReadOnlyToolDisabled(t *testing.T) {
  440. cfg := &Config{
  441. Options: &Options{
  442. DisabledTools: []string{
  443. "glob",
  444. "grep",
  445. "ls",
  446. "sourcegraph",
  447. "view",
  448. },
  449. },
  450. }
  451. cfg.SetupAgents()
  452. coderAgent, ok := cfg.Agents[AgentCoder]
  453. require.True(t, ok)
  454. assert.Equal(t, []string{"agent", "bash", "job_output", "job_kill", "download", "edit", "multiedit", "lsp_diagnostics", "lsp_references", "fetch", "agentic_fetch", "todos", "write"}, coderAgent.AllowedTools)
  455. taskAgent, ok := cfg.Agents[AgentTask]
  456. require.True(t, ok)
  457. assert.Equal(t, []string{}, taskAgent.AllowedTools)
  458. }
  459. func TestConfig_configureProvidersWithDisabledProvider(t *testing.T) {
  460. knownProviders := []catwalk.Provider{
  461. {
  462. ID: "openai",
  463. APIKey: "$OPENAI_API_KEY",
  464. APIEndpoint: "https://api.openai.com/v1",
  465. Models: []catwalk.Model{{
  466. ID: "test-model",
  467. }},
  468. },
  469. }
  470. cfg := &Config{
  471. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  472. "openai": {
  473. Disable: true,
  474. },
  475. }),
  476. }
  477. cfg.setDefaults("/tmp", "")
  478. env := env.NewFromMap(map[string]string{
  479. "OPENAI_API_KEY": "test-key",
  480. })
  481. resolver := NewEnvironmentVariableResolver(env)
  482. err := cfg.configureProviders(env, resolver, knownProviders)
  483. require.NoError(t, err)
  484. require.Equal(t, cfg.Providers.Len(), 1)
  485. prov, exists := cfg.Providers.Get("openai")
  486. require.True(t, exists)
  487. require.True(t, prov.Disable)
  488. }
  489. func TestConfig_configureProvidersCustomProviderValidation(t *testing.T) {
  490. t.Run("custom provider with missing API key is allowed, but not known providers", func(t *testing.T) {
  491. cfg := &Config{
  492. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  493. "custom": {
  494. BaseURL: "https://api.custom.com/v1",
  495. Models: []catwalk.Model{{
  496. ID: "test-model",
  497. }},
  498. },
  499. "openai": {
  500. APIKey: "$MISSING",
  501. },
  502. }),
  503. }
  504. cfg.setDefaults("/tmp", "")
  505. env := env.NewFromMap(map[string]string{})
  506. resolver := NewEnvironmentVariableResolver(env)
  507. err := cfg.configureProviders(env, resolver, []catwalk.Provider{})
  508. require.NoError(t, err)
  509. require.Equal(t, cfg.Providers.Len(), 1)
  510. _, exists := cfg.Providers.Get("custom")
  511. require.True(t, exists)
  512. })
  513. t.Run("custom provider with missing BaseURL is removed", func(t *testing.T) {
  514. cfg := &Config{
  515. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  516. "custom": {
  517. APIKey: "test-key",
  518. Models: []catwalk.Model{{
  519. ID: "test-model",
  520. }},
  521. },
  522. }),
  523. }
  524. cfg.setDefaults("/tmp", "")
  525. env := env.NewFromMap(map[string]string{})
  526. resolver := NewEnvironmentVariableResolver(env)
  527. err := cfg.configureProviders(env, resolver, []catwalk.Provider{})
  528. require.NoError(t, err)
  529. require.Equal(t, cfg.Providers.Len(), 0)
  530. _, exists := cfg.Providers.Get("custom")
  531. require.False(t, exists)
  532. })
  533. t.Run("custom provider with no models is removed", func(t *testing.T) {
  534. cfg := &Config{
  535. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  536. "custom": {
  537. APIKey: "test-key",
  538. BaseURL: "https://api.custom.com/v1",
  539. Models: []catwalk.Model{},
  540. },
  541. }),
  542. }
  543. cfg.setDefaults("/tmp", "")
  544. env := env.NewFromMap(map[string]string{})
  545. resolver := NewEnvironmentVariableResolver(env)
  546. err := cfg.configureProviders(env, resolver, []catwalk.Provider{})
  547. require.NoError(t, err)
  548. require.Equal(t, cfg.Providers.Len(), 0)
  549. _, exists := cfg.Providers.Get("custom")
  550. require.False(t, exists)
  551. })
  552. t.Run("custom provider with unsupported type is removed", func(t *testing.T) {
  553. cfg := &Config{
  554. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  555. "custom": {
  556. APIKey: "test-key",
  557. BaseURL: "https://api.custom.com/v1",
  558. Type: "unsupported",
  559. Models: []catwalk.Model{{
  560. ID: "test-model",
  561. }},
  562. },
  563. }),
  564. }
  565. cfg.setDefaults("/tmp", "")
  566. env := env.NewFromMap(map[string]string{})
  567. resolver := NewEnvironmentVariableResolver(env)
  568. err := cfg.configureProviders(env, resolver, []catwalk.Provider{})
  569. require.NoError(t, err)
  570. require.Equal(t, cfg.Providers.Len(), 0)
  571. _, exists := cfg.Providers.Get("custom")
  572. require.False(t, exists)
  573. })
  574. t.Run("valid custom provider is kept and ID is set", func(t *testing.T) {
  575. cfg := &Config{
  576. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  577. "custom": {
  578. APIKey: "test-key",
  579. BaseURL: "https://api.custom.com/v1",
  580. Type: catwalk.TypeOpenAI,
  581. Models: []catwalk.Model{{
  582. ID: "test-model",
  583. }},
  584. },
  585. }),
  586. }
  587. cfg.setDefaults("/tmp", "")
  588. env := env.NewFromMap(map[string]string{})
  589. resolver := NewEnvironmentVariableResolver(env)
  590. err := cfg.configureProviders(env, resolver, []catwalk.Provider{})
  591. require.NoError(t, err)
  592. require.Equal(t, cfg.Providers.Len(), 1)
  593. customProvider, exists := cfg.Providers.Get("custom")
  594. require.True(t, exists)
  595. require.Equal(t, "custom", customProvider.ID)
  596. require.Equal(t, "test-key", customProvider.APIKey)
  597. require.Equal(t, "https://api.custom.com/v1", customProvider.BaseURL)
  598. })
  599. t.Run("custom anthropic provider is supported", func(t *testing.T) {
  600. cfg := &Config{
  601. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  602. "custom-anthropic": {
  603. APIKey: "test-key",
  604. BaseURL: "https://api.anthropic.com/v1",
  605. Type: catwalk.TypeAnthropic,
  606. Models: []catwalk.Model{{
  607. ID: "claude-3-sonnet",
  608. }},
  609. },
  610. }),
  611. }
  612. cfg.setDefaults("/tmp", "")
  613. env := env.NewFromMap(map[string]string{})
  614. resolver := NewEnvironmentVariableResolver(env)
  615. err := cfg.configureProviders(env, resolver, []catwalk.Provider{})
  616. require.NoError(t, err)
  617. require.Equal(t, cfg.Providers.Len(), 1)
  618. customProvider, exists := cfg.Providers.Get("custom-anthropic")
  619. require.True(t, exists)
  620. require.Equal(t, "custom-anthropic", customProvider.ID)
  621. require.Equal(t, "test-key", customProvider.APIKey)
  622. require.Equal(t, "https://api.anthropic.com/v1", customProvider.BaseURL)
  623. require.Equal(t, catwalk.TypeAnthropic, customProvider.Type)
  624. })
  625. t.Run("disabled custom provider is removed", func(t *testing.T) {
  626. cfg := &Config{
  627. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  628. "custom": {
  629. APIKey: "test-key",
  630. BaseURL: "https://api.custom.com/v1",
  631. Type: catwalk.TypeOpenAI,
  632. Disable: true,
  633. Models: []catwalk.Model{{
  634. ID: "test-model",
  635. }},
  636. },
  637. }),
  638. }
  639. cfg.setDefaults("/tmp", "")
  640. env := env.NewFromMap(map[string]string{})
  641. resolver := NewEnvironmentVariableResolver(env)
  642. err := cfg.configureProviders(env, resolver, []catwalk.Provider{})
  643. require.NoError(t, err)
  644. require.Equal(t, cfg.Providers.Len(), 0)
  645. _, exists := cfg.Providers.Get("custom")
  646. require.False(t, exists)
  647. })
  648. }
  649. func TestConfig_configureProvidersEnhancedCredentialValidation(t *testing.T) {
  650. t.Run("VertexAI provider removed when credentials missing with existing config", func(t *testing.T) {
  651. knownProviders := []catwalk.Provider{
  652. {
  653. ID: catwalk.InferenceProviderVertexAI,
  654. APIKey: "",
  655. APIEndpoint: "",
  656. Models: []catwalk.Model{{
  657. ID: "gemini-pro",
  658. }},
  659. },
  660. }
  661. cfg := &Config{
  662. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  663. "vertexai": {
  664. BaseURL: "custom-url",
  665. },
  666. }),
  667. }
  668. cfg.setDefaults("/tmp", "")
  669. env := env.NewFromMap(map[string]string{
  670. "GOOGLE_GENAI_USE_VERTEXAI": "false",
  671. })
  672. resolver := NewEnvironmentVariableResolver(env)
  673. err := cfg.configureProviders(env, resolver, knownProviders)
  674. require.NoError(t, err)
  675. require.Equal(t, cfg.Providers.Len(), 0)
  676. _, exists := cfg.Providers.Get("vertexai")
  677. require.False(t, exists)
  678. })
  679. t.Run("Bedrock provider removed when AWS credentials missing with existing config", func(t *testing.T) {
  680. knownProviders := []catwalk.Provider{
  681. {
  682. ID: catwalk.InferenceProviderBedrock,
  683. APIKey: "",
  684. APIEndpoint: "",
  685. Models: []catwalk.Model{{
  686. ID: "anthropic.claude-sonnet-4-20250514-v1:0",
  687. }},
  688. },
  689. }
  690. cfg := &Config{
  691. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  692. "bedrock": {
  693. BaseURL: "custom-url",
  694. },
  695. }),
  696. }
  697. cfg.setDefaults("/tmp", "")
  698. env := env.NewFromMap(map[string]string{})
  699. resolver := NewEnvironmentVariableResolver(env)
  700. err := cfg.configureProviders(env, resolver, knownProviders)
  701. require.NoError(t, err)
  702. require.Equal(t, cfg.Providers.Len(), 0)
  703. _, exists := cfg.Providers.Get("bedrock")
  704. require.False(t, exists)
  705. })
  706. t.Run("provider removed when API key missing with existing config", func(t *testing.T) {
  707. knownProviders := []catwalk.Provider{
  708. {
  709. ID: "openai",
  710. APIKey: "$MISSING_API_KEY",
  711. APIEndpoint: "https://api.openai.com/v1",
  712. Models: []catwalk.Model{{
  713. ID: "test-model",
  714. }},
  715. },
  716. }
  717. cfg := &Config{
  718. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  719. "openai": {
  720. BaseURL: "custom-url",
  721. },
  722. }),
  723. }
  724. cfg.setDefaults("/tmp", "")
  725. env := env.NewFromMap(map[string]string{})
  726. resolver := NewEnvironmentVariableResolver(env)
  727. err := cfg.configureProviders(env, resolver, knownProviders)
  728. require.NoError(t, err)
  729. require.Equal(t, cfg.Providers.Len(), 0)
  730. _, exists := cfg.Providers.Get("openai")
  731. require.False(t, exists)
  732. })
  733. t.Run("known provider should still be added if the endpoint is missing the client will use default endpoints", func(t *testing.T) {
  734. knownProviders := []catwalk.Provider{
  735. {
  736. ID: "openai",
  737. APIKey: "$OPENAI_API_KEY",
  738. APIEndpoint: "$MISSING_ENDPOINT",
  739. Models: []catwalk.Model{{
  740. ID: "test-model",
  741. }},
  742. },
  743. }
  744. cfg := &Config{
  745. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  746. "openai": {
  747. APIKey: "test-key",
  748. },
  749. }),
  750. }
  751. cfg.setDefaults("/tmp", "")
  752. env := env.NewFromMap(map[string]string{
  753. "OPENAI_API_KEY": "test-key",
  754. })
  755. resolver := NewEnvironmentVariableResolver(env)
  756. err := cfg.configureProviders(env, resolver, knownProviders)
  757. require.NoError(t, err)
  758. require.Equal(t, cfg.Providers.Len(), 1)
  759. _, exists := cfg.Providers.Get("openai")
  760. require.True(t, exists)
  761. })
  762. }
  763. func TestConfig_defaultModelSelection(t *testing.T) {
  764. t.Run("default behavior uses the default models for given provider", func(t *testing.T) {
  765. knownProviders := []catwalk.Provider{
  766. {
  767. ID: "openai",
  768. APIKey: "abc",
  769. DefaultLargeModelID: "large-model",
  770. DefaultSmallModelID: "small-model",
  771. Models: []catwalk.Model{
  772. {
  773. ID: "large-model",
  774. DefaultMaxTokens: 1000,
  775. },
  776. {
  777. ID: "small-model",
  778. DefaultMaxTokens: 500,
  779. },
  780. },
  781. },
  782. }
  783. cfg := &Config{}
  784. cfg.setDefaults("/tmp", "")
  785. env := env.NewFromMap(map[string]string{})
  786. resolver := NewEnvironmentVariableResolver(env)
  787. err := cfg.configureProviders(env, resolver, knownProviders)
  788. require.NoError(t, err)
  789. large, small, err := cfg.defaultModelSelection(knownProviders)
  790. require.NoError(t, err)
  791. require.Equal(t, "large-model", large.Model)
  792. require.Equal(t, "openai", large.Provider)
  793. require.Equal(t, int64(1000), large.MaxTokens)
  794. require.Equal(t, "small-model", small.Model)
  795. require.Equal(t, "openai", small.Provider)
  796. require.Equal(t, int64(500), small.MaxTokens)
  797. })
  798. t.Run("should error if no providers configured", func(t *testing.T) {
  799. knownProviders := []catwalk.Provider{
  800. {
  801. ID: "openai",
  802. APIKey: "$MISSING_KEY",
  803. DefaultLargeModelID: "large-model",
  804. DefaultSmallModelID: "small-model",
  805. Models: []catwalk.Model{
  806. {
  807. ID: "large-model",
  808. DefaultMaxTokens: 1000,
  809. },
  810. {
  811. ID: "small-model",
  812. DefaultMaxTokens: 500,
  813. },
  814. },
  815. },
  816. }
  817. cfg := &Config{}
  818. cfg.setDefaults("/tmp", "")
  819. env := env.NewFromMap(map[string]string{})
  820. resolver := NewEnvironmentVariableResolver(env)
  821. err := cfg.configureProviders(env, resolver, knownProviders)
  822. require.NoError(t, err)
  823. _, _, err = cfg.defaultModelSelection(knownProviders)
  824. require.Error(t, err)
  825. })
  826. t.Run("should error if model is missing", func(t *testing.T) {
  827. knownProviders := []catwalk.Provider{
  828. {
  829. ID: "openai",
  830. APIKey: "abc",
  831. DefaultLargeModelID: "large-model",
  832. DefaultSmallModelID: "small-model",
  833. Models: []catwalk.Model{
  834. {
  835. ID: "not-large-model",
  836. DefaultMaxTokens: 1000,
  837. },
  838. {
  839. ID: "small-model",
  840. DefaultMaxTokens: 500,
  841. },
  842. },
  843. },
  844. }
  845. cfg := &Config{}
  846. cfg.setDefaults("/tmp", "")
  847. env := env.NewFromMap(map[string]string{})
  848. resolver := NewEnvironmentVariableResolver(env)
  849. err := cfg.configureProviders(env, resolver, knownProviders)
  850. require.NoError(t, err)
  851. _, _, err = cfg.defaultModelSelection(knownProviders)
  852. require.Error(t, err)
  853. })
  854. t.Run("should configure the default models with a custom provider", func(t *testing.T) {
  855. knownProviders := []catwalk.Provider{
  856. {
  857. ID: "openai",
  858. APIKey: "$MISSING", // will not be included in the config
  859. DefaultLargeModelID: "large-model",
  860. DefaultSmallModelID: "small-model",
  861. Models: []catwalk.Model{
  862. {
  863. ID: "not-large-model",
  864. DefaultMaxTokens: 1000,
  865. },
  866. {
  867. ID: "small-model",
  868. DefaultMaxTokens: 500,
  869. },
  870. },
  871. },
  872. }
  873. cfg := &Config{
  874. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  875. "custom": {
  876. APIKey: "test-key",
  877. BaseURL: "https://api.custom.com/v1",
  878. Models: []catwalk.Model{
  879. {
  880. ID: "model",
  881. DefaultMaxTokens: 600,
  882. },
  883. },
  884. },
  885. }),
  886. }
  887. cfg.setDefaults("/tmp", "")
  888. env := env.NewFromMap(map[string]string{})
  889. resolver := NewEnvironmentVariableResolver(env)
  890. err := cfg.configureProviders(env, resolver, knownProviders)
  891. require.NoError(t, err)
  892. large, small, err := cfg.defaultModelSelection(knownProviders)
  893. require.NoError(t, err)
  894. require.Equal(t, "model", large.Model)
  895. require.Equal(t, "custom", large.Provider)
  896. require.Equal(t, int64(600), large.MaxTokens)
  897. require.Equal(t, "model", small.Model)
  898. require.Equal(t, "custom", small.Provider)
  899. require.Equal(t, int64(600), small.MaxTokens)
  900. })
  901. t.Run("should fail if no model configured", func(t *testing.T) {
  902. knownProviders := []catwalk.Provider{
  903. {
  904. ID: "openai",
  905. APIKey: "$MISSING", // will not be included in the config
  906. DefaultLargeModelID: "large-model",
  907. DefaultSmallModelID: "small-model",
  908. Models: []catwalk.Model{
  909. {
  910. ID: "not-large-model",
  911. DefaultMaxTokens: 1000,
  912. },
  913. {
  914. ID: "small-model",
  915. DefaultMaxTokens: 500,
  916. },
  917. },
  918. },
  919. }
  920. cfg := &Config{
  921. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  922. "custom": {
  923. APIKey: "test-key",
  924. BaseURL: "https://api.custom.com/v1",
  925. Models: []catwalk.Model{},
  926. },
  927. }),
  928. }
  929. cfg.setDefaults("/tmp", "")
  930. env := env.NewFromMap(map[string]string{})
  931. resolver := NewEnvironmentVariableResolver(env)
  932. err := cfg.configureProviders(env, resolver, knownProviders)
  933. require.NoError(t, err)
  934. _, _, err = cfg.defaultModelSelection(knownProviders)
  935. require.Error(t, err)
  936. })
  937. t.Run("should use the default provider first", func(t *testing.T) {
  938. knownProviders := []catwalk.Provider{
  939. {
  940. ID: "openai",
  941. APIKey: "set",
  942. DefaultLargeModelID: "large-model",
  943. DefaultSmallModelID: "small-model",
  944. Models: []catwalk.Model{
  945. {
  946. ID: "large-model",
  947. DefaultMaxTokens: 1000,
  948. },
  949. {
  950. ID: "small-model",
  951. DefaultMaxTokens: 500,
  952. },
  953. },
  954. },
  955. }
  956. cfg := &Config{
  957. Providers: csync.NewMapFrom(map[string]ProviderConfig{
  958. "custom": {
  959. APIKey: "test-key",
  960. BaseURL: "https://api.custom.com/v1",
  961. Models: []catwalk.Model{
  962. {
  963. ID: "large-model",
  964. DefaultMaxTokens: 1000,
  965. },
  966. },
  967. },
  968. }),
  969. }
  970. cfg.setDefaults("/tmp", "")
  971. env := env.NewFromMap(map[string]string{})
  972. resolver := NewEnvironmentVariableResolver(env)
  973. err := cfg.configureProviders(env, resolver, knownProviders)
  974. require.NoError(t, err)
  975. large, small, err := cfg.defaultModelSelection(knownProviders)
  976. require.NoError(t, err)
  977. require.Equal(t, "large-model", large.Model)
  978. require.Equal(t, "openai", large.Provider)
  979. require.Equal(t, int64(1000), large.MaxTokens)
  980. require.Equal(t, "small-model", small.Model)
  981. require.Equal(t, "openai", small.Provider)
  982. require.Equal(t, int64(500), small.MaxTokens)
  983. })
  984. }
  985. func TestConfig_configureSelectedModels(t *testing.T) {
  986. t.Run("should override defaults", func(t *testing.T) {
  987. knownProviders := []catwalk.Provider{
  988. {
  989. ID: "openai",
  990. APIKey: "abc",
  991. DefaultLargeModelID: "large-model",
  992. DefaultSmallModelID: "small-model",
  993. Models: []catwalk.Model{
  994. {
  995. ID: "larger-model",
  996. DefaultMaxTokens: 2000,
  997. },
  998. {
  999. ID: "large-model",
  1000. DefaultMaxTokens: 1000,
  1001. },
  1002. {
  1003. ID: "small-model",
  1004. DefaultMaxTokens: 500,
  1005. },
  1006. },
  1007. },
  1008. }
  1009. cfg := &Config{
  1010. Models: map[SelectedModelType]SelectedModel{
  1011. "large": {
  1012. Model: "larger-model",
  1013. },
  1014. },
  1015. }
  1016. cfg.setDefaults("/tmp", "")
  1017. env := env.NewFromMap(map[string]string{})
  1018. resolver := NewEnvironmentVariableResolver(env)
  1019. err := cfg.configureProviders(env, resolver, knownProviders)
  1020. require.NoError(t, err)
  1021. err = cfg.configureSelectedModels(knownProviders)
  1022. require.NoError(t, err)
  1023. large := cfg.Models[SelectedModelTypeLarge]
  1024. small := cfg.Models[SelectedModelTypeSmall]
  1025. require.Equal(t, "larger-model", large.Model)
  1026. require.Equal(t, "openai", large.Provider)
  1027. require.Equal(t, int64(2000), large.MaxTokens)
  1028. require.Equal(t, "small-model", small.Model)
  1029. require.Equal(t, "openai", small.Provider)
  1030. require.Equal(t, int64(500), small.MaxTokens)
  1031. })
  1032. t.Run("should be possible to use multiple providers", func(t *testing.T) {
  1033. knownProviders := []catwalk.Provider{
  1034. {
  1035. ID: "openai",
  1036. APIKey: "abc",
  1037. DefaultLargeModelID: "large-model",
  1038. DefaultSmallModelID: "small-model",
  1039. Models: []catwalk.Model{
  1040. {
  1041. ID: "large-model",
  1042. DefaultMaxTokens: 1000,
  1043. },
  1044. {
  1045. ID: "small-model",
  1046. DefaultMaxTokens: 500,
  1047. },
  1048. },
  1049. },
  1050. {
  1051. ID: "anthropic",
  1052. APIKey: "abc",
  1053. DefaultLargeModelID: "a-large-model",
  1054. DefaultSmallModelID: "a-small-model",
  1055. Models: []catwalk.Model{
  1056. {
  1057. ID: "a-large-model",
  1058. DefaultMaxTokens: 1000,
  1059. },
  1060. {
  1061. ID: "a-small-model",
  1062. DefaultMaxTokens: 200,
  1063. },
  1064. },
  1065. },
  1066. }
  1067. cfg := &Config{
  1068. Models: map[SelectedModelType]SelectedModel{
  1069. "small": {
  1070. Model: "a-small-model",
  1071. Provider: "anthropic",
  1072. MaxTokens: 300,
  1073. },
  1074. },
  1075. }
  1076. cfg.setDefaults("/tmp", "")
  1077. env := env.NewFromMap(map[string]string{})
  1078. resolver := NewEnvironmentVariableResolver(env)
  1079. err := cfg.configureProviders(env, resolver, knownProviders)
  1080. require.NoError(t, err)
  1081. err = cfg.configureSelectedModels(knownProviders)
  1082. require.NoError(t, err)
  1083. large := cfg.Models[SelectedModelTypeLarge]
  1084. small := cfg.Models[SelectedModelTypeSmall]
  1085. require.Equal(t, "large-model", large.Model)
  1086. require.Equal(t, "openai", large.Provider)
  1087. require.Equal(t, int64(1000), large.MaxTokens)
  1088. require.Equal(t, "a-small-model", small.Model)
  1089. require.Equal(t, "anthropic", small.Provider)
  1090. require.Equal(t, int64(300), small.MaxTokens)
  1091. })
  1092. t.Run("should override the max tokens only", func(t *testing.T) {
  1093. knownProviders := []catwalk.Provider{
  1094. {
  1095. ID: "openai",
  1096. APIKey: "abc",
  1097. DefaultLargeModelID: "large-model",
  1098. DefaultSmallModelID: "small-model",
  1099. Models: []catwalk.Model{
  1100. {
  1101. ID: "large-model",
  1102. DefaultMaxTokens: 1000,
  1103. },
  1104. {
  1105. ID: "small-model",
  1106. DefaultMaxTokens: 500,
  1107. },
  1108. },
  1109. },
  1110. }
  1111. cfg := &Config{
  1112. Models: map[SelectedModelType]SelectedModel{
  1113. "large": {
  1114. MaxTokens: 100,
  1115. },
  1116. },
  1117. }
  1118. cfg.setDefaults("/tmp", "")
  1119. env := env.NewFromMap(map[string]string{})
  1120. resolver := NewEnvironmentVariableResolver(env)
  1121. err := cfg.configureProviders(env, resolver, knownProviders)
  1122. require.NoError(t, err)
  1123. err = cfg.configureSelectedModels(knownProviders)
  1124. require.NoError(t, err)
  1125. large := cfg.Models[SelectedModelTypeLarge]
  1126. require.Equal(t, "large-model", large.Model)
  1127. require.Equal(t, "openai", large.Provider)
  1128. require.Equal(t, int64(100), large.MaxTokens)
  1129. })
  1130. }