Merge pull request #446 from thedadams/ui-prompt-env-fixes-2

fix: clear prompt env vars when running the UI

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
 		}