feat: add stacked credential contexts (#849)

g-linville · web-flow · commit 780e07eaeb40 · 2024-09-18T09:19:30.000-04:00
Signed-off-by: Grant Linville &lt;grant@acorn.io&gt;
diff --git a/docs/docs/03-tools/04-credential-tools.md b/docs/docs/03-tools/04-credential-tools.md
@@ -222,3 +222,55 @@ import os
 
 print("myCred expires at " + os.getenv("GPTSCRIPT_CREDENTIAL_EXPIRATION", ""))
 ```
+
+## Stacked Credential Contexts (Advanced)
+
+When setting the `--credential-context` argument in GPTScript, you can specify multiple contexts separated by commas.
+We refer to this as "stacked credential contexts", or just stacked contexts for short. This allows you to specify an order
+of priority for credential contexts. This is best explained by example.
+
+### Example: stacked contexts when running a script that uses a credential
+
+Let's say you have two contexts, `one` and `two`, and you specify them like this:
+
+```bash
+gptscript --credential-context one,two my-script.gpt
+```
+
+```
+Credential: my-credential-tool.gpt as myCred
+
+<tool stuff here>
+```
+
+When GPTScript runs, it will first look for a credential called `myCred` in the `one` context.
+If it doesn't find it there, it will look for it in the `two` context. If it also doesn't find it there,
+it will run the `my-credential-tool.gpt` tool to get the credential. It will then store the new credential into the `one`
+context, since that has the highest priority.
+
+### Example: stacked contexts when listing credentials
+
+```bash
+gptscript --credential-context one,two credentials
+```
+
+When you list credentials like this, GPTScript will print out the information for all credentials in contexts one and two,
+with one exception. If there is a credential name that exists in both contexts, GPTScript will only print the information
+for the credential in the context with the highest priority, which in this case is `one`.
+
+(To see all credentials in all contexts, you can still use the `--all-contexts` flag, and it will show all credentials,
+regardless of whether the same name appears in another context.)
+
+### Example: stacked contexts when showing credentials
+
+```bash
+gptscript --credential-context one,two credential show myCred
+```
+
+When you show a credential like this, GPTScript will first look for `myCred` in the `one` context. If it doesn't find it
+there, it will look for it in the `two` context. If it doesn't find it in either context, it will print an error message.
+
+:::note
+You cannot specify stacked contexts when doing `gptscript credential delete`. GPTScript will return an error if
+more than one context is specified for this command.
+:::
diff --git a/docs/docs/04-command-line-reference/gptscript.md b/docs/docs/04-command-line-reference/gptscript.md
@@ -18,7 +18,7 @@ gptscript [flags] PROGRAM_FILE [INPUT...]
       --color                               Use color in output (default true) ($GPTSCRIPT_COLOR)
       --config string                       Path to GPTScript config file ($GPTSCRIPT_CONFIG)
       --confirm                             Prompt before running potentially dangerous commands ($GPTSCRIPT_CONFIRM)
-      --credential-context string           Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings          Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
       --credential-override strings         Credentials to override (ex: --credential-override github.com/example/cred-tool:API_TOKEN=1234) ($GPTSCRIPT_CREDENTIAL_OVERRIDE)
       --debug                               Enable debug logging ($GPTSCRIPT_DEBUG)
       --debug-messages                      Enable logging of chat completion calls ($GPTSCRIPT_DEBUG_MESSAGES)
diff --git a/docs/docs/04-command-line-reference/gptscript_credential.md b/docs/docs/04-command-line-reference/gptscript_credential.md
@@ -20,7 +20,7 @@ gptscript credential [flags]
 ### Options inherited from parent commands
 
 ```
-      --credential-context string   Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings   Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
 ```
 
 ### SEE ALSO
diff --git a/docs/docs/04-command-line-reference/gptscript_credential_delete.md b/docs/docs/04-command-line-reference/gptscript_credential_delete.md
@@ -18,7 +18,7 @@ gptscript credential delete <credential name> [flags]
 ### Options inherited from parent commands
 
 ```
-      --credential-context string   Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings   Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
 ```
 
 ### SEE ALSO
diff --git a/docs/docs/04-command-line-reference/gptscript_credential_show.md b/docs/docs/04-command-line-reference/gptscript_credential_show.md
@@ -18,7 +18,7 @@ gptscript credential show <credential name> [flags]
 ### Options inherited from parent commands
 
 ```
-      --credential-context string   Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings   Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
 ```
 
 ### SEE ALSO
diff --git a/docs/docs/04-command-line-reference/gptscript_eval.md b/docs/docs/04-command-line-reference/gptscript_eval.md
@@ -30,7 +30,7 @@ gptscript eval [flags]
       --color                           Use color in output (default true) ($GPTSCRIPT_COLOR)
       --config string                   Path to GPTScript config file ($GPTSCRIPT_CONFIG)
       --confirm                         Prompt before running potentially dangerous commands ($GPTSCRIPT_CONFIRM)
-      --credential-context string       Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings      Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
       --credential-override strings     Credentials to override (ex: --credential-override github.com/example/cred-tool:API_TOKEN=1234) ($GPTSCRIPT_CREDENTIAL_OVERRIDE)
       --debug                           Enable debug logging ($GPTSCRIPT_DEBUG)
       --debug-messages                  Enable logging of chat completion calls ($GPTSCRIPT_DEBUG_MESSAGES)
diff --git a/docs/docs/04-command-line-reference/gptscript_fmt.md b/docs/docs/04-command-line-reference/gptscript_fmt.md
@@ -24,7 +24,7 @@ gptscript fmt [flags]
       --color                           Use color in output (default true) ($GPTSCRIPT_COLOR)
       --config string                   Path to GPTScript config file ($GPTSCRIPT_CONFIG)
       --confirm                         Prompt before running potentially dangerous commands ($GPTSCRIPT_CONFIRM)
-      --credential-context string       Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings      Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
       --credential-override strings     Credentials to override (ex: --credential-override github.com/example/cred-tool:API_TOKEN=1234) ($GPTSCRIPT_CREDENTIAL_OVERRIDE)
       --debug                           Enable debug logging ($GPTSCRIPT_DEBUG)
       --debug-messages                  Enable logging of chat completion calls ($GPTSCRIPT_DEBUG_MESSAGES)
diff --git a/docs/docs/04-command-line-reference/gptscript_getenv.md b/docs/docs/04-command-line-reference/gptscript_getenv.md
@@ -23,7 +23,7 @@ gptscript getenv [flags] KEY [DEFAULT]
       --color                           Use color in output (default true) ($GPTSCRIPT_COLOR)
       --config string                   Path to GPTScript config file ($GPTSCRIPT_CONFIG)
       --confirm                         Prompt before running potentially dangerous commands ($GPTSCRIPT_CONFIRM)
-      --credential-context string       Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings      Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
       --credential-override strings     Credentials to override (ex: --credential-override github.com/example/cred-tool:API_TOKEN=1234) ($GPTSCRIPT_CREDENTIAL_OVERRIDE)
       --debug                           Enable debug logging ($GPTSCRIPT_DEBUG)
       --debug-messages                  Enable logging of chat completion calls ($GPTSCRIPT_DEBUG_MESSAGES)
diff --git a/docs/docs/04-command-line-reference/gptscript_parse.md b/docs/docs/04-command-line-reference/gptscript_parse.md
@@ -24,7 +24,7 @@ gptscript parse [flags]
       --color                           Use color in output (default true) ($GPTSCRIPT_COLOR)
       --config string                   Path to GPTScript config file ($GPTSCRIPT_CONFIG)
       --confirm                         Prompt before running potentially dangerous commands ($GPTSCRIPT_CONFIRM)
-      --credential-context string       Context name in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT) (default "default")
+      --credential-context strings      Context name(s) in which to store credentials ($GPTSCRIPT_CREDENTIAL_CONTEXT)
       --credential-override strings     Credentials to override (ex: --credential-override github.com/example/cred-tool:API_TOKEN=1234) ($GPTSCRIPT_CREDENTIAL_OVERRIDE)
       --debug                           Enable debug logging ($GPTSCRIPT_DEBUG)
       --debug-messages                  Enable logging of chat completion calls ($GPTSCRIPT_DEBUG_MESSAGES)
diff --git a/integration/cred_test.go b/integration/cred_test.go
@@ -45,3 +45,57 @@ func TestCredentialExpirationEnv(t *testing.T) {
 		}
 	}
 }
+
+// TestStackedCredentialContexts tests creating, using, listing, showing, and deleting credentials when there are multiple contexts.
+func TestStackedCredentialContexts(t *testing.T) {
+	// First, test credential creation. We will create a credential called testcred in two different contexts called one and two.
+	_, err := RunScript("scripts/cred_stacked.gpt", "--sub-tool", "testcred_one", "--credential-context", "one,two")
+	require.NoError(t, err)
+
+	_, err = RunScript("scripts/cred_stacked.gpt", "--sub-tool", "testcred_two", "--credential-context", "two")
+	require.NoError(t, err)
+
+	// Next, we try running the testcred_one tool. It should print the value of "testcred" in whichever context it finds the cred first.
+	out, err := RunScript("scripts/cred_stacked.gpt", "--sub-tool", "testcred_one", "--credential-context", "one,two")
+	require.NoError(t, err)
+	require.Contains(t, out, "one")
+	require.NotContains(t, out, "two")
+
+	out, err = RunScript("scripts/cred_stacked.gpt", "--sub-tool", "testcred_one", "--credential-context", "two,one")
+	require.NoError(t, err)
+	require.Contains(t, out, "two")
+	require.NotContains(t, out, "one")
+
+	// Next, list credentials and specify both contexts. We should get the credential from the first specified context.
+	out, err = GPTScriptExec("--credential-context", "one,two", "cred")
+	require.NoError(t, err)
+	require.Contains(t, out, "one")
+	require.NotContains(t, out, "two")
+
+	out, err = GPTScriptExec("--credential-context", "two,one", "cred")
+	require.NoError(t, err)
+	require.Contains(t, out, "two")
+	require.NotContains(t, out, "one")
+
+	// Next, try showing the credentials.
+	out, err = GPTScriptExec("--credential-context", "one,two", "cred", "show", "testcred")
+	require.NoError(t, err)
+	require.Contains(t, out, "one")
+	require.NotContains(t, out, "two")
+
+	out, err = GPTScriptExec("--credential-context", "two,one", "cred", "show", "testcred")
+	require.NoError(t, err)
+	require.Contains(t, out, "two")
+	require.NotContains(t, out, "one")
+
+	// Make sure we get an error if we try to delete a credential with multiple contexts specified.
+	_, err = GPTScriptExec("--credential-context", "one,two", "cred", "delete", "testcred")
+	require.Error(t, err)
+
+	// Now actually delete the credentials.
+	_, err = GPTScriptExec("--credential-context", "one", "cred", "delete", "testcred")
+	require.NoError(t, err)
+
+	_, err = GPTScriptExec("--credential-context", "two", "cred", "delete", "testcred")
+	require.NoError(t, err)
+}
diff --git a/integration/scripts/cred_stacked.gpt b/integration/scripts/cred_stacked.gpt
@@ -0,0 +1,36 @@
+name: testcred_one
+credential: cred_one as testcred
+
+#!python3
+
+import os
+
+print(os.environ.get("VALUE"))
+
+---
+name: testcred_two
+credential: cred_two as testcred
+
+#!python3
+
+import os
+
+print(os.environ.get("VALUE"))
+
+---
+name: cred_one
+
+#!python3
+
+import json
+
+print(json.dumps({"env": {"VALUE": "one"}}))
+
+---
+name: cred_two
+
+#!python3
+
+import json
+
+print(json.dumps({"env": {"VALUE": "two"}}))
diff --git a/pkg/cli/credential.go b/pkg/cli/credential.go
@@ -9,11 +9,10 @@ import (
 	"time"
 
 	cmd2 "github.com/gptscript-ai/cmd"
-	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/config"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
+	"github.com/gptscript-ai/gptscript/pkg/gptscript"
 	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
-	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/spf13/cobra"
 )
 
@@ -43,27 +42,26 @@ func (c *Credential) Run(cmd *cobra.Command, _ []string) error {
 		return fmt.Errorf("failed to read CLI config: %w", err)
 	}
 
-	ctx := c.root.CredentialContext
-	if c.AllContexts {
-		ctx = credentials.AllCredentialContexts
-	}
-
 	opts, err := c.root.NewGPTScriptOpts()
 	if err != nil {
 		return err
 	}
-	opts.Cache = cache.Complete(opts.Cache)
-	opts.Runner = runner.Complete(opts.Runner)
+	opts = gptscript.Complete(opts)
 	if opts.Runner.RuntimeManager == nil {
 		opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
 	}
 
+	ctxs := opts.CredentialContexts
+	if c.AllContexts {
+		ctxs = []string{credentials.AllCredentialContexts}
+	}
+
 	if err = opts.Runner.RuntimeManager.SetUpCredentialHelpers(cmd.Context(), cfg); err != nil {
 		return err
 	}
 
 	// Initialize the credential store and get all the credentials.
-	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, ctx, opts.Cache.CacheDir)
+	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, ctxs, opts.Cache.CacheDir)
 	if err != nil {
 		return fmt.Errorf("failed to get credentials store: %w", err)
 	}
@@ -77,7 +75,7 @@ func (c *Credential) Run(cmd *cobra.Command, _ []string) error {
 	defer w.Flush()
 
 	// Sort credentials and print column names, depending on the options.
-	if c.AllContexts {
+	if c.AllContexts || len(c.root.CredentialContext) > 1 {
 		// Sort credentials by context
 		sort.Slice(creds, func(i, j int) bool {
 			if creds[i].Context == creds[j].Context {
@@ -114,7 +112,7 @@ func (c *Credential) Run(cmd *cobra.Command, _ []string) error {
 		}
 
 		var fields []any
-		if c.AllContexts {
+		if c.AllContexts || len(c.root.CredentialContext) > 1 {
 			fields = []any{cred.Context, cred.ToolName, expires}
 		} else {
 			fields = []any{cred.ToolName, expires}
diff --git a/pkg/cli/credential_delete.go b/pkg/cli/credential_delete.go
@@ -3,11 +3,10 @@ package cli
 import (
 	"fmt"
 
-	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/config"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
+	"github.com/gptscript-ai/gptscript/pkg/gptscript"
 	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
-	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/spf13/cobra"
 )
 
@@ -34,8 +33,7 @@ func (c *Delete) Run(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to read CLI config: %w", err)
 	}
 
-	opts.Cache = cache.Complete(opts.Cache)
-	opts.Runner = runner.Complete(opts.Runner)
+	opts = gptscript.Complete(opts)
 	if opts.Runner.RuntimeManager == nil {
 		opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
 	}
@@ -44,7 +42,7 @@ func (c *Delete) Run(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
+	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, opts.CredentialContexts, opts.Cache.CacheDir)
 	if err != nil {
 		return fmt.Errorf("failed to get credentials store: %w", err)
 	}
diff --git a/pkg/cli/credential_show.go b/pkg/cli/credential_show.go
@@ -5,11 +5,10 @@ import (
 	"os"
 	"text/tabwriter"
 
-	"github.com/gptscript-ai/gptscript/pkg/cache"
 	"github.com/gptscript-ai/gptscript/pkg/config"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
+	"github.com/gptscript-ai/gptscript/pkg/gptscript"
 	"github.com/gptscript-ai/gptscript/pkg/repos/runtimes"
-	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/spf13/cobra"
 )
 
@@ -36,8 +35,7 @@ func (c *Show) Run(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to read CLI config: %w", err)
 	}
 
-	opts.Cache = cache.Complete(opts.Cache)
-	opts.Runner = runner.Complete(opts.Runner)
+	opts = gptscript.Complete(opts)
 	if opts.Runner.RuntimeManager == nil {
 		opts.Runner.RuntimeManager = runtimes.Default(opts.Cache.CacheDir)
 	}
@@ -46,7 +44,7 @@ func (c *Show) Run(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, c.root.CredentialContext, opts.Cache.CacheDir)
+	store, err := credentials.NewStore(cfg, opts.Runner.RuntimeManager, opts.CredentialContexts, opts.Cache.CacheDir)
 	if err != nil {
 		return fmt.Errorf("failed to get credentials store: %w", err)
 	}
diff --git a/pkg/cli/gptscript.go b/pkg/cli/gptscript.go
@@ -64,7 +64,7 @@ type GPTScript struct {
 	Chdir                    string   `usage:"Change current working directory" short:"C"`
 	Daemon                   bool     `usage:"Run tool as a daemon" local:"true" hidden:"true"`
 	Ports                    string   `usage:"The port range to use for ephemeral daemon ports (ex: 11000-12000)" hidden:"true"`
-	CredentialContext        string   `usage:"Context name in which to store credentials" default:"default"`
+	CredentialContext        []string `usage:"Context name(s) in which to store credentials"`
 	CredentialOverride       []string `usage:"Credentials to override (ex: --credential-override github.com/example/cred-tool:API_TOKEN=1234)"`
 	ChatState                string   `usage:"The chat state to continue, or null to start a new chat and return the state" local:"true"`
 	ForceChat                bool     `usage:"Force an interactive chat session if even the top level tool is not a chat tool" local:"true"`
@@ -142,7 +142,7 @@ func (r *GPTScript) NewGPTScriptOpts() (gptscript.Options, error) {
 		},
 		Quiet:                r.Quiet,
 		Env:                  os.Environ(),
-		CredentialContext:    r.CredentialContext,
+		CredentialContexts:   r.CredentialContext,
 		Workspace:            r.Workspace,
 		DisablePromptServer:  r.UI,
 		DefaultModelProvider: r.DefaultModelProvider,
diff --git a/pkg/credentials/store.go b/pkg/credentials/store.go
diff --git a/pkg/credentials/util.go b/pkg/credentials/util.go
diff --git a/pkg/gptscript/gptscript.go b/pkg/gptscript/gptscript.go
diff --git a/pkg/sdkserver/credentials.go b/pkg/sdkserver/credentials.go
diff --git a/pkg/sdkserver/routes.go b/pkg/sdkserver/routes.go
diff --git a/pkg/sdkserver/types.go b/pkg/sdkserver/types.go
