fix: shield --from flag supports stealth addys

CHANGELOG.md

@@ -20,6 +20,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 - fix `balances --verbose` to no longer sync privacy protocols twice unnecessarily.
 - `shield` and `unshield` keep a live progress timer during protocol sync (including Railgun WASM), matching `balances`.
+- the `--from` flag on `shield` correctly supports `<stealth address>` or `sN` (stealth address by index)
 
 ## [0.0.4] — 2026-08-25
 
@@ -37,3 +38,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 ### Fixed
 
 - Unshielding to a custom / ephemeral recipient (`--to` an address that is not a stored HD account) no longer fails or mis-routes funds.
+

README.md

@@ -418,7 +418,7 @@ Move funds from a **public** account into a private protocol.
 | `--protocol <railgun\|privacy-pools\|tornado>` | Required unless `DEFAULT_PRIVACY_PROTOCOL` is set to one of those values. |
 | `--wallet <name>` | Wallet. |
 | `--password <password>` | Unlock password. |
-| `--from <address-or-index>` | Sender public account (address or HD index). |
+| `--from <address-or-index>` | Sender public account address, HD index, or stealth selector (`s0`). |
 | `--from-priv` | With `--broadcast`: derive private key by index from mnemonic if account not yet in stored public list. |
 | `--token <address\|eth>` | Token (default: `eth`). |
 | `--amount-wei <n>` | Amount in base units. |
@@ -446,6 +446,7 @@ When a shield needs more than one on-chain call, the CLI uses EIP-7702 Simple770
 
 ```bash
 kohaku shield --protocol tornado --wallet testWallet --from 0 --amount-formatted 0.1 --broadcast
+kohaku shield --protocol tornado --wallet testWallet --from s0 --amount-formatted 0.1 --broadcast
 kohaku shield --protocol tornado --wallet testWallet --from 0 --amount-max --broadcast
 kohaku shield --protocol railgun --wallet testWallet --from 0 --token 0x1c7D4B196Cb0C7B01d743Fbc6116a902379C7238 --amount-formatted 10 --broadcast
 kohaku shield --protocol tornado --wallet testWallet --from 0 --amount-formatted 0.1 --without-tor

src/commands/shield.ts

@@ -4,16 +4,17 @@ import chalk from "chalk";
 import type { AssetAmount } from "@kohaku-eth/plugins";
 import type { Command } from "commander";
 import { formatUnits, getAddress, isAddress, parseUnits } from "viem";
-import { Mnemonic } from "derive-railgun-keys";
 
 import { makeHost } from "../host/makeHost";
 import {
   buildShieldCallList,
   formatAccountSelector,
   formatPublicAccountBalanceLabel,
   listPublicAccountsWithBalance,
+  parseFromIndex,
   partitionShieldTxs,
   resolveShieldApprovalCalls,
+  resolveShieldSender,
   shieldTransactionConfirmMessage,
   summarizeMultiShieldPlan,
   toShieldTxs,
@@ -42,7 +43,6 @@ import {
 } from "../utils/fee-preview.js";
 import { resolveAddressOrName } from "../utils/resolve-name.js";
 import {
-  addressFromPrivateKey,
   makeWalletClient,
   sendTransactionAndWait,
   simulateCallOrThrow,
@@ -71,7 +71,6 @@ import {
   resolveWalletPassword,
 } from "../utils/wallets-util";
 import { readSeedKeystore } from "../utils/mnemonic";
-import { makePublicAccountsStorage } from "../utils/public-accounts";
 import {
   assertPpErc20TokenWhitelisted,
   createProtocolPlugin,
@@ -135,13 +134,6 @@ function etherscanTxUrl(chainId: bigint, txHash: string): string {
   return `https://${host}/tx/${txHash}`;
 }
 
