agentic-coding-workflow/internal/acp/supervisor.go

// supervisor.go 管理 ACP agent 子进程池。每 session 一个 Process，spawn / kill /
// monitor 三类操作通过 mu 保护 procs map。
//
// 生命周期 invariants（spec §6.8）：
//   - procs 写锁仅在 spawn / Kill / monitorProcess delete 时持有
//   - done channel 严格 close-once（monitor 唯一持有）
//   - KilledByUs atomic 防 race（Kill / monitor 并发）
//   - onExit 不能再 Lock supervisor.mu（避免死锁）
package acp

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"

	"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
	"github.com/yan1h/agent-coding-workflow/internal/audit"
	"github.com/yan1h/agent-coding-workflow/internal/infra/crypto"
	"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
	"github.com/yan1h/agent-coding-workflow/internal/mcp"
	"github.com/yan1h/agent-coding-workflow/internal/workspace"
)

// SupervisorConfig 是从 config.AcpConfig 派生的 supervisor 子集。
type SupervisorConfig struct {
	SpawnTimeout       time.Duration
	KillGrace          time.Duration
	StderrBufferLines  int
	StderrTailBytes    int
	EventMaxPayload    int
	EventTruncateField int
	ShutdownGrace      time.Duration
}

// Supervisor 是 ACP 子进程总管。
type Supervisor struct {
	mu    sync.RWMutex
	procs map[uuid.UUID]*Process

	repo   Repository
	audit  audit.Recorder
	notify *notify.Dispatcher
	wtSvc     workspace.WorktreeService
	wsRepo    workspace.Repository
	mcpTokens mcp.TokenService // spec §6.3 — onExit revokes system token
	crypto    *crypto.Encryptor
	cfg    SupervisorConfig
	log    *slog.Logger
}

// NewSupervisor 构造空 Supervisor。Start 不需要 — supervisor 是被动的（spawn 时
// 起 goroutine，无主循环）。Stop 在 ShutdownAll 中实现。
func NewSupervisor(repo Repository, rec audit.Recorder, disp *notify.Dispatcher,
	wtSvc workspace.WorktreeService, wsRepo workspace.Repository,
	mcpTokens mcp.TokenService,
	enc *crypto.Encryptor, cfg SupervisorConfig, log *slog.Logger) *Supervisor {
	if log == nil {
		log = slog.Default()
	}
	return &Supervisor{
		procs:  map[uuid.UUID]*Process{},
		repo:   repo,
		audit:  rec,
		notify: disp,
		wtSvc:  wtSvc,
		wsRepo: wsRepo,
		mcpTokens: mcpTokens,
		crypto: enc,
		cfg:    cfg,
		log:    log,
	}
}

// Process 是单个 agent 子进程的运行时句柄。
type Process struct {
	SessionID   uuid.UUID
	WorkspaceID uuid.UUID
	UserID      uuid.UUID
	IsMain      bool       // 主 worktree 占用 → release 走不同路径
	WorktreeID  *uuid.UUID // 子 worktree 时填，用于 release

	Cmd    *exec.Cmd
	Stdin  io.WriteCloser
	Stdout io.ReadCloser
	Stderr io.ReadCloser
	Relay  *Relay

	StderrBuf  *stderrRing  // 排障 ring buffer
	StartedAt  time.Time
	KilledByUs atomic.Bool   // true → exited，false → crashed
	done       chan struct{} // monitor 退出时 close
}

// stderrRing 是固定行数 ring buffer，monitor 退出时取 tail 写 last_error。
type stderrRing struct {
	mu    sync.Mutex
	lines []string
	cap   int
}

func newStderrRing(cap int) *stderrRing { return &stderrRing{cap: cap} }

func (r *stderrRing) Append(line string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.lines) >= r.cap {
		r.lines = r.lines[1:]
	}
	r.lines = append(r.lines, line)
}

func (r *stderrRing) Tail(maxBytes int) string {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := ""
	for i := len(r.lines) - 1; i >= 0 && len(out) < maxBytes; i-- {
		out = r.lines[i] + "\n" + out
	}
	if len(out) > maxBytes {
		out = out[len(out)-maxBytes:]
	}
	return out
}

