//go:build integration package mcp_test import ( "context" "testing" "time" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgxpool" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" tcpg "github.com/testcontainers/testcontainers-go/modules/postgres" "github.com/yan1h/agent-coding-workflow/internal/infra/db" "github.com/yan1h/agent-coding-workflow/internal/infra/errs" "github.com/yan1h/agent-coding-workflow/internal/mcp" ) func setupRepo(t *testing.T) (mcp.Repository, *pgxpool.Pool) { t.Helper() ctx := context.Background() container, err := tcpg.Run(ctx, "postgres:16-alpine", tcpg.WithDatabase("acw"), tcpg.WithUsername("acw"), tcpg.WithPassword("acw"), ) require.NoError(t, err) t.Cleanup(func() { _ = container.Terminate(ctx) }) dsn, err := container.ConnectionString(ctx, "sslmode=disable") require.NoError(t, err) require.NoError(t, db.Migrate(ctx, dsn, "file://../../migrations")) pool, err := db.NewPool(ctx, dsn) require.NoError(t, err) t.Cleanup(pool.Close) return mcp.NewPostgresRepository(pool), pool } // mustInsertUser 直插一条 user,返 UID。沿用 acp 测试模式。 func mustInsertUser(t *testing.T, ctx context.Context, pool *pgxpool.Pool, email string, isAdmin bool) uuid.UUID { t.Helper() uid := uuid.New() _, err := pool.Exec(ctx, ` INSERT INTO users (id, email, password_hash, display_name, is_admin) VALUES ($1, $2, 'argon2id$dummy', $3, $4)`, uid, email, email, isAdmin) require.NoError(t, err) return uid } func TestPgRepo_AdminToken_CRUD(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) uid := mustInsertUser(t, ctx, pool, "admin1@local", true) tok, hash, err := mcp.NewToken() require.NoError(t, err) expires := time.Now().Add(90 * 24 * time.Hour) created, err := repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash, UserID: uid, Name: "test admin token", Issuer: mcp.IssuerAdmin, Scope: mcp.Scope{Tools: []string{"list_issues"}, ProjectIDs: nil}, ExpiresAt: &expires, CreatedBy: uid, }) require.NoError(t, err) assert.NotEqual(t, uuid.Nil, created.ID) assert.Equal(t, mcp.IssuerAdmin, created.Issuer) assert.Equal(t, []string{"list_issues"}, created.Scope.Tools) // GetByHash got, err := repo.GetByHash(ctx, mcp.HashToken(tok)) require.NoError(t, err) assert.Equal(t, created.ID, got.ID) // GetByID got2, err := repo.GetByID(ctx, created.ID) require.NoError(t, err) assert.Equal(t, created.ID, got2.ID) // List (admin call: callerUserID nil) all, err := repo.List(ctx, nil, false) require.NoError(t, err) require.Len(t, all, 1) assert.Equal(t, created.ID, all[0].ID) // UpdateLastUsed require.NoError(t, repo.UpdateLastUsed(ctx, created.ID)) got3, err := repo.GetByID(ctx, created.ID) require.NoError(t, err) require.NotNil(t, got3.LastUsedAt) assert.WithinDuration(t, time.Now(), *got3.LastUsedAt, 5*time.Second) // Revoke revokedID, err := repo.Revoke(ctx, created.ID) require.NoError(t, err) assert.Equal(t, created.ID, revokedID) // 二次 Revoke → NotFound _, err = repo.Revoke(ctx, created.ID) ae, ok := errs.As(err) require.True(t, ok) assert.Equal(t, errs.CodeMcpTokenNotFound, ae.Code) } func TestPgRepo_GetByHash_NotFound(t *testing.T) { t.Parallel() ctx := context.Background() repo, _ := setupRepo(t) _, err := repo.GetByHash(ctx, mcp.HashToken("mcpt_nope")) ae, ok := errs.As(err) require.True(t, ok) assert.Equal(t, errs.CodeMcpTokenInvalid, ae.Code) } func TestPgRepo_CheckConstraint_SystemRequiresSession(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) uid := mustInsertUser(t, ctx, pool, "u@local", false) _, hash, err := mcp.NewToken() require.NoError(t, err) // system token + 没 acp_session_id → CHECK 拒绝 _, err = repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash, UserID: uid, Name: "x", Issuer: mcp.IssuerSystem, CreatedBy: uid, }) require.Error(t, err) // PG 抛 23514 check_violation;包装在 errs.CodeInternal 里 } func TestPgRepo_CheckConstraint_AdminMustNotHaveSession(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) uid := mustInsertUser(t, ctx, pool, "u2@local", false) _, hash, err := mcp.NewToken() require.NoError(t, err) fakeSessionID := uuid.New() // admin token + 有 acp_session_id → CHECK 拒绝 _, err = repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash, UserID: uid, Name: "y", Issuer: mcp.IssuerAdmin, ACPSessionID: &fakeSessionID, CreatedBy: uid, }) require.Error(t, err) } func TestPgRepo_List_FilterByCaller(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) admin := mustInsertUser(t, ctx, pool, "admin@local", true) user1 := mustInsertUser(t, ctx, pool, "u1@local", false) user2 := mustInsertUser(t, ctx, pool, "u2@local", false) mkAdmin := func(uid uuid.UUID, name string) { _, hash, _ := mcp.NewToken() _, err := repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash, UserID: uid, Name: name, Issuer: mcp.IssuerAdmin, CreatedBy: admin, }) require.NoError(t, err) } mkAdmin(user1, "u1.t1") mkAdmin(user1, "u1.t2") mkAdmin(user2, "u2.t1") // admin 视角:看全部 all, err := repo.List(ctx, nil, false) require.NoError(t, err) assert.Len(t, all, 3) // user1 视角:仅自己持有的 admin-issued mine, err := repo.List(ctx, &user1, false) require.NoError(t, err) assert.Len(t, mine, 2) for _, t2 := range mine { assert.Equal(t, user1, t2.UserID) } } func TestPgRepo_List_ActiveOnly(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) admin := mustInsertUser(t, ctx, pool, "a@local", true) user := mustInsertUser(t, ctx, pool, "u@local", false) // 一个有效,一个已撤销,一个已过期 expFuture := time.Now().Add(24 * time.Hour) expPast := time.Now().Add(-1 * time.Hour) mk := func(name string, exp *time.Time) uuid.UUID { _, hash, _ := mcp.NewToken() tok, err := repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash, UserID: user, Name: name, Issuer: mcp.IssuerAdmin, ExpiresAt: exp, CreatedBy: admin, }) require.NoError(t, err) return tok.ID } mk("active", &expFuture) revID := mk("revoked", &expFuture) mk("expired", &expPast) _, err := repo.Revoke(ctx, revID) require.NoError(t, err) all, err := repo.List(ctx, nil, false) require.NoError(t, err) assert.Len(t, all, 3) active, err := repo.List(ctx, nil, true) require.NoError(t, err) require.Len(t, active, 1) assert.Equal(t, "active", active[0].Name) } func TestPgRepo_PurgeExpired(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) admin := mustInsertUser(t, ctx, pool, "p@local", true) user := mustInsertUser(t, ctx, pool, "p2@local", false) old := time.Now().Add(-30 * 24 * time.Hour) expFuture := time.Now().Add(24 * time.Hour) // 一个早已过期的 row(手工写 expires_at = 30d ago) _, hash, _ := mcp.NewToken() tk, err := repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash, UserID: user, Name: "old", Issuer: mcp.IssuerAdmin, ExpiresAt: &old, CreatedBy: admin, }) require.NoError(t, err) // 一个新 token _, hash2, _ := mcp.NewToken() _, err = repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash2, UserID: user, Name: "new", Issuer: mcp.IssuerAdmin, ExpiresAt: &expFuture, CreatedBy: admin, }) require.NoError(t, err) // 删 7 天前 retention cutoff := time.Now().Add(-7 * 24 * time.Hour) n, err := repo.PurgeExpired(ctx, cutoff) require.NoError(t, err) assert.Equal(t, int64(1), n) // 老 token 应已删 _, err = repo.GetByID(ctx, tk.ID) ae, ok := errs.As(err) require.True(t, ok) assert.Equal(t, errs.CodeMcpTokenNotFound, ae.Code) } // mustInsertACPSetup 准备测 system token 需要的 fixtures: // project + workspace + agent_kind + acp_sessions row(status 任选)。 // 返回 sessionID + projectID(用于 system token 的 ACPSessionID + scope.project_ids)。 func mustInsertACPSetup(t *testing.T, ctx context.Context, pool *pgxpool.Pool, ownerUID uuid.UUID, status string) (sessionID, projectID uuid.UUID) { t.Helper() // project pid := uuid.New() _, err := pool.Exec(ctx, ` INSERT INTO projects (id, slug, name, owner_id, visibility) VALUES ($1, $2, $2, $3, 'private')`, pid, "p"+pid.String()[:8], ownerUID) require.NoError(t, err) // workspace wid := uuid.New() _, err = pool.Exec(ctx, ` INSERT INTO workspaces (id, project_id, slug, name, git_remote_url, default_branch, main_path, sync_status) VALUES ($1, $2, $3, $3, 'git@local:x.git', 'main', '/tmp/'||$1::text, 'idle')`, wid, pid, "w"+wid.String()[:8]) require.NoError(t, err) // agent_kind akid := uuid.New() _, err = pool.Exec(ctx, ` INSERT INTO acp_agent_kinds (id, name, display_name, binary_path, args, enabled, created_by) VALUES ($1, $2, $2, '/bin/echo', '{}', TRUE, $3)`, akid, "ak"+akid.String()[:8], ownerUID) require.NoError(t, err) // session sid := uuid.New() _, err = pool.Exec(ctx, ` INSERT INTO acp_sessions (id, workspace_id, project_id, agent_kind_id, user_id, branch, cwd_path, is_main_worktree, status) VALUES ($1, $2, $3, $4, $5, 'main', '/tmp/x', TRUE, $6)`, sid, wid, pid, akid, ownerUID, status) require.NoError(t, err) return sid, pid } func TestPgRepo_RevokeBySession(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) uid := mustInsertUser(t, ctx, pool, "rs@local", false) sid, pid := mustInsertACPSetup(t, ctx, pool, uid, "running") // 插一个 system token _, hash, _ := mcp.NewToken() tk, err := repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: hash, UserID: uid, Name: "system-x", Issuer: mcp.IssuerSystem, Scope: mcp.Scope{ProjectIDs: []uuid.UUID{pid}}, ACPSessionID: &sid, CreatedBy: uid, }) require.NoError(t, err) // 撤销 revoked, err := repo.RevokeBySession(ctx, sid) require.NoError(t, err) require.Len(t, revoked, 1) assert.Equal(t, tk.ID, revoked[0]) // row 上 revoked_at 写入 got, err := repo.GetByID(ctx, tk.ID) require.NoError(t, err) require.NotNil(t, got.RevokedAt) // 重复调用 → 0 行(已撤销不再返) again, err := repo.RevokeBySession(ctx, sid) require.NoError(t, err) assert.Empty(t, again) } func TestPgRepo_ReapStaleSystemTokens(t *testing.T) { t.Parallel() ctx := context.Background() repo, pool := setupRepo(t) uid := mustInsertUser(t, ctx, pool, "reap@local", false) // 一个 crashed session 的 system token sidCrashed, pidA := mustInsertACPSetup(t, ctx, pool, uid, "crashed") _, h1, _ := mcp.NewToken() tk1, err := repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: h1, UserID: uid, Name: "crashed-tok", Issuer: mcp.IssuerSystem, Scope: mcp.Scope{ProjectIDs: []uuid.UUID{pidA}}, ACPSessionID: &sidCrashed, CreatedBy: uid, }) require.NoError(t, err) // 一个 running session 的 system token(不该被 reap) sidRunning, pidB := mustInsertACPSetup(t, ctx, pool, uid, "running") _, h2, _ := mcp.NewToken() tk2, err := repo.Create(ctx, &mcp.MCPToken{ ID: uuid.New(), TokenHash: h2, UserID: uid, Name: "running-tok", Issuer: mcp.IssuerSystem, Scope: mcp.Scope{ProjectIDs: []uuid.UUID{pidB}}, ACPSessionID: &sidRunning, CreatedBy: uid, }) require.NoError(t, err) reaped, err := repo.ReapStaleSystemTokens(ctx) require.NoError(t, err) require.Len(t, reaped, 1) assert.Equal(t, tk1.ID, reaped[0]) got1, _ := repo.GetByID(ctx, tk1.ID) require.NotNil(t, got1.RevokedAt) got2, _ := repo.GetByID(ctx, tk2.ID) assert.Nil(t, got2.RevokedAt) }