fix interrupt handling: propagate context to subprocess tree
- Pass caller context into Executor.Execute so cancellation reaches the subprocess immediately - Use Setpgid + SIGKILL on process group to kill grandchild processes - Add context to ToolExecutor signature and thread it through dispatch - Reset agent state to idle after drainAgent - Rollback incomplete session turn on interrupt
This commit is contained in:
parent
c2cbc3e806
commit
70a50d4ed8
|
|
@ -89,6 +89,13 @@ func (cb *ContextBuilder) Messages() []backend.Message {
|
|||
return cb.messages
|
||||
}
|
||||
|
||||
// Truncate discards all messages after index i.
|
||||
func (cb *ContextBuilder) Truncate(i int) {
|
||||
if i < len(cb.messages) {
|
||||
cb.messages = cb.messages[:i]
|
||||
}
|
||||
}
|
||||
|
||||
// BoundedHistory returns a context-window-safe slice of messages.
|
||||
func (cb *ContextBuilder) BoundedHistory() []backend.Message {
|
||||
return cb.buildBounded(false)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import (
|
|||
|
||||
const maxRateLimitRetries = 3
|
||||
|
||||
type ToolExecutor func(name string, args json.RawMessage) (string, error)
|
||||
type ToolExecutor func(ctx context.Context, name string, args json.RawMessage) (string, error)
|
||||
type OutputFn func(msg OutputMsg)
|
||||
|
||||
type OutputMsg struct {
|
||||
|
|
@ -99,7 +99,7 @@ func Run(ctx context.Context, cfg Config, state State) error {
|
|||
var result string
|
||||
var isErr bool
|
||||
if cfg.Exec != nil {
|
||||
out, err := cfg.Exec(tc.Name, tc.Arguments)
|
||||
out, err := cfg.Exec(ctx, tc.Name, tc.Arguments)
|
||||
if err != nil {
|
||||
result = fmt.Sprintf("error: %v", err)
|
||||
isErr = true
|
||||
|
|
|
|||
|
|
@ -57,7 +57,17 @@ func (s *Session) MarkComplete() error {
|
|||
return nil
|
||||
}
|
||||
|
||||
// AppendUserMessage adds a follow-up user message and resets the complete
|
||||
// Rollback removes any trailing non-user messages from history, discarding
|
||||
// an incomplete assistant turn caused by an interruption.
|
||||
func (s *Session) Rollback() {
|
||||
s.complete = false
|
||||
msgs := s.ctx.Messages()
|
||||
i := len(msgs)
|
||||
for i > 0 && msgs[i-1].Role != "user" {
|
||||
i--
|
||||
}
|
||||
s.ctx.Truncate(i)
|
||||
}
|
||||
// flag so the loop will run again on the next call to Loop.Run.
|
||||
func (s *Session) AppendUserMessage(content string) {
|
||||
s.complete = false
|
||||
|
|
|
|||
12
exec/exec.go
12
exec/exec.go
|
|
@ -12,6 +12,7 @@ import (
|
|||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"9fans.net/go/plan9/client"
|
||||
|
|
@ -266,7 +267,7 @@ func loadLayeredConfig(name string) (*sandbox.Config, error) {
|
|||
}
|
||||
|
||||
// Execute runs code in a sandbox and returns combined stdout+stderr.
|
||||
func (e *Executor) Execute(code, language string, timeout int, sandboxName string, trusted bool) (string, error) {
|
||||
func (e *Executor) Execute(ctx context.Context, code, language string, timeout int, sandboxName string, trusted bool) (string, error) {
|
||||
start := time.Now()
|
||||
|
||||
if timeout <= 0 {
|
||||
|
|
@ -314,7 +315,7 @@ func (e *Executor) Execute(code, language string, timeout int, sandboxName strin
|
|||
defer os.RemoveAll(workDir)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var cmd *exec.Cmd
|
||||
|
|
@ -323,6 +324,13 @@ func (e *Executor) Execute(code, language string, timeout int, sandboxName strin
|
|||
wrapped := sandbox.WrapCommand(cfg, []string{"bash", "-c", code}, workDir)
|
||||
cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...)
|
||||
cmd.Dir = workDir
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||||
cmd.Cancel = func() error {
|
||||
if cmd.Process != nil {
|
||||
syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
default:
|
||||
return "", fmt.Errorf("unsupported language: %s (supported: bash)", language)
|
||||
}
|
||||
|
|
|
|||
14
main.go
14
main.go
|
|
@ -264,7 +264,7 @@ func buildAgentEnv(cfg *config.Config, builtinExec *execpkg.Executor) agentEnv {
|
|||
builtinNames[t.Name] = struct{}{}
|
||||
}
|
||||
|
||||
execFn := func(name string, args json.RawMessage) (string, error) {
|
||||
execFn := func(ctx context.Context, name string, args json.RawMessage) (string, error) {
|
||||
if server, ok := serverOf[name]; ok {
|
||||
raw, err := mcpExec.Execute(server, name, args)
|
||||
if err != nil {
|
||||
|
|
@ -273,7 +273,7 @@ func buildAgentEnv(cfg *config.Config, builtinExec *execpkg.Executor) agentEnv {
|
|||
return extractMCPText(raw), nil
|
||||
}
|
||||
if _, ok := builtinNames[name]; ok {
|
||||
return dispatchBuiltinExec(name, builtinExec, args)
|
||||
return dispatchBuiltinExec(ctx, name, builtinExec, args)
|
||||
}
|
||||
return "", fmt.Errorf("unknown tool: %s", name)
|
||||
}
|
||||
|
|
@ -537,6 +537,7 @@ func (m *model) apply(am agentMsg) {
|
|||
}
|
||||
|
||||
// drainAgent cancels the in-flight goroutine and drains remaining events.
|
||||
// It also rolls back any incomplete assistant turn from the session history.
|
||||
func (m *model) drainAgent() {
|
||||
if m.cancel == nil {
|
||||
return
|
||||
|
|
@ -555,6 +556,11 @@ func (m *model) drainAgent() {
|
|||
}
|
||||
m.agentCh = nil
|
||||
}
|
||||
if m.session != nil {
|
||||
m.session.Rollback()
|
||||
}
|
||||
m.state = agentIdle
|
||||
m.currentTool = ""
|
||||
}
|
||||
|
||||
func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
|
|
@ -893,7 +899,7 @@ func extractMCPText(raw json.RawMessage) string {
|
|||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func dispatchBuiltinExec(name string, e *execpkg.Executor, args json.RawMessage) (string, error) {
|
||||
func dispatchBuiltinExec(ctx context.Context, name string, e *execpkg.Executor, args json.RawMessage) (string, error) {
|
||||
var a struct {
|
||||
Code string `json:"code"`
|
||||
Language string `json:"language"`
|
||||
|
|
@ -954,5 +960,5 @@ func dispatchBuiltinExec(name string, e *execpkg.Executor, args json.RawMessage)
|
|||
code = a.Code
|
||||
}
|
||||
|
||||
return e.Execute(code, a.Language, a.Timeout, a.Sandbox, trusted)
|
||||
return e.Execute(ctx, code, a.Language, a.Timeout, a.Sandbox, trusted)
|
||||
}
|
||||
|
|
|
|||
Reference in New Issue