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

// session_service.go 实现 SessionService。Create 走完整事务（worktree
// acquire + supervisor.Spawn）；Get/List/Terminate 是配套的简单读写。
package acp

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"

	"github.com/yan1h/agent-coding-workflow/internal/audit"
	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	"github.com/yan1h/agent-coding-workflow/internal/mcp"
	"github.com/yan1h/agent-coding-workflow/internal/project"
	"github.com/yan1h/agent-coding-workflow/internal/workspace"
)

// SessionServiceConfig 是 service 的限流 + MCP 配置（从 AcpConfig / MCPConfig 派生）。
type SessionServiceConfig struct {
	MaxActiveGlobal  int
	MaxActivePerUser int
	SystemTokenTTL   time.Duration
	MCPPublicURL     string
}

// ProjectAccess 是 acp 模块对 project 模块的窄接口；adapter 在 app.go 实现。
type ProjectAccess interface {
	GetProjectByWorkspace(ctx context.Context, wsID uuid.UUID) (*project.Project, error)
	GetIssueByNumber(ctx context.Context, projectID uuid.UUID, number int) (*project.Issue, error)
	GetRequirementByNumber(ctx context.Context, projectID uuid.UUID, number int) (*project.Requirement, error)
}

type sessionService struct {
	repo      Repository
	wsSvc     workspace.WorkspaceService
	wtSvc     workspace.WorktreeService
	wsRepo    workspace.Repository
	pa        ProjectAccess
	sup       *Supervisor
	audit     audit.Recorder
	cfg       SessionServiceConfig
	mcpTokens mcp.TokenService // spec §6.1 — system token issuance
}

// NewSessionService 构造 SessionService。
func NewSessionService(repo Repository, wsSvc workspace.WorkspaceService,
	wtSvc workspace.WorktreeService, wsRepo workspace.Repository,
	pa ProjectAccess, sup *Supervisor, rec audit.Recorder,
	cfg SessionServiceConfig, mcpTokens mcp.TokenService) SessionService {
	return &sessionService{
		repo: repo, wsSvc: wsSvc, wtSvc: wtSvc, wsRepo: wsRepo,
		pa: pa, sup: sup, audit: rec, cfg: cfg, mcpTokens: mcpTokens,
	}
}

// ResolveBranch 是 spec §8.1 的纯函数实现：根据 CreateSessionInput 决定
// (branch, isMainWorktree)。允许独立单测，不依赖 service 状态。
//
// 规则（按优先级）：
//  1. in.Branch 显式：直接用；isMain = (== ws.DefaultBranch)
//  2. in.IssueNumber：拉 issue，校验 workspace 关联，命名 "issue/<slug>-<num>"
//  3. in.RequirementNumber：拉 requirement，校验 workspace，命名 "req/<slug>-<num>"
//  4. 都没给：fallback "acp-auto/<8-hex>"，isMain = false
func ResolveBranch(ctx context.Context, ws *workspace.Workspace, sessionID uuid.UUID,
	in CreateSessionInput, projectSlug string, pa ProjectAccess) (string, bool, error) {

	if in.Branch != nil && *in.Branch != "" {
		br := *in.Branch
		return br, br == ws.DefaultBranch, nil
	}

	if in.IssueNumber != nil {
		iss, err := pa.GetIssueByNumber(ctx, ws.ProjectID, *in.IssueNumber)
		if err != nil {
			return "", false, err
		}
		if iss.WorkspaceID != nil && *iss.WorkspaceID != ws.ID {
			return "", false, errs.New(errs.CodeInvalidInput,
				fmt.Sprintf("issue #%d already linked to other workspace", *in.IssueNumber))
		}
		return fmt.Sprintf("issue/%s-%d", projectSlug, iss.Number), false, nil
	}

	if in.RequirementNumber != nil {
		req, err := pa.GetRequirementByNumber(ctx, ws.ProjectID, *in.RequirementNumber)
		if err != nil {
			return "", false, err
		}
		if req.WorkspaceID != nil && *req.WorkspaceID != ws.ID {
			return "", false, errs.New(errs.CodeInvalidInput,
				fmt.Sprintf("requirement #%d already linked to other workspace", *in.RequirementNumber))
		}
		return fmt.Sprintf("req/%s-%d", projectSlug, req.Number), false, nil
	}

	short := strings.ReplaceAll(sessionID.String(), "-", "")
	if len(short) > 8 {
		short = short[:8]
	}
	return fmt.Sprintf("acp-auto/%s", short), false, nil
}

