Estimate compute limit for Orca/OKX

apps/api/src/index.ts

@@ -42,7 +42,7 @@ const providers: Record<string, Protocol> = {
     relay: new RelayProvider(),
     orca: new OrcaWhirlpoolProvider(solanaRpc),
     panora: new PanoraProvider({ rpcUrl: process.env.APTOS_URL }),
-    okx: new OkxProvider(),
+    okx: new OkxProvider(solanaRpc),
 };
 
 app.get("/", (_, res) => {

packages/swapper/src/chain/solana/tx_builder.ts

@@ -1,6 +1,7 @@
 import {
     Connection,
     ComputeBudgetProgram,
+    MessageV0,
     PublicKey,
     TransactionInstruction,
     VersionedTransaction,
@@ -96,6 +97,35 @@ export function findComputeUnitLimit(instructions: TransactionInstruction[]): st
     return undefined;
 }
 
+export async function estimateComputeUnitLimit(
+    connection: Connection,
+    transaction: Transaction | VersionedTransaction,
+): Promise<number | undefined> {
+    try {
+        const versionedTx =
+            transaction instanceof VersionedTransaction
+                ? transaction
+                : new VersionedTransaction(
+                      MessageV0.compile({
+                          payerKey: transaction.feePayer!,
+                          instructions: transaction.instructions,
+                          recentBlockhash: transaction.recentBlockhash!,
+                      }),
+                  );
+        const response = await connection.simulateTransaction(versionedTx, {
+            replaceRecentBlockhash: true,
+            sigVerify: false,
+        });
+        if (response.value.err) {
+            return undefined;
+        }
+        if (response.value.unitsConsumed && response.value.unitsConsumed > 0) {
+            return Math.ceil(response.value.unitsConsumed * COMPUTE_UNIT_MULTIPLIER);
+        }
+    } catch (error) { void error; }
+    return undefined;
+}
+
 export function serializeTransaction(transaction: VersionedTransaction | Transaction): string {
     const serialized =
         transaction instanceof VersionedTransaction

packages/swapper/src/debug.ts

@@ -0,0 +1,14 @@
+/* eslint-disable no-console */
+export function timedPromise<T>(label: string, promise: Promise<T>): Promise<T> {
+    const start = performance.now();
+    return promise.then(
+        (result) => {
+            console.log(`${label}: ${(performance.now() - start).toFixed(0)}ms`);
+            return result;
+        },
+        (err) => {
+            console.log(`${label} FAILED: ${(performance.now() - start).toFixed(0)}ms`);
+            throw err;
+        },
+    );
+}

packages/swapper/src/okx/integration.test.ts

@@ -19,7 +19,7 @@ describe("OKX live integration", () => {
     jest.setTimeout(60_000);
 
     itIntegration("fetches a live quote and builds quote data", async () => {
-        const provider = new OkxProvider();
+        const provider = new OkxProvider(process.env.SOLANA_URL || "https://solana-rpc.publicnode.com");
         const quote = await provider.get_quote(REQUEST_TEMPLATE);
 
         expect(BigInt(quote.output_value) > BigInt(0)).toBe(true);

packages/swapper/src/okx/provider.test.ts

@@ -14,14 +14,9 @@ function createRequest(slippageBps = 100) {
 function createProvider() {
     const getQuote = jest.fn();
     const getSwapData = jest.fn();
-    const getGasLimit = jest.fn().mockResolvedValue({
-        code: "0",
-        msg: "",
-        data: [{ gasLimit: "500000" }],
-    });
-    const client = { dex: { getQuote, getSwapData, getGasLimit } } as unknown as OKXDexClient;
-    const provider = new OkxProvider(client);
-    return { provider, getQuote, getSwapData, getGasLimit };
+    const client = { dex: { getQuote, getSwapData } } as unknown as OKXDexClient;
+    const provider = new OkxProvider("https://localhost:8899", client);
+    return { provider, getQuote, getSwapData };
 }
 
 const mockRoute = {
@@ -95,19 +90,13 @@ describe("OkxProvider", () => {
             expect(swapParams.slippagePercent).toBe("1");
         });
 
-        it("estimates compute unit limit via getGasLimit", async () => {
-            const { provider, getSwapData, getGasLimit } = createProvider();
+        it("returns undefined gasLimit when simulation fails", async () => {
+            const { provider, getSwapData } = createProvider();
             getSwapData.mockResolvedValue(mockSwapResponse());
 
             const result = await provider.get_quote_data(mockQuote());
 
-            expect(getGasLimit).toHaveBeenCalledWith({
-                chainIndex: "501",
-                fromAddress: "SenderAddress",
-                toAddress: "JUP4Fb2cqiRUcaTHdrPC8h2gNsA2ETXiPDD33WcGuJB",
-                extJson: { inputData: SOL_MINT },
-            });
-            expect(result.gasLimit).toBe("550000");
+            expect(result.gasLimit).toBeUndefined();
         });
 
         it("falls back to 1% slippage when slippage_bps is 0", async () => {
@@ -121,9 +110,8 @@ describe("OkxProvider", () => {
             expect(swapParams.maxAutoSlippagePercent).toBeUndefined();
         });
 
-        it("handles getGasLimit failure gracefully", async () => {
-            const { provider, getSwapData, getGasLimit } = createProvider();
-            getGasLimit.mockRejectedValue(new Error("API error"));
+        it("handles simulation failure gracefully", async () => {
+            const { provider, getSwapData } = createProvider();
             getSwapData.mockResolvedValue(mockSwapResponse());
 
             const result = await provider.get_quote_data(mockQuote());

packages/swapper/src/okx/provider.ts

@@ -1,10 +1,13 @@
 import { AssetId, Chain, Quote, QuoteRequest, SwapQuoteData, SwapQuoteDataType } from "@gemwallet/types";
 import { OKXDexClient } from "@okx-dex/okx-dex-sdk";
-import type { QuoteData, SwapParams, TransactionData } from "@okx-dex/okx-dex-sdk";
+import type { QuoteData, SwapParams } from "@okx-dex/okx-dex-sdk";
 import bs58 from "bs58";
 
+import { Connection, VersionedTransaction } from "@solana/web3.js";
+
 import { BigIntMath } from "../bigint_math";
-import { COMPUTE_UNIT_MULTIPLIER } from "../chain/solana/tx_builder";
+import { DEFAULT_COMMITMENT } from "../chain/solana/constants";
+import { estimateComputeUnitLimit as simulateComputeUnits } from "../chain/solana/tx_builder";
 import { SwapperException } from "../error";
 import { Protocol } from "../protocol";
 import { getReferrerAddresses } from "../referrer";
@@ -84,8 +87,11 @@ function minOutputValue(route: QuoteData, slippageBps: number): string {
 
 export class OkxProvider implements Protocol {
     private readonly client: OKXDexClient;
+    private readonly connection: Connection;
+
+    constructor(solanaRpcEndpoint: string, client?: OKXDexClient) {
+        this.connection = new Connection(solanaRpcEndpoint, { commitment: DEFAULT_COMMITMENT });
 
-    constructor(client?: OKXDexClient) {
         if (client) {
             this.client = client;
             return;
@@ -115,22 +121,13 @@ export class OkxProvider implements Protocol {
         });
     }
 
-    private async estimateComputeUnitLimit(tx: TransactionData): Promise<string | undefined> {
+    private async estimateComputeUnitLimit(txData: string): Promise<string | undefined> {
         try {
-            const result = await this.client.dex.getGasLimit({
-                chainIndex: SOLANA_CHAIN_INDEX,
-                fromAddress: tx.from,
-                toAddress: tx.to,
-                extJson: { inputData: tx.data },
-            });
-
-            if (result.code === "0") {
-                const raw = Number(result.data[0]?.gasLimit);
-                if (raw > 0) {
-                    return Math.ceil(raw * COMPUTE_UNIT_MULTIPLIER).toString();
-                }
-            }
-        } catch { }
+            const bytes = bs58.decode(txData);
+            const tx = VersionedTransaction.deserialize(bytes);
+            const estimate = await simulateComputeUnits(this.connection, tx);
+            return estimate?.toString();
+        } catch (error) { void error; }
         return undefined;
     }
 
@@ -192,7 +189,7 @@ export class OkxProvider implements Protocol {
             throw new SwapperException({ type: "invalid_route" });
         }
 
-        const gasLimit = await this.estimateComputeUnitLimit(swapData.tx);
+        const gasLimit = await this.estimateComputeUnitLimit(swapData.tx.data);
 
         let serializedBase64: string;
         try {

packages/swapper/src/orca/provider.ts

@@ -33,6 +33,9 @@ import { getMintAddress, parsePublicKey, resolveTokenProgram } from "../chain/so
 import { DEFAULT_COMMITMENT } from "../chain/solana/constants";
 import {
     addComputeBudgetInstructions,
+    createComputeUnitLimitInstruction,
+    DEFAULT_COMPUTE_UNIT_LIMIT,
+    estimateComputeUnitLimit,
     getRecentBlockhash,
     getRecentPriorityFee,
     serializeTransaction,
@@ -150,23 +153,22 @@ export class OrcaWhirlpoolProvider implements Protocol {
 
         const signer = this.createPassthroughSigner(userPublicKey);
 
-        const swapPromise = swapInstructions(
-            this.rpc,
-            { inputAmount: amount, mint: inputMintAddress },
-            poolAddress,
-            slippageBps,
-            signer,
-        );
-        const priorityFeePromise = this.getPriorityFee();
-        const blockhashPromise = getRecentBlockhash(this.connection, DEFAULT_COMMITMENT);
-
-        const [{ instructions }, priorityFee] = await Promise.all([swapPromise, priorityFeePromise]);
+        const [{ instructions }, priorityFee, { blockhash, lastValidBlockHeight }, referralInstruction] =
+            await Promise.all([
+                swapInstructions(
+                    this.rpc,
+                    { inputAmount: amount, mint: inputMintAddress },
+                    poolAddress,
+                    slippageBps,
+                    signer,
+                ),
+                this.getPriorityFee(),
+                getRecentBlockhash(this.connection, DEFAULT_COMMITMENT),
+                this.buildReferralInstruction(quote, userPublicKey),
+            ]);
 
         const legacyInstructions = instructions.map((instruction) => this.toLegacyInstruction(instruction));
-
-        const computeBudgetInstructions = addComputeBudgetInstructions([], undefined, priorityFee);
-
-        const referralInstruction = await this.buildReferralInstruction(quote, userPublicKey);
+        const computeBudgetInstructions = addComputeBudgetInstructions([], DEFAULT_COMPUTE_UNIT_LIMIT, priorityFee);
 
         const transaction = new Transaction();
         transaction.add(...computeBudgetInstructions, ...legacyInstructions);
@@ -176,18 +178,19 @@ export class OrcaWhirlpoolProvider implements Protocol {
         }
 
         transaction.feePayer = userPublicKey;
-
-        const { blockhash, lastValidBlockHeight } = await blockhashPromise;
-
         setTransactionBlockhash(transaction, blockhash, lastValidBlockHeight);
 
+        const estimatedLimit = await estimateComputeUnitLimit(this.connection, transaction);
+        transaction.instructions[0] = createComputeUnitLimitInstruction(estimatedLimit ?? DEFAULT_COMPUTE_UNIT_LIMIT);
+
         const serialized = serializeTransaction(transaction);
 
         return {
             to: "",
             value: "0",
             data: serialized,
             dataType: SwapQuoteDataType.Contract,
+            gasLimit: estimatedLimit?.toString(),
         };
     }