fix(tee-prover): simplify TeeProofGenerationDataResponse

This PR addresses Alex's code review comment:
#3017 (comment)
and partially addresses this one too:
#3017 (comment)

core/bin/zksync_tee_prover/src/api_client.rs

@@ -77,7 +77,7 @@ impl TeeApiClient {
     pub async fn get_job(
         &self,
         tee_type: TeeType,
-    ) -> Result<Option<Box<TeeVerifierInput>>, TeeProverError> {
+    ) -> Result<Box<TeeVerifierInput>, TeeProverError> {
         let request = TeeProofGenerationDataRequest { tee_type };
         let response = self
             .post::<_, TeeProofGenerationDataResponse, _>("/tee/proof_inputs", request)

core/bin/zksync_tee_prover/src/tee_prover.rs

@@ -1,5 +1,6 @@
 use std::fmt;
 
+use reqwest::StatusCode;
 use secp256k1::{ecdsa::Signature, Message, PublicKey, Secp256k1};
 use zksync_basic_types::H256;
 use zksync_node_framework::{
@@ -90,8 +91,8 @@ impl TeeProver {
     }
 
     async fn step(&self, public_key: &PublicKey) -> Result<Option<L1BatchNumber>, TeeProverError> {
-        match self.api_client.get_job(self.config.tee_type).await? {
-            Some(job) => {
+        match self.api_client.get_job(self.config.tee_type).await {
+            Ok(job) => {
                 let (signature, batch_number, root_hash) = self.verify(*job)?;
                 self.api_client
                     .submit_proof(
@@ -104,10 +105,15 @@ impl TeeProver {
                     .await?;
                 Ok(Some(batch_number))
             }
-            None => {
-                tracing::trace!("There are currently no pending batches to be proven");
-                Ok(None)
-            }
+            Err(err) => match err {
+                TeeProverError::Request(req_err)
+                    if req_err.status() == Some(StatusCode::NOT_FOUND) =>
+                {
+                    tracing::trace!("There are currently no pending batches to be proven");
+                    Ok(None)
+                }
+                _ => Err(err),
+            },
         }
     }
 }

core/lib/dal/src/tee_proof_generation_dal.rs

@@ -33,7 +33,7 @@ impl TeeProofGenerationDal<'_, '_> {
         tee_type: TeeType,
         processing_timeout: Duration,
         min_batch_number: Option<L1BatchNumber>,
-    ) -> DalResult<Option<L1BatchNumber>> {
+    ) -> DalResult<L1BatchNumber> {
         let processing_timeout = pg_interval_from_duration(processing_timeout);
         let min_batch_number = min_batch_number.map_or(0, |num| i64::from(num.0));
         let batch_number = sqlx::query!(
@@ -105,11 +105,10 @@ impl TeeProofGenerationDal<'_, '_> {
         .with_arg("tee_type", &tee_type)
         .with_arg("processing_timeout", &processing_timeout)
         .with_arg("l1_batch_number", &min_batch_number)
-        .fetch_optional(self.storage)
-        .await?
-        .map(|row| L1BatchNumber(row.l1_batch_number as u32));
+        .fetch_one(self.storage)
+        .await?;
 
-        Ok(batch_number)
+        Ok(L1BatchNumber(batch_number.l1_batch_number as u32))
     }
 
     pub async fn unlock_batch(

core/lib/prover_interface/src/api.rs

@@ -31,7 +31,7 @@ pub enum ProofGenerationDataResponse {
 }
 
 #[derive(Debug, Serialize, Deserialize)]
-pub struct TeeProofGenerationDataResponse(pub Option<Box<TeeVerifierInput>>);
+pub struct TeeProofGenerationDataResponse(pub Box<TeeVerifierInput>);
 
 #[derive(Debug, Serialize, Deserialize)]
 pub enum SubmitProofResponse {

core/node/proof_data_handler/src/lib.rs

@@ -10,7 +10,7 @@ use zksync_dal::{ConnectionPool, Core};
 use zksync_object_store::ObjectStore;
 use zksync_prover_interface::api::{
     ProofGenerationDataRequest, RegisterTeeAttestationRequest, SubmitProofRequest,
-    SubmitTeeProofRequest, TeeProofGenerationDataRequest, TeeProofGenerationDataResponse,
+    SubmitTeeProofRequest, TeeProofGenerationDataRequest,
 };
 use zksync_types::{commitment::L1BatchCommitmentMode, L2ChainId};
 
@@ -109,7 +109,6 @@ fn create_proof_processing_router(
                         .await;
 
                     match result {
-                        Ok(Json(TeeProofGenerationDataResponse(None))) => (StatusCode::NO_CONTENT, Json("No new TeeVerifierInputs are available yet")).into_response(),
                         Ok(data) => (StatusCode::OK, data).into_response(),
                         Err(e) => e.into_response(),
                     }

core/node/proof_data_handler/src/tee_request_processor.rs

@@ -51,20 +51,16 @@ impl TeeRequestProcessor {
         let mut missing_range: Option<(L1BatchNumber, L1BatchNumber)> = None;
 
         let result = loop {
-            let l1_batch_number = match self
+            let l1_batch_number = self
                 .lock_batch_for_proving(request.tee_type, min_batch_number)
-                .await?
-            {
-                Some(number) => number,
-                None => break Ok(Json(TeeProofGenerationDataResponse(None))),
-            };
+                .await?;
 
             match self
                 .tee_verifier_input_for_existing_batch(l1_batch_number)
                 .await
             {
                 Ok(input) => {
-                    break Ok(Json(TeeProofGenerationDataResponse(Some(Box::new(input)))));
+                    break Ok(Json(TeeProofGenerationDataResponse(Box::new(input))));
                 }
                 Err(RequestProcessorError::ObjectStore(ObjectStoreError::KeyNotFound(_))) => {
                     missing_range = match missing_range {
@@ -156,7 +152,7 @@ impl TeeRequestProcessor {
         &self,
         tee_type: TeeType,
         min_batch_number: Option<L1BatchNumber>,
-    ) -> Result<Option<L1BatchNumber>, RequestProcessorError> {
+    ) -> Result<L1BatchNumber, RequestProcessorError> {
         let result = self
             .pool
             .connection()

core/node/proof_data_handler/src/tests.rs

@@ -31,7 +31,7 @@ async fn request_tee_proof_inputs() {
         L2ChainId::default(),
     );
     let test_cases = vec![
-        (json!({ "tee_type": "sgx" }), StatusCode::NO_CONTENT),
+        (json!({ "tee_type": "sgx" }), StatusCode::NOT_FOUND),
         (
             json!({ "tee_type": "Sgx" }),
             StatusCode::UNPROCESSABLE_ENTITY,