Skip to content

Commit 9993a87

Browse files
authored
feat(auth,web): workspace picker at OAuth consent for multi-workspace users (#236)
* feat(auth,web): workspace picker at OAuth consent for multi-workspace users The consent page gains a workspace selector (shown only when the user has 2+ org memberships, defaulting to the oldest — today's implicit behavior). The choice is stored via a new POST /oauth2/workspace-choice endpoint and threaded into token issuance through the oauth-provider plugin's postLogin.consentReferenceId hook: the returned ws:<slug> keys consent records per (user, client, workspace) — so changing the choice naturally re-triggers consent — and rides the grant into access/refresh token rows, where customAccessTokenClaims overrides the oldest-membership default. Single-workspace users and existing grants are untouched (referenceId stays null). Multi-workspace users re-consent once, seeing the picker. Closes #231 * fix(auth,web): consent picker defaults to the server-resolved workspace The page always POSTs the selection on Allow, so a client-side default (first org by org createdAt) could silently overwrite a stored choice or shift the token's workspace for a user who touched nothing. New GET /oauth2/workspace-choice returns the AS's own resolution (stored choice if still a live membership, else oldest membership); the picker defaults to it, keeping org ordering as display order only. Also documents why the choice row is keyed by user rather than per client/grant: the consentReferenceId hook receives no client or request context, and per-grant scoping already holds via the consent-row keying.
1 parent 42b866d commit 9993a87

10 files changed

Lines changed: 830 additions & 5 deletions

File tree

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
-- Per-grant workspace choice (issue #231, auth side): stores a
2+
-- multi-workspace user's last explicit workspace pick for the OAuth 2.1
3+
-- authorization server. One row per user, keyed by user_id. Not a
4+
-- @better-auth/oauth-provider table -- read/written by
5+
-- src/workspace-choice.ts (see src/schema.ts:oauthWorkspaceChoice). Timestamps
6+
-- are integer epoch ms, matching every other hand-synced table in this
7+
-- schema (see migrations/20260717000000_oauth_provider.sql).
8+
9+
CREATE TABLE oauth_workspace_choice (
10+
user_id TEXT PRIMARY KEY,
11+
workspace TEXT NOT NULL,
12+
created_at INTEGER NOT NULL,
13+
updated_at INTEGER NOT NULL
14+
);

apps/auth/src/auth.ts

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,11 @@ import { sendAuthEmail } from "./email";
2323
import { localDemoEnabled, localDemoPlugin } from "./local-demo";
2424
import * as schema from "./schema";
2525
import { authTrustedOrigins, isTrustedOrigin } from "./trusted-origins";
26+
import {
27+
applyWorkspaceChoice,
28+
resolveWorkspaceChoiceReferenceId,
29+
workspaceChoicePlugin,
30+
} from "./workspace-choice";
2631
import {
2732
resolveDashApiKey,
2833
resolveGitHubCredentials,
@@ -257,11 +262,33 @@ function buildAuth(
257262
// auditable and can't silently drift. Enforced only when Better
258263
// Auth's core rate limiter is on (see rateLimit below).
259264
rateLimit: { register: { window: 60, max: 5 } },
265+
// Issue #231 (auth side): lets a multi-workspace user's consent (and
266+
// the resulting tokens) be scoped to a specific workspace instead of
267+
// always the oldest membership. The plugin recomputes this at
268+
// authorize-time and filters its `oauth_consent` lookup by the
269+
// returned string — a changed choice naturally re-triggers consent.
270+
// `undefined` for 0/1-membership users preserves today's
271+
// null-referenceId behavior (see src/workspace-choice.ts).
272+
postLogin: {
273+
// Never used: `shouldRedirect` below always returns false, so
274+
// `/oauth2/authorize` never redirects here. The workspace picker
275+
// itself lives on /oauth/consent (issue #231's web-side half);
276+
// this is just the required sibling field the plugin's types
277+
// demand alongside `consentReferenceId`.
278+
page: `${webOrigin}/oauth/consent`,
279+
consentReferenceId: ({ user }) => resolveWorkspaceChoiceReferenceId(db, user?.id),
280+
shouldRedirect: () => false,
281+
},
260282
// member ⋈ organization, oldest membership wins for `workspace`; all
261-
// slugs ride along in `workspaces` for a future workspace picker
262-
// (design doc: "deferred"). Zero memberships still issues a token
263-
// (workspace: null) — the MCP worker is responsible for the 403.
264-
customAccessTokenClaims: async ({ user }) => resolveWorkspaceClaims(db, user?.id),
283+
// slugs ride along in `workspaces`. Issue #231 (auth side): when
284+
// `referenceId` is one of ours (`ws:<slug>`, see
285+
// postLogin.consentReferenceId above) and `<slug>` is still one of
286+
// the user's workspaces, it overrides the oldest-membership default
287+
// — the user's per-grant choice wins. Zero memberships still issues
288+
// a token (workspace: null) — the MCP worker is responsible for the
289+
// 403.
290+
customAccessTokenClaims: async ({ user, referenceId }) =>
291+
applyWorkspaceChoice(await resolveWorkspaceClaims(db, user?.id), referenceId),
265292
}),
266293
// D5/Phase 4: bearer() lets the CLI present the device-flow session token
267294
// as `Authorization: Bearer <token>` so apps/api's session verification
@@ -293,6 +320,10 @@ function buildAuth(
293320
verificationUri: `${webOrigin}/device`,
294321
validateClient: (clientId) => clientId === UPLOADS_CLI_CLIENT_ID,
295322
}),
323+
// Issue #231 (auth side): POST /oauth2/workspace-choice, letting a
324+
// signed-in multi-workspace user record which workspace an OAuth grant
325+
// should operate on (read back by postLogin.consentReferenceId above).
326+
workspaceChoicePlugin(db),
296327
// Hosted dashboard (`@better-auth/infra`). Omit when the API key is unset.
297328
...(dashApiKey ? [dash({ apiKey: dashApiKey })] : []),
298329
// This endpoint is omitted entirely unless the lifecycle runner supplies

apps/auth/src/schema.ts

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -374,6 +374,26 @@ export const oauthConsent = sqliteTable(
374374
(t) => [index("idx_oauth_consent_user_client").on(t.userId, t.clientId)],
375375
);
376376

377+
/**
378+
* Per-grant workspace choice (issue #231, auth side): the user's last
379+
* explicit workspace pick for the OAuth 2.1 authorization server, one row per
380+
* user. Custom table (not part of `@better-auth/oauth-provider`) — read by
381+
* `resolveWorkspaceChoiceReferenceId` (src/workspace-choice.ts) to build the
382+
* `postLogin.consentReferenceId` the plugin ties to `oauth_consent` rows, and
383+
* written by `POST /api/auth/oauth2/workspace-choice`. Single-workspace users
384+
* never get a row (the hook returns `undefined` for them, see
385+
* src/workspace-choice.ts) so this table only ever holds multi-workspace
386+
* users' picks.
387+
*
388+
* Paired migration: `migrations/20260718000000_oauth_workspace_choice.sql`.
389+
*/
390+
export const oauthWorkspaceChoice = sqliteTable("oauth_workspace_choice", {
391+
userId: text("user_id").primaryKey(),
392+
workspace: text("workspace").notNull(),
393+
createdAt: timestampCol("created_at"),
394+
updatedAt: timestampCol("updated_at"),
395+
});
396+
377397
/**
378398
* Drizzle relations for Better Auth `experimental.joins` (adapter needs these
379399
* on the same schema object as the tables). No SQL/migration impact.
@@ -451,3 +471,4 @@ export type AuthMember = typeof member.$inferSelect;
451471
export type AuthInvitation = typeof invitation.$inferSelect;
452472
export type AuthDeviceCode = typeof deviceCode.$inferSelect;
453473
export type AuthOauthClient = typeof oauthClient.$inferSelect;
474+
export type AuthOauthWorkspaceChoice = typeof oauthWorkspaceChoice.$inferSelect;
Lines changed: 290 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,290 @@
1+
/**
2+
* Issue #231 (auth side): per-grant workspace choice. Covers the
3+
* `POST /oauth2/workspace-choice` endpoint (driven against the real Better
4+
* Auth handler via src/index.ts's `app`, same pattern as device.test.ts),
5+
* the `resolveWorkspaceChoiceReferenceId` postLogin hook, and the
6+
* `applyWorkspaceChoice` claims override — see src/workspace-choice.ts.
7+
*/
8+
import { eq } from "drizzle-orm";
9+
import { drizzle } from "drizzle-orm/d1";
10+
import { describe, expect, it } from "vitest";
11+
import type { AuthEnv } from "./auth";
12+
import { app } from "./index";
13+
import * as schema from "./schema";
14+
import { createFakeD1 } from "./test/fake-d1";
15+
import { applyWorkspaceChoice, resolveWorkspaceChoiceReferenceId } from "./workspace-choice";
16+
17+
function dbEnv(overrides: Partial<AuthEnv> = {}): AuthEnv {
18+
return {
19+
DB: createFakeD1(),
20+
WEB_ORIGIN: "https://uploads.sh",
21+
BETTER_AUTH_URL: "https://auth.uploads.sh",
22+
ENVIRONMENT: "development",
23+
BETTER_AUTH_SECRET_DEV: "test-signing-secret-at-least-32-chars-long",
24+
...overrides,
25+
};
26+
}
27+
28+
/** Seed a user + an active session, returning the raw session token to present as a bearer (pattern: device.test.ts). */
29+
async function seedSignedInUser(
30+
env: AuthEnv,
31+
orgSlugs: { slug: string; createdAt: Date }[] = [],
32+
): Promise<{ userId: string; sessionToken: string }> {
33+
const orm = drizzle(env.DB, { schema });
34+
const userId = crypto.randomUUID();
35+
await orm.insert(schema.user).values({
36+
id: userId,
37+
name: "Ada Lovelace",
38+
email: `ada-${userId}@example.com`,
39+
emailVerified: true,
40+
createdAt: new Date(),
41+
updatedAt: new Date(),
42+
role: "user",
43+
});
44+
const sessionToken = `sess-${crypto.randomUUID()}`;
45+
await orm.insert(schema.session).values({
46+
id: crypto.randomUUID(),
47+
userId,
48+
token: sessionToken,
49+
expiresAt: new Date(Date.now() + 60 * 60 * 1000),
50+
createdAt: new Date(),
51+
updatedAt: new Date(),
52+
});
53+
for (const { slug, createdAt } of orgSlugs) {
54+
const orgId = crypto.randomUUID();
55+
await orm.insert(schema.organization).values({ id: orgId, name: slug, slug, createdAt });
56+
await orm.insert(schema.member).values({
57+
id: crypto.randomUUID(),
58+
organizationId: orgId,
59+
userId,
60+
role: "member",
61+
createdAt,
62+
});
63+
}
64+
return { userId, sessionToken };
65+
}
66+
67+
function requestWorkspaceChoice(env: AuthEnv, body: unknown, sessionToken?: string) {
68+
return app.request(
69+
"/api/auth/oauth2/workspace-choice",
70+
{
71+
method: "POST",
72+
headers: {
73+
"content-type": "application/json",
74+
...(sessionToken ? { Authorization: `Bearer ${sessionToken}` } : {}),
75+
},
76+
body: JSON.stringify(body),
77+
},
78+
env,
79+
);
80+
}
81+
82+
describe("POST /oauth2/workspace-choice", () => {
83+
it("401s when unauthenticated", async () => {
84+
const res = await requestWorkspaceChoice(dbEnv(), { workspace: "acme" });
85+
expect(res.status).toBe(401);
86+
});
87+
88+
it("400s with code invalid_workspace for a non-membership slug", async () => {
89+
const env = dbEnv();
90+
const { sessionToken } = await seedSignedInUser(env, [
91+
{ slug: "acme", createdAt: new Date("2026-01-01T00:00:00Z") },
92+
]);
93+
const res = await requestWorkspaceChoice(env, { workspace: "not-a-member-org" }, sessionToken);
94+
expect(res.status).toBe(400);
95+
const body = (await res.json()) as { code?: string };
96+
expect(body.code).toBe("invalid_workspace");
97+
});
98+
99+
it("400s with code invalid_workspace for a missing/malformed body", async () => {
100+
const env = dbEnv();
101+
const { sessionToken } = await seedSignedInUser(env, [
102+
{ slug: "acme", createdAt: new Date("2026-01-01T00:00:00Z") },
103+
]);
104+
const res = await requestWorkspaceChoice(env, {}, sessionToken);
105+
expect(res.status).toBe(400);
106+
const body = (await res.json()) as { code?: string };
107+
expect(body.code).toBe("invalid_workspace");
108+
});
109+
110+
it("upserts the choice and returns { status: true } on a valid membership", async () => {
111+
const env = dbEnv();
112+
const { userId, sessionToken } = await seedSignedInUser(env, [
113+
{ slug: "acme", createdAt: new Date("2026-01-01T00:00:00Z") },
114+
{ slug: "beta", createdAt: new Date("2026-02-01T00:00:00Z") },
115+
]);
116+
117+
const res1 = await requestWorkspaceChoice(env, { workspace: "acme" }, sessionToken);
118+
expect(res1.status).toBe(200);
119+
expect(await res1.json()).toEqual({ status: true });
120+
121+
const orm = drizzle(env.DB, { schema });
122+
const [row1] = await orm
123+
.select()
124+
.from(schema.oauthWorkspaceChoice)
125+
.where(eq(schema.oauthWorkspaceChoice.userId, userId));
126+
expect(row1?.workspace).toBe("acme");
127+
128+
// Re-picking updates the same row rather than inserting a second one.
129+
const res2 = await requestWorkspaceChoice(env, { workspace: "beta" }, sessionToken);
130+
expect(res2.status).toBe(200);
131+
132+
const rows = await orm
133+
.select()
134+
.from(schema.oauthWorkspaceChoice)
135+
.where(eq(schema.oauthWorkspaceChoice.userId, userId));
136+
expect(rows).toHaveLength(1);
137+
expect(rows[0]?.workspace).toBe("beta");
138+
});
139+
});
140+
141+
function requestWorkspaceChoiceGet(env: AuthEnv, sessionToken?: string) {
142+
return app.request(
143+
"/api/auth/oauth2/workspace-choice",
144+
{
145+
method: "GET",
146+
headers: sessionToken ? { Authorization: `Bearer ${sessionToken}` } : {},
147+
},
148+
env,
149+
);
150+
}
151+
152+
describe("GET /oauth2/workspace-choice", () => {
153+
it("401s when unauthenticated", async () => {
154+
const res = await requestWorkspaceChoiceGet(dbEnv());
155+
expect(res.status).toBe(401);
156+
});
157+
158+
it("returns { workspace: null } for a user with zero memberships", async () => {
159+
const env = dbEnv();
160+
const { sessionToken } = await seedSignedInUser(env);
161+
const res = await requestWorkspaceChoiceGet(env, sessionToken);
162+
expect(res.status).toBe(200);
163+
expect(await res.json()).toEqual({ workspace: null });
164+
});
165+
166+
it("resolves to the oldest membership when no choice is stored", async () => {
167+
const env = dbEnv();
168+
const { sessionToken } = await seedSignedInUser(env, [
169+
{ slug: "beta", createdAt: new Date("2026-02-01T00:00:00Z") },
170+
{ slug: "acme", createdAt: new Date("2026-01-01T00:00:00Z") },
171+
]);
172+
const res = await requestWorkspaceChoiceGet(env, sessionToken);
173+
expect(await res.json()).toEqual({ workspace: "acme" });
174+
});
175+
176+
it("resolves to the stored choice while it is a live membership, else falls back", async () => {
177+
const env = dbEnv();
178+
const { userId, sessionToken } = await seedSignedInUser(env, [
179+
{ slug: "acme", createdAt: new Date("2026-01-01T00:00:00Z") },
180+
{ slug: "beta", createdAt: new Date("2026-02-01T00:00:00Z") },
181+
]);
182+
const orm = drizzle(env.DB, { schema });
183+
await orm.insert(schema.oauthWorkspaceChoice).values({
184+
userId,
185+
workspace: "beta",
186+
createdAt: new Date(),
187+
updatedAt: new Date(),
188+
});
189+
190+
const res = await requestWorkspaceChoiceGet(env, sessionToken);
191+
expect(await res.json()).toEqual({ workspace: "beta" });
192+
193+
// Stale choice (membership gone) falls back to the oldest live one.
194+
await orm
195+
.update(schema.oauthWorkspaceChoice)
196+
.set({ workspace: "departed-org" })
197+
.where(eq(schema.oauthWorkspaceChoice.userId, userId));
198+
const res2 = await requestWorkspaceChoiceGet(env, sessionToken);
199+
expect(await res2.json()).toEqual({ workspace: "acme" });
200+
});
201+
});
202+
203+
describe("resolveWorkspaceChoiceReferenceId", () => {
204+
it("returns undefined for an undefined user", async () => {
205+
const db = drizzle(createFakeD1(), { schema });
206+
expect(await resolveWorkspaceChoiceReferenceId(db, undefined)).toBeUndefined();
207+
});
208+
209+
it("returns undefined for a user with zero memberships", async () => {
210+
const env = dbEnv();
211+
const { userId } = await seedSignedInUser(env, []);
212+
const db = drizzle(env.DB, { schema });
213+
expect(await resolveWorkspaceChoiceReferenceId(db, userId)).toBeUndefined();
214+
});
215+
216+
it("returns undefined for a user with exactly one membership", async () => {
217+
const env = dbEnv();
218+
const { userId } = await seedSignedInUser(env, [
219+
{ slug: "solo-org", createdAt: new Date("2026-01-01T00:00:00Z") },
220+
]);
221+
const db = drizzle(env.DB, { schema });
222+
expect(await resolveWorkspaceChoiceReferenceId(db, userId)).toBeUndefined();
223+
});
224+
225+
it("returns ws:<oldest-slug> for a multi-workspace user with no stored choice", async () => {
226+
const env = dbEnv();
227+
const { userId } = await seedSignedInUser(env, [
228+
{ slug: "newer-org", createdAt: new Date("2026-02-01T00:00:00Z") },
229+
{ slug: "older-org", createdAt: new Date("2026-01-01T00:00:00Z") },
230+
]);
231+
const db = drizzle(env.DB, { schema });
232+
expect(await resolveWorkspaceChoiceReferenceId(db, userId)).toBe("ws:older-org");
233+
});
234+
235+
it("returns ws:<stored-choice> when the stored choice is still a live membership", async () => {
236+
const env = dbEnv();
237+
const { userId } = await seedSignedInUser(env, [
238+
{ slug: "newer-org", createdAt: new Date("2026-02-01T00:00:00Z") },
239+
{ slug: "older-org", createdAt: new Date("2026-01-01T00:00:00Z") },
240+
]);
241+
const db = drizzle(env.DB, { schema });
242+
const now = new Date();
243+
await db
244+
.insert(schema.oauthWorkspaceChoice)
245+
.values({ userId, workspace: "newer-org", createdAt: now, updatedAt: now });
246+
247+
expect(await resolveWorkspaceChoiceReferenceId(db, userId)).toBe("ws:newer-org");
248+
});
249+
250+
it("falls back to the oldest membership when the stored choice is no longer a live membership", async () => {
251+
const env = dbEnv();
252+
const { userId } = await seedSignedInUser(env, [
253+
{ slug: "newer-org", createdAt: new Date("2026-02-01T00:00:00Z") },
254+
{ slug: "older-org", createdAt: new Date("2026-01-01T00:00:00Z") },
255+
]);
256+
const db = drizzle(env.DB, { schema });
257+
const now = new Date();
258+
await db.insert(schema.oauthWorkspaceChoice).values({
259+
userId,
260+
workspace: "stale-org-no-longer-a-member",
261+
createdAt: now,
262+
updatedAt: now,
263+
});
264+
265+
expect(await resolveWorkspaceChoiceReferenceId(db, userId)).toBe("ws:older-org");
266+
});
267+
});
268+
269+
describe("applyWorkspaceChoice", () => {
270+
const claims = { workspace: "older-org", workspaces: ["older-org", "newer-org"] };
271+
272+
it("overrides workspace when referenceId names a workspace the user belongs to", () => {
273+
expect(applyWorkspaceChoice(claims, "ws:newer-org")).toEqual({
274+
workspace: "newer-org",
275+
workspaces: ["older-org", "newer-org"],
276+
});
277+
});
278+
279+
it("ignores a referenceId for a workspace the user does not belong to", () => {
280+
expect(applyWorkspaceChoice(claims, "ws:not-a-member-org")).toEqual(claims);
281+
});
282+
283+
it("ignores an undefined referenceId", () => {
284+
expect(applyWorkspaceChoice(claims, undefined)).toEqual(claims);
285+
});
286+
287+
it("ignores a referenceId that isn't one of ours (no ws: prefix)", () => {
288+
expect(applyWorkspaceChoice(claims, "some-other-reference")).toEqual(claims);
289+
});
290+
});

0 commit comments

Comments
 (0)