aegis-defi

Instructions

Implementation Reference

• src/mcp-server/server.ts:160-278 (handler)

  The handler function for the `assess_risk` tool. It orchestrates contract scanning, transaction simulation, and token safety checks to generate a final risk assessment and attestation.

  server.tool(
    "assess_risk",
    "Comprehensive risk assessment combining contract scanning, transaction simulation, and token checks. This is the recommended all-in-one safety check before any DeFi interaction. Returns a go/no-go recommendation.",
    {
      action: z.enum(["swap", "approve", "transfer", "interact"]).describe("Type of action being assessed"),
      targetContract: z.string().describe("The contract being interacted with"),
      chainId: z.number().default(1).describe("Chain ID"),
      from: z.string().describe("The agent's wallet address"),
      transactionData: z.string().optional().describe("Calldata for the transaction (hex)"),
      value: z.string().default("0").describe("ETH value (in wei)"),
      tokenAddress: z.string().optional().describe("Token address if this involves a token swap"),
    },
    async ({ action, targetContract, chainId, from, transactionData, value, tokenAddress }) => {
      const checks: Record<string, any> = {};
  
      // 1. Scan the contract
      const fetched = await fetchContractSource(targetContract, chainId);
      if (fetched.source) {
        checks.contractScan = scanContractSource(fetched.source);
      } else if (fetched.bytecode) {
        checks.contractScan = scanBytecode(fetched.bytecode);
      }
  
      // 2. Simulate the transaction if we have calldata
      if (transactionData) {
        checks.simulation = await simulateTransaction({
          chainId,
          from: from as Address,
          to: targetContract as Address,
          data: transactionData as Hex,
          value: BigInt(value),
        });
        // Convert bigint for JSON serialization
        if (checks.simulation) {
          checks.simulation.gasUsed = checks.simulation.gasUsed.toString();
        }
      }
  
      // 3. Check token safety if relevant
      if (tokenAddress) {
        checks.tokenSafety = await checkTokenSellability(
          chainId,
          tokenAddress as Address,
          from as Address,
        );
      }
  
      // 4. Compute overall risk
      let overallRisk = 0;
      const risks: string[] = [];
  
      if (checks.contractScan) {
        overallRisk = Math.max(overallRisk, checks.contractScan.riskScore);
        if (checks.contractScan.riskScore >= 70) risks.push("contract_high_risk");
      }
  
      if (checks.simulation && !checks.simulation.success) {
        overallRisk = Math.max(overallRisk, 80);
        risks.push("transaction_reverts");
      }
  
      if (checks.simulation?.gasAnomaly) {
        overallRisk = Math.max(overallRisk, 60);
        risks.push("gas_anomaly");
      }
  
      if (checks.tokenSafety && !checks.tokenSafety.canSell) {
        overallRisk = Math.max(overallRisk, 90);
        risks.push("cannot_sell_token");
      }
  
      const decision = overallRisk >= 70 ? "BLOCK" : overallRisk >= 40 ? "WARN" : "ALLOW";
  
      // If not blocked, sign an attestation for on-chain verification
      let attestation: Record<string, string> | undefined;
      if (decision !== "BLOCK") {
        try {
          const selector = transactionData ? transactionData.slice(0, 10) as Hex : "0x00000000" as Hex;
          const att = await signAttestation({
            agent: from as Address,
            target: targetContract as Address,
            selector,
            riskScore: overallRisk,
          });
          attestation = {
            attestationId: att.attestationId,
            agent: att.agent,
            target: att.target,
            selector: att.selector,
            riskScore: att.riskScore.toString(),
            expiresAt: att.expiresAt.toString(),
            signature: att.signature,
          };
        } catch {
          // Attester key not configured - attestation unavailable
          // MCP-only mode still works (agent gets the risk assessment)
        }
      }
  
      return {
        content: [{
          type: "text" as const,
          text: JSON.stringify({
            decision,
            overallRiskScore: overallRisk,
            riskFactors: risks,
            action,
            checks,
            attestation,
            recommendation: decision === "BLOCK"
              ? "DO NOT proceed with this transaction. High risk of fund loss."
              : decision === "WARN"
              ? "Proceed with caution. Some risk indicators detected."
              : "Transaction appears safe. Proceed normally.",
          }, null, 2),
        }],
      };
    },
  );