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

package mcp_test

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"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"
)

// ===== fake repo =====

type fakeRepo struct {
	mu     sync.Mutex
	tokens map[uuid.UUID]*mcp.Token
	byHash map[string]uuid.UUID
}

func newFakeRepo() *fakeRepo {
	return &fakeRepo{
		tokens: map[uuid.UUID]*mcp.Token{},
		byHash: map[string]uuid.UUID{},
	}
}

func (r *fakeRepo) Create(ctx context.Context, t *mcp.Token) (*mcp.Token, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	c := *t
	c.CreatedAt = time.Now()
	r.tokens[c.ID] = &c
	r.byHash[string(c.TokenHash)] = c.ID
	return &c, nil
}
func (r *fakeRepo) GetByHash(ctx context.Context, h []byte) (*mcp.Token, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	id, ok := r.byHash[string(h)]
	if !ok {
		return nil, errs.New(errs.CodeMcpTokenInvalid, "not found")
	}
	c := *r.tokens[id]
	return &c, nil
}
func (r *fakeRepo) GetByID(ctx context.Context, id uuid.UUID) (*mcp.Token, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	t, ok := r.tokens[id]
	if !ok {
		return nil, errs.New(errs.CodeMcpTokenNotFound, "not found")
	}
	c := *t
	return &c, nil
}
func (r *fakeRepo) List(ctx context.Context, caller *uuid.UUID, activeOnly bool) ([]*mcp.Token, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	var out []*mcp.Token
	for _, t := range r.tokens {
		if caller != nil && (t.UserID != *caller || t.Issuer != mcp.IssuerAdmin) {
			continue
		}
		if activeOnly && !t.IsActive(time.Now()) {
			continue
		}
		c := *t
		out = append(out, &c)
	}
	return out, nil
}
func (r *fakeRepo) UpdateLastUsed(ctx context.Context, id uuid.UUID) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	if t, ok := r.tokens[id]; ok {
		now := time.Now()
		t.LastUsedAt = &now
	}
	return nil
}
func (r *fakeRepo) Revoke(ctx context.Context, id uuid.UUID) (uuid.UUID, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	t, ok := r.tokens[id]
	if !ok || t.RevokedAt != nil {
		return uuid.Nil, errs.New(errs.CodeMcpTokenNotFound, "not found")
	}
	now := time.Now()
	t.RevokedAt = &now
	return id, nil
}
func (r *fakeRepo) RevokeBySession(ctx context.Context, sid uuid.UUID) ([]uuid.UUID, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	var revoked []uuid.UUID
	now := time.Now()
	for _, t := range r.tokens {
		if t.Issuer == mcp.IssuerSystem && t.ACPSessionID != nil && *t.ACPSessionID == sid && t.RevokedAt == nil {
			t.RevokedAt = &now
			revoked = append(revoked, t.ID)
		}
	}
	return revoked, nil
}
func (r *fakeRepo) ReapStaleSystemTokens(ctx context.Context) ([]uuid.UUID, error) {
	return nil, nil
}
func (r *fakeRepo) PurgeExpired(ctx context.Context, before time.Time) (int64, error) {
	return 0, nil
}
func (r *fakeRepo) DeleteByID(ctx context.Context, id uuid.UUID) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	delete(r.tokens, id)
	return nil
}

func (r *fakeRepo) WithTx(tx pgx.Tx) mcp.Repository {
	return r
}

// ===== fake audit =====

type fakeAudit struct {
	mu      sync.Mutex
	entries []audit.Entry
}

func (a *fakeAudit) Record(ctx context.Context, e audit.Entry) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.entries = append(a.entries, e)
	return nil
}
func (a *fakeAudit) actions() []string {
	a.mu.Lock()
	defer a.mu.Unlock()
	out := make([]string, 0, len(a.entries))
	for _, e := range a.entries {
		out = append(out, e.Action)
	}
	return out
}

// ===== tests =====

