19c0e04 · crush · git.iko.soy

git.iko.soy/crush

feat: make it possible for api key to be programatic

commit19c0e04c7bb8be287931346bc53243c805958e4e authorKujtim Hoxha <kujtimii.h@gmail.com> dateSat 28 Jun 2025, 23:21 (1 year ago) parent03871a5

8 files changed+154−65

compare with main

unifiedsplit

cspell.json+1−39 internal/config/config.go+3−3 internal/config/shell.go+74 internal/llm/agent/agent.go+2−2 internal/llm/provider/anthropic.go+22−10 internal/llm/provider/gemini.go+22−1 internal/llm/provider/openai.go+22−9 internal/llm/provider/provider.go+8−1

cspell.json+1−39

@@ -1,39 +1 @@

1−{

2− "flagWords": [],

3− "words": [

4− "afero",

5− "alecthomas",

6− "bubbletea",

7− "charmbracelet",

8− "charmtone",

9− "Charple",

10− "crush",

11− "diffview",

12− "Emph",

13− "filepicker",

14− "Focusable",

15− "fsext",

16− "GROQ",

17− "Guac",

18− "imageorient",

19− "Lanczos",

20− "lipgloss",

21− "lsps",

22− "lucasb",

23− "nfnt",

24− "oksvg",

25− "Preproc",

26− "rasterx",

27− "rivo",

28− "Sourcegraph",

29− "srwiley",

30− "Strikethrough",

31− "termenv",

32− "textinput",

33− "trashhalo",

34− "uniseg",

35− "Unticked"

36− ],

37− "version": "0.2",

38− "language": "en"

39−}

1+{"language":"en","flagWords":[],"version":"0.2","words":["afero","alecthomas","bubbletea","charmbracelet","charmtone","Charple","crush","diffview","Emph","filepicker","Focusable","fsext","GROQ","Guac","imageorient","Lanczos","lipgloss","lsps","lucasb","nfnt","oksvg","Preproc","rasterx","rivo","Sourcegraph","srwiley","Strikethrough","termenv","textinput","trashhalo","uniseg","Unticked","genai"]}

internal/config/config.go+3−3

@@ -592,11 +592,11 @@ func mergeProviderConfigs(base, global, local *Config) {

592592 if cfg == nil {

593593 continue

594594 }

595− for providerName, globalProvider := range cfg.Providers {

595+ for providerName, p := range cfg.Providers {

596596 if _, ok := base.Providers[providerName]; !ok {

597− base.Providers[providerName] = globalProvider

597+ base.Providers[providerName] = p

598598 } else {

599− base.Providers[providerName] = mergeProviderConfig(providerName, base.Providers[providerName], globalProvider)

599+ base.Providers[providerName] = mergeProviderConfig(providerName, base.Providers[providerName], p)

600600 }

601601 }

602602 }

internal/config/shell.gonew file+74

@@ -0,0 +1,74 @@

1+package config

3+import (

4+ "context"

5+ "fmt"

6+ "os"

7+ "strings"

8+ "time"

10+ "github.com/charmbracelet/crush/internal/llm/tools/shell"

11+ "github.com/charmbracelet/crush/internal/logging"

12+)

13+

14+// ExecuteCommand executes a shell command and returns the output

15+// This is a shared utility that can be used by both provider config and tools

16+func ExecuteCommand(ctx context.Context, command string, workingDir string) (string, error) {

17+ if workingDir == "" {

18+ workingDir = WorkingDirectory()

19+ }

20+

21+ persistentShell := shell.GetPersistentShell(workingDir)

22+

23+ stdout, stderr, err := persistentShell.Exec(ctx, command)

24+ if err != nil {

25+ logging.Debug("Command execution failed", "command", command, "error", err, "stderr", stderr)

26+ return "", fmt.Errorf("command execution failed: %w", err)

27+ }

28+

29+ return strings.TrimSpace(stdout), nil

30+}

31+

32+// ResolveAPIKey resolves an API key that can be either:

33+// - A direct string value

34+// - An environment variable (prefixed with $)

35+// - A shell command (wrapped in $(...))

36+func ResolveAPIKey(apiKey string) (string, error) {

37+ if !strings.HasPrefix(apiKey, "$") {

38+ return apiKey, nil

39+ }

40+

41+ if strings.HasPrefix(apiKey, "$(") && strings.HasSuffix(apiKey, ")") {

42+ command := strings.TrimSuffix(strings.TrimPrefix(apiKey, "$("), ")")

43+ logging.Debug("Resolving API key from command", "command", command)

44+ return resolveCommandAPIKey(command)

45+ }

46+

47+ envVar := strings.TrimPrefix(apiKey, "$")

48+ if value := os.Getenv(envVar); value != "" {

49+ logging.Debug("Resolved environment variable", "envVar", envVar, "value", value)

50+ return value, nil

51+ }

52+

53+ logging.Debug("Environment variable not found", "envVar", envVar)

54+

55+ return "", fmt.Errorf("environment variable %s not found", envVar)

56+}

57+

58+// resolveCommandAPIKey executes a command to get an API key, with caching support

59+func resolveCommandAPIKey(command string) (string, error) {

60+ ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)

61+ defer cancel()

62+

63+ logging.Debug("Executing command for API key", "command", command)

64+

65+ workingDir := WorkingDirectory()

66+

67+ result, err := ExecuteCommand(ctx, command, workingDir)

68+ if err != nil {

69+ return "", fmt.Errorf("failed to execute API key command: %w", err)

70+ }

71+ logging.Debug("Command executed successfully", "command", command, "result", result)

72+ return result, nil

73+}

