diff --git a/internal/mcp/repository_test.go b/internal/mcp/repository_test.go index f841dff..43e329a 100644 --- a/internal/mcp/repository_test.go +++ b/internal/mcp/repository_test.go @@ -271,3 +271,120 @@ func TestPgRepo_PurgeExpired(t *testing.T) { 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) +}