BM-1547: move lift to join

bento/crates/workflow/src/tasks/join.rs

@@ -9,7 +9,7 @@ use crate::{
     tasks::{RECUR_RECEIPT_PATH, deserialize_obj, serialize_obj},
 };
 use anyhow::{Context, Result};
-use risc0_zkvm::{ReceiptClaim, SuccinctReceipt};
+use risc0_zkvm::{ReceiptClaim, SegmentReceipt, SuccinctReceipt};
 use uuid::Uuid;
 use workflow_common::JoinReq;
 
@@ -23,15 +23,52 @@ pub async fn join(agent: &Agent, job_id: &Uuid, request: &JoinReq) -> Result<()>
     let left_path_key = format!("{recur_receipts_prefix}:{}", request.left);
     let right_path_key = format!("{recur_receipts_prefix}:{}", request.right);
 
-    let (left_receipt, right_receipt): (Vec<u8>, Vec<u8>) =
+    let (left_receipt_data, right_receipt_data): (Vec<u8>, Vec<u8>) =
         conn.mget::<_, (Vec<u8>, Vec<u8>)>(&[&left_path_key, &right_path_key]).await.with_context(
             || format!("failed to get receipts for keys: {left_path_key}, {right_path_key}"),
         )?;
 
-    let left_receipt: SuccinctReceipt<ReceiptClaim> =
-        deserialize_obj(&left_receipt).context("Failed to deserialize left receipt")?;
-    let right_receipt: SuccinctReceipt<ReceiptClaim> =
-        deserialize_obj(&right_receipt).context("Failed to deserialize right receipt")?;
+    // Handle each receipt independently - they could be mixed types
+    let prover = agent.prover.as_ref().context("Missing prover from resolve task")?;
+    let left_idx = request.left;
+    let right_idx = request.right;
+
+    let (left_receipt, right_receipt) = tokio::try_join!(
+        async {
+            match deserialize_obj::<SegmentReceipt>(&left_receipt_data) {
+                Ok(segment) => {
+                    let receipt = prover
+                        .lift(&segment)
+                        .with_context(|| format!("Failed to lift segment {left_idx}"))?;
+                    tracing::debug!("lifting complete {job_id} - {left_idx}");
+                    Ok::<_, anyhow::Error>(receipt)
+                }
+                Err(_) => {
+                    let receipt: SuccinctReceipt<ReceiptClaim> =
+                        deserialize_obj(&left_receipt_data)
+                            .context("Failed to deserialize left receipt")?;
+                    Ok(receipt)
+                }
+            }
+        },
+        async {
+            match deserialize_obj::<SegmentReceipt>(&right_receipt_data) {
+                Ok(segment) => {
+                    let receipt = prover
+                        .lift(&segment)
+                        .with_context(|| format!("Failed to lift segment {right_idx}"))?;
+                    tracing::debug!("lifting complete {job_id} - {right_idx}");
+                    Ok::<_, anyhow::Error>(receipt)
+                }
+                Err(_) => {
+                    let receipt: SuccinctReceipt<ReceiptClaim> =
+                        deserialize_obj(&right_receipt_data)
+                            .context("Failed to deserialize right receipt")?;
+                    Ok(receipt)
+                }
+            }
+        }
+    )?;
 
     tracing::trace!("Joining {job_id} - {} + {} -> {}", request.left, request.right, request.idx);
 

bento/crates/workflow/src/tasks/join_povw.rs

@@ -9,7 +9,7 @@ use crate::{
     tasks::{RECUR_RECEIPT_PATH, deserialize_obj, serialize_obj},
 };
 use anyhow::{Context, Result};
-use risc0_zkvm::{ReceiptClaim, SuccinctReceipt, WorkClaim};
+use risc0_zkvm::{ReceiptClaim, SegmentReceipt, SuccinctReceipt, WorkClaim};
 use uuid::Uuid;
 use workflow_common::JoinReq;
 
@@ -29,14 +29,50 @@ pub async fn join_povw(agent: &Agent, job_id: &Uuid, request: &JoinReq) -> Resul
             format!("failed to get receipts for keys: {left_receipt_key}, {right_receipt_key}")
         })?;
 
