fix: clear prompt env vars when running the UI

When running the UI all prompts except the first prompt for an OpenAI
key should go through the UI. Therefore, the prompt environment
variables shouldn't go to the UI tool.

Signed-off-by: Donnie Adams <[email protected]>

Author: thedadams

pkg/cli/gptscript.go

@@ -467,6 +467,12 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
 		}, gptOpt.Env, toolInput, r.SaveChatStateFile)
 	}
 
+	if r.UI {
+		// If the UI is running, then all prompts should go through the SDK and the UI.
+		// Not clearing ExtraEnv here would mean that the prompts would go through the terminal.
+		gptScript.ExtraEnv = nil
+	}
+
 	s, err := gptScript.Run(cmd.Context(), prg, gptOpt.Env, toolInput)
 	if err != nil {
 		return err

pkg/gptscript/gptscript.go

@@ -110,10 +110,9 @@ func New(opts *Options) (*GPTScript, error) {
 		closeServer()
 		return nil, err
 	}
-	opts.Env = append(opts.Env, extraEnv...)
-	oaiClient.SetEnvs(opts.Env)
+	oaiClient.SetEnvs(extraEnv)
 
-	remoteClient := remote.New(runner, opts.Env, cacheClient, cliCfg, opts.CredentialContext)
+	remoteClient := remote.New(runner, extraEnv, cacheClient, cliCfg, opts.CredentialContext)
 	if err := registry.AddClient(remoteClient); err != nil {
 		closeServer()
 		return nil, err

pkg/prompt/server.go

@@ -14,13 +14,20 @@ import (
 )
 
 func NewServer(ctx context.Context, envs []string) ([]string, error) {
+	var extraEnvs []string
 	for _, env := range envs {
-		_, v, ok := strings.Cut(env, types.PromptTokenEnvVar+"=")
-		if ok && v != "" {
-			return nil, nil
+		for _, prefix := range []string{types.PromptURLEnvVar, types.PromptTokenEnvVar} {
+			v, ok := strings.CutPrefix(env, prefix+"=")
+			if ok && v != "" {
+				extraEnvs = append(extraEnvs, env)
+			}
 		}
 	}
 
+	if len(extraEnvs) == 2 {
+		return extraEnvs, nil
+	}
+
 	l, err := net.Listen("tcp", "127.0.0.1:0")
 	if err != nil {
 		return nil, err

pkg/sdkserver/routes.go

@@ -166,7 +166,7 @@ func (s *server) execHandler(w http.ResponseWriter, r *http.Request) {
 	// Don't overwrite the PromptURLEnvVar if it is already set in the environment.
 	var promptTokenAlreadySet bool
 	for _, env := range reqObject.Env {
-		if strings.HasPrefix(env, types.PromptTokenEnvVar+"=") {
+		if v, ok := strings.CutPrefix(env, types.PromptTokenEnvVar+"="); ok && v != "" {
 			promptTokenAlreadySet = true
 			break
 		}