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
|
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)
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
12
exec/exec.go
12
exec/exec.go
|
|
@ -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
14
main.go
|
|
@ -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)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Reference in New Issue