onboarding more reliable (#720)

apps/extension/package.json

@@ -25,6 +25,7 @@
     "@penumbra-zone/wasm": "workspace:*",
     "@tanstack/react-query": "^5.25.0",
     "buffer": "^6.0.3",
+    "exponential-backoff": "^3.1.1",
     "framer-motion": "^11.0.8",
     "immer": "^10.0.4",
     "node-fetch": "^3.3.2",

apps/extension/public/manifest.json

@@ -26,6 +26,10 @@
       "run_at": "document_start"
     }
   ],
+  "options_ui": {
+    "page": "page.html",
+    "open_in_tab": true
+  },
   "background": {
     "service_worker": "service-worker.js"
   },

apps/extension/src/hooks/onboarding.ts

@@ -19,6 +19,6 @@ export const useOnboardingSave = () => {
     await setPassword(plaintextPassword);
 
     await addWallet({ label: 'Wallet #1', seedPhrase });
-    void chrome.runtime.sendMessage(ServicesMessage.ClearCache);
+    void chrome.runtime.sendMessage(ServicesMessage.OnboardComplete);
   };
 };

apps/extension/src/routes/popup/home/index.tsx

@@ -1,11 +1,10 @@
-import { redirect } from 'react-router-dom';
 import { SelectAccount } from '@penumbra-zone/ui';
-import { PopupPath } from '../paths';
 import { IndexHeader } from './index-header';
 import { useStore } from '../../../state';
 import { BlockSync } from './block-sync';
-import { localExtStorage, sessionExtStorage } from '@penumbra-zone/storage';
+import { localExtStorage } from '@penumbra-zone/storage';
 import { addrByIndexSelector } from '../../../state/wallets';
