Skip to content

Commit 09ed990

Browse files
waleedlatif1waleedlatif1
authored
feat(mothership): erase private memory when a member leaves or an organization is deleted (#9503)
* feat(mothership): erase private memory when a member leaves or an organization is deleted Account deletion already queues the worker's memory cleanup. Removing a member, removing an external collaborator from an organization's workspaces, moving a member to another organization and deleting an organization now queue it too, in the same transaction, with the worker's new organization_member and organization targets. * test(mothership): prove offboarding queues memory erasure with a real database The unit tests only checked mocked outbox calls. The integration test runs member removal, external removal and organization cleanup against Postgres: each committed change queues its erasure, a rolled-back one queues nothing. * test(mothership): cover a successful external collaborator removal
1 parent b3f9298 commit 09ed990

4 files changed

Lines changed: 168 additions & 1 deletion

File tree

‎apps/sim/lib/billing/organizations/membership.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ import { isRetryableTransactionError } from '@/lib/db/transaction'
6161
import type { DbOrTx, DbTransaction } from '@/lib/db/types'
6262
import { acquireInvitationMutationLocks } from '@/lib/invitations/locks'
6363
import { type StoppedRun, stopLiveRunsTx } from '@/lib/mothership/async-runs/repository'
64+
import { enqueueMemoryCleanup } from '@/lib/mothership/memories/cleanup-outbox'
6465
import { requireMemberManagementAuthority } from '@/lib/organizations/members/authority'
6566
import {
6667
revokePersonalApiKeysTx,
@@ -1192,6 +1193,13 @@ export async function transferUserBetweenOrganizations(
11921193
if (deleted.length === 0) {
11931194
throw new Error('Member could not be transferred because their role changed')
11941195
}
1196+
await enqueueMemoryCleanup(tx, [
1197+
{
1198+
kind: 'organization_member',
1199+
organizationId: params.sourceOrganizationId,
1200+
userId: params.userId,
1201+
},
1202+
])
11951203

11961204
const cancelledInvitations = invitationIds.length
11971205
? await tx
@@ -1407,6 +1415,7 @@ export async function removeUserFromOrganization(
14071415
'The membership changed before removal. Refresh and try again.'
14081416
)
14091417
}
1418+
await enqueueMemoryCleanup(tx, [{ kind: 'organization_member', organizationId, userId }])
14101419

14111420
if (!skipBillingLogic) {
14121421
await enqueueOutboxEvent(tx, MEMBER_BILLING_RECONCILIATION_EVENT_TYPE, {
@@ -1679,6 +1688,7 @@ export async function removeExternalUserFromOrganizationWorkspaces(params: {
16791688
)
16801689

16811690
await setOrgMemberUsageLimit(organizationId, userId, null, undefined, tx)
1691+
await enqueueMemoryCleanup(tx, [{ kind: 'organization_member', organizationId, userId }])
16821692

16831693
const cancelledInvitations = invitationIds.length
16841694
? await tx

‎apps/sim/lib/mothership/generated/memory-notes-access.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,14 @@ export const MemoryCleanupInput = z.strictObject({
5353
z.strictObject({ kind: z.literal("private_owner"), userId: z.string().min(1).max(200) }),
5454
z.strictObject({ kind: z.literal("workspace"), workspaceId: z.uuid() }),
5555
z.strictObject({ kind: z.literal("project"), projectId: z.uuid() }),
56+
// Erase private memory held in an organization: all of it when the organization is
57+
// deleted, one person's when they leave it (D106).
58+
z.strictObject({ kind: z.literal("organization"), organizationId: z.string().min(1).max(200) }),
59+
z.strictObject({
60+
kind: z.literal("organization_member"),
61+
organizationId: z.string().min(1).max(200),
62+
userId: z.string().min(1).max(200),
63+
}),
5664
]),
5765
});
5866
export type MemoryCleanupInput = z.infer<typeof MemoryCleanupInput>;

‎apps/sim/lib/mothership/memories/offboarding.integration.ts‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
import { db } from '@sim/db'
2+
import { member, organization, outboxEvent, permissions, user, workspace } from '@sim/db/schema'
3+
import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure'
4+
import { deleteWorkspaceFixture, insertWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
5+
import { generateId } from '@sim/utils/id'
6+
import { and, eq, inArray, sql } from 'drizzle-orm'
7+
import { afterAll, beforeAll, expect, it } from 'vitest'
8+
9+
readTestDatabaseUrl()
10+
11+
const ownerId = generateId()
12+
const memberId = generateId()
13+
const externalId = generateId()
14+
const organizationId = generateId()
15+
const membershipId = generateId()
16+
const workspaceId = generateId()
17+
let runtime: {
18+
membership: typeof import('@/lib/billing/organizations/membership')
19+
resources: typeof import('@/lib/organizations/resource-cleanup')
20+
event: typeof import('@/lib/mothership/memories/cleanup-event')
21+
}
22+
23+
beforeAll(async () => {
24+
// The cleanup outbox reads its worker credential and destination when it loads.
25+
process.env.SIM_AGENT_API_URL = 'http://127.0.0.1:9'
26+
process.env.COPILOT_API_KEY = 'memory-offboarding-fixture-key'
27+
const [membership, resources, event] = await Promise.all([
28+
import('@/lib/billing/organizations/membership'),
29+
import('@/lib/organizations/resource-cleanup'),
30+
import('@/lib/mothership/memories/cleanup-event'),
31+
])
32+
runtime = { membership, resources, event }
33+
await db.insert(user).values(
34+
[ownerId, memberId, externalId].map((id) => ({
35+
id,
36+
name: 'Offboarding fixture',
37+
email: `${id}@memory.test`,
38+
emailVerified: true,
39+
createdAt: new Date(),
40+
updatedAt: new Date(),
41+
}))
42+
)
43+
await db
44+
.insert(organization)
45+
.values({ id: organizationId, name: 'Offboarding', slug: organizationId })
46+
await db.insert(member).values([
47+
{ id: generateId(), organizationId, userId: ownerId, role: 'owner', createdAt: new Date() },
48+
{ id: membershipId, organizationId, userId: memberId, role: 'member', createdAt: new Date() },
49+
])
50+
// The external collaborator works in one of the organization's workspaces without membership.
51+
await insertWorkspaceFixture(db, {
52+
id: workspaceId,
53+
name: 'Offboarding',
54+
ownerId,
55+
billedAccountUserId: ownerId,
56+
organizationId,
57+
workspaceMode: 'organization',
58+
})
59+
await db.insert(permissions).values({
60+
id: generateId(),
61+
userId: externalId,
62+
entityType: 'workspace',
63+
entityId: workspaceId,
64+
permissionType: 'write',
65+
})
66+
}, 60_000)
67+
68+
afterAll(async () => {
69+
await db
70+
.delete(outboxEvent)
71+
.where(
72+
and(
73+
eq(outboxEvent.eventType, runtime.event.MEMORY_CLEANUP_EVENT),
74+
sql`${outboxEvent.payload}->'request'->'target'->>'organizationId' = ${organizationId}`
75+
)
76+
)
77+
await deleteWorkspaceFixture(db, eq(workspace.id, workspaceId))
78+
await db.delete(organization).where(eq(organization.id, organizationId))
79+
await db.delete(user).where(inArray(user.id, [ownerId, memberId, externalId]))
80+
})
81+
82+
/** The memory erasures queued for this organization, as `kind:userId` (or `kind`). */
83+
async function queued(): Promise<string[]> {
84+
const rows = await db
85+
.select({ payload: outboxEvent.payload })
86+
.from(outboxEvent)
87+
.where(
88+
and(
89+
eq(outboxEvent.eventType, runtime.event.MEMORY_CLEANUP_EVENT),
90+
sql`${outboxEvent.payload}->'request'->'target'->>'organizationId' = ${organizationId}`
91+
)
92+
)
93+
return rows
94+
.map(({ payload }) => {
95+
const target = (payload as { request: { target: { kind: string; userId?: string } } }).request
96+
.target
97+
return target.userId ? `${target.kind}:${target.userId}` : target.kind
98+
})
99+
.sort()
100+
}
101+
102+
it('queues private memory erasure with each committed departure and organization deletion, never on rollback', async () => {
103+
// A removal whose actor lacks authority rolls back with nothing queued.
104+
await expect(
105+
runtime.membership.removeUserFromOrganization({
106+
organizationId,
107+
userId: memberId,
108+
memberId: membershipId,
109+
actorUserId: externalId,
110+
skipBillingLogic: true,
111+
onError: 'throw',
112+
})
113+
).rejects.toThrow()
114+
expect(await queued()).toEqual([])
115+
116+
await runtime.membership.removeUserFromOrganization({
117+
organizationId,
118+
userId: memberId,
119+
memberId: membershipId,
120+
skipBillingLogic: true,
121+
onError: 'throw',
122+
})
123+
expect(
124+
await runtime.membership.removeExternalUserFromOrganizationWorkspaces({
125+
organizationId,
126+
userId: externalId,
127+
})
128+
).toMatchObject({ success: true, workspaceAccessRevoked: 1 })
129+
expect(await queued()).toEqual(
130+
[`organization_member:${externalId}`, `organization_member:${memberId}`].sort()
131+
)
132+
133+
// Organization deletion queues the erasure in the delete transaction: a rollback drops it.
134+
await expect(
135+
db.transaction(async (tx) => {
136+
await runtime.resources.enqueueOrganizationResourceCleanup(tx, organizationId)
137+
throw new Error('rollback fixture')
138+
})
139+
).rejects.toThrow('rollback fixture')
140+
expect(await queued()).not.toContain('organization')
141+
await db.transaction(async (tx) => {
142+
await runtime.resources.enqueueOrganizationResourceCleanup(tx, organizationId)
143+
})
144+
expect(await queued()).toContain('organization')
145+
}, 30_000)

‎apps/sim/lib/organizations/resource-cleanup.ts‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { env } from '@/lib/core/config/env'
88
import { enqueueOutboxEvent, type OutboxHandlerRegistry } from '@/lib/core/outbox/service'
99
import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
1010
import type { DbTransaction } from '@/lib/db/types'
11+
import { enqueueMemoryCleanup } from '@/lib/mothership/memories/cleanup-outbox'
1112
import { ORGANIZATION_RESOURCE_CLEANUP_EVENT } from '@/lib/organizations/resource-cleanup-event'
1213
import { deleteFile } from '@/lib/uploads/core/storage-service'
1314

@@ -35,7 +36,8 @@ const organizationResourceCleanupSchema = z.discriminatedUnion('kind', [
3536
* Captures canonical bindings before the org FK cascade removes them. The org
3637
* row lock excludes new direct resources; a rebound resource is checked again
3738
* by the worker. Each keyset page becomes its own bounded outbox event in the
38-
* caller's delete transaction, so rollback cannot leave cleanup work behind.
39+
* caller's delete transaction, so rollback cannot leave cleanup work behind. The
40+
* members' private Mothership memory in the organization is erased the same way.
3941
*/
4042
export async function enqueueOrganizationResourceCleanup(
4143
tx: DbTransaction,
@@ -105,6 +107,8 @@ export async function enqueueOrganizationResourceCleanup(
105107
chatCursor = chats[chats.length - 1].id
106108
if (chats.length < CLEANUP_BATCH_SIZE) break
107109
}
110+
111+
await enqueueMemoryCleanup(tx, [{ kind: 'organization', organizationId }])
108112
}
109113

110114
export const organizationResourceCleanupOutboxHandlers = {

0 commit comments

Comments
 (0)