feat: add credential management (#65)

Signed-off-by: Grant Linville <grant@acorn.io>

Author: g-linville

credentials.go

@@ -0,0 +1,27 @@
+package gptscript
+
+import "time"
+
+type CredentialType string
+
+const (
+	CredentialTypeTool          CredentialType = "tool"
+	CredentialTypeModelProvider CredentialType = "modelProvider"
+)
+
+type Credential struct {
+	Context      string            `json:"context"`
+	ToolName     string            `json:"toolName"`
+	Type         CredentialType    `json:"type"`
+	Env          map[string]string `json:"env"`
+	Ephemeral    bool              `json:"ephemeral,omitempty"`
+	ExpiresAt    *time.Time        `json:"expiresAt"`
+	RefreshToken string            `json:"refreshToken"`
+}
+
+type CredentialRequest struct {
+	Content     string   `json:"content"`
+	AllContexts bool     `json:"allContexts"`
+	Context     []string `json:"context"`
+	Name        string   `json:"name"`
+}

go.mod

@@ -2,15 +2,20 @@ module github.com/gptscript-ai/go-gptscript
 
 go 1.23.0
 
-require github.com/getkin/kin-openapi v0.124.0
+require (
+	github.com/getkin/kin-openapi v0.124.0
+	github.com/stretchr/testify v1.8.4
+)
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/go-openapi/jsonpointer v0.20.2 // indirect
 	github.com/go-openapi/swag v0.22.8 // indirect
 	github.com/invopop/yaml v0.2.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect
 	github.com/perimeterx/marshmallow v1.1.5 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

gptscript.go

@@ -336,6 +336,67 @@ func (g *GPTScript) PromptResponse(ctx context.Context, resp PromptResponse) err
 	return err
 }
 
+type ListCredentialsOptions struct {
+	CredentialContexts []string
+	AllContexts        bool
+}
+
+func (g *GPTScript) ListCredentials(ctx context.Context, opts ListCredentialsOptions) ([]Credential, error) {
+	req := CredentialRequest{}
+	if opts.AllContexts {
+		req.AllContexts = true
+	} else if len(opts.CredentialContexts) > 0 {
+		req.Context = opts.CredentialContexts
+	} else {
+		req.Context = []string{"default"}
+	}
+
+	out, err := g.runBasicCommand(ctx, "credentials", req)
+	if err != nil {
+		return nil, err
+	}
+
+	var creds []Credential
+	if err = json.Unmarshal([]byte(out), &creds); err != nil {
+		return nil, err
+	}
+	return creds, nil
+}
+
+func (g *GPTScript) CreateCredential(ctx context.Context, cred Credential) error {
+	credJSON, err := json.Marshal(cred)
+	if err != nil {
+		return fmt.Errorf("failed to marshal credential: %w", err)
+	}
+
+	_, err = g.runBasicCommand(ctx, "credentials/create", CredentialRequest{Content: string(credJSON)})
+	return err
+}
+
+func (g *GPTScript) RevealCredential(ctx context.Context, credCtxs []string, name string) (Credential, error) {
+	out, err := g.runBasicCommand(ctx, "credentials/reveal", CredentialRequest{
+		Context: credCtxs,
+		Name:    name,
+	})
+	if err != nil {
+		return Credential{}, err
+	}
+
+	var cred Credential
+	if err = json.Unmarshal([]byte(out), &cred); err != nil {
+		return Credential{}, err
+	}
+	return cred, nil
+}
+
+func (g *GPTScript) DeleteCredential(ctx context.Context, credCtx, name string) error {
+	_, err := g.runBasicCommand(ctx, "credentials/delete", CredentialRequest{
+		Context: []string{credCtx}, // Only one context can be specified for delete operations
+		Name:    name,
+	})
+	return err
+}
+
 func (g *GPTScript) runBasicCommand(ctx context.Context, requestPath string, body any) (string, error) {
 	run := &Run{
 		url:          g.globalOpts.URL,

gptscript_test.go

@@ -2,14 +2,18 @@ package gptscript
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"math/rand"
 	"os"
 	"path/filepath"
 	"runtime"
+	"strconv"
 	"strings"
 	"testing"
 
 	"github.com/getkin/kin-openapi/openapi3"
+	"github.com/stretchr/testify/require"
 )
 
 var g *GPTScript
@@ -1448,3 +1452,44 @@ func TestLoadTools(t *testing.T) {
 		t.Errorf("Unexpected name: %s", prg.Name)
 	}
 }
+
+func TestCredentials(t *testing.T) {
+	// We will test in the following order of create, list, reveal, delete.
+	name := "test-" + strconv.Itoa(rand.Int())
+	if len(name) > 20 {
+		name = name[:20]
+	}
+
+	// Create
+	err := g.CreateCredential(context.Background(), Credential{
+		Context:      "testing",
+		ToolName:     name,
+		Type:         CredentialTypeTool,
+		Env:          map[string]string{"ENV": "testing"},
+		RefreshToken: "my-refresh-token",
+	})
+	require.NoError(t, err)
+
+	// List
+	creds, err := g.ListCredentials(context.Background(), ListCredentialsOptions{
+		CredentialContexts: []string{"testing"},
+	})
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, len(creds), 1)
+
+	// Reveal
+	cred, err := g.RevealCredential(context.Background(), []string{"testing"}, name)
+	require.NoError(t, err)
+	require.Contains(t, cred.Env, "ENV")
+	require.Equal(t, cred.Env["ENV"], "testing")
+	require.Equal(t, cred.RefreshToken, "my-refresh-token")
+
+	// Delete
+	err = g.DeleteCredential(context.Background(), "testing", name)
+	require.NoError(t, err)
+
+	// Delete again and make sure we get a NotFoundError
+	err = g.DeleteCredential(context.Background(), "testing", name)
+	require.Error(t, err)
+	require.True(t, errors.As(err, &ErrNotFound{}))
+}

