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

//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.Token{
		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.Token{
		ID: uuid.New(), TokenHash: hash, UserID: uid,
		Name: "x", Issuer: mcp.IssuerSystem, CreatedBy: uid,
	})
	require.Error(t, err)
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeInternal, ae.Code, "CHECK violation should wrap to 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.Token{
		ID: uuid.New(), TokenHash: hash, UserID: uid,
		Name: "y", Issuer: mcp.IssuerAdmin, ACPSessionID: &fakeSessionID,
		CreatedBy: uid,
	})
	require.Error(t, err)
	ae, ok := errs.As(err)
	require.True(t, ok)
	assert.Equal(t, errs.CodeInternal, ae.Code, "CHECK violation should wrap to CodeInternal")
}

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.Token{
			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.Token{
			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.Token{
		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.Token{
		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.Token{
		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.Token{
		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.Token{
		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)
}