rvoip_transaction_core/client/

invite.rs

/// # INVITE Client Transaction Implementation
///
/// This module implements the INVITE client transaction state machine as defined in
/// [RFC 3261 Section 17.1.1](https://datatracker.ietf.org/doc/html/rfc3261#section-17.1.1).
///
/// ## State Machine
///
/// The INVITE client transaction follows this state machine:
///
/// ```text
///                                         |INVITE from TU
///                Timer A fires            |INVITE sent
///                Reset A,                 V
///                INVITE sent +-----------+
///                  +---------|           |-------------------+
///                  |         |  Calling  |  Timer B fires    |
///                  +-------->|           |  or Transport Err.|
///                            +-----------+  inform TU        |
///                               |  |                         |
///                               |  |1xx                      |
///                               |  |from                     |
///                               |  |TU                       V
///                            1xx|  |                   +-----------+
///                            from|  |                   |           |
///                            TL  |  |                   | Terminated|
///                               |  |                    |           |
///                               |  |                    +-----------+
///                               |  |
///                               |  |
///                               |  |                      300-699 from TL
///                               |  |                      ACK sent, resp. to TU
///                               |  |                         +----+
///                               |  |                         |    |
///                               |  |                         V    |
///                 +-----------+ |  |      Timer D fires  +-----------+
///                 |           |<|--|-------------------+|           |
///                 |Proceeding |--+  |                   | Completed |
///                 |           |<----+                   |           |
///                 +-----------+                         +-----------+
///                   |      |                            ^      |
///                   |      |                            |      |
///                   |      +-------------+              |      |
///                   |                    |              |      |
///                   |                    |              |      |
///                   |                    |              |      |
///    300-699 from TL|                    |2xx from TL   |      |2xx from TL
///     resp. to TU   |                    |resp. to TU   |      |resp. to TU
///                   |                    |              |      |
///                   V                    V              |      |
///             +-----------+        +-----------+        |      |
///             |           |        |           |<-------+      |
///             | Completed |        | Terminated|<--------------+
///             |           |        |           |
///             +-----------+        +-----------+
/// ```
///
/// ## Timers
///
/// INVITE client transactions use the following timers:
///
/// - **Timer A**: Initial value T1, doubles on each retransmission. Controls request retransmissions in Calling state.
/// - **Timer B**: Typically 64*T1. Controls transaction timeout. When it fires, the transaction terminates with an error.
/// - **Timer D**: Typically 32s, but can be shorter. Controls how long to wait for retransmissions of the final response.

use std::fmt;
use std::future::Future;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{mpsc, Mutex};
use tokio::task::JoinHandle;
use tracing::{debug, error, trace, warn};

use rvoip_sip_core::prelude::*;
use rvoip_sip_transport::Transport;

use crate::error::{Error, Result};
use crate::transaction::{
    Transaction, TransactionAsync, TransactionState, TransactionKind, TransactionKey, TransactionEvent,
    InternalTransactionCommand, AtomicTransactionState
};
use crate::timer::{TimerSettings, TimerFactory, TimerManager, TimerType};
use crate::client::{
    ClientTransaction, ClientTransactionData,
};
use crate::utils;
use crate::transaction::logic::TransactionLogic;
use crate::transaction::runner::{run_transaction_loop, HasCommandSender, AsRefKey};
// Add imports for our utility modules
use crate::transaction::timer_utils;
use crate::transaction::validators;
use crate::transaction::common_logic;

/// Client INVITE transaction implementation as defined in RFC 3261 Section 17.1.1.
///
/// This struct implements the state machine for client INVITE transactions, which are used
/// for initiating SIP sessions. INVITE transactions have unique behavior, including:
///
/// - Automatic ACK generation for non-2xx responses
/// - Special handling for 2xx responses (transaction terminates without sending ACK)
/// - Unique timer requirements (Timers A, B, D)
///
/// Key behaviors:
/// - In Calling state: Retransmits INVITE periodically until response or timeout
/// - In Proceeding state: Waits for a final response
/// - In Completed state: Has received a final non-2xx response, sent ACK, waiting for retransmissions
/// - In Terminated state: Transaction is finished
#[derive(Debug, Clone)]
pub struct ClientInviteTransaction {
    data: Arc<ClientTransactionData>,
    logic: Arc<ClientInviteLogic>,
}

/// Holds JoinHandles and dynamic state for timers specific to Client INVITE transactions.
///
/// Used by the transaction runner to manage the various timers required by the
/// INVITE client transaction state machine as defined in RFC 3261.
#[derive(Default, Debug)]
struct ClientInviteTimerHandles {
    /// Handle for Timer A, which controls INVITE retransmissions
    timer_a: Option<JoinHandle<()>>,
    
    /// Current interval for Timer A, which doubles after each firing
    current_timer_a_interval: Option<Duration>, // For backoff
    
    /// Handle for Timer B, which controls transaction timeout
    timer_b: Option<JoinHandle<()>>,
    
    /// Handle for Timer D, which controls how long to wait in Completed state
    timer_d: Option<JoinHandle<()>>,
}

/// Implements the TransactionLogic for Client INVITE transactions.
///
/// This struct contains the core logic for the INVITE client transaction state machine,
/// implementing the behavior defined in RFC 3261 Section 17.1.1.
#[derive(Debug, Clone, Default)]
struct ClientInviteLogic {
    _data_marker: std::marker::PhantomData<ClientTransactionData>,
    timer_factory: TimerFactory,
}

impl ClientInviteLogic {
    // Helper method to start Timer A (retransmission timer) using timer utils
    async fn start_timer_a(
        &self,
        data: &Arc<ClientTransactionData>,
        timer_handles: &mut ClientInviteTimerHandles,
        command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) {
        let tx_id = &data.id;
        let timer_config = &data.timer_config;
        
        // Start Timer A (retransmission) with initial interval T1
        let initial_interval_a = timer_handles.current_timer_a_interval.unwrap_or(timer_config.t1);
        timer_handles.current_timer_a_interval = Some(initial_interval_a);
        
        // Use timer_utils to start the timer
        let timer_manager = self.timer_factory.timer_manager();
        match timer_utils::start_transaction_timer(
            &timer_manager,
            tx_id,
            "A",
            TimerType::A,
            initial_interval_a,
            command_tx
        ).await {
            Ok(handle) => {
                timer_handles.timer_a = Some(handle);
                trace!(id=%tx_id, interval=?initial_interval_a, "Started Timer A for Calling state");
            },
            Err(e) => {
                error!(id=%tx_id, error=%e, "Failed to start Timer A");
            }
        }
    }
    