// Create 执行完整 session 创建事务：
//  1. workspace 鉴权（wsSvc.Get）
//  2. agent kind 校验（enabled?）
//  3. 全局 + 单用户限流
//  4. 分支解析（ResolveBranch + project 查询）
//  5. worktree 占用（main → repo.AcquireMainWorktree；
//     子分支 → wtSvc.List 找已存在或 Create + Acquire，holder = "session:<sid>"）
//  6. issueID/requirementID 反查（best-effort）
//  7. INSERT acp_sessions
//  8. audit.Record
//  9. async supervisor.Spawn（失败回滚 worktree）
//  10. async initial_prompt 转发（等握手完成）
func (s *sessionService) Create(ctx context.Context, c Caller, in CreateSessionInput) (*Session, error) {
	// 0. 互斥校验：branch / issue_number / requirement_number 三选一或全空（spec §8.1）。
	// service 层统一拦截，HTTP / MCP 入口共用。
	specified := 0
	if in.Branch != nil && *in.Branch != "" {
		specified++
	}
	if in.IssueNumber != nil {
		specified++
	}
	if in.RequirementNumber != nil {
		specified++
	}
	if specified > 1 {
		return nil, errs.New(errs.CodeInvalidInput,
			"branch / issue_number / requirement_number are mutually exclusive; specify at most one")
	}

	// 1. workspace + authz
	wsCaller := workspace.Caller{UserID: c.UserID, IsAdmin: c.IsAdmin}
	ws, err := s.wsSvc.Get(ctx, wsCaller, in.WorkspaceID)
	if err != nil {
		return nil, err
	}

	// 2. agent kind
	kind, err := s.repo.GetAgentKindByID(ctx, in.AgentKindID)
	if err != nil {
		return nil, err
	}
	if !kind.Enabled {
		return nil, errs.New(errs.CodeAcpAgentKindDisabled, "agent kind disabled")
	}

	// 3. 限流
	nGlobal, _ := s.repo.CountActiveSessions(ctx)
	if int(nGlobal) >= s.cfg.MaxActiveGlobal {
		return nil, errs.New(errs.CodeAcpSessionQuotaExceeded, "global session quota exceeded")
	}
	nUser, _ := s.repo.CountActiveSessionsByUser(ctx, c.UserID)
	if int(nUser) >= s.cfg.MaxActivePerUser {
		return nil, errs.New(errs.CodeAcpSessionQuotaExceeded, "per-user session quota exceeded")
	}

	// 4. 分支解析
	pj, err := s.pa.GetProjectByWorkspace(ctx, ws.ID)
	if err != nil {
		return nil, err
	}
	sessionID := uuid.New()
	branch, isMain, err := ResolveBranch(ctx, ws, sessionID, in, pj.Slug, s.pa)
	if err != nil {
		return nil, err
	}

	// 5. worktree 占用
	var (
		cwdPath    string
		worktreeID *uuid.UUID
	)
	sessCaller := workspace.Caller{UserID: c.UserID, IsAdmin: c.IsAdmin, SessionID: &sessionID}
	if isMain {
		ok, err := s.repo.AcquireMainWorktree(ctx, sessionID, ws.ID)
		if err != nil {
			return nil, err
		}
		if !ok {
			return nil, errs.New(errs.CodeAcpSessionBranchConflict, "main worktree already in use")
		}
		cwdPath = ws.MainPath
	} else {
		wt, err := s.findOrCreateWorktree(ctx, sessCaller, ws.ID, branch)
		if err != nil {
			return nil, err
		}
		if _, err := s.wtSvc.Acquire(ctx, sessCaller, wt.ID); err != nil {
			return nil, err
		}
		cwdPath = wt.Path
		wid := wt.ID
		worktreeID = &wid
	}

	// 6. PM 关联（best-effort：拿不到 ID 不阻塞 session 创建）
	var issueID, reqID *uuid.UUID
	if in.IssueNumber != nil {
		if iss, _ := s.pa.GetIssueByNumber(ctx, ws.ProjectID, *in.IssueNumber); iss != nil {
			id := iss.ID
			issueID = &id
		}
	}
	if in.RequirementNumber != nil {
		if req, _ := s.pa.GetRequirementByNumber(ctx, ws.ProjectID, *in.RequirementNumber); req != nil {
			id := req.ID
			reqID = &id
		}
	}

	// 7. INSERT acp_sessions + system MCP token（spec §6.1 atomic guarantee）
	sess := &Session{
		ID: sessionID, WorkspaceID: ws.ID, ProjectID: ws.ProjectID,
		AgentKindID: kind.ID, UserID: c.UserID,
		IssueID: issueID, RequirementID: reqID,
		Branch: branch, CwdPath: cwdPath, IsMainWorktree: isMain,
		Status: SessionStarting,
	}
	var (
		out      *Session
		tokenRes mcp.IssueResult
	)
	txErr := s.repo.InTx(ctx, func(txCtx context.Context, tx pgx.Tx) error {
		inserted, err := s.repo.WithTx(tx).InsertSession(txCtx, sess)
		if err != nil {
			return err
		}
		out = inserted

		res, err := s.mcpTokens.IssueSystemTokenWithTx(txCtx, tx, mcp.IssueSystemTokenInput{
			UserID:       c.UserID,
			ACPSessionID: inserted.ID,
			ProjectIDs:   []uuid.UUID{ws.ProjectID},
			ExpiresIn:    s.cfg.SystemTokenTTL,
		})
		if err != nil {
			return err
		}
		tokenRes = res
		return nil
	})
	if txErr != nil {
		s.releaseOnFailure(ctx, sess, sessCaller, worktreeID)
		return nil, txErr
	}

	// 8. audit
	s.recordAudit(ctx, c, "acp.session.create", out.ID.String(), map[string]any{
		"agent_kind":     kind.Name,
		"branch":         branch,
		"cwd_path":       cwdPath,
		"issue_id":       issueID,
		"requirement_id": reqID,
	})

	// 9. async spawn（失败时回滚 worktree + revoke MCP token）
	extraEnv := map[string]string{
		"ACW_MCP_TOKEN": tokenRes.Plaintext,
		"ACW_MCP_URL":   s.cfg.MCPPublicURL,
	}
	go func() {
		spawnCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
		defer cancel()
		if _, err := s.sup.Spawn(spawnCtx, out, kind, extraEnv); err != nil {
			// 守卫式标记：Spawn 早期失败（进程未注册、onExit 不会触发）时把
			// starting 收尾为 crashed；handshake 失败路径 onExit 已写终态，此处 no-op。
			if _, merr := s.repo.MarkSessionFailedIfActive(spawnCtx, out.ID, err.Error()); merr != nil {
				s.sup.log.Error("acp.spawn_failed.mark_session",
					"session_id", out.ID, "err", merr.Error())
			}
			_ = s.mcpTokens.RevokeBySession(spawnCtx, out.ID)
			s.releaseOnFailure(spawnCtx, out, sessCaller, worktreeID)
		}
	}()

	// 10. initial prompt：握手成功后异步转发
	if in.InitialPrompt != nil && *in.InitialPrompt != "" {
		go s.sendInitialPrompt(out.ID, *in.InitialPrompt)
	}

	return out, nil
}

