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:
Levi Neely 2026-04-04 18:18:02 +02:00
parent c2cbc3e806
commit 70a50d4ed8
5 changed files with 40 additions and 9 deletions

View File

@ -89,6 +89,13 @@ func (cb *ContextBuilder) Messages() []backend.Message {
return cb.messages 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. // BoundedHistory returns a context-window-safe slice of messages.
func (cb *ContextBuilder) BoundedHistory() []backend.Message { func (cb *ContextBuilder) BoundedHistory() []backend.Message {
return cb.buildBounded(false) return cb.buildBounded(false)

View File

@ -13,7 +13,7 @@ import (
const maxRateLimitRetries = 3 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 OutputFn func(msg OutputMsg)
type OutputMsg struct { type OutputMsg struct {
@ -99,7 +99,7 @@ func Run(ctx context.Context, cfg Config, state State) error {
var result string var result string
var isErr bool var isErr bool
if cfg.Exec != nil { if cfg.Exec != nil {
out, err := cfg.Exec(tc.Name, tc.Arguments) out, err := cfg.Exec(ctx, tc.Name, tc.Arguments)
if err != nil { if err != nil {
result = fmt.Sprintf("error: %v", err) result = fmt.Sprintf("error: %v", err)
isErr = true isErr = true

View File

@ -57,7 +57,17 @@ func (s *Session) MarkComplete() error {
return nil 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. // flag so the loop will run again on the next call to Loop.Run.
func (s *Session) AppendUserMessage(content string) { func (s *Session) AppendUserMessage(content string) {
s.complete = false s.complete = false

View File

@ -12,6 +12,7 @@ import (
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
"syscall"
"time" "time"
"9fans.net/go/plan9/client" "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. // 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() start := time.Now()
if timeout <= 0 { if timeout <= 0 {
@ -314,7 +315,7 @@ func (e *Executor) Execute(code, language string, timeout int, sandboxName strin
defer os.RemoveAll(workDir) 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() defer cancel()
var cmd *exec.Cmd 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) wrapped := sandbox.WrapCommand(cfg, []string{"bash", "-c", code}, workDir)
cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...) cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...)
cmd.Dir = workDir 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: default:
return "", fmt.Errorf("unsupported language: %s (supported: bash)", language) return "", fmt.Errorf("unsupported language: %s (supported: bash)", language)
} }

14
main.go
View File

@ -264,7 +264,7 @@ func buildAgentEnv(cfg *config.Config, builtinExec *execpkg.Executor) agentEnv {
builtinNames[t.Name] = struct{}{} 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 { if server, ok := serverOf[name]; ok {
raw, err := mcpExec.Execute(server, name, args) raw, err := mcpExec.Execute(server, name, args)
if err != nil { if err != nil {
@ -273,7 +273,7 @@ func buildAgentEnv(cfg *config.Config, builtinExec *execpkg.Executor) agentEnv {
return extractMCPText(raw), nil return extractMCPText(raw), nil
} }
if _, ok := builtinNames[name]; ok { if _, ok := builtinNames[name]; ok {
return dispatchBuiltinExec(name, builtinExec, args) return dispatchBuiltinExec(ctx, name, builtinExec, args)
} }
return "", fmt.Errorf("unknown tool: %s", name) 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. // 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() { func (m *model) drainAgent() {
if m.cancel == nil { if m.cancel == nil {
return return
@ -555,6 +556,11 @@ func (m *model) drainAgent() {
} }
m.agentCh = nil 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) { 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") 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 { var a struct {
Code string `json:"code"` Code string `json:"code"`
Language string `json:"language"` Language string `json:"language"`
@ -954,5 +960,5 @@ func dispatchBuiltinExec(name string, e *execpkg.Executor, args json.RawMessage)
code = a.Code 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)
} }