fix(credentials): restore reliable OAuth connections (#6753)

* fix(credentials): restore reliable OAuth connections

* fix(credentials): cover OAuth handoff edge cases

* fix(credentials): preserve reauthorization outcomes
This commit is contained in:
Theodore Li
2026-08-15 20:43:34 -04:00
committed by GitHub
parent cbbcca970c
commit 623c30f05c
33 changed files with 536 additions and 139 deletions
+6 -3
View File
@@ -249,14 +249,17 @@ describe('createHandoffManager', () => {
const manager = createHandoffManager(deps, makeCallbacks())
await manager.beginConnect('google-email', {
workspaceId: 'workspace-1',
draftId: 'draft-1',
chatAttemptId: 'attempt-1',
})
const state = new URL(vi.mocked(deps.openExternal).mock.calls[0][0]).searchParams.get(
'state'
) as string
const landing = new URL(vi.mocked(deps.openExternal).mock.calls[0][0])
const state = landing.searchParams.get('state') as string
expect(landing.searchParams.get('draftId')).toBe('draft-1')
expect(manager.consumeConnect(state)).toEqual({
workspaceId: 'workspace-1',
draftId: 'draft-1',
chatAttemptId: 'attempt-1',
})
expect(manager.consumeConnect(state)).toBeNull()
+2
View File
@@ -67,6 +67,7 @@ export interface HandoffManagerDeps {
export interface ConnectScope {
workspaceId?: string
credentialId?: string
draftId?: string
chatAttemptId?: string
}
@@ -292,6 +293,7 @@ export function createHandoffManager(
...(userId ? { user: userId } : {}),
...(scope.workspaceId ? { workspaceId: scope.workspaceId } : {}),
...(scope.credentialId ? { credentialId: scope.credentialId } : {}),
...(scope.draftId ? { draftId: scope.draftId } : {}),
},
scope
)
+4 -1
View File
@@ -335,21 +335,24 @@ describe('registerIpcHandlers', () => {
expect(await handler?.(appEvent, 'slack')).toBe(true)
expect(deps.beginOAuthConnect).toHaveBeenCalledWith('slack', {})
// Chip-initiated connects carry workspace/credential scope; malformed
// Connects carry workspace/credential or exact-draft scope; malformed
// scopes (wrong types, unsafe ids) are rejected before the handoff.
expect(
await handler?.(appEvent, 'slack', {
workspaceId: 'ws1',
credentialId: 'cred_1',
draftId: 'draft_1',
chatAttemptId: 'attempt_1',
})
).toBe(true)
expect(deps.beginOAuthConnect).toHaveBeenCalledWith('slack', {
workspaceId: 'ws1',
credentialId: 'cred_1',
draftId: 'draft_1',
chatAttemptId: 'attempt_1',
})
expect(await handler?.(appEvent, 'slack', { workspaceId: 'ws/../evil' })).toBe(false)
expect(await handler?.(appEvent, 'slack', { draftId: '../wrong' })).toBe(false)
expect(await handler?.(appEvent, 'slack', { chatAttemptId: '../wrong' })).toBe(false)
expect(await handler?.(appEvent, 'slack', 'not-an-object')).toBe(false)
})
+7 -1
View File
@@ -95,6 +95,7 @@ function parseDesktopScope(raw: unknown): string | null {
export interface OAuthConnectScope {
workspaceId?: string
credentialId?: string
draftId?: string
chatAttemptId?: string
}
@@ -110,9 +111,10 @@ export function parseOAuthConnectScope(raw: unknown): OAuthConnectScope | undefi
if (typeof raw !== 'object') {
return undefined
}
const { workspaceId, credentialId, chatAttemptId } = raw as {
const { workspaceId, credentialId, draftId, chatAttemptId } = raw as {
workspaceId?: unknown
credentialId?: unknown
draftId?: unknown
chatAttemptId?: unknown
}
if (
@@ -127,6 +129,9 @@ export function parseOAuthConnectScope(raw: unknown): OAuthConnectScope | undefi
) {
return undefined
}
if (draftId !== undefined && (typeof draftId !== 'string' || !ID_PATTERN.test(draftId))) {
return undefined
}
if (
chatAttemptId !== undefined &&
(typeof chatAttemptId !== 'string' || !ID_PATTERN.test(chatAttemptId))
@@ -136,6 +141,7 @@ export function parseOAuthConnectScope(raw: unknown): OAuthConnectScope | undefi
return {
...(workspaceId !== undefined ? { workspaceId } : {}),
...(credentialId !== undefined ? { credentialId } : {}),
...(draftId !== undefined ? { draftId } : {}),
...(chatAttemptId !== undefined ? { chatAttemptId } : {}),
}
}
@@ -11,7 +11,7 @@ import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
import { CredentialConnectionProviderMismatchError } from '@/lib/credentials/application/connection-target'
import { createCredentialConnection } from '@/lib/credentials/application/create-credential-connection'
import { launchCredentialConnection } from '@/lib/credentials/application/launch-credential-connection'
import { OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM } from '@/lib/credentials/draft-processor'
import { OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM } from '@/lib/credentials/draft-constants'
const logger = createLogger('OAuth2Authorize')
+15 -2
View File
@@ -50,8 +50,16 @@ export default async function DesktopConnectPage({ searchParams }: DesktopConnec
const port = parseLoopbackPort(typeof params.port === 'string' ? params.port : '')
const workspaceId = isValidOpaqueId(params.workspaceId) ? params.workspaceId : undefined
const credentialId = isValidOpaqueId(params.credentialId) ? params.credentialId : undefined
const draftId = isValidOpaqueId(params.draftId) ? params.draftId : undefined
const expectedUserId = isValidOpaqueId(params.user) ? params.user : undefined
if (!isValidOAuthProviderId(providerId) || !isValidHandoffState(state) || port === null) {
const hasInvalidDraftId = params.draftId !== undefined && draftId === undefined
if (
!isValidOAuthProviderId(providerId) ||
!isValidHandoffState(state) ||
port === null ||
hasInvalidDraftId ||
(workspaceId !== undefined && draftId !== undefined)
) {
return <InvalidRequest />
}
@@ -65,6 +73,7 @@ export default async function DesktopConnectPage({ searchParams }: DesktopConnec
buildDesktopConnectPath(providerId, state, port, {
workspaceId,
credentialId,
draftId,
user: expectedUserId,
})
)}`
@@ -86,6 +95,7 @@ export default async function DesktopConnectPage({ searchParams }: DesktopConnec
returnTo={buildDesktopConnectPath(providerId, state, port, {
workspaceId,
credentialId,
draftId,
user: expectedUserId,
})}
/>
@@ -113,6 +123,9 @@ export default async function DesktopConnectPage({ searchParams }: DesktopConnec
}
return (
<ConnectLauncher providerId={providerId} completePath={buildConnectCompletePath(state, port)} />
<ConnectLauncher
providerId={providerId}
completePath={buildConnectCompletePath(state, port, draftId)}
/>
)
}
@@ -2,6 +2,7 @@
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM } from '@/lib/credentials/draft-constants'
import {
buildConnectCompletePath,
buildConnectLoopbackUrl,
@@ -36,19 +37,23 @@ describe('sanitizeOAuthErrorSlug', () => {
describe('URL builders', () => {
it('buildDesktopConnectPath round-trips provider, state, and port', () => {
const path = buildDesktopConnectPath('google-email', STATE, 49152)
const path = buildDesktopConnectPath('google-email', STATE, 49152, {
draftId: 'draft-1',
})
const url = new URL(path, 'https://sim.ai')
expect(url.pathname).toBe('/desktop/connect')
expect(url.searchParams.get('provider')).toBe('google-email')
expect(url.searchParams.get('state')).toBe(STATE)
expect(url.searchParams.get('port')).toBe('49152')
expect(url.searchParams.get('draftId')).toBe('draft-1')
})
it('buildConnectCompletePath carries state and port', () => {
const url = new URL(buildConnectCompletePath(STATE, 49152), 'https://sim.ai')
it('buildConnectCompletePath carries state, port, and the exact credential draft', () => {
const url = new URL(buildConnectCompletePath(STATE, 49152, 'draft-1'), 'https://sim.ai')
expect(url.pathname).toBe('/desktop/connect/complete')
expect(url.searchParams.get('state')).toBe(STATE)
expect(url.searchParams.get('port')).toBe('49152')
expect(url.searchParams.get(OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM)).toBe('draft-1')
})
it('buildConnectLoopbackUrl targets the 127.0.0.1 connect callback, error optional', () => {
+6 -1
View File
@@ -1,3 +1,5 @@
import { OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM } from '@/lib/credentials/draft-constants'
/**
* OAuth providerIds are kebab-case service slugs (e.g. "google-email"). The
* value is only used to start a better-auth oauth2.link flow, which validates
@@ -36,6 +38,7 @@ export function isValidOpaqueId(value: unknown): value is string {
export interface ConnectScope {
workspaceId?: string
credentialId?: string
draftId?: string
/** The account the desktop app is signed in as; the flow is pinned to it. */
user?: string
}
@@ -53,6 +56,7 @@ export function buildDesktopConnectPath(
const params = new URLSearchParams({ provider: providerId, state, port: String(port) })
if (scope.workspaceId) params.set('workspaceId', scope.workspaceId)
if (scope.credentialId) params.set('credentialId', scope.credentialId)
if (scope.draftId) params.set('draftId', scope.draftId)
if (scope.user) params.set('user', scope.user)
return `/desktop/connect?${params.toString()}`
}
@@ -61,8 +65,9 @@ export function buildDesktopConnectPath(
* The same-origin path better-auth redirects the browser to after the OAuth
* callback — the complete page then bounces to the desktop app's loopback.
*/
export function buildConnectCompletePath(state: string, port: number): string {
export function buildConnectCompletePath(state: string, port: number, draftId?: string): string {
const params = new URLSearchParams({ state, port: String(port) })
if (draftId) params.set(OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM, draftId)
return `/desktop/connect/complete?${params.toString()}`
}
@@ -256,6 +256,7 @@ export function ConnectOAuthModal(props: ConnectOAuthModalProps) {
setSubmitError(null)
try {
let connectorType: string | undefined
let draftId: string | undefined
if (isConnect) {
const trimmed = displayName.trim()
@@ -264,12 +265,13 @@ export function ConnectOAuthModal(props: ConnectOAuthModalProps) {
return
}
await createDraft.mutateAsync({
const draft = await createDraft.mutateAsync({
workspaceId,
providerId,
displayName: trimmed,
description: description.trim() || undefined,
})
draftId = draft.draftId
const preCount = credentials.filter(
(c) => c.type === 'oauth' && c.providerId === providerId
@@ -279,6 +281,15 @@ export function ConnectOAuthModal(props: ConnectOAuthModalProps) {
displayName: trimmed,
providerId,
preCount,
baselineCredentials: credentials
.filter(
(credential) => credential.type === 'oauth' && credential.providerId === providerId
)
.map((credential) => ({
id: credential.id,
accountId: credential.accountId,
updatedAt: credential.updatedAt,
})),
workspaceId,
requestedAt: Date.now(),
}
@@ -319,6 +330,7 @@ export function ConnectOAuthModal(props: ConnectOAuthModalProps) {
await connectOAuthService.mutateAsync({
providerId,
callbackURL: callbackURL.toString(),
draftId,
})
handleClose()
} catch (err: unknown) {
@@ -111,7 +111,7 @@ export function ConnectedCredentialDetail({
const handleReconnectOAuth = async () => {
if (!credential || credential.type !== 'oauth' || !credential.providerId || !workspaceId) return
try {
await createDraft.mutateAsync({
const draft = await createDraft.mutateAsync({
workspaceId,
providerId: credential.providerId,
displayName: credential.displayName,
@@ -135,6 +135,7 @@ export function ConnectedCredentialDetail({
await connectOAuthService.mutateAsync({
providerId: credential.providerId,
callbackURL: window.location.href,
draftId: draft.draftId,
})
} catch (error: unknown) {
toast.error("Couldn't start reconnect", {
@@ -573,6 +573,7 @@ function ConnectorCard({
providerId: providerId!,
preCount: credentials?.length ?? 0,
workspaceId,
reconnect: true,
requestedAt: Date.now(),
})
}
@@ -607,6 +608,7 @@ function ConnectorCard({
providerId: providerId!,
preCount: credentials?.length ?? 0,
workspaceId,
reconnect: true,
requestedAt: Date.now(),
})
}
@@ -471,6 +471,7 @@ export function CredentialSelector({
providerId: effectiveProviderId,
preCount: credentials.filter((c) => c.type !== 'service_account').length,
workspaceId,
reconnect: true,
requestedAt: Date.now(),
})
setShowOAuthModal(true)
@@ -251,6 +251,7 @@ export function ToolCredentialSelector({
providerId: effectiveProviderId,
preCount: credentials.length,
workspaceId,
reconnect: true,
requestedAt: Date.now(),
})
setShowOAuthModal(true)
@@ -572,6 +572,7 @@ const WorkflowContent = React.memo(
providerId: detail.providerId,
preCount: 0,
workspaceId,
reconnect: true,
requestedAt: Date.now(),
})
+1 -1
View File
@@ -106,7 +106,7 @@ export function useWorkspaceCredential(credentialId?: string, enabled = true) {
export function useCreateCredentialDraft() {
return useMutation({
mutationFn: async (payload: ContractBodyInput<typeof createCredentialDraftContract>) => {
await requestJson(createCredentialDraftContract, { body: payload })
return requestJson(createCredentialDraftContract, { body: payload })
},
})
}
@@ -8,6 +8,7 @@ import {
type OAuthConnection,
} from '@/lib/api/contracts/oauth-connections'
import { client } from '@/lib/auth/auth-client'
import { OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM } from '@/lib/credentials/draft-constants'
import { getDesktopBridge } from '@/lib/desktop'
import { OAUTH_PROVIDERS, type OAuthServiceConfig } from '@/lib/oauth'
@@ -138,6 +139,7 @@ export function useOAuthConnections() {
interface ConnectServiceParams {
providerId: string
callbackURL: string
draftId?: string
}
/**
@@ -148,22 +150,25 @@ export function useConnectOAuthService() {
const queryClient = useQueryClient()
return useMutation({
mutationFn: async ({ providerId, callbackURL }: ConnectServiceParams) => {
mutationFn: async ({ providerId, callbackURL, draftId }: ConnectServiceParams) => {
if (providerId === 'trello') {
const returnUrl = encodeURIComponent(callbackURL)
window.location.href = `/api/auth/trello/authorize?returnUrl=${returnUrl}`
const draftQuery = draftId ? `&draftId=${encodeURIComponent(draftId)}` : ''
window.location.href = `/api/auth/trello/authorize?returnUrl=${returnUrl}${draftQuery}`
return { success: true }
}
if (providerId === 'instagram') {
const returnUrl = encodeURIComponent(callbackURL)
window.location.href = `/api/auth/instagram/authorize?returnUrl=${returnUrl}`
const draftQuery = draftId ? `&draftId=${encodeURIComponent(draftId)}` : ''
window.location.href = `/api/auth/instagram/authorize?returnUrl=${returnUrl}${draftQuery}`
return { success: true }
}
if (providerId === 'shopify') {
const returnUrl = encodeURIComponent(callbackURL)
window.location.href = `/api/auth/shopify/authorize?returnUrl=${returnUrl}`
const draftQuery = draftId ? `&draftId=${encodeURIComponent(draftId)}` : ''
window.location.href = `/api/auth/shopify/authorize?returnUrl=${returnUrl}${draftQuery}`
return { success: true }
}
@@ -175,16 +180,24 @@ export function useConnectOAuthService() {
// which refreshes caches and shows the connected toast.
const desktopBridge = getDesktopBridge()
if (desktopBridge?.beginOAuthConnect) {
const opened = await desktopBridge.beginOAuthConnect(providerId)
const opened = await desktopBridge.beginOAuthConnect(
providerId,
draftId ? { draftId } : undefined
)
if (!opened) {
throw new Error('Could not open your browser to connect this account.')
}
return { success: true }
}
const stateCallbackUrl = new URL(callbackURL)
if (draftId) {
stateCallbackUrl.searchParams.set(OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM, draftId)
}
await client.oauth2.link({
providerId,
callbackURL,
callbackURL: stateCallbackUrl.toString(),
})
return { success: true }
+99
View File
@@ -0,0 +1,99 @@
/**
* @vitest-environment node
*/
import { beforeEach, describe, expect, it, vi } from 'vitest'
const mocks = vi.hoisted(() => ({
requestJson: vi.fn(),
requireWorkspaceCredentialListResponse: vi.fn(),
}))
vi.mock('@sim/emcn', () => ({
toast: { error: vi.fn(), success: vi.fn() },
}))
vi.mock('next/navigation', () => ({
useParams: vi.fn(),
useRouter: vi.fn(),
}))
vi.mock('@/lib/api/client/request', () => ({ requestJson: mocks.requestJson }))
vi.mock('@/hooks/queries/utils/fetch-workspace-credentials', () => ({
requireWorkspaceCredentialListResponse: mocks.requireWorkspaceCredentialListResponse,
}))
import type { OAuthReturnContext } from '@/lib/credentials/client-state'
import { resolveOAuthMessage } from '@/hooks/use-oauth-return'
const context: OAuthReturnContext = {
origin: 'integrations',
displayName: 'New Gmail',
providerId: 'google-email',
preCount: 1,
baselineCredentials: [
{
id: 'credential-existing',
accountId: 'account-1',
updatedAt: '2026-08-14T17:00:00.000Z',
},
],
workspaceId: 'workspace-1',
requestedAt: Date.now(),
}
const existingCredential = {
id: 'credential-existing',
workspaceId: 'workspace-1',
type: 'oauth' as const,
displayName: 'Existing Gmail',
description: null,
providerId: 'google-email',
accountId: 'account-1',
envKey: null,
envOwnerUserId: null,
createdBy: 'user-1',
createdAt: '2026-08-01T00:00:00.000Z',
updatedAt: '2026-08-14T18:00:00.000Z',
}
describe('resolveOAuthMessage', () => {
beforeEach(() => {
vi.clearAllMocks()
mocks.requestJson.mockResolvedValue({})
})
it('recognizes an idempotent already-connected account from its reconnect timestamp', async () => {
mocks.requireWorkspaceCredentialListResponse.mockReturnValue([existingCredential])
await expect(resolveOAuthMessage(context)).resolves.toEqual({
kind: 'success',
text: 'This account is already connected as "Existing Gmail".',
})
})
it('keeps explicit update-access flows on the reconnect success path without a baseline', async () => {
await expect(
resolveOAuthMessage({
...context,
baselineCredentials: undefined,
reconnect: true,
})
).resolves.toEqual({
kind: 'success',
text: '"New Gmail" reconnected successfully.',
})
expect(mocks.requestJson).not.toHaveBeenCalled()
})
it('does not report success when the credential list is unchanged', async () => {
mocks.requireWorkspaceCredentialListResponse.mockReturnValue([
{
...existingCredential,
updatedAt: context.baselineCredentials?.[0].updatedAt,
},
])
await expect(resolveOAuthMessage(context)).resolves.toEqual({
kind: 'error',
text: 'We couldn’t verify the "New Gmail" connection. Try again.',
})
})
})
+49 -10
View File
@@ -31,9 +31,14 @@ const OAUTH_CREDENTIAL_UPDATED_EVENT = 'oauth-credentials-updated'
const SETTINGS_RETURN_URL_KEY = 'settings-return-url'
const CONTEXT_MAX_AGE_MS = 15 * 60 * 1000
async function resolveOAuthMessage(ctx: OAuthReturnContext): Promise<string> {
export interface OAuthResultMessage {
kind: 'success' | 'error'
text: string
}
export async function resolveOAuthMessage(ctx: OAuthReturnContext): Promise<OAuthResultMessage> {
if (ctx.reconnect) {
return `"${ctx.displayName}" reconnected successfully.`
return { kind: 'success', text: `"${ctx.displayName}" reconnected successfully.` }
}
try {
@@ -44,14 +49,48 @@ async function resolveOAuthMessage(ctx: OAuthReturnContext): Promise<string> {
const forProvider = oauthCredentials.filter((c) => c.providerId === ctx.providerId)
if (forProvider.length > ctx.preCount) {
return `"${ctx.displayName}" credential connected successfully.`
return {
kind: 'success',
text: `"${ctx.displayName}" credential connected successfully.`,
}
}
const existing = forProvider[0]
return `This account is already connected as "${existing?.displayName || ctx.displayName}".`
const baselineCredentials = new Map(
ctx.baselineCredentials?.map((credential) => [credential.id, credential]) ?? []
)
const reauthorizedCredential = forProvider.find((credential) => {
const baseline = baselineCredentials.get(credential.id)
return (
baseline !== undefined &&
(baseline.accountId !== credential.accountId ||
(baseline.updatedAt !== undefined && baseline.updatedAt !== credential.updatedAt))
)
})
if (reauthorizedCredential) {
return {
kind: 'success',
text: `This account is already connected as "${reauthorizedCredential.displayName}".`,
}
}
} catch {
return `"${ctx.displayName}" credential connected successfully.`
return {
kind: 'error',
text: `We couldn’t verify the "${ctx.displayName}" connection. Try again.`,
}
}
return {
kind: 'error',
text: `We couldn’t verify the "${ctx.displayName}" connection. Try again.`,
}
}
function showOAuthResultMessage(result: OAuthResultMessage): void {
if (result.kind === 'success') {
toast.success(result.text)
return
}
toast.error(result.text)
}
function dispatchCredentialUpdate(ctx: { providerId: string; workspaceId: string }) {
@@ -171,7 +210,7 @@ export function useOAuthReturnRouter() {
consumeOAuthReturnContext()
void (async () => {
const message = await resolveOAuthMessage(ctx)
toast.success(message)
showOAuthResultMessage(message)
dispatchCredentialUpdate(ctx)
})()
return
@@ -213,7 +252,7 @@ export function useOAuthReturnForWorkflow(workflowId: string) {
void (async () => {
const message = await resolveOAuthMessage(ctx)
toast.success(message)
showOAuthResultMessage(message)
dispatchCredentialUpdate(ctx)
})()
}, [workflowId])
@@ -233,7 +272,7 @@ export function useOAuthReturnForKBConnectors(knowledgeBaseId: string) {
void (async () => {
const message = await resolveOAuthMessage(ctx)
toast.success(message)
showOAuthResultMessage(message)
dispatchCredentialUpdate(ctx)
})()
}, [knowledgeBaseId])
@@ -277,7 +316,7 @@ export function useDesktopOAuthConnectListener() {
if (ctx) {
void (async () => {
const message = await resolveOAuthMessage(ctx)
toast.success(message)
showOAuthResultMessage(message)
dispatchCredentialUpdate(ctx)
})()
return
@@ -402,6 +402,7 @@ export const createCredentialDraftContract = defineRouteContract({
mode: 'json',
schema: z.object({
success: z.literal(true),
draftId: z.string().min(1),
}),
},
})
+27 -13
View File
@@ -97,7 +97,8 @@ import {
import { PlatformEvents } from '@/lib/core/telemetry'
import { getBaseUrl, isLocalhostUrl, parseOriginList } from '@/lib/core/utils/urls'
import {
loadOAuthCredentialDraftBinding,
captureOAuthCredentialDraftBinding,
consumeOAuthCredentialDraftBinding,
processCredentialDraft,
} from '@/lib/credentials/draft-processor'
import { sendEmail } from '@/lib/messaging/email/mailer'
@@ -404,9 +405,22 @@ export const auth = betterAuth({
},
account: {
create: {
before: async (account) => {
before: async (account, context) => {
const modifiedAccount = { ...account }
if (context?.path.startsWith('/oauth2/callback/')) {
try {
await captureOAuthCredentialDraftBinding(context, () => getOAuthState())
} catch (error) {
logger.error('[account.create.before] Failed to read OAuth credential draft state', {
userId: account.userId,
providerId: account.providerId,
error,
})
throw error
}
}
if (account.accessToken && isSalesforceOAuthProviderId(account.providerId)) {
const instanceUrl = await fetchSalesforceInstanceUrl(
account.providerId,
@@ -433,7 +447,7 @@ export const auth = betterAuth({
return { data: modifiedAccount }
},
after: async (account) => {
after: async (account, context) => {
/**
* Migrate credentials from stale account rows to the newly created one.
*
@@ -526,18 +540,18 @@ export const auth = betterAuth({
}
}
const credentialDraftBinding = await loadOAuthCredentialDraftBinding(() =>
getOAuthState()
)
if (credentialDraftBinding.status === 'unavailable') {
logger.error('[account.create.after] Failed to read OAuth credential draft state', {
userId: account.userId,
providerId: account.providerId,
error: credentialDraftBinding.error,
})
const isOAuth2Callback = context?.path.startsWith('/oauth2/callback/') === true
const credentialDraftBinding = context
? consumeOAuthCredentialDraftBinding(context)
: undefined
if (isOAuth2Callback && !credentialDraftBinding) {
throw new Error(
'OAuth credential draft binding was not captured before account creation'
)
}
if (credentialDraftBinding.status === 'available') {
if (credentialDraftBinding) {
try {
await processCredentialDraft({
draftId: credentialDraftBinding.draftId,
@@ -97,7 +97,6 @@ describe('createCredentialConnection', () => {
providerId: 'google-email',
credentialId: undefined,
displayName: 'Work Gmail',
displayNameDefinesIntent: true,
})
expect(result).toEqual({
authorizationUrl: 'https://sim.ai/api/auth/oauth2/authorize?draftId=draft-1',
@@ -127,7 +126,6 @@ describe('createCredentialConnection', () => {
providerId: 'google-email',
credentialId: 'credential-1',
displayName: 'Existing Gmail',
displayNameDefinesIntent: false,
})
})
})
@@ -53,7 +53,6 @@ export const createCredentialConnection = defineAuthorizedWorkspaceUseCase({
providerId: target.providerId,
credentialId: target.credentialId,
displayName,
displayNameDefinesIntent: input.providerId !== undefined && displayName !== undefined,
})
const authorizationUrl = new URL('/api/auth/oauth2/authorize', getBaseUrl())
authorizationUrl.searchParams.set('draftId', draft.id)
@@ -59,7 +59,7 @@ describe('saveCredentialDraft', () => {
})
it('authorizes workspace access before resolving reconnect credential access', async () => {
await saveCredentialDraft.execute({
const result = await saveCredentialDraft.execute({
principal,
input: {
workspaceId: 'workspace-1',
@@ -69,6 +69,7 @@ describe('saveCredentialDraft', () => {
},
})
expect(result).toEqual({ success: true, draftId: 'draft-1' })
expect(mocks.resolvePermission.mock.invocationCallOrder[0]).toBeLessThan(
mocks.getActor.mock.invocationCallOrder[0]
)
@@ -77,7 +78,6 @@ describe('saveCredentialDraft', () => {
userId: 'user-1',
workspaceId: 'workspace-1',
credentialId: 'credential-1',
displayNameDefinesIntent: false,
})
)
})
@@ -53,15 +53,14 @@ export const saveCredentialDraft = defineAuthorizedWorkspaceUseCase({
}
},
async execute({ principal, input }) {
await createConnectDraft({
const draft = await createConnectDraft({
userId: requirePrincipalSubjectUserId(principal),
workspaceId: input.workspaceId,
providerId: input.providerId,
displayName: input.displayName,
description: input.description,
credentialId: input.credentialId,
displayNameDefinesIntent: input.credentialId === undefined,
})
return { success: true as const }
return { success: true as const, draftId: draft.id }
},
})
+5
View File
@@ -52,6 +52,11 @@ interface OAuthReturnBase {
displayName: string
providerId: string
preCount: number
baselineCredentials?: Array<{
id: string
accountId: string | null
updatedAt?: string
}>
workspaceId: string
reconnect?: boolean
requestedAt: number
+25 -20
View File
@@ -1,7 +1,7 @@
/**
* @vitest-environment node
*/
import { dbChainMockFns, drizzleOrmMock, resetDbChainMock, schemaMock } from '@sim/testing'
import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
const { mockGenerateId } = vi.hoisted(() => ({
@@ -19,16 +19,15 @@ describe('createConnectDraft', () => {
mockGenerateId.mockReturnValue('new-draft-id')
})
it('refreshes the expiry without changing an active connection intent', async () => {
it('supersedes an active connection intent with a new exact draft id', async () => {
const expiresAt = new Date('2026-08-13T20:15:00.000Z')
dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'active-draft-id', expiresAt }])
dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'new-draft-id', expiresAt }])
const result = await createConnectDraft({
userId: 'user-1',
workspaceId: 'workspace-1',
providerId: 'google-email',
displayName: 'Work Gmail',
displayNameDefinesIntent: true,
})
expect(dbChainMockFns.values).toHaveBeenCalledWith(
@@ -37,16 +36,20 @@ describe('createConnectDraft', () => {
const conflict = dbChainMockFns.onConflictDoUpdate.mock.calls[0]?.[0] as
| { set?: Record<string, unknown>; setWhere?: unknown }
| undefined
expect(conflict?.set).not.toHaveProperty('id')
expect(conflict?.set).not.toHaveProperty('displayName')
expect(conflict?.set).not.toHaveProperty('credentialId')
expect(conflict?.setWhere).toBeDefined()
expect(result).toEqual({ id: 'active-draft-id', expiresAt })
expect(conflict?.set).toEqual(
expect.objectContaining({
id: 'new-draft-id',
displayName: 'Work Gmail',
credentialId: null,
})
)
expect(conflict?.setWhere).toBeUndefined()
expect(result).toEqual({ id: 'new-draft-id', expiresAt })
})
it('refreshes a reconnect target when its mutable display name changes', async () => {
it('supersedes a reconnect target when its mutable display name changes', async () => {
const expiresAt = new Date('2026-08-13T20:15:00.000Z')
dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'active-draft-id', expiresAt }])
dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'new-draft-id', expiresAt }])
await expect(
createConnectDraft({
@@ -56,15 +59,20 @@ describe('createConnectDraft', () => {
credentialId: 'credential-1',
displayName: 'Renamed Gmail',
})
).resolves.toEqual({ id: 'active-draft-id', expiresAt })
).resolves.toEqual({ id: 'new-draft-id', expiresAt })
expect(drizzleOrmMock.eq).not.toHaveBeenCalledWith(
schemaMock.pendingCredentialDraft.displayName,
'Renamed Gmail'
expect(dbChainMockFns.onConflictDoUpdate).toHaveBeenCalledWith(
expect.objectContaining({
set: expect.objectContaining({
id: 'new-draft-id',
displayName: 'Renamed Gmail',
credentialId: 'credential-1',
}),
})
)
})
it('fails fast when an active draft has a different connection intent', async () => {
it('fails when the draft upsert does not return a row', async () => {
dbChainMockFns.returning.mockResolvedValueOnce([])
await expect(
@@ -75,9 +83,6 @@ describe('createConnectDraft', () => {
credentialId: 'credential-1',
displayName: 'Existing Gmail',
})
).rejects.toMatchObject({
code: 'conflict',
message: 'A different OAuth connection flow is already active for this provider',
})
).rejects.toThrow('Failed to create OAuth credential draft')
})
})
+15 -16
View File
@@ -2,8 +2,7 @@ import { db } from '@sim/db'
import { credential, pendingCredentialDraft, user } from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { generateId } from '@sim/utils/id'
import { and, eq, gt, isNull, lt } from 'drizzle-orm'
import { OrchestrationError } from '@/lib/core/orchestration/types'
import { and, eq, gt, lt } from 'drizzle-orm'
import { defaultCredentialDisplayName } from '@/lib/credentials/display-name'
import { CREDENTIAL_DRAFT_TTL_MS } from '@/lib/credentials/draft-constants'
import { credentialProviderMatchesService, getAllOAuthServices } from '@/lib/oauth/utils'
@@ -30,8 +29,6 @@ export async function createConnectDraft(params: {
/** Reconnect only: the credential's actual name, so audit records stay accurate. */
displayName?: string
description?: string
/** Whether an explicitly requested name distinguishes this new-connection intent. */
displayNameDefinesIntent?: boolean
}): Promise<CreatedConnectDraft> {
const { userId, workspaceId, providerId, credentialId } = params
@@ -70,12 +67,6 @@ export async function createConnectDraft(params: {
and(eq(pendingCredentialDraft.userId, userId), lt(pendingCredentialDraft.expiresAt, now))
)
const id = generateId()
const sameTarget = credentialId
? eq(pendingCredentialDraft.credentialId, credentialId)
: isNull(pendingCredentialDraft.credentialId)
const sameIntent = params.displayNameDefinesIntent
? and(sameTarget, eq(pendingCredentialDraft.displayName, displayName))
: sameTarget
const [draft] = await db
.insert(pendingCredentialDraft)
.values({
@@ -95,16 +86,24 @@ export async function createConnectDraft(params: {
pendingCredentialDraft.providerId,
pendingCredentialDraft.workspaceId,
],
set: { expiresAt, createdAt: now },
setWhere: sameIntent,
/**
* A new launch supersedes an abandoned or failed launch for this provider.
* Rotating the primary key invalidates the earlier callback's exact draft
* binding, so only the newest OAuth flow can materialize its intent.
*/
set: {
id,
displayName,
description: params.description?.trim() || null,
credentialId: credentialId ?? null,
expiresAt,
createdAt: now,
},
})
.returning({ id: pendingCredentialDraft.id, expiresAt: pendingCredentialDraft.expiresAt })
if (!draft) {
throw new OrchestrationError(
'conflict',
'A different OAuth connection flow is already active for this provider'
)
throw new Error('Failed to create OAuth credential draft')
}
logger.info('Created OAuth connect credential draft', {
@@ -1,2 +1,3 @@
export const CREDENTIAL_DRAFT_TTL_MS = 15 * 60 * 1000
export const CREDENTIAL_DRAFT_TTL_SECONDS = CREDENTIAL_DRAFT_TTL_MS / 1000
export const OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM = 'credentialDraftId'
+77 -2
View File
@@ -1,7 +1,14 @@
/**
* @vitest-environment node
*/
import { auditMock, auditMockFns, queueTableRows, resetDbChainMock, schemaMock } from '@sim/testing'
import {
auditMock,
auditMockFns,
dbChainMockFns,
queueTableRows,
resetDbChainMock,
schemaMock,
} from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'
const mocks = vi.hoisted(() => ({
@@ -12,7 +19,75 @@ vi.mock('@sim/audit', () => auditMock)
vi.mock('@/lib/oauth/terminal-errors', () => ({ clearDeadFlag: mocks.clearDeadFlag }))
vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: vi.fn() }))
import { handleReconnectCredential } from '@/lib/credentials/draft-hooks'
import {
handleCreateCredentialFromDraft,
handleReconnectCredential,
} from '@/lib/credentials/draft-hooks'
describe('handleCreateCredentialFromDraft', () => {
beforeEach(() => {
vi.clearAllMocks()
resetDbChainMock()
})
it('treats a wrapped workspace-account conflict as an existing connection', async () => {
const now = new Date('2026-08-14T18:00:00.000Z')
const cause = Object.assign(new Error('duplicate key'), {
code: '23505',
constraint_name: 'credential_workspace_account_unique',
})
dbChainMockFns.values.mockRejectedValueOnce(new Error('Failed query', { cause }))
queueTableRows(schemaMock.credential, [
{ id: 'credential-existing', displayName: 'Existing Gmail' },
])
await handleCreateCredentialFromDraft({
draft: {
workspaceId: 'workspace-1',
displayName: 'New Gmail',
description: null,
},
accountId: 'account-1',
providerId: 'google-email',
userId: 'user-1',
now,
})
expect(dbChainMockFns.update).toHaveBeenCalledWith(schemaMock.credential)
expect(dbChainMockFns.set).toHaveBeenCalledWith({ updatedAt: now })
expect(mocks.clearDeadFlag).toHaveBeenCalledWith('account-1')
expect(auditMockFns.mockRecordAudit).toHaveBeenCalledWith(
expect.objectContaining({
action: 'credential.reconnected',
resourceId: 'credential-existing',
resourceName: 'Existing Gmail',
})
)
})
it('rethrows a different unique constraint failure', async () => {
const cause = Object.assign(new Error('duplicate key'), {
code: '23505',
constraint_name: 'credential_display_name_unique',
})
const wrapped = new Error('Failed query', { cause })
dbChainMockFns.values.mockRejectedValueOnce(wrapped)
await expect(
handleCreateCredentialFromDraft({
draft: {
workspaceId: 'workspace-1',
displayName: 'New Gmail',
description: null,
},
accountId: 'account-1',
providerId: 'google-email',
userId: 'user-1',
now: new Date('2026-08-14T18:00:00.000Z'),
})
).rejects.toBe(wrapped)
})
})
describe('handleReconnectCredential', () => {
beforeEach(() => {
+84 -45
View File
@@ -2,6 +2,7 @@ import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
import { db } from '@sim/db'
import * as schema from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { getPostgresConstraintName, getPostgresErrorCode } from '@sim/utils/errors'
import { generateId } from '@sim/utils/id'
import { and, eq, sql } from 'drizzle-orm'
import { deleteOrphanedOAuthAccount } from '@/lib/credentials/deletion'
@@ -36,65 +37,103 @@ export async function handleCreateCredentialFromDraft(params: {
createdAt: now,
updatedAt: now,
})
} catch (insertError: unknown) {
if (
getPostgresErrorCode(insertError) !== '23505' ||
getPostgresConstraintName(insertError) !== 'credential_workspace_account_unique'
) {
throw insertError
}
await db.insert(schema.credentialMember).values({
id: generateId(),
credentialId,
userId,
role: 'admin',
status: 'active',
joinedAt: now,
invitedBy: userId,
createdAt: now,
updatedAt: now,
})
const [existingCredential] = await db
.select({ id: schema.credential.id, displayName: schema.credential.displayName })
.from(schema.credential)
.where(
and(
eq(schema.credential.workspaceId, draft.workspaceId),
eq(schema.credential.accountId, accountId)
)
)
.limit(1)
logger.info('Created credential from draft', {
credentialId,
displayName: draft.displayName,
if (!existingCredential) {
throw new Error(
`Credential account conflict was reported without an existing credential for account ${accountId}`
)
}
logger.info('OAuth account is already connected to this workspace', {
credentialId: existingCredential.id,
displayName: existingCredential.displayName,
providerId,
accountId,
})
await db
.update(schema.credential)
.set({ updatedAt: now })
.where(eq(schema.credential.id, existingCredential.id))
await clearDeadFlag(accountId)
recordAudit({
workspaceId: draft.workspaceId,
actorId: userId,
action: AuditAction.CREDENTIAL_CREATED,
action: AuditAction.CREDENTIAL_RECONNECTED,
resourceType: AuditResourceType.CREDENTIAL,
resourceId: credentialId,
resourceName: draft.displayName,
description: `Created OAuth credential "${draft.displayName}"`,
metadata: { providerId, accountId },
})
captureServerEvent(
userId,
'credential_connected',
{
credential_type: 'oauth',
provider_id: providerId,
workspace_id: draft.workspaceId,
},
{
groups: { workspace: draft.workspaceId },
setOnce: { first_credential_connected_at: now.toISOString() },
}
)
} catch (insertError: unknown) {
const code =
insertError && typeof insertError === 'object' && 'code' in insertError
? (insertError as { code: string }).code
: undefined
if (code !== '23505') {
throw insertError
}
logger.info('Credential already exists, skipping draft', {
providerId,
accountId,
resourceId: existingCredential.id,
resourceName: existingCredential.displayName,
description: `Reconnected OAuth credential "${existingCredential.displayName}"`,
metadata: { accountId },
})
return
}
await db.insert(schema.credentialMember).values({
id: generateId(),
credentialId,
userId,
role: 'admin',
status: 'active',
joinedAt: now,
invitedBy: userId,
createdAt: now,
updatedAt: now,
})
logger.info('Created credential from draft', {
credentialId,
displayName: draft.displayName,
providerId,
accountId,
})
await clearDeadFlag(accountId)
recordAudit({
workspaceId: draft.workspaceId,
actorId: userId,
action: AuditAction.CREDENTIAL_CREATED,
resourceType: AuditResourceType.CREDENTIAL,
resourceId: credentialId,
resourceName: draft.displayName,
description: `Created OAuth credential "${draft.displayName}"`,
metadata: { providerId, accountId },
})
captureServerEvent(
userId,
'credential_connected',
{
credential_type: 'oauth',
provider_id: providerId,
workspace_id: draft.workspaceId,
},
{
groups: { workspace: draft.workspaceId },
setOnce: { first_credential_connected_at: now.toISOString() },
}
)
}
/**
@@ -21,6 +21,8 @@ vi.mock('@/lib/credentials/draft-hooks', () => ({
}))
import {
captureOAuthCredentialDraftBinding,
consumeOAuthCredentialDraftBinding,
loadOAuthCredentialDraftBinding,
parseCredentialDraftIdFromCallbackUrl,
processCredentialDraft,
@@ -150,3 +152,31 @@ describe('loadOAuthCredentialDraftBinding', () => {
expect(binding.status).toBe('unavailable')
})
})
describe('OAuth credential draft account hook binding', () => {
it('carries an exact binding from the account before-hook to the deferred after-hook', async () => {
const context = {}
await captureOAuthCredentialDraftBinding(context, async () => ({
callbackURL: 'https://sim.test/oauth/credential-connected?credentialDraftId=draft-exact',
}))
expect(consumeOAuthCredentialDraftBinding(context)).toEqual({
status: 'available',
draftId: 'draft-exact',
})
expect(consumeOAuthCredentialDraftBinding(context)).toBeUndefined()
})
it('fails before account creation when OAuth state cannot be captured', async () => {
const context = {}
const stateError = new Error('OAuth state is unavailable')
await expect(
captureOAuthCredentialDraftBinding(context, async () => {
throw stateError
})
).rejects.toBe(stateError)
expect(consumeOAuthCredentialDraftBinding(context)).toBeUndefined()
})
})
+28 -3
View File
@@ -2,6 +2,7 @@ import { db } from '@sim/db'
import * as schema from '@sim/db/schema'
import { createLogger } from '@sim/logger'
import { and, desc, eq, sql } from 'drizzle-orm'
import { OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM } from '@/lib/credentials/draft-constants'
import {
handleCreateCredentialFromDraft,
handleReconnectCredential,
@@ -9,16 +10,21 @@ import {
const logger = createLogger('CredentialDraftProcessor')
export const OAUTH_CREDENTIAL_DRAFT_CALLBACK_PARAM = 'credentialDraftId'
interface OAuthStateWithCallbackUrl {
callbackURL?: unknown
}
type OAuthCredentialDraftBinding =
export type OAuthCredentialDraftBinding =
| { status: 'available'; draftId?: string }
| { status: 'unavailable'; error: unknown }
type AvailableOAuthCredentialDraftBinding = Extract<
OAuthCredentialDraftBinding,
{ status: 'available' }
>
const oauthCredentialDraftBindings = new WeakMap<object, AvailableOAuthCredentialDraftBinding>()
/** Extracts a draft binding from Better Auth state and rejects malformed callback state. */
export function parseCredentialDraftIdFromCallbackUrl(callbackUrl: unknown): string | undefined {
if (callbackUrl === undefined) return undefined
@@ -43,6 +49,25 @@ export async function loadOAuthCredentialDraftBinding(
}
}
/** Captures request-scoped OAuth state before Better Auth defers its account after-hook. */
export async function captureOAuthCredentialDraftBinding(
context: object,
loadOAuthState: () => Promise<OAuthStateWithCallbackUrl | null | undefined>
): Promise<void> {
const binding = await loadOAuthCredentialDraftBinding(loadOAuthState)
if (binding.status === 'unavailable') throw binding.error
oauthCredentialDraftBindings.set(context, binding)
}
/** Consumes the binding captured for the same Better Auth account-hook context. */
export function consumeOAuthCredentialDraftBinding(
context: object
): AvailableOAuthCredentialDraftBinding | undefined {
const binding = oauthCredentialDraftBindings.get(context)
oauthCredentialDraftBindings.delete(context)
return binding
}
interface ProcessCredentialDraftParams {
draftId?: string
userId: string
+2 -1
View File
@@ -641,11 +641,12 @@ export interface DesktopOAuthConnectResult {
* Optional scope for an OAuth connect handoff. Chip-initiated connects carry
* the workspace (the browser flow creates the workspace connect draft
* server-side) and, for reconnects, the credential to rebind. Modal-initiated
* connects omit both — the app already created the draft.
* connects carry the exact draft the app already created.
*/
export interface DesktopOAuthConnectScope {
workspaceId?: string
credentialId?: string
draftId?: string
/** Mothership credential-chip attempt to echo on desktop completion. */
chatAttemptId?: string
}