// findOrCreateWorktree 在指定 workspace 内查找 branch 同名 worktree，
// 不存在则创建。caller 已通过 wsSvc.Get 完成鉴权。
func (s *sessionService) findOrCreateWorktree(ctx context.Context, c workspace.Caller, wsID uuid.UUID, branch string) (*workspace.Worktree, error) {
	wts, err := s.wtSvc.List(ctx, c, wsID)
	if err != nil {
		return nil, err
	}
	for _, wt := range wts {
		if wt.Branch == branch {
			return wt, nil
		}
	}
	return s.wtSvc.Create(ctx, c, wsID, workspace.CreateWorktreeInput{Branch: branch})
}

// releaseOnFailure 在 INSERT / Spawn 失败时回滚 worktree 占用。幂等且静默。
func (s *sessionService) releaseOnFailure(ctx context.Context, sess *Session, c workspace.Caller, worktreeID *uuid.UUID) {
	if sess.IsMainWorktree {
		_, _ = s.repo.ReleaseMainWorktree(ctx, sess.WorkspaceID, sess.ID)
		return
	}
	if worktreeID != nil {
		_, _ = s.wtSvc.Release(ctx, c, *worktreeID)
	}
}

// sendInitialPrompt 等待 supervisor 完成 handshake（最多 30s），随后转发一条
// session/prompt 给 agent。轮询 200ms 一次拉 relay + agent_session_id；任一缺失
// 继续等待。超时后静默放弃。
func (s *sessionService) sendInitialPrompt(sid uuid.UUID, text string) {
	deadline := time.After(30 * time.Second)
	tick := time.NewTicker(200 * time.Millisecond)
	defer tick.Stop()
	for {
		select {
		case <-deadline:
			return
		case <-tick.C:
		}
		relay := s.sup.GetRelay(sid)
		if relay == nil {
			continue
		}
		sess, _ := s.repo.GetSessionByID(context.Background(), sid)
		if sess == nil || sess.AgentSessionID == nil {
			continue
		}
		idJSON, _ := json.Marshal("client-init")
		params, _ := json.Marshal(map[string]any{
			"sessionId": *sess.AgentSessionID,
			"prompt":    []map[string]any{{"type": "text", "text": text}},
		})
		msg := &Message{
			JSONRPC: "2.0", ID: idJSON, Method: "session/prompt", Params: params,
		}
		_ = relay.SendClient(msg)
		return
	}
}

