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"
	procgrp "github.com/yan1h/agent-coding-workflow/internal/infra/proc"
	"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
	// AgentHomesDir 是 agent CLI 受管 home 的根目录（<DataDir>/agent-homes）。
	// 空串时跳过配置物化与重定向 env 注入（部分测试）。
	AgentHomesDir string
}

// 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
	permSvc   *PermissionService // 工具调用权限审批；可为 nil（部分测试）
	cfg       SupervisorConfig
	log       *slog.Logger
}

// SetPermissionService 注入权限审批服务。与 Supervisor 互相依赖，故装配时回填
// （permSvc 需要 Supervisor 作为 RelayLocator，Supervisor 需要 permSvc 注入 handler）。
func (s *Supervisor) SetPermissionService(p *PermissionService) { s.permSvc = p }

// 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 走不同路径

	Cmd    *exec.Cmd
	group  procgrp.Group // 进程树句柄（整树终止），Spawn 时 Prepare+Adopt
	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

	// relayCancel 取消 relay 的会话级 context。该 context 由 Spawn 用
	// context.Background() 派生，独立于调用方的 spawn ctx（spawn ctx 在 Spawn
	// 返回后立即取消），故 relay 的 reader/writer/persist/permission Decide 在整个
	// 会话生命周期内都使用未取消的 ctx。onExit 调用 relayCancel 完成 relay teardown。
	relayCancel context.CancelFunc
}

// 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) {
	// 受管 home：物化 DB 中的配置文件并注入重定向 env（优先级最低，可被
	// AgentKind env / extraEnv 覆盖）。AgentHomesDir 未配置时跳过（部分测试）。
	envMap := map[string]string{}
	if s.cfg.AgentHomesDir != "" {
		home, err := s.materializeAgentHome(ctx, kind)
		if err != nil {
			return nil, fmt.Errorf("materialize agent home: %w", err)
		}
		envMap = agentHomeEnv(home, kind.ClientType)
	}

	kindEnv, err := decryptEnv(s.crypto, kind.EncryptedEnv)
	if err != nil {
		return nil, err
	}
	for k, v := range kindEnv {
		envMap[k] = v
	}

	// 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

	// 进程树管理：Prepare 须在 Start 前（Unix Setpgid），Adopt 须在 Start 后
	// （Windows Job Object）。终止时整树清理，避免 agent 派生的子孙进程残留。
	group := procgrp.NewGroup()
	group.Prepare(cmd)

	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)
	}
	if err := group.Adopt(cmd); err != nil {
		_ = cmd.Process.Kill()
		_ = cmd.Wait() // 回收进程资源，避免 Unix 僵尸进程
		return nil, fmt.Errorf("adopt process group: %w", err)
	}

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

	// relay 的会话级 context：独立于调用方的 spawn ctx（生产调用方在 Spawn 返回后
	// 立即 cancel spawn ctx）。relay 的 reader/writer/persist/permission Decide 必须
	// 在整个会话生命周期内使用未取消的 ctx，否则 prompt 被丢、事件停止持久化、权限
	// 审批被强制拒绝。teardown 由 onExit 调 relayCancel + Relay.Close（关 r.done）完成。
	relayCtx, relayCancel := context.WithCancel(context.Background())

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

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

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

	// 握手仍走调用方的 spawn ctx，保留 SpawnTimeout 截止语义。
	if err := s.handshake(ctx, sess, kind, relay, proc, extraEnv); 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 是整个握手的硬上限。
// session/new 同时下发 mcpServers：ACW 自家 MCP（extraEnv 中的 per-session
// token）+ AgentKind 第三方 MCP，按 initialize 声明的 mcpCapabilities 过滤。
func (s *Supervisor) handshake(ctx context.Context, sess *Session, kind *AgentKind,
	relay *Relay, proc *Process, extraEnv map[string]string) error {
	hsCtx, cancel := context.WithTimeout(ctx, s.cfg.SpawnTimeout)
	defer cancel()

	initResp, err := relay.Call(hsCtx, "initialize", handlers.BuildInitializeParams("dev"))
	if err != nil {
		return fmt.Errorf("initialize: %w", err)
	}
	var ir handlers.InitializeResult
	if err := json.Unmarshal(initResp.Result, &ir); err != nil {
		// 能力解析失败不阻断握手：按无 http/sse 能力处理（仅基线 stdio 可下发）。
		s.log.Warn("acp.handshake.parse_initialize_result", "session_id", sess.ID, "err", err.Error())
	}
	httpOK, sseOK := ir.McpCapabilities()

	specs, err := decryptMCPServers(s.crypto, kind.EncryptedMCPServers)
	if err != nil {
		return err
	}
	mcpServers := buildMcpServers(specs, extraEnv[EnvMCPURL], extraEnv[EnvMCPToken],
		httpOK, sseOK, s.log, sess.ID)

	resp, err := relay.Call(hsCtx, "session/new", handlers.BuildSessionNewParams(sess.CwdPath, mcpServers))
	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(-pgid) + grace + SIGKILL(-pgid)；Windows 关闭
	// Job 句柄触发 KILL_ON_JOB_CLOSE。实现下沉至 internal/infra/proc。
	proc.group.TerminateGroup(proc.done, 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()

	// 释放进程树句柄（Windows 关闭 Job handle；Unix no-op）。done 已 close、
	// 进程已 Wait 返回，此处 Close 不影响 onExit-must-not-relock 不变量。
	proc.group.Close()

	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 按 holder
	// 释放（holder = "session:<sid>"，见 workspace.Caller.Holder），与启动 reaper
	// 一致，不需要知道 worktree ID。
	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 s.wtSvc != nil {
		if _, err := s.wtSvc.ReleaseByHolder(ctx, "session:"+proc.SessionID.String()); err != nil {
			s.log.Error("acp.on_exit.release_worktree",
				"session_id", proc.SessionID, "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

	// 终止会话的全部挂起权限请求：标 expired 并唤醒阻塞的 Decide goroutine，
	// 避免 agent 端永久挂起。须在 Relay.Close 之前（Close 后 FanoutControl 静默丢弃）。
	if s.permSvc != nil {
		s.permSvc.CancelSession(ctx, proc.SessionID)
	}

	// 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 的会话级 context：停止 reader/writer，并让仍阻塞在 Decide / 已派生
	// 的 handleAgentRequest goroutine 收到 ctx.Done 后退出。在 Relay.Close 之前调，
	// 二者共同完成 teardown（relayCancel 停 ctx 派生的 goroutine，Close 关 r.done +
	// 广播 session_terminated）。relayCancel 在早期 spawn 失败路径也可能为 nil。
	if proc.relayCancel != nil {
		proc.relayCancel()
	}

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