func TestIssueAdmin_HappyPath(t *testing.T) {
	t.Parallel()
	repo := newFakeRepo()
	au := &fakeAudit{}
	svc := mcp.NewTokenService(repo, au, nil)

	uid := uuid.New()
	in := mcp.IssueAdminInput{
		UserID:    uid,
		Name:      "test",
		Scope:     mcp.Scope{Tools: []string{"list_issues"}},
		ExpiresAt: time.Now().Add(time.Hour),
		CreatedBy: uid,
	}
	res, err := svc.IssueAdmin(context.Background(), in)
	require.NoError(t, err)
	assert.NotEmpty(t, res.Plaintext)
	assert.Equal(t, []string{"mcp.token.create"}, au.actions())
}

func TestIssueAdmin_NameRequired(t *testing.T) {
	t.Parallel()
	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}, nil)

	_, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
		UserID:    uuid.New(),
		Name:      "",
		ExpiresAt: time.Now().Add(time.Hour),
		CreatedBy: uuid.New(),
	})
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeMcpTokenNameRequired, ae.Code)
}

func TestIssueAdmin_ExpiresRequired(t *testing.T) {
	t.Parallel()
	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}, nil)

	_, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
		UserID:    uuid.New(),
		Name:      "x",
		CreatedBy: uuid.New(),
	})
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeMcpTokenExpiresRequired, ae.Code)
}

func TestIssueAdmin_ExpiresMustBeFuture(t *testing.T) {
	t.Parallel()
	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}, nil)

	_, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
		UserID:    uuid.New(),
		Name:      "x",
		ExpiresAt: time.Now().Add(-time.Hour), // past
		CreatedBy: uuid.New(),
	})
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeMcpTokenExpiresRequired, ae.Code)
}

func TestIssueSystemToken_Defaults(t *testing.T) {
	t.Parallel()
	repo := newFakeRepo()
	au := &fakeAudit{}
	svc := mcp.NewTokenService(repo, au, nil)

	uid := uuid.New()
	sid := uuid.New()
	pid := uuid.New()
	res, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{
		UserID:       uid,
		ACPSessionID: sid,
		ProjectIDs:   []uuid.UUID{pid},
	})
	require.NoError(t, err)

	got, err := repo.GetByID(context.Background(), res.ID)
	require.NoError(t, err)
	assert.Equal(t, mcp.IssuerSystem, got.Issuer)
	require.NotNil(t, got.ACPSessionID)
	assert.Equal(t, sid, *got.ACPSessionID)
	assert.NotNil(t, got.ExpiresAt)
	assert.WithinDuration(t, time.Now().Add(24*time.Hour), *got.ExpiresAt, time.Minute)
	// 默认工具白名单包含 PM 写入但不含 list_recent_messages
	assert.Contains(t, got.Scope.Tools, "create_issue")
	assert.NotContains(t, got.Scope.Tools, "list_recent_messages")
	assert.Equal(t, []string{"mcp.token.create_system"}, au.actions())
}

func TestAuthenticate_HappyPath(t *testing.T) {
	t.Parallel()
	repo := newFakeRepo()
	svc := mcp.NewTokenService(repo, &fakeAudit{}, nil)

	res, err := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
		UserID:    uuid.New(),
		Name:      "x",
		ExpiresAt: time.Now().Add(time.Hour),
		CreatedBy: uuid.New(),
	})
	require.NoError(t, err)

	tok, err := svc.Authenticate(context.Background(), res.Plaintext)
	require.NoError(t, err)
	assert.Equal(t, res.ID, tok.ID)
}

func TestAuthenticate_Empty(t *testing.T) {
	t.Parallel()
	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}, nil)

	_, err := svc.Authenticate(context.Background(), "")
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code)
}