-    // Deserialize POVW receipts
+    // Handle each receipt independently - they could be mixed types
+    let prover = agent.prover.as_ref().context("Missing prover from POVW join task")?;
+
     let (left_receipt, right_receipt): (
         SuccinctReceipt<WorkClaim<ReceiptClaim>>,
         SuccinctReceipt<WorkClaim<ReceiptClaim>>,
-    ) = (
-        deserialize_obj::<SuccinctReceipt<WorkClaim<ReceiptClaim>>>(&left_receipt_bytes)?,
-        deserialize_obj::<SuccinctReceipt<WorkClaim<ReceiptClaim>>>(&right_receipt_bytes)?,
-    );
+    ) = tokio::try_join!(
+        async {
+            match deserialize_obj::<SegmentReceipt>(&left_receipt_bytes) {
+                Ok(segment_receipt) => {
+                    // Successfully deserialized as segment receipt, now lift to POVW
+                    let povw_receipt = prover
+                        .lift_povw(&segment_receipt)
+                        .context("Failed to lift left segment to POVW")?;
+                    Ok::<_, anyhow::Error>(povw_receipt)
+                }
+                Err(_) => {
+                    // Failed to deserialize as segment, try as already-lifted POVW receipt
+                    let povw_receipt: SuccinctReceipt<WorkClaim<ReceiptClaim>> =
+                        deserialize_obj(&left_receipt_bytes)
+                            .context("Failed to deserialize left POVW receipt")?;
+                    Ok(povw_receipt)
+                }
+            }
+        },
+        async {
+            match deserialize_obj::<SegmentReceipt>(&right_receipt_bytes) {
+                Ok(segment_receipt) => {
+                    // Successfully deserialized as segment receipt, now lift to POVW
+                    let povw_receipt = prover
+                        .lift_povw(&segment_receipt)
+                        .context("Failed to lift right segment to POVW")?;
+                    Ok::<_, anyhow::Error>(povw_receipt)
+                }
+                Err(_) => {
+                    // Failed to deserialize as segment, try as already-lifted POVW receipt
+                    let povw_receipt: SuccinctReceipt<WorkClaim<ReceiptClaim>> =
+                        deserialize_obj(&right_receipt_bytes)
+                            .context("Failed to deserialize right POVW receipt")?;
+                    Ok(povw_receipt)
+                }
+            }
+        }
+    )?;
 
     tracing::debug!("Starting POVW join of receipts {} and {}", request.left, request.right);
 

bento/crates/workflow/src/tasks/prove.rs

@@ -9,7 +9,7 @@ use crate::{
     tasks::{RECUR_RECEIPT_PATH, SEGMENTS_PATH, deserialize_obj, serialize_obj},
 };
 use anyhow::{Context, Result};
-use risc0_zkvm::{ReceiptClaim, SuccinctReceipt, WorkClaim};
+use risc0_zkvm::Segment;
 use uuid::Uuid;
 use workflow_common::ProveReq;
 
@@ -25,52 +25,37 @@ pub async fn prover(agent: &Agent, job_id: &Uuid, task_id: &str, request: &Prove
         .get::<_, Vec<u8>>(&segment_key)
         .await
         .with_context(|| format!("segment data not found for segment key: {segment_key}"))?;
-    let segment =
-        deserialize_obj(&segment_vec).context("Failed to deserialize segment data from redis")?;
 
-    let segment_receipt = agent
-        .prover
-        .as_ref()
-        .context("Missing prover from prove task")?
-        .prove_segment(&agent.verifier_ctx, &segment)
-        .context("Failed to prove segment")?;
+    // Try to deserialize as segment first, then prove it
+    let segment_receipt = match deserialize_obj::<Segment>(&segment_vec) {
+        Ok(segment) => {
+            // Successfully deserialized as segment, now prove it
+            agent
+                .prover
+                .as_ref()
+                .context("Missing prover from prove task")?
+                .prove_segment(&agent.verifier_ctx, &segment)
+                .context("Failed to prove segment")?
+        }
+        Err(_) => {
+            // Failed to deserialize as segment, try as already-proven receipt
+            deserialize_obj(&segment_vec)
+                .context("Failed to deserialize segment data from redis")?
+        }
+    };
 
     tracing::debug!("Completed proof: {job_id} - {index}");
-
-    tracing::debug!("lifting {job_id} - {index}");
-
     let output_key = format!("{job_prefix}:{RECUR_RECEIPT_PATH}:{task_id}");
 
-    if agent.is_povw_enabled() {
-        let lift_receipt: SuccinctReceipt<WorkClaim<ReceiptClaim>> = agent
-            .prover
-            .as_ref()
-            .context("Missing prover from resolve task")?
-            .lift_povw(&segment_receipt)
-            .with_context(|| format!("Failed to POVW lift segment {index}"))?;
-
-        tracing::debug!("lifting complete {job_id} - {index}");
-
-        // Write out lifted POVW receipt
-        let lift_asset =
-            serialize_obj(&lift_receipt).expect("Failed to serialize the POVW segment");
-        redis::set_key_with_expiry(&mut conn, &output_key, lift_asset, Some(agent.args.redis_ttl))
-            .await?;
-    } else {
-        let lift_receipt: SuccinctReceipt<ReceiptClaim> = agent
-            .prover
-            .as_ref()
-            .context("Missing prover from resolve task")?
-            .lift(&segment_receipt)
-            .with_context(|| format!("Failed to lift segment {index}"))?;
-
-        tracing::debug!("lifting complete {job_id} - {index}");
-
-        // Write out lifted regular receipt
-        let lift_asset = serialize_obj(&lift_receipt).expect("Failed to serialize the segment");
-        redis::set_key_with_expiry(&mut conn, &output_key, lift_asset, Some(agent.args.redis_ttl))
-            .await?;
-    }
+    let serialized_receipt =
+        serialize_obj(&segment_receipt).context("Failed to serialize segment receipt")?;
+    redis::set_key_with_expiry(
+        &mut conn,
+        &output_key,
+        serialized_receipt,
+        Some(agent.args.redis_ttl),
+    )
+    .await?;
 
     Ok(())
 }

