Skip to content

Commit 3889f78

Browse files
FedericoAmuraFedericoAmura
committed
feat: generalize context usage to states variables
1 parent 06a6da5 commit 3889f78

File tree

8 files changed

+108
-87
lines changed

8 files changed

+108
-87
lines changed

‎packages/automation/src/lib/context/index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
export * from './machine-context';
1+
export * from './machine-context';

‎packages/automation/src/lib/context/machine-context.spec.ts‎

Lines changed: 31 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,9 @@ describe('MachineContext', () => {
99
token: '0x123...',
1010
},
1111
values: {
12-
amount: 100
12+
amount: 100,
1313
},
14-
existingArray: [1, 2, 3]
14+
existingArray: [1, 2, 3],
1515
};
1616

1717
beforeEach(() => {
@@ -53,13 +53,15 @@ describe('MachineContext', () => {
5353

5454
it('should create intermediate objects when setting deep paths', () => {
5555
context.set('a.b.c', 'value');
56-
expect(context.get()).toEqual(expect.objectContaining({
57-
a: {
58-
b: {
59-
c: 'value'
60-
}
61-
}
62-
}));
56+
expect(context.get()).toEqual(
57+
expect.objectContaining({
58+
a: {
59+
b: {
60+
c: 'value',
61+
},
62+
},
63+
})
64+
);
6365
});
6466

6567
it('should override existing values', () => {
@@ -115,19 +117,21 @@ describe('MachineContext', () => {
115117

116118
describe('array indexing', () => {
117119
beforeEach(() => {
118-
context = new MachineContext(deepCopy({
119-
simple: ['a', 'b', 'c'],
120-
complex: [
121-
{ id: 1, value: { foo: 'bar' } },
122-
{ id: 2, value: { foo: 'baz' } }
123-
],
124-
nested: {
125-
arrays: [
126-
[1, 2],
127-
[3, 4]
128-
]
129-
}
130-
}));
120+
context = new MachineContext(
121+
deepCopy({
122+
simple: ['a', 'b', 'c'],
123+
complex: [
124+
{ id: 1, value: { foo: 'bar' } },
125+
{ id: 2, value: { foo: 'baz' } },
126+
],
127+
nested: {
128+
arrays: [
129+
[1, 2],
130+
[3, 4],
131+
],
132+
},
133+
})
134+
);
131135
});
132136

133137
it('should access array elements using index notation', () => {
@@ -152,7 +156,11 @@ describe('MachineContext', () => {
152156

153157
it('should create arrays when setting with index notation', () => {
154158
context.set('new[2].foo', 'bar');
155-
expect(context.get('new')).toEqual([undefined, undefined, { foo: 'bar' }]);
159+
expect(context.get('new')).toEqual([
160+
undefined,
161+
undefined,
162+
{ foo: 'bar' },
163+
]);
156164
});
157165

158166
it('should handle array notation with dot notation mixed', () => {

‎packages/automation/src/lib/context/machine-context.ts‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,22 +47,25 @@ export class MachineContext {
4747
if (isLast) {
4848
current[index] = value;
4949
} else {
50-
current[index] = current[index] ?? (isNaN(parseInt(parts[i + 1])) ? {} : []);
50+
current[index] =
51+
current[index] ?? (isNaN(parseInt(parts[i + 1])) ? {} : []);
5152
current = current[index];
5253
}
5354
} else {
5455
if (isLast) {
5556
current[key] = value;
5657
} else {
57-
current[key] = current[key] ?? (isNaN(parseInt(parts[i + 1])) ? {} : []);
58+
current[key] =
59+
current[key] ?? (isNaN(parseInt(parts[i + 1])) ? {} : []);
5860
current = current[key];
5961
}
6062
}
6163
} else {
6264
if (isLast) {
6365
current[key] = value;
6466
} else {
65-
current = current[key] = current[key] ?? (isNaN(parseInt(parts[i + 1])) ? {} : []);
67+
current = current[key] =
68+
current[key] ?? (isNaN(parseInt(parts[i + 1])) ? {} : []);
6669
}
6770
}
6871
}
@@ -80,7 +83,11 @@ export class MachineContext {
8083
}
8184
}
8285

83-
public setFromData(location: string | string[], data?: Record<string, any>, path?: string | string[]) {
86+
public setFromData(
87+
location: string | string[],
88+
data?: Record<string, any>,
89+
path?: string | string[]
90+
) {
8491
if (!data) return;
8592

8693
const value = getFromObject(data, path);

‎packages/automation/src/lib/litActions.ts‎

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -15,24 +15,16 @@ export const signWithLitActionCode = `(async () => {
1515
Lit.Actions.setResponse({ response: signature });
1616
})();`;
1717

18-
interface ExecuteLitActionBase {
18+
interface ExecuteLitAction {
1919
litNodeClient: LitNodeClient;
2020
pkpPublicKey: string;
2121
authSigner: ethers.Wallet;
2222
ipfsId?: string;
23-
code: string;
24-
jsParams: Record<string, any>;
23+
code?: string;
24+
jsParams?: Record<string, any>;
2525
}
2626

27-
interface ExecuteCodeLitAction extends ExecuteLitActionBase {
28-
code: string;
29-
}
30-
31-
interface ExecuteIPFSLitAction extends ExecuteLitActionBase {
32-
ipfsId: string;
33-
}
34-
35-
type ExecuteLitAction = ExecuteCodeLitAction | ExecuteIPFSLitAction;
27+
const ONE_MINUTE = 1 * 60 * 1000;
3628

3729
export async function executeLitAction({
3830
litNodeClient,
@@ -49,7 +41,7 @@ export async function executeLitAction({
4941
await EthWalletProvider.authenticate({
5042
signer: authSigner,
5143
litNodeClient: litNodeClient,
52-
expiration: new Date(Date.now() + 1000 * 60 * 10).toISOString(), // 10 minutes
44+
expiration: new Date(Date.now() + ONE_MINUTE).toISOString(),
5345
}),
5446
],
5547
resourceAbilityRequests: [

‎packages/automation/src/lib/state-machine.spec.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -190,14 +190,14 @@ describe('StateMachine', () => {
190190
token: '0x123...',
191191
},
192192
values: {
193-
amount: 100
194-
}
193+
amount: 100,
194+
},
195195
};
196196

197197
beforeEach(() => {
198198
machine = new StateMachine({
199199
...stateMachineParams,
200-
context: initialContext
200+
context: initialContext,
201201
});
202202
});
203203

‎packages/automation/src/lib/state-machine.ts‎

Lines changed: 35 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,12 @@ import {
1414
import { signWithLitActionCode, executeLitAction } from './litActions';
1515
import { State, StateParams } from './states';
1616
import { Check, Transition } from './transitions';
17-
import { getChain } from './utils/chain';
17+
import { getEvmChain } from './utils/chain';
1818
import { getBalanceTransitionCheck, getERC20Balance } from './utils/erc20';
1919

2020
import {
2121
BaseStateMachineParams,
22+
ContextOrLiteral,
2223
PKPInfo,
2324
StateMachineDefinition,
2425
TransitionDefinition,
@@ -135,15 +136,15 @@ export class StateMachine {
135136
onEnterFunctions.push(async () => {
136137
console.log(
137138
`MachineContext at state ${state.key} enter: `,
138-
stateMachine.context.get(contextAction.log?.path),
139+
stateMachine.context.get(contextAction.log?.path)
139140
);
140141
});
141142
}
142143
if (contextAction.log?.atExit) {
143144
onExitFunctions.push(async () => {
144145
console.log(
145146
`MachineContext at state ${state.key} exit: `,
146-
stateMachine.context.get(contextAction.log?.path),
147+
stateMachine.context.get(contextAction.log?.path)
147148
);
148149
});
149150
}
@@ -162,8 +163,8 @@ export class StateMachine {
162163
litNodeClient: litNodeClientInstance,
163164
pkpPublicKey: stateMachine.pkp!.publicKey,
164165
authSigner: signer,
165-
ipfsId: litAction.ipfsId,
166-
code: litAction.code,
166+
ipfsId: stateMachine.resolveValue(litAction.ipfsId),
167+
code: stateMachine.resolveValue(litAction.code),
167168
jsParams: litAction.jsParams,
168169
});
169170

@@ -175,36 +176,38 @@ export class StateMachine {
175176
if (transaction) {
176177
onEnterFunctions.push(async () => {
177178
const yellowstoneMachineSigner = new ethers.Wallet(
178-
transaction.pkpOwnerKey,
179+
stateMachine.resolveValue(transaction.pkpOwnerKey),
179180
new ethers.providers.JsonRpcProvider(LIT_RPC.CHRONICLE_YELLOWSTONE)
180181
);
181182

182-
const chain = getChain(transaction);
183+
const chainId = stateMachine.resolveValue(transaction.evmChainId);
184+
const chain = getEvmChain(chainId);
183185
const chainProvider = new ethers.providers.JsonRpcProvider(
184186
chain.rpcUrls[0],
185187
chain.chainId
186188
);
187189

188190
const contract = new ethers.Contract(
189-
transaction.contractAddress,
191+
stateMachine.resolveValue(transaction.contractAddress),
190192
transaction.contractABI,
191193
chainProvider
192194
);
193195

194-
const txParams = (transaction.params || []).map(param =>
195-
'contextPath' in param
196-
? stateMachine.context.get(param.contextPath)
197-
: param
196+
const txParams = (transaction.params || []).map(
197+
stateMachine.resolveValue.bind(stateMachine)
198+
);
199+
const txMethod = stateMachine.resolveValue(transaction.method);
200+
const txData = await contract.populateTransaction[txMethod](
201+
...txParams
198202
);
199-
const txData = await contract.populateTransaction[transaction.method](...txParams);
200203
const gasLimit = await chainProvider.estimateGas({
201-
to: transaction.contractAddress,
204+
to: stateMachine.resolveValue(transaction.contractAddress),
202205
data: txData.data,
203-
from: transaction.pkpEthAddress,
206+
from: stateMachine.resolveValue(transaction.pkpEthAddress),
204207
});
205208
const gasPrice = await chainProvider.getGasPrice();
206209
const nonce = await chainProvider.getTransactionCount(
207-
transaction.pkpEthAddress
210+
stateMachine.resolveValue(transaction.pkpEthAddress)
208211
);
209212

210213
const rawTx = {
@@ -213,7 +216,7 @@ export class StateMachine {
213216
gasLimit: gasLimit.toHexString(),
214217
gasPrice: gasPrice.toHexString(),
215218
nonce,
216-
to: transaction.contractAddress,
219+
to: stateMachine.resolveValue(transaction.contractAddress),
217220
};
218221
const rawTxHash = ethers.utils.keccak256(
219222
ethers.utils.serializeTransaction(rawTx)
@@ -222,7 +225,7 @@ export class StateMachine {
222225
// Sign with the PKP in a LitAction
223226
const litActionResponse = await executeLitAction({
224227
litNodeClient: litNodeClientInstance,
225-
pkpPublicKey: transaction.pkpPublicKey,
228+
pkpPublicKey: stateMachine.resolveValue(transaction.pkpPublicKey),
226229
authSigner: yellowstoneMachineSigner,
227230
code: signWithLitActionCode,
228231
jsParams: {
@@ -282,7 +285,8 @@ export class StateMachine {
282285

283286
if (evmContractEvent) {
284287
const transitionIndex = checks.length;
285-
const chain = getChain(evmContractEvent);
288+
const chainId = stateMachine.resolveValue(evmContractEvent.evmChainId);
289+
const chain = getEvmChain(chainId);
286290

287291
listeners.push(
288292
new EVMContractEventListener(
@@ -302,12 +306,12 @@ export class StateMachine {
302306
| ContractEventData
303307
| undefined;
304308

305-
evmContractEvent.contextUpdates?.forEach(contextUpdate =>
309+
evmContractEvent.contextUpdates?.forEach((contextUpdate) =>
306310
stateMachine.context.setFromData(
307311
contextUpdate.contextPath,
308312
eventData,
309-
contextUpdate.dataPath,
310-
),
313+
contextUpdate.dataPath
314+
)
311315
);
312316

313317
return eventData?.event.event === evmContractEvent.eventName;
@@ -317,7 +321,8 @@ export class StateMachine {
317321
if (balances) {
318322
balances.forEach((balance) => {
319323
const transitionIndex = checks.length;
320-
const chain = getChain(balance);
324+
const chainId = stateMachine.resolveValue(balance.evmChainId);
325+
const chain = getEvmChain(chainId);
321326

322327
const chainProvider = new ethers.providers.JsonRpcProvider(
323328
chain.rpcUrls[0],
@@ -471,6 +476,13 @@ export class StateMachine {
471476
this.debug && console.log('State machine stopped');
472477
}
473478

479+
public resolveValue<T = unknown>(value: ContextOrLiteral<T> | T): T {
480+
if (value && typeof value === 'object' && 'contextPath' in value) {
481+
return this.context.get(value.contextPath) as T;
482+
}
483+
return value;
484+
}
485+
474486
/**
475487
* Validates whether a PKP (Private Key Pair) is configured in the state machine.
476488
* If a PKP is not present, it initiates the minting of a new PKP through the

0 commit comments

Comments
 (0)