improvement(mship): questions improvement (#6385)

* credentials continue

* fixes

Author: Sg312

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

@@ -8,9 +8,11 @@ import {
   buildRedeemScript,
   type ConnectHandoffCallback,
   createAuthFlow,
+  createConnectFlow,
   createHandoffManager,
   type HandoffCallback,
   type HandoffCallbacks,
+  type HandoffManager,
   type HandoffManagerDeps,
 } from '@/main/handoff'
 import type { EventRecorder } from '@/main/observability'
@@ -241,6 +243,24 @@ describe('createHandoffManager', () => {
     expect(manager.consume(state, 'login')).toBe(false)
     expect(manager.consume(state, 'connect')).toBe(true)
   })
+
+  it('returns the chat attempt correlated with the accepted connect state', async () => {
+    const deps = makeDeps()
+    const manager = createHandoffManager(deps, makeCallbacks())
+    await manager.beginConnect('google-email', {
+      workspaceId: 'workspace-1',
+      chatAttemptId: 'attempt-1',
+    })
+    const state = new URL(vi.mocked(deps.openExternal).mock.calls[0][0]).searchParams.get(
+      'state'
+    ) as string
+
+    expect(manager.consumeConnect(state)).toEqual({
+      workspaceId: 'workspace-1',
+      chatAttemptId: 'attempt-1',
+    })
+    expect(manager.consumeConnect(state)).toBeNull()
+  })
 })
 
 describe('connect handoff account pinning', () => {
@@ -273,6 +293,50 @@ describe('connect handoff account pinning', () => {
   })
 })
 