// GetRelay returns the relay for an active session, or nil if not running.
// WS handler 用此判断是否进 live 模式。
func (s *Supervisor) GetRelay(sid uuid.UUID) *Relay {
	s.mu.RLock()
	defer s.mu.RUnlock()
	p, ok := s.procs[sid]
	if !ok {
		return nil
	}
	return p.Relay
}

// Spawn 启动一个子进程并完成 ACP handshake。返回的 *Process 已注册到 procs map；
// handshake 失败时调用 Kill 回滚并返回错误。调用方（SessionService）负责前置的
// worktree 占用与 acp_sessions 行 INSERT；本方法仅负责 OS 层资源与协议层握手。
func (s *Supervisor) Spawn(ctx context.Context, sess *Session, kind *AgentKind, extraEnv map[string]string) (*Process, error) {
	envMap, err := decryptEnv(s.crypto, kind.EncryptedEnv)
	if err != nil {
		return nil, err
	}

	// System override: extraEnv overrides agent_kinds same-name keys (spec §6.1 #11)
	for k, v := range extraEnv {
		envMap[k] = v
	}

	// strip env：只保留 PATH + 解密 overrides，避免泄漏 master key、git 凭据等。
	osEnv := []string{"PATH=" + os.Getenv("PATH")}
	for k, v := range envMap {
		osEnv = append(osEnv, fmt.Sprintf("%s=%s", k, v))
	}

	cmd := exec.Command(kind.BinaryPath, kind.Args...)
	cmd.Dir = sess.CwdPath
	cmd.Env = osEnv

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("stdin pipe: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("stdout pipe: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return nil, fmt.Errorf("stderr pipe: %w", err)
	}

	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("start: %w", err)
	}

	relay := NewRelay(
		sess.ID,
		NewDecoder(stdout, s.cfg.EventMaxPayload),
		NewEncoder(stdin),
		s.repo,
		handlers.NewFsHandler(),
		handlers.NewPermissionHandler(),
		RelayConfig{
			EventMaxPayload:    s.cfg.EventMaxPayload,
			EventTruncateField: s.cfg.EventTruncateField,
			WSSendBuffer:       256,
		},
		s.log,
	)

	proc := &Process{
		SessionID:   sess.ID,
		WorkspaceID: sess.WorkspaceID,
		UserID:      sess.UserID,
		IsMain:      sess.IsMainWorktree,
		Cmd:         cmd,
		Stdin:       stdin,
		Stdout:      stdout,
		Stderr:      stderr,
		Relay:       relay,
		StderrBuf:   newStderrRing(s.cfg.StderrBufferLines),
		StartedAt:   time.Now(),
		done:        make(chan struct{}),
	}

	s.mu.Lock()
	s.procs[sess.ID] = proc
	s.mu.Unlock()

	go s.drainStderr(proc)
	go s.monitorProcess(proc)
	go relay.Run(ctx, handlers.SessionContext{
		SessionID:      sess.ID,
		WorkspaceID:    sess.WorkspaceID,
		UserID:         sess.UserID,
		CwdPath:        sess.CwdPath,
		AgentSessionID: sess.AgentSessionID,
	})

	if err := s.handshake(ctx, sess, relay, proc); err != nil {
		s.log.Error("acp.handshake_failed", "session_id", sess.ID, "err", err.Error())
		s.Kill(ctx, sess.ID, s.cfg.KillGrace)
		return nil, err
	}
	return proc, nil
}

