Skip to content

Commit d16a3f3

Browse files
RhysSullivanRhysSullivan
authored
fix WorkOS list boundaries (#809)
1 parent 6f6e250 commit d16a3f3

5 files changed

Lines changed: 249 additions & 27 deletions

File tree

‎apps/cloud/src/auth/api-keys.node.test.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ describe("ApiKeyService.WorkOS", () => {
109109
stubWorkOS({
110110
listUserApiKeys: () =>
111111
Effect.succeed({
112+
object: "list" as const,
112113
data: [
113114
{
114115
id: "api_key_listed",
@@ -124,6 +125,10 @@ describe("ApiKeyService.WorkOS", () => {
124125
},
125126
},
126127
],
128+
listMetadata: {
129+
before: null,
130+
after: null,
131+
},
127132
}),
128133
createUserApiKey: () =>
129134
Effect.succeed({

‎apps/cloud/src/auth/workos.node.test.ts‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
import { describe, expect, it } from "@effect/vitest";
2+
3+
import { collectRawWorkOSList, collectWorkOSList } from "./workos";
4+
5+
describe("collectWorkOSList", () => {
6+
it("collects memberships beyond the first WorkOS page", async () => {
7+
const autoPaginationCalls: string[] = [];
8+
9+
const response = await collectWorkOSList({
10+
object: "list",
11+
data: [{ id: "om_first_page" }],
12+
listMetadata: {
13+
before: null,
14+
after: "om_next_page",
15+
},
16+
autoPagination: async () => {
17+
autoPaginationCalls.push("called");
18+
return [{ id: "om_first_page" }, { id: "om_second_page" }];
19+
},
20+
});
21+
22+
expect(response.data).toEqual([{ id: "om_first_page" }, { id: "om_second_page" }]);
23+
expect(response.listMetadata).toEqual({ before: null, after: null });
24+
expect(autoPaginationCalls).toEqual(["called"]);
25+
});
26+
27+
it("keeps the first page when WorkOS reports no next page", async () => {
28+
let autoPaginationCalls = 0;
29+
30+
const response = await collectWorkOSList({
31+
object: "list",
32+
data: [{ id: "om_only_page" }],
33+
listMetadata: {
34+
before: null,
35+
after: null,
36+
},
37+
autoPagination: async () => {
38+
autoPaginationCalls += 1;
39+
return [{ id: "om_unexpected_page" }];
40+
},
41+
});
42+
43+
expect(response.data).toEqual([{ id: "om_only_page" }]);
44+
expect(response.listMetadata).toEqual({ before: null, after: null });
45+
expect(autoPaginationCalls).toBe(0);
46+
});
47+
});
48+
49+
describe("collectRawWorkOSList", () => {
50+
it("collects raw WorkOS lists using snake-case cursors", async () => {
51+
const requestedCursors: Array<string | undefined> = [];
52+
53+
const response = await collectRawWorkOSList(async (after) => {
54+
requestedCursors.push(after);
55+
return after
56+
? {
57+
data: [{ id: "api_key_second_page" }],
58+
list_metadata: {
59+
before: null,
60+
after: null,
61+
},
62+
}
63+
: {
64+
data: [{ id: "api_key_first_page" }],
65+
list_metadata: {
66+
before: null,
67+
after: "api_key_second_page",
68+
},
69+
};
70+
});
71+
72+
expect(response.data).toEqual([{ id: "api_key_first_page" }, { id: "api_key_second_page" }]);
73+
expect(response.listMetadata).toEqual({ before: null, after: null });
74+
expect(requestedCursors).toEqual([undefined, "api_key_second_page"]);
75+
});
76+
77+
it("collects raw WorkOS lists using camel-case cursors", async () => {
78+
const response = await collectRawWorkOSList(async (after) =>
79+
after
80+
? {
81+
data: [{ id: "second" }],
82+
listMetadata: {
83+
before: null,
84+
after: null,
85+
},
86+
}
87+
: {
88+
data: [{ id: "first" }],
89+
listMetadata: {
90+
before: null,
91+
after: "second",
92+
},
93+
},
94+
);
95+
96+
expect(response.data).toEqual([{ id: "first" }, { id: "second" }]);
97+
});
98+
});

‎apps/cloud/src/auth/workos.ts‎

Lines changed: 130 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
// ---------------------------------------------------------------------------
44

55
import { env } from "cloudflare:workers";
6-
import { Context, Data, Effect, Layer } from "effect";
6+
import { Context, Data, Effect, Layer, Option, Schema } from "effect";
77
import { GeneratePortalLinkIntent, WorkOS } from "@workos-inc/node/worker";
88
import { WorkOSError, tryPromiseService, withServiceLogging } from "./errors";
99

@@ -24,6 +24,88 @@ type RawWorkOS = WorkOS & {
2424
) => Promise<{ readonly data: unknown }>;
2525
};
2626

27+
type WorkOSListMetadata = {
28+
readonly before?: string | null;
29+
readonly after?: string | null;
30+
};
31+
32+
type WorkOSAutoPaginatable<Resource> = {
33+
readonly object: "list";
34+
readonly data: Resource[];
35+
readonly listMetadata: WorkOSListMetadata;
36+
readonly autoPagination: () => Promise<Resource[]>;
37+
};
38+
39+
export type WorkOSCollectedList<Resource> = {
40+
readonly object: "list";
41+
readonly data: Resource[];
42+
readonly listMetadata: {
43+
readonly before: string | null;
44+
readonly after: string | null;
45+
};
46+
};
47+
48+
const RawWorkOSListMetadata = Schema.Struct({
49+
before: Schema.optional(Schema.NullOr(Schema.String)),
50+
after: Schema.optional(Schema.NullOr(Schema.String)),
51+
});
52+
53+
const RawWorkOSListResponse = Schema.Struct({
54+
data: Schema.Array(Schema.Unknown),
55+
listMetadata: Schema.optional(RawWorkOSListMetadata),
56+
list_metadata: Schema.optional(RawWorkOSListMetadata),
57+
});
58+
59+
const decodeRawWorkOSListResponse = Schema.decodeUnknownOption(RawWorkOSListResponse);
60+
61+
const completedListMetadata = {
62+
before: null,
63+
after: null,
64+
} as const;
65+
66+
const nextCursorFromRawList = (response: typeof RawWorkOSListResponse.Type): string | null =>
67+
response.listMetadata?.after ?? response.list_metadata?.after ?? null;
68+
69+
export const collectWorkOSList = async <Resource>(
70+
response: WorkOSAutoPaginatable<Resource>,
71+
): Promise<WorkOSCollectedList<Resource>> => {
72+
const data = response.listMetadata.after ? await response.autoPagination() : response.data;
73+
return {
74+
object: "list",
75+
data,
76+
listMetadata: completedListMetadata,
77+
};
78+
};
79+
80+
export const collectRawWorkOSList = async (
81+
loadPage: (after?: string) => Promise<unknown>,
82+
): Promise<WorkOSCollectedList<unknown>> => {
83+
const first = Option.getOrNull(decodeRawWorkOSListResponse(await loadPage()));
84+
if (!first) {
85+
return {
86+
object: "list",
87+
data: [],
88+
listMetadata: completedListMetadata,
89+
};
90+
}
91+
92+
const data = [...first.data];
93+
let after = nextCursorFromRawList(first);
94+
95+
while (after) {
96+
const next = Option.getOrNull(decodeRawWorkOSListResponse(await loadPage(after)));
97+
if (!next) break;
98+
data.push(...next.data);
99+
after = nextCursorFromRawList(next);
100+
}
101+
102+
return {
103+
object: "list",
104+
data,
105+
listMetadata: completedListMetadata,
106+
};
107+
};
108+
27109
class WorkOSAuthConfigurationError extends Data.TaggedError("WorkOSAuthConfigurationError")<{
28110
readonly message: string;
29111
}> {}
@@ -132,11 +214,13 @@ const make = Effect.gen(function* () {
132214

133215
/** List organization memberships for a user. */
134216
listUserMemberships: (userId: string) =>
135-
use((wos) =>
136-
wos.userManagement.listOrganizationMemberships({
137-
userId,
138-
statuses: ["active", "pending"],
139-
}),
217+
use(async (wos) =>
218+
collectWorkOSList(
219+
await wos.userManagement.listOrganizationMemberships({
220+
userId,
221+
statuses: ["active", "pending"],
222+
}),
223+
),
140224
),
141225

142226
/**
@@ -183,10 +267,16 @@ const make = Effect.gen(function* () {
183267
listUserApiKeys: (userId: string, organizationId: string) =>
184268
use(async (wos) => {
185269
const raw = wos as RawWorkOS;
186-
const response = await raw.get(`/user_management/users/${userId}/api_keys`, {
187-
query: { organization_id: organizationId },
270+
return collectRawWorkOSList(async (after) => {
271+
const response = await raw.get(`/user_management/users/${userId}/api_keys`, {
272+
query: {
273+
organization_id: organizationId,
274+
limit: 100,
275+
...(after ? { after } : {}),
276+
},
277+
});
278+
return response.data;
188279
});
189-
return response.data;
190280
}),
191281

192282
createUserApiKey: (params: { userId: string; organizationId: string; name: string }) =>
@@ -203,12 +293,25 @@ const make = Effect.gen(function* () {
203293

204294
/** List organization memberships with user details. */
205295
listOrgMembers: (organizationId: string) =>
206-
use((wos) =>
207-
wos.userManagement.listOrganizationMemberships({
296+
use(async (wos) =>
297+
collectWorkOSList(
298+
await wos.userManagement.listOrganizationMemberships({
299+
organizationId,
300+
statuses: ["active", "pending"],
301+
}),
302+
),
303+
),
304+
305+
/** Get a user's membership in an organization. */
306+
getUserOrgMembership: (organizationId: string, userId: string) =>
307+
use(async (wos) => {
308+
const response = await wos.userManagement.listOrganizationMemberships({
208309
organizationId,
310+
userId,
209311
statuses: ["active", "pending"],
210-
}),
211-
),
312+
});
313+
return response.data[0] ?? null;
314+
}),
212315

213316
/** Get a user by ID. */
214317
getUser: (userId: string) => use((wos) => wos.userManagement.getUser(userId)),
@@ -229,7 +332,13 @@ const make = Effect.gen(function* () {
229332
* API level, so we filter after.
230333
*/
231334
listPendingInvitations: (organizationId: string) =>
232-
use((wos) => wos.userManagement.listInvitations({ organizationId })).pipe(
335+
use(async (wos) =>
336+
collectWorkOSList(
337+
await wos.userManagement.listInvitations({
338+
organizationId,
339+
}),
340+
),
341+
).pipe(
233342
Effect.map((response) => ({
234343
...response,
235344
data: response.data.filter((i) => i.state === "pending"),
@@ -238,7 +347,13 @@ const make = Effect.gen(function* () {
238347

239348
/** List invitations for an email address (across all orgs). */
240349
listInvitationsByEmail: (email: string) =>
241-
use((wos) => wos.userManagement.listInvitations({ email })),
350+
use(async (wos) =>
351+
collectWorkOSList(
352+
await wos.userManagement.listInvitations({
353+
email,
354+
}),
355+
),
356+
),
242357

243358
/** Accept an invitation; returns the (now accepted) invitation. */
244359
acceptInvitation: (invitationId: string) =>

0 commit comments

Comments
 (0)