    // Helper method to start Timer B (transaction timeout) using timer utils
    async fn start_timer_b(
        &self,
        data: &Arc<ClientTransactionData>,
        timer_handles: &mut ClientInviteTimerHandles,
        command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) {
        let tx_id = &data.id;
        let timer_config = &data.timer_config;
        
        // Start Timer B (transaction timeout)
        let interval_b = timer_config.transaction_timeout;
        
        // Use timer_utils to start the timer
        let timer_manager = self.timer_factory.timer_manager();
        match timer_utils::start_transaction_timer(
            &timer_manager,
            tx_id,
            "B", 
            TimerType::B,
            interval_b,
            command_tx
        ).await {
            Ok(handle) => {
                timer_handles.timer_b = Some(handle);
                trace!(id=%tx_id, interval=?interval_b, "Started Timer B for Calling state");
            },
            Err(e) => {
                error!(id=%tx_id, error=%e, "Failed to start Timer B");
            }
        }
    }
    
    // Helper method to start Timer D (wait for response retransmissions) using timer utils with transition
    async fn start_timer_d(
        &self,
        data: &Arc<ClientTransactionData>,
        timer_handles: &mut ClientInviteTimerHandles,
        command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) {
        let tx_id = &data.id;
        let timer_config = &data.timer_config;
        
        // Start Timer D that automatically transitions to Terminated state when it fires
        let interval_d = timer_config.wait_time_d;
        
        // Use timer_utils to start the timer with transition
        let timer_manager = self.timer_factory.timer_manager();
        match timer_utils::start_timer_with_transition(
            &timer_manager,
            tx_id,
            "D",
            TimerType::D,
            interval_d,
            command_tx,
            TransactionState::Terminated
        ).await {
            Ok(handle) => {
                timer_handles.timer_d = Some(handle);
                trace!(id=%tx_id, interval=?interval_d, "Started Timer D for Completed state");
            },
            Err(e) => {
                error!(id=%tx_id, error=%e, "Failed to start Timer D");
            }
        }
    }
    
    /// Handle initial request sending in Calling state
    ///
    /// This method is called when the transaction enters the Calling state.
    /// It sends the initial INVITE request and starts Timers A and B according to RFC 3261 Section 17.1.1.2.
    async fn handle_calling_state(
        &self,
        data: &Arc<ClientTransactionData>,
        timer_handles: &mut ClientInviteTimerHandles,
        command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) -> Result<()> {
        let tx_id = &data.id;
        
        // Send the initial request
        debug!(id=%tx_id, "ClientInviteLogic: Sending initial request in Calling state");
        let request_guard = data.request.lock().await;
        if let Err(e) = data.transport.send_message(
            Message::Request(request_guard.clone()),
            data.remote_addr
        ).await {
            error!(id=%tx_id, error=%e, "Failed to send initial request from Calling state");
            common_logic::send_transport_error_event(tx_id, &data.events_tx).await;
            // If send fails, command a transition to Terminated
            let _ = command_tx.send(InternalTransactionCommand::TransitionTo(TransactionState::Terminated)).await;
            return Err(Error::transport_error(e, "Failed to send initial request"));
        }
        drop(request_guard); // Release lock

        // Start timers for Calling state
        timer_handles.current_timer_a_interval = Some(data.timer_config.t1);
        self.start_timer_a(data, timer_handles, command_tx.clone()).await;
        self.start_timer_b(data, timer_handles, command_tx).await;
        
        Ok(())
    }

    /// Handle Timer A (retransmission) trigger
    ///
    /// This method is called when Timer A fires. According to RFC 3261 Section 17.1.1.2,
    /// when Timer A fires in the Calling state, the client should retransmit the INVITE
    /// request and restart Timer A with a doubled interval (capped by T2).
    async fn handle_timer_a_trigger(
        &self,
        data: &Arc<ClientTransactionData>,
        current_state: TransactionState,
        timer_handles: &mut ClientInviteTimerHandles,
        command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) -> Result<Option<TransactionState>> {
        let tx_id = &data.id;
        let timer_config = &data.timer_config;
        
        match current_state {
            TransactionState::Calling => {
                debug!(id=%tx_id, "Timer A triggered, retransmitting INVITE request");
                
                // Retransmit the request
                let request_guard = data.request.lock().await;
                if let Err(e) = data.transport.send_message(
                    Message::Request(request_guard.clone()),
                    data.remote_addr
                ).await {
                    error!(id=%tx_id, error=%e, "Failed to retransmit request");
                    common_logic::send_transport_error_event(tx_id, &data.events_tx).await;
                    return Ok(Some(TransactionState::Terminated));
                }
                
                // Update and restart Timer A with increased interval using the utility function
                let current_interval = timer_handles.current_timer_a_interval.unwrap_or(timer_config.t1);
                let new_interval = timer_utils::calculate_backoff_interval(current_interval, timer_config);
                timer_handles.current_timer_a_interval = Some(new_interval);
                
                // Start new Timer A with the increased interval
                self.start_timer_a(data, timer_handles, command_tx).await;
            },
            _ => {
                trace!(id=%tx_id, state=?current_state, "Timer A fired in invalid state, ignoring");
            }
        }
        
        Ok(None)
    }
    
    /// Handle Timer B (transaction timeout) trigger
    ///
    /// This method is called when Timer B fires. According to RFC 3261 Section 17.1.1.2,
    /// when Timer B fires in the Calling state, the client should inform the TU that the
    /// transaction has timed out and transition to the Terminated state.
    async fn handle_timer_b_trigger(
        &self,
        data: &Arc<ClientTransactionData>,
        current_state: TransactionState,
        _command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) -> Result<Option<TransactionState>> {
        let tx_id = &data.id;
        
        match current_state {
            TransactionState::Calling => {
                warn!(id=%tx_id, "Timer B (Timeout) fired in state {:?}", current_state);
                
                // Notify TU about timeout using common logic
                common_logic::send_transaction_timeout_event(tx_id, &data.events_tx).await;
                
                // Return state transition
                return Ok(Some(TransactionState::Terminated));
            },
            _ => {
                trace!(id=%tx_id, state=?current_state, "Timer B fired in invalid state, ignoring");
            }
        }
        
        Ok(None)
    }
    
