Merge pull request #43 from szhygulin/feat/tron-preflight-status

TRON: pre-flight simulation, status polling, verify-decode preview

Author: szhygulin

src/index.ts

@@ -904,7 +904,7 @@ async function main() {
     "get_transaction_status",
     {
       description:
-        "Poll a transaction's status via the chain's RPC. Returns pending / success / failed, or unknown if the node hasn't seen it yet.",
+        "Poll a transaction's status via the chain's RPC (EVM) or TronGrid (TRON). Returns pending / success / failed, or unknown if the node hasn't seen it yet. Pass chain='tron' with the bare hex txID for TRON.",
       inputSchema: getTransactionStatusInput.shape,
     },
     handler(getTransactionStatus)

src/modules/execution/index.ts

@@ -22,6 +22,7 @@ import {
   tronPathForAccountIndex,
 } from "../../signing/tron-usb-signer.js";
 import { broadcastTronTx } from "../tron/broadcast.js";
+import { getTronTransactionStatus } from "../tron/status.js";
 import { assertTransactionSafe } from "../../signing/pre-sign-check.js";
 import {
   eip1559PreSignHash,
@@ -673,6 +674,9 @@ export async function sendTransaction(args: SendTransactionArgs): Promise<{
 }
 
 export async function getTransactionStatus(args: GetTransactionStatusArgs) {
+  if (args.chain === "tron") {
+    return getTronTransactionStatus(args.txHash);
+  }
   const client = getClient(args.chain as SupportedChain);
   try {
     const receipt = await client.getTransactionReceipt({ hash: args.txHash as `0x${string}` });
@@ -750,7 +754,10 @@ export function getTxVerification(args: GetTxVerificationArgs): UnsignedTx | Uns
  * client-rendered SPA output. One MCP tool = one auditable code path.
  */
 export async function verifyTxDecode(args: GetTxVerificationArgs): Promise<VerifyDecodeResult> {
-  if (hasTronHandle(args.handle)) return notApplicableForTron();
+  if (hasTronHandle(args.handle)) {
+    const tronTx = consumeTronHandle(args.handle);
+    return notApplicableForTron(tronTx);
+  }
   if (!hasHandle(args.handle)) {
     throw new Error(
       `Unknown or expired tx handle '${args.handle}'. Prepared transactions live for ` +

src/modules/execution/schemas.ts

@@ -137,8 +137,15 @@ export const previewSendInput = z.object({
 });
 
 export const getTransactionStatusInput = z.object({
-  chain: chainEnum,
-  txHash: z.string().regex(/^0x[a-fA-F0-9]{64}$/),
+  chain: z
+    .enum([...SUPPORTED_CHAINS, "tron"] as unknown as [string, ...string[]])
+    .describe("EVM chain or 'tron'."),
+  txHash: z
+    .string()
+    .regex(/^(0x)?[a-fA-F0-9]{64}$/)
+    .describe(
+      "32-byte tx hash as hex. EVM txs are conventionally 0x-prefixed; TRON tx IDs are bare hex — both are accepted."
+    ),
 });
 
 export const getTxVerificationInput = z.object({

src/modules/tron/actions.ts

@@ -94,6 +94,112 @@ interface TrongridTriggerResponse {
   };
 }
 
+interface TrongridConstantResponse {
+  result?: { result?: boolean; message?: string; code?: string };
+  energy_used?: number;
+  constant_result?: string[];
+}
+
+/**
+ * Mainnet energy price in sun-per-energy. Hardcoded at the October 2024
+ * governance value (420 sun/energy). If governance changes it, the estimate
+ * drifts; fee_limit (the cap) is still enforced by the network, so drift
+ * here just affects the preview string. A dynamic read via
+ * /wallet/getchainparameters is possible but not worth the extra round-trip
+ * for a preview-only number.
+ */
+const ENERGY_PRICE_SUN = 420n;
+
+/** Well-known solidity revert selector: Error(string) = 0x08c379a0. */
+const ERROR_STRING_SELECTOR = "08c379a0";
+
+/**
+ * Decode the revert payload from a triggerconstantcontract constant_result.
+ * The network returns ABI-encoded revert data: 4-byte Error(string) selector
+ * plus an ABI-encoded string. We crudely extract the string bytes without
+ * pulling in viem — this helper lives on the TRON path which is otherwise
+ * viem-free.
+ */
+function decodeRevertString(constantResult: string[] | undefined): string | undefined {
+  if (!constantResult || constantResult.length === 0) return undefined;
+  const hex = constantResult[0].replace(/^0x/, "");
+  if (!hex.startsWith(ERROR_STRING_SELECTOR)) return undefined;
+  const body = hex.slice(ERROR_STRING_SELECTOR.length);
+  // body = 32-byte offset (ignored) + 32-byte length + string bytes padded to 32.
+  if (body.length < 128) return undefined;
+  const lengthHex = body.slice(64, 128);
+  const length = parseInt(lengthHex, 16);
+  if (!Number.isFinite(length) || length <= 0 || length > body.length / 2) return undefined;
+  const stringHex = body.slice(128, 128 + length * 2);
+  try {
+    return Buffer.from(stringHex, "hex").toString("utf8");
+  } catch {
+    return undefined;
+  }
+}
+
+/**
+ * Dry-run a smart-contract call via /wallet/triggerconstantcontract. This is
+ * TRON's eth_call analogue — it executes the call against current state
+ * without building a broadcastable tx. We use it as a pre-flight before
+ * /wallet/triggersmartcontract so we refuse to hand out a handle for a tx
+ * that would revert on-chain (insufficient balance, paused token, blocked
+ * recipient). Returns the energy estimate for the preview.
+ */
+async function preflightConstantContract(
+  body: Record<string, unknown>,
+  apiKey: string | undefined
+): Promise<{ energyUsed: bigint }> {
+  const res = await trongridPost<TrongridConstantResponse>(
+    "/wallet/triggerconstantcontract",
+    body,
+    apiKey
+  );
+  if (res.result?.result === false) {
+    throw new Error(
+      `TronGrid pre-flight rejected the call: ${res.result.message ?? "unknown validation error"}`
+    );
+  }
+  const revert = decodeRevertString(res.constant_result);
+  if (revert) {
+    throw new Error(
+      `TronGrid pre-flight reverted: ${revert}. This tx would fail on-chain — refusing to prepare a handle.`
+    );
+  }
+  const energyUsed = BigInt(res.energy_used ?? 0);
+  return { energyUsed };
+}
+
+interface TrongridGetAccountResponse {
+  latest_withdraw_time?: number;
+}
+
+const CLAIM_COOLDOWN_MS = 24 * 60 * 60 * 1000;
+
+async function readClaimCooldownRemaining(
+  owner: string,
+  apiKey: string | undefined
+): Promise<number | null> {
+  const res = await trongridPost<TrongridGetAccountResponse>(
+    "/wallet/getaccount",
+    { address: owner, visible: true },
+    apiKey
+  );
+  const last = res.latest_withdraw_time;
+  if (!last) return null;
+  const elapsed = Date.now() - last;
+  if (elapsed >= CLAIM_COOLDOWN_MS) return 0;
+  return CLAIM_COOLDOWN_MS - elapsed;
+}
+
+function formatDuration(ms: number): string {
+  const totalMin = Math.ceil(ms / 60_000);
+  const h = Math.floor(totalMin / 60);
+  const m = totalMin % 60;
+  if (h === 0) return `${m}m`;
+  return `${h}h ${m}m`;
+}
+
 // ----- Native TRX send -----
 
 export interface BuildTronNativeSendArgs {
@@ -216,6 +322,12 @@ export async function buildTronTokenSend(
     visible: true,
   };
   const apiKey = resolveTronApiKey(readUserConfig());
+  // Pre-flight dry-run via triggerconstantcontract. Catches the broad class of
+  // prepare-succeeds-then-broadcast-reverts failures: insufficient token
+  // balance, USDT blocklist, paused contract. Also gives us the energy
+  // estimate for the preview (vs. the fee_limit cap).
+  const { energyUsed } = await preflightConstantContract(body, apiKey);
+  const estimatedEnergySun = energyUsed * ENERGY_PRICE_SUN;
   const res = await trongridPost<TrongridTriggerResponse>(
     "/wallet/triggersmartcontract",
     body,
@@ -258,6 +370,8 @@ export async function buildTronTokenSend(
       },
     },
     feeLimitSun: feeLimitSun.toString(),
+    estimatedEnergyUsed: energyUsed.toString(),
+    estimatedEnergyCostSun: estimatedEnergySun.toString(),
   };
   return issueTronHandle(tx);
 }
@@ -559,8 +673,22 @@ export async function buildTronClaimRewards(
     apiKey
   );
   if (res.Error) {
-    // Common case: "WithdrawBalance not allowed, need 24 hours since last Withdraw"
-    // — TRON enforces a 24h rate limit on claims. Surface TronGrid's message verbatim.
+    // Most common failure: TRON's 24h-between-claims rate limit. TronGrid
+    // returns "WithdrawBalance not allowed, need 24 hours since last Withdraw"
+    // without telling you when the cooldown expires. Read the account's
+    // latest_withdraw_time and translate to "claim again in X hours Y min"
+    // so the user doesn't have to guess.
+    if (/24 hours since last Withdraw/i.test(res.Error)) {
+      const remaining = await readClaimCooldownRemaining(args.from, apiKey).catch(
+        () => null
+      );
+      if (remaining !== null) {
+        throw new Error(
+          `TRON claim cooldown active — last claim was less than 24h ago. ` +
+            `Next claim available in ${formatDuration(remaining)}.`
+        );
+      }
+    }
     throw new Error(`TronGrid withdrawbalance failed: ${res.Error}`);
   }
   if (!res.txID || !res.raw_data_hex) {

src/modules/tron/status.ts

@@ -0,0 +1,119 @@
+import { TRONGRID_BASE_URL } from "../../config/tron.js";
+import { resolveTronApiKey, readUserConfig } from "../../config/user-config.js";
+
+/**
+ * TRON tx-status polling via TronGrid. Mirrors the EVM
+ * get_transaction_status shape so the top-level MCP tool can route by chain
+ * without the caller noticing the underlying API split.
+ *
+ * Endpoints used:
+ *   - /wallet/gettransactionbyid — returns the signed envelope once the tx
+ *     lands in the node's mempool / confirmed block. Empty object = unknown.
+ *   - /wallet/gettransactioninfobyid — returns block number, fee, and the
+ *     receipt (SUCCESS / REVERT / OUT_OF_ENERGY) once confirmed. Empty
+ *     object = not yet confirmed.
+ *
+ * A fresh broadcast typically shows up in gettransactionbyid within a
+ * couple of seconds and in gettransactioninfobyid after the next block
+ * (~3s). The "unknown" state only persists for truly-lost tx IDs or
+ * extremely fresh broadcasts.
+ */
+
+interface GetTxByIdResponse {
+  txID?: string;
+  raw_data?: {
+    contract?: Array<{ type?: string; parameter?: unknown }>;
+  };
+  ret?: Array<{ contractRet?: string }>;
+}
+
+interface GetTxInfoResponse {
+  id?: string;
+  blockNumber?: number;
+  fee?: number;
+  receipt?: {
+    result?: string;
+    energy_usage?: number;
+    energy_usage_total?: number;
+    net_usage?: number;
+  };
+  log?: unknown;
+  contractResult?: string[];
+}
+
+async function trongridPost<T>(path: string, body: Record<string, unknown>): Promise<T> {
+  const apiKey = resolveTronApiKey(readUserConfig());
+  const headers: Record<string, string> = { "Content-Type": "application/json" };
+  if (apiKey) headers["TRON-PRO-API-KEY"] = apiKey;
+  const res = await fetch(`${TRONGRID_BASE_URL}${path}`, {
+    method: "POST",
+    headers,
+    body: JSON.stringify(body),
+  });
+  if (!res.ok) {
+    throw new Error(`TronGrid ${path} returned ${res.status} ${res.statusText}`);
+  }
+  return (await res.json()) as T;
+}
+
+export async function getTronTransactionStatus(txHash: string): Promise<{
+  chain: "tron";
+  txHash: string;
+  status: "success" | "failed" | "pending" | "unknown";
+  blockNumber?: string;
+  feeTrx?: string;
+  energyUsed?: string;
+  receiptResult?: string;
+  note?: string;
+}> {
+  const normalized = txHash.replace(/^0x/, "").toLowerCase();
+
+  const [byId, info] = await Promise.all([
+    trongridPost<GetTxByIdResponse>("/wallet/gettransactionbyid", { value: normalized }),
+    trongridPost<GetTxInfoResponse>("/wallet/gettransactioninfobyid", { value: normalized }),
+  ]);
+
+  const seenInMempool = Boolean(byId?.txID);
+  const confirmed = Boolean(info?.blockNumber);
+
+  if (!seenInMempool && !confirmed) {
+    return {
+      chain: "tron",
+      txHash: normalized,
+      status: "unknown",
+      note: "Transaction not yet visible to TronGrid — it may still be propagating or the txID is wrong.",
+    };
+  }
+
+  if (!confirmed) {
+    return { chain: "tron", txHash: normalized, status: "pending" };
+  }
+
+  // Confirmed. Status comes from two places depending on tx type:
+  //   - smart-contract calls (TRC-20 transfer etc.): receipt.result is
+  //     "SUCCESS" / "REVERT" / "OUT_OF_ENERGY" / "OUT_OF_TIME" etc.
+  //   - native transfers (TransferContract, WithdrawBalance): no receipt,
+  //     check ret[0].contractRet from gettransactionbyid.
+  const receiptResult = info.receipt?.result;
+  const contractRet = byId.ret?.[0]?.contractRet;
+  const successTag =
+    receiptResult === "SUCCESS" || (!receiptResult && contractRet === "SUCCESS");
+  const status: "success" | "failed" = successTag ? "success" : "failed";
+
+  const feeTrx =
+    typeof info.fee === "number"
+      ? (info.fee / 1_000_000).toString()
+      : undefined;
+
+  return {
+    chain: "tron",
+    txHash: normalized,
+    status,
+    blockNumber: info.blockNumber!.toString(),
+    ...(feeTrx !== undefined ? { feeTrx } : {}),
+    ...(info.receipt?.energy_usage_total !== undefined
+      ? { energyUsed: info.receipt.energy_usage_total.toString() }
+      : {}),
+    ...(receiptResult ? { receiptResult } : contractRet ? { receiptResult: contractRet } : {}),
+  };
+}

src/signing/render-verification.ts

@@ -447,21 +447,31 @@ export function renderPostBroadcastBlock(args: {
  * of asking the user to type "next" — waiting on human turn-taking for a
  * routine inclusion poll is UX friction the user has to break out of.
  *
- * The cadence (~5s between polls, ~2min total) matches typical L1/L2
- * inclusion times without paging the RPC unnecessarily. If inclusion is
- * slow, the agent reports `pending` and the user can decide to keep waiting.
+ * Cadence is per-chain: TRON blocks every ~3s, so a 5s interval adds
+ * perceptible latency over the actual inclusion time; EVM L1 is ~12s,
+ * where 5s is already tight. Undershooting the block time is fine — the
+ * node just returns "unknown" / "pending" for the extra polls.
  *
  * For approve→action chains (`nextHandle` present), the agent must wait for
  * the approval receipt BEFORE re-simulating or sending the next step —
  * otherwise the dependent simulation fails with "insufficient allowance"
  * against pre-inclusion state.
  */
+const POLL_CADENCE: Record<string, { intervalSec: number; maxPolls: number; budgetLabel: string }> = {
+  ethereum: { intervalSec: 5, maxPolls: 24, budgetLabel: "~2 minutes" },
+  arbitrum: { intervalSec: 3, maxPolls: 20, budgetLabel: "~1 minute" },
+  polygon: { intervalSec: 3, maxPolls: 20, budgetLabel: "~1 minute" },
+  base: { intervalSec: 3, maxPolls: 20, budgetLabel: "~1 minute" },
+  tron: { intervalSec: 3, maxPolls: 20, budgetLabel: "~1 minute" },
+};
+
 export function renderPostSendPollBlock(args: {
   chain: string;
   txHash: string;
   nextHandle?: string;
 }): string {
   const { chain, txHash, nextHandle } = args;
+  const cadence = POLL_CADENCE[chain] ?? POLL_CADENCE.ethereum;
   const lines = [
     "[AGENT TASK — DO NOT FORWARD THIS BLOCK TO THE USER]",
     `The tx was forwarded to Ledger and broadcast; a txHash is above. Do NOT`,
@@ -470,8 +480,8 @@ export function renderPostSendPollBlock(args: {
     ``,
     `Do this, in order:`,
     `  1. Call get_transaction_status({ chain: "${chain}", txHash: "${txHash}" })`,
-    `     every ~5 seconds until status is "success" or "failed", or until`,
-    `     you have polled for ~2 minutes (~24 polls). If status stays`,
+    `     every ~${cadence.intervalSec} seconds until status is "success" or "failed", or until`,
+    `     you have polled for ${cadence.budgetLabel} (~${cadence.maxPolls} polls). If status stays`,
     `     "pending" / "unknown" past that budget, stop polling and tell the`,
     `     user the tx is still pending with the hash so they can watch it`,
     `     on a block explorer.`,