feat: add ability to gracefully stop runs (#950)

Various parts of a tool execution cannot be stopped gracefully. For example,
non-streamed HTTP request can't be stopped gracefully. However, commands and
chat completions can be gracefully stopped by a user and the result returned.
An "ABORTED BY USER" message is added to such messages. Additionally, aborted
chat completion responses are not stored in the cache.

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

Author: thedadams

.vscode/launch.json

@@ -15,12 +15,12 @@
       ]
     },
     {
-      "name": "Launch Server",
+      "name": "Clicky Serves",
       "type": "go",
       "request": "launch",
       "mode": "debug",
       "program": "main.go",
-      "args": ["--server"]
+      "args": ["--debug", "--listen-address", "127.0.0.1:63774", "sys.sdkserver"]
     }
   ]
 }

go.mod

@@ -15,7 +15,7 @@ require (
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
 	github.com/google/uuid v1.6.0
 	github.com/gptscript-ai/broadcaster v0.0.0-20240625175512-c43682019b86
-	github.com/gptscript-ai/chat-completion-client v0.0.0-20250128181713-57857b74f9f1
+	github.com/gptscript-ai/chat-completion-client v0.0.0-20250224164718-139cb4507b1d
 	github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb
 	github.com/gptscript-ai/go-gptscript v0.9.6-0.20250204133419-744b25b84a61
 	github.com/gptscript-ai/tui v0.0.0-20250204145344-33cd15de4cee

go.sum

@@ -197,8 +197,8 @@ github.com/gorilla/css v1.0.0 h1:BQqNyPTi50JCFMTw/b67hByjMVXZRwGha6wxVGkeihY=
 github.com/gorilla/css v1.0.0/go.mod h1:Dn721qIggHpt4+EFCcTLTU/vk5ySda2ReITrtgBl60c=
 github.com/gptscript-ai/broadcaster v0.0.0-20240625175512-c43682019b86 h1:m9yLtIEd0z1ia8qFjq3u0Ozb6QKwidyL856JLJp6nbA=
 github.com/gptscript-ai/broadcaster v0.0.0-20240625175512-c43682019b86/go.mod h1:lK3K5EZx4dyT24UG3yCt0wmspkYqrj4D/8kxdN3relk=
-github.com/gptscript-ai/chat-completion-client v0.0.0-20250128181713-57857b74f9f1 h1:D8VmhL68Fm6YI7fue4wkzd1TqODn//LtcJtPvWk8BQ8=
-github.com/gptscript-ai/chat-completion-client v0.0.0-20250128181713-57857b74f9f1/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
+github.com/gptscript-ai/chat-completion-client v0.0.0-20250224164718-139cb4507b1d h1:p5uqZufDIMQzAALblZFkr8fwbnZbFXbBCR1ZMAFylXk=
+github.com/gptscript-ai/chat-completion-client v0.0.0-20250224164718-139cb4507b1d/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
 github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb h1:ky2J2CzBOskC7Jgm2VJAQi2x3p7FVGa+2/PcywkFJuc=
 github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb/go.mod h1:DJAo1xTht1LDkNYFNydVjTHd576TC7MlpsVRl3oloVw=
 github.com/gptscript-ai/go-gptscript v0.9.6-0.20250204133419-744b25b84a61 h1:QxLjsLOYlsVLPwuRkP0Q8EcAoZT1s8vU2ZBSX0+R6CI=

pkg/builtin/builtin.go

@@ -269,18 +269,14 @@ func ListTools() (result []types.Tool) {
 
 	sort.Strings(keys)
 	for _, key := range keys {
-		t, _ := Builtin(key)
+		t, _ := DefaultModel(key, "")
 		result = append(result, t)
 	}
 
 	return
 }
 
-func Builtin(name string) (types.Tool, bool) {
-	return BuiltinWithDefaultModel(name, "")
-}
-
-func BuiltinWithDefaultModel(name, defaultModel string) (types.Tool, bool) {
+func DefaultModel(name, defaultModel string) (types.Tool, bool) {
 	// Legacy syntax not used anymore
 	name = strings.TrimSuffix(name, "?")
 	t, ok := tools[name]
@@ -332,7 +328,7 @@ func SysFind(_ context.Context, _ []string, input string, _ chan<- string) (stri
 	return strings.Join(result, "\n"), nil
 }
 
-func SysExec(_ context.Context, env []string, input string, progress chan<- string) (string, error) {
+func SysExec(ctx context.Context, env []string, input string, progress chan<- string) (string, error) {
 	var params struct {
 		Command   string `json:"command,omitempty"`
 		Directory string `json:"directory,omitempty"`
@@ -345,14 +341,20 @@ func SysExec(_ context.Context, env []string, input string, progress chan<- stri
 		params.Directory = "."
 	}
 
+	commandCtx, _ := engine.FromContext(ctx)
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	commandCtx.OnUserCancel(ctx, cancel)
+
 	log.Debugf("Running %s in %s", params.Command, params.Directory)
 
 	var cmd *exec.Cmd
-
 	if runtime.GOOS == "windows" {
-		cmd = exec.Command("cmd.exe", "/c", params.Command)
+		cmd = exec.CommandContext(ctx, "cmd.exe", "/c", params.Command)
 	} else {
-		cmd = exec.Command("/bin/sh", "-c", params.Command)
+		cmd = exec.CommandContext(ctx, "/bin/sh", "-c", params.Command)
 	}
 
 	var (
@@ -371,7 +373,8 @@ func SysExec(_ context.Context, env []string, input string, progress chan<- stri
 	cmd.Dir = params.Directory
 	cmd.Stdout = combined
 	cmd.Stderr = combined
-	if err := cmd.Run(); err != nil {
+	if err := cmd.Run(); err != nil && (ctx.Err() == nil || commandCtx.Ctx.Err() != nil) {
+		// If the command failed and the context hasn't been canceled, then return the error.
 		return fmt.Sprintf("ERROR: %s\nOUTPUT:\n%s", err, &out), nil
 	}
 	return out.String(), nil
@@ -420,7 +423,6 @@ func getWorkspaceEnvFileContents(envs []string) ([]string, error) {
 	}
 
 	return envContents, nil
-
 }
 
 func getWorkspaceDir(envs []string) (string, error) {
@@ -665,6 +667,7 @@ func DiscardProgress() (progress chan<- string, closeFunc func()) {
 	ch := make(chan string)
 	go func() {
 		for range ch {
+			continue
 		}
 	}()
 	return ch, func() {

pkg/cache/cache.go

@@ -105,6 +105,13 @@ func (c *Client) Store(ctx context.Context, key, value any) error {
 		return nil
 	}
 
+	select {
+	// If the context has been canceled, then don't try to save.
+	case <-ctx.Done():
+		return nil
+	default:
+	}
+
 	if c.noop || IsNoCache(ctx) {
 		keyValue, err := c.cacheKey(key)
 		if err == nil {

pkg/chat/chat.go

@@ -17,7 +17,7 @@ type Prompter interface {
 }
 
 type Chatter interface {
-	Chat(ctx context.Context, prevState runner.ChatState, prg types.Program, env []string, input string) (resp runner.ChatResponse, err error)
+	Chat(ctx context.Context, prevState runner.ChatState, prg types.Program, env []string, input string, opts runner.RunOptions) (resp runner.ChatResponse, err error)
 }
 
 type GetProgram func() (types.Program, error)
@@ -74,7 +74,7 @@ func Start(ctx context.Context, prevState runner.ChatState, chatter Chatter, prg
 			}
 		}
 
-		resp, err = chatter.Chat(ctx, prevState, prog, env, input)
+		resp, err = chatter.Chat(ctx, prevState, prog, env, input, runner.RunOptions{})
 		if err != nil {
 			return err
 		}

pkg/cli/eval.go

@@ -10,6 +10,7 @@ import (
 	"github.com/gptscript-ai/gptscript/pkg/gptscript"
 	"github.com/gptscript-ai/gptscript/pkg/input"
 	"github.com/gptscript-ai/gptscript/pkg/loader"
+	"github.com/gptscript-ai/gptscript/pkg/runner"
 	"github.com/gptscript-ai/gptscript/pkg/types"
 	"github.com/spf13/cobra"
 )
@@ -56,13 +57,13 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	runner, err := gptscript.New(cmd.Context(), opts)
+	g, err := gptscript.New(cmd.Context(), opts)
 	if err != nil {
 		return err
 	}
 
 	prg, err := loader.ProgramFromSource(cmd.Context(), tool.String(), "", loader.Options{
-		Cache: runner.Cache,
+		Cache: g.Cache,
 	})
 	if err != nil {
 		return err
@@ -74,14 +75,14 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
 	}
 
 	if e.Chat {
-		return chat.Start(cmd.Context(), nil, runner, func() (types.Program, error) {
+		return chat.Start(cmd.Context(), nil, g, func() (types.Program, error) {
 			return loader.ProgramFromSource(cmd.Context(), tool.String(), "", loader.Options{
-				Cache: runner.Cache,
+				Cache: g.Cache,
 			})
 		}, os.Environ(), toolInput, "")
 	}
 
-	toolOutput, err := runner.Run(cmd.Context(), prg, opts.Env, toolInput)
+	toolOutput, err := g.Run(cmd.Context(), prg, opts.Env, toolInput, runner.RunOptions{})
 	if err != nil {
 		return err
 	}

pkg/cli/gptscript.go

@@ -469,7 +469,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
 
 	// This chat in a stateless mode
 	if r.SaveChatStateFile == "-" || r.SaveChatStateFile == "stdout" {
-		resp, err := gptScript.Chat(cmd.Context(), chatState, prg, gptOpt.Env, toolInput)
+		resp, err := gptScript.Chat(cmd.Context(), chatState, prg, gptOpt.Env, toolInput, runner.RunOptions{})
 		if err != nil {
 			return err
 		}
@@ -511,7 +511,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
 		gptScript.ExtraEnv = nil
 	}
 
-	s, err := gptScript.Run(cmd.Context(), prg, gptOpt.Env, toolInput)
+	s, err := gptScript.Run(cmd.Context(), prg, gptOpt.Env, toolInput, runner.RunOptions{})
 	if err != nil {
 		return err
 	}

pkg/engine/cmd.go

@@ -119,10 +119,14 @@ func (e *Engine) runCommand(ctx Context, tool types.Tool, input string, toolCate
 		instructions = append(instructions, inputContext.Content)
 	}
 
-	var extraEnv = []string{
+	extraEnv := []string{
 		strings.TrimSpace("GPTSCRIPT_CONTEXT=" + strings.Join(instructions, "\n")),
 	}
-	cmd, stop, err := e.newCommand(ctx.Ctx, extraEnv, tool, input, true)
+
+	commandCtx, cancel := context.WithCancel(ctx.Ctx)
+	defer cancel()
+
+	cmd, stop, err := e.newCommand(commandCtx, extraEnv, tool, input, true)
 	if err != nil {
 		if toolCategory == NoCategory && ctx.Parent != nil {
 			return fmt.Sprintf("ERROR: got (%v) while parsing command", err), nil
@@ -155,18 +159,22 @@ func (e *Engine) runCommand(ctx Context, tool types.Tool, input string, toolCate
 	cmd.Stdout = io.MultiWriter(stdout, stdoutAndErr, progressOut)
 	cmd.Stderr = io.MultiWriter(stdoutAndErr, progressOut, os.Stderr)
 	result = stdout
+	defer func() {
+		combinedOutput = stdoutAndErr.String()
+	}()
+
+	ctx.OnUserCancel(commandCtx, cancel)
 
-	if err := cmd.Run(); err != nil {
+	if err := cmd.Run(); err != nil && (commandCtx.Err() == nil || ctx.Ctx.Err() != nil) {
+		// If the command failed and the context hasn't been canceled, then return the error.
 		if toolCategory == NoCategory && ctx.Parent != nil {
 			// If this is a sub-call, then don't return the error; return the error as a message so that the LLM can retry.
 			return fmt.Sprintf("ERROR: got (%v) while running tool, OUTPUT: %s", err, stdoutAndErr), nil
 		}
 		log.Errorf("failed to run tool [%s] cmd %v: %v", tool.Parameters.Name, cmd.Args, err)
-		combinedOutput = stdoutAndErr.String()
 		return "", fmt.Errorf("ERROR: %s: %w", stdoutAndErr, err)
 	}
 
-	combinedOutput = stdoutAndErr.String()
 	return result.String(), IsChatFinishMessage(result.String())
 }
 

pkg/engine/daemon.go

@@ -229,7 +229,7 @@ func (e *Engine) startDaemon(tool types.Tool) (string, error) {
 	return url, fmt.Errorf("timeout waiting for 200 response from GET %s", url)
 }
 
-func (e *Engine) runDaemon(ctx context.Context, prg *types.Program, tool types.Tool, input string) (cmdRet *Return, cmdErr error) {
+func (e *Engine) runDaemon(ctx Context, tool types.Tool, input string) (cmdRet *Return, cmdErr error) {
 	url, err := e.startDaemon(tool)
 	if err != nil {
 		return nil, err
@@ -238,5 +238,5 @@ func (e *Engine) runDaemon(ctx context.Context, prg *types.Program, tool types.T
 	tool.Instructions = strings.Join(append([]string{
 		types.CommandPrefix + url,
 	}, strings.Split(tool.Instructions, "\n")[1:]...), "\n")
-	return e.runHTTP(ctx, prg, tool, input)
+	return e.runHTTP(ctx, tool, input)
 }