74+

internal/llm/agent/agent.go+2−2

@@ -253,10 +253,10 @@ func (a *agent) IsBusy() bool {

253253 if cancelFunc, ok := value.(context.CancelFunc); ok {

254254 if cancelFunc != nil {

255255 busy = true

256− return false // Stop iterating

256+ return false

257257 }

258258 }

259− return true // Continue iterating

259+ return true

260260 })

261261 return busy

262262}

internal/llm/provider/anthropic.go+22−10

@@ -21,12 +21,20 @@ import (

2121

2222type anthropicClient struct {

2323 providerOptions providerClientOptions

24+ useBedrock bool

2425 client anthropic.Client

2526}

2627

2728type AnthropicClient ProviderClient

2829

2930func newAnthropicClient(opts providerClientOptions, useBedrock bool) AnthropicClient {

31+ return &anthropicClient{

32+ providerOptions: opts,

33+ client: createAnthropicClient(opts, useBedrock),

34+ }

35+}

36+

37+func createAnthropicClient(opts providerClientOptions, useBedrock bool) anthropic.Client {

3038 anthropicClientOptions := []option.RequestOption{}

3139 if opts.apiKey != "" {

3240 anthropicClientOptions = append(anthropicClientOptions, option.WithAPIKey(opts.apiKey))

@@ -34,12 +42,7 @@ func newAnthropicClient(opts providerClientOptions, useBedrock bool) AnthropicCl

3442 if useBedrock {

3543 anthropicClientOptions = append(anthropicClientOptions, bedrock.WithLoadDefaultConfig(context.Background()))

3644 }

37−

38− client := anthropic.NewClient(anthropicClientOptions...)

39− return &anthropicClient{

40− providerOptions: opts,

41− client: client,

42− }

45+ return anthropic.NewClient(anthropicClientOptions...)

4346}

4447

4548func (a *anthropicClient) convertMessages(messages []message.Message) (anthropicMessages []anthropic.MessageParam) {

@@ -385,12 +388,21 @@ func (a *anthropicClient) stream(ctx context.Context, messages []message.Message

385388}

386389

387390func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, error) {

388− var apierr *anthropic.Error

389− if !errors.As(err, &apierr) {

391+ var apiErr *anthropic.Error

392+ if !errors.As(err, &apiErr) {

390393 return false, 0, err

391394 }

392395

393− if apierr.StatusCode != 429 && apierr.StatusCode != 529 {

396+ if apiErr.StatusCode == 401 {

397+ a.providerOptions.apiKey, err = config.ResolveAPIKey(a.providerOptions.config.APIKey)

398+ if err != nil {

399+ return false, 0, fmt.Errorf("failed to resolve API key: %w", err)

400+ }

401+ a.client = createAnthropicClient(a.providerOptions, a.useBedrock)

402+ return true, 0, nil

403+ }

404+

405+ if apiErr.StatusCode != 429 && apiErr.StatusCode != 529 {

394406 return false, 0, err

395407 }

396408

@@ -399,7 +411,7 @@ func (a *anthropicClient) shouldRetry(attempts int, err error) (bool, int64, err

399411 }

400412

401413 retryMs := 0

402− retryAfterValues := apierr.Response.Header.Values("Retry-After")

414+ retryAfterValues := apiErr.Response.Header.Values("Retry-After")

403415

404416 backoffMs := 2000 * (1 << (attempts - 1))

405417 jitterMs := int(float64(backoffMs) * 0.2)

internal/llm/provider/gemini.go+22−1

@@ -25,7 +25,7 @@ type geminiClient struct {

2525type GeminiClient ProviderClient

2626

2727func newGeminiClient(opts providerClientOptions) GeminiClient {

28− client, err := genai.NewClient(context.Background(), &genai.ClientConfig{APIKey: opts.apiKey, Backend: genai.BackendGeminiAPI})

28+ client, err := createGeminiClient(opts)

2929 if err != nil {

3030 logging.Error("Failed to create Gemini client", "error", err)

3131 return nil

@@ -37,6 +37,14 @@ func newGeminiClient(opts providerClientOptions) GeminiClient {

3737 }

3838}

3939

40+func createGeminiClient(opts providerClientOptions) (*genai.Client, error) {

41+ client, err := genai.NewClient(context.Background(), &genai.ClientConfig{APIKey: opts.apiKey, Backend: genai.BackendGeminiAPI})

42+ if err != nil {

43+ return nil, err

44+ }

45+ return client, nil

46+}

47+

4048func (g *geminiClient) convertMessages(messages []message.Message) []*genai.Content {

4149 var history []*genai.Content

4250 for _, msg := range messages {

@@ -414,6 +422,19 @@ func (g *geminiClient) shouldRetry(attempts int, err error) (bool, int64, error)

414422 errMsg := err.Error()

415423 isRateLimit := contains(errMsg, "rate limit", "quota exceeded", "too many requests")

416424

425+ // Check for token expiration (401 Unauthorized)

426+ if contains(errMsg, "unauthorized", "invalid api key", "api key expired") {

427+ g.providerOptions.apiKey, err = config.ResolveAPIKey(g.providerOptions.config.APIKey)

428+ if err != nil {

429+ return false, 0, fmt.Errorf("failed to resolve API key: %w", err)

430+ }

431+ g.client, err = createGeminiClient(g.providerOptions)

432+ if err != nil {

433+ return false, 0, fmt.Errorf("failed to create Gemini client after API key refresh: %w", err)

434+ }

435+ return true, 0, nil

436+ }

437+

417438 // Check for common rate limit error messages

418439

419440 if !isRateLimit {

internal/llm/provider/openai.go+22−9

@@ -26,6 +26,13 @@ type openaiClient struct {

2626type OpenAIClient ProviderClient

2727

2828func newOpenAIClient(opts providerClientOptions) OpenAIClient {

29+ return &openaiClient{

30+ providerOptions: opts,

31+ client: createOpenAIClient(opts),

32+ }

33+}

34+

35+func createOpenAIClient(opts providerClientOptions) openai.Client {

2936 openaiClientOptions := []option.RequestOption{}

3037 if opts.apiKey != "" {

3138 openaiClientOptions = append(openaiClientOptions, option.WithAPIKey(opts.apiKey))

@@ -40,11 +47,7 @@ func newOpenAIClient(opts providerClientOptions) OpenAIClient {

4047 }

4148 }

4249

43− client := openai.NewClient(openaiClientOptions...)

44− return &openaiClient{

45− providerOptions: opts,

46− client: client,

47− }

50+ return openai.NewClient(openaiClientOptions...)

4851}

4952

5053func (o *openaiClient) convertMessages(messages []message.Message) (openaiMessages []openai.ChatCompletionMessageParamUnion) {

@@ -339,12 +342,22 @@ func (o *openaiClient) stream(ctx context.Context, messages []message.Message, t

339342}

340343

341344func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error) {

342− var apierr *openai.Error

343− if !errors.As(err, &apierr) {

345+ var apiErr *openai.Error

346+ if !errors.As(err, &apiErr) {

344347 return false, 0, err

345348 }

346349

347− if apierr.StatusCode != 429 && apierr.StatusCode != 500 {

350+ // Check for token expiration (401 Unauthorized)

351+ if apiErr.StatusCode == 401 {

352+ o.providerOptions.apiKey, err = config.ResolveAPIKey(o.providerOptions.config.APIKey)

353+ if err != nil {

354+ return false, 0, fmt.Errorf("failed to resolve API key: %w", err)

355+ }

356+ o.client = createOpenAIClient(o.providerOptions)

357+ return true, 0, nil

358+ }

359+

360+ if apiErr.StatusCode != 429 && apiErr.StatusCode != 500 {

348361 return false, 0, err

349362 }

350363

@@ -353,7 +366,7 @@ func (o *openaiClient) shouldRetry(attempts int, err error) (bool, int64, error)

353366 }

354367

355368 retryMs := 0

356− retryAfterValues := apierr.Response.Header.Values("Retry-After")

369+ retryAfterValues := apiErr.Response.Header.Values("Retry-After")

357370

358371 backoffMs := 2000 * (1 << (attempts - 1))

359372 jitterMs := int(float64(backoffMs) * 0.2)

internal/llm/provider/provider.go+8−1

@@ -60,6 +60,7 @@ type Provider interface {

6060

6161type providerClientOptions struct {

6262 baseURL string

63+ config config.ProviderConfig

6364 apiKey string

6465 modelType config.ModelType

6566 model func(config.ModelType) config.Model

@@ -134,9 +135,15 @@ func WithMaxTokens(maxTokens int64) ProviderClientOption {

134135}

135136

136137func NewProvider(cfg config.ProviderConfig, opts ...ProviderClientOption) (Provider, error) {

138+ resolvedAPIKey, err := config.ResolveAPIKey(cfg.APIKey)

139+ if err != nil {

140+ return nil, fmt.Errorf("failed to resolve API key for provider %s: %w", cfg.ID, err)

141+ }

142+

137143 clientOptions := providerClientOptions{

138144 baseURL: cfg.BaseURL,

139− apiKey: cfg.APIKey,

145+ config: cfg,

146+ apiKey: resolvedAPIKey,

140147 extraHeaders: cfg.ExtraHeaders,

141148 model: func(tp config.ModelType) config.Model {

142149 return config.GetModel(tp)