@@ -19,6 +19,13 @@ import (
1919 "github.com/pkg/browser"
2020)
2121
22+ type authType string
23+
24+ const (
25+ authTypePAT authType = "Personal Access Token (PAT)"
26+ authTypeOAuth authType = "OAuth"
27+ )
28+
2229type input struct {
2330 OAuthInfo oauthInfo `json:"oauthInfo"`
2431 PromptInfo * promptInfo `json:"promptInfo,omitempty"`
@@ -68,18 +75,28 @@ type cred struct {
6875 RefreshToken string `json:"refreshToken"`
6976}
7077
78+ type urls struct {
79+ authorizeURL string
80+ refreshURL string
81+ tokenURL string
82+ }
83+
7184func normalizeForEnv (appName string ) string {
7285 return strings .ToUpper (strings .ReplaceAll (appName , "-" , "_" ))
7386}
7487
75- func getURLs (appName string ) ( string , string , string ) {
88+ func getURLs (appName string ) urls {
7689 var (
7790 normalizedAppName = normalizeForEnv (appName )
7891 authorizeURL = os .Getenv (fmt .Sprintf ("GPTSCRIPT_OAUTH_%s_AUTH_URL" , normalizedAppName ))
7992 refreshURL = os .Getenv (fmt .Sprintf ("GPTSCRIPT_OAUTH_%s_REFRESH_URL" , normalizedAppName ))
8093 tokenURL = os .Getenv (fmt .Sprintf ("GPTSCRIPT_OAUTH_%s_TOKEN_URL" , normalizedAppName ))
8194 )
82- return authorizeURL , refreshURL , tokenURL
95+ return urls {
96+ authorizeURL : authorizeURL ,
97+ refreshURL : refreshURL ,
98+ tokenURL : tokenURL ,
99+ }
83100}
84101
85102func main () {
@@ -118,22 +135,35 @@ func mainErr() (err error) {
118135 }
119136 }()
120137
121- authorizeURL , refreshURL , tokenURL := getURLs (in .OAuthInfo .Integration )
122- if authorizeURL == "" || refreshURL == "" || tokenURL == "" {
123- // The URLs aren't set for this credential. Check to see if we should prompt the user for other tokens
124- if in .PromptInfo == nil {
125- fmt .Printf ("All the following environment variables must be set: GPTSCRIPT_OAUTH_%s_AUTH_URL, GPTSCRIPT_OAUTH_%[1]s_REFRESH_URL, GPTSCRIPT_OAUTH_%[1]s_TOKEN_URL" , normalizeForEnv (in .OAuthInfo .Integration ))
126- fmt .Printf ("Or the promptInfo configuration must be provided for token prompting." )
127- os .Exit (1 )
138+ urls := getURLs (in .OAuthInfo .Integration )
139+
140+ if in .PromptInfo != nil && os .Getenv ("GPTSCRIPT_EXISTING_CREDENTIAL" ) == "" {
141+ if urls .authorizeURL == "" && urls .refreshURL == "" && urls .tokenURL == "" {
142+ credJSON , err = promptForTokens (ctx , gs , in .OAuthInfo .Integration , in .PromptInfo )
143+ if err != nil {
144+ return fmt .Errorf ("main: failed to prompt for tokens: %w" , err )
145+ }
146+ return nil
128147 }
129148
130- credJSON , err = promptForTokens (ctx , gs , in . OAuthInfo . Integration , in . PromptInfo )
149+ authType , err := promptForSelect (ctx , gs )
131150 if err != nil {
132- return fmt .Errorf ("main: failed to prompt for tokens: %w" , err )
151+ return fmt .Errorf ("main: failed to prompt for auth type: %w" , err )
152+ }
153+
154+ if authType == authTypePAT {
155+ credJSON , err = promptForTokens (ctx , gs , in .OAuthInfo .Integration , in .PromptInfo )
156+ if err != nil {
157+ return fmt .Errorf ("main: failed to prompt for tokens: %w" , err )
158+ }
159+ return nil
133160 }
134- return nil
135161 }
136162
163+ return promptForOauth (gs , & urls , & in , & credJSON )
164+ }
165+
166+ func promptForOauth (gs * gptscript.GPTScript , urls * urls , in * input , credJSON * []byte ) error {
137167 // Refresh existing credential if there is one.
138168 existing := os .Getenv ("GPTSCRIPT_EXISTING_CREDENTIAL" )
139169 if existing != "" {
@@ -142,7 +172,7 @@ func mainErr() (err error) {
142172 return fmt .Errorf ("main: failed to unmarshal existing credential: %w" , err )
143173 }
144174
145- u , err := url .Parse (refreshURL )
175+ u , err := url .Parse (urls . refreshURL )
146176 if err != nil {
147177 return fmt .Errorf ("main: failed to parse refresh URL: %w" , err )
148178 }
@@ -195,7 +225,7 @@ func mainErr() (err error) {
195225 out .ExpiresAt = & expiresAt
196226 }
197227
198- credJSON , err = json .Marshal (out )
228+ * credJSON , err = json .Marshal (out )
199229 if err != nil {
200230 return fmt .Errorf ("main: failed to marshal refreshed credential: %w" , err )
201231 }
@@ -217,7 +247,7 @@ func mainErr() (err error) {
217247 h .Write ([]byte (verifier ))
218248 challenge := hex .EncodeToString (h .Sum (nil ))
219249
220- u , err := url .Parse (authorizeURL )
250+ u , err := url .Parse (urls . authorizeURL )
221251 if err != nil {
222252 return fmt .Errorf ("main: failed to parse authorize URL: %w" , err )
223253 }
@@ -274,7 +304,7 @@ func mainErr() (err error) {
274304 t := time .NewTicker (2 * time .Second )
275305 for range t .C {
276306 now := time .Now ()
277- oauthResp , retry , err := makeTokenRequest (tokenURL , state , verifier )
307+ oauthResp , retry , err := makeTokenRequest (urls . tokenURL , state , verifier )
278308 if err != nil {
279309 if ! retry {
280310 return err
@@ -301,7 +331,7 @@ func mainErr() (err error) {
301331 out .ExpiresAt = & expiresAt
302332 }
303333
304- credJSON , err = json .Marshal (out )
334+ * credJSON , err = json .Marshal (out )
305335 if err != nil {
306336 return fmt .Errorf ("main: failed to marshal token credential: %w" , err )
307337 }
@@ -357,6 +387,40 @@ func generateString() (string, error) {
357387 return string (b ), nil
358388}
359389
390+ func promptForSelect (ctx context.Context , g * gptscript.GPTScript ) (authType , error ) {
391+ fieldName := "Authentication Method"
392+
393+ fields := gptscript.Fields {gptscript.Field {Name : fieldName , Description : "The authentication method to use for this tool." , Options : []string {string (authTypePAT ), string (authTypeOAuth )}}}
394+ sysPromptIn , err := json .Marshal (sysPromptInput {
395+ Message : "This tool has personal access token (PAT) and OAuth support. Select the authentication method you would like to use for this tool." ,
396+ Fields : fields ,
397+ })
398+ if err != nil {
399+ return "" , fmt .Errorf ("promptForSelect: error marshalling sys prompt input: %w" , err )
400+ }
401+
402+ run , err := g .Run (ctx , "sys.prompt" , gptscript.Options {
403+ Input : string (sysPromptIn ),
404+ })
405+ if err != nil {
406+ return "" , fmt .Errorf ("promptForSelect: failed to run sys.prompt: %w" , err )
407+ }
408+
409+ out , err := run .Text ()
410+ if err != nil {
411+ return "" , fmt .Errorf ("promptForSelect: failed to get prompt response: %w" , err )
412+ }
413+
414+ m := make (map [string ]string )
415+ if err = json .Unmarshal ([]byte (out ), & m ); err != nil {
416+ return "" , fmt .Errorf ("promptForSelect: failed to unmarshal prompt response: %w" , err )
417+ }
418+
419+ authType := authType (m [fieldName ])
420+
421+ return authType , nil
422+ }
423+
360424func promptForTokens (ctx context.Context , g * gptscript.GPTScript , integration string , prompt * promptInfo ) ([]byte , error ) {
361425 if prompt .Metadata == nil {
362426 prompt .Metadata = make (map [string ]string )
0 commit comments