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