Skip to content

Commit 5ba1185

Browse files
Sg312Sg312
committed
fixes
1 parent 817ca0d commit 5ba1185

9 files changed

Lines changed: 240 additions & 31 deletions

File tree

‎apps/desktop/src/main/handoff.test.ts‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,11 @@ import {
88
buildRedeemScript,
99
type ConnectHandoffCallback,
1010
createAuthFlow,
11+
createConnectFlow,
1112
createHandoffManager,
1213
type HandoffCallback,
1314
type HandoffCallbacks,
15+
type HandoffManager,
1416
type HandoffManagerDeps,
1517
} from '@/main/handoff'
1618
import type { EventRecorder } from '@/main/observability'
@@ -241,6 +243,24 @@ describe('createHandoffManager', () => {
241243
expect(manager.consume(state, 'login')).toBe(false)
242244
expect(manager.consume(state, 'connect')).toBe(true)
243245
})
246+
247+
it('returns the chat attempt correlated with the accepted connect state', async () => {
248+
const deps = makeDeps()
249+
const manager = createHandoffManager(deps, makeCallbacks())
250+
await manager.beginConnect('google-email', {
251+
workspaceId: 'workspace-1',
252+
chatAttemptId: 'attempt-1',
253+
})
254+
const state = new URL(vi.mocked(deps.openExternal).mock.calls[0][0]).searchParams.get(
255+
'state'
256+
) as string
257+
258+
expect(manager.consumeConnect(state)).toEqual({
259+
workspaceId: 'workspace-1',
260+
chatAttemptId: 'attempt-1',
261+
})
262+
expect(manager.consumeConnect(state)).toBeNull()
263+
})
244264
})
245265

246266
describe('connect handoff account pinning', () => {
@@ -273,6 +293,50 @@ describe('connect handoff account pinning', () => {
273293
})
274294
})
275295