func (s *sessionService) Get(ctx context.Context, c Caller, id uuid.UUID) (*Session, error) {
	sess, err := s.repo.GetSessionByID(ctx, id)
	if err != nil {
		return nil, err
	}
	if !c.IsAdmin && sess.UserID != c.UserID {
		return nil, errs.New(errs.CodeAcpSessionNotFound, "session not found")
	}
	return sess, nil
}

func (s *sessionService) List(ctx context.Context, c Caller, f SessionListFilter) ([]*Session, error) {
	if f.All && !c.IsAdmin {
		return nil, errs.New(errs.CodeForbidden, "admin only")
	}
	if f.All {
		return s.repo.ListAllSessions(ctx, f.RequirementID)
	}
	return s.repo.ListSessionsByUser(ctx, c.UserID, f.RequirementID)
}

func (s *sessionService) Terminate(ctx context.Context, c Caller, id uuid.UUID) error {
	sess, err := s.Get(ctx, c, id)
	if err != nil {
		return err
	}
	if !sess.Status.IsActive() {
		return nil
	}
	s.sup.Kill(ctx, id, 5*time.Second)
	return nil
}

func (s *sessionService) recordAudit(ctx context.Context, c Caller, action, targetID string, meta map[string]any) {
	if s.audit == nil {
		return
	}
	uid := c.UserID
	_ = s.audit.Record(context.WithoutCancel(ctx), audit.Entry{
		UserID: &uid, Action: action, TargetType: "acp_session", TargetID: targetID, Metadata: meta,
	})
}