Skip to content

Commit abaa3d8

Browse files
mrubensmrubens
authored
Store the organization id in credentials (#5002)
* Store the organization id in credentials * Better organization logic * Fix tests * Update cloud settings defaults * Fix organization_id handling in Clerk API calls Address review feedback by properly handling 3 cases for organization_id: 1. Have an org id: send organization_id=THE_ORG_ID 2. Have a personal account: send organization_id= (empty string) 3. Don't know if you have an org id (old credentials): don't send organization_id param at all Changes: - Updated clerkCreateSessionToken() to check credentials.organizationId !== undefined - Updated fetchUserInfo() to handle all 3 cases consistently - Added fallback logic for old credentials without organization context - Improved logging for better debugging of organization context * DRY up organization loading code in AuthService Extract common organization membership processing logic into reusable helper methods: - findOrganizationMembership(): Find specific org membership by ID - findPrimaryOrganizationMembership(): Get first/primary org membership - setUserOrganizationInfo(): Set organization info on user object This eliminates duplication between the two clerkGetOrganizationMemberships() call sites that were doing very similar organization data processing.
1 parent 5bc3af1 commit abaa3d8

File tree

8 files changed

+250
-31
lines changed

8 files changed

+250
-31
lines changed

‎packages/cloud/src/AuthService.ts‎

Lines changed: 98 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ export interface AuthServiceEvents {
2121
const authCredentialsSchema = z.object({
2222
clientToken: z.string().min(1, "Client token cannot be empty"),
2323
sessionId: z.string().min(1, "Session ID cannot be empty"),
24+
organizationId: z.string().nullable().optional(),
2425
})
2526

2627
type AuthCredentials = z.infer<typeof authCredentialsSchema>
@@ -220,7 +221,16 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
220221

221222
try {
222223
const parsedJson = JSON.parse(credentialsJson)
223-
return authCredentialsSchema.parse(parsedJson)
224+
const credentials = authCredentialsSchema.parse(parsedJson)
225+
226+
// Migration: If no organizationId but we have userInfo, add it
227+
if (credentials.organizationId === undefined && this.userInfo?.organizationId) {
228+
credentials.organizationId = this.userInfo.organizationId
229+
await this.storeCredentials(credentials)
230+
this.log("[auth] Migrated credentials with organizationId")
231+
}
232+
233+
return credentials
224234
} catch (error) {
225235
if (error instanceof z.ZodError) {
226236
this.log("[auth] Invalid credentials format:", error.errors)
@@ -269,8 +279,13 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
269279
*
270280
* @param code The authorization code from the callback
271281
* @param state The state parameter from the callback
282+
* @param organizationId The organization ID from the callback (null for personal accounts)
272283
*/
273-
public async handleCallback(code: string | null, state: string | null): Promise<void> {
284+
public async handleCallback(
285+
code: string | null,
286+
state: string | null,
287+
organizationId?: string | null,
288+
): Promise<void> {
274289
if (!code || !state) {
275290
vscode.window.showInformationMessage("Invalid Roo Code Cloud sign in url")
276291
return
@@ -287,6 +302,9 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
287302

288303
const credentials = await this.clerkSignIn(code)
289304

305+
// Set organizationId (null for personal accounts)
306+
credentials.organizationId = organizationId || null
307+
290308
await this.storeCredentials(credentials)
291309

292310
vscode.window.showInformationMessage("Successfully authenticated with Roo Code Cloud")
@@ -417,6 +435,15 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
417435
return this.userInfo
418436
}
419437

438+
/**
439+
* Get the stored organization ID from credentials
440+
*
441+
* @returns The stored organization ID, null for personal accounts or if no credentials exist
442+
*/
443+
public getStoredOrganizationId(): string | null {
444+
return this.credentials?.organizationId || null
445+
}
446+
420447
private async clerkSignIn(ticket: string): Promise<AuthCredentials> {
421448
const formData = new URLSearchParams()
422449
formData.append("strategy", "ticket")
@@ -454,6 +481,17 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
454481
const formData = new URLSearchParams()
455482
formData.append("_is_native", "1")
456483

484+
// Handle 3 cases for organization_id:
485+
// 1. Have an org id: organization_id=THE_ORG_ID
486+
// 2. Have a personal account: organization_id= (empty string)
487+
// 3. Don't know if you have an org id (old style credentials): don't send organization_id param at all
488+
const organizationId = this.getStoredOrganizationId()
489+
if (this.credentials?.organizationId !== undefined) {
490+
// We have organization context info (either org id or personal account)
491+
formData.append("organization_id", organizationId || "")
492+
}
493+
// If organizationId is undefined, don't send the param at all (old credentials)
494+
457495
const response = await fetch(`${getClerkBaseUrl()}/v1/client/sessions/${this.credentials!.sessionId}/tokens`, {
458496
method: "POST",
459497
headers: {
@@ -505,29 +543,74 @@ export class AuthService extends EventEmitter<AuthServiceEvents> {
505543

506544
userInfo.picture = userData.image_url
507545

508-
// Fetch organization memberships separately
546+
// Fetch organization info if user is in organization context
509547
try {
510-
const orgMemberships = await this.clerkGetOrganizationMemberships()
511-
if (orgMemberships && orgMemberships.length > 0) {
512-
// Get the first (or active) organization membership
513-
const primaryOrgMembership = orgMemberships[0]
514-
const organization = primaryOrgMembership?.organization
515-
516-
if (organization) {
517-
userInfo.organizationId = organization.id
518-
userInfo.organizationName = organization.name
519-
userInfo.organizationRole = primaryOrgMembership.role
520-
userInfo.organizationImageUrl = organization.image_url
548+
const storedOrgId = this.getStoredOrganizationId()
549+
550+
if (this.credentials?.organizationId !== undefined) {
551+
// We have organization context info
552+
if (storedOrgId !== null) {
553+
// User is in organization context - fetch user's memberships and filter
554+
const orgMemberships = await this.clerkGetOrganizationMemberships()
555+
const userMembership = this.findOrganizationMembership(orgMemberships, storedOrgId)
556+
557+
if (userMembership) {
558+
this.setUserOrganizationInfo(userInfo, userMembership)
559+
this.log("[auth] User in organization context:", {
560+
id: userMembership.organization.id,
561+
name: userMembership.organization.name,
562+
role: userMembership.role,
563+
})
564+
} else {
565+
this.log("[auth] Warning: User not found in stored organization:", storedOrgId)
566+
}
567+
} else {
568+
this.log("[auth] User in personal account context - not setting organization info")
569+
}
570+
} else {
571+
// Old credentials without organization context - fetch organization info to determine context
572+
const orgMemberships = await this.clerkGetOrganizationMemberships()
573+
const primaryOrgMembership = this.findPrimaryOrganizationMembership(orgMemberships)
574+
575+
if (primaryOrgMembership) {
576+
this.setUserOrganizationInfo(userInfo, primaryOrgMembership)
577+
this.log("[auth] Legacy credentials: Found organization membership:", {
578+
id: primaryOrgMembership.organization.id,
579+
name: primaryOrgMembership.organization.name,
580+
role: primaryOrgMembership.role,
581+
})
582+
} else {
583+
this.log("[auth] Legacy credentials: No organization memberships found")
521584
}
522585
}
523586
} catch (error) {
524-
this.log("[auth] Failed to fetch organization memberships:", error)
587+
this.log("[auth] Failed to fetch organization info:", error)
525588
// Don't throw - organization info is optional
526589
}
527590

528591
return userInfo
529592
}
530593

594+
private findOrganizationMembership(
595+
memberships: CloudOrganizationMembership[],
596+
organizationId: string,
597+
): CloudOrganizationMembership | undefined {
598+
return memberships?.find((membership) => membership.organization.id === organizationId)
599+
}
600+
601+
private findPrimaryOrganizationMembership(
602+
memberships: CloudOrganizationMembership[],
603+
): CloudOrganizationMembership | undefined {
604+
return memberships && memberships.length > 0 ? memberships[0] : undefined
605+
}
606+
607+
private setUserOrganizationInfo(userInfo: CloudUserInfo, membership: CloudOrganizationMembership): void {
608+
userInfo.organizationId = membership.organization.id
609+
userInfo.organizationName = membership.organization.name
610+
userInfo.organizationRole = membership.role
611+
userInfo.organizationImageUrl = membership.organization.image_url
612+
}
613+
531614
private async clerkGetOrganizationMemberships(): Promise<CloudOrganizationMembership[]> {
532615
const response = await fetch(`${getClerkBaseUrl()}/v1/me/organization_memberships`, {
533616
headers: {

‎packages/cloud/src/CloudService.ts‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,14 +118,28 @@ export class CloudService {
118118
return userInfo?.organizationRole || null
119119
}
120120

121+
public hasStoredOrganizationId(): boolean {
122+
this.ensureInitialized()
123+
return this.authService!.getStoredOrganizationId() !== null
124+
}
125+
126+
public getStoredOrganizationId(): string | null {
127+
this.ensureInitialized()
128+
return this.authService!.getStoredOrganizationId()
129+
}
130+
121131
public getAuthState(): string {
122132
this.ensureInitialized()
123133
return this.authService!.getState()
124134
}
125135

126-
public async handleAuthCallback(code: string | null, state: string | null): Promise<void> {
136+
public async handleAuthCallback(
137+
code: string | null,
138+
state: string | null,
139+
organizationId?: string | null,
140+
): Promise<void> {
127141
this.ensureInitialized()
128-
return this.authService!.handleCallback(code, state)
142+
return this.authService!.handleCallback(code, state, organizationId)
129143
}
130144

131145
// SettingsService

‎packages/cloud/src/__tests__/AuthService.spec.ts‎

Lines changed: 52 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -328,7 +328,7 @@ describe("AuthService", () => {
328328

329329
expect(mockContext.secrets.store).toHaveBeenCalledWith(
330330
"clerk-auth-credentials",
331-
JSON.stringify({ clientToken: "Bearer token-123", sessionId: "session-123" }),
331+
JSON.stringify({ clientToken: "Bearer token-123", sessionId: "session-123", organizationId: null }),
332332
)
333333
expect(mockShowInfo).toHaveBeenCalledWith("Successfully authenticated with Roo Code Cloud")
334334
})
@@ -633,9 +633,55 @@ describe("AuthService", () => {
633633
expect(authService.getUserInfo()).toBeNull()
634634
})
635635

636-
it("should parse user info correctly", async () => {
637-
// Set up with credentials
638-
const credentials = { clientToken: "test-token", sessionId: "test-session" }
636+
it("should parse user info correctly for personal accounts", async () => {
637+
// Set up with credentials for personal account (no organizationId)
638+
const credentials = { clientToken: "test-token", sessionId: "test-session", organizationId: null }
639+
mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials))
640+
await authService.initialize()
641+
642+
// Clear previous mock calls
643+
mockFetch.mockClear()
644+
645+
// Mock successful responses
646+
mockFetch
647+
.mockResolvedValueOnce({
648+
ok: true,
649+
json: () => Promise.resolve({ jwt: "jwt-token" }),
650+
})
651+
.mockResolvedValueOnce({
652+
ok: true,
653+
json: () =>
654+
Promise.resolve({
655+
response: {
656+
first_name: "Jane",
657+
last_name: "Smith",
658+
image_url: "https://example.com/jane.jpg",
659+
primary_email_address_id: "email-2",
660+
email_addresses: [
661+
{ id: "email-1", email_address: "[email protected]" },
662+
{ id: "email-2", email_address: "[email protected]" },
663+
],
664+
},
665+
}),
666+
})
667+
668+
const timerCallback = vi.mocked(RefreshTimer).mock.calls[0][0].callback
669+
await timerCallback()
670+
671+
// Wait for async operations to complete
672+
await new Promise((resolve) => setTimeout(resolve, 0))
673+
674+
const userInfo = authService.getUserInfo()
675+
expect(userInfo).toEqual({
676+
name: "Jane Smith",
677+
email: "[email protected]",
678+
picture: "https://example.com/jane.jpg",
679+
})
680+
})
681+
682+
it("should parse user info correctly for organization accounts", async () => {
683+
// Set up with credentials for organization account
684+
const credentials = { clientToken: "test-token", sessionId: "test-session", organizationId: "org_1" }
639685
mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials))
640686
await authService.initialize()
641687

@@ -699,8 +745,8 @@ describe("AuthService", () => {
699745
})
700746

701747
it("should handle missing user info fields", async () => {
702-
// Set up with credentials
703-
const credentials = { clientToken: "test-token", sessionId: "test-session" }
748+
// Set up with credentials for personal account (no organizationId)
749+
const credentials = { clientToken: "test-token", sessionId: "test-session", organizationId: null }
704750
mockContext.secrets.get.mockResolvedValue(JSON.stringify(credentials))
705751
await authService.initialize()
706752

‎packages/cloud/src/__tests__/CloudService.test.ts‎

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ describe("CloudService", () => {
4040
getState: ReturnType<typeof vi.fn>
4141
getSessionToken: ReturnType<typeof vi.fn>
4242
handleCallback: ReturnType<typeof vi.fn>
43+
getStoredOrganizationId: ReturnType<typeof vi.fn>
4344
on: ReturnType<typeof vi.fn>
4445
off: ReturnType<typeof vi.fn>
4546
once: ReturnType<typeof vi.fn>
@@ -88,6 +89,7 @@ describe("CloudService", () => {
8889
getState: vi.fn().mockReturnValue("logged-out"),
8990
getSessionToken: vi.fn(),
9091
handleCallback: vi.fn(),
92+
getStoredOrganizationId: vi.fn().mockReturnValue(null),
9193
on: vi.fn(),
9294
off: vi.fn(),
9395
once: vi.fn(),
@@ -255,7 +257,41 @@ describe("CloudService", () => {
255257

256258
it("should delegate handleAuthCallback to AuthService", async () => {
257259
await cloudService.handleAuthCallback("code", "state")
258-
expect(mockAuthService.handleCallback).toHaveBeenCalledWith("code", "state")
260+
expect(mockAuthService.handleCallback).toHaveBeenCalledWith("code", "state", undefined)
261+
})
262+
263+
it("should delegate handleAuthCallback with organizationId to AuthService", async () => {
264+
await cloudService.handleAuthCallback("code", "state", "org_123")
265+
expect(mockAuthService.handleCallback).toHaveBeenCalledWith("code", "state", "org_123")
266+
})
267+
268+
it("should return stored organization ID from AuthService", () => {
269+
mockAuthService.getStoredOrganizationId.mockReturnValue("org_456")
270+
271+
const result = cloudService.getStoredOrganizationId()
272+
expect(mockAuthService.getStoredOrganizationId).toHaveBeenCalled()
273+
expect(result).toBe("org_456")
274+
})
275+
276+
it("should return null when no stored organization ID available", () => {
277+
mockAuthService.getStoredOrganizationId.mockReturnValue(null)
278+
279+
const result = cloudService.getStoredOrganizationId()
280+
expect(result).toBe(null)
281+
})
282+
283+
it("should return true when stored organization ID exists", () => {
284+
mockAuthService.getStoredOrganizationId.mockReturnValue("org_789")
285+
286+
const result = cloudService.hasStoredOrganizationId()
287+
expect(result).toBe(true)
288+
})
289+
290+
it("should return false when no stored organization ID exists", () => {
291+
mockAuthService.getStoredOrganizationId.mockReturnValue(null)
292+
293+
const result = cloudService.hasStoredOrganizationId()
294+
expect(result).toBe(false)
259295
})
260296
})
261297

‎packages/types/src/cloud.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ export const ORGANIZATION_ALLOW_ALL: OrganizationAllowList = {
125125
export const ORGANIZATION_DEFAULT: OrganizationSettings = {
126126
version: 0,
127127
cloudSettings: {
128+
recordTaskMessages: true,
128129
enableTaskSharing: true,
129130
taskShareExpirationDays: 30,
130131
},

‎src/activate/handleUri.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,13 @@ export const handleUri = async (uri: vscode.Uri) => {
3838
case "/auth/clerk/callback": {
3939
const code = query.get("code")
4040
const state = query.get("state")
41-
await CloudService.instance.handleAuthCallback(code, state)
41+
const organizationId = query.get("organizationId")
42+
43+
await CloudService.instance.handleAuthCallback(
44+
code,
45+
state,
46+
organizationId === "null" ? null : organizationId,
47+
)
4248
break
4349
}
4450
default:

0 commit comments

Comments
 (0)