296+
describe('connect completion correlation', () => {
297+
function makeConnectManager(scope: { chatAttemptId?: string }): HandoffManager {
298+
return {
299+
begin: vi.fn(async () => true),
300+
beginConnect: vi.fn(async () => true),
301+
consume: vi.fn(() => true),
302+
consumeConnect: vi.fn(() => scope),
303+
clear: vi.fn(),
304+
}
305+
}
306+
307+
it('echoes the accepted handoff chat attempt to the renderer', () => {
308+
const notifyRenderer = vi.fn()
309+
const flow = createConnectFlow({
310+
handoff: makeConnectManager({ chatAttemptId: 'attempt-1' }),
311+
events: makeEvents(),
312+
focusMainWindow: vi.fn(),
313+
notifyRenderer,
314+
})
315+
316+
flow.handleCallback({ state: VALID_STATE })
317+
318+
expect(notifyRenderer).toHaveBeenCalledWith({ ok: true, chatAttemptId: 'attempt-1' })
319+
})
320+
321+
it('marks ordinary integrations-page completions as explicitly uncorrelated', () => {
322+
const notifyRenderer = vi.fn()
323+
const flow = createConnectFlow({
324+
handoff: makeConnectManager({}),
325+
events: makeEvents(),
326+
focusMainWindow: vi.fn(),
327+
notifyRenderer,
328+
})
329+
330+
flow.handleCallback({ state: VALID_STATE, error: 'oauth_failed' })
331+
332+
expect(notifyRenderer).toHaveBeenCalledWith({
333+
ok: false,
334+
error: 'oauth_failed',
335+
chatAttemptId: null,
336+
})
337+
})
338+
})
339+
276340
describe('createAuthFlow window failures', () => {
277341
function makeAuthDeps(ensureMainWindow: () => Promise<never>) {
278342
const events = makeEvents()

‎apps/desktop/src/main/handoff.ts‎

Lines changed: 54 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -67,12 +67,14 @@ export interface HandoffManagerDeps {
6767
export interface ConnectScope {
6868
workspaceId?: string
6969
credentialId?: string
70+
chatAttemptId?: string
7071
}
7172

7273
export interface HandoffManager {
7374
begin(): Promise<boolean>
7475
beginConnect(providerId: string, scope?: ConnectScope): Promise<boolean>
7576
consume(state: string, kind: HandoffKind): boolean
77+
consumeConnect(state: string): ConnectScope | null
7678
clear(): void
7779
}
7880

@@ -92,7 +94,12 @@ export function createHandoffManager(
9294
const now = deps.now ?? Date.now
9395
let loopbackServer: Server | null = null
9496
let loopbackTimer: NodeJS.Timeout | undefined
95-
let pending: { state: string; createdAt: number; kind: HandoffKind } | null = null
97+
let pending: {
98+
state: string
99+
createdAt: number
100+
kind: HandoffKind
101+
connectScope?: ConnectScope
102+
} | null = null
96103

97104
const stopLoopback = () => {
98105
clearTimeout(loopbackTimer)
@@ -214,10 +221,23 @@ export function createHandoffManager(
214221
pending = null
215222
}
216223

224+
const consumePending = (state: string, kind: HandoffKind): NonNullable<typeof pending> | null => {
225+
if (!pending || pending.kind !== kind) return null
226+
if (now() - pending.createdAt > HANDOFF_TTL_MS) {
227+
clear()
228+
return null
229+
}
230+
if (!safeCompare(pending.state, state)) return null
231+
const consumed = pending
232+
clear()
233+
return consumed
234+
}
235+
217236
const beginFlow = async (
218237
kind: HandoffKind,
219238
landingPath: string,
220-
params: Record<string, string>
239+
params: Record<string, string>,
240+
connectScope?: ConnectScope
221241
): Promise<boolean> => {
222242
const state = generateShortId(STATE_LENGTH)
223243
// startLoopback() already tore down any prior server; if this bind fails,
@@ -228,7 +248,12 @@ export function createHandoffManager(
228248
clear()
229249
return false
230250
}
231-
pending = { state, createdAt: now(), kind }
251+
pending = {
252+
state,
253+
createdAt: now(),
254+
kind,
255+
...(connectScope ? { connectScope: { ...connectScope } } : {}),
256+
}
232257
const landing = new URL(landingPath, deps.origin())
233258
for (const [key, value] of Object.entries(params)) {
234259
landing.searchParams.set(key, value)
@@ -259,26 +284,24 @@ export function createHandoffManager(
259284
// unknown (offline, signed out): the page then falls back to its normal
260285
// login redirect rather than blocking a connect on a failed probe.
261286
const userId = await deps.currentUserId()
262-
return beginFlow('connect', '/desktop/connect', {
263-
provider: providerId,
264-
...(userId ? { user: userId } : {}),
265-
...(scope.workspaceId ? { workspaceId: scope.workspaceId } : {}),
266-
...(scope.credentialId ? { credentialId: scope.credentialId } : {}),
267-
})
287+
return beginFlow(
288+
'connect',
289+
'/desktop/connect',
290+
{
291+
provider: providerId,
292+
...(userId ? { user: userId } : {}),
293+
...(scope.workspaceId ? { workspaceId: scope.workspaceId } : {}),
294+
...(scope.credentialId ? { credentialId: scope.credentialId } : {}),
295+
},
296+
scope
297+
)
268298
},
269299
consume(state: string, kind: HandoffKind) {
270-
if (!pending || pending.kind !== kind) {
271-
return false
272-
}
273-
if (now() - pending.createdAt > HANDOFF_TTL_MS) {
274-
clear()
275-
return false
276-
}
277-
if (!safeCompare(pending.state, state)) {
278-
return false
279-
}
280-
clear()
281-
return true
300+
return consumePending(state, kind) !== null
301+
},
302+
consumeConnect(state: string) {
303+
const consumed = consumePending(state, 'connect')
304+
return consumed ? { ...(consumed.connectScope ?? {}) } : null
282305
},
283306
clear,
284307
}
@@ -443,6 +466,8 @@ export function createAuthFlow(deps: AuthFlowDeps): AuthFlow {
443466
export interface ConnectHandoffResult {
444467
ok: boolean
445468
error?: string
469+
/** Exact Mothership chat attempt, or null for ordinary integration flows. */
470+
chatAttemptId: string | null
446471
}
447472

448473
export interface ConnectFlowDeps {
@@ -476,19 +501,24 @@ export function createConnectFlow(deps: ConnectFlowDeps): ConnectFlow {
476501
return opened
477502
},
478503
handleCallback(callback: ConnectHandoffCallback) {
479-
if (!deps.handoff.consume(callback.state, 'connect')) {
504+
const scope = deps.handoff.consumeConnect(callback.state)
505+
if (!scope) {
480506
deps.events.record('connect_handoff_state_fail')
481507
return
482508
}
483509
if (callback.error === undefined) {
484510
deps.events.record('connect_handoff_ok')
485511
deps.focusMainWindow()
486-
deps.notifyRenderer({ ok: true })
512+
deps.notifyRenderer({ ok: true, chatAttemptId: scope.chatAttemptId ?? null })
487513
return
488514
}
489515
deps.events.record('connect_handoff_error', { error: callback.error })
490516
deps.focusMainWindow()
491-
deps.notifyRenderer({ ok: false, error: callback.error })
517+
deps.notifyRenderer({
518+
ok: false,
519+
error: callback.error,
520+
chatAttemptId: scope.chatAttemptId ?? null,
521+
})
492522
},
493523
}
494524
}

‎apps/desktop/src/main/ipc.test.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -337,14 +337,20 @@ describe('registerIpcHandlers', () => {
337337

338338
// Chip-initiated connects carry workspace/credential scope; malformed
339339
// scopes (wrong types, unsafe ids) are rejected before the handoff.
340-
expect(await handler?.(appEvent, 'slack', { workspaceId: 'ws1', credentialId: 'cred_1' })).toBe(
341-
true
342-
)
340+
expect(
341+
await handler?.(appEvent, 'slack', {
342+
workspaceId: 'ws1',
343+
credentialId: 'cred_1',
344+
chatAttemptId: 'attempt_1',
345+
})
346+
).toBe(true)
343347
expect(deps.beginOAuthConnect).toHaveBeenCalledWith('slack', {
344348
workspaceId: 'ws1',
345349
credentialId: 'cred_1',
350+
chatAttemptId: 'attempt_1',
346351
})
347352
expect(await handler?.(appEvent, 'slack', { workspaceId: 'ws/../evil' })).toBe(false)
353+
expect(await handler?.(appEvent, 'slack', { chatAttemptId: '../wrong' })).toBe(false)
348354
expect(await handler?.(appEvent, 'slack', 'not-an-object')).toBe(false)
349355
})
350356

‎apps/desktop/src/main/ipc.ts‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ function parseDesktopScope(raw: unknown): string | null {
9090
export interface OAuthConnectScope {
9191
workspaceId?: string
9292
credentialId?: string
93+
chatAttemptId?: string
9394
}
9495

9596
/**
@@ -104,7 +105,11 @@ export function parseOAuthConnectScope(raw: unknown): OAuthConnectScope | undefi
104105
if (typeof raw !== 'object') {
105106
return undefined
106107
}
107-
const { workspaceId, credentialId } = raw as { workspaceId?: unknown; credentialId?: unknown }
108+
const { workspaceId, credentialId, chatAttemptId } = raw as {
109+
workspaceId?: unknown
110+
credentialId?: unknown
111+
chatAttemptId?: unknown
112+
}
108113
if (
109114
workspaceId !== undefined &&
110115
(typeof workspaceId !== 'string' || !ID_PATTERN.test(workspaceId))
@@ -117,9 +122,16 @@ export function parseOAuthConnectScope(raw: unknown): OAuthConnectScope | undefi
117122
) {
118123
return undefined
119124
}
125+
if (
126+
chatAttemptId !== undefined &&
127+
(typeof chatAttemptId !== 'string' || !ID_PATTERN.test(chatAttemptId))
128+
) {
129+
return undefined
130+
}
120131
return {
121132
...(workspaceId !== undefined ? { workspaceId } : {}),
122133
...(credentialId !== undefined ? { credentialId } : {}),
134+
...(chatAttemptId !== undefined ? { chatAttemptId } : {}),
123135
}
124136
}
125137

‎apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/special-tags/use-oauth-chip-connection.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
44
import { useParams } from 'next/navigation'
55
import {
66
addOAuthChatAttemptToAuthorizeUrl,
7+
clearActiveDesktopOAuthChatAttempt,
78
createOAuthChatAttempt,
89
getOAuthCredentialBaseline,
910
hasOAuthCredentialChanged,
@@ -315,9 +316,13 @@ export function useOAuthChipConnection({
315316
.beginOAuthConnect(providerId, {
316317
workspaceId: url.searchParams.get('workspaceId') ?? workspaceId,
317318
credentialId: url.searchParams.get('credentialId') ?? undefined,
319+
chatAttemptId: attempt.id,
318320
})
319321
.then((opened) => {
320-
if (!opened) setOAuthChatAttemptStatus(attempt.id, 'failed')
322+
if (!opened) {
323+
clearActiveDesktopOAuthChatAttempt(attempt.id)
324+
setOAuthChatAttemptStatus(attempt.id, 'failed')
325+
}
321326
})
322327
return
323328
}

‎apps/sim/hooks/use-oauth-return.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import {
1919
OAUTH_CHAT_ATTEMPT_MAX_AGE_MS,
2020
OAUTH_CHAT_ATTEMPT_PARAM,
2121
readOAuthChatAttempt,
22-
resolveActiveDesktopOAuthChatAttempt,
22+
resolveDesktopOAuthChatAttempt,
2323
setOAuthChatAttemptStatus,
2424
} from '@/lib/credentials/oauth-chat-attempt'
2525
import { getDesktopBridge } from '@/lib/desktop'
@@ -266,7 +266,7 @@ export function useDesktopOAuthConnectListener() {
266266
const rawCtx = readOAuthReturnContext()
267267
if (rawCtx) consumeOAuthReturnContext()
268268
const ctx = rawCtx && Date.now() - rawCtx.requestedAt <= CONTEXT_MAX_AGE_MS ? rawCtx : null
269-
const chatAttempt = resolveActiveDesktopOAuthChatAttempt(result.ok ? 'connected' : 'failed')
269+
const chatAttempt = resolveDesktopOAuthChatAttempt(result, result.ok ? 'connected' : 'failed')
270270

271271
if (!result.ok) {
272272
toast.error('The account connection didn’t finish. Try connecting again.')

0 commit comments

Comments
 (0)