From 04fc99b079d2ccccc801ee54eddb30cd1d40a43b Mon Sep 17 00:00:00 2001
From: Donnie Adams <donnie@acorn.io>
Date: Wed, 5 Jun 2024 20:41:28 -0400
Subject: [PATCH] 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 <donnie@acorn.io>
---
 pkg/cli/gptscript.go       |  6 ++++++
 pkg/gptscript/gptscript.go |  5 ++---
 pkg/prompt/server.go       | 13 ++++++++++---
 pkg/sdkserver/routes.go    |  2 +-
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/pkg/cli/gptscript.go b/pkg/cli/gptscript.go
index a907cdc4..c0eb1dee 100644
--- a/pkg/cli/gptscript.go
+++ b/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
diff --git a/pkg/gptscript/gptscript.go b/pkg/gptscript/gptscript.go
index afc94e30..b314d725 100644
--- a/pkg/gptscript/gptscript.go
+++ b/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
diff --git a/pkg/prompt/server.go b/pkg/prompt/server.go
index aa56f33f..28903480 100644
--- a/pkg/prompt/server.go
+++ b/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
diff --git a/pkg/sdkserver/routes.go b/pkg/sdkserver/routes.go
index 7aabefc2..bcdfbf2b 100644
--- a/pkg/sdkserver/routes.go
+++ b/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
 		}