fix: credentials: block output while running credential tools (#261)

Signed-off-by: Grant Linville <[email protected]>

Author: g-linville

pkg/builtin/builtin.go

@@ -21,7 +21,6 @@ import (
 	"github.com/BurntSushi/locker"
 	"github.com/google/shlex"
 	"github.com/gptscript-ai/gptscript/pkg/confirm"
-	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 	"github.com/jaytaylor/html2text"
 )
@@ -647,15 +646,7 @@ func SysDownload(ctx context.Context, env []string, input string) (_ string, err
 	return params.Location, nil
 }
 
-func SysPrompt(ctx context.Context, _ []string, input string) (_ string, err error) {
-	monitor := ctx.Value(runner.MonitorKey{})
-	if monitor == nil {
-		return "", errors.New("no monitor in context")
-	}
-
-	unpause := monitor.(runner.Monitor).Pause()
-	defer unpause()
-
+func SysPrompt(_ context.Context, _ []string, input string) (_ string, err error) {
 	var params struct {
 		Message   string `json:"message,omitempty"`
 		Fields    string `json:"fields,omitempty"`

pkg/context/context.go

@@ -0,0 +1,19 @@
+package context
+
+import (
+	"context"
+)
+
+type pauseKey struct{}
+
+func AddPauseFuncToCtx(ctx context.Context, pauseF func() func()) context.Context {
+	return context.WithValue(ctx, pauseKey{}, pauseF)
+}
+
+func GetPauseFuncFromCtx(ctx context.Context) func() func() {
+	pauseF, ok := ctx.Value(pauseKey{}).(func() func())
+	if !ok {
+		return func() func() { return func() {} }
+	}
+	return pauseF
+}

pkg/engine/cmd.go

@@ -15,12 +15,13 @@ import (
 	"sync/atomic"
 
 	"github.com/google/shlex"
+	context2 "github.com/gptscript-ai/gptscript/pkg/context"
 	"github.com/gptscript-ai/gptscript/pkg/env"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 	"github.com/gptscript-ai/gptscript/pkg/version"
 )
 
-func (e *Engine) runCommand(ctx context.Context, tool types.Tool, input string) (cmdOut string, cmdErr error) {
+func (e *Engine) runCommand(ctx context.Context, tool types.Tool, input string, isCredential bool) (cmdOut string, cmdErr error) {
 	id := fmt.Sprint(atomic.AddInt64(&completionID, 1))
 
 	defer func() {
@@ -64,6 +65,12 @@ func (e *Engine) runCommand(ctx context.Context, tool types.Tool, input string)
 	cmd.Stderr = io.MultiWriter(all, os.Stderr)
 	cmd.Stdout = io.MultiWriter(all, output)
 
+	if isCredential {
+		pause := context2.GetPauseFuncFromCtx(ctx)
+		unpause := pause()
+		defer unpause()
+	}
+
 	if err := cmd.Run(); err != nil {
 		_, _ = os.Stderr.Write(output.Bytes())
 		log.Errorf("failed to run tool [%s] cmd %v: %v", tool.Parameters.Name, cmd.Args, err)

pkg/engine/engine.go

@@ -171,7 +171,7 @@ func (e *Engine) Start(ctx Context, input string) (*Return, error) {
 		} else if tool.IsOpenAPI() {
 			return e.runOpenAPI(tool, input)
 		}
-		s, err := e.runCommand(ctx.WrappedContext(), tool, input)
+		s, err := e.runCommand(ctx.WrappedContext(), tool, input, ctx.IsCredential)
 		if err != nil {
 			return nil, err
 		}

pkg/runner/runner.go

@@ -11,6 +11,7 @@ import (
 	"time"
 
 	"github.com/gptscript-ai/gptscript/pkg/config"
+	context2 "github.com/gptscript-ai/gptscript/pkg/context"
 	"github.com/gptscript-ai/gptscript/pkg/credentials"
 	"github.com/gptscript-ai/gptscript/pkg/engine"
 	"github.com/gptscript-ai/gptscript/pkg/types"
@@ -27,8 +28,6 @@ type Monitor interface {
 	Stop(output string, err error)
 }
 
-type MonitorKey struct{}
-
 type Options struct {
 	MonitorFactory     MonitorFactory        `usage:"-"`
 	RuntimeManager     engine.RuntimeManager `usage:"-"`
@@ -182,10 +181,7 @@ func (r *Runner) call(callCtx engine.Context, monitor Monitor, env []string, inp
 		Content:     input,
 	})
 
-	// The sys.prompt tool is a special case where we need to pass the monitor to the builtin function.
-	if callCtx.Tool.BuiltinFunc != nil && callCtx.Tool.ID == "sys.prompt" {
-		callCtx.Ctx = context.WithValue(callCtx.Ctx, MonitorKey{}, monitor)
-	}
+	callCtx.Ctx = context2.AddPauseFuncToCtx(callCtx.Ctx, monitor.Pause)
 
 	result, err := e.Start(callCtx, input)
 	if err != nil {