feat: add ability to gracefully stop runs

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

@@ -2,6 +2,8 @@ module github.com/gptscript-ai/gptscript
 
 go 1.23.1
 
+replace github.com/gptscript-ai/chat-completion-client => github.com/thedadams/chat-completion-client v0.0.0-20250221020242-9dc6a5b82cd2
+
 require (
 	github.com/AlecAivazis/survey/v2 v2.3.7
 	github.com/BurntSushi/locker v0.0.0-20171006230638-a6e239ea1c69

go.sum

@@ -197,8 +197,6 @@ 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/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=
@@ -376,6 +374,8 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/thedadams/chat-completion-client v0.0.0-20250221020242-9dc6a5b82cd2 h1:YrShgu4eoW4UvdpQ9YMDfj/hqkvGzuR5HGNIcrnvrhc=
+github.com/thedadams/chat-completion-client v0.0.0-20250221020242-9dc6a5b82cd2/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
 github.com/therootcompany/xz v1.0.1 h1:CmOtsn1CbtmyYiusbfmhmkpAAETj0wBIH6kCYaX+xzw=
 github.com/therootcompany/xz v1.0.1/go.mod h1:3K3UH1yCKgBneZYhuQUvJ9HPD19UEXEI0BWbMn8qNMY=
 github.com/tidwall/gjson v1.17.1 h1:wlYEnwqAHgzmhNUFfw7Xalt2JzQvsMx2Se4PcoFCT/U=

pkg/builtin/builtin.go

@@ -277,10 +277,10 @@ func ListTools() (result []types.Tool) {
 }
 
 func Builtin(name string) (types.Tool, bool) {
-	return BuiltinWithDefaultModel(name, "")
+	return BuiltinDefaultModel(name, "")
 }
 
-func BuiltinWithDefaultModel(name, defaultModel string) (types.Tool, bool) {
+func BuiltinDefaultModel(name, defaultModel string) (types.Tool, bool) {
 	// Legacy syntax not used anymore
 	name = strings.TrimSuffix(name, "?")
 	t, ok := tools[name]
@@ -332,7 +332,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 +345,26 @@ func SysExec(_ context.Context, env []string, input string, progress chan<- stri
 		params.Directory = "."
 	}
 
+	commandCtx, _ := engine.FromContext(ctx)
+
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-commandCtx.UserCancel:
+			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 +383,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 {
+		// 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

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,32 @@ 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()
+	}()
 
-	if err := cmd.Run(); err != nil {
+	go func() {
+		select {
+		case <-commandCtx.Done():
+			// If the commandCtx was canceled, then nothing to do.
+			return
+		case <-ctx.UserCancel:
+			// If the user has canceled the run, then stop the command by canceling the commandCtx.
+			stop()
+			cancel()
+		}
+	}()
+
+	if err := cmd.Run(); err != nil && commandCtx.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)
 }