opts.go

@@ -70,6 +70,7 @@ type Options struct {
 	IncludeEvents       bool     `json:"includeEvents"`
 	Prompt              bool     `json:"prompt"`
 	CredentialOverrides []string `json:"credentialOverrides"`
+	CredentialContexts  []string `json:"credentialContext"` // json tag is left singular to match SDKServer
 	Location            string   `json:"location"`
 	ForceSequential     bool     `json:"forceSequential"`
 }

run.go

@@ -17,6 +17,14 @@ import (
 
 var errAbortRun = errors.New("run aborted")
 
+type ErrNotFound struct {
+	Message string
+}
+
+func (e ErrNotFound) Error() string {
+	return e.Message
+}
+
 type Run struct {
 	url, token, requestPath, toolPath string
 	tools                             []ToolDef
@@ -36,6 +44,7 @@ type Run struct {
 	output, errput    string
 	events            chan Frame
 	lock              sync.Mutex
+	responseCode      int
 }
 
 // Text returns the text output of the gptscript. It blocks until the output is ready.
@@ -60,6 +69,11 @@ func (r *Run) State() RunState {
 // Err returns the error that caused the gptscript to fail, if any.
 func (r *Run) Err() error {
 	if r.err != nil {
+		if r.responseCode == http.StatusNotFound {
+			return ErrNotFound{
+				Message: fmt.Sprintf("run encountered an error: %s", r.errput),
+			}
+		}
 		return fmt.Errorf("run encountered an error: %w with error output: %s", r.err, r.errput)
 	}
 	return nil
@@ -245,6 +259,7 @@ func (r *Run) request(ctx context.Context, payload any) (err error) {
 		return r.err
 	}
 
+	r.responseCode = resp.StatusCode
 	if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusBadRequest {
 		r.state = Error
 		r.err = fmt.Errorf("run encountered an error")
@@ -335,6 +350,15 @@ func (r *Run) request(ctx context.Context, payload any) (err error) {
 
 						done, _ = out["done"].(bool)
 						r.rawOutput = out
+					case []any:
+						b, err := json.Marshal(out)
+						if err != nil {
+							r.state = Error
+							r.err = fmt.Errorf("failed to process stdout: %w", err)
+							return
+						}
+
+						r.output = string(b)
 					default:
 						r.state = Error
 						r.err = fmt.Errorf("failed to process stdout, invalid type: %T", out)

run_test.go

@@ -2,8 +2,13 @@ package gptscript
 
 import (
 	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"os"
 	"runtime"
 	"testing"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestRestartingErrorRun(t *testing.T) {
@@ -42,3 +47,44 @@ func TestRestartingErrorRun(t *testing.T) {
 		t.Errorf("executing run with input of 0 should not fail: %v", err)
 	}
 }
+
+func TestStackedContexts(t *testing.T) {
+	const name = "testcred"
+
+	wd, err := os.Getwd()
+	require.NoError(t, err)
+
+	bytes := make([]byte, 32)
+	_, err = rand.Read(bytes)
+	require.NoError(t, err)
+
+	context1 := hex.EncodeToString(bytes)[:16]
+	context2 := hex.EncodeToString(bytes)[16:]
+
+	run, err := g.Run(context.Background(), wd+"/test/credential.gpt", Options{
+		CredentialContexts: []string{context1, context2},
+	})
+	require.NoError(t, err)
+
+	_, err = run.Text()
+	require.NoError(t, err)
+
+	// The credential should exist in context1 now.
+	cred, err := g.RevealCredential(context.Background(), []string{context1, context2}, name)
+	require.NoError(t, err)
+	require.Equal(t, cred.Context, context1)
+
+	// Now change the context order and run the script again.
+	run, err = g.Run(context.Background(), wd+"/test/credential.gpt", Options{
+		CredentialContexts: []string{context2, context1},
+	})
+	require.NoError(t, err)
+
+	_, err = run.Text()
+	require.NoError(t, err)
+
+	// Now make sure the credential exists in context1 still.
+	cred, err = g.RevealCredential(context.Background(), []string{context2, context1}, name)
+	require.NoError(t, err)
+	require.Equal(t, cred.Context, context1)
+}

test/credential.gpt

@@ -0,0 +1,13 @@
+name: echocred
+credential: mycredentialtool as testcred
+
+#!/usr/bin/env bash
+
+echo $VALUE
+
+---
+name: mycredentialtool
+
+#!sys.echo
+
+{"env":{"VALUE":"hello"}}

test/global-tools.gpt

@@ -4,7 +4,7 @@ Runbook 3
 
 ---
 Name: tool_1
-Global Tools: github.com/gptscript-ai/knowledge, github.com/drpebcak/duckdb, github.com/gptscript-ai/browser, github.com/gptscript-ai/browser-search/google, github.com/gptscript-ai/browser-search/google-question-answerer
+Global Tools: github.com/drpebcak/duckdb, github.com/gptscript-ai/browser, github.com/gptscript-ai/browser-search/google, github.com/gptscript-ai/browser-search/google-question-answerer
 
 Say "Hello!"