feat: return and monitor TXs from the same method

Signed-off-by: Simon Gellis <simongellis@gmail.com>

Author: SupernaviX

firefly-balius/src/monitor.rs

@@ -1,6 +1,6 @@
 use std::collections::HashMap;
 
-use balius_sdk::WorkerResult;
+use balius_sdk::{wit, Json, WorkerResult};
 use serde::{Deserialize, Serialize};
 
 use crate::kv;
@@ -27,3 +27,34 @@ pub enum FinalizationCondition {
     // Treat the transaction as finalized after the given number of blocks have reached the chain.
     AfterBlocks(u64),
 }
+
+/// A new transaction which FireFly will build, submit, and monitor for you.
+pub struct NewMonitoredTx(
+    pub Box<dyn balius_sdk::txbuilder::TxExpr>,
+    pub FinalizationCondition,
+);
+
+impl TryInto<wit::Response> for NewMonitoredTx {
+    type Error = balius_sdk::Error;
+
+    fn try_into(self) -> Result<wit::Response, Self::Error> {
+        let balius_sdk::wit::Response::PartialTx(tx) = balius_sdk::NewTx(self.0).try_into()? else {
+            return Err(balius_sdk::Error::Internal("Unexpected response".into()));
+        };
+
+        let new_tx = SerializedNewTx::FireFlyCardanoNewTx {
+            bytes: hex::encode(tx),
+            condition: self.1,
+        };
+        Json(new_tx).try_into()
+    }
+}
+
+#[derive(Serialize)]
+#[serde(tag = "type")]
+enum SerializedNewTx {
+    FireFlyCardanoNewTx {
+        bytes: String,
+        condition: FinalizationCondition,
+    },
+}

firefly-cardanoconnect/src/contracts.rs

@@ -4,12 +4,12 @@ use std::{
 };
 
 use anyhow::{bail, Result};
-use balius_runtime::{ledgers::Ledger, Response};
+use balius_runtime::ledgers::Ledger;
 use dashmap::{DashMap, Entry};
-pub use runtime::ContractEvent;
-use runtime::ContractRuntime;
+pub use runtime::{ContractEvent, NewTx};
+use runtime::{ContractRuntime, InvokeResponse};
 use serde::Deserialize;
-use serde_json::{json, Value};
+use serde_json::Value;
 use tokio::fs;
 use tracing::{error, warn};
 
@@ -81,13 +81,13 @@ impl ContractManager {
         contract: &str,
         method: &str,
         params: Value,
-    ) -> Result<Option<Vec<u8>>> {
+    ) -> Result<Option<NewTx>> {
         let Some(runtime) = self.runtimes.get(contract) else {
             bail!("unrecognized contract {contract}");
         };
         let response = runtime.invoke(method, params).await?;
         match response {
-            Response::PartialTx(bytes) => Ok(Some(bytes)),
+            InvokeResponse::NewTx(tx) => Ok(Some(tx)),
             _ => Ok(None),
         }
     }
@@ -98,20 +98,14 @@ impl ContractManager {
         };
         let response = runtime.invoke(method, params).await?;
         match response {
-            Response::Acknowledge => Ok(json!({})),
-            Response::Cbor(bytes) => Ok(json!({ "cbor": hex::encode(bytes) })),
-            Response::Json(bytes) => Ok(serde_json::from_slice(&bytes)?),
-            Response::PartialTx(_) => bail!("Cannot build transactions from query"),
+            InvokeResponse::Json(value) => Ok(value),
+            InvokeResponse::NewTx(_) => bail!("Cannot build transactions from query"),
         }
     }
 
