251 lines
6.2 KiB
Go
251 lines
6.2 KiB
Go
package process
|
||
|
||
import (
|
||
"bufio"
|
||
"context"
|
||
"cursor-api-proxy/internal/env"
|
||
"fmt"
|
||
"os/exec"
|
||
"strings"
|
||
"sync"
|
||
"syscall"
|
||
"time"
|
||
)
|
||
|
||
type RunResult struct {
|
||
Code int
|
||
Stdout string
|
||
Stderr string
|
||
}
|
||
|
||
type RunOptions struct {
|
||
Cwd string
|
||
TimeoutMs int
|
||
MaxMode bool
|
||
ConfigDir string
|
||
Ctx context.Context
|
||
}
|
||
|
||
type RunStreamingOptions struct {
|
||
RunOptions
|
||
OnLine func(line string)
|
||
}
|
||
|
||
// ─── Global child process registry ──────────────────────────────────────────
|
||
|
||
var (
|
||
activeMu sync.Mutex
|
||
activeChildren []*exec.Cmd
|
||
)
|
||
|
||
func registerChild(c *exec.Cmd) {
|
||
activeMu.Lock()
|
||
activeChildren = append(activeChildren, c)
|
||
activeMu.Unlock()
|
||
}
|
||
|
||
func unregisterChild(c *exec.Cmd) {
|
||
activeMu.Lock()
|
||
for i, ch := range activeChildren {
|
||
if ch == c {
|
||
activeChildren = append(activeChildren[:i], activeChildren[i+1:]...)
|
||
break
|
||
}
|
||
}
|
||
activeMu.Unlock()
|
||
}
|
||
|
||
func KillAllChildProcesses() {
|
||
activeMu.Lock()
|
||
all := make([]*exec.Cmd, len(activeChildren))
|
||
copy(all, activeChildren)
|
||
activeChildren = nil
|
||
activeMu.Unlock()
|
||
for _, c := range all {
|
||
killProcessGroup(c)
|
||
}
|
||
}
|
||
|
||
// ─── Spawn ────────────────────────────────────────────────────────────────
|
||
|
||
func spawnChild(cmdStr string, args []string, opts *RunOptions, maxModeFn func(scriptPath, configDir string)) *exec.Cmd {
|
||
envSrc := env.OsEnvToMap()
|
||
resolved := env.ResolveAgentCommand(cmdStr, args, envSrc, opts.Cwd)
|
||
|
||
if opts.MaxMode && maxModeFn != nil {
|
||
maxModeFn(resolved.AgentScriptPath, opts.ConfigDir)
|
||
}
|
||
|
||
envMap := make(map[string]string, len(resolved.Env))
|
||
for k, v := range resolved.Env {
|
||
envMap[k] = v
|
||
}
|
||
if opts.ConfigDir != "" {
|
||
envMap["CURSOR_CONFIG_DIR"] = opts.ConfigDir
|
||
} else if resolved.ConfigDir != "" {
|
||
if _, exists := envMap["CURSOR_CONFIG_DIR"]; !exists {
|
||
envMap["CURSOR_CONFIG_DIR"] = resolved.ConfigDir
|
||
}
|
||
}
|
||
|
||
envSlice := make([]string, 0, len(envMap))
|
||
for k, v := range envMap {
|
||
envSlice = append(envSlice, k+"="+v)
|
||
}
|
||
|
||
ctx := opts.Ctx
|
||
if ctx == nil {
|
||
ctx = context.Background()
|
||
}
|
||
|
||
// 使用 WaitDelay 確保 context cancel 後子程序 goroutine 能及時退出
|
||
c := exec.CommandContext(ctx, resolved.Command, resolved.Args...)
|
||
c.Dir = opts.Cwd
|
||
c.Env = envSlice
|
||
// 設定新的 process group,使 kill 能傳遞給所有子孫程序
|
||
c.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
|
||
// WaitDelay:context cancel 後額外等待這麼久再強制關閉 pipes
|
||
c.WaitDelay = 5 * time.Second
|
||
// Cancel 函式:殺死整個 process group
|
||
c.Cancel = func() error {
|
||
return killProcessGroup(c)
|
||
}
|
||
return c
|
||
}
|
||
|
||
// MaxModeFn is set by the agent package to avoid import cycle.
|
||
var MaxModeFn func(agentScriptPath, configDir string)
|
||
|
||
func Run(cmdStr string, args []string, opts RunOptions) (RunResult, error) {
|
||
ctx := opts.Ctx
|
||
var cancel context.CancelFunc
|
||
if opts.TimeoutMs > 0 {
|
||
if ctx == nil {
|
||
ctx, cancel = context.WithTimeout(context.Background(), time.Duration(opts.TimeoutMs)*time.Millisecond)
|
||
} else {
|
||
ctx, cancel = context.WithTimeout(ctx, time.Duration(opts.TimeoutMs)*time.Millisecond)
|
||
}
|
||
defer cancel()
|
||
opts.Ctx = ctx
|
||
} else if ctx == nil {
|
||
opts.Ctx = context.Background()
|
||
}
|
||
|
||
c := spawnChild(cmdStr, args, &opts, MaxModeFn)
|
||
var stdoutBuf, stderrBuf strings.Builder
|
||
c.Stdout = &stdoutBuf
|
||
c.Stderr = &stderrBuf
|
||
|
||
if err := c.Start(); err != nil {
|
||
// context 已取消或命令找不到時
|
||
if opts.Ctx != nil && opts.Ctx.Err() != nil {
|
||
return RunResult{Code: -1}, nil
|
||
}
|
||
if strings.Contains(err.Error(), "exec: ") || strings.Contains(err.Error(), "no such file") {
|
||
return RunResult{}, fmt.Errorf("command not found: %s. Install Cursor CLI (agent) or set CURSOR_AGENT_BIN to its path", cmdStr)
|
||
}
|
||
return RunResult{}, err
|
||
}
|
||
registerChild(c)
|
||
defer unregisterChild(c)
|
||
|
||
err := c.Wait()
|
||
code := 0
|
||
if err != nil {
|
||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||
code = exitErr.ExitCode()
|
||
if code == 0 {
|
||
code = -1
|
||
}
|
||
} else {
|
||
// context cancelled or killed — return -1 but no error
|
||
return RunResult{Code: -1, Stdout: stdoutBuf.String(), Stderr: stderrBuf.String()}, nil
|
||
}
|
||
}
|
||
return RunResult{
|
||
Code: code,
|
||
Stdout: stdoutBuf.String(),
|
||
Stderr: stderrBuf.String(),
|
||
}, nil
|
||
}
|
||
|
||
type StreamResult struct {
|
||
Code int
|
||
Stderr string
|
||
}
|
||
|
||
func RunStreaming(cmdStr string, args []string, opts RunStreamingOptions) (StreamResult, error) {
|
||
ctx := opts.Ctx
|
||
var cancel context.CancelFunc
|
||
if opts.TimeoutMs > 0 {
|
||
if ctx == nil {
|
||
ctx, cancel = context.WithTimeout(context.Background(), time.Duration(opts.TimeoutMs)*time.Millisecond)
|
||
} else {
|
||
ctx, cancel = context.WithTimeout(ctx, time.Duration(opts.TimeoutMs)*time.Millisecond)
|
||
}
|
||
defer cancel()
|
||
opts.RunOptions.Ctx = ctx
|
||
} else if opts.RunOptions.Ctx == nil {
|
||
opts.RunOptions.Ctx = context.Background()
|
||
}
|
||
|
||
c := spawnChild(cmdStr, args, &opts.RunOptions, MaxModeFn)
|
||
stdoutPipe, err := c.StdoutPipe()
|
||
if err != nil {
|
||
return StreamResult{}, err
|
||
}
|
||
stderrPipe, err := c.StderrPipe()
|
||
if err != nil {
|
||
return StreamResult{}, err
|
||
}
|
||
|
||
if err := c.Start(); err != nil {
|
||
if strings.Contains(err.Error(), "exec: ") || strings.Contains(err.Error(), "no such file") {
|
||
return StreamResult{}, fmt.Errorf("command not found: %s. Install Cursor CLI (agent) or set CURSOR_AGENT_BIN to its path", cmdStr)
|
||
}
|
||
return StreamResult{}, err
|
||
}
|
||
registerChild(c)
|
||
defer unregisterChild(c)
|
||
|
||
var stderrBuf strings.Builder
|
||
var wg sync.WaitGroup
|
||
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
scanner := bufio.NewScanner(stdoutPipe)
|
||
scanner.Buffer(make([]byte, 10*1024*1024), 10*1024*1024)
|
||
for scanner.Scan() {
|
||
line := scanner.Text()
|
||
if strings.TrimSpace(line) != "" {
|
||
opts.OnLine(line)
|
||
}
|
||
}
|
||
}()
|
||
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
scanner := bufio.NewScanner(stderrPipe)
|
||
scanner.Buffer(make([]byte, 10*1024*1024), 10*1024*1024)
|
||
for scanner.Scan() {
|
||
stderrBuf.WriteString(scanner.Text())
|
||
stderrBuf.WriteString("\n")
|
||
}
|
||
}()
|
||
|
||
wg.Wait()
|
||
err = c.Wait()
|
||
code := 0
|
||
if err != nil {
|
||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||
code = exitErr.ExitCode()
|
||
if code == 0 {
|
||
code = -1
|
||
}
|
||
}
|
||
}
|
||
return StreamResult{Code: code, Stderr: stderrBuf.String()}, nil
|
||
}
|