    /// Handle Timer D (wait for retransmissions) trigger
    ///
    /// This method is called when Timer D fires. According to RFC 3261 Section 17.1.1.2,
    /// when Timer D fires in the Completed state, the client should transition to the
    /// Terminated state. Timer D ensures that any retransmissions of the final response
    /// are properly ACKed before the transaction terminates.
    async fn handle_timer_d_trigger(
        &self,
        data: &Arc<ClientTransactionData>,
        current_state: TransactionState,
        _command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) -> Result<Option<TransactionState>> {
        let tx_id = &data.id;
        
        match current_state {
            TransactionState::Completed => {
                debug!(id=%tx_id, "Timer D fired in Completed state, terminating");
                // Timer D automatically transitions to Terminated (handled by timer_utils)
                Ok(None)
            },
            _ => {
                trace!(id=%tx_id, state=?current_state, "Timer D fired in invalid state, ignoring");
                Ok(None)
            }
        }
    }
    
    /// Create and send an ACK for a non-2xx response
    ///
    /// According to RFC 3261 Section 17.1.1.3, the client transaction must generate an
    /// ACK request when it receives a final response (3xx-6xx) to an INVITE request.
    /// This method creates that ACK and sends it to the same address as the original INVITE.
    async fn create_and_send_ack_for_response(
        &self,
        data: &Arc<ClientTransactionData>,
        response: &Response,
    ) -> Result<()> {
        let tx_id = &data.id;
        
        // Create ACK from the original INVITE
        let request_guard = data.request.lock().await;
        match utils::create_ack_from_invite(&request_guard, response) {
            Ok(ack) => {
                // Send the ACK request
                if let Err(e) = data.transport.send_message(
                    Message::Request(ack),
                    data.remote_addr
                ).await {
                    error!(id=%tx_id, error=%e, "Failed to send ACK");
                    common_logic::send_transport_error_event(tx_id, &data.events_tx).await;
                    return Err(Error::transport_error(e, "Failed to send ACK"));
                }
                Ok(())
            },
            Err(e) => {
                error!(id=%tx_id, error=%e, "Failed to create ACK for response");
                Err(e)
            }
        }
    }
    
    /// Process a SIP response
    ///
    /// This method implements the core logic for handling incoming responses
    /// according to RFC 3261 Section 17.1.1.2. It validates the response,
    /// determines the appropriate action based on the current state and the
    /// response status, and returns any state transition that should occur.
    async fn process_response(
        &self,
        data: &Arc<ClientTransactionData>,
        response: Response,
        current_state: TransactionState,
        timer_handles: &mut ClientInviteTimerHandles,
        command_tx: mpsc::Sender<InternalTransactionCommand>,
    ) -> Result<Option<TransactionState>> {
        let tx_id = &data.id;
        
        // Get the original method from the request to validate the response
        let request_guard = data.request.lock().await;
        let original_method = validators::get_method_from_request(&request_guard);
        drop(request_guard);
        
        // Validate that the response matches our transaction
        if let Err(e) = validators::validate_response_matches_transaction(&response, tx_id, &original_method) {
            warn!(id=%tx_id, error=%e, "Response validation failed");
            return Ok(None);
        }
        
        // Get status information
        let (is_provisional, is_success, is_failure) = validators::categorize_response_status(&response);
        
        match current_state {
            TransactionState::Calling => {
                // Receive any response -> cancel retransmission
                if let Some(handle) = timer_handles.timer_a.take() {
                    handle.abort();
                }
                
                if is_provisional {
                    // 1xx -> Proceeding
                    common_logic::send_provisional_response_event(tx_id, response, &data.events_tx).await;
                    return Ok(Some(TransactionState::Proceeding));
                } else if is_success {
                    // 2xx -> Terminated (RFC 3261 17.1.1.2)
                    common_logic::send_success_response_event(tx_id, response, &data.events_tx, data.remote_addr).await;
                    return Ok(Some(TransactionState::Terminated));
                } else {
                    // 3xx-6xx -> Completed
                    // Send ACK for non-2xx final response
                    if let Err(e) = self.create_and_send_ack_for_response(data, &response).await {
                        error!(id=%tx_id, error=%e, "Failed to ACK failure response");
                        return Ok(Some(TransactionState::Terminated));
                    }
                    
                    common_logic::send_failure_response_event(tx_id, response, &data.events_tx).await;
                    return Ok(Some(TransactionState::Completed));
                }
            },
            TransactionState::Proceeding => {
                if is_provisional {
                    // Additional 1xx in Proceeding state
                    common_logic::send_provisional_response_event(tx_id, response, &data.events_tx).await;
                    return Ok(None); // Stay in Proceeding
                } else if is_success {
                    // 2xx responses transition directly to Terminated (RFC 3261 17.1.1.2)
                    common_logic::send_success_response_event(tx_id, response, &data.events_tx, data.remote_addr).await;
                    return Ok(Some(TransactionState::Terminated));
                } else {
                    // 3xx-6xx transition to Completed
                    // Send ACK for non-2xx final response
                    if let Err(e) = self.create_and_send_ack_for_response(data, &response).await {
                        error!(id=%tx_id, error=%e, "Failed to ACK failure response");
                        return Ok(Some(TransactionState::Terminated));
                    }
                    
                    common_logic::send_failure_response_event(tx_id, response, &data.events_tx).await;
                    return Ok(Some(TransactionState::Completed));
                }
            },
            TransactionState::Completed => {
                if is_failure {
                    // Retransmission of final error response, resend ACK
                    debug!(id=%tx_id, "Received retransmission of error response in Completed state, resending ACK");
                    
                    // Just best effort for retransmissions; don't fail the transaction on ACK error
                    let _ = self.create_and_send_ack_for_response(data, &response).await;
                }
                // Stay in Completed state
                return Ok(None);
            },
            _ => {
                warn!(id=%tx_id, state=?current_state, "Received response in unexpected state");
                return Ok(None);
            }
        }
    }
}

#[async_trait::async_trait]
impl TransactionLogic<ClientTransactionData, ClientInviteTimerHandles> for ClientInviteLogic {
    fn kind(&self) -> TransactionKind {
        TransactionKind::InviteClient
    }

    fn initial_state(&self) -> TransactionState {
        TransactionState::Initial
    }