-    pub async fn handle_submit(&self, contract: &str, method: &str, tx_id: &str) {
-        let params = serde_json::json!({
-            "method": method,
-            "hash": tx_id,
-        });
+    pub async fn handle_submit(&self, contract: &str, tx_id: &str, tx: NewTx) -> Result<()> {
         let runtime = self.get_contract_runtime(contract).await;
-        let _: Result<_, _> = runtime.invoke("__tx_submitted", params.clone()).await;
+        runtime.handle_submit(tx_id, tx).await
     }
 
     pub async fn listen(&self, listener: &Listener) -> ContractListener {

firefly-cardanoconnect/src/contracts/runtime.rs

@@ -76,7 +76,7 @@ impl ContractRuntime {
         rx.await?
     }
 
-    pub async fn invoke(&self, method: &str, params: Value) -> Result<Response> {
+    pub async fn invoke(&self, method: &str, params: Value) -> Result<InvokeResponse> {
         let (tx, rx) = oneshot::channel();
         if self
             .tx
@@ -92,6 +92,22 @@ impl ContractRuntime {
         rx.await?
     }
 
+    pub async fn handle_submit(&self, tx_id: &str, new_tx: NewTx) -> Result<()> {
+        let (tx, rx) = oneshot::channel();
+        if self
+            .tx
+            .send(ContractRuntimeCommand::HandleSubmit {
+                tx_id: tx_id.to_string(),
+                new_tx,
+                done: tx,
+            })
+            .is_err()
+        {
+            bail!("runtime for contract {} has stopped", self.contract);
+        }
+        rx.await?
+    }
+
     pub async fn apply(&self, rollbacks: &[BlockInfo], block: &BlockInfo) -> Result<()> {
         let (tx, rx) = oneshot::channel();
         if self
@@ -139,7 +155,12 @@ enum ContractRuntimeCommand {
     Invoke {
         method: String,
         params: Value,
-        done: oneshot::Sender<Result<Response>>,
+        done: oneshot::Sender<Result<InvokeResponse>>,
+    },
+    HandleSubmit {
+        tx_id: String,
+        new_tx: NewTx,
+        done: oneshot::Sender<Result<()>>,
     },
     Apply {
         rollbacks: Vec<BlockInfo>,
@@ -194,6 +215,13 @@ impl ContractRuntimeWorker {
                 } => {
                     let _ = done.send(self.invoke(&method, params).await);
                 }
+                ContractRuntimeCommand::HandleSubmit {
+                    tx_id,
+                    new_tx,
+                    done,
+                } => {
+                    let _ = done.send(self.handle_submit(&tx_id, new_tx).await);
+                }
                 ContractRuntimeCommand::Apply {
                     rollbacks,
                     block,
@@ -241,14 +269,53 @@ impl ContractRuntimeWorker {
         Ok(())
     }
 
-    async fn invoke(&mut self, method: &str, params: Value) -> Result<Response> {
+    async fn invoke(&mut self, method: &str, params: Value) -> Result<InvokeResponse> {
         let params = serde_json::to_vec(&params)?;
         let Some(runtime) = self.runtime.as_mut() else {
             bail!("Contract {} failed to initialize", self.contract);
         };
-        Ok(runtime
+        let response = runtime
             .handle_request(&self.contract, method, params)
-            .await?)
+            .await?;
+        match response {
+            Response::Acknowledge => Ok(InvokeResponse::Json(json!({}))),
+            Response::Cbor(bytes) => Ok(InvokeResponse::Json(json!({ "cbor": bytes }))),
+            Response::PartialTx(bytes) => Ok(InvokeResponse::NewTx(NewTx {
+                bytes,
+                method: method.to_string(),
+                condition: None,
+            })),
+            Response::Json(bytes) => {
+                let value: Value = serde_json::from_slice(&bytes)?;
+                if let Ok(RawNewTx::FireFlyCardanoNewTx { bytes, condition }) =
+                    serde_json::from_value(value.clone())
+                {
+                    Ok(InvokeResponse::NewTx(NewTx {
+                        bytes: hex::decode(bytes)?,
+                        method: method.to_string(),
+                        condition: Some(condition),
+                    }))
+                } else {
+                    Ok(InvokeResponse::Json(value))
+                }
+            }
+        }
+    }
+
+    async fn handle_submit(&mut self, tx_id: &str, new_tx: NewTx) -> Result<()> {
+        if let Some(condition) = new_tx.condition {
+            let mut monitored_txs: HashMap<String, FinalizationCondition> =
+                self.get_value("__monitored_txs").await?.unwrap_or_default();
+            monitored_txs.insert(tx_id.to_string(), condition);
+            self.set_value("__monitored_txs", monitored_txs).await?;
+        }
+
+        let params = json!({
+            "method": new_tx.method,
+            "hash": tx_id,
+        });
+        let _: Result<_, _> = self.invoke("__tx_submitted", params).await;
+        Ok(())
     }
 
     async fn apply(&mut self, rollbacks: &[BlockInfo], block: &BlockInfo) -> Result<()> {
@@ -419,14 +486,34 @@ pub struct ContractEvent {
     pub data: serde_json::Value,
 }
 
+pub enum InvokeResponse {
+    Json(Value),
+    NewTx(NewTx),
+}
+
+pub struct NewTx {
+    pub bytes: Vec<u8>,
+    method: String,
+    condition: Option<FinalizationCondition>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+enum FinalizationCondition {
+    AfterBlocks(u64),
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 struct RawEvent {
     pub tx_hash: Vec<u8>,
     pub signature: String,
     pub data: serde_json::Value,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
-enum FinalizationCondition {
-    AfterBlocks(u64),
+#[derive(Deserialize)]
+#[serde(tag = "type")]
+enum RawNewTx {
+    FireFlyCardanoNewTx {
+        bytes: String,
+        condition: FinalizationCondition,
+    },
 }

firefly-cardanoconnect/src/operations/manager.rs

@@ -85,9 +85,9 @@ impl OperationsManager {
             }
         };
         if let Some(tx) = value {
-            let tx_id = self.submit_transaction(from, tx).await?;
+            let tx_id = self.submit_transaction(from, &tx.bytes).await?;
             op.tx_id = Some(tx_id.clone());
-            self.contracts.handle_submit(contract, method, &tx_id).await;
+            self.contracts.handle_submit(contract, &tx_id, tx).await?;
         }
 
         op.status = OperationStatus::Succeeded;
@@ -114,8 +114,8 @@ impl OperationsManager {
         Ok(())
     }
 
-    async fn submit_transaction(&self, address: &str, tx: Vec<u8>) -> ApiResult<String> {
-        let mut transaction: Tx = minicbor::decode(&tx)?;
+    async fn submit_transaction(&self, address: &str, tx: &[u8]) -> ApiResult<String> {
+        let mut transaction: Tx = minicbor::decode(tx)?;
         self.signer
             .sign(address.to_string(), &mut transaction)
             .await?;

wasm/simple-tx/src/lib.rs

@@ -5,10 +5,11 @@ use balius_sdk::{
         AddressPattern, BuildError, FeeChangeReturn, OutputBuilder, TxBuilder, UtxoPattern,
         UtxoSource,
     },
-    Ack, Config, FnHandler, NewTx, Params, Worker, WorkerResult,
+    Ack, Config, FnHandler, Params, Worker, WorkerResult,
 };
 use firefly_balius::{
-    balius_sdk::{self, Json}, kv, CoinSelectionInput, FinalityMonitor, FinalizationCondition, SubmittedTx, WorkerExt as _
+    balius_sdk::{self, Json},
+    kv, CoinSelectionInput, FinalizationCondition, NewMonitoredTx, SubmittedTx, WorkerExt as _,
 };
 use pallas_addresses::Address;
 use serde::{Deserialize, Serialize};
@@ -30,7 +31,7 @@ struct CurrentState {
 }
 
 /// This function builds a transaction to send ADA from one address to another.
-fn send_ada(_: Config<()>, req: Params<SendAdaRequest>) -> WorkerResult<NewTx> {
+fn send_ada(_: Config<()>, req: Params<SendAdaRequest>) -> WorkerResult<NewMonitoredTx> {
     let from_address =
         Address::from_bech32(&req.from_address).map_err(|_| BuildError::MalformedAddress)?;
 
@@ -62,17 +63,18 @@ fn send_ada(_: Config<()>, req: Params<SendAdaRequest>) -> WorkerResult<NewTx> {
         )
         .with_output(FeeChangeReturn(address_source));
 
-    // Return that TX. The framework will sign and submit it.
-    Ok(NewTx(Box::new(tx)))
+    // Return that TX. The framework will sign, submit, and monitor it.
+    // By returning a `NewMonitoredTx`, we tell the framework that we want it to monitor this transaction.
+    // This enables the TransactionApproved, TransactionRolledBack, and TransactionFinalized events from before.
+    // Note that we decide the transaction has been finalized after 4 blocks have reached the chain.
+    Ok(NewMonitoredTx(
+        Box::new(tx),
+        FinalizationCondition::AfterBlocks(4),
+    ))
 }
 
 /// This function is called when a TX produced by this contract is submitted to the blockchain, but before it has reached a block.
 fn handle_submit(_: Config<()>, tx: SubmittedTx) -> WorkerResult<Ack> {
-    // Tell the framework that we want it to monitor this transaction.
-    // This enables the TransactionApproved, TransactionRolledBack, and TransactionFinalized events from before.
-    // Note that we decide the transaction has been finalized after 4 blocks have reached the chain.
-    FinalityMonitor.monitor_tx(&tx.hash, FinalizationCondition::AfterBlocks(4))?;
-
     // Keep track of which TXs have been submitted.
     let mut state: CurrentState = kv::get("current_state")?.unwrap_or_default();
     state.submitted_txs.insert(tx.hash);