agentic-coding-workflow/internal/mcp/token_service.go

package mcp

import (
	"context"
	"strings"
	"time"

	"github.com/google/uuid"

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

// TokenService 是 MCP token 的业务编排层。三类入口：
//   - IssueAdmin：admin 后台手动签发，必须给 expires_at
//   - IssueSystemToken：ACP supervisor 起 session 时调，绑定 acp_session_id
//   - Authenticate：HTTP 中间件调，验 plaintext + 返回 user_id + scope
//
// 撤销与列表通过 Repository 直接调即可，不在 service 上加薄壳。
type TokenService interface {
	IssueAdmin(ctx context.Context, in IssueAdminInput) (IssueResult, error)
	IssueSystemToken(ctx context.Context, in IssueSystemTokenInput) (IssueResult, error)
	Authenticate(ctx context.Context, plaintext string) (*Token, error)
	Revoke(ctx context.Context, id, byUserID uuid.UUID) error
	RevokeBySession(ctx context.Context, sessionID uuid.UUID) error
}

// IssueAdminInput 是 admin 签发时的入参。UserID 缺省 = 当前 admin 自己（由调用方填）。
// Name 必填；ExpiresAt 必填（spec §10 决策 #14：防长期泄露）。
type IssueAdminInput struct {
	UserID    uuid.UUID
	Name      string
	Scope     Scope
	ExpiresAt time.Time
	CreatedBy uuid.UUID // 当前 admin 用户 ID
}

// IssueSystemTokenInput 是 ACP session 启动时的入参。
type IssueSystemTokenInput struct {
	UserID       uuid.UUID
	ACPSessionID uuid.UUID
	ProjectIDs   []uuid.UUID
	Tools        []string      // 缺省 nil = 系统默认（不含 chat）
	ExpiresIn    time.Duration // 默认 24h
}

// IssueResult 是签发返回。Plaintext 仅此一次返回；调用方须立即转交给客户端。
type IssueResult struct {
	ID        uuid.UUID
	Plaintext string
	ExpiresAt *time.Time
}

// defaultSystemTools 是 system token 的默认工具白名单（spec §6.2）。
// 不含 list_recent_messages：agent 不该看用户私人对话历史。
var defaultSystemTools = []string{
	"list_projects", "list_workspaces", "list_requirements", "list_issues",
	"get_issue", "get_requirement",
	"create_issue", "update_issue", "close_issue", "reopen_issue", "assign_issue",
	"create_requirement", "update_requirement", "close_requirement", "reopen_requirement", "set_requirement_phase",
}

// NewTokenService 构造 service。now 用于测试注入；nil 时默认 time.Now。
func NewTokenService(repo Repository, audit audit.Recorder, now func() time.Time) TokenService {
	if now == nil {
		now = time.Now
	}
	return &tokenService{repo: repo, audit: audit, now: now}
}

type tokenService struct {
	repo  Repository
	audit audit.Recorder
	now   func() time.Time
}

// IssueAdmin 签发一个长期 admin token。Name + ExpiresAt 必填；scope 可空（= 全权）。
// 写一条 audit `mcp.token.create`。
func (s *tokenService) IssueAdmin(ctx context.Context, in IssueAdminInput) (IssueResult, error) {
	if strings.TrimSpace(in.Name) == "" {
		return IssueResult{}, errs.New(errs.CodeMcpTokenNameRequired, "name is required")
	}
	if in.ExpiresAt.IsZero() {
		return IssueResult{}, errs.New(errs.CodeMcpTokenExpiresRequired, "expires_at is required for admin token")
	}
	if !in.ExpiresAt.After(s.now()) {
		return IssueResult{}, errs.New(errs.CodeMcpTokenExpiresRequired, "expires_at must be in the future")
	}

	plain, hash, err := NewToken()
	if err != nil {
		return IssueResult{}, errs.Wrap(err, errs.CodeInternal, "generate token plaintext")
	}

	id := uuid.New()
	created, err := s.repo.Create(ctx, &Token{
		ID:        id,
		TokenHash: hash,
		UserID:    in.UserID,
		Name:      in.Name,
		Issuer:    IssuerAdmin,
		Scope:     in.Scope,
		ExpiresAt: &in.ExpiresAt,
		CreatedBy: in.CreatedBy,
	})
	if err != nil {
		return IssueResult{}, err
	}

	_ = s.audit.Record(ctx, audit.Entry{
		UserID:     &in.CreatedBy,
		Action:     "mcp.token.create",
		TargetType: "mcp_token",
		TargetID:   created.ID.String(),
		Metadata: map[string]any{
			"token_id":   created.ID.String(),
			"user_id":    in.UserID.String(),
			"name":       in.Name,
			"scope":      in.Scope,
			"expires_at": in.ExpiresAt,
		},
	})

	return IssueResult{
		ID:        created.ID,
		Plaintext: plain,
		ExpiresAt: created.ExpiresAt,
	}, nil
}

// IssueSystemToken 在 ACP supervisor 起 session 时调用。
//
// 重要约束：调用必须在已经创建 acp_sessions row 的事务**之后或同一事务内**进行 ——
// 因为 mcp_tokens.acp_session_id 是 NOT NULL CHECK + FK，session row 不存在则
// FK 校验失败。spec §6.1 明确建议同事务提交。本 service 不直接管事务边界（由
// ACP SessionService 在 Create 流程内编排），但调用顺序由调用方保证。
//
// TODO(Part 3): 当前 pgRepo 持有 *pgxpool.Pool 并忽略 ctx 中可能存在的 pgx tx，
// 因此严格意义上 IssueSystemToken 与 acp_sessions row 的写入并不在同一事务里 ——
// 调用顺序虽然能避开 FK 错误，但不是真正的原子提交。要兑现 spec §6.1 同事务
// 保证，Part 3 需要为 Repository 增加 `WithTx(tx pgx.Tx) Repository` 构造，
// 并在 acp.SessionService.Create 中以同一事务执行 acp_sessions 插入与本次签发。
//
// 默认 ExpiresIn = 24h；默认 Tools 为内置 PM 工具白名单（spec §6.2）。
// 写一条 audit `mcp.token.create_system`。
func (s *tokenService) IssueSystemToken(ctx context.Context, in IssueSystemTokenInput) (IssueResult, error) {
	if in.ACPSessionID == uuid.Nil {
		return IssueResult{}, errs.New(errs.CodeInvalidInput, "acp_session_id is required for system token")
	}
	if in.UserID == uuid.Nil {
		return IssueResult{}, errs.New(errs.CodeInvalidInput, "user_id is required for system token")
	}
	if in.ExpiresIn <= 0 {
		in.ExpiresIn = 24 * time.Hour
	}
	tools := in.Tools
	if tools == nil {
		// 用 copy 避免共享 defaultSystemTools 给调用方修改
		tools = append([]string{}, defaultSystemTools...)
	}
	expires := s.now().Add(in.ExpiresIn)

	plain, hash, err := NewToken()
	if err != nil {
		return IssueResult{}, errs.Wrap(err, errs.CodeInternal, "generate system token")
	}

	id := uuid.New()
	created, err := s.repo.Create(ctx, &Token{
		ID:           id,
		TokenHash:    hash,
		UserID:       in.UserID,
		Name:         "acp:" + in.ACPSessionID.String(),
		Issuer:       IssuerSystem,
		Scope:        Scope{Tools: tools, ProjectIDs: in.ProjectIDs},
		ACPSessionID: &in.ACPSessionID,
		ExpiresAt:    &expires,
		CreatedBy:    in.UserID, // system 签发时 created_by 即 owner（无 admin 操作者）
	})
	if err != nil {
		return IssueResult{}, err
	}

	_ = s.audit.Record(ctx, audit.Entry{
		UserID:     &in.UserID,
		Action:     "mcp.token.create_system",
		TargetType: "mcp_token",
		TargetID:   created.ID.String(),
		Metadata: map[string]any{
			"token_id":       created.ID.String(),
			"acp_session_id": in.ACPSessionID.String(),
			"user_id":        in.UserID.String(),
		},
	})

	return IssueResult{
		ID:        created.ID,
		Plaintext: plain,
		ExpiresAt: created.ExpiresAt,
	}, nil
}
// Authenticate 校验 plaintext token：
//  1. 计算 hash 反查 row
//  2. 检查 revoked_at IS NULL
//  3. 检查 expires_at IS NULL OR > now()
//  4. 异步触发 last_used_at 更新（不阻塞调用方；使用独立后台 ctx）
//
// 失败映射：
//   - row 不存在        → CodeMcpTokenInvalid (401)
//   - revoked_at 非空   → CodeMcpTokenRevoked (401)
//   - expires_at 已过   → CodeMcpTokenExpired (401)
//
// 返回的 *Token 不含 plaintext；调用方从中读 user_id + scope。
func (s *tokenService) Authenticate(ctx context.Context, plaintext string) (*Token, error) {
	if strings.TrimSpace(plaintext) == "" {
		return nil, errs.New(errs.CodeMcpTokenInvalid, "missing token")
	}
	tok, err := s.repo.GetByHash(ctx, HashToken(plaintext))
	if err != nil {
		return nil, err // 已在 repo 层映射为 CodeMcpTokenInvalid
	}
	if tok.RevokedAt != nil {
		return nil, errs.New(errs.CodeMcpTokenRevoked, "token revoked")
	}
	if tok.ExpiresAt != nil && !tok.ExpiresAt.After(s.now()) {
		return nil, errs.New(errs.CodeMcpTokenExpired, "token expired")
	}

	// 异步更新 last_used_at，不阻塞调用方；使用独立 ctx 防 request 取消
	go func(id uuid.UUID) {
		bg, cancel := context.WithTimeout(context.Background(), 2*time.Second)
		defer cancel()
		_ = s.repo.UpdateLastUsed(bg, id)
	}(tok.ID)

	return tok, nil
}
// Revoke 由 admin 后台或持有者主动调用。
// 写一条 audit `mcp.token.revoke`。
func (s *tokenService) Revoke(ctx context.Context, id, byUserID uuid.UUID) error {
	revokedID, err := s.repo.Revoke(ctx, id)
	if err != nil {
		return err
	}
	_ = s.audit.Record(ctx, audit.Entry{
		UserID:     &byUserID,
		Action:     "mcp.token.revoke",
		TargetType: "mcp_token",
		TargetID:   revokedID.String(),
		Metadata: map[string]any{
			"token_id":    revokedID.String(),
			"by_user_id":  byUserID.String(),
		},
	})
	return nil
}

// RevokeBySession 由 ACP SessionService 在 session 终止时调用。
// 撤销该 session 的所有 system token，逐条写 audit `mcp.token.revoke_by_session`。
//
// 错误处理：底层 repo 失败 → 返错，让调用方决定回滚或继续；
// 部分 row audit 失败 → 仅记日志，不影响其他 row（audit 失败不阻塞业务）。
func (s *tokenService) RevokeBySession(ctx context.Context, sessionID uuid.UUID) error {
	revoked, err := s.repo.RevokeBySession(ctx, sessionID)
	if err != nil {
		return err
	}
	for _, tokenID := range revoked {
		_ = s.audit.Record(ctx, audit.Entry{
			Action:     "mcp.token.revoke_by_session",
			TargetType: "mcp_token",
			TargetID:   tokenID.String(),
			Metadata: map[string]any{
				"token_id":       tokenID.String(),
				"acp_session_id": sessionID.String(),
			},
		})
	}
	return nil
}