feat(server): improve subscription sync stability (#14703)

This commit is contained in:
DarkSky
2026-03-23 21:13:00 +08:00
committed by GitHub
parent dcf041a3f2
commit 5d124ee55b
31 changed files with 1538 additions and 66 deletions
@@ -8,7 +8,12 @@ import ava from 'ava';
import { nanoid } from 'nanoid';
import Sinon from 'sinon';
import { EventBus, JobQueue } from '../../base';
import {
EventBus,
JobQueue,
RequestMutex,
SpaceAccessDenied,
} from '../../base';
import { ConfigModule } from '../../base/config';
import { AuthService } from '../../core/auth';
import { QuotaModule } from '../../core/quota';
@@ -16,10 +21,14 @@ import { StorageModule, WorkspaceBlobStorage } from '../../core/storage';
import {
ContextCategories,
CopilotSessionModel,
Models,
WorkspaceMemberStatus,
WorkspaceModel,
WorkspaceRole,
} from '../../models';
import { CopilotModule } from '../../plugins/copilot';
import { CopilotContextService } from '../../plugins/copilot/context';
import { CopilotContextResolver } from '../../plugins/copilot/context/resolver';
import { CopilotCronJobs } from '../../plugins/copilot/cron';
import {
CopilotEmbeddingJob,
@@ -69,6 +78,7 @@ type Context = {
module: TestingModule;
db: PrismaClient;
event: EventBus;
models: Models;
workspace: WorkspaceModel;
workspaceStorage: WorkspaceBlobStorage;
copilotSession: CopilotSessionModel;
@@ -125,6 +135,11 @@ test.before(async t => {
tapModule: builder => {
// use real JobQueue for testing
builder.overrideProvider(JobQueue).useClass(JobQueue);
builder.overrideProvider(RequestMutex).useValue({
acquire: async () => ({
async [Symbol.asyncDispose]() {},
}),
});
builder.overrideProvider(OpenAIProvider).useClass(MockCopilotProvider);
builder.overrideProvider(SubscriptionService).useClass(
class {
@@ -139,6 +154,7 @@ test.before(async t => {
const auth = module.get(AuthService);
const db = module.get(PrismaClient);
const event = module.get(EventBus);
const models = module.get(Models);
const workspace = module.get(WorkspaceModel);
const workspaceStorage = module.get(WorkspaceBlobStorage);
const copilotSession = module.get(CopilotSessionModel);
@@ -160,6 +176,7 @@ test.before(async t => {
t.context.auth = auth;
t.context.db = db;
t.context.event = event;
t.context.models = models;
t.context.workspace = workspace;
t.context.workspaceStorage = workspaceStorage;
t.context.copilotSession = copilotSession;
@@ -197,7 +214,7 @@ test.beforeEach(async t => {
});
test.after.always(async t => {
await t.context.module.close();
await t.context.module?.close();
});
// ==================== prompt ====================
@@ -241,6 +258,64 @@ test('should be able to manage prompt', async t => {
t.is(await prompt.get(promptName), null, 'should not have the prompt');
});
test('should reject context file uploads after workspace write access is revoked', async t => {
const { auth, context, models, prompt, session, storage, workspace } =
t.context;
const contextResolver = await t.context.module.resolve(
CopilotContextResolver
);
const owner = await auth.signUp(`test-${randomUUID()}@affine.pro`, '123456');
const member = await auth.signUp(`test-${randomUUID()}@affine.pro`, '123456');
const ws = await workspace.create(owner.id);
await models.workspaceUser.set(ws.id, member.id, WorkspaceRole.Collaborator, {
status: WorkspaceMemberStatus.Accepted,
});
await prompt.set(promptName, 'test', [
{ role: 'system', content: 'hello {{word}}' },
]);
const sessionId = await session.create({
userId: member.id,
workspaceId: ws.id,
docId: randomUUID(),
promptName,
pinned: false,
});
const contextSession = await context.create(sessionId);
await models.workspaceUser.set(ws.id, member.id, WorkspaceRole.External);
Sinon.stub(context, 'canEmbedding').get(() => true);
Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient());
const put = Sinon.stub(storage, 'put').resolves();
const buffer = Buffer.from('test pdf');
await t.throwsAsync(
contextResolver.addContextFile(
{ id: member.id } as any,
{
req: {
headers: {
'content-length': String(buffer.length),
},
},
} as any,
{ contextId: contextSession.id },
{
filename: 'sample.pdf',
mimetype: 'application/pdf',
createReadStream: () => Readable.from(buffer),
} as any
),
{
instanceOf: SpaceAccessDenied,
}
);
t.false(put.called);
});
test('should be able to render prompt', async t => {
const { prompt } = t.context;
@@ -2,10 +2,12 @@ import {
acceptInviteByInviteIdMutation,
approveWorkspaceTeamMemberMutation,
createInviteLinkMutation,
deleteBlobMutation,
getInviteInfoQuery,
getMembersByWorkspaceIdQuery,
inviteByEmailsMutation,
leaveWorkspaceMutation,
releaseDeletedBlobsMutation,
revokeMemberPermissionMutation,
WorkspaceInviteLinkExpireTime,
WorkspaceMemberStatus,
@@ -13,6 +15,11 @@ import {
import { faker } from '@faker-js/faker';
import { Models } from '../../../models';
import { FeatureConfigs } from '../../../models/common/feature';
import {
SubscriptionPlan,
SubscriptionRecurring,
} from '../../../plugins/payment/types';
import { Mockers } from '../../mocks';
import { app, e2e } from '../test';
@@ -81,6 +88,175 @@ e2e('should invite a user', async t => {
t.is(getInviteInfo2.status, WorkspaceMemberStatus.Accepted);
});
e2e('should re-check seat when accepting an email invitation', async t => {
const { owner, workspace } = await createWorkspace();
const member = await app.create(Mockers.User);
await app.create(Mockers.TeamWorkspace, {
id: workspace.id,
quantity: 4,
});
await app.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id,
userId: (await app.create(Mockers.User)).id,
});
await app.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id,
userId: (await app.create(Mockers.User)).id,
});
await app.login(owner);
const invite = await app.gql({
query: inviteByEmailsMutation,
variables: {
emails: [member.email],
workspaceId: workspace.id,
},
});
await app.eventBus.emitAsync('workspace.members.allocateSeats', {
workspaceId: workspace.id,
quantity: 4,
});
await app.models.workspaceFeature.remove(workspace.id, 'team_plan_v1');
await app.login(member);
await t.throwsAsync(
app.gql({
query: acceptInviteByInviteIdMutation,
variables: {
workspaceId: workspace.id,
inviteId: invite.inviteMembers[0].inviteId!,
},
})
);
const { getInviteInfo } = await app.gql({
query: getInviteInfoQuery,
variables: {
inviteId: invite.inviteMembers[0].inviteId!,
},
});
t.is(getInviteInfo.status, WorkspaceMemberStatus.Pending);
});
e2e.serial(
'should block accepting pending invitations in readonly mode and recover after blob cleanup',
async t => {
const { owner, workspace } = await createWorkspace();
const member = await app.create(Mockers.User);
const freeStorageQuota = FeatureConfigs.free_plan_v1.configs.storageQuota;
const lifetimeStorageQuota =
FeatureConfigs.lifetime_pro_plan_v1.configs.storageQuota;
FeatureConfigs.free_plan_v1.configs.storageQuota = 1;
FeatureConfigs.lifetime_pro_plan_v1.configs.storageQuota = 2;
t.teardown(() => {
FeatureConfigs.free_plan_v1.configs.storageQuota = freeStorageQuota;
FeatureConfigs.lifetime_pro_plan_v1.configs.storageQuota =
lifetimeStorageQuota;
});
await app.models.userFeature.switchQuota(
owner.id,
'lifetime_pro_plan_v1',
'test setup'
);
await app.login(owner);
const invite = await app.gql({
query: inviteByEmailsMutation,
variables: {
emails: [member.email],
workspaceId: workspace.id,
},
});
await app.models.blob.upsert({
workspaceId: workspace.id,
key: 'overflow-blob',
mime: 'application/octet-stream',
size: 2,
status: 'completed',
uploadId: null,
});
await app.eventBus.emitAsync('user.subscription.canceled', {
userId: owner.id,
plan: SubscriptionPlan.Pro,
recurring: SubscriptionRecurring.Lifetime,
});
t.true(
await app.models.workspaceFeature.has(
workspace.id,
'quota_exceeded_readonly_workspace_v1'
)
);
await app.login(member);
await t.throwsAsync(
app.gql({
query: acceptInviteByInviteIdMutation,
variables: {
workspaceId: workspace.id,
inviteId: invite.inviteMembers[0].inviteId!,
},
})
);
const { getInviteInfo: pendingInvite } = await app.gql({
query: getInviteInfoQuery,
variables: {
inviteId: invite.inviteMembers[0].inviteId!,
},
});
t.is(pendingInvite.status, WorkspaceMemberStatus.Pending);
await app.login(owner);
await app.gql({
query: deleteBlobMutation,
variables: {
workspaceId: workspace.id,
key: 'overflow-blob',
permanently: false,
},
});
await app.gql({
query: releaseDeletedBlobsMutation,
variables: {
workspaceId: workspace.id,
},
});
t.false(
await app.models.workspaceFeature.has(
workspace.id,
'quota_exceeded_readonly_workspace_v1'
)
);
await app.login(member);
await app.gql({
query: acceptInviteByInviteIdMutation,
variables: {
workspaceId: workspace.id,
inviteId: invite.inviteMembers[0].inviteId!,
},
});
const { getInviteInfo: acceptedInvite } = await app.gql({
query: getInviteInfoQuery,
variables: {
inviteId: invite.inviteMembers[0].inviteId!,
},
});
t.is(acceptedInvite.status, WorkspaceMemberStatus.Accepted);
}
);
e2e('should leave a workspace', async t => {
const { owner, workspace } = await createWorkspace();
const u2 = await app.create(Mockers.User);
@@ -1,10 +1,18 @@
import {
getInviteInfoQuery,
inviteByEmailsMutation,
publishPageMutation,
revokeMemberPermissionMutation,
revokePublicPageMutation,
WorkspaceMemberStatus,
} from '@affine/graphql';
import { QuotaService } from '../../../core/quota/service';
import { WorkspaceRole } from '../../../models';
import {
SubscriptionPlan,
SubscriptionRecurring,
} from '../../../plugins/payment/types';
import { Mockers } from '../../mocks';
import { app, e2e } from '../test';
@@ -54,6 +62,42 @@ const getInvitationInfo = async (inviteId: string) => {
return result.getInviteInfo;
};
const publishDoc = async (workspaceId: string, docId: string) => {
const { publishDoc } = await app.gql({
query: publishPageMutation,
variables: {
workspaceId,
pageId: docId,
},
});
return publishDoc;
};
const revokePublicDoc = async (workspaceId: string, docId: string) => {
const { revokePublicDoc } = await app.gql({
query: revokePublicPageMutation,
variables: {
workspaceId,
pageId: docId,
},
});
return revokePublicDoc;
};
const revokeMember = async (workspaceId: string, userId: string) => {
const { revokeMember } = await app.gql({
query: revokeMemberPermissionMutation,
variables: {
workspaceId,
userId,
},
});
return revokeMember;
};
e2e('should set new invited users to AllocatingSeat', async t => {
const { owner, workspace } = await createTeamWorkspace();
await app.login(owner);
@@ -165,3 +209,136 @@ e2e('should set all rests to NeedMoreSeat', async t => {
WorkspaceMemberStatus.NeedMoreSeat
);
});
e2e(
'should cleanup non-accepted members when team workspace is downgraded',
async t => {
const { workspace } = await createTeamWorkspace();
const pending = await app.create(Mockers.User);
await app.create(Mockers.WorkspaceUser, {
userId: pending.id,
workspaceId: workspace.id,
status: WorkspaceMemberStatus.Pending,
});
const allocating = await app.create(Mockers.User);
await app.create(Mockers.WorkspaceUser, {
userId: allocating.id,
workspaceId: workspace.id,
status: WorkspaceMemberStatus.AllocatingSeat,
source: 'Email',
});
const underReview = await app.create(Mockers.User);
await app.create(Mockers.WorkspaceUser, {
userId: underReview.id,
workspaceId: workspace.id,
status: WorkspaceMemberStatus.UnderReview,
});
await app.eventBus.emitAsync('workspace.subscription.canceled', {
workspaceId: workspace.id,
plan: SubscriptionPlan.Team,
recurring: SubscriptionRecurring.Monthly,
});
const [members] = await app.models.workspaceUser.paginate(workspace.id, {
first: 20,
offset: 0,
});
t.deepEqual(
members.map(member => member.status),
[
WorkspaceMemberStatus.Accepted,
WorkspaceMemberStatus.Accepted,
WorkspaceMemberStatus.Accepted,
]
);
t.false(await app.models.workspace.isTeamWorkspace(workspace.id));
}
);
e2e(
'should demote accepted admins and keep workspace writable when downgrade stays within owner quota',
async t => {
const { workspace, owner, admin } = await createTeamWorkspace();
await app.eventBus.emitAsync('workspace.subscription.canceled', {
workspaceId: workspace.id,
plan: SubscriptionPlan.Team,
recurring: SubscriptionRecurring.Monthly,
});
t.false(await app.models.workspace.isTeamWorkspace(workspace.id));
t.false(
await app.models.workspaceFeature.has(
workspace.id,
'quota_exceeded_readonly_workspace_v1'
)
);
t.is(
(await app.models.workspaceUser.get(workspace.id, admin.id))?.type,
WorkspaceRole.Collaborator
);
await app.login(owner);
await t.notThrowsAsync(publishDoc(workspace.id, 'doc-1'));
}
);
e2e(
'should enter readonly mode on over-quota team downgrade and recover through cleanup actions',
async t => {
const { workspace, owner, admin } = await createTeamWorkspace(20);
const extraMembers = await Promise.all(
Array.from({ length: 8 }).map(async () => {
const member = await app.create(Mockers.User);
await app.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id,
userId: member.id,
});
return member;
})
);
await app.login(owner);
await publishDoc(workspace.id, 'published-doc');
await app.eventBus.emitAsync('workspace.subscription.canceled', {
workspaceId: workspace.id,
plan: SubscriptionPlan.Team,
recurring: SubscriptionRecurring.Monthly,
});
t.false(await app.models.workspace.isTeamWorkspace(workspace.id));
t.true(
await app.models.workspaceFeature.has(
workspace.id,
'quota_exceeded_readonly_workspace_v1'
)
);
t.is(
(await app.models.workspaceUser.get(workspace.id, admin.id))?.type,
WorkspaceRole.Collaborator
);
await t.throwsAsync(publishDoc(workspace.id, 'blocked-doc'));
await t.notThrowsAsync(revokePublicDoc(workspace.id, 'published-doc'));
const quota = await app
.get(QuotaService)
.getWorkspaceQuotaWithUsage(workspace.id);
for (const member of extraMembers.slice(0, quota.overcapacityMemberCount)) {
await revokeMember(workspace.id, member.id);
}
t.false(
await app.models.workspaceFeature.has(
workspace.id,
'quota_exceeded_readonly_workspace_v1'
)
);
}
);
@@ -65,6 +65,28 @@ test('should transfer workespace owner', async t => {
const owner2 = await models.workspaceUser.getOwner(workspace.id);
t.is(owner2.id, user2.id);
const oldOwnerRole = await models.workspaceUser.get(workspace.id, user.id);
t.is(oldOwnerRole?.type, WorkspaceRole.Collaborator);
});
test('should keep old owner as admin when transferring a team workspace', async t => {
const [user, user2] = await module.create(Mockers.User, 2);
const workspace = await module.create(Mockers.Workspace, {
owner: { id: user.id },
});
await module.create(Mockers.TeamWorkspace, {
id: workspace.id,
quantity: 10,
});
await module.create(Mockers.WorkspaceUser, {
workspaceId: workspace.id,
userId: user2.id,
});
await models.workspaceUser.setOwner(workspace.id, user2.id);
const oldOwnerRole = await models.workspaceUser.get(workspace.id, user.id);
t.is(oldOwnerRole?.type, WorkspaceRole.Admin);
});
test('should throw if transfer owner to non-active member', async t => {
@@ -14,6 +14,7 @@ import { ServerFeature } from '../../core/config/types';
import { Models } from '../../models';
import { OAuthProviderName } from '../../plugins/oauth/config';
import { OAuthProviderFactory } from '../../plugins/oauth/factory';
import { GithubOAuthProvider } from '../../plugins/oauth/providers/github';
import { GoogleOAuthProvider } from '../../plugins/oauth/providers/google';
import { OIDCProvider } from '../../plugins/oauth/providers/oidc';
import { OAuthService } from '../../plugins/oauth/service';
@@ -38,6 +39,10 @@ test.before(async t => {
clientId: 'google-client-id',
clientSecret: 'google-client-secret',
},
github: {
clientId: 'github-client-id',
clientSecret: 'github-client-secret',
},
oidc: {
clientId: '',
clientSecret: '',
@@ -293,7 +298,7 @@ test('should be able to get registered oauth providers', async t => {
const providers = oauth.availableOAuthProviders();
t.deepEqual(providers, [OAuthProviderName.Google]);
t.deepEqual(providers, [OAuthProviderName.Google, OAuthProviderName.GitHub]);
});
test('should throw if code is missing in callback uri', async t => {
@@ -441,6 +446,24 @@ function mockOAuthProvider(
return clientNonce;
}
function mockGithubOAuthProvider(
app: TestingApp,
clientNonce: string = randomUUID()
) {
const provider = app.get(GithubOAuthProvider);
const oauth = app.get(OAuthService);
Sinon.stub(oauth, 'isValidState').resolves(true);
Sinon.stub(oauth, 'getOAuthState').resolves({
provider: OAuthProviderName.GitHub,
clientNonce,
});
Sinon.stub(provider, 'getToken').resolves({ accessToken: '1' });
return { provider, clientNonce };
}
function mockOidcProvider(
provider: OIDCProvider,
{
@@ -645,6 +668,76 @@ test('should be able to fullfil user with oauth sign in', async t => {
t.is(account!.user.id, u3.id);
});
test('github oauth should resolve private email from emails api', async t => {
const { app, db } = t.context;
const email = 'github-private@affine.pro';
const { clientNonce, provider } = mockGithubOAuthProvider(app);
const fetchJson = Sinon.stub(provider as any, 'fetchJson');
fetchJson.onFirstCall().resolves({
login: 'github-user',
email: null,
avatar_url: 'avatar',
name: 'DarkSky',
});
fetchJson.onSecondCall().resolves([
{ email: 'unverified@affine.pro', primary: true, verified: false },
{ email, primary: false, verified: true },
]);
await app
.POST('/api/oauth/callback')
.send({ code: '1', state: '1', client_nonce: clientNonce })
.expect(HttpStatus.OK);
const sessionUser = await currentUser(app);
t.truthy(sessionUser);
t.is(sessionUser!.email, email);
const user = await db.user.findFirst({
select: {
email: true,
connectedAccounts: true,
},
where: {
email,
},
});
t.truthy(user);
t.is(user!.connectedAccounts[0].provider, OAuthProviderName.GitHub);
t.is(user!.connectedAccounts[0].providerAccountId, 'github-user');
});
test('github oauth should reject responses without a verified email', async t => {
const { app } = t.context;
const provider = app.get(GithubOAuthProvider);
const fetchJson = Sinon.stub(provider as any, 'fetchJson');
fetchJson.onFirstCall().resolves({
login: 'github-user',
email: null,
avatar_url: 'avatar',
name: 'DarkSky',
});
fetchJson
.onSecondCall()
.resolves([
{ email: 'private@affine.pro', primary: true, verified: false },
]);
const error = await t.throwsAsync(
provider.getUser(
{ accessToken: 'token' },
{ token: 'state', provider: OAuthProviderName.GitHub }
)
);
t.true(error instanceof InvalidOauthResponse);
});
test('oidc should accept email from id token when userinfo email is missing', async t => {
const { app } = t.context;
@@ -11,6 +11,7 @@ import { ConfigFactory, ConfigModule } from '../../base/config';
import { CurrentUser } from '../../core/auth';
import { AuthService } from '../../core/auth/service';
import { EarlyAccessType, FeatureService } from '../../core/features';
import { SubscriptionCronJobs } from '../../plugins/payment/cron';
import { SubscriptionService } from '../../plugins/payment/service';
import { StripeFactory } from '../../plugins/payment/stripe';
import {
@@ -871,6 +872,34 @@ test('should be able to cancel subscription', async t => {
t.truthy(subInDB.canceledAt);
});
test('should reconcile canceled stripe subscriptions and revoke local entitlement', async t => {
const { app, db, event, service, stripe, u1 } = t.context;
const cron = app.get(SubscriptionCronJobs);
await service.saveStripeSubscription(sub);
event.emit.resetHistory();
stripe.subscriptions.retrieve.resolves({
...sub,
status: SubscriptionStatus.Canceled,
} as any);
await cron.reconcileStripeSubscriptions();
const subInDB = await db.subscription.findFirst({
where: { targetId: u1.id, stripeSubscriptionId: sub.id },
});
t.is(subInDB, null);
t.true(
event.emit.calledWith('user.subscription.canceled', {
userId: u1.id,
plan: SubscriptionPlan.Pro,
recurring: SubscriptionRecurring.Monthly,
})
);
});
test('should be able to resume subscription', async t => {
const { service, db, u1, stripe } = t.context;