[AI] Implement BALANCE_OF as HyperFormula custom function via context

Author: StephenBrown2

packages/desktop-client/src/hooks/useFormulaExecution.ts

@@ -327,8 +327,8 @@ export function useFormulaExecution(
         // Create HyperFormula instance
         hfInstance = HyperFormula.buildEmpty({
           licenseKey: 'gpl-v3',
-          localeLang: typeof locale === 'string' ? locale : 'en-US',
           language: 'enUS',
+          localeLang: typeof locale === 'string' ? locale : 'en-US',
         });
 
         // Add a sheet and set the formula in cell A1

packages/desktop-client/src/hooks/useTransactionFormulaExecution.ts

@@ -1,7 +1,21 @@
 import { useEffect, useState } from 'react';
 
+import { HyperFormula } from 'hyperformula';
+import enUS from 'hyperformula/i18n/languages/enUS';
+
+import {
+  CustomFunctionsPlugin,
+  customFunctionsTranslations,
+} from '@actual-app/core/server/rules/customFunctions';
+
 import { useLocale } from './useLocale';
 
+HyperFormula.registerLanguage('enUS', enUS);
+HyperFormula.registerFunctionPlugin(
+  CustomFunctionsPlugin,
+  customFunctionsTranslations,
+);
+
 type TransactionContext = {
   amount?: number;
   date?: string;
@@ -26,8 +40,7 @@ export function useTransactionFormulaExecution(
   useEffect(() => {
     let cancelled = false;
 
-    async function executeFormula() {
-      const { HyperFormula } = await import('hyperformula');
+    function executeFormula() {
       let hfInstance: ReturnType<typeof HyperFormula.buildEmpty> | null = null;
 
       if (!formula || !formula.startsWith('=')) {
@@ -40,7 +53,12 @@ export function useTransactionFormulaExecution(
         // Create HyperFormula instance
         hfInstance = HyperFormula.buildEmpty({
           licenseKey: 'gpl-v3',
+          language: 'enUS',
           localeLang: typeof locale === 'string' ? locale : 'en-US',
+          context: {
+            // No server prefetch in preview
+            balanceOfPrefetch: new Map(),
+          },
         });
 
         // Add a sheet
@@ -59,7 +77,12 @@ export function useTransactionFormulaExecution(
         };
 
         for (const [key, value] of Object.entries(fieldValues)) {
-          if (value !== undefined && value !== null) {
+          if (
+            value !== undefined &&
+            value !== null &&
+            typeof value !== 'object' &&
+            !key.startsWith('_')
+          ) {
             // Set the value in a cell
             hfInstance.setCellContents({ sheet: sheetId, col: 0, row }, [
               [value],

packages/loot-core/src/server/rules/action.ts

@@ -10,7 +10,6 @@ import { currentDay, format, parseDate } from '../../shared/months';
 import { FIELD_TYPES } from '../../shared/rules';
 import type { TransactionForRules } from '../transactions/transaction-rules';
 
-import { substituteBalanceOfLiterals } from './balanceOfFormula';
 import {
   CustomFunctionsPlugin,
   customFunctionsTranslations,
@@ -304,6 +303,10 @@ export class Action {
       hfInstance = HyperFormula.buildEmpty({
         licenseKey: 'gpl-v3',
         language: 'enUS',
+        context: {
+          balanceOfPrefetch:
+            transaction['_balanceOfPrefetched'] ?? new Map(),
+        },
       });
 
       const sheetName = hfInstance.addSheet('Sheet1');
@@ -341,13 +344,8 @@ export class Action {
         hfInstance.addNamedExpression(key, cellValue);
       }
 
-      const evaluatedFormula = substituteBalanceOfLiterals(
-        formula,
-        transaction._balanceOfPrefetched,
-      );
-
       hfInstance.setCellContents({ sheet: sheetId, col: 0, row: 0 }, [
-        [evaluatedFormula],
+        [formula],
       ]);
 
       const cellAddress = { sheet: sheetId, col: 0, row: 0 };

packages/loot-core/src/server/rules/balanceOfFormula.test.ts

@@ -6,7 +6,6 @@ import {
   decodeBalanceOfQuotedLiteral,
   extractBalanceOfLiterals,
   resolveAccountIdForBalanceOf,
-  substituteBalanceOfLiterals,
 } from './balanceOfFormula';
 
 describe('balanceOfFormula', () => {
@@ -25,19 +24,6 @@ describe('balanceOfFormula', () => {
     expect(decodeBalanceOfQuotedLiteral(String.raw`\"x\"`)).toBe('"x"');
   });
 
-  it('substituteBalanceOfLiterals replaces calls with cent literals', () => {
-    const map = new Map([
-      ['Checking', 42],
-      ['id-1', 99],
-    ]);
-    expect(substituteBalanceOfLiterals('=BALANCE_OF("Checking")+1', map)).toBe(
-      '=42+1',
-    );
-    expect(substituteBalanceOfLiterals('=BALANCE_OF("Missing")', map)).toBe(
-      '=0',
-    );
-  });
-
   it('resolveAccountIdForBalanceOf prefers map key then name', () => {
     const id = 'acc-1';
     const a1: db.DbAccount = {

packages/loot-core/src/server/rules/balanceOfFormula.ts

@@ -54,21 +54,3 @@ export function resolveAccountIdForBalanceOf(
   }
   return null;
 }
-
-/**
- * Replace each BALANCE_OF("…") with a cent literal so HyperFormula never needs
- * runtime DB (prefetch map is keyed by decoded string literals).
- */
-export function substituteBalanceOfLiterals(
-  formula: string,
-  map: Map<string, number> | null | undefined,
-): string {
-  return formula.replace(
-    /BALANCE_OF\s*\(\s*"((?:[^"\\]|\\.)*)"\s*\)/gi,
-    (_match, inner: string) => {
-      const key = decodeBalanceOfQuotedLiteral(inner);
-      const cents = map?.get(key) ?? 0;
-      return String(cents);
-    },
-  );
-}

packages/loot-core/src/server/rules/customFunctions.ts

@@ -26,36 +26,55 @@ export class CustomFunctionsPlugin extends FunctionPlugin {
       },
     );
   }
+
+  balanceOf(ast: ProcedureAst, state: InterpreterState) {
+    return this.runFunction(
+      ast.args,
+      state,
+      this.metadata('BALANCE_OF'),
+      (accountKey: string) => {
+        const ctx = this.config.context as
+          | { balanceOfPrefetch?: Map<string, number> }
+          | undefined;
+        return ctx?.balanceOfPrefetch?.get(accountKey) ?? 0;
+      },
+    );
+  }
 }
 
 CustomFunctionsPlugin.implementedFunctions = {
-  INTEGER_TO_AMOUNT: {
-    method: 'integerToAmount',
+  BALANCE_OF: {
+    method: 'balanceOf',
+    parameters: [{ argumentType: FunctionArgumentType.STRING }],
+  },
+  FIXED: {
+    method: 'fixed',
     parameters: [
       { argumentType: FunctionArgumentType.NUMBER },
       {
         argumentType: FunctionArgumentType.NUMBER,
         optionalArg: true,
-        defaultValue: 2,
+        defaultValue: 0,
       },
     ],
   },
-  FIXED: {
-    method: 'fixed',
+  INTEGER_TO_AMOUNT: {
+    method: 'integerToAmount',
     parameters: [
       { argumentType: FunctionArgumentType.NUMBER },
       {
         argumentType: FunctionArgumentType.NUMBER,
         optionalArg: true,
-        defaultValue: 0,
+        defaultValue: 2,
       },
     ],
   },
 };
 
 export const customFunctionsTranslations = {
   enUS: {
-    INTEGER_TO_AMOUNT: 'INTEGER_TO_AMOUNT',
+    BALANCE_OF: 'BALANCE_OF',
     FIXED: 'FIXED',
+    INTEGER_TO_AMOUNT: 'INTEGER_TO_AMOUNT',
   },
 };