// handshake 发送 initialize → session/new，并把返回的 agent session id + pid
// 持久化到 acp_sessions 行（status → running）。SpawnTimeout 是整个握手的硬上限。
func (s *Supervisor) handshake(ctx context.Context, sess *Session, relay *Relay, proc *Process) error {
	hsCtx, cancel := context.WithTimeout(ctx, s.cfg.SpawnTimeout)
	defer cancel()

	if _, err := relay.Call(hsCtx, "initialize", handlers.BuildInitializeParams("dev")); err != nil {
		return fmt.Errorf("initialize: %w", err)
	}
	resp, err := relay.Call(hsCtx, "session/new", handlers.BuildSessionNewParams(sess.CwdPath))
	if err != nil {
		return fmt.Errorf("session/new: %w", err)
	}
	var snr handlers.SessionNewResult
	if err := json.Unmarshal(resp.Result, &snr); err != nil {
		return fmt.Errorf("parse session/new result: %w", err)
	}
	if snr.SessionID == "" {
		return fmt.Errorf("agent did not return sessionId")
	}
	pid := int32(proc.Cmd.Process.Pid)
	if err := s.repo.UpdateSessionRunning(ctx, sess.ID, snr.SessionID, pid); err != nil {
		return fmt.Errorf("update session running: %w", err)
	}
	sess.AgentSessionID = &snr.SessionID
	sess.PID = &pid
	sess.Status = SessionRunning
	return nil
}

// BuildSpawnEnv exposes env merge logic for testing.
func (s *Supervisor) BuildSpawnEnv(kind *AgentKind, extraEnv map[string]string) ([]string, error) {
	envMap, err := decryptEnv(s.crypto, kind.EncryptedEnv)
	if err != nil {
		return nil, err
	}
	for k, v := range extraEnv {
		envMap[k] = v
	}
	out := []string{"PATH=" + os.Getenv("PATH")}
	for k, v := range envMap {
		out = append(out, fmt.Sprintf("%s=%s", k, v))
	}
	return out, nil
}

// Kill 终止指定 session 的子进程。在 Unix 平台上先发 SIGTERM 等待 grace 时间，
// 仍未退出再 Process.Kill；Windows 没有可靠 SIGTERM 实现，直接 Process.Kill。
// 始终阻塞直到 monitorProcess 关闭 done channel。Kill 是幂等的——目标 session
// 不存在时立即返回。
func (s *Supervisor) Kill(ctx context.Context, sid uuid.UUID, grace time.Duration) {
	s.mu.RLock()
	proc, ok := s.procs[sid]
	s.mu.RUnlock()
	if !ok {
		return
	}

	proc.KilledByUs.Store(true)

	// 平台差异：Unix 走 SIGTERM + grace + 强 Kill；Windows 直接 Kill。
	// 实现见 supervisor_unix.go / supervisor_windows.go。
	s.unixGracefulTerm(proc, grace)

	select {
	case <-proc.done:
	case <-ctx.Done():
	}
}

// ShutdownAll 并行 Kill 所有正在运行的子进程。bounded by cfg.ShutdownGrace。
func (s *Supervisor) ShutdownAll(ctx context.Context) {
	s.mu.RLock()
	ids := make([]uuid.UUID, 0, len(s.procs))
	for id := range s.procs {
		ids = append(ids, id)
	}
	s.mu.RUnlock()

	shutdownCtx, cancel := context.WithTimeout(ctx, s.cfg.ShutdownGrace)
	defer cancel()

	var wg sync.WaitGroup
	for _, id := range ids {
		wg.Add(1)
		go func(sid uuid.UUID) {
			defer wg.Done()
			s.Kill(shutdownCtx, sid, s.cfg.KillGrace)
		}(id)
	}

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()
	select {
	case <-done:
	case <-shutdownCtx.Done():
		s.log.Warn("acp.shutdown_all.timeout")
	}
}

// drainStderr 持续读取子进程 stderr，按行追加到 ring buffer。EOF 或 read error
// 时安静退出（cmd.Wait() 已在 monitorProcess 处理）。
func (s *Supervisor) drainStderr(proc *Process) {
	scanner := bufio.NewScanner(proc.Stderr)
	// 默认 token 大小 64KB，足够 ACP agent 单行 log。
	for scanner.Scan() {
		proc.StderrBuf.Append(scanner.Text())
	}
}

// monitorProcess 等待子进程退出，更新 procs map 并触发 onExit。这是子进程
// 状态唯一的合法 close 路径（done channel）。
func (s *Supervisor) monitorProcess(proc *Process) {
	defer close(proc.done)
	waitErr := proc.Cmd.Wait()

	exitCode := int32(0)
	if exitErr, ok := waitErr.(*exec.ExitError); ok {
		exitCode = int32(exitErr.ExitCode())
	}

	s.mu.Lock()
	delete(s.procs, proc.SessionID)
	s.mu.Unlock()

	status := SessionExited
	if !proc.KilledByUs.Load() {
		status = SessionCrashed
	}
	s.onExit(context.Background(), proc, status, exitCode, waitErr)
}

