From be62dbc9743037fba8d7da10fc9bc63cdf63b7ee Mon Sep 17 00:00:00 2001
From: FedericoAmura <federicoamura@gmail.com>
Date: Tue, 3 Dec 2024 21:26:10 +0100
Subject: [PATCH] feat: add lit action run state machine capability as a state

---
 packages/automation/src/lib/state-machine.ts | 86 +++++++++++++++++++-
 packages/automation/src/lib/types.ts         |  9 ++
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/packages/automation/src/lib/state-machine.ts b/packages/automation/src/lib/state-machine.ts
index 268603967..3dbe12a0b 100644
--- a/packages/automation/src/lib/state-machine.ts
+++ b/packages/automation/src/lib/state-machine.ts
@@ -1,6 +1,19 @@
 import { ethers } from 'ethers';
 
+import {
+  // createSiweMessageWithRecaps,
+  // generateAuthSig,
+  LitActionResource,
+  // LitPKPResource,
+} from '@lit-protocol/auth-helpers';
+import {
+  LIT_ABILITY,
+  LIT_EVM_CHAINS,
+  LIT_RPC,
+  LIT_NETWORK,
+} from '@lit-protocol/constants';
 import { LitContracts } from '@lit-protocol/contracts-sdk';
+import { EthWalletProvider } from '@lit-protocol/lit-auth-client';
 import { LitNodeClient } from '@lit-protocol/lit-node-client';
 
 import {
@@ -34,6 +47,15 @@ import type {
 
 export type MachineStatus = 'running' | 'stopped';
 
+const ethPrivateKey = process.env['ETHEREUM_PRIVATE_KEY'];
+if (!ethPrivateKey) {
+  throw new Error('ethPrivateKey not defined');
+}
+const yellowstoneSigner = new ethers.Wallet(
+  ethPrivateKey,
+  new ethers.providers.JsonRpcProvider(LIT_RPC.CHRONICLE_YELLOWSTONE)
+);
+
 /**
  * A StateMachine class that manages states and transitions between them.
  */
@@ -66,7 +88,12 @@ export class StateMachine {
         ? litNodeClient
         : new LitNodeClient(litNodeClient);
     const litContractsInstance =
-      'connect' in litContracts ? litContracts : new LitContracts(litContracts);
+      'connect' in litContracts
+        ? litContracts
+        : new LitContracts({
+            signer: yellowstoneSigner,
+            ...litContracts,
+          });
 
     if (
       litNodeClientInstance.config.litNetwork !== litContractsInstance.network
@@ -83,7 +110,62 @@ export class StateMachine {
     });
 
     machineConfig.states.forEach((state) => {
-      stateMachine.addState(state);
+      const { litAction } = state;
+
+      const stateConfig: StateParams = {
+        key: state.key,
+      };
+
+      if (litAction) {
+        let pkpPublicKey: string = litAction.pkpPublicKey;
+
+        stateConfig.onEnter = async () => {
+          const yellowstoneSigner = new ethers.Wallet(
+            litAction.pkpOwnerKey,
+            new ethers.providers.JsonRpcProvider(LIT_RPC.CHRONICLE_YELLOWSTONE)
+          );
+
+          if (!pkpPublicKey) {
+            console.log(`No PKP for LitAction, minting one...`);
+            const mintingReceipt =
+              await litContractsInstance.pkpNftContractUtils.write.mint();
+            const pkp = mintingReceipt.pkp;
+            pkpPublicKey = pkp.publicKey;
+            console.log(`Minted PKP: ${pkp}`);
+          }
+
+          const pkpSessionSigs = await litNodeClientInstance.getPkpSessionSigs({
+            pkpPublicKey,
+            capabilityAuthSigs: [],
+            authMethods: [
+              await EthWalletProvider.authenticate({
+                signer: yellowstoneSigner,
+                litNodeClient: litNodeClientInstance,
+                expiration: new Date(Date.now() + 1000 * 60 * 10).toISOString(), // 10 minutes
+              }),
+            ],
+            resourceAbilityRequests: [
+              {
+                resource: new LitActionResource('*'),
+                ability: LIT_ABILITY.LitActionExecution,
+              },
+            ],
+          });
+
+          // Run a LitAction
+          const executeJsResponse = await litNodeClientInstance.executeJs({
+            sessionSigs: pkpSessionSigs,
+            ipfsId: litAction.ipfsId,
+            code: litAction.code,
+            jsParams: litAction.jsParams,
+          });
+
+          // TODO send user this result with a webhook maybe
+          console.log(`============ executeJsResponse:`, executeJsResponse);
+        };
+      }
+
+      stateMachine.addState(stateConfig);
     });
 
     machineConfig.transitions.forEach((transition) => {
diff --git a/packages/automation/src/lib/types.ts b/packages/automation/src/lib/types.ts
index 64dcca263..59725e7eb 100644
--- a/packages/automation/src/lib/types.ts
+++ b/packages/automation/src/lib/types.ts
@@ -7,8 +7,17 @@ import { BaseTransitionParams } from './transitions';
 
 export type Address = `0x${string}`;
 
+export interface LitActionStateDefinition {
+  pkpOwnerKey: string;
+  pkpPublicKey: string;
+  ipfsId?: string; // TODO separate into another without code
+  code: string;
+  jsParams: Record<string, any>;
+}
+
 export interface StateDefinition {
   key: string;
+  litAction?: LitActionStateDefinition;
 }
 
 export interface OnEvmChainEvent {