    fn timer_settings<'a>(data: &'a Arc<ClientTransactionData>) -> &'a TimerSettings {
        &data.timer_config
    }

    fn cancel_all_specific_timers(&self, timer_handles: &mut ClientInviteTimerHandles) {
        if let Some(handle) = timer_handles.timer_a.take() {
            handle.abort();
        }
        if let Some(handle) = timer_handles.timer_b.take() {
            handle.abort();
        }
        if let Some(handle) = timer_handles.timer_d.take() {
            handle.abort();
        }
        timer_handles.current_timer_a_interval = None;
    }

    async fn on_enter_state(
        &self,
        data: &Arc<ClientTransactionData>,
        new_state: TransactionState,
        previous_state: TransactionState,
        timer_handles: &mut ClientInviteTimerHandles,
        command_tx: mpsc::Sender<InternalTransactionCommand>, 
    ) -> Result<()> {
        let tx_id = &data.id;

        match new_state {
            TransactionState::Calling => {
                self.handle_calling_state(data, timer_handles, command_tx).await?;
            }
            TransactionState::Proceeding => {
                trace!(id=%tx_id, "Entered Proceeding state. Timer A discontinued, Timer B continues.");
                // Timer A is discontinued in Proceeding
                if let Some(handle) = timer_handles.timer_a.take() {
                    handle.abort();
                }
                // Timer B continues
            }
            TransactionState::Completed => {
                // Start Timer D (wait for response retransmissions)
                self.start_timer_d(data, timer_handles, command_tx).await;
            }
            TransactionState::Terminated => {
                trace!(id=%tx_id, "Entered Terminated state. Specific timers should have been cancelled by runner.");
                // Unregister from timer manager when terminated
                let timer_manager = self.timer_factory.timer_manager();
                timer_utils::unregister_transaction(&timer_manager, tx_id).await;
            }
            _ => { // Initial state, or others not directly part of the main flow.
                trace!(id=%tx_id, "Entered unhandled state {:?} in on_enter_state", new_state);
            }
        }
        Ok(())
    }

    async fn handle_timer(
        &self,
        data: &Arc<ClientTransactionData>,
        timer_name: &str,
        current_state: TransactionState,
        timer_handles: &mut ClientInviteTimerHandles,
    ) -> Result<Option<TransactionState>> {
        let tx_id = &data.id;
        
        if timer_name == "A" {
            // Clear the timer handle since it fired
            timer_handles.timer_a.take();
        }
        
        // Send timer triggered event using common logic
        common_logic::send_timer_triggered_event(tx_id, timer_name, &data.events_tx).await;
        
        // Use the command_tx from data to set up timers
        let self_command_tx = data.cmd_tx.clone();
        
        match timer_name {
            "A" => self.handle_timer_a_trigger(data, current_state, timer_handles, self_command_tx).await,
            "B" => self.handle_timer_b_trigger(data, current_state, self_command_tx).await,
            "D" => self.handle_timer_d_trigger(data, current_state, self_command_tx).await,
            _ => {
                warn!(id=%tx_id, timer_name=%timer_name, "Unknown timer triggered for ClientInvite");
                Ok(None)
            }
        }
    }

    async fn process_message(
        &self,
        data: &Arc<ClientTransactionData>,
        message: Message,
        current_state: TransactionState,
        timer_handles: &mut ClientInviteTimerHandles,
    ) -> Result<Option<TransactionState>> {
        let tx_id = &data.id;
        
        // Use the validators utility to extract and validate the response
        match validators::extract_response(&message, tx_id) {
            Ok(response) => {
                // Store the response
                {
                    let mut last_response = data.last_response.lock().await;
                    *last_response = Some(response.clone());
                }
                
                // Use the command_tx from data for timer operations
                let self_command_tx = data.cmd_tx.clone();
                
                // Process the response with real timer_handles
                self.process_response(data, response, current_state, timer_handles, self_command_tx).await
            },
            Err(e) => {
                warn!(id=%tx_id, error=%e, "Received non-response message");
                Ok(None)
            }
        }
    }
}

impl ClientInviteTransaction {
    /// Create a new client INVITE transaction.
    ///
    /// This method creates a new INVITE client transaction with the specified parameters.
    /// It validates that the request is an INVITE, initializes the transaction data, and
    /// spawns the transaction runner task.
    ///
    /// # Arguments
    ///
    /// * `id` - The unique identifier for this transaction
    /// * `request` - The INVITE request that initiates this transaction
    /// * `remote_addr` - The address to which the request should be sent
    /// * `transport` - The transport layer to use for sending messages
    /// * `events_tx` - The channel for sending events to the Transaction User
    /// * `timer_config_override` - Optional custom timer settings
    ///
    /// # Returns
    ///
    /// A Result containing the new ClientInviteTransaction or an error
    pub fn new(
        id: TransactionKey,
        request: Request,
        remote_addr: SocketAddr,
        transport: Arc<dyn Transport>,
        events_tx: mpsc::Sender<TransactionEvent>,
        timer_config_override: Option<TimerSettings>,
    ) -> Result<Self> {
        println!("Creating new ClientInviteTransaction: {}", id);
        
        if request.method() != Method::Invite {
            return Err(Error::Other("Request must be INVITE for INVITE client transaction".to_string()));
        }

        let timer_config = timer_config_override.unwrap_or_default();
        let (cmd_tx, local_cmd_rx) = mpsc::channel(32);

        let data = Arc::new(ClientTransactionData {
            id: id.clone(),
            state: Arc::new(AtomicTransactionState::new(TransactionState::Initial)),
            request: Arc::new(Mutex::new(request.clone())),
            last_response: Arc::new(Mutex::new(None)),
            remote_addr,
            transport,
            events_tx,
            cmd_tx: cmd_tx.clone(),
            event_loop_handle: Arc::new(Mutex::new(None)),
            timer_config: timer_config.clone(),
        });

        let logic = Arc::new(ClientInviteLogic {
            _data_marker: std::marker::PhantomData,
            timer_factory: TimerFactory::new(Some(timer_config), Arc::new(TimerManager::new(None))),
        });

        let data_for_runner = data.clone();
        let logic_for_runner = logic.clone();
        
        // Create a clone for logging in the spawn function
        let id_for_logging = id.clone();
        
        // Spawn the generic event loop runner
        let event_loop_handle = tokio::spawn(async move {
            println!("Starting event loop for INVITE Client transaction: {}", id_for_logging);
            run_transaction_loop(data_for_runner, logic_for_runner, local_cmd_rx).await;
            println!("Event loop for INVITE Client transaction ended: {}", id_for_logging);
        });

        // Store the handle for cleanup
        if let Ok(mut handle_guard) = data.event_loop_handle.try_lock() {
            *handle_guard = Some(event_loop_handle);
        }
        
        println!("Created ClientInviteTransaction: {}", id);
        
        Ok(Self { data, logic })
    }
}

