move p2p logic outside of py-sdk worker for sharing with validator + orchestrator

Author: JannikSt

…rime-protocol-py/src/worker/constants.rs crates/prime-protocol-py/src/constants.rscrates/prime-protocol-py/src/worker/constants.rs renamed to crates/prime-protocol-py/src/constants.rs crates/prime-protocol-py/src/worker/constants.rs renamed to crates/prime-protocol-py/src/constants.rs

@@ -3,8 +3,10 @@ use crate::validator::ValidatorClient;
 use crate::worker::WorkerClient;
 use pyo3::prelude::*;
 
+mod constants;
 mod error;
 mod orchestrator;
+mod p2p_handler;
 mod utils;
 mod validator;
 mod worker;

crates/prime-protocol-py/src/lib.rs

@@ -1,5 +1,5 @@
 use crate::error::{PrimeProtocolError, Result};
-use crate::worker::p2p::Message;
+use crate::p2p_handler::Message;
 use alloy::primitives::{Address, Signature};
 use rand::Rng;
 use shared::security::request_signer::sign_message;

…tes/prime-protocol-py/src/worker/auth.rs …rime-protocol-py/src/p2p_handler/auth.rscrates/prime-protocol-py/src/worker/auth.rs renamed to crates/prime-protocol-py/src/p2p_handler/auth.rs crates/prime-protocol-py/src/worker/auth.rs renamed to crates/prime-protocol-py/src/p2p_handler/auth.rs

@@ -1,6 +1,6 @@
 use crate::error::Result;
