fix: select one discovered HTP session for Hexagon inference (#907)

* fix: select one discovered Hexagon device

* fix: resolve saved Hexagon devices before model initialization

* chore: cover exact Hexagon selection in Settings and benchmarks

* chore: isolate device-rule fetching in store-dependent unit tests

---------

Co-authored-by: a-ghorbani <ghorbani59@gmail.com>

Author: pocketpal-dev-team[bot]

__mocks__/external/llama.rn.ts

@@ -50,6 +50,8 @@ export const loadLlamaModelInfo = jest.fn();
 
 export const initLlama = jest.fn();
 
+export const getBackendDevicesInfo = jest.fn().mockResolvedValue([]);
+
 export const toggleNativeLog = jest.fn().mockResolvedValue(undefined);
 
 export const addNativeLogListener = jest.fn().mockReturnValue({

src/__automation__/screens/__tests__/BenchmarkRunnerScreen.test.tsx

@@ -1,4 +1,5 @@
 import React from 'react';
+import {Platform} from 'react-native';
 import {act} from 'react-test-renderer';
 
 import {fireEvent, render, waitFor} from '../../../../jest/test-utils';
@@ -495,15 +496,68 @@ describe('BenchmarkRunnerScreen', () => {
       stubHexagonLogs();
       getDeviceOptions.mockResolvedValueOnce([
         {id: 'cpu', label: 'CPU', devices: ['CPU']},
-        {id: 'hexagon', label: 'Hexagon', devices: ['HTP*']},
+        {id: 'hexagon', label: 'Hexagon', devices: ['HTP0']},
       ]);
       const cfg: BenchConfig = {...VALID_CONFIG, backends: ['hexagon']};
       await runMatrix(cfg, setStatus, setLastCell);
       const [paramsArg] = (initLlama as jest.Mock).mock.calls[0];
-      expect(paramsArg.devices).toEqual(['HTP*']);
+      expect(paramsArg.devices).toEqual(['HTP0']);
       expect(paramsArg.n_gpu_layers).toBe(99);
     });
 
+    it.each([['HTP0', 'HTP1', 'HTP2', 'HTP3', 'HTP4', 'HTP5'], ['HTP3']])(
+      'uses real discovery for Hexagon sessions %j',
+      async (...names) => {
+        const originalOS = Platform.OS;
+        Platform.OS = 'android';
+        const {getBackendDevicesInfo} = require('llama.rn');
+        const actualSelection = jest.requireActual(
+          '../../../utils/deviceSelection',
+        );
+        getBackendDevicesInfo.mockResolvedValue(
+          names.map(deviceName => ({
+            deviceName,
+            type: 'accel',
+            backend: 'HTP',
+          })),
+        );
+        getDeviceOptions.mockImplementationOnce(
+          actualSelection.getDeviceOptions,
+        );
+        stubHexagonLogs();
+        const saved = JSON.parse(JSON.stringify(modelStore.contextInitParams));
+        try {
+          const cfg: BenchConfig = {
+            ...VALID_CONFIG,
+            backends: ['hexagon'],
+            settings_axes: [{name: 'flash_attn_type', values: ['on']}],
+          };
+          await runMatrix(cfg, setStatus, setLastCell);
+          const lastWrite =
+            RNFS.writeFile.mock.calls[RNFS.writeFile.mock.calls.length - 1];
+          const report = JSON.parse(lastWrite[1]);
+          expect(initLlama).toHaveBeenCalledTimes(1);
+          expect(initLlama.mock.calls[0][0]).toMatchObject({
+            devices: [names[0]],
+            n_gpu_layers: 99,
+            flash_attn_type: 'on',
+          });
+          expect(report.runs[0]).toMatchObject({
+            status: 'ok',
+            effective_init_params: {
+              devices: [names[0]],
+              n_gpu_layers: 99,
+              flash_attn_type: 'on',
+            },
+          });
+          expect(modelStore.contextInitParams).toEqual(saved);
+        } finally {
+          Platform.OS = originalOS;
+          getBackendDevicesInfo.mockReset().mockResolvedValue([]);
+        }
+      },
+    );
+
     // -------------------------------------------------------------------------
     // Per-cell context release (sole release site is the per-cell finally)
     // -------------------------------------------------------------------------
@@ -1326,10 +1380,10 @@ describe('BenchmarkRunnerScreen', () => {
         filePath: '/mock/path/m.gguf',
         base: DEFAULT_BENCH_BASE_PARAMS,
         overrides: {} as any,
-        devices: ['HTP*'],
+        devices: ['HTP0'],
         n_gpu_layers: 99,
       });
-      expect(params.devices).toEqual(['HTP*']);
+      expect(params.devices).toEqual(['HTP0']);
       expect(params.n_gpu_layers).toBe(99);
       expect(params.model).toBe('/mock/path/m.gguf');
     });

src/hooks/__tests__/useTheme.test.tsx

@@ -6,6 +6,10 @@ import {observer} from 'mobx-react';
 
 jest.unmock('../useTheme');
 jest.unmock('../../store');
+
+jest.mock('../../services/deviceRules/rules', () => ({
+  fetchRules: jest.fn().mockResolvedValue(null),
+}));
 import {useTheme} from '../useTheme';
 
 import {uiStore} from '../../store';

src/screens/SettingsScreen/__tests__/SettingsScreen.test.tsx

@@ -1,4 +1,5 @@
 import React from 'react';
+import {getBackendDevicesInfo} from 'llama.rn';
 import {Platform, Keyboard} from 'react-native';
 import {runInAction} from 'mobx';
 
@@ -33,6 +34,47 @@ describe('SettingsScreen', () => {
     jest.useRealTimers();
   });
 
+  it.each(['HTP*', 'HTP99'])(
+    'shows saved %s as Hexagon and writes its discovered name when selected',
+    async savedName => {
+      const originalOS = Platform.OS;
+      const originalSettings = {...modelStore.contextInitParams};
+      Platform.OS = 'android';
+      (getBackendDevicesInfo as jest.Mock).mockResolvedValue([
+        {deviceName: 'HTP3', type: 'accel', backend: 'HTP'},
+        {deviceName: 'HTP4', type: 'accel', backend: 'HTP'},
+      ]);
+      runInAction(() => {
+        modelStore.contextInitParams.devices = [savedName];
+        modelStore.contextInitParams.flash_attn_type = 'off';
+      });
+      try {
+        const {getByTestId} = render(<SettingsScreen />, {
+          withSafeArea: true,
+          withNavigation: true,
+        });
+        await waitFor(() => {
+          expect(
+            getByTestId('device-option-hexagon').props.accessibilityState
+              .checked,
+          ).toBe(true);
+        });
+        act(() => {
+          fireEvent.press(getByTestId('device-option-cpu'));
+          fireEvent.press(getByTestId('device-option-hexagon'));
+        });
+        expect(modelStore.setDevices).toHaveBeenLastCalledWith(['HTP3']);
+        expect(modelStore.setFlashAttnType).not.toHaveBeenCalled();
+      } finally {
+        Platform.OS = originalOS;
+        runInAction(() => {
+          modelStore.contextInitParams = originalSettings;
+        });
+        (getBackendDevicesInfo as jest.Mock).mockReset().mockResolvedValue([]);
+      }
+    },
+  );
+
   it('renders settings screen correctly', async () => {
     const {getByText, getByDisplayValue} = render(<SettingsScreen />, {
       withSafeArea: true,

src/store/ModelStore.ts

@@ -31,6 +31,7 @@ import {
   unpairedDraftCandidate,
 } from './draftResolution';
 import {checkGpuSupport} from '../utils/deviceCapabilities';
+import {resolveDeviceSelection} from '../utils/deviceSelection';
 import {
   deepMerge,
   getSHA256Hash,
@@ -573,7 +574,7 @@ class ModelStore {
       cache_type_k: this.contextInitParams.cache_type_k,
       cache_type_v: this.contextInitParams.cache_type_v,
       n_gpu_layers: this.contextInitParams.n_gpu_layers ?? 99,
-      devices: this.contextInitParams.devices, // NEW
+      devices: this.contextInitParams.devices?.slice(),
       kv_unified: this.contextInitParams.kv_unified ?? true, // NEW (default true!)
       n_parallel: this.contextInitParams.n_parallel ?? 1, // NEW (1 for blocking mode only)
       use_mlock: this.contextInitParams.use_mlock,
@@ -617,6 +618,14 @@ class ModelStore {
       }
     }
 
+    Object.assign(
+      params,
+      await resolveDeviceSelection({
+        devices: params.devices,
+        n_gpu_layers: params.n_gpu_layers,
+      }),
+    );
+
     // Remove undefined values from the params object
     return Object.fromEntries(
       Object.entries(params).filter(([_, value]) => value !== undefined),

src/store/__tests__/ModelStore.registerOnboardingPalModel.test.ts

@@ -1,5 +1,9 @@
 jest.unmock('../../store');
 
+jest.mock('../../services/deviceRules/rules', () => ({
+  fetchRules: jest.fn().mockResolvedValue(null),
+}));
+
 import {downloadManager} from '../../services/downloads';
 import {modelStore} from '..';
 import {ModelOrigin} from '../../utils/types';