agentic-coding-workflow/internal/jobs/runners/acp_session_reaper.go

package runners

import (
	"context"
	"log/slog"
	"time"

	"github.com/google/uuid"

	"github.com/yan1h/agent-coding-workflow/internal/infra/notify"
)

// ReaperSession is the minimal session projection the reaper needs.
type ReaperSession struct {
	ID                  uuid.UUID
	UserID              uuid.UUID
	StartedAt           time.Time
	LastActivityAt      time.Time
	MaxWallClockSeconds *int32
}

// AcpSessionReaperRepo is the narrow repo surface this runner needs.
// internal/acp.Repository satisfies it via an adapter in app.go.
type AcpSessionReaperRepo interface {
	// ListSessionsForReaper returns active sessions exceeding their wall-clock
	// cap (per-session or the global wallClockStartedBefore cutoff) or idle since
	// idleSince. A zero wallClockStartedBefore disables the global wall-clock cut.
	ListSessionsForReaper(ctx context.Context, wallClockStartedBefore, idleSince time.Time) ([]ReaperSession, error)
	// MarkSessionTerminatedReason records why a session was reaped (no-op if a
	// reason is already set, protecting against clobbering a budget-gate reason).
	MarkSessionTerminatedReason(ctx context.Context, id uuid.UUID, reason string) error
}

// SessionKiller terminates a session subprocess. internal/acp.Supervisor.Kill
// satisfies it directly.
type SessionKiller interface {
	Kill(ctx context.Context, sid uuid.UUID, grace time.Duration)
}

// AcpSessionReaperConfig controls the reaper's cutoffs.
type AcpSessionReaperConfig struct {
	// WallClockTimeout is the global maximum session age (0 = rely only on the
	// per-session budget_max_wall_clock_seconds).
	WallClockTimeout time.Duration
	// IdleTimeout is the max time with no new activity before reaping.
	IdleTimeout time.Duration
	// KillGrace is passed to SessionKiller.Kill.
	KillGrace time.Duration
}

// AcpSessionReaper terminates ACP sessions that exceed wall-clock or are idle.
type AcpSessionReaper struct {
	repo   AcpSessionReaperRepo
	killer SessionKiller
	disp   *notify.Dispatcher
	cfg    AcpSessionReaperConfig
	log    *slog.Logger
}

// NewAcpSessionReaper builds the reaper.
func NewAcpSessionReaper(repo AcpSessionReaperRepo, killer SessionKiller, disp *notify.Dispatcher,
	cfg AcpSessionReaperConfig, log *slog.Logger) *AcpSessionReaper {
	if log == nil {
		log = slog.Default()
	}
	if cfg.KillGrace <= 0 {
		cfg.KillGrace = 5 * time.Second
	}
	return &AcpSessionReaper{repo: repo, killer: killer, disp: disp, cfg: cfg, log: log}
}

// Run reaps once. Periodic; idempotent (terminated_reason is only set once).
func (r *AcpSessionReaper) Run(ctx context.Context) {
	now := time.Now()
	var wallBefore time.Time
	if r.cfg.WallClockTimeout > 0 {
		wallBefore = now.Add(-r.cfg.WallClockTimeout)
	}
	idleSince := now.Add(-r.cfg.IdleTimeout)
	if r.cfg.IdleTimeout <= 0 {
		// No idle timeout configured -> make the idle cutoff unreachable so only
		// wall-clock rules apply (a zero time would reap everything).
		idleSince = time.Time{}
	}

	sessions, err := r.repo.ListSessionsForReaper(ctx, wallBefore, idleSince)
	if err != nil {
		r.log.Error("acp_session_reaper.list", "err", err.Error())
		return
	}
	for _, s := range sessions {
		reason := r.classify(now, s)
		if reason == "" {
			continue
		}
		if merr := r.repo.MarkSessionTerminatedReason(ctx, s.ID, reason); merr != nil {
			r.log.Error("acp_session_reaper.mark", "session_id", s.ID, "err", merr.Error())
			// continue: still attempt the kill so the resource is freed.
		}
		// Kill must not be fatal on error (idempotent; session may have already
		// exited between listing and now).
		r.killer.Kill(ctx, s.ID, r.cfg.KillGrace)
		r.log.Info("acp_session_reaper.killed", "session_id", s.ID, "reason", reason)
		r.dispatch(ctx, s, reason)
	}
}

// classify decides which reaper reason applies, preferring wall-clock over idle.
func (r *AcpSessionReaper) classify(now time.Time, s ReaperSession) string {
	// Per-session wall-clock cap.
	if s.MaxWallClockSeconds != nil {
		if now.Sub(s.StartedAt) >= time.Duration(*s.MaxWallClockSeconds)*time.Second {
			return "reaper_wall_clock"
		}
	}
	// Global wall-clock cap.
	if r.cfg.WallClockTimeout > 0 && now.Sub(s.StartedAt) >= r.cfg.WallClockTimeout {
		return "reaper_wall_clock"
	}
	// Idle cap.
	if r.cfg.IdleTimeout > 0 {
		activity := s.LastActivityAt
		if activity.IsZero() {
			activity = s.StartedAt
		}
		if now.Sub(activity) >= r.cfg.IdleTimeout {
			return "reaper_idle"
		}
	}
	return ""
}

func (r *AcpSessionReaper) dispatch(ctx context.Context, s ReaperSession, reason string) {
	if r.disp == nil {
		return
	}
	_ = r.disp.Dispatch(ctx, notify.Message{
		UserID:   s.UserID,
		Topic:    "acp.session_reaped",
		Severity: notify.SeverityWarning,
		Title:    "ACP session terminated by reaper",
		Body:     "An ACP session was terminated: " + reason + ".",
		Metadata: map[string]any{
			"session_id": s.ID.String(),
			"reason":     reason,
		},
	})
}