+import { needsLogin } from '../popup-needs';
 
 export interface PopupLoaderData {
   fullSyncHeight: number;
@@ -14,22 +13,8 @@ export interface PopupLoaderData {
 // Because Zustand initializes default empty (prior to persisted storage synced),
 // We need to manually check storage for accounts & password in the loader.
 // Will redirect to onboarding or password check if necessary.
-export const popupIndexLoader = async (): Promise<Response | PopupLoaderData> => {
-  const wallets = await localExtStorage.get('wallets');
-
-  if (!wallets.length) {
-    await chrome.tabs.create({ url: chrome.runtime.getURL(`page.html`) });
-    window.close();
-  }
-
-  const password = await sessionExtStorage.get('passwordKey');
-
-  if (!password) return redirect(PopupPath.LOGIN);
-
-  const fullSyncHeight = await localExtStorage.get('fullSyncHeight');
-
-  return { fullSyncHeight };
-};
+export const popupIndexLoader = async (): Promise<Response | PopupLoaderData> =>
+  (await needsLogin()) ?? { fullSyncHeight: await localExtStorage.get('fullSyncHeight') };
 
 export const PopupIndex = () => {
   const getAddrByIndex = useStore(addrByIndexSelector);

apps/extension/src/routes/popup/login.tsx

@@ -6,6 +6,9 @@ import { useStore } from '../../state';
 import { passwordSelector } from '../../state/password';
 import { FormEvent, useState } from 'react';
 import { PopupPath } from './paths';
+import { needsOnboard } from './popup-needs';
+
+export const popupLoginLoader = () => needsOnboard();
 
 export const Login = () => {
   const navigate = usePopupNav();

apps/extension/src/routes/popup/popup-needs.ts

@@ -0,0 +1,18 @@
+import { redirect } from 'react-router-dom';
+import { PopupPath } from './paths';
+import { localExtStorage, sessionExtStorage } from '@penumbra-zone/storage';
+
+export const needsLogin = async (): Promise<Response | undefined> => {
+  const password = await sessionExtStorage.get('passwordKey');
+  if (password) return;
+
+  return redirect(PopupPath.LOGIN);
+};
+
+export const needsOnboard = async () => {
+  const wallets = await localExtStorage.get('wallets');
+  if (wallets.length) return;
+
+  void chrome.runtime.openOptionsPage();
+  window.close();
+};

apps/extension/src/routes/popup/router.tsx

@@ -1,6 +1,6 @@
 import { createHashRouter, RouteObject } from 'react-router-dom';
 import { PopupIndex, popupIndexLoader } from './home';
-import { Login } from './login';
+import { Login, popupLoginLoader } from './login';
 import { PopupPath } from './paths';
 import { PopupLayout } from './popup-layout';
 import { Settings } from './settings';
@@ -20,6 +20,7 @@ export const popupRoutes: RouteObject[] = [
       {
         path: PopupPath.LOGIN,
         element: <Login />,
+        loader: popupLoginLoader,
       },
       {
         path: PopupPath.SETTINGS,

apps/extension/src/service-worker.ts

@@ -35,18 +35,33 @@ import { approveTransaction } from './approve-transaction';
 
 // all rpc implementations, local and proxy
 import { rpcImpls } from './impls';
+import { backOff } from 'exponential-backoff';
 
-const services = new Services({
-  idbVersion: IDB_VERSION,
-  grpcEndpoint: await localExtStorage.get('grpcEndpoint'),
-  getWallet: async () => {
-    const wallets = await localExtStorage.get('wallets');
-    if (!wallets[0]) throw new Error('No wallets connected');
-    const { fullViewingKey, id } = wallets[0];
-    return { walletId: id, fullViewingKey };
+// prevent spamming the focus-stealing openOptionsPage
+let openOptionsOnce: undefined | Promise<void>;
+const startServices = async () => {
+  const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
+
+  const wallet0 = (await localExtStorage.get('wallets'))[0];
+  if (!wallet0) openOptionsOnce ??= chrome.runtime.openOptionsPage();
+
+  const services = new Services({
+    idbVersion: IDB_VERSION,
+    grpcEndpoint,
+    walletId: wallet0?.id,
+    fullViewingKey: wallet0?.fullViewingKey,
+  });
+  await services.initialize();
+  return services;
+};
+
+const services = await backOff(startServices, {
+  retry: (e, attemptNumber) => {
+    if (process.env['NODE_ENV'] === 'development')
+      console.warn("Prax couldn't start ", attemptNumber, e);
+    return true;
   },
 });
-await services.initialize();
 
 let custodyClient: PromiseClient<typeof CustodyService> | undefined;
 let stakingClient: PromiseClient<typeof StakingService> | undefined;

packages/router/src/grpc/view-protocol-server/app-parameters.test.ts

@@ -1,4 +1,4 @@
-import { beforeEach, describe, expect, test, vi } from 'vitest';
+import { Mock, beforeEach, describe, expect, test, vi } from 'vitest';
 import {
   AppParametersRequest,
   AppParametersResponse,
@@ -15,15 +15,30 @@ describe('AppParameters request handler', () => {
   let mockServices: MockServices;
   let mockIndexedDb: IndexedDbMock;
   let mockCtx: HandlerContext;
+  let apSubNext: Mock;
 
   beforeEach(() => {
     vi.resetAllMocks();
 
+    apSubNext = vi.fn();
+
+    const mockAppParametersSubscription = {
+      next: apSubNext,
+      [Symbol.asyncIterator]: () => mockAppParametersSubscription,
+    };
+
     mockIndexedDb = {
       getAppParams: vi.fn(),
+      subscribe: (table: string) => {
+        if (table === 'APP_PARAMETERS') return mockAppParametersSubscription;
+        throw new Error('Table not supported');
+      },
     };
+
     mockServices = {
-      getWalletServices: vi.fn(() => Promise.resolve({ indexedDb: mockIndexedDb })),
+      getWalletServices: vi.fn(() =>
+        Promise.resolve({ indexedDb: mockIndexedDb }),
+      ) as MockServices['getWalletServices'],
     };
     mockCtx = createHandlerContext({
       service: ViewService,
@@ -46,9 +61,12 @@ describe('AppParameters request handler', () => {
     expect(appParameterResponse.parameters?.equals(testData)).toBeTruthy();
   });
 
-  test('should fail to get appParameters when idb has none', async () => {
+  test('should wait for appParameters when idb has none', async () => {
     mockIndexedDb.getAppParams?.mockResolvedValue(undefined);
-    await expect(appParameters(new AppParametersRequest(), mockCtx)).rejects.toThrow();
+    apSubNext.mockResolvedValueOnce({
+      value: { value: new AppParametersRequest(), table: 'APP_PARAMETERS' },
+    });
+    await expect(appParameters(new AppParametersRequest(), mockCtx)).resolves.toBeTruthy();
   });
 });
 

packages/router/src/grpc/view-protocol-server/app-parameters.ts

@@ -1,10 +1,16 @@
+import { AppParameters } from '@buf/penumbra-zone_penumbra.bufbuild_es/penumbra/core/app/v1/app_pb';
 import type { Impl } from '.';
 import { servicesCtx } from '../../ctx';
 
 export const appParameters: Impl['appParameters'] = async (_, ctx) => {
   const services = ctx.values.get(servicesCtx);
   const { indexedDb } = await services.getWalletServices();
+
+  const subscription = indexedDb.subscribe('APP_PARAMETERS');
   const parameters = await indexedDb.getAppParams();
-  if (!parameters) throw new Error('App parameters not available');
-  return { parameters };
+  if (parameters) return { parameters };
+  for await (const update of subscription)
+    return { parameters: AppParameters.fromJson(update.value) };
+
+  throw new Error('App parameters not available');
 };

packages/services/package.json

@@ -12,5 +12,8 @@
     "@penumbra-zone/storage": "workspace:*",
     "@penumbra-zone/wasm": "workspace:*",
     "exponential-backoff": "^3.1.1"
+  },
+  "devDependencies": {
+    "@penumbra-zone/polyfills": "workspace:*"
   }
 }

packages/services/src/index.ts

@@ -1,32 +1,53 @@
 import { BlockProcessor, RootQuerier } from '@penumbra-zone/query';
-import { IndexedDb, syncLastBlockWithLocal } from '@penumbra-zone/storage';
+import { IndexedDb, localExtStorage, syncLastBlockWithLocal } from '@penumbra-zone/storage';
 import { ViewServer } from '@penumbra-zone/wasm';
 import {
   ServicesInterface,
   ServicesMessage,
   WalletServices,
 } from '@penumbra-zone/types/src/services';
 import type { JsonValue } from '@bufbuild/protobuf';
+import '@penumbra-zone/polyfills/Promise.withResolvers';
 
 export interface ServicesConfig {
-  grpcEndpoint: string;
-  idbVersion: number;
-  getWallet(): Promise<{ walletId: string; fullViewingKey: string }>;
+  readonly idbVersion: number;
+  readonly grpcEndpoint?: string;
+  readonly walletId?: string;
+  readonly fullViewingKey?: string;
 }
 
+const isCompleteServicesConfig = (c: Partial<ServicesConfig>): c is Required<ServicesConfig> =>
+  c.grpcEndpoint != null && c.idbVersion != null && c.walletId != null && c.fullViewingKey != null;
+
 export class Services implements ServicesInterface {
   private walletServicesPromise: Promise<WalletServices> | undefined;
+  private config: Promise<Required<ServicesConfig>>;
+
+  constructor(initConfig: ServicesConfig) {
+    const {
+      promise: completeConfig,
+      resolve: resolveConfig,
+      reject: rejectConfig,
+    } = Promise.withResolvers<Required<ServicesConfig>>();
+    this.config = completeConfig;
+
+    if (isCompleteServicesConfig(initConfig)) resolveConfig(initConfig);
 
-  constructor(private readonly config: ServicesConfig) {
     // Attach a listener to allow extension documents to control services.
     // Note that you can't activate this handler from another part of the background script.
-    chrome.runtime.onMessage.addListener((req: JsonValue, sender, respond) => {
+    chrome.runtime.onMessage.addListener((req: JsonValue, sender, respond: () => void) => {
+      const emptyResponse = () => respond();
       if (sender.origin !== origin || typeof req !== 'string') return false;
       switch (req in ServicesMessage && (req as ServicesMessage)) {
         case false:
           return false;
         case ServicesMessage.ClearCache:
-          void this.clearCache().then(() => respond());
+          void this.clearCache().then(emptyResponse);
+          return true;
+        case ServicesMessage.OnboardComplete:
+          void this.completeConfig(initConfig)
+            .then(resolveConfig, rejectConfig)
+            .then(emptyResponse);
           return true;
       }
     });
@@ -40,7 +61,8 @@ export class Services implements ServicesInterface {
   }
 
   public async initialize(): Promise<void> {
-    this._querier = new RootQuerier({ grpcEndpoint: this.config.grpcEndpoint });
+    const { grpcEndpoint } = await this.config;
+    this._querier = new RootQuerier({ grpcEndpoint });
 
     // initialize walletServices separately without exponential backoff to bubble up errors immediately
     await this.getWalletServices();
@@ -61,7 +83,7 @@ export class Services implements ServicesInterface {
   }
 
   private async initializeWalletServices(): Promise<WalletServices> {
-    const { walletId, fullViewingKey } = await this.config.getWallet();
+    const { walletId, fullViewingKey, idbVersion: dbVersion } = await this.config;
     const params = await this.querier.app.appParams();
     if (!params.sctParams?.epochDuration) throw new Error('Epoch duration unknown');
     const {
@@ -71,7 +93,7 @@ export class Services implements ServicesInterface {
 
     const indexedDb = await IndexedDb.initialize({
       chainId,
-      dbVersion: this.config.idbVersion,
+      dbVersion,
       walletId,
     });
 
@@ -93,6 +115,19 @@ export class Services implements ServicesInterface {
     return { viewServer, blockProcessor, indexedDb, querier: this.querier };
   }
 
+  private async completeConfig(initConfig: ServicesConfig): Promise<Required<ServicesConfig>> {
+    const grpcEndpoint = await localExtStorage.get('grpcEndpoint');
+    const wallet0 = (await localExtStorage.get('wallets'))[0];
+    if (!wallet0) throw Error('No wallets found');
+    const { id: walletId, fullViewingKey } = wallet0;
+    return {
+      ...initConfig,
+      grpcEndpoint,
+      walletId,
+      fullViewingKey,
+    };
+  }
+
   private async clearCache() {
     const ws = await this.getWalletServices();
 

packages/storage/src/indexed-db/index.ts

@@ -281,7 +281,9 @@ export class IndexedDb implements IndexedDbInterface {
   async getAppParams(): Promise<AppParameters | undefined> {
     const json = await this.db.get('APP_PARAMETERS', 'params');
     if (!json) return undefined;
-    return AppParameters.fromJson(json);
+    const appParams = AppParameters.fromJson(json);
+    if (!appParams.chainId) return undefined;
+    return appParams;
   }
 
   async saveAppParams(app: AppParameters): Promise<void> {

packages/types/src/services.ts

@@ -17,5 +17,6 @@ export interface ServicesInterface {
 }
 
 export enum ServicesMessage {
+  OnboardComplete = 'OnboardComplete',
   ClearCache = 'ClearCache',
 }