impl ClientTransaction for ClientInviteTransaction {
    fn initiate(&self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
        let data = self.data.clone();
        let kind = self.kind(); // Get kind for the error message
        let tx_id = self.data.id.clone(); // Get ID for logging
        
        Box::pin(async move {
            println!("ClientInviteTransaction::initiate called for {}", tx_id);
            let current_state = data.state.get();
            println!("Current state is {:?}", current_state);
            
            if current_state != TransactionState::Initial {
                println!("Invalid state transition: {:?} -> Calling", current_state);
                return Err(Error::invalid_state_transition(
                    kind,
                    current_state,
                    TransactionState::Calling,
                    Some(data.id.clone()),
                ));
            }

            println!("Sending TransitionTo(Calling) command for {}", tx_id);
            match data.cmd_tx.send(InternalTransactionCommand::TransitionTo(TransactionState::Calling)).await {
                Ok(_) => {
                    println!("Successfully sent TransitionTo command for {}", tx_id);
                    // Wait a small amount of time to allow the transaction runner to process the command
                    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
                    
                    // Verify state change
                    let new_state = data.state.get();
                    println!("State after sending command: {:?}", new_state);
                    if new_state != TransactionState::Calling {
                        println!("WARNING: State didn't change to Calling, still: {:?}", new_state);
                    }
                    
                    Ok(())
                },
                Err(e) => {
                    println!("Failed to send command: {}", e);
                    Err(Error::Other(format!("Failed to send command: {}", e)))
                }
            }
        })
    }

    fn process_response(&self, response: Response) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
        let data = self.data.clone();
        Box::pin(async move {
            trace!(id=%data.id, method=%response.status(), "Received response");
            
            data.cmd_tx.send(InternalTransactionCommand::ProcessMessage(Message::Response(response))).await
                .map_err(|e| Error::Other(format!("Failed to send command: {}", e)))?;
            
            Ok(())
        })
    }

    // Implement the original_request method
    fn original_request(&self) -> Pin<Box<dyn Future<Output = Option<Request>> + Send + '_>> {
        let request_arc = self.data.request.clone();
        Box::pin(async move {
            let req = request_arc.lock().await;
            Some(req.clone()) // Clone the request out of the Mutex guard
        })
    }

    // Add the required last_response implementation for ClientTransaction
    fn last_response<'a>(&'a self) -> Pin<Box<dyn Future<Output = Option<Response>> + Send + 'a>> {
        // Create a future that just returns the last response
        let last_response = self.data.last_response.clone();
        Box::pin(async move {
            last_response.lock().await.clone()
        })
    }
}

impl Transaction for ClientInviteTransaction {
    fn id(&self) -> &TransactionKey {
        &self.data.id
    }

    fn kind(&self) -> TransactionKind {
        TransactionKind::InviteClient
    }

    fn state(&self) -> TransactionState {
        self.data.state.get()
    }
    
    fn remote_addr(&self) -> SocketAddr {
        self.data.remote_addr
    }
    
    fn matches(&self, message: &Message) -> bool {
        // Follow the same approach as in non_invite.rs, checking:
        // 1. Top Via branch parameter matches transaction ID's branch
        // 2. CSeq method matches original request's method
        if !message.is_response() { return false; }
        
        let response = match message {
            Message::Response(r) => r,
            _ => return false,
        };

        // Check Via headers
        if let Some(TypedHeader::Via(via_header_vec)) = response.header(&HeaderName::Via) {
            if let Some(via_header) = via_header_vec.0.first() {
                if via_header.branch() != Some(self.data.id.branch()) {
                    return false;
                }
            } else {
                return false; // No Via value in the vector
            }
        } else {
            return false; // No Via header or not of TypedHeader::Via type
        }

        // Check CSeq method
        let original_request_method = self.data.id.method().clone();
        if let Some(TypedHeader::CSeq(cseq_header)) = response.header(&HeaderName::CSeq) {
            if cseq_header.method != original_request_method {
                return false;
            }
        } else {
            return false; // No CSeq header or not of TypedHeader::CSeq type
        }
        
        true
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}

impl TransactionAsync for ClientInviteTransaction {
    fn process_event<'a>(
        &'a self,
        event_type: &'a str, // e.g. "response" from TransactionManager when it routes a message
        message: Option<Message>
    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
        Box::pin(async move {
            match event_type {
                "response" => {
                    if let Some(msg) = message {
                        self.data.cmd_tx.send(InternalTransactionCommand::ProcessMessage(msg)).await
                            .map_err(|e| Error::Other(format!("Failed to send ProcessMessage command: {}", e)))?;
                    } else {
                        return Err(Error::Other("Expected Message for 'response' event type".to_string()));
                    }
                },
                _ => return Err(Error::Other(format!("Unhandled event type in TransactionAsync::process_event: {}", event_type))),
            }
            Ok(())
        })
    }