+describe('connect completion correlation', () => {
+  function makeConnectManager(scope: { chatAttemptId?: string }): HandoffManager {
+    return {
+      begin: vi.fn(async () => true),
+      beginConnect: vi.fn(async () => true),
+      consume: vi.fn(() => true),
+      consumeConnect: vi.fn(() => scope),
+      clear: vi.fn(),
+    }
+  }
+
+  it('echoes the accepted handoff chat attempt to the renderer', () => {
+    const notifyRenderer = vi.fn()
+    const flow = createConnectFlow({
+      handoff: makeConnectManager({ chatAttemptId: 'attempt-1' }),
+      events: makeEvents(),
+      focusMainWindow: vi.fn(),
+      notifyRenderer,
+    })
+
+    flow.handleCallback({ state: VALID_STATE })
+
+    expect(notifyRenderer).toHaveBeenCalledWith({ ok: true, chatAttemptId: 'attempt-1' })
+  })
+
+  it('marks ordinary integrations-page completions as explicitly uncorrelated', () => {
+    const notifyRenderer = vi.fn()
+    const flow = createConnectFlow({
+      handoff: makeConnectManager({}),
+      events: makeEvents(),
+      focusMainWindow: vi.fn(),
+      notifyRenderer,
+    })
+
+    flow.handleCallback({ state: VALID_STATE, error: 'oauth_failed' })
+
+    expect(notifyRenderer).toHaveBeenCalledWith({
+      ok: false,
+      error: 'oauth_failed',
+      chatAttemptId: null,
+    })
+  })
+})
+
 describe('createAuthFlow window failures', () => {
   function makeAuthDeps(ensureMainWindow: () => Promise<never>) {
     const events = makeEvents()

apps/desktop/src/main/handoff.ts

@@ -67,12 +67,14 @@ export interface HandoffManagerDeps {
 export interface ConnectScope {
   workspaceId?: string
   credentialId?: string
+  chatAttemptId?: string
 }
 
 export interface HandoffManager {
   begin(): Promise<boolean>
   beginConnect(providerId: string, scope?: ConnectScope): Promise<boolean>
   consume(state: string, kind: HandoffKind): boolean
+  consumeConnect(state: string): ConnectScope | null
   clear(): void
 }
 
@@ -92,7 +94,12 @@ export function createHandoffManager(
   const now = deps.now ?? Date.now
   let loopbackServer: Server | null = null
   let loopbackTimer: NodeJS.Timeout | undefined
-  let pending: { state: string; createdAt: number; kind: HandoffKind } | null = null
+  let pending: {
+    state: string
+    createdAt: number
+    kind: HandoffKind
+    connectScope?: ConnectScope
+  } | null = null
 
   const stopLoopback = () => {
     clearTimeout(loopbackTimer)
@@ -214,10 +221,23 @@ export function createHandoffManager(
     pending = null
   }
 
+  const consumePending = (state: string, kind: HandoffKind): NonNullable<typeof pending> | null => {
+    if (!pending || pending.kind !== kind) return null
+    if (now() - pending.createdAt > HANDOFF_TTL_MS) {
+      clear()
+      return null
+    }
+    if (!safeCompare(pending.state, state)) return null
+    const consumed = pending
+    clear()
+    return consumed
+  }
+
   const beginFlow = async (
     kind: HandoffKind,
     landingPath: string,
-    params: Record<string, string>
+    params: Record<string, string>,
+    connectScope?: ConnectScope
   ): Promise<boolean> => {
     const state = generateShortId(STATE_LENGTH)
     // startLoopback() already tore down any prior server; if this bind fails,
@@ -228,7 +248,12 @@ export function createHandoffManager(
       clear()
       return false
     }
-    pending = { state, createdAt: now(), kind }
+    pending = {
+      state,
+      createdAt: now(),
+      kind,
+      ...(connectScope ? { connectScope: { ...connectScope } } : {}),
+    }
     const landing = new URL(landingPath, deps.origin())
     for (const [key, value] of Object.entries(params)) {
       landing.searchParams.set(key, value)
@@ -259,26 +284,24 @@ export function createHandoffManager(
       // unknown (offline, signed out): the page then falls back to its normal
       // login redirect rather than blocking a connect on a failed probe.
       const userId = await deps.currentUserId()
-      return beginFlow('connect', '/desktop/connect', {
-        provider: providerId,
-        ...(userId ? { user: userId } : {}),
-        ...(scope.workspaceId ? { workspaceId: scope.workspaceId } : {}),
-        ...(scope.credentialId ? { credentialId: scope.credentialId } : {}),
-      })
+      return beginFlow(
+        'connect',
+        '/desktop/connect',
+        {
+          provider: providerId,
+          ...(userId ? { user: userId } : {}),
+          ...(scope.workspaceId ? { workspaceId: scope.workspaceId } : {}),
+          ...(scope.credentialId ? { credentialId: scope.credentialId } : {}),
+        },
+        scope
+      )
     },
     consume(state: string, kind: HandoffKind) {
-      if (!pending || pending.kind !== kind) {
-        return false
-      }
-      if (now() - pending.createdAt > HANDOFF_TTL_MS) {
-        clear()
-        return false
-      }
-      if (!safeCompare(pending.state, state)) {
-        return false
-      }
-      clear()
-      return true
+      return consumePending(state, kind) !== null
+    },
+    consumeConnect(state: string) {
+      const consumed = consumePending(state, 'connect')
+      return consumed ? { ...(consumed.connectScope ?? {}) } : null
     },
     clear,
   }
@@ -443,6 +466,8 @@ export function createAuthFlow(deps: AuthFlowDeps): AuthFlow {
 export interface ConnectHandoffResult {
   ok: boolean
   error?: string
+  /** Exact Mothership chat attempt, or null for ordinary integration flows. */
+  chatAttemptId: string | null
 }
 
 export interface ConnectFlowDeps {
@@ -476,19 +501,24 @@ export function createConnectFlow(deps: ConnectFlowDeps): ConnectFlow {
       return opened
     },
     handleCallback(callback: ConnectHandoffCallback) {
-      if (!deps.handoff.consume(callback.state, 'connect')) {
+      const scope = deps.handoff.consumeConnect(callback.state)
+      if (!scope) {
         deps.events.record('connect_handoff_state_fail')
         return
       }
       if (callback.error === undefined) {
         deps.events.record('connect_handoff_ok')
         deps.focusMainWindow()
-        deps.notifyRenderer({ ok: true })
+        deps.notifyRenderer({ ok: true, chatAttemptId: scope.chatAttemptId ?? null })
         return
       }
       deps.events.record('connect_handoff_error', { error: callback.error })
       deps.focusMainWindow()
-      deps.notifyRenderer({ ok: false, error: callback.error })
+      deps.notifyRenderer({
+        ok: false,
+        error: callback.error,
+        chatAttemptId: scope.chatAttemptId ?? null,
+      })
     },
   }
 }

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

@@ -337,14 +337,20 @@ describe('registerIpcHandlers', () => {
 
     // Chip-initiated connects carry workspace/credential scope; malformed
     // scopes (wrong types, unsafe ids) are rejected before the handoff.
-    expect(await handler?.(appEvent, 'slack', { workspaceId: 'ws1', credentialId: 'cred_1' })).toBe(
-      true
-    )
+    expect(
+      await handler?.(appEvent, 'slack', {
+        workspaceId: 'ws1',
+        credentialId: 'cred_1',
+        chatAttemptId: 'attempt_1',
+      })
+    ).toBe(true)
     expect(deps.beginOAuthConnect).toHaveBeenCalledWith('slack', {
       workspaceId: 'ws1',
       credentialId: 'cred_1',
+      chatAttemptId: 'attempt_1',
     })
     expect(await handler?.(appEvent, 'slack', { workspaceId: 'ws/../evil' })).toBe(false)
+    expect(await handler?.(appEvent, 'slack', { chatAttemptId: '../wrong' })).toBe(false)
     expect(await handler?.(appEvent, 'slack', 'not-an-object')).toBe(false)
   })
 

apps/desktop/src/main/ipc.ts

@@ -90,6 +90,7 @@ function parseDesktopScope(raw: unknown): string | null {
 export interface OAuthConnectScope {
   workspaceId?: string
   credentialId?: string
+  chatAttemptId?: string
 }
 
 /**
@@ -104,7 +105,11 @@ export function parseOAuthConnectScope(raw: unknown): OAuthConnectScope | undefi
   if (typeof raw !== 'object') {
     return undefined
   }
-  const { workspaceId, credentialId } = raw as { workspaceId?: unknown; credentialId?: unknown }
+  const { workspaceId, credentialId, chatAttemptId } = raw as {
+    workspaceId?: unknown
+    credentialId?: unknown
+    chatAttemptId?: unknown
+  }
   if (
     workspaceId !== undefined &&
     (typeof workspaceId !== 'string' || !ID_PATTERN.test(workspaceId))
@@ -117,9 +122,16 @@ export function parseOAuthConnectScope(raw: unknown): OAuthConnectScope | undefi
   ) {
     return undefined
   }
+  if (
+    chatAttemptId !== undefined &&
+    (typeof chatAttemptId !== 'string' || !ID_PATTERN.test(chatAttemptId))
+  ) {
+    return undefined
+  }
   return {
     ...(workspaceId !== undefined ? { workspaceId } : {}),
     ...(credentialId !== undefined ? { credentialId } : {}),
+    ...(chatAttemptId !== undefined ? { chatAttemptId } : {}),
   }
 }
 

apps/sim/app/api/auth/trello/authorize/route.ts

@@ -6,6 +6,7 @@ import { parseRequest } from '@/lib/api/server'
 import { getSession } from '@/lib/auth'
 import { env } from '@/lib/core/config/env'
 import { getBaseUrl } from '@/lib/core/utils/urls'
+import { isSameOrigin } from '@/lib/core/utils/validation'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { getCanonicalScopesForProvider } from '@/lib/oauth/utils'
 
@@ -14,6 +15,7 @@ const logger = createLogger('TrelloAuthorize')
 export const dynamic = 'force-dynamic'
 
 const TRELLO_STATE_COOKIE = 'trello_oauth_state'
+const TRELLO_RETURN_URL_COOKIE = 'trello_return_url'
 const TRELLO_STATE_COOKIE_PATH = '/api/auth/trello'
 const TRELLO_STATE_COOKIE_MAX_AGE_SECONDS = 60 * 10
 
@@ -26,6 +28,7 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
 
     const parsed = await parseRequest(authorizeTrelloContract, request, {})
     if (!parsed.success) return parsed.response
+    const { returnUrl: requestedReturnUrl } = parsed.data.query
 
     const apiKey = env.TRELLO_API_KEY
 
@@ -57,6 +60,20 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
       maxAge: TRELLO_STATE_COOKIE_MAX_AGE_SECONDS,
       path: TRELLO_STATE_COOKIE_PATH,
     })
+    if (requestedReturnUrl && isSameOrigin(requestedReturnUrl)) {
+      response.cookies.set(TRELLO_RETURN_URL_COOKIE, requestedReturnUrl, {
+        httpOnly: true,
+        secure: process.env.NODE_ENV === 'production',
+        sameSite: 'lax',
+        maxAge: TRELLO_STATE_COOKIE_MAX_AGE_SECONDS,
+        path: TRELLO_STATE_COOKIE_PATH,
+      })
+    } else {
+      response.cookies.delete({
+        name: TRELLO_RETURN_URL_COOKIE,
+        path: TRELLO_STATE_COOKIE_PATH,
+      })
+    }
     return response
   } catch (error) {
     logger.error('Error initiating Trello authorization:', error)