feat: add dry run swap and integration mainnet

package.json

@@ -53,7 +53,8 @@
     "test": "turbo run test",
     "test:cov": "turbo run test:cov",
     "test:watch": "turbo run test:watch",
-    "test:integration": "turbo run test:integration",
+    "test:integration:testnet": "turbo run test:integration:testnet",
+    "test:integration:mainnet": "turbo run test:integration:mainnet",
     "test:e2e": "turbo run test:e2e",
     "clean": "turbo run clean",
     "build": "turbo run build",

packages/core/src/defi/swap.ts

@@ -165,7 +165,6 @@ function getCrossChainMultilocations(args: SwapTokenArgs) {
  */
 function formatCrossChainAmount(args: SwapTokenArgs): string {
   const decimals = getAssetDecimals(args.from as TNodeDotKsmWithRelayChains, args.currencyFrom)
-
   if (!decimals) {
     throw new Error(`Failed to get decimals for ${args.currencyFrom} on ${args.from}`)
   }
@@ -243,6 +242,31 @@ async function validateSwapFees({
 }): Promise<void> {
   const fees = await builder.getXcmFees()
 
+  // Check origin balance sufficiency
+  if (!fees.origin.sufficient) {
+    throw new Error(`Unable to swap due to insufficient balance`)
+  }
+
+  // Check destination balance sufficiency
+  if (!fees.destination.sufficient) {
+    throw new Error(`Unable to swap due to insufficient destination balance`)
+  }
+
+  // Check each hop for sufficiency and errors
+  if (fees.hops && fees.hops.length > 0) {
+    for (const hop of fees.hops) {
+      if (!hop.result.sufficient) {
+        throw new Error(
+          `Insufficient balance for hop on ${hop.chain}: ${hop.result.dryRunError || "Unknown error"}`
+        )
+      }
+
+      if (hop.result.dryRunError) {
+        throw new Error(`Dry run error on ${hop.chain}: ${hop.result.dryRunError}`)
+      }
+    }
+  }
+
   if (fees.failureChain || fees.failureReason) {
     throw new Error(
       `Failed to calculate ${swapType} swap fees: ${fees.failureChain || fees.failureReason}`

packages/llm/src/types/constants.ts

@@ -1,6 +1,6 @@
 // DEX swap constants
 export const DEFAULT_DEX = "HydrationDex"
 export const UNKNOWN_ERROR_MESSAGE = "Unknown error occurred"
-export const SWAP_TIMEOUT_MS = 100000 // 100 seconds
+export const SWAP_TIMEOUT_MS = 600000
 export const MAX_RETRIES = 2
 export const RETRY_DELAY_MS = 2000 // 2 seconds

packages/llm/src/types/defi/swap.ts

@@ -16,8 +16,8 @@ import { ToolNames } from "../common"
  * Cross-chain swap (one-click):
  * ```typescript
  * {
- *   from: "polkadot",
- *   to: "hydra",
+ *   from: "Polkadot",
+ *   to: "Hydration",
  *   currencyFrom: "DOT",
  *   currencyTo: "HDX",
  *   amount: "10000000000",
@@ -42,13 +42,13 @@ export const swapTokensToolSchema = z.object({
     .string()
     .optional()
     .describe(
-      "The source chain ID where the swap originates (e.g., 'polkadot', 'kusama', 'hydra'). Required for cross-chain swaps."
+      "The source chain ID where the swap originates (e.g., 'Polkadot', 'Kusama', 'Hydra', 'AssetHubPolkadot'). Required for cross-chain swaps."
     ),
   to: z
     .string()
     .optional()
     .describe(
-      "The destination chain ID where the swap completes (e.g., 'polkadot', 'kusama', 'hydra'). Required for cross-chain swaps."
+      "The destination chain ID where the swap completes (e.g., 'Polkadot', 'Kusama', 'Hydra', 'AssetHubPolkadot'). Required for cross-chain swaps."
     ),
   currencyFrom: z
     .string()
@@ -78,8 +78,8 @@ export const swapTokensToolSchema = z.object({
  * ```typescript
  * const swapTool: SwapTokensTool = swapTokensTool(apis, signer);
  * const result = await swapTool.invoke({
- *   from: "polkadot",
- *   to: "hydra",
+ *   from: "Polkadot",
+ *   to: "Hydration",
  *   currencyFrom: "DOT",
  *   currencyTo: "HDX",
  *   amount: "10000000000"

packages/sdk/package.json

@@ -37,7 +37,8 @@
     "test:cov": "vitest run --coverage",
     "test:watch": "vitest",
     "test:e2e": "vitest run --config ./vitest.config.e2e.ts --sequence.concurrent",
-    "test:integration": "vitest run --config ./vitest.config.integration.ts --sequence.concurrent"
+    "test:integration:testnet": "vitest run --config ./vitest.config.testnet.ts --sequence.concurrent",
+    "test:integration:mainnet": "vitest run --config ./vitest.config.mainnet.ts --sequence.concurrent"
   },
   "dependencies": {
     "@langchain/core": "^0.3.40",

packages/sdk/tests/integration-tests/sdk.itest.ts

@@ -177,13 +177,13 @@ describe('PolkadotAgentKit Integration with OllamaAgent', () => {
 
   it('should call xcm_transfer_native_asset tool for West Asset Hub to West People Chain transfer', async () => {
 
-    const userQuery = `transfer 0.5 WND to ${RECIPIENT4} from AssetHubWestend to PeopleWestend via XCM`;
+    const userQuery = `transfer 0.2 WND to ${RECIPIENT4} from AssetHubWestend to PeopleWestend via XCM`;
 
     // Get balances before transfer
     const balanceAgentBefore = await getBalance(agentKit.getApi('west_asset_hub'), agentKit.getCurrentAddress());
     const balanceRecipientBefore = await getBalance(agentKit.getApi('west_people'), RECIPIENT4);
 
-    const amount = parseUnits("0.5", getDecimalsByChainId('west_asset_hub'));
+    const amount = parseUnits("0.2", getDecimalsByChainId('west_asset_hub'));
 
     const result = await ollamaAgent.ask(userQuery);
     console.log('XCM Transfer Query Result (West Asset Hub → West People Chain):', result);
@@ -198,7 +198,7 @@ describe('PolkadotAgentKit Integration with OllamaAgent', () => {
 
     expect(xcmTransferCall).toBeDefined();
     expect(xcmTransferCall.action.toolInput).toMatchObject({
-      amount: '0.5',
+      amount: '0.2',
       to: RECIPIENT4,
       sourceChain: 'AssetHubWestend',  
       destChain: 'PeopleWestend'      
@@ -220,4 +220,6 @@ describe('PolkadotAgentKit Integration with OllamaAgent', () => {
     expect(balanceAgentAfter.data.free).toBeLessThan(balanceAgentBefore.data.free - amount - feeXCM.fee);
 
   }, 1500000);
+
+
 });

packages/sdk/tests/integration-tests/sdk.mainnet.itest.ts

@@ -0,0 +1,55 @@
+import { describe, it, expect, beforeAll, afterAll } from 'vitest';
+import { PolkadotAgentKit } from '../../src/api';
+import { OllamaAgent } from './ollamaAgent';
+import dotenv from 'dotenv';
+dotenv.config({ path: '../../.env' });
+let agentKit: PolkadotAgentKit;
+let ollamaAgent: OllamaAgent;
+
+beforeAll(async () => {
+
+
+    if (process.env.AGENT_PRIVATE_KEY_MAINNET) {
+        agentKit = new PolkadotAgentKit({ privateKey: process.env.AGENT_PRIVATE_KEY_MAINNET, keyType: 'Sr25519', chains: ['polkadot'] });
+        await agentKit.initializeApi();
+        ollamaAgent = new OllamaAgent(agentKit);
+        await ollamaAgent.init();
+    } else {
+        throw new Error('AGENT_PRIVATE_KEY_MAINNET is not set');
+    }
+}, 1500000);
+
+afterAll(async () => {
+    await agentKit.disconnect();
+});
+
+
+describe('PolkadotAgentKit Integration with OllamaAgent', () => {
+
+
+    it('should call swap_tokens tool for Polkadot to Polkadot Asset Hub swap', async () => {
+        const userQuery = `swap 0.1 DOT from Polkadot to USDt on Polkadot Asset Hub`;
+        const result = await ollamaAgent.ask(userQuery);
+        console.log('Swap Tokens Query Result (Polkadot → Polkadot Asset Hub):', result);
+
+        expect(result.output).toBeDefined();
+        expect(result.intermediateSteps).toBeDefined();
+        expect(result.intermediateSteps.length).toBeGreaterThan(0);
+
+        const swapCall = result.intermediateSteps.find((step: any) =>
+            step.action?.tool === 'swap_tokens'
+        );
+
+        expect(swapCall).toBeDefined();
+        expect(swapCall.action.toolInput).toMatchObject({
+            amount: '0.1',
+            currencyFrom: 'DOT',
+            currencyTo: 'USDt',
+            from: 'Polkadot',
+            to: 'AssetHubPolkadot'
+        });
+
+
+    }, 1500000);
+
+});

packages/sdk/tests/integration-tests/utils.ts

@@ -1,12 +1,12 @@
 import type { Api, KnownChainId } from "@polkadot-agent-kit/common"
-import {ASSETS_PROMPT, SWAP_PROMPT, NOMINATION_PROMPT, IDENTITY_PROMPT, BIFROST_PROMPT, DYNAMIC_CHAIN_INITIALIZATION_PROMPT} from "@polkadot-agent-kit/llm"
+import {ASSETS_PROMPT, SWAP_PROMPT, NOMINATION_PROMPT, IDENTITY_PROMPT, BIFROST_PROMPT} from "@polkadot-agent-kit/llm"
 
 import { UnsafeTransactionType } from "@polkadot-agent-kit/common"
 export const RECIPIENT= '5CcqKCNDxrYYkPNWys8yrjHJVTzd69i66VTgtewrSbJiVqoR';
 export const RECIPIENT2 = '5D7jcv6aYbhbYGVY8k65oemM6FVNoyBfoVkuJ5cbFvbefftr';
 export const RECIPIENT3 = '5FdxcDTshU5yhHrC91NneaJ64XCE2jwxnMCv8bfxQbwhkWMG';
 export const RECIPIENT4 = '5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY';
-export const SYSTEM_PROMPT = ASSETS_PROMPT + SWAP_PROMPT + NOMINATION_PROMPT + IDENTITY_PROMPT + BIFROST_PROMPT + DYNAMIC_CHAIN_INITIALIZATION_PROMPT;
+export const SYSTEM_PROMPT = ASSETS_PROMPT + SWAP_PROMPT + NOMINATION_PROMPT + IDENTITY_PROMPT + BIFROST_PROMPT;
 
 export function sleep(ms: number) {
   return new Promise(resolve => setTimeout(resolve, ms));

packages/sdk/vitest.config.mainnet.ts

@@ -0,0 +1,13 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    globals: true,
+    environment: 'node',
+    include: ['./tests/integration-tests/sdk.mainnet.itest.ts'],
+    exclude: ['./tests/integration-tests/sdk.itest.ts'],
+    testTimeout: 600000,
+    hookTimeout: 600000,
+    onConsoleLog: () => true,
+  }
+})

packages/sdk/vitest.config.testnet.ts

@@ -0,0 +1,13 @@
+import { defineConfig } from 'vitest/config'
+
+export default defineConfig({
+  test: {
+    globals: true,
+    environment: 'node',
+    include: ['./tests/integration-tests/sdk.itest.ts'],
+    exclude: ['./tests/integration-tests/sdk.mainnet.itest.ts'],
+    testTimeout: 600000,
+    hookTimeout: 600000,
+    onConsoleLog: () => true,
+  }
+})

turbo.json

@@ -13,7 +13,10 @@
       "test:watch": {
         "cache": false
       },
-      "test:integration": {
+      "test:integration:testnet": {
+        "cache": false
+      },
+      "test:integration:mainnet": {
         "cache": false
       },
       "test:e2e": {},