    fn send_command<'a>(
        &'a self,
        cmd: InternalTransactionCommand
    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
        let cmd_tx = self.data.cmd_tx.clone();
        Box::pin(async move {
            cmd_tx.send(cmd).await
                .map_err(|e| Error::Other(format!("Failed to send command via TransactionAsync: {}", e)))
        })
    }

    fn original_request<'a>(
        &'a self
    ) -> Pin<Box<dyn Future<Output = Option<Request>> + Send + 'a>> {
        let request_mutex = self.data.request.clone();
        Box::pin(async move {
            Some(request_mutex.lock().await.clone())
        })
    }

    fn last_response<'a>(
        &'a self
    ) -> Pin<Box<dyn Future<Output = Option<Response>> + Send + 'a>> {
        let response_mutex = self.data.last_response.clone();
        Box::pin(async move {
            response_mutex.lock().await.clone()
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::transaction::runner::{AsRefState, AsRefKey, HasTransactionEvents, HasTransport, HasCommandSender};
    use rvoip_sip_core::builder::{SimpleRequestBuilder, SimpleResponseBuilder};
    use rvoip_sip_core::types::status::StatusCode;
    use rvoip_sip_core::Response as SipCoreResponse;
    use std::collections::VecDeque;
    use std::str::FromStr;
    use tokio::sync::Notify;
    use tokio::time::timeout as TokioTimeout;

    // A simple mock transport for these unit tests
    #[derive(Debug, Clone)]
    struct UnitTestMockTransport {
        sent_messages: Arc<Mutex<VecDeque<(Message, SocketAddr)>>>,
        local_addr: SocketAddr,
        // Notifier for when a message is sent, to help synchronize tests
        message_sent_notifier: Arc<Notify>,
    }

    impl UnitTestMockTransport {
        fn new(local_addr_str: &str) -> Self {
            Self {
                sent_messages: Arc::new(Mutex::new(VecDeque::new())),
                local_addr: SocketAddr::from_str(local_addr_str).unwrap(),
                message_sent_notifier: Arc::new(Notify::new()),
            }
        }

        async fn get_sent_message(&self) -> Option<(Message, SocketAddr)> {
            self.sent_messages.lock().await.pop_front()
        }

        async fn wait_for_message_sent(&self, duration: Duration) -> std::result::Result<(), tokio::time::error::Elapsed> {
            TokioTimeout(duration, self.message_sent_notifier.notified()).await
        }
    }

    #[async_trait::async_trait]
    impl Transport for UnitTestMockTransport {
        fn local_addr(&self) -> std::result::Result<SocketAddr, rvoip_sip_transport::Error> {
            Ok(self.local_addr)
        }

        async fn send_message(&self, message: Message, destination: SocketAddr) -> std::result::Result<(), rvoip_sip_transport::Error> {
            self.sent_messages.lock().await.push_back((message.clone(), destination));
            self.message_sent_notifier.notify_one(); // Notify that a message was "sent"
            Ok(())
        }

        async fn close(&self) -> std::result::Result<(), rvoip_sip_transport::Error> {
            Ok(())
        }

        fn is_closed(&self) -> bool {
            false
        }
    }

    struct TestSetup {
        transaction: ClientInviteTransaction,
        mock_transport: Arc<UnitTestMockTransport>,
        tu_events_rx: mpsc::Receiver<TransactionEvent>,
    }

    async fn setup_test_environment(target_uri_str: &str) -> TestSetup {
        let local_addr = "127.0.0.1:5090";
        let mock_transport = Arc::new(UnitTestMockTransport::new(local_addr));
        let (tu_events_tx, tu_events_rx) = mpsc::channel(100);

        let req_uri = Uri::from_str(target_uri_str).unwrap();
        let builder = SimpleRequestBuilder::new(Method::Invite, &req_uri.to_string())
            .expect("Failed to create SimpleRequestBuilder")
            .from("Alice", "sip:test@test.com", Some("fromtag"))
            .to("Bob", "sip:bob@target.com", None)
            .call_id("callid-invite-test")
            .cseq(1);
        
        let via_branch = format!("z9hG4bK.{}", uuid::Uuid::new_v4().as_simple());
        let builder = builder.via(mock_transport.local_addr.to_string().as_str(), "UDP", Some(&via_branch));

        let request = builder.build();
        
        let remote_addr = SocketAddr::from_str("127.0.0.1:5070").unwrap();
        let tx_key = TransactionKey::from_request(&request).expect("Failed to create tx key from request");

        let settings = TimerSettings {
            t1: Duration::from_millis(50),
            transaction_timeout: Duration::from_millis(200),
            wait_time_d: Duration::from_millis(100),
            ..Default::default()
        };

        let transaction = ClientInviteTransaction::new(
            tx_key,
            request,
            remote_addr,
            mock_transport.clone() as Arc<dyn Transport>,
            tu_events_tx,
            Some(settings),
        ).unwrap();

        TestSetup {
            transaction,
            mock_transport,
            tu_events_rx,
        }
    }
    
    fn build_simple_response(status_code: StatusCode, original_request: &Request) -> SipCoreResponse {
        let response_builder = SimpleResponseBuilder::response_from_request(
            original_request,
            status_code,
            Some(status_code.reason_phrase())
        );
        
        let response_builder = if original_request.to().unwrap().tag().is_none() {
             response_builder.to(
                original_request.to().unwrap().address().display_name().unwrap_or_default(),
                &original_request.to().unwrap().address().uri().to_string(),
                Some("totag-server")
            )
        } else {
            response_builder
        };
        
        response_builder.build()
    }

    #[tokio::test]
    async fn test_invite_client_creation_and_initial_state() {
        let setup = setup_test_environment("sip:bob@target.com").await;
        assert_eq!(setup.transaction.state(), TransactionState::Initial);
        assert!(setup.transaction.data.event_loop_handle.lock().await.is_some());
    }

    #[tokio::test]
    async fn test_invite_client_initiate_sends_request_and_starts_timers() {
        let mut setup = setup_test_environment("sip:bob@target.com").await;
        
        setup.transaction.initiate().await.expect("initiate should succeed");

        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.expect("Message should be sent quickly");

        tokio::time::sleep(Duration::from_millis(20)).await;
        assert_eq!(setup.transaction.state(), TransactionState::Calling, "State should be Calling after initiate");

        let sent_msg_info = setup.mock_transport.get_sent_message().await;
        assert!(sent_msg_info.is_some(), "Request should have been sent");
        if let Some((msg, dest)) = sent_msg_info {
            assert!(msg.is_request());
            assert_eq!(msg.method(), Some(Method::Invite));
            assert_eq!(dest, setup.transaction.remote_addr());
        }
        
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.expect("Timer A retransmission failed to occur");
        let retransmitted_msg_info = setup.mock_transport.get_sent_message().await;
        assert!(retransmitted_msg_info.is_some(), "Request should have been retransmitted by Timer A");
        if let Some((msg, _)) = retransmitted_msg_info {
            assert!(msg.is_request());
            assert_eq!(msg.method(), Some(Method::Invite));
        }
    }

    #[tokio::test]
    async fn test_invite_client_provisional_response() {
        let mut setup = setup_test_environment("sip:bob@target.com").await;
        setup.transaction.initiate().await.expect("initiate failed");
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        setup.mock_transport.get_sent_message().await;

        // Wait for and ignore the StateChanged event
        match TokioTimeout(Duration::from_millis(100), setup.tu_events_rx.recv()).await {
            Ok(Some(TransactionEvent::StateChanged { .. })) => {
                // Expected StateChanged event, continue
            },
            Ok(Some(other_event)) => panic!("Unexpected first event: {:?}", other_event),
            _ => panic!("Expected StateChanged event"),
        }

        tokio::time::sleep(Duration::from_millis(20)).await;
        assert_eq!(setup.transaction.state(), TransactionState::Calling);

        let original_request_clone = setup.transaction.data.request.lock().await.clone();
        let prov_response = build_simple_response(StatusCode::Ringing, &original_request_clone);
        
        setup.transaction.process_response(prov_response.clone()).await.expect("process_response failed");

        // Wait for ProvisionalResponse event
        match TokioTimeout(Duration::from_millis(100), setup.tu_events_rx.recv()).await {
            Ok(Some(TransactionEvent::ProvisionalResponse { transaction_id, response, .. })) => {
                assert_eq!(transaction_id, *setup.transaction.id());
                assert_eq!(response.status_code(), StatusCode::Ringing.as_u16());
            },
            Ok(Some(other_event)) => panic!("Unexpected event: {:?}", other_event),
            Ok(None) => panic!("Event channel closed"),
            Err(_) => panic!("Timeout waiting for ProvisionalResponse event"),
        }
        
        // Check for StateChanged from Calling to Proceeding
        match TokioTimeout(Duration::from_millis(100), setup.tu_events_rx.recv()).await {
            Ok(Some(TransactionEvent::StateChanged { transaction_id, previous_state, new_state })) => {
                assert_eq!(transaction_id, *setup.transaction.id());
                assert_eq!(previous_state, TransactionState::Calling);
                assert_eq!(new_state, TransactionState::Proceeding);
            },
            Ok(Some(other_event)) => panic!("Unexpected event: {:?}", other_event),
            _ => panic!("Expected StateChanged event"),
        }
        
        tokio::time::sleep(Duration::from_millis(20)).await;
        assert_eq!(setup.transaction.state(), TransactionState::Proceeding, "State should be Proceeding");
    }

    #[tokio::test]
    async fn test_invite_client_success_response() {
        let mut setup = setup_test_environment("sip:bob@target.com").await;
        setup.transaction.initiate().await.expect("initiate failed");
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        setup.mock_transport.get_sent_message().await;

        // Wait for and ignore the StateChanged event
        match TokioTimeout(Duration::from_millis(100), setup.tu_events_rx.recv()).await {
            Ok(Some(TransactionEvent::StateChanged { .. })) => {
                // Expected StateChanged event, continue
            },
            Ok(Some(other_event)) => panic!("Unexpected first event: {:?}", other_event),
            _ => panic!("Expected StateChanged event"),
        }

        let original_request_clone = setup.transaction.data.request.lock().await.clone();
        let success_response = build_simple_response(StatusCode::Ok, &original_request_clone);
        
        setup.transaction.process_response(success_response.clone()).await.expect("process_response failed");

        // Success response should go directly to Terminated in INVITE
        let mut success_response_received = false;
        let mut calling_to_terminated_received = false;
        let mut transaction_terminated_received = false;

        // Collect all events until we get terminated or timeout
        for _ in 0..5 {  // Give it 5 iterations max
            match TokioTimeout(Duration::from_millis(150), setup.tu_events_rx.recv()).await {
                Ok(Some(TransactionEvent::SuccessResponse { transaction_id, response, .. })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    assert_eq!(response.status_code(), StatusCode::Ok.as_u16());
                    success_response_received = true;
                },
                Ok(Some(TransactionEvent::StateChanged { transaction_id, previous_state, new_state })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    if previous_state == TransactionState::Calling && new_state == TransactionState::Terminated {
                        calling_to_terminated_received = true;
                    } else {
                        panic!("Unexpected state transition: {:?} -> {:?}", previous_state, new_state);
                    }
                },
                Ok(Some(TransactionEvent::TransactionTerminated { transaction_id, .. })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    transaction_terminated_received = true;
                    break;  // We got the terminal event, can stop waiting
                },
                Ok(Some(TransactionEvent::TimerTriggered { .. })) => {
                    // Timer events can happen, ignore them
                    continue;
                },
                Ok(Some(other_event)) => panic!("Unexpected event: {:?}", other_event),
                Ok(None) => panic!("Event channel closed"),
                Err(_) => {
                    // If we timed out but already got the necessary events, we're good
                    if success_response_received && calling_to_terminated_received {
                        break;
                    }
                    continue;
                }
            }
            
            // If we got all the necessary events, we can stop waiting
            if success_response_received && calling_to_terminated_received && transaction_terminated_received {
                break;
            }
        }

        // Check that we got all the expected events
        assert!(success_response_received, "SuccessResponse event not received");
        assert!(calling_to_terminated_received, "StateChanged Calling->Terminated event not received");
        
        // The transaction should already be in Terminated state
        tokio::time::sleep(Duration::from_millis(20)).await;
        assert_eq!(setup.transaction.state(), TransactionState::Terminated, "State should be Terminated after 2xx response");
    }
    
    #[tokio::test]
    async fn test_invite_client_failure_response_and_ack() {
        let mut setup = setup_test_environment("sip:bob@target.com").await;
        setup.transaction.initiate().await.expect("initiate failed");
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        setup.mock_transport.get_sent_message().await;

        // Wait for and ignore the StateChanged event
        match TokioTimeout(Duration::from_millis(100), setup.tu_events_rx.recv()).await {
            Ok(Some(TransactionEvent::StateChanged { .. })) => {
                // Expected StateChanged event, continue
            },
            Ok(Some(other_event)) => panic!("Unexpected first event: {:?}", other_event),
            _ => panic!("Expected StateChanged event"),
        }

        let original_request_clone = setup.transaction.data.request.lock().await.clone();
        let failure_response = build_simple_response(StatusCode::NotFound, &original_request_clone);
        
        setup.transaction.process_response(failure_response.clone()).await.expect("process_response failed");

        // First check for ACK being sent
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        let sent_msg_info = setup.mock_transport.get_sent_message().await;
        assert!(sent_msg_info.is_some(), "ACK request should have been sent");
        if let Some((msg, dest)) = sent_msg_info {
            assert!(msg.is_request());
            assert_eq!(msg.method(), Some(Method::Ack));
            assert_eq!(dest, setup.transaction.remote_addr());
        }

        // Now check events
        let mut failure_response_received = false;
        let mut calling_to_completed_received = false;
        let mut completed_to_terminated_received = false;
        let mut transaction_terminated_received = false;

        // Collect all events until we get terminated or timeout
        for _ in 0..8 {  // More iterations because more events to process
            match TokioTimeout(Duration::from_millis(150), setup.tu_events_rx.recv()).await {
                Ok(Some(TransactionEvent::FailureResponse { transaction_id, response, .. })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    assert_eq!(response.status_code(), StatusCode::NotFound.as_u16());
                    failure_response_received = true;
                },
                Ok(Some(TransactionEvent::StateChanged { transaction_id, previous_state, new_state })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    if previous_state == TransactionState::Calling && new_state == TransactionState::Completed {
                        calling_to_completed_received = true;
                    } else if previous_state == TransactionState::Completed && new_state == TransactionState::Terminated {
                        completed_to_terminated_received = true;
                    } else {
                        panic!("Unexpected state transition: {:?} -> {:?}", previous_state, new_state);
                    }
                },
                Ok(Some(TransactionEvent::TransactionTerminated { transaction_id, .. })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    transaction_terminated_received = true;
                    break;  // We got the terminal event, can stop waiting
                },
                Ok(Some(TransactionEvent::TimerTriggered { .. })) => {
                    // Timer events can happen, ignore them
                    continue;
                },
                Ok(Some(other_event)) => panic!("Unexpected event: {:?}", other_event),
                Ok(None) => panic!("Event channel closed"),
                Err(_) => {
                    // If we timed out but already got the necessary events, we're good
                    if failure_response_received && calling_to_completed_received && 
                       (completed_to_terminated_received || transaction_terminated_received) {
                        break;
                    }
                    continue;
                }
            }
            
            // If we got all the necessary events, we can stop waiting
            if failure_response_received && calling_to_completed_received && 
               completed_to_terminated_received && transaction_terminated_received {
                break;
            }
        }

        // Check that we got all the expected events
        assert!(failure_response_received, "FailureResponse event not received");
        assert!(calling_to_completed_received, "StateChanged Calling->Completed event not received");
        
        // Eventually the transaction should transition to Terminated via Timer D
        tokio::time::sleep(Duration::from_millis(150)).await;
        assert_eq!(setup.transaction.state(), TransactionState::Terminated, "State should be Terminated after Timer D");
    }
    
    #[tokio::test]
    async fn test_invite_client_timer_b_timeout() {
        let mut setup = setup_test_environment("sip:bob@target.com").await;
        setup.transaction.initiate().await.expect("initiate failed");

        // Wait for and ignore the StateChanged event
        match TokioTimeout(Duration::from_millis(100), setup.tu_events_rx.recv()).await {
            Ok(Some(TransactionEvent::StateChanged { .. })) => {
                // Expected StateChanged event, continue
            },
            Ok(Some(other_event)) => panic!("Unexpected first event: {:?}", other_event),
            _ => panic!("Expected StateChanged event"),
        }

        let mut timeout_event_received = false;
        let mut terminated_event_received = false;
        let mut timer_b_received = false;

        // Loop to catch multiple events, specifically TimerTriggered for A/B, then TransactionTimeout, then TransactionTerminated.
        for _ in 0..8 { 
            match TokioTimeout(Duration::from_millis(150), setup.tu_events_rx.recv()).await { 
                Ok(Some(TransactionEvent::TransactionTimeout { transaction_id, .. })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    timeout_event_received = true;
                },
                Ok(Some(TransactionEvent::TransactionTerminated { transaction_id, .. })) => {
                    assert_eq!(transaction_id, *setup.transaction.id());
                    terminated_event_received = true;
                },
                Ok(Some(TransactionEvent::TimerTriggered { ref timer, .. })) => { 
                    if timer == "A" { 
                        debug!("Timer A triggered during B timeout test, continuing...");
                        continue; 
                    } else if timer == "B" {
                        timer_b_received = true;
                        continue;
                    }
                    panic!("Unexpected TimerTriggered event: {:?}", timer);
                },
                Ok(Some(TransactionEvent::StateChanged { .. })) => {
                    // State transitions can happen, ignore them
                    continue;
                },
                Ok(Some(other_event)) => {
                    panic!("Unexpected event: {:?}", other_event);
                },
                Ok(None) => panic!("Event channel closed prematurely"),
                Err(_) => { 
                    if !timeout_event_received || !terminated_event_received {
                        debug!("TokioTimeout while waiting for B events, may be normal if timers are still running");
                    } else {
                        break;
                    }
                }
            }
            if timeout_event_received && terminated_event_received { break; }
        }
        
        assert!(timeout_event_received, "TransactionTimeout event not received");
        assert!(terminated_event_received, "TransactionTerminated event not received");
        
        tokio::time::sleep(Duration::from_millis(20)).await;
        assert_eq!(setup.transaction.state(), TransactionState::Terminated, "State should be Terminated after Timer B");
    }
    
    #[tokio::test]
    async fn test_invite_client_retransmit_request() {
        let mut setup = setup_test_environment("sip:bob@target.com").await;
        setup.transaction.initiate().await.expect("initiate failed");
        
        // Clear the initial request
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        setup.mock_transport.get_sent_message().await;
        
        // Now we should get a retransmission due to Timer A
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        let msg_info = setup.mock_transport.get_sent_message().await;
        assert!(msg_info.is_some(), "Request should have been retransmitted");
        if let Some((msg, _)) = msg_info {
            assert!(msg.is_request());
            assert_eq!(msg.method(), Some(Method::Invite));
        }
        
        // And we should get a second retransmission with larger interval
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(200)).await.unwrap();
        let msg_info2 = setup.mock_transport.get_sent_message().await;
        assert!(msg_info2.is_some(), "Request should have been retransmitted a second time");
        if let Some((msg, _)) = msg_info2 {
            assert!(msg.is_request());
            assert_eq!(msg.method(), Some(Method::Invite));
        }
    }
    
    #[tokio::test]
    async fn test_invite_client_retransmit_response_ack() {
        let mut setup = setup_test_environment("sip:bob@target.com").await;
        setup.transaction.initiate().await.expect("initiate failed");
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        let _ = setup.mock_transport.get_sent_message().await;

        // Wait for initial StateChanged event
        match TokioTimeout(Duration::from_millis(100), setup.tu_events_rx.recv()).await {
            Ok(Some(TransactionEvent::StateChanged { .. })) => {},
            _ => panic!("Expected StateChanged event"),
        }

        // Send a failure response to get to Completed state
        let original_request_clone = setup.transaction.data.request.lock().await.clone();
        let failure_response = build_simple_response(StatusCode::NotFound, &original_request_clone);
        setup.transaction.process_response(failure_response.clone()).await.expect("process_response failed");
        
        // Wait for the ACK
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        let _ = setup.mock_transport.get_sent_message().await;
        
        // Now let's simulate a retransmitted response from the server
        setup.transaction.process_response(failure_response.clone()).await.expect("process_response for retransmit failed");
        
        // We should see another ACK for this retransmitted response
        setup.mock_transport.wait_for_message_sent(Duration::from_millis(100)).await.unwrap();
        let msg_info = setup.mock_transport.get_sent_message().await;
        assert!(msg_info.is_some(), "ACK should have been sent for retransmitted response");
        if let Some((msg, _)) = msg_info {
            assert!(msg.is_request());
            assert_eq!(msg.method(), Some(Method::Ack));
        }
    }
}