-use crate::worker::auth::AuthenticationManager;
-use crate::worker::p2p::{Message, MessageType};
+use crate::p2p_handler::auth::AuthenticationManager;
+use crate::p2p_handler::{Message, MessageType};
 use std::collections::HashMap;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -50,7 +50,7 @@ impl MessageProcessor {
                 message_result = async {
                     let mut rx = self.message_queue_rx.lock().await;
                     tokio::time::timeout(
-                        crate::worker::constants::MESSAGE_QUEUE_TIMEOUT,
+                        crate::constants::MESSAGE_QUEUE_TIMEOUT,
                         rx.recv()
                     ).await
                 } => {
@@ -63,6 +63,8 @@ impl MessageProcessor {
                         Err(_) => continue, // Timeout, continue loop
                     };
 
+                    log::debug!("Received message: {:?}", message);
+
                     if let Err(e) = self.process_message(message).await {
                         log::error!("Failed to process message: {}", e);
                     }

…tocol-py/src/worker/message_processor.rs …-py/src/p2p_handler/message_processor.rscrates/prime-protocol-py/src/worker/message_processor.rs renamed to crates/prime-protocol-py/src/p2p_handler/message_processor.rs crates/prime-protocol-py/src/worker/message_processor.rs renamed to crates/prime-protocol-py/src/p2p_handler/message_processor.rs

@@ -10,7 +10,10 @@ use tokio::sync::{
 };
 use tokio_util::sync::CancellationToken;
 
-use crate::worker::constants::{MESSAGE_QUEUE_CHANNEL_SIZE, P2P_CHANNEL_SIZE};
+use crate::constants::{MESSAGE_QUEUE_CHANNEL_SIZE, P2P_CHANNEL_SIZE};
+
+pub(crate) mod auth;
+pub(crate) mod message_processor;
 
 // Type alias for the complex return type of Service::new
 type ServiceNewResult = Result<(

…ates/prime-protocol-py/src/worker/p2p.rs …prime-protocol-py/src/p2p_handler/mod.rscrates/prime-protocol-py/src/worker/p2p.rs renamed to crates/prime-protocol-py/src/p2p_handler/mod.rs crates/prime-protocol-py/src/worker/p2p.rs renamed to crates/prime-protocol-py/src/p2p_handler/mod.rs

@@ -8,7 +8,7 @@ use shared::web3::contracts::structs::compute_pool::PoolStatus;
 use shared::web3::wallet::{Wallet, WalletProvider};
 use url::Url;
 
-use crate::worker::constants::{BLOCKCHAIN_OPERATION_TIMEOUT, DEFAULT_COMPUTE_UNITS};
+use crate::constants::{BLOCKCHAIN_OPERATION_TIMEOUT, DEFAULT_COMPUTE_UNITS};
 
 /// Configuration for blockchain operations
 pub struct BlockchainConfig {
@@ -109,7 +109,7 @@ impl BlockchainService {
                         self.config.compute_pool_id,
                         pool.status
                     );
-                    tokio::time::sleep(crate::worker::constants::POOL_STATUS_CHECK_INTERVAL).await;
+                    tokio::time::sleep(crate::constants::POOL_STATUS_CHECK_INTERVAL).await;
                 }
                 Err(e) => {
                     return Err(anyhow::anyhow!("Failed to get pool info: {}", e));

crates/prime-protocol-py/src/worker/blockchain.rs

@@ -1,11 +1,9 @@
+use crate::constants::{DEFAULT_FUNDING_RETRY_COUNT, MESSAGE_QUEUE_TIMEOUT, P2P_SHUTDOWN_TIMEOUT};
 use crate::error::{PrimeProtocolError, Result};
-use crate::worker::auth::AuthenticationManager;
+use crate::p2p_handler::auth::AuthenticationManager;
+use crate::p2p_handler::message_processor::MessageProcessor;
 use crate::worker::blockchain::{BlockchainConfig, BlockchainService};
-use crate::worker::constants::{
-    DEFAULT_FUNDING_RETRY_COUNT, MESSAGE_QUEUE_TIMEOUT, P2P_SHUTDOWN_TIMEOUT,
-};
-use crate::worker::message_processor::MessageProcessor;
-use crate::worker::p2p::{Message, MessageType, Service as P2PService};
+use crate::worker::p2p_handler::{Message, MessageType, Service as P2PService};
 use p2p::{Keypair, PeerId};
 use std::sync::Arc;
 use tokio::sync::mpsc::{Receiver, Sender};

crates/prime-protocol-py/src/worker/client.rs

@@ -1,19 +1,15 @@
 use pyo3::prelude::*;
 
-mod auth;
 mod blockchain;
 mod client;
-mod constants;
 mod discovery;
-mod message_processor;
-mod p2p;
 
+use crate::constants::DEFAULT_P2P_PORT;
+use crate::p2p_handler;
+use crate::p2p_handler::Message;
 pub(crate) use client::WorkerClientCore;
 use tokio_util::sync::CancellationToken;
 
-use crate::worker::p2p::Message;
-use constants::DEFAULT_P2P_PORT;
-
 /// Prime Protocol Worker Client - for compute nodes that execute tasks
 #[pyclass]
 pub(crate) struct WorkerClient {
@@ -87,7 +83,7 @@ impl WorkerClient {
         let rt = self.ensure_runtime()?;
 
         let message = Message {
-            message_type: p2p::MessageType::General { data },
+            message_type: p2p_handler::MessageType::General { data },
             peer_id,
             multiaddrs,
             sender_address: None, // Will be filled from our wallet automatically
@@ -192,19 +188,19 @@ fn to_py_runtime_err(msg: &str) -> PyErr {
 
 fn message_to_pyobject(message: Message) -> PyObject {
     let message_data = match message.message_type {
-        p2p::MessageType::General { data } => {
+        p2p_handler::MessageType::General { data } => {
             serde_json::json!({
                 "type": "general",
                 "data": data,
             })
         }
-        p2p::MessageType::AuthenticationInitiation { challenge } => {
+        p2p_handler::MessageType::AuthenticationInitiation { challenge } => {
             serde_json::json!({
                 "type": "auth_initiation",
                 "challenge": challenge,
             })
         }
-        p2p::MessageType::AuthenticationResponse {
+        p2p_handler::MessageType::AuthenticationResponse {
             challenge,
             signature,
         } => {
@@ -214,7 +210,7 @@ fn message_to_pyobject(message: Message) -> PyObject {
                 "signature": signature,
             })
         }
-        p2p::MessageType::AuthenticationSolution { signature } => {
+        p2p_handler::MessageType::AuthenticationSolution { signature } => {
             serde_json::json!({
                 "type": "auth_solution",
                 "signature": signature,

crates/prime-protocol-py/src/worker/mod.rs

@@ -166,7 +166,6 @@ async fn upload_to_single_discovery(
     node_data: &serde_json::Value,
     wallet: &Wallet,
 ) -> Result<()> {
-
     let signed_request = sign_request_with_nonce(endpoint, wallet, Some(node_data))
         .await
         .map_err(|e| anyhow::anyhow!("{}", e))?;