mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-18 15:42:26 +02:00
feat: manage task process lifetimes and preserve turn history
This commit is contained in:
Ed1s0nZ
+21
-132
@@ -1,7 +1,6 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
@@ -9,7 +8,6 @@ import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -836,128 +834,13 @@ func (e *Executor) executeSystemCommand(ctx context.Context, args map[string]int
|
||||
zap.Bool("isBackground", isBackground),
|
||||
)
|
||||
|
||||
// 如果是后台命令，使用特殊处理来获取实际的后台进程PID
|
||||
if isBackground {
|
||||
// 移除命令末尾的 & 符号
|
||||
commandWithoutAmpersand := strings.TrimSuffix(strings.TrimSpace(command), "&")
|
||||
commandWithoutAmpersand = strings.TrimSpace(commandWithoutAmpersand)
|
||||
|
||||
// 构建新命令：后台作业重定向标准流后 echo $pid（与 RedirectBackgroundJobStdio 一致）。
|
||||
pidCommand := RedirectBackgroundJobStdio(commandWithoutAmpersand+" &") + " pid=$!; echo $pid"
|
||||
|
||||
// 创建新命令来获取PID
|
||||
var pidCmd *exec.Cmd
|
||||
if workDir != "" {
|
||||
pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand)
|
||||
pidCmd.Dir = workDir
|
||||
} else {
|
||||
pidCmd = exec.CommandContext(ctx, shell, "-c", pidCommand)
|
||||
}
|
||||
ConfigureShellCmdForAgentExecute(pidCmd)
|
||||
|
||||
// 获取stdout管道
|
||||
stdout, err := pidCmd.StdoutPipe()
|
||||
job := strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(command), "&"))
|
||||
session, err := StartManagedBackground(ctx, shell, job, workDir)
|
||||
if err != nil {
|
||||
e.logger.Error("创建stdout管道失败",
|
||||
zap.String("command", command),
|
||||
zap.Error(err),
|
||||
)
|
||||
// 如果创建管道失败，使用shell进程的PID作为fallback
|
||||
if err := pidCmd.Start(); err != nil {
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("后台命令启动失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
pid := pidCmd.Process.Pid
|
||||
go pidCmd.Wait() // 在后台等待，避免僵尸进程
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d (可能不准确，获取PID失败)\n\n注意: 后台进程将继续运行，不会等待其完成。", command, pid),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: fmt.Sprintf("后台命令启动失败: %v", err)}}, IsError: true}, nil
|
||||
}
|
||||
|
||||
// 启动命令
|
||||
if err := pidCmd.Start(); err != nil {
|
||||
stdout.Close()
|
||||
e.logger.Error("后台命令启动失败",
|
||||
zap.String("command", command),
|
||||
zap.Error(err),
|
||||
)
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("后台命令启动失败: %v", err),
|
||||
},
|
||||
},
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// 读取第一行输出（PID）
|
||||
reader := bufio.NewReader(stdout)
|
||||
pidLine, err := reader.ReadString('\n')
|
||||
stdout.Close()
|
||||
|
||||
var actualPid int
|
||||
if err != nil && err != io.EOF {
|
||||
e.logger.Warn("读取后台进程PID失败",
|
||||
zap.String("command", command),
|
||||
zap.Error(err),
|
||||
)
|
||||
// 如果读取失败，使用shell进程的PID
|
||||
actualPid = pidCmd.Process.Pid
|
||||
} else {
|
||||
// 解析PID
|
||||
pidStr := strings.TrimSpace(pidLine)
|
||||
if parsedPid, err := strconv.Atoi(pidStr); err == nil {
|
||||
actualPid = parsedPid
|
||||
} else {
|
||||
e.logger.Warn("解析后台进程PID失败",
|
||||
zap.String("command", command),
|
||||
zap.String("pidLine", pidStr),
|
||||
zap.Error(err),
|
||||
)
|
||||
// 如果解析失败，使用shell进程的PID
|
||||
actualPid = pidCmd.Process.Pid
|
||||
}
|
||||
}
|
||||
|
||||
// 在goroutine中等待shell进程，避免僵尸进程
|
||||
go func() {
|
||||
if err := pidCmd.Wait(); err != nil {
|
||||
e.logger.Debug("后台命令shell进程执行完成",
|
||||
zap.String("command", command),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
}()
|
||||
|
||||
e.logger.Info("后台命令已启动",
|
||||
zap.String("command", command),
|
||||
zap.Int("actualPid", actualPid),
|
||||
)
|
||||
|
||||
return &mcp.ToolResult{
|
||||
Content: []mcp.Content{
|
||||
{
|
||||
Type: "text",
|
||||
Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程ID: %d\n\n注意: 后台进程将继续运行，不会等待其完成。", command, actualPid),
|
||||
},
|
||||
},
|
||||
IsError: false,
|
||||
}, nil
|
||||
return &mcp.ToolResult{Content: []mcp.Content{{Type: "text", Text: fmt.Sprintf("后台命令已启动\n命令: %s\n进程组ID: %d\n\n后台进程由本轮任务托管，任务结束时自动清理。", command, session.rootPID)}}}, nil
|
||||
}
|
||||
|
||||
// 非后台命令：等待输出
|
||||
@@ -1041,7 +924,7 @@ func combinedOutputCancellableWithLimit(ctx context.Context, cmd *exec.Cmd, maxB
|
||||
cmd.Stdout = stdoutBuf
|
||||
cmd.Stderr = stderrBuf
|
||||
|
||||
session, err := StartShellSession(cmd)
|
||||
session, err := StartShellSessionContext(ctx, cmd)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -1248,7 +1131,7 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
|
||||
_ = stdoutPipe.Close()
|
||||
return "", err
|
||||
}
|
||||
session, err := StartShellSession(cmd)
|
||||
session, err := StartShellSessionContext(ctx, cmd)
|
||||
if err != nil {
|
||||
_ = stdoutPipe.Close()
|
||||
_ = stderrPipe.Close()
|
||||
@@ -1265,6 +1148,8 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
|
||||
}()
|
||||
defer close(stopWatch)
|
||||
|
||||
readStop := make(chan struct{})
|
||||
defer close(readStop)
|
||||
chunks := make(chan string, 64)
|
||||
var wg sync.WaitGroup
|
||||
readFn := func(r io.Reader) {
|
||||
@@ -1273,7 +1158,11 @@ func streamCommandOutput(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallba
|
||||
for {
|
||||
n, readErr := r.Read(buf)
|
||||
if n > 0 {
|
||||
chunks <- string(buf[:n])
|
||||
select {
|
||||
case chunks <- string(buf[:n]):
|
||||
case <-readStop:
|
||||
return
|
||||
}
|
||||
}
|
||||
if readErr != nil {
|
||||
return
|
||||
@@ -1422,24 +1311,24 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
||||
}
|
||||
|
||||
_ = prepareShellCmdSession(cmd)
|
||||
ptmx, err := pty.Start(cmd)
|
||||
var ptmx *os.File
|
||||
session, err := startShellSessionContext(ctx, cmd, func() error {
|
||||
var startErr error
|
||||
ptmx, startErr = pty.Start(cmd)
|
||||
return startErr
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer func() { _ = ptmx.Close() }()
|
||||
|
||||
rootPID := 0
|
||||
if cmd.Process != nil {
|
||||
rootPID = cmd.Process.Pid
|
||||
}
|
||||
|
||||
// ctx 取消时尽快终止子进程
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
_ = ptmx.Close() // 触发读退出
|
||||
terminateProcessGroup(rootPID, cmd)
|
||||
session.Terminate()
|
||||
case <-done:
|
||||
}
|
||||
}()
|
||||
@@ -1484,7 +1373,7 @@ func runCommandWithPTY(ctx context.Context, cmd *exec.Cmd, cb ToolOutputCallback
|
||||
}
|
||||
flush()
|
||||
|
||||
waitErr := cmd.Wait()
|
||||
waitErr := session.Wait()
|
||||
return finalizeBoundedOutput(outBuilder, maxBytes, tee), waitErr
|
||||
}
|
||||
|
||||
|
||||
@@ -54,7 +54,13 @@ func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T)
|
||||
executor, _ := setupTestExecutor(t)
|
||||
// 子进程先向 stdout 写无换行字符再长时间 sleep；若与 echo $pid 共享管道且未重定向子进程 stdout，
|
||||
// ReadString('\n') 会阻塞到子进程退出。后台包装须将子进程标准流与 PID 行分离。
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
|
||||
scope := NewProcessScope()
|
||||
t.Cleanup(func() {
|
||||
if err := scope.Close(); err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
})
|
||||
ctx, cancel := context.WithTimeout(WithProcessScope(context.Background(), scope), 4*time.Second)
|
||||
defer cancel()
|
||||
args := map[string]interface{}{
|
||||
"command": `(sh -c 'printf x; sleep 120') &`,
|
||||
|
||||
@@ -39,3 +39,14 @@ func terminateProcessGroup(rootPID int, cmd *exec.Cmd) {
|
||||
func terminateCmdTree(cmd *exec.Cmd) {
|
||||
terminateProcessGroup(0, cmd)
|
||||
}
|
||||
|
||||
// stopProcessGroup gives the whole job a grace period to release resources.
|
||||
func stopProcessGroup(pid int, cmd *exec.Cmd) {
|
||||
if pid > 0 {
|
||||
_ = syscall.Kill(-pid, syscall.SIGTERM)
|
||||
}
|
||||
}
|
||||
|
||||
func processGroupExists(pid int) bool {
|
||||
return pid > 0 && syscall.Kill(-pid, 0) != syscall.ESRCH
|
||||
}
|
||||
|
||||
@@ -3,9 +3,11 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"strconv"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
|
||||
func prepareShellCmdSession(cmd *exec.Cmd) error {
|
||||
@@ -29,7 +31,9 @@ func terminateProcessGroup(rootPID int, cmd *exec.Cmd) {
|
||||
if pid <= 0 {
|
||||
return
|
||||
}
|
||||
tk := exec.Command("taskkill", "/F", "/T", "/PID", strconv.Itoa(pid))
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
tk := exec.CommandContext(ctx, "taskkill", "/F", "/T", "/PID", strconv.Itoa(pid))
|
||||
if err := tk.Run(); err != nil {
|
||||
if cmd != nil && cmd.Process != nil {
|
||||
_ = cmd.Process.Kill()
|
||||
@@ -41,3 +45,12 @@ func terminateProcessGroup(rootPID int, cmd *exec.Cmd) {
|
||||
func terminateCmdTree(cmd *exec.Cmd) {
|
||||
terminateProcessGroup(0, cmd)
|
||||
}
|
||||
|
||||
func stopProcessGroup(pid int, cmd *exec.Cmd) {
|
||||
// Windows has no portable SIGTERM equivalent for arbitrary console jobs.
|
||||
terminateProcessGroup(pid, cmd)
|
||||
}
|
||||
|
||||
// Windows taskkill /T is best effort; unlike a Unix PGID it has no persistent
|
||||
// group handle to query after the root exits. Job Objects are needed for that.
|
||||
func processGroupExists(pid int) bool { return false }
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
package security
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"cyberstrike-ai/internal/processguard"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
var ErrProcessScopeClosed = errors.New("task is ending; new processes are not allowed")
|
||||
var ErrBackgroundNeedsTask = errors.New("background commands require a managed task")
|
||||
|
||||
type processScopeKey struct{}
|
||||
|
||||
// ProcessScope owns local commands for one task run, including background work.
|
||||
// Ownership is carried by context values, so MCP's WithoutCancel retains it.
|
||||
// Start and Seal serialize under the same lock: no process can escape cleanup
|
||||
// by starting between the final snapshot and task completion.
|
||||
type ProcessScope struct {
|
||||
ID string
|
||||
mu sync.Mutex
|
||||
closed bool
|
||||
sessions map[*ShellSession]struct{}
|
||||
guard processguard.Group
|
||||
guardErr error
|
||||
closeMu sync.Mutex
|
||||
}
|
||||
|
||||
func NewProcessScope() *ProcessScope {
|
||||
return &ProcessScope{ID: uuid.NewString(), sessions: make(map[*ShellSession]struct{})}
|
||||
}
|
||||
|
||||
func WithProcessScope(ctx context.Context, scope *ProcessScope) context.Context {
|
||||
return context.WithValue(ctx, processScopeKey{}, scope)
|
||||
}
|
||||
|
||||
func ProcessScopeFromContext(ctx context.Context) *ProcessScope {
|
||||
if ctx == nil {
|
||||
return nil
|
||||
}
|
||||
scope, _ := ctx.Value(processScopeKey{}).(*ProcessScope)
|
||||
return scope
|
||||
}
|
||||
|
||||
func (s *ProcessScope) Seal() {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.closed = true
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func startShellSessionContext(ctx context.Context, cmd *exec.Cmd, start func() error) (*ShellSession, error) {
|
||||
scope := ProcessScopeFromContext(ctx)
|
||||
if scope != nil {
|
||||
scope.mu.Lock()
|
||||
defer scope.mu.Unlock()
|
||||
if scope.closed {
|
||||
return nil, ErrProcessScopeClosed
|
||||
}
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := prepareShellCmdSession(cmd); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var launch *processguard.Launch
|
||||
if scope != nil {
|
||||
if scope.guard == nil && scope.guardErr == nil {
|
||||
scope.guard, scope.guardErr = processguard.New(scope.ID)
|
||||
}
|
||||
if scope.guardErr != nil {
|
||||
return nil, scope.guardErr
|
||||
}
|
||||
var err error
|
||||
launch, err = scope.guard.Prepare(cmd)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer launch.Dispose()
|
||||
}
|
||||
// Bound Go's output-copy goroutines when descendants inherit a pipe.
|
||||
if cmd.WaitDelay == 0 {
|
||||
cmd.WaitDelay = 2 * time.Second
|
||||
}
|
||||
if err := start(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if launch != nil {
|
||||
if err := launch.Commit(); err != nil {
|
||||
terminateProcessGroup(cmd.Process.Pid, cmd)
|
||||
_ = cmd.Wait()
|
||||
if scope != nil {
|
||||
_ = scope.guard.Release(cmd.Process.Pid)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
session := &ShellSession{Cmd: cmd, rootPID: cmd.Process.Pid, scope: scope, done: make(chan struct{})}
|
||||
if scope != nil {
|
||||
scope.sessions[session] = struct{}{}
|
||||
}
|
||||
return session, nil
|
||||
}
|
||||
|
||||
// Close seals the scope, asks every process group to exit, then escalates to
|
||||
// SIGKILL. It waits for command reaping, with one shared deadline, not N timeouts.
|
||||
// Failed entries remain owned, permitting a later Close to retry cleanup.
|
||||
func (s *ProcessScope) Close() error {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.closeMu.Lock()
|
||||
defer s.closeMu.Unlock()
|
||||
s.mu.Lock()
|
||||
s.closed = true
|
||||
sessions := make([]*ShellSession, 0, len(s.sessions))
|
||||
for session := range s.sessions {
|
||||
sessions = append(sessions, session)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if len(sessions) == 0 {
|
||||
return s.closeGuard()
|
||||
}
|
||||
for _, session := range sessions {
|
||||
session.signal(false)
|
||||
}
|
||||
if waitShellSessions(sessions, 3*time.Second) {
|
||||
return s.closeGuard()
|
||||
}
|
||||
for _, session := range sessions {
|
||||
session.Terminate()
|
||||
}
|
||||
guardErr := s.closeGuard()
|
||||
if waitShellSessions(sessions, 3*time.Second) {
|
||||
return guardErr
|
||||
}
|
||||
remaining := make([]int, 0, len(sessions))
|
||||
for _, session := range sessions {
|
||||
if !session.tryComplete() {
|
||||
remaining = append(remaining, session.rootPID)
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("task %s: process cleanup timed out (process groups %v)", s.ID, remaining)
|
||||
}
|
||||
|
||||
func waitShellSessions(sessions []*ShellSession, timeout time.Duration) bool {
|
||||
deadline := time.NewTimer(timeout)
|
||||
defer deadline.Stop()
|
||||
ticker := time.NewTicker(10 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
complete := true
|
||||
for _, session := range sessions {
|
||||
if !session.tryComplete() {
|
||||
complete = false
|
||||
}
|
||||
}
|
||||
if complete {
|
||||
return true
|
||||
}
|
||||
select {
|
||||
case <-deadline.C:
|
||||
return false
|
||||
case <-ticker.C:
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
// StartManagedBackground returns promptly while retaining task ownership. The
|
||||
// shell executes the job in the foreground internally, keeping a waitable root
|
||||
// alive; tool completion must not cancel the job's lifetime.
|
||||
func StartManagedBackground(ctx context.Context, shell, command, dir string) (*ShellSession, error) {
|
||||
if ProcessScopeFromContext(ctx) == nil {
|
||||
return nil, ErrBackgroundNeedsTask
|
||||
}
|
||||
cmd := exec.Command(shell, "-c", PrepareShellCommandForExecute(command))
|
||||
cmd.Dir = dir
|
||||
ConfigureShellCmdForAgentExecute(cmd)
|
||||
// Nil output streams use /dev/null; background output cannot hold tool pipes.
|
||||
session, err := StartShellSessionContext(ctx, cmd)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
go func() { _ = session.Wait() }()
|
||||
return session, nil
|
||||
}
|
||||
|
||||
func (s *ProcessScope) closeGuard() error {
|
||||
if s.guard == nil {
|
||||
return nil
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
return s.guard.Close(ctx)
|
||||
}
|
||||
|
||||
func (s *ProcessScope) IsolationBackend() string {
|
||||
if s == nil {
|
||||
return "none"
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.guard != nil {
|
||||
return s.guard.Name()
|
||||
}
|
||||
if s.guardErr != nil {
|
||||
return "unavailable"
|
||||
}
|
||||
return "pending"
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
//go:build !windows
|
||||
|
||||
package security
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/cloudwego/eino/adk/filesystem"
|
||||
)
|
||||
|
||||
func readTestPID(t *testing.T, path string) int {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
b, err := os.ReadFile(path)
|
||||
if err == nil {
|
||||
if pid, err := strconv.Atoi(strings.TrimSpace(string(b))); err == nil && pid > 0 {
|
||||
return pid
|
||||
}
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("process did not write PID to %s", path)
|
||||
return 0
|
||||
}
|
||||
|
||||
func requireProcessGone(t *testing.T, pid int) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if syscall.Kill(pid, 0) == syscall.ESRCH {
|
||||
return
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("process %d survived task cleanup", pid)
|
||||
}
|
||||
|
||||
func TestProcessScope_BackgroundSurvivesToolButEndsWithTask(t *testing.T) {
|
||||
executor, _ := setupTestExecutor(t)
|
||||
scope := NewProcessScope()
|
||||
t.Cleanup(func() { _ = scope.Close() })
|
||||
taskCtx := WithProcessScope(context.Background(), scope)
|
||||
ctx, cancel := context.WithCancel(context.WithoutCancel(taskCtx))
|
||||
defer cancel()
|
||||
pidFile := filepath.Join(t.TempDir(), "pid")
|
||||
result, err := executor.executeSystemCommand(ctx, map[string]interface{}{
|
||||
"command": fmt.Sprintf("echo $$ > %q; sleep 300 &", pidFile),
|
||||
})
|
||||
if err != nil || result.IsError {
|
||||
t.Fatalf("background launch: %v, %+v", err, result)
|
||||
}
|
||||
pid := readTestPID(t, pidFile)
|
||||
cancel() // MCP completes and cancels its per-tool context.
|
||||
if err := syscall.Kill(pid, 0); err != nil {
|
||||
t.Fatalf("tool completion killed task background process: %v", err)
|
||||
}
|
||||
if err := scope.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
requireProcessGone(t, pid)
|
||||
if _, err := StartManagedBackground(taskCtx, "sh", "sleep 300", ""); !errors.Is(err, ErrProcessScopeClosed) {
|
||||
t.Fatalf("closed task accepted a new process: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessScope_EinoBackgroundReturnsPromptlyAndIsOwned(t *testing.T) {
|
||||
for _, useFlag := range []bool{false, true} {
|
||||
t.Run(fmt.Sprint(useFlag), func(t *testing.T) {
|
||||
scope := NewProcessScope()
|
||||
t.Cleanup(func() { _ = scope.Close() })
|
||||
ctx := WithProcessScope(context.Background(), scope)
|
||||
pidFile := filepath.Join(t.TempDir(), "pid")
|
||||
command := fmt.Sprintf("echo $$ > %q; sleep 300", pidFile)
|
||||
if !useFlag {
|
||||
command += " &"
|
||||
}
|
||||
stream, err := NewEinoStreamingShell().ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: command, RunInBackendGround: useFlag})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer stream.Close()
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
for {
|
||||
_, err := stream.Recv()
|
||||
if err != nil {
|
||||
done <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, io.EOF) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("background launch waited for job completion")
|
||||
}
|
||||
pid := readTestPID(t, pidFile)
|
||||
if err := scope.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
requireProcessGone(t, pid)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessScope_ForceKillsIgnoringTERMAndGrandchild(t *testing.T) {
|
||||
scope := NewProcessScope()
|
||||
t.Cleanup(func() { _ = scope.Close() })
|
||||
ctx := WithProcessScope(context.Background(), scope)
|
||||
pidFile := filepath.Join(t.TempDir(), "child")
|
||||
session, err := StartManagedBackground(ctx, "sh", fmt.Sprintf("trap '' TERM; sleep 300 & echo $! > %q; wait", pidFile), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
childPID := readTestPID(t, pidFile)
|
||||
if err := scope.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
requireProcessGone(t, session.rootPID)
|
||||
requireProcessGone(t, childPID)
|
||||
if session.Cmd.ProcessState == nil {
|
||||
t.Fatal("root process was not reaped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessScope_ConcurrentStartAndClose(t *testing.T) {
|
||||
scope := NewProcessScope()
|
||||
ctx := WithProcessScope(context.Background(), scope)
|
||||
t.Cleanup(func() { _ = scope.Close() })
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
var sessions []*ShellSession
|
||||
begin := make(chan struct{})
|
||||
for i := 0; i < 24; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-begin
|
||||
session, err := StartManagedBackground(ctx, "sh", "sleep 300", "")
|
||||
if err != nil {
|
||||
if !errors.Is(err, ErrProcessScopeClosed) {
|
||||
t.Errorf("start: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
mu.Lock()
|
||||
sessions = append(sessions, session)
|
||||
mu.Unlock()
|
||||
}()
|
||||
}
|
||||
close(begin)
|
||||
if err := scope.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wg.Wait()
|
||||
for _, session := range sessions {
|
||||
requireProcessGone(t, session.rootPID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessScope_UnmanagedBackgroundRejected(t *testing.T) {
|
||||
if _, err := StartManagedBackground(context.Background(), "sh", "sleep 300", ""); !errors.Is(err, ErrBackgroundNeedsTask) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProcessScope_ForegroundExitKillsLeftoverChild(t *testing.T) {
|
||||
scope := NewProcessScope()
|
||||
t.Cleanup(func() { _ = scope.Close() })
|
||||
ctx := WithProcessScope(context.Background(), scope)
|
||||
pidFile := filepath.Join(t.TempDir(), "child")
|
||||
// A shell that exits with a redirected child must not lose that child.
|
||||
cmd := exec.CommandContext(ctx, "sh", "-c", fmt.Sprintf("sleep 300 </dev/null >/dev/null 2>&1 & echo $! > %q", pidFile))
|
||||
if _, err := combinedOutputCancellable(ctx, cmd); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := readTestPID(t, pidFile)
|
||||
if err := scope.Close(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
requireProcessGone(t, pid)
|
||||
}
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/cloudwego/eino/adk/filesystem"
|
||||
@@ -49,7 +50,7 @@ func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesy
|
||||
}
|
||||
|
||||
sr, w := schema.Pipe[*filesystem.ExecuteResponse](100)
|
||||
if input.RunInBackendGround {
|
||||
if input.RunInBackendGround || IsBackgroundShellCommand(input.Command) {
|
||||
go runShellInBackground(ctx, input.Command, w)
|
||||
return sr, nil
|
||||
}
|
||||
@@ -60,45 +61,18 @@ func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesy
|
||||
func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) {
|
||||
defer w.Close()
|
||||
|
||||
command = PrepareShellCommandForExecute(command)
|
||||
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
|
||||
applyDefaultTerminalEnv(cmd)
|
||||
attachNonInteractiveStdin(cmd)
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
command = strings.TrimSpace(command)
|
||||
if IsBackgroundShellCommand(command) {
|
||||
command = strings.TrimSpace(strings.TrimSuffix(command, "&"))
|
||||
}
|
||||
session, err := StartManagedBackground(ctx, "/bin/sh", command, "")
|
||||
if err != nil {
|
||||
_ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err))
|
||||
_ = w.Send(nil, err)
|
||||
return
|
||||
}
|
||||
stderr, err := cmd.StderrPipe()
|
||||
if err != nil {
|
||||
_ = stdout.Close()
|
||||
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
|
||||
return
|
||||
}
|
||||
session, err := StartShellSession(cmd)
|
||||
if err != nil {
|
||||
_ = stdout.Close()
|
||||
_ = stderr.Close()
|
||||
_ = w.Send(nil, fmt.Errorf("failed to start command: %w", err))
|
||||
return
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
drainShellPipes(stdout, stderr)
|
||||
_ = session.Wait()
|
||||
close(done)
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-ctx.Done():
|
||||
TerminateShellCmdSession(session)
|
||||
}
|
||||
|
||||
exitCode := 0
|
||||
_ = w.Send(&filesystem.ExecuteResponse{
|
||||
Output: "command started in background\n",
|
||||
Output: fmt.Sprintf("command started in background (process group %d); cleaned up when this task ends\n", session.rootPID),
|
||||
ExitCode: &exitCode,
|
||||
}, nil)
|
||||
}
|
||||
@@ -136,7 +110,7 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream
|
||||
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
|
||||
return
|
||||
}
|
||||
session, err := StartShellSession(cmd)
|
||||
session, err := StartShellSessionContext(ctx, cmd)
|
||||
if err != nil {
|
||||
_ = stdoutPipe.Close()
|
||||
_ = stderrPipe.Close()
|
||||
@@ -154,6 +128,8 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream
|
||||
}()
|
||||
defer close(stopWatch)
|
||||
|
||||
readStop := make(chan struct{})
|
||||
defer close(readStop)
|
||||
chunks := make(chan string, 64)
|
||||
var wg sync.WaitGroup
|
||||
readFn := func(r io.Reader) {
|
||||
@@ -162,7 +138,11 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream
|
||||
for {
|
||||
n, readErr := r.Read(buf)
|
||||
if n > 0 {
|
||||
chunks <- string(buf[:n])
|
||||
select {
|
||||
case chunks <- string(buf[:n]):
|
||||
case <-readStop:
|
||||
return
|
||||
}
|
||||
}
|
||||
if readErr != nil {
|
||||
return
|
||||
@@ -186,6 +166,7 @@ func streamShellForeground(ctx context.Context, command string, w *schema.Stream
|
||||
hadOutput = true
|
||||
if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) {
|
||||
TerminateShellCmdSession(session)
|
||||
go func() { _ = session.Wait() }()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,47 +1,102 @@
|
||||
package security
|
||||
|
||||
import "os/exec"
|
||||
import (
|
||||
"context"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ShellSession 在 Start 时记录根 shell 的进程组 ID，取消/超时时可杀整组（即使 cmd.Process 已失效）。
|
||||
// ShellSession caches the process group while its command is alive. Signals
|
||||
// and Wait completion synchronize to avoid signalling already-released sessions.
|
||||
type ShellSession struct {
|
||||
Cmd *exec.Cmd
|
||||
rootPID int
|
||||
Cmd *exec.Cmd
|
||||
rootPID int
|
||||
scope *ProcessScope
|
||||
done chan struct{}
|
||||
waitOnce sync.Once
|
||||
waitErr error
|
||||
signalMu sync.Mutex
|
||||
finished bool
|
||||
waited bool
|
||||
}
|
||||
|
||||
// StartShellSession 配置独立进程组并启动 shell，缓存 rootPID（Unix 下即 PGID）。
|
||||
func StartShellSession(cmd *exec.Cmd) (*ShellSession, error) {
|
||||
if err := prepareShellCmdSession(cmd); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pid := 0
|
||||
if cmd.Process != nil {
|
||||
pid = cmd.Process.Pid
|
||||
}
|
||||
return &ShellSession{Cmd: cmd, rootPID: pid}, nil
|
||||
return StartShellSessionContext(context.Background(), cmd)
|
||||
}
|
||||
|
||||
func StartShellSessionContext(ctx context.Context, cmd *exec.Cmd) (*ShellSession, error) {
|
||||
return startShellSessionContext(ctx, cmd, cmd.Start)
|
||||
}
|
||||
|
||||
// Wait 等待 shell 退出。
|
||||
func (s *ShellSession) Wait() error {
|
||||
if s == nil || s.Cmd == nil {
|
||||
return nil
|
||||
}
|
||||
return s.Cmd.Wait()
|
||||
s.waitOnce.Do(func() {
|
||||
s.waitErr = s.Cmd.Wait()
|
||||
s.signalMu.Lock()
|
||||
s.waited = true
|
||||
terminateProcessGroup(s.rootPID, s.Cmd)
|
||||
s.signalMu.Unlock()
|
||||
// Usually the group disappears immediately. Retain ownership if the
|
||||
// kernel cannot confirm exit; task cleanup will retry and report it.
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for !s.tryComplete() && time.Now().Before(deadline) {
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
})
|
||||
return s.waitErr
|
||||
}
|
||||
|
||||
// Terminate 终止 shell 及其进程组。
|
||||
func (s *ShellSession) Terminate() {
|
||||
func (s *ShellSession) signal(force bool) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
terminateProcessGroup(s.rootPID, s.Cmd)
|
||||
s.signalMu.Lock()
|
||||
defer s.signalMu.Unlock()
|
||||
if s.finished {
|
||||
return
|
||||
}
|
||||
if force {
|
||||
terminateProcessGroup(s.rootPID, s.Cmd)
|
||||
} else {
|
||||
stopProcessGroup(s.rootPID, s.Cmd)
|
||||
}
|
||||
}
|
||||
|
||||
// TerminateShellSession 终止由 StartShellSession 启动的会话。
|
||||
func (s *ShellSession) Terminate() { s.signal(true) }
|
||||
|
||||
func TerminateShellSession(session *ShellSession) {
|
||||
if session != nil {
|
||||
session.Terminate()
|
||||
}
|
||||
}
|
||||
|
||||
// tryComplete confirms group exit after Wait reaped the direct child. Never
|
||||
// release ownership merely because a signal was sent successfully.
|
||||
func (s *ShellSession) tryComplete() bool {
|
||||
s.signalMu.Lock()
|
||||
defer s.signalMu.Unlock()
|
||||
if s.finished {
|
||||
return true
|
||||
}
|
||||
if !s.waited || processGroupExists(s.rootPID) {
|
||||
return false
|
||||
}
|
||||
s.finished = true
|
||||
if s.scope != nil {
|
||||
s.scope.mu.Lock()
|
||||
if s.scope.guard != nil {
|
||||
if err := s.scope.guard.Release(s.rootPID); err != nil {
|
||||
s.scope.mu.Unlock()
|
||||
s.finished = false
|
||||
return false
|
||||
}
|
||||
}
|
||||
delete(s.scope.sessions, s)
|
||||
s.scope.mu.Unlock()
|
||||
}
|
||||
close(s.done)
|
||||
return true
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user