From d063c792b123649c72b49494dbaaf77ea88f961b Mon Sep 17 00:00:00 2001
From: Tomas Martykan <t.martykan@gmail.com>
Date: Mon, 3 Feb 2025 13:26:37 +0100
Subject: [PATCH] chore(suite-desktop-connect-popup): typed call per method

---
 packages/connect/src/events/call.ts           |  1 +
 .../src/connectPopupThunks.ts                 | 47 +++++++++++++------
 .../middlewares/wallet/discoveryMiddleware.ts |  4 +-
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/packages/connect/src/events/call.ts b/packages/connect/src/events/call.ts
index 871c198630e..40e00da0d3d 100644
--- a/packages/connect/src/events/call.ts
+++ b/packages/connect/src/events/call.ts
@@ -43,6 +43,7 @@ type CallApi = {
 
 export type CallMethodUnion = CallApi[keyof CallApi];
 export type CallMethodPayload = Parameters<CallMethodUnion>[0];
+export type CallMethodParams<M extends keyof CallApi> = Parameters<CallApi[M]>[0];
 export type CallMethodResponse<M extends keyof CallApi> = UnwrappedResponse<ReturnType<CallApi[M]>>;
 export type CallMethodAnyResponse = ReturnType<CallMethodUnion>;
 
diff --git a/packages/suite-desktop-connect-popup/src/connectPopupThunks.ts b/packages/suite-desktop-connect-popup/src/connectPopupThunks.ts
index c0d88fd9c8c..18854c9118f 100644
--- a/packages/suite-desktop-connect-popup/src/connectPopupThunks.ts
+++ b/packages/suite-desktop-connect-popup/src/connectPopupThunks.ts
@@ -1,25 +1,31 @@
-import { createThunk } from '@suite-common/redux-utils';
+import { AsyncThunkAction } from '@reduxjs/toolkit';
+
+import { CustomThunkAPI, createThunk } from '@suite-common/redux-utils';
 import { selectSelectedDevice } from '@suite-common/wallet-core';
-import TrezorConnect, { ERRORS } from '@trezor/connect';
+import TrezorConnect, { CallMethodParams, CallMethodResponse, ERRORS } from '@trezor/connect';
 import { serializeError } from '@trezor/connect/src/constants/errors';
 import { desktopApi } from '@trezor/suite-desktop-api';
 import { createDeferred } from '@trezor/utils';
 
 const CONNECT_POPUP_MODULE = '@common/connect-popup';
 
-export const connectPopupCallThunk = createThunk<
-    Promise<{
-        id: number;
-        success: boolean;
-        payload: any;
-    }>,
-    {
-        id: number;
-        method: string;
-        payload: any;
-        processName?: string;
-        origin?: string;
-    }
+type ConnectPopupCallThunkResponse<M extends keyof typeof TrezorConnect> = Promise<{
+    id: number;
+    success: boolean;
+    payload: CallMethodResponse<M>;
+}>;
+
+type ConnectPopupCallThunkParams<M extends keyof typeof TrezorConnect> = {
+    id: number;
+    processName?: string;
+    origin?: string;
+    method: M;
+    payload: Omit<CallMethodParams<M>, 'method'>;
+};
+
+export const connectPopupCallThunkInner = createThunk<
+    ConnectPopupCallThunkResponse<keyof typeof TrezorConnect>,
+    ConnectPopupCallThunkParams<keyof typeof TrezorConnect>
 >(
     `${CONNECT_POPUP_MODULE}/callThunk`,
     async ({ id, method, payload, processName, origin }, { dispatch, getState, extra }) => {
@@ -86,11 +92,22 @@ export const connectPopupCallThunk = createThunk<
     },
 );
 
+// Typed thunk that takes the method as a generic parameter
+// Original thunk is exposed as well for using .fulfilled, .rejected, etc.
+export const connectPopupCallThunk = <M extends keyof typeof TrezorConnect>(
+    params: ConnectPopupCallThunkParams<M>,
+): AsyncThunkAction<
+    ConnectPopupCallThunkResponse<M>,
+    ConnectPopupCallThunkParams<M>,
+    CustomThunkAPI
+> => connectPopupCallThunkInner(params) as any;
+
 export const connectPopupInitThunk = createThunk(
     `${CONNECT_POPUP_MODULE}/initPopupThunk`,
     async (_, { dispatch }) => {
         if (desktopApi.available && (await desktopApi.connectPopupEnabled())) {
             desktopApi.on('connect-popup/call', async params => {
+                // @ts-expect-error: params in desktopApi are not fully typed
                 const response = await dispatch(connectPopupCallThunk(params)).unwrap();
                 desktopApi.connectPopupResponse(response);
             });
diff --git a/packages/suite/src/middlewares/wallet/discoveryMiddleware.ts b/packages/suite/src/middlewares/wallet/discoveryMiddleware.ts
index cd1ace96332..fc0951d02e0 100644
--- a/packages/suite/src/middlewares/wallet/discoveryMiddleware.ts
+++ b/packages/suite/src/middlewares/wallet/discoveryMiddleware.ts
@@ -15,7 +15,7 @@ import {
 } from '@suite-common/wallet-core';
 import * as discoveryActions from '@suite-common/wallet-core';
 import { UI } from '@trezor/connect';
-import { connectPopupCallThunk } from '@trezor/suite-desktop-connect-popup';
+import { connectPopupCallThunkInner } from '@trezor/suite-desktop-connect-popup';
 
 import * as walletSettingsActions from 'src/actions/settings/walletSettingsActions';
 import { MODAL, ROUTER, SUITE } from 'src/actions/suite/constants';
@@ -153,7 +153,7 @@ export const prepareDiscoveryMiddleware = createMiddlewareWithExtraDeps(
         if (
             becomesConnected ||
             action.type === SUITE.APP_CHANGED ||
-            connectPopupCallThunk.fulfilled.match(action) ||
+            connectPopupCallThunkInner.fulfilled.match(action) ||
             deviceActions.selectDevice.match(action) ||
             authorizeDeviceThunk.fulfilled.match(action) ||
             walletSettingsActions.changeNetworks.match(action) ||