func TestAuthenticate_Revoked(t *testing.T) {
	t.Parallel()
	repo := newFakeRepo()
	svc := mcp.NewTokenService(repo, &fakeAudit{}, nil)

	res, _ := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
		UserID:    uuid.New(),
		Name:      "x",
		ExpiresAt: time.Now().Add(time.Hour),
		CreatedBy: uuid.New(),
	})
	require.NoError(t, svc.Revoke(context.Background(), res.ID, uuid.New()))

	_, err := svc.Authenticate(context.Background(), res.Plaintext)
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeMcpTokenRevoked, ae.Code)
}

func TestAuthenticate_Expired(t *testing.T) {
	t.Parallel()
	repo := newFakeRepo()
	svc := mcp.NewTokenService(repo, &fakeAudit{}, nil)

	res, _ := svc.IssueAdmin(context.Background(), mcp.IssueAdminInput{
		UserID:    uuid.New(),
		Name:      "x",
		ExpiresAt: time.Now().Add(time.Hour),
		CreatedBy: uuid.New(),
	})
	// 直接改 fake repo 中 expires_at 为已过去
	past := time.Now().Add(-time.Hour)
	repo.tokens[res.ID].ExpiresAt = &past

	_, err := svc.Authenticate(context.Background(), res.Plaintext)
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeMcpTokenExpired, ae.Code)
}

func TestAuthenticate_TamperedToken(t *testing.T) {
	t.Parallel()
	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}, nil)

	_, err := svc.Authenticate(context.Background(), "mcpt_tampered_value")
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code)
}

func TestRevokeBySession_AuditPerToken(t *testing.T) {
	t.Parallel()
	repo := newFakeRepo()
	au := &fakeAudit{}
	svc := mcp.NewTokenService(repo, au, nil)

	uid := uuid.New()
	sid := uuid.New()
	pid := uuid.New()
	for i := 0; i < 3; i++ {
		_, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{
			UserID: uid, ACPSessionID: sid, ProjectIDs: []uuid.UUID{pid},
		})
		require.NoError(t, err)
	}

	require.NoError(t, svc.RevokeBySession(context.Background(), sid))

	// 3 个 system 创建 + 3 个 by_session 撤销 = 6 条 audit
	actions := au.actions()
	assert.Len(t, actions, 6)
	createCnt, revokeCnt := 0, 0
	for _, a := range actions {
		switch a {
		case "mcp.token.create_system":
			createCnt++
		case "mcp.token.revoke_by_session":
			revokeCnt++
		}
	}
	assert.Equal(t, 3, createCnt)
	assert.Equal(t, 3, revokeCnt)
}

func TestIssueSystemToken_NilACPSessionID(t *testing.T) {
	t.Parallel()
	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}, nil)

	_, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{
		UserID:       uuid.New(),
		ACPSessionID: uuid.Nil,
	})
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeInvalidInput, ae.Code)
}

func TestIssueSystemToken_NilUserID(t *testing.T) {
	t.Parallel()
	svc := mcp.NewTokenService(newFakeRepo(), &fakeAudit{}, nil)

	_, err := svc.IssueSystemToken(context.Background(), mcp.IssueSystemTokenInput{
		UserID:       uuid.Nil,
		ACPSessionID: uuid.New(),
	})
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeInvalidInput, ae.Code)
}

// 编译期保证：类型实现接口
var _ = errors.New

func TestIssueSystemTokenWithTx_DelegatesToRepoWithTx(t *testing.T) {
	t.Parallel()
	repo := newFakeRepo()
	au := &fakeAudit{}
	svc := mcp.NewTokenService(repo, au, nil)

	uid := uuid.New()
	sid := uuid.New()
	pid := uuid.New()
	res, err := svc.IssueSystemTokenWithTx(context.Background(), nil, mcp.IssueSystemTokenInput{
		UserID: uid, ACPSessionID: sid, ProjectIDs: []uuid.UUID{pid},
	})
	require.NoError(t, err)
	assert.NotEmpty(t, res.Plaintext)
	assert.Equal(t, []string{"mcp.token.create_system"}, au.actions())

	got, err := repo.GetByID(context.Background(), res.ID)
	require.NoError(t, err)
	assert.Equal(t, mcp.IssuerSystem, got.Issuer)
}