bento/crates/workflow/src/tasks/resolve.rs

@@ -10,7 +10,7 @@ use crate::{
 };
 use anyhow::{Context, Result};
 use risc0_zkvm::sha::Digestible;
-use risc0_zkvm::{ReceiptClaim, SuccinctReceipt, Unknown};
+use risc0_zkvm::{ReceiptClaim, SegmentReceipt, SuccinctReceipt, Unknown};
 use uuid::Uuid;
 use workflow_common::{KECCAK_RECEIPT_PATH, ResolveReq};
 
@@ -29,76 +29,86 @@ pub async fn resolver(agent: &Agent, job_id: &Uuid, request: &ResolveReq) -> Res
     })?;
 
     tracing::debug!("Root receipt size: {} bytes", receipt.len());
-    let mut conditional_receipt: SuccinctReceipt<ReceiptClaim> = deserialize_obj(&receipt)?;
+
+    // Try to deserialize as segment receipt first, then lift it if needed
+    let mut conditional_receipt: SuccinctReceipt<ReceiptClaim> =
+        match deserialize_obj::<SegmentReceipt>(&receipt) {
+            Ok(segment_receipt) => {
+                // Successfully deserialized as segment receipt, now lift it
+                agent
+                    .prover
+                    .as_ref()
+                    .context("Missing prover from resolve task")?
+                    .lift(&segment_receipt)
+                    .context("Failed to lift segment receipt")?
+            }
+            Err(_) => {
+                // Failed to deserialize as segment receipt, try as already-lifted succinct receipt
+                deserialize_obj(&receipt).context("Failed to deserialize receipt")?
+            }
+        };
 
     let mut assumptions_len: Option<u64> = None;