// onExit 写终止状态、释放 worktree、触发通知与 audit、最后关闭 relay。
// 不能再 acquire s.mu（monitorProcess 已 delete map 后调用本函数，但 Kill
// 也持锁等 done，会形成 monitor → Kill → onExit 的链式 deadlock）。
func (s *Supervisor) onExit(ctx context.Context, proc *Process, status SessionStatus, exitCode int32, waitErr error) {
	stderrTail := proc.StderrBuf.Tail(s.cfg.StderrTailBytes)
	var lastErr *string
	if status == SessionCrashed && stderrTail != "" {
		lastErr = &stderrTail
	}
	if err := s.repo.UpdateSessionFinished(ctx, proc.SessionID, status, &exitCode, lastErr); err != nil {
		s.log.Error("acp.on_exit.update_session", "session_id", proc.SessionID, "err", err.Error())
	}

	// Release worktree。主 worktree 走 acp 表自己的 CAS；子 worktree 走 workspace
	// service 的 holder 校验。当前 wtSvc.Release 接受 workspace.Caller，子 worktree
	// 释放需要绕过 holder 检查（spec 期望 holder = "session:<sid>"，与 user holder
	// 不同），用 IsAdmin: true 强制释放。F4 reaper 会替换为更严格的 byHolder 调用。
	if proc.IsMain {
		if _, err := s.repo.ReleaseMainWorktree(ctx, proc.WorkspaceID, proc.SessionID); err != nil {
			s.log.Error("acp.on_exit.release_main",
				"session_id", proc.SessionID, "err", err.Error())
		}
	} else if proc.WorktreeID != nil && s.wtSvc != nil {
		if _, err := s.wtSvc.Release(ctx, workspace.Caller{IsAdmin: true}, *proc.WorktreeID); err != nil {
			s.log.Error("acp.on_exit.release_worktree",
				"session_id", proc.SessionID, "worktree_id", *proc.WorktreeID, "err", err.Error())
		}
	}

	if status == SessionCrashed {
		if s.notify != nil {
			_ = s.notify.Dispatch(ctx, notify.Message{
				UserID:   proc.UserID,
				Topic:    "acp.session_crashed",
				Severity: notify.SeverityError,
				Title:    "ACP session crashed",
				Body:     fmt.Sprintf("Session %s exited with code %d.", proc.SessionID, exitCode),
				Metadata: map[string]any{
					"session_id": proc.SessionID.String(),
					"exit_code":  exitCode,
				},
			})
		}
		if s.audit != nil {
			uid := proc.UserID
			_ = s.audit.Record(ctx, audit.Entry{
				UserID:     &uid,
				Action:     "acp.session.crashed",
				TargetType: "acp_session",
				TargetID:   proc.SessionID.String(),
				Metadata: map[string]any{
					"exit_code":  exitCode,
					"last_error": stderrTail,
				},
			})
		}
	} else {
		if s.audit != nil {
			uid := proc.UserID
			_ = s.audit.Record(ctx, audit.Entry{
				UserID:     &uid,
				Action:     "acp.session.terminate",
				TargetType: "acp_session",
				TargetID:   proc.SessionID.String(),
				Metadata:   map[string]any{"exit_code": exitCode},
			})
		}
	}
	_ = waitErr // 已通过 KilledByUs / ExitError 编码到 status / exitCode
	// Revoke session's system MCP tokens (spec §6.3)
	if s.mcpTokens != nil {
		if err := s.mcpTokens.RevokeBySession(ctx, proc.SessionID); err != nil {
			s.log.Error("acp.on_exit.mcp_revoke",
				"session_id", proc.SessionID, "err", err.Error())
		}
	}

	// 最后关 relay：广播 session_terminated 给 WS subscribers。
	proc.Relay.Close(string(status), &exitCode)
}