From 1d0bdf9d7c5e006d493ce9786ee1306775f523a6 Mon Sep 17 00:00:00 2001
From: Nathan Klick <nathan@swirldslabs.com>
Date: Mon, 17 Feb 2025 11:13:05 -0600
Subject: [PATCH] chore: ensure staking updates support multi-cluster
 deployments

Signed-off-by: Nathan Klick <nathan@swirldslabs.com>
---
 src/commands/node/tasks.ts | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/commands/node/tasks.ts b/src/commands/node/tasks.ts
index 4a4d5d97b..e408a73d2 100644
--- a/src/commands/node/tasks.ts
+++ b/src/commands/node/tasks.ts
@@ -565,6 +565,7 @@ export class NodeCommandTasks {
     accountId: string,
     nodeAlias: NodeAlias,
     stakeAmount: number = HEDERA_NODE_DEFAULT_STAKE_AMOUNT,
+    context?: string,
   ) {
     try {
       const deploymentName = this.configManager.getFlag<DeploymentName>(flags.deployment);
@@ -573,6 +574,7 @@ export class NodeCommandTasks {
         this.parent.getClusterRefs(),
         deploymentName,
         this.configManager.getFlag<boolean>(flags.forcePortForward),
+        context,
       );
       const client = this.accountManager._nodeClient;
       const treasuryKey = await this.accountManager.getTreasuryAccountKeys(namespace);
@@ -1274,11 +1276,12 @@ export class NodeCommandTasks {
         let nodeIndex = 0;
         for (const nodeAlias of ctx.config.nodeAliases) {
           const accountId = accountMap.get(nodeAlias);
+          const context = helpers.extractContextFromConsensusNodes(nodeAlias, ctx.config.consensusNodes);
           const stakeAmount =
             stakeAmountParsed.length > 0 ? stakeAmountParsed[nodeIndex] : HEDERA_NODE_DEFAULT_STAKE_AMOUNT;
           subTasks.push({
             title: `Adding stake for node: ${chalk.yellow(nodeAlias)}`,
-            task: async () => await self._addStake(ctx.config.namespace, accountId, nodeAlias, +stakeAmount),
+            task: async () => await self._addStake(ctx.config.namespace, accountId, nodeAlias, +stakeAmount, context),
           });
           nodeIndex++;
         }
@@ -1297,13 +1300,16 @@ export class NodeCommandTasks {
   stakeNewNode() {
     const self = this;
     return new Task('Stake new node', async (ctx: any, task: ListrTaskWrapper<any, any, any>) => {
+      const context = helpers.extractContextFromConsensusNodes(ctx.config.nodeAlias, ctx.config.consensusNodes);
       await self.accountManager.refreshNodeClient(
         ctx.config.namespace,
         ctx.config.nodeAlias,
         this.parent.getClusterRefs(),
         this.configManager.getFlag<DeploymentName>(flags.deployment),
+        context,
+        this.configManager.getFlag<boolean>(flags.forcePortForward),
       );
-      await this._addStake(ctx.config.namespace, ctx.newNode.accountId, ctx.config.nodeAlias);
+      await this._addStake(ctx.config.namespace, ctx.newNode.accountId, ctx.config.nodeAlias, undefined, context);
     });
   }