-function parseFromIndex(fromValue: string): number | null {
-  if (!/^\d+$/.test(fromValue)) return null;
-  const parsed = Number(fromValue);
-  if (!Number.isInteger(parsed) || parsed < 0) return null;
-  return parsed;
-}
-
 function findAccountWithBalance(
   fromValue: string,
   accounts: PublicAccountWithBalance[]
@@ -290,7 +282,10 @@ export function registerShieldCommand(program: Command): void {
     )
     .option("--wallet <name>", cliOptions.walletPickList)
     .option("--password <password>", cliOptions.password)
-    .option("--from <address-or-index>", "Public sender address or public-account index")
+    .option(
+      "--from <address-or-index>",
+      "Public sender address, HD index, or stealth selector (s0)"
+    )
     .option(
       "--from-priv",
       "With --broadcast: derive --from index from mnemonic when missing from public accounts (not required for dry-run)"
@@ -455,8 +450,6 @@ export function registerShieldCommand(program: Command): void {
         return;
       }
 
-      const publicStorage = makePublicAccountsStorage(walletDir, mnemonic, password);
-
       const withBalances = await listPublicAccountsWithBalance(
         rpcUrl,
         walletDir,
@@ -676,40 +669,21 @@ export function registerShieldCommand(program: Command): void {
         }
       }
 
-      const fromIndex = parseFromIndex(fromValue);
-      let senderPrivateKey: string | undefined;
       let senderAddress: string;
-      if (fromIndex !== null) {
-        const account = publicStorage.getAccount(fromIndex);
-        if (account) {
-          senderPrivateKey = account.priv;
-          senderAddress = account.address;
-        } else if (opts.fromPriv || dryRun) {
-          senderPrivateKey = Mnemonic.to0xPrivateKeyByIndex(mnemonic, fromIndex);
-          senderAddress = addressFromPrivateKey(senderPrivateKey);
-        } else {
-          cliError(
-            `Public account index ${fromIndex} not found. Use --from-priv with --broadcast to derive from mnemonic, or omit --broadcast for a dry-run.`
-          );
-          return;
-        }
-      } else if (isAddress(fromValue)) {
-        senderAddress = getAddress(fromValue);
-        const match = publicStorage
-          .getAccounts()
-          .find((x) => x.address.toLowerCase() === senderAddress.toLowerCase());
-        if (match) {
-          senderPrivateKey = match.priv;
-        } else if (dryRun) {
-          senderPrivateKey = undefined;
-        } else {
-          cliError(
-            `Address ${senderAddress} is not in this wallet's public accounts. Use --broadcast with --from-priv and an index, or omit --broadcast to preview txs for this address.`
-          );
-          return;
-        }
-      } else {
-        cliError("--from must be either a valid address or a non-negative index.");
+      let senderPrivateKey: string | undefined;
+      try {
+        const resolved = resolveShieldSender({
+          fromValue,
+          walletDir,
+          mnemonic,
+          password,
+          dryRun,
+          allowDeriveFromMnemonic: !!opts.fromPriv,
+        });
+        senderAddress = resolved.senderAddress;
+        senderPrivateKey = resolved.senderPrivateKey;
+      } catch (e) {
+        cliErrorFromCaught(e);
         return;
       }
 
@@ -1005,7 +979,7 @@ export function registerShieldCommand(program: Command): void {
 
         if (!senderPrivateKey) {
           cliError(
-            "Cannot sign: no private key for this --from (use a saved public account or --from-priv with --broadcast)."
+            "Cannot sign: no private key for this --from (use a saved public/stealth account or --from-priv with --broadcast)."
           );
           return;
         }

tests/shield-txs.test.ts

@@ -1,15 +1,22 @@
 import assert from "node:assert/strict";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
 import { describe, it } from "node:test";
-import { encodeFunctionData, getAddress } from "viem";
+import { encodeFunctionData, getAddress, type Hex } from "viem";
+import { privateKeyToAccount } from "viem/accounts";
 
 import {
   parseFromIndex,
   partitionShieldTxs,
+  resolveShieldSender,
   toShieldTxs,
   tryDecodeErc20Approve,
   type ShieldCall,
 } from "../src/lib/shield-flow.js";
+import { makeStealthAccountsStorage } from "../src/lib/stealth/storage.js";
 import { ERC20_ABI } from "../src/utils/tokens-util.js";
+import { addressFromPrivateKey } from "../src/utils/viem-tx.js";
 
 const TOKEN = "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48";
 const POOL_A = getAddress("0x12d66f87a04a9e220743712ce6d9bb1b5616b8fc");
@@ -32,6 +39,21 @@ function depositCall(to: string, value = 0n): ShieldCall {
   return { to, data: "0xdead", value };
 }
 
+const MNEMONIC =
+  "test test test test test test test test test test test junk";
+const STEALTH_PRIV =
+  "0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" as Hex;
+const STEALTH_ADDRESS = privateKeyToAccount(STEALTH_PRIV).address;
+
+function withWalletDir(fn: (walletDir: string) => void): void {
+  const dir = mkdtempSync(join(tmpdir(), "kohaku-shield-from-"));
+  try {
+    fn(dir);
+  } finally {
+    rmSync(dir, { recursive: true, force: true });
+  }
+}
+
 describe("parseFromIndex", () => {
   it("parses a non-negative decimal HD index", () => {
     assert.equal(parseFromIndex("0"), 0);
@@ -50,6 +72,84 @@ describe("parseFromIndex", () => {
   });
 });
 
+describe("resolveShieldSender stealth --from", () => {
+  const senderOpts = (walletDir: string, fromValue: string, dryRun = false) => ({
+    fromValue,
+    walletDir,
+    mnemonic: MNEMONIC,
+    password: "pw",
+    dryRun,
+    allowDeriveFromMnemonic: false,
+  });
+
+  it("resolves --from s0 and stealth:0 to the stored stealth key", () => {
+    withWalletDir((walletDir) => {
+      makeStealthAccountsStorage(walletDir, "pw").upsertAccount({
+        address: STEALTH_ADDRESS,
+        priv: STEALTH_PRIV,
+        ephemeralPublicKey:
+          "0x02aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+        schemeId: 1,
+        lastUpdated: 1,
+        ethBalance: "0",
+        erc20Balances: {},
+      });
+
+      for (const fromValue of ["s0", "S0", "stealth:0"]) {
+        const resolved = resolveShieldSender(senderOpts(walletDir, fromValue));
+        assert.equal(resolved.senderAddress, getAddress(STEALTH_ADDRESS));
+        assert.equal(resolved.senderPrivateKey, STEALTH_PRIV);
+      }
+    });
+  });
+
+  it("resolves --from <stealth address> to the stored stealth key", () => {
+    withWalletDir((walletDir) => {
+      makeStealthAccountsStorage(walletDir, "pw").upsertAccount({
+        address: STEALTH_ADDRESS,
+        priv: STEALTH_PRIV,
+        ephemeralPublicKey:
+          "0x02aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
+        schemeId: 1,
+        lastUpdated: 1,
+        ethBalance: "0",
+        erc20Balances: {},
+      });
+
+      for (const fromValue of [STEALTH_ADDRESS, STEALTH_ADDRESS.toLowerCase()]) {
+        const resolved = resolveShieldSender(
+          senderOpts(walletDir, fromValue, false)
+        );
+        assert.equal(resolved.senderAddress, getAddress(STEALTH_ADDRESS));
+        assert.equal(resolved.senderPrivateKey, STEALTH_PRIV);
+      }
+    });
+  });
+
+  it("rejects a missing stealth selector", () => {
+    withWalletDir((walletDir) => {
+      assert.throws(
+        () => resolveShieldSender(senderOpts(walletDir, "s0")),
+        /Stealth account s0 not found/
+      );
+    });
+  });
+
+  it("still derives an HD index on dry-run when the public account is missing", () => {
+    withWalletDir((walletDir) => {
+      const resolved = resolveShieldSender(senderOpts(walletDir, "0", true));
+      assert.equal(
+        resolved.senderAddress,
+        addressFromPrivateKey(
+          // index 0 of the well-known test mnemonic
+          "0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"
+        )
+      );
+      assert.ok(resolved.senderPrivateKey);
+    });
+  });
+});
+
 describe("tryDecodeErc20Approve", () => {
   it("decodes spender and amount from approve calldata", () => {
     const data = encodeFunctionData({