Skip to content

Commit 0d20be1

Browse files
thedadamsthedadams
authored
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]>
1 parent 86e85f0 commit 0d20be1

26 files changed

+268
-151
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

+1-1
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ require (
1515
github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
1616
github.com/google/uuid v1.6.0
1717
github.com/gptscript-ai/broadcaster v0.0.0-20240625175512-c43682019b86
18-
github.com/gptscript-ai/chat-completion-client v0.0.0-20250128181713-57857b74f9f1
18+
github.com/gptscript-ai/chat-completion-client v0.0.0-20250224164718-139cb4507b1d
1919
github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb
2020
github.com/gptscript-ai/go-gptscript v0.9.6-0.20250204133419-744b25b84a61
2121
github.com/gptscript-ai/tui v0.0.0-20250204145344-33cd15de4cee

‎go.sum

+2-2
Original file line numberDiff line numberDiff line change
@@ -197,8 +197,8 @@ 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=
200+
github.com/gptscript-ai/chat-completion-client v0.0.0-20250224164718-139cb4507b1d h1:p5uqZufDIMQzAALblZFkr8fwbnZbFXbBCR1ZMAFylXk=
201+
github.com/gptscript-ai/chat-completion-client v0.0.0-20250224164718-139cb4507b1d/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
202202
github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb h1:ky2J2CzBOskC7Jgm2VJAQi2x3p7FVGa+2/PcywkFJuc=
203203
github.com/gptscript-ai/cmd v0.0.0-20240802230653-326b7baf6fcb/go.mod h1:DJAo1xTht1LDkNYFNydVjTHd576TC7MlpsVRl3oloVw=
204204
github.com/gptscript-ai/go-gptscript v0.9.6-0.20250204133419-744b25b84a61 h1:QxLjsLOYlsVLPwuRkP0Q8EcAoZT1s8vU2ZBSX0+R6CI=

‎pkg/builtin/builtin.go

+15-12
Original file line numberDiff line numberDiff line change
@@ -269,18 +269,14 @@ func ListTools() (result []types.Tool) {
269269

270270
sort.Strings(keys)
271271
for _, key := range keys {
272-
t, _ := Builtin(key)
272+
t, _ := DefaultModel(key, "")
273273
result = append(result, t)
274274
}
275275

276276
return
277277
}
278278

279-
func Builtin(name string) (types.Tool, bool) {
280-
return BuiltinWithDefaultModel(name, "")
281-
}
282-
283-
func BuiltinWithDefaultModel(name, defaultModel string) (types.Tool, bool) {
279+
func DefaultModel(name, defaultModel string) (types.Tool, bool) {
284280
// Legacy syntax not used anymore
285281
name = strings.TrimSuffix(name, "?")
286282
t, ok := tools[name]
@@ -332,7 +328,7 @@ func SysFind(_ context.Context, _ []string, input string, _ chan<- string) (stri
332328
return strings.Join(result, "\n"), nil
333329
}
334330

335-
func SysExec(_ context.Context, env []string, input string, progress chan<- string) (string, error) {
331+
func SysExec(ctx context.Context, env []string, input string, progress chan<- string) (string, error) {
336332
var params struct {
337333
Command string `json:"command,omitempty"`
338334
Directory string `json:"directory,omitempty"`
@@ -345,14 +341,20 @@ func SysExec(_ context.Context, env []string, input string, progress chan<- stri
345341
params.Directory = "."
346342
}
347343

344+
commandCtx, _ := engine.FromContext(ctx)
345+
346+
ctx, cancel := context.WithCancel(ctx)
347+
defer cancel()
348+
349+
commandCtx.OnUserCancel(ctx, cancel)
350+
348351
log.Debugf("Running %s in %s", params.Command, params.Directory)
349352

350353
var cmd *exec.Cmd
351-
352354
if runtime.GOOS == "windows" {
353-
cmd = exec.Command("cmd.exe", "/c", params.Command)
355+
cmd = exec.CommandContext(ctx, "cmd.exe", "/c", params.Command)
354356
} else {
355-
cmd = exec.Command("/bin/sh", "-c", params.Command)
357+
cmd = exec.CommandContext(ctx, "/bin/sh", "-c", params.Command)
356358
}
357359

358360
var (
@@ -371,7 +373,8 @@ func SysExec(_ context.Context, env []string, input string, progress chan<- stri
371373
cmd.Dir = params.Directory
372374
cmd.Stdout = combined
373375
cmd.Stderr = combined
374-
if err := cmd.Run(); err != nil {
376+
if err := cmd.Run(); err != nil && (ctx.Err() == nil || commandCtx.Ctx.Err() != nil) {
377+
// If the command failed and the context hasn't been canceled, then return the error.
375378
return fmt.Sprintf("ERROR: %s\nOUTPUT:\n%s", err, &out), nil
376379
}
377380
return out.String(), nil
@@ -420,7 +423,6 @@ func getWorkspaceEnvFileContents(envs []string) ([]string, error) {
420423
}
421424

422425
return envContents, nil
423-
424426
}
425427

426428
func getWorkspaceDir(envs []string) (string, error) {
@@ -665,6 +667,7 @@ func DiscardProgress() (progress chan<- string, closeFunc func()) {
665667
ch := make(chan string)
666668
go func() {
667669
for range ch {
670+
continue
668671
}
669672
}()
670673
return ch, func() {

‎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

+13-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,22 @@ 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+
}()
165+
166+
ctx.OnUserCancel(commandCtx, cancel)
158167

159-
if err := cmd.Run(); err != nil {
168+
if err := cmd.Run(); err != nil && (commandCtx.Err() == nil || ctx.Ctx.Err() != nil) {
169+
// If the command failed and the context hasn't been canceled, then return the error.
160170
if toolCategory == NoCategory && ctx.Parent != nil {
161171
// If this is a sub-call, then don't return the error; return the error as a message so that the LLM can retry.
162172
return fmt.Sprintf("ERROR: got (%v) while running tool, OUTPUT: %s", err, stdoutAndErr), nil
163173
}
164174
log.Errorf("failed to run tool [%s] cmd %v: %v", tool.Parameters.Name, cmd.Args, err)
165-
combinedOutput = stdoutAndErr.String()
166175
return "", fmt.Errorf("ERROR: %s: %w", stdoutAndErr, err)
167176
}
168177

169-
combinedOutput = stdoutAndErr.String()
170178
return result.String(), IsChatFinishMessage(result.String())
171179
}
172180

‎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)