-    if conditional_receipt.claim.clone().as_value()?.output.is_some() {
-        if let Some(guest_output) =
+    if conditional_receipt.claim.clone().as_value()?.output.is_some()
+        && let Some(guest_output) =
             conditional_receipt.claim.clone().as_value()?.output.as_value()?
-        {
-            if !guest_output.assumptions.is_empty() {
-                let assumptions = guest_output
-                    .assumptions
-                    .as_value()
-                    .context("Failed unwrap the assumptions of the guest output")?
-                    .iter();
-
-                tracing::debug!("Resolving {} assumption(s)", assumptions.len());
-                assumptions_len =
-                    Some(assumptions.len().try_into().context("Failed to convert to u64")?);
-
-                let mut union_claim = String::new();
-                if let Some(idx) = request.union_max_idx {
-                    let union_root_receipt_key =
-                        format!("{job_prefix}:{KECCAK_RECEIPT_PATH}:{idx}");
-                    tracing::debug!(
-                        "Deserializing union_root_receipt_key: {union_root_receipt_key}"
-                    );
-                    let union_receipt: Vec<u8> = conn.get(&union_root_receipt_key).await?;
-                    let union_receipt: SuccinctReceipt<Unknown> =
-                        deserialize_obj(&union_receipt)
-                            .context("Failed to deserialize to SuccinctReceipt<Unknown> type")?;
-                    union_claim = union_receipt.claim.digest().to_string();
-
-                    // Resolve union receipt
-                    tracing::debug!("Resolving union claim digest: {union_claim}");
-                    conditional_receipt = agent
-                        .prover
-                        .as_ref()
-                        .context("Missing prover from resolve task")?
-                        .resolve(&conditional_receipt, &union_receipt)
-                        .context("Failed to resolve the union receipt")?;
-                }
-
-                for assumption in assumptions {
-                    let assumption_claim = assumption.as_value()?.claim.to_string();
-                    if assumption_claim.eq(&union_claim) {
-                        tracing::debug!("Skipping already resolved union claim: {union_claim}");
-                        continue;
-                    }
-                    let assumption_key = format!("{receipts_key}:{assumption_claim}");
-                    tracing::debug!("Deserializing assumption with key: {assumption_key}");
-                    let assumption_bytes: Vec<u8> = conn
-                        .get(&assumption_key)
-                        .await
-                        .context("corroborating receipt not found: key {assumption_key}")?;
-
-                    let assumption_receipt: SuccinctReceipt<Unknown> =
-                        deserialize_obj(&assumption_bytes).with_context(|| {
-                            format!("could not deserialize assumption receipt: {assumption_key}")
-                        })?;
-
-                    // Resolve
-                    conditional_receipt = agent
-                        .prover
-                        .as_ref()
-                        .context("Missing prover from resolve task")?
-                        .resolve(&conditional_receipt, &assumption_receipt)
-                        .context("Failed to resolve the conditional receipt")?;
-                }
-                tracing::debug!("Resolve complete for job_id: {job_id}");
+        && !guest_output.assumptions.is_empty()
+    {
+        let assumptions = guest_output
+            .assumptions
+            .as_value()
+            .context("Failed unwrap the assumptions of the guest output")?
+            .iter();
+
+        tracing::debug!("Resolving {} assumption(s)", assumptions.len());
+        assumptions_len = Some(assumptions.len().try_into().context("Failed to convert to u64")?);
+
+        let mut union_claim = String::new();
+        if let Some(idx) = request.union_max_idx {
+            let union_root_receipt_key = format!("{job_prefix}:{KECCAK_RECEIPT_PATH}:{idx}");
+            tracing::debug!("Deserializing union_root_receipt_key: {union_root_receipt_key}");
+            let union_receipt: Vec<u8> = conn.get(&union_root_receipt_key).await?;
+            let union_receipt: SuccinctReceipt<Unknown> = deserialize_obj(&union_receipt)
+                .context("Failed to deserialize to SuccinctReceipt<Unknown> type")?;
+            union_claim = union_receipt.claim.digest().to_string();
+
+            // Resolve union receipt
+            tracing::debug!("Resolving union claim digest: {union_claim}");
+            conditional_receipt = agent
+                .prover
+                .as_ref()
+                .context("Missing prover from resolve task")?
+                .resolve(&conditional_receipt, &union_receipt)
+                .context("Failed to resolve the union receipt")?;
+        }
+
+        for assumption in assumptions {
+            let assumption_claim = assumption.as_value()?.claim.to_string();
+            if assumption_claim.eq(&union_claim) {
+                tracing::debug!("Skipping already resolved union claim: {union_claim}");
+                continue;
             }
+            let assumption_key = format!("{receipts_key}:{assumption_claim}");
+            tracing::debug!("Deserializing assumption with key: {assumption_key}");
+            let assumption_bytes: Vec<u8> = conn
+                .get(&assumption_key)
+                .await
+                .context("corroborating receipt not found: key {assumption_key}")?;
+
+            let assumption_receipt: SuccinctReceipt<Unknown> = deserialize_obj(&assumption_bytes)
+                .with_context(|| {
+                format!("could not deserialize assumption receipt: {assumption_key}")
+            })?;
+
+            // Resolve
+            conditional_receipt = agent
+                .prover
+                .as_ref()
+                .context("Missing prover from resolve task")?
+                .resolve(&conditional_receipt, &assumption_receipt)
+                .context("Failed to resolve the conditional receipt")?;
         }
+        tracing::debug!("Resolve complete for job_id: {job_id}");
     }
 
     // Write out the resolved receipt