Skip to content

Commit 34e46fd

Browse files
thedadamsthedadams
committed
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]>
1 parent ed71575 commit 34e46fd

26 files changed

+276
-145
lines changed

‎.vscode/launch.json

+2-2
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,12 @@
1515
]
1616
},
1717
{
18-
"name": "Launch Server",
18+
"name": "Clicky Serves",
1919
"type": "go",
2020
"request": "launch",
2121
"mode": "debug",
2222
"program": "main.go",
23-
"args": ["--server"]
23+
"args": ["--debug", "--listen-address", "127.0.0.1:63774", "sys.sdkserver"]
2424
}
2525
]
2626
}

‎go.mod

+2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ module github.com/gptscript-ai/gptscript
22

33
go 1.23.1
44

5+
replace github.com/gptscript-ai/chat-completion-client => github.com/thedadams/chat-completion-client v0.0.0-20250221020242-9dc6a5b82cd2
6+
57
require (
68
github.com/AlecAivazis/survey/v2 v2.3.7
79
github.com/BurntSushi/locker v0.0.0-20171006230638-a6e239ea1c69

‎go.sum

+2-2
Original file line numberDiff line numberDiff line change
@@ -197,8 +197,6 @@ github.com/gorilla/css v1.0.0 h1:BQqNyPTi50JCFMTw/b67hByjMVXZRwGha6wxVGkeihY=
197197
github.com/gorilla/css v1.0.0/go.mod h1:Dn721qIggHpt4+EFCcTLTU/vk5ySda2ReITrtgBl60c=
198198
github.com/gptscript-ai/broadcaster v0.0.0-20240625175512-c43682019b86 h1:m9yLtIEd0z1ia8qFjq3u0Ozb6QKwidyL856JLJp6nbA=
199199
github.com/gptscript-ai/broadcaster v0.0.0-20240625175512-c43682019b86/go.mod h1:lK3K5EZx4dyT24UG3yCt0wmspkYqrj4D/8kxdN3relk=
200-
github.com/gptscript-ai/chat-completion-client v0.0.0-20250128181713-57857b74f9f1 h1:D8VmhL68Fm6YI7fue4wkzd1TqODn//LtcJtPvWk8BQ8=
201-
github.com/gptscript-ai/chat-completion-client v0.0.0-20250128181713-57857b74f9f1/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
202200
github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb h1:ky2J2CzBOskC7Jgm2VJAQi2x3p7FVGa+2/PcywkFJuc=
203201
github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb/go.mod h1:DJAo1xTht1LDkNYFNydVjTHd576TC7MlpsVRl3oloVw=
204202
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
376374
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
377375
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
378376
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
377+
github.com/thedadams/chat-completion-client v0.0.0-20250221020242-9dc6a5b82cd2 h1:YrShgu4eoW4UvdpQ9YMDfj/hqkvGzuR5HGNIcrnvrhc=
378+
github.com/thedadams/chat-completion-client v0.0.0-20250221020242-9dc6a5b82cd2/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
379379
github.com/therootcompany/xz v1.0.1 h1:CmOtsn1CbtmyYiusbfmhmkpAAETj0wBIH6kCYaX+xzw=
380380
github.com/therootcompany/xz v1.0.1/go.mod h1:3K3UH1yCKgBneZYhuQUvJ9HPD19UEXEI0BWbMn8qNMY=
381381
github.com/tidwall/gjson v1.17.1 h1:wlYEnwqAHgzmhNUFfw7Xalt2JzQvsMx2Se4PcoFCT/U=

‎pkg/builtin/builtin.go

+20-7
Original file line numberDiff line numberDiff line change
@@ -277,10 +277,10 @@ func ListTools() (result []types.Tool) {
277277
}
278278

279279
func Builtin(name string) (types.Tool, bool) {
280-
return BuiltinWithDefaultModel(name, "")
280+
return BuiltinDefaultModel(name, "")
281281
}
282282

283-
func BuiltinWithDefaultModel(name, defaultModel string) (types.Tool, bool) {
283+
func BuiltinDefaultModel(name, defaultModel string) (types.Tool, bool) {
284284
// Legacy syntax not used anymore
285285
name = strings.TrimSuffix(name, "?")
286286
t, ok := tools[name]
@@ -332,7 +332,7 @@ func SysFind(_ context.Context, _ []string, input string, _ chan<- string) (stri
332332
return strings.Join(result, "\n"), nil
333333
}
334334

335-
func SysExec(_ context.Context, env []string, input string, progress chan<- string) (string, error) {
335+
func SysExec(ctx context.Context, env []string, input string, progress chan<- string) (string, error) {
336336
var params struct {
337337
Command string `json:"command,omitempty"`
338338
Directory string `json:"directory,omitempty"`
@@ -345,14 +345,26 @@ func SysExec(_ context.Context, env []string, input string, progress chan<- stri
345345
params.Directory = "."
346346
}
347347

348+
commandCtx, _ := engine.FromContext(ctx)
349+
350+
ctx, cancel := context.WithCancel(ctx)
351+
defer cancel()
352+
353+
go func() {
354+
select {
355+
case <-ctx.Done():
356+
case <-commandCtx.UserCancel:
357+
cancel()
358+
}
359+
}()
360+
348361
log.Debugf("Running %s in %s", params.Command, params.Directory)
349362

350363
var cmd *exec.Cmd
351-
352364
if runtime.GOOS == "windows" {
353-
cmd = exec.Command("cmd.exe", "/c", params.Command)
365+
cmd = exec.CommandContext(ctx, "cmd.exe", "/c", params.Command)
354366
} else {
355-
cmd = exec.Command("/bin/sh", "-c", params.Command)
367+
cmd = exec.CommandContext(ctx, "/bin/sh", "-c", params.Command)
356368
}
357369

358370
var (
@@ -371,7 +383,8 @@ func SysExec(_ context.Context, env []string, input string, progress chan<- stri
371383
cmd.Dir = params.Directory
372384
cmd.Stdout = combined
373385
cmd.Stderr = combined
374-
if err := cmd.Run(); err != nil {
386+
if err := cmd.Run(); err != nil && ctx.Err() == nil {
387+
// If the command failed and the context hasn't been canceled, then return the error.
375388
return fmt.Sprintf("ERROR: %s\nOUTPUT:\n%s", err, &out), nil
376389
}
377390
return out.String(), nil

‎pkg/cache/cache.go

+7
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,13 @@ func (c *Client) Store(ctx context.Context, key, value any) error {
105105
return nil
106106
}
107107

108+
select {
109+
// If the context has been canceled, then don't try to save.
110+
case <-ctx.Done():
111+
return nil
112+
default:
113+
}
114+
108115
if c.noop || IsNoCache(ctx) {
109116
keyValue, err := c.cacheKey(key)
110117
if err == nil {

‎pkg/chat/chat.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ type Prompter interface {
1717
}
1818

1919
type Chatter interface {
20-
Chat(ctx context.Context, prevState runner.ChatState, prg types.Program, env []string, input string) (resp runner.ChatResponse, err error)
20+
Chat(ctx context.Context, prevState runner.ChatState, prg types.Program, env []string, input string, opts runner.RunOptions) (resp runner.ChatResponse, err error)
2121
}
2222

2323
type GetProgram func() (types.Program, error)
@@ -74,7 +74,7 @@ func Start(ctx context.Context, prevState runner.ChatState, chatter Chatter, prg
7474
}
7575
}
7676

77-
resp, err = chatter.Chat(ctx, prevState, prog, env, input)
77+
resp, err = chatter.Chat(ctx, prevState, prog, env, input, runner.RunOptions{})
7878
if err != nil {
7979
return err
8080
}

‎pkg/cli/eval.go

+6-5
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"github.com/gptscript-ai/gptscript/pkg/gptscript"
1111
"github.com/gptscript-ai/gptscript/pkg/input"
1212
"github.com/gptscript-ai/gptscript/pkg/loader"
13+
"github.com/gptscript-ai/gptscript/pkg/runner"
1314
"github.com/gptscript-ai/gptscript/pkg/types"
1415
"github.com/spf13/cobra"
1516
)
@@ -56,13 +57,13 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
5657
return err
5758
}
5859

59-
runner, err := gptscript.New(cmd.Context(), opts)
60+
g, err := gptscript.New(cmd.Context(), opts)
6061
if err != nil {
6162
return err
6263
}
6364

6465
prg, err := loader.ProgramFromSource(cmd.Context(), tool.String(), "", loader.Options{
65-
Cache: runner.Cache,
66+
Cache: g.Cache,
6667
})
6768
if err != nil {
6869
return err
@@ -74,14 +75,14 @@ func (e *Eval) Run(cmd *cobra.Command, args []string) error {
7475
}
7576

7677
if e.Chat {
77-
return chat.Start(cmd.Context(), nil, runner, func() (types.Program, error) {
78+
return chat.Start(cmd.Context(), nil, g, func() (types.Program, error) {
7879
return loader.ProgramFromSource(cmd.Context(), tool.String(), "", loader.Options{
79-
Cache: runner.Cache,
80+
Cache: g.Cache,
8081
})
8182
}, os.Environ(), toolInput, "")
8283
}
8384

84-
toolOutput, err := runner.Run(cmd.Context(), prg, opts.Env, toolInput)
85+
toolOutput, err := g.Run(cmd.Context(), prg, opts.Env, toolInput, runner.RunOptions{})
8586
if err != nil {
8687
return err
8788
}

‎pkg/cli/gptscript.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -469,7 +469,7 @@ func (r *GPTScript) Run(cmd *cobra.Command, args []string) (retErr error) {
469469

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

514-
s, err := gptScript.Run(cmd.Context(), prg, gptOpt.Env, toolInput)
514+
s, err := gptScript.Run(cmd.Context(), prg, gptOpt.Env, toolInput, runner.RunOptions{})
515515
if err != nil {
516516
return err
517517
}

‎pkg/engine/cmd.go

+23-5
Original file line numberDiff line numberDiff line change
@@ -119,10 +119,14 @@ func (e *Engine) runCommand(ctx Context, tool types.Tool, input string, toolCate
119119
instructions = append(instructions, inputContext.Content)
120120
}
121121

122-
var extraEnv = []string{
122+
extraEnv := []string{
123123
strings.TrimSpace("GPTSCRIPT_CONTEXT=" + strings.Join(instructions, "\n")),
124124
}
125-
cmd, stop, err := e.newCommand(ctx.Ctx, extraEnv, tool, input, true)
125+
126+
commandCtx, cancel := context.WithCancel(ctx.Ctx)
127+
defer cancel()
128+
129+
cmd, stop, err := e.newCommand(commandCtx, extraEnv, tool, input, true)
126130
if err != nil {
127131
if toolCategory == NoCategory && ctx.Parent != nil {
128132
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
155159
cmd.Stdout = io.MultiWriter(stdout, stdoutAndErr, progressOut)
156160
cmd.Stderr = io.MultiWriter(stdoutAndErr, progressOut, os.Stderr)
157161
result = stdout
162+
defer func() {
163+
combinedOutput = stdoutAndErr.String()
164+
}()
158165

159-
if err := cmd.Run(); err != nil {
166+
go func() {
167+
select {
168+
case <-commandCtx.Done():
169+
// If the commandCtx was canceled, then nothing to do.
170+
return
171+
case <-ctx.UserCancel:
172+
// If the user has canceled the run, then stop the command by canceling the commandCtx.
173+
stop()
174+
cancel()
175+
}
176+
}()
177+
178+
if err := cmd.Run(); err != nil && commandCtx.Err() == nil {
179+
// If the command failed and the context hasn't been canceled, then return the error.
160180
if toolCategory == NoCategory && ctx.Parent != nil {
161181
// If this is a sub-call, then don't return the error; return the error as a message so that the LLM can retry.
162182
return fmt.Sprintf("ERROR: got (%v) while running tool, OUTPUT: %s", err, stdoutAndErr), nil
163183
}
164184
log.Errorf("failed to run tool [%s] cmd %v: %v", tool.Parameters.Name, cmd.Args, err)
165-
combinedOutput = stdoutAndErr.String()
166185
return "", fmt.Errorf("ERROR: %s: %w", stdoutAndErr, err)
167186
}
168187

169-
combinedOutput = stdoutAndErr.String()
170188
return result.String(), IsChatFinishMessage(result.String())
171189
}
172190

‎pkg/engine/daemon.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ func (e *Engine) startDaemon(tool types.Tool) (string, error) {
229229
return url, fmt.Errorf("timeout waiting for 200 response from GET %s", url)
230230
}
231231

232-
func (e *Engine) runDaemon(ctx context.Context, prg *types.Program, tool types.Tool, input string) (cmdRet *Return, cmdErr error) {
232+
func (e *Engine) runDaemon(ctx Context, tool types.Tool, input string) (cmdRet *Return, cmdErr error) {
233233
url, err := e.startDaemon(tool)
234234
if err != nil {
235235
return nil, err
@@ -238,5 +238,5 @@ func (e *Engine) runDaemon(ctx context.Context, prg *types.Program, tool types.T
238238
tool.Instructions = strings.Join(append([]string{
239239
types.CommandPrefix + url,
240240
}, strings.Split(tool.Instructions, "\n")[1:]...), "\n")
241-
return e.runHTTP(ctx, prg, tool, input)
241+
return e.runHTTP(ctx, tool, input)
242242
}

0 commit comments

Comments
 (0)