Skip to content

Commit 6055f37

Browse files
thedadamsthedadams
authored
Merge pull request #561 from thedadams/deferred-cred-helper-build
feat: build credential helpers just in time
2 parents ead5eaf + 638f645 commit 6055f37

File tree

14 files changed

+94
-52
lines changed

14 files changed

+94
-52
lines changed

‎pkg/cli/credential.go

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import (
1212
"github.com/gptscript-ai/gptscript/pkg/cache"
1313
"github.com/gptscript-ai/gptscript/pkg/config"
1414
"github.com/gptscript-ai/gptscript/pkg/credentials"
15+
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
16+
"github.com/gptscript-ai/gptscript/pkg/runner"
1517
"github.com/spf13/cobra"
1618
)
1719

@@ -35,7 +37,7 @@ func (c *Credential) Customize(cmd *cobra.Command) {
3537
cmd.AddCommand(cmd2.Command(&Show{root: c.root}))
3638
}
3739

38-
func (c *Credential) Run(_ *cobra.Command, _ []string) error {
40+
func (c *Credential) Run(cmd *cobra.Command, _ []string) error {
3941
cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
4042
if err != nil {
4143
return fmt.Errorf("failed to read CLI config: %w", err)
@@ -51,14 +53,22 @@ func (c *Credential) Run(_ *cobra.Command, _ []string) error {
5153
return err
5254
}
5355
opts.Cache = cache.Complete(opts.Cache)
56+
opts.Runner = runner.Complete(opts.Runner)
57+
if opts.Runner.RuntimeManager == nil {
58+
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
59+
}
60+
61+
if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
62+
return err
63+
}
5464

5565
// Initialize the credential store and get all the credentials.
56-
store, err := credentials.NewStore(cfg, ctx, opts.Cache.CacheDir)
66+
store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, ctx, opts.Cache.CacheDir)
5767
if err != nil {
5868
return fmt.Errorf("failed to get credentials store: %w", err)
5969
}
6070

61-
creds, err := store.List()
71+
creds, err := store.List(cmd.Context())
6272
if err != nil {
6373
return fmt.Errorf("failed to list credentials: %w", err)
6474
}

‎pkg/cli/credential_delete.go

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ import (
66
"github.com/gptscript-ai/gptscript/pkg/cache"
77
"github.com/gptscript-ai/gptscript/pkg/config"
88
"github.com/gptscript-ai/gptscript/pkg/credentials"
9+
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
10+
"github.com/gptscript-ai/gptscript/pkg/runner"
911
"github.com/spf13/cobra"
1012
)
1113

@@ -21,24 +23,33 @@ func (c *Delete) Customize(cmd *cobra.Command) {
2123
cmd.Args = cobra.ExactArgs(1)
2224
}
2325

24-
func (c *Delete) Run(_ *cobra.Command, args []string) error {
26+
func (c *Delete) Run(cmd *cobra.Command, args []string) error {
2527
opts, err := c.root.NewGPTScriptOpts()
2628
if err != nil {
2729
return err
2830
}
29-
opts.Cache = cache.Complete(opts.Cache)
3031

3132
cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
3233
if err != nil {
3334
return fmt.Errorf("failed to read CLI config: %w", err)
3435
}
3536

36-
store, err := credentials.NewStore(cfg, c.root.CredentialContext, opts.Cache.CacheDir)
37+
opts.Cache = cache.Complete(opts.Cache)
38+
opts.Runner = runner.Complete(opts.Runner)
39+
if opts.Runner.RuntimeManager == nil {
40+
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
41+
}
42+
43+
if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
44+
return err
45+
}
46+
47+
store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
3748
if err != nil {
3849
return fmt.Errorf("failed to get credentials store: %w", err)
3950
}
4051

41-
if err = store.Remove(args[0]); err != nil {
52+
if err = store.Remove(cmd.Context(), args[0]); err != nil {
4253
return fmt.Errorf("failed to remove credential: %w", err)
4354
}
4455
return nil

‎pkg/cli/credential_show.go

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ import (
88
"github.com/gptscript-ai/gptscript/pkg/cache"
99
"github.com/gptscript-ai/gptscript/pkg/config"
1010
"github.com/gptscript-ai/gptscript/pkg/credentials"
11+
"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
12+
"github.com/gptscript-ai/gptscript/pkg/runner"
1113
"github.com/spf13/cobra"
1214
)
1315

@@ -23,24 +25,33 @@ func (c *Show) Customize(cmd *cobra.Command) {
2325
cmd.Args = cobra.ExactArgs(1)
2426
}
2527

26-
func (c *Show) Run(_ *cobra.Command, args []string) error {
28+
func (c *Show) Run(cmd *cobra.Command, args []string) error {
2729
opts, err := c.root.NewGPTScriptOpts()
2830
if err != nil {
2931
return err
3032
}
31-
opts.Cache = cache.Complete(opts.Cache)
3233

3334
cfg, err := config.ReadCLIConfig(c.root.ConfigFile)
3435
if err != nil {
3536
return fmt.Errorf("failed to read CLI config: %w", err)
3637
}
3738

38-
store, err := credentials.NewStore(cfg, c.root.CredentialContext, opts.Cache.CacheDir)
39+
opts.Cache = cache.Complete(opts.Cache)
40+
opts.Runner = runner.Complete(opts.Runner)
41+
if opts.Runner.RuntimeManager == nil {
42+
opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
43+
}
44+
45+
if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg, opts.Env); err != nil {
46+
return err
47+
}
48+
49+
store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
3950
if err != nil {
4051
return fmt.Errorf("failed to get credentials store: %w", err)
4152
}
4253

43-
cred, exists, err := store.Get(args[0])
54+
cred, exists, err := store.Get(cmd.Context(), args[0])
4455
if err != nil {
4556
return fmt.Errorf("failed to get credential: %w", err)
4657
}

‎pkg/cli/eval.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
5656
return err
5757
}
5858

59-
runner, err := gptscript.New(opts)
59+
runner, err := gptscript.New(cmd.Context(), opts)
6060
if err != nil {
6161
return err
6262
}

‎pkg/cli/gptscript.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
380380

381381
ctx := cmd.Context()
382382

383-
gptScript, err := gptscript.New(gptOpt)
383+
gptScript, err := gptscript.New(ctx, gptOpt)
384384
if err != nil {
385385
return err
386386
}

‎pkg/credentials/noop.go

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,21 @@
11
package credentials
22

3+
import "context"
4+
35
type NoopStore struct{}
46

5-
func (s NoopStore) Get(_ string) (*Credential, bool, error) {
7+
func (s NoopStore) Get(context.Context, string) (*Credential, bool, error) {
68
return nil, false, nil
79
}
810

9-
func (s NoopStore) Add(_ Credential) error {
11+
func (s NoopStore) Add(context.Context, Credential) error {
1012
return nil
1113
}
1214

13-
func (s NoopStore) Remove(_ string) error {
15+
func (s NoopStore) Remove(context.Context, string) error {
1416
return nil
1517
}
1618

17-
func (s NoopStore) List() ([]Credential, error) {
19+
func (s NoopStore) List(context.Context) ([]Credential, error) {
1820
return nil, nil
1921
}

‎pkg/credentials/store.go

Lines changed: 27 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package credentials
22

33
import (
4+
"context"
45
"fmt"
56
"path/filepath"
67
"regexp"
@@ -10,32 +11,38 @@ import (
1011
"github.com/gptscript-ai/gptscript/pkg/config"
1112
)
1213

14+
type CredentialBuilder interface {
15+
EnsureCredentialHelpers(ctx context.Context) error
16+
}
17+
1318
type CredentialStore interface {
14-
Get(toolName string) (*Credential, bool, error)
15-
Add(cred Credential) error
16-
Remove(toolName string) error
17-
List() ([]Credential, error)
19+
Get(ctx context.Context, toolName string) (*Credential, bool, error)
20+
Add(ctx context.Context, cred Credential) error
21+
Remove(ctx context.Context, toolName string) error
22+
List(ctx context.Context) ([]Credential, error)
1823
}
1924

2025
type Store struct {
2126
credCtx string
27+
credBuilder CredentialBuilder
2228
credHelperDirs CredentialHelperDirs
2329
cfg *config.CLIConfig
2430
}
2531

26-
func NewStore(cfg *config.CLIConfig, credCtx, cacheDir string) (CredentialStore, error) {
32+
func NewStore(cfg *config.CLIConfig, credentialBuilder CredentialBuilder, credCtx, cacheDir string) (CredentialStore, error) {
2733
if err := validateCredentialCtx(credCtx); err != nil {
2834
return nil, err
2935
}
3036
return Store{
3137
credCtx: credCtx,
38+
credBuilder: credentialBuilder,
3239
credHelperDirs: GetCredentialHelperDirs(cacheDir),
3340
cfg: cfg,
3441
}, nil
3542
}
3643

37-
func (s Store) Get(toolName string) (*Credential, bool, error) {
38-
store, err := s.getStore()
44+
func (s Store) Get(ctx context.Context, toolName string) (*Credential, bool, error) {
45+
store, err := s.getStore(ctx)
3946
if err != nil {
4047
return nil, false, err
4148
}
@@ -57,9 +64,9 @@ func (s Store) Get(toolName string) (*Credential, bool, error) {
5764
return &cred, true, nil
5865
}
5966

60-
func (s Store) Add(cred Credential) error {
67+
func (s Store) Add(ctx context.Context, cred Credential) error {
6168
cred.Context = s.credCtx
62-
store, err := s.getStore()
69+
store, err := s.getStore(ctx)
6370
if err != nil {
6471
return err
6572
}
@@ -70,16 +77,16 @@ func (s Store) Add(cred Credential) error {
7077
return store.Store(auth)
7178
}
7279

73-
func (s Store) Remove(toolName string) error {
74-
store, err := s.getStore()
80+
func (s Store) Remove(ctx context.Context, toolName string) error {
81+
store, err := s.getStore(ctx)
7582
if err != nil {
7683
return err
7784
}
7885
return store.Erase(toolNameWithCtx(toolName, s.credCtx))
7986
}
8087

81-
func (s Store) List() ([]Credential, error) {
82-
store, err := s.getStore()
88+
func (s Store) List(ctx context.Context) ([]Credential, error) {
89+
store, err := s.getStore(ctx)
8390
if err != nil {
8491
return nil, err
8592
}
@@ -106,17 +113,21 @@ func (s Store) List() ([]Credential, error) {
106113
return creds, nil
107114
}
108115

109-
func (s *Store) getStore() (credentials.Store, error) {
110-
return s.getStoreByHelper(config.GPTScriptHelperPrefix + s.cfg.CredentialsStore)
116+
func (s *Store) getStore(ctx context.Context) (credentials.Store, error) {
117+
return s.getStoreByHelper(ctx, config.GPTScriptHelperPrefix+s.cfg.CredentialsStore)
111118
}
112119

113-
func (s *Store) getStoreByHelper(helper string) (credentials.Store, error) {
120+
func (s *Store) getStoreByHelper(ctx context.Context, helper string) (credentials.Store, error) {
114121
if helper == "" || helper == config.GPTScriptHelperPrefix+"file" {
115122
return credentials.NewFileStore(s.cfg), nil
116123
}
117124

118125
// If the helper is referencing one of the credential helper programs, then reference the full path.
119126
if strings.HasPrefix(helper, "gptscript-credential-") {
127+
if err := s.credBuilder.EnsureCredentialHelpers(ctx); err != nil {
128+
return nil, err
129+
}
130+
120131
helper = filepath.Join(s.credHelperDirs.BinDir, helper)
121132
}
122133

‎pkg/gptscript/gptscript.go

Lines changed: 5 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ func complete(opts ...Options) Options {
7777
return result
7878
}
7979

80-
func New(o ...Options) (*GPTScript, error) {
80+
func New(ctx context.Context, o ...Options) (*GPTScript, error) {
8181
opts := complete(o...)
8282
registry := llm.NewRegistry()
8383

@@ -91,23 +91,20 @@ func New(o ...Options) (*GPTScript, error) {
9191
return nil, err
9292
}
9393

94-
credStore, err := credentials.NewStore(cliCfg, opts.CredentialContext, cacheClient.CacheDir())
95-
if err != nil {
96-
return nil, err
97-
}
98-
9994
if opts.Runner.RuntimeManager == nil {
10095
opts.Runner.RuntimeManager = runtimes.Default(cacheClient.CacheDir())
10196
}
10297

10398
if err := opts.Runner.RuntimeManager.SetUpCredentialHelpers(context.Background(), cliCfg, opts.Env); err != nil {
10499
return nil, err
105100
}
106-
if err := opts.Runner.RuntimeManager.EnsureCredentialHelpers(context.Background()); err != nil {
101+
102+
credStore, err := credentials.NewStore(cliCfg, opts.Runner.RuntimeManager, opts.CredentialContext, cacheClient.CacheDir())
103+
if err != nil {
107104
return nil, err
108105
}
109106

110-
oaiClient, err := openai.NewClient(credStore, opts.OpenAI, openai.Options{
107+
oaiClient, err := openai.NewClient(ctx, credStore, opts.OpenAI, openai.Options{
111108
Cache: cacheClient,
112109
SetSeed: true,
113110
})

‎pkg/openai/client.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,15 +93,15 @@ func complete(opts ...Options) (Options, error) {
9393
return result, err
9494
}
9595

96-
func NewClient(credStore credentials.CredentialStore, opts ...Options) (*Client, error) {
96+
func NewClient(ctx context.Context, credStore credentials.CredentialStore, opts ...Options) (*Client, error) {
9797
opt, err := complete(opts...)
9898
if err != nil {
9999
return nil, err
100100
}
101101

102102
// If the API key is not set, try to get it from the cred store
103103
if opt.APIKey == "" && opt.BaseURL == "" {
104-
cred, exists, err := credStore.Get(BuiltinCredName)
104+
cred, exists, err := credStore.Get(ctx, BuiltinCredName)
105105
if err != nil {
106106
return nil, err
107107
}

‎pkg/prompt/credential.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ import (
99
)
1010

1111
func GetModelProviderCredential(ctx context.Context, credStore credentials.CredentialStore, credName, env, message string, envs []string) (string, error) {
12-
cred, exists, err := credStore.Get(credName)
12+
cred, exists, err := credStore.Get(ctx, credName)
1313
if err != nil {
1414
return "", err
1515
}
@@ -25,7 +25,7 @@ func GetModelProviderCredential(ctx context.Context, credStore credentials.Crede
2525
}
2626

2727
k = gjson.Get(result, "key").String()
28-
if err := credStore.Add(credentials.Credential{
28+
if err := credStore.Add(ctx, credentials.Credential{
2929
ToolName: credName,
3030
Type: credentials.CredentialTypeModelProvider,
3131
Env: map[string]string{

‎pkg/remote/remote.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ func (c *Client) clientFromURL(ctx context.Context, apiURL string) (*openai.Clie
116116
}
117117
}
118118

119-
return openai.NewClient(c.credStore, openai.Options{
119+
return openai.NewClient(ctx, c.credStore, openai.Options{
120120
BaseURL: apiURL,
121121
Cache: c.cache,
122122
APIKey: key,
@@ -163,7 +163,7 @@ func (c *Client) load(ctx context.Context, toolName string) (*openai.Client, err
163163
url += "/v1"
164164
}
165165

166-
client, err = openai.NewClient(c.credStore, openai.Options{
166+
client, err = openai.NewClient(ctx, c.credStore, openai.Options{
167167
BaseURL: url,
168168
Cache: c.cache,
169169
CacheKey: prg.EntryToolID,

0 commit comments

Comments
 (0)