rialo_types/

websocket_op.rs

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

// ============================================================================
// WebSocket Operations
// ============================================================================

// ============================================================================
// WebSocket Types
// ============================================================================

use std::fmt;
#[cfg(feature = "non-pdk")]
use std::str::FromStr;

use borsh::{BorshDeserialize, BorshSerialize};
#[cfg(feature = "non-pdk")]
use clap::Subcommand;
use serde::{Deserialize, Serialize};

use crate::rex_info::{RexId, RexUrl, RexValue};

/// Estimated size in bytes for system messages (OversizedWarning).
/// This accounts for the struct fields and some padding.
pub const SYSTEM_MESSAGE_SIZE: usize = 128;

/// Read mode for WebSocket buffer.
///
/// Determines how messages are retrieved from the buffer.
#[derive(
    Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize,
)]
pub enum WebSocketReadMode {
    /// Return only the most recent message (default, backward compatible).
    #[default]
    Latest,
    /// Return all accumulated messages.
    All,
    /// Return messages newer than the given index.
    ///
    /// **Note:** Index wraparound at `u64::MAX` is not handled. See DataBuffer
    /// documentation for limitations. At 1 million messages per second,
    /// wraparound would take ~584,000 years.
    FromIndex(u64),
}

/// Message content type - either user data or system notification.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub enum MessageContent {
    /// Normal message data from WebSocket.
    Data(Vec<u8>),
    /// System-generated message for oversized content that was rejected.
    OversizedWarning {
        /// Size of the rejected message in bytes.
        original_size: usize,
        /// Buffer size limit in bytes.
        limit: usize,
        /// When the oversized message was received (RFC3339 with microseconds).
        original_timestamp: String,
    },
}

impl MessageContent {
    /// Returns the byte size of this content for buffer accounting.
    ///
    /// For data messages, this is the actual data length.
    /// For system messages, this is a fixed size constant.
    pub fn byte_size(&self) -> usize {
        match self {
            MessageContent::Data(data) => data.len(),
            MessageContent::OversizedWarning { .. } => SYSTEM_MESSAGE_SIZE,
        }
    }

    /// Returns true if this is a system message.
    pub fn is_system(&self) -> bool {
        matches!(self, MessageContent::OversizedWarning { .. })
    }

    /// Returns the data bytes if this is a Data message, None otherwise.
    pub fn as_data(&self) -> Option<&[u8]> {
        match self {
            MessageContent::Data(data) => Some(data),
            MessageContent::OversizedWarning { .. } => None,
        }
    }
}

/// A single buffered message with metadata.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct BufferedMessage {
    /// Unique, monotonically increasing index (never reused, persists across reconnects).
    pub index: u64,
    /// The message content.
    pub content: MessageContent,
    /// When this message was received/created (RFC3339 with microseconds).
    pub received_at: String,
}

impl BufferedMessage {
    /// Returns the size of this message in bytes (for buffer accounting).
    pub fn byte_size(&self) -> usize {
        self.content.byte_size()
    }

    /// Returns true if this is a system message.
    pub fn is_system_message(&self) -> bool {
        self.content.is_system()
    }

    /// Returns the data bytes if this contains a Data message, None otherwise.
    pub fn as_data(&self) -> Option<&[u8]> {
        self.content.as_data()
    }
}

/// Response for WebSocket read operations.
///
/// This struct captures the read result with messages and buffer state information.
///
/// # Gap Detection
///
/// Consumers can detect gaps by comparing message indices. Each `BufferedMessage`
/// contains an `index` field that increases monotonically. If there's a gap between
/// indices, some messages were evicted from the buffer.
///
/// # Convenience Methods
///
/// For common access patterns, use the convenience methods instead of accessing
/// fields directly:
///
/// ```ignore
/// // Instead of:
/// let data = response.messages.first().unwrap().content.as_data().unwrap();
///
/// // Use:
/// let data = response.first_data().unwrap();
/// ```
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
pub struct WebSocketReadResponse {
    /// All requested messages (based on read mode).
    pub messages: Vec<BufferedMessage>,
    /// Current highest index in buffer (for tracking).
    pub latest_index: Option<u64>,
    /// Current lowest index in buffer (helpful for gap detection).
    pub oldest_index: Option<u64>,
}

impl WebSocketReadResponse {
    /// Returns true if there are any messages in the response.
    pub fn has_messages(&self) -> bool {
        !self.messages.is_empty()
    }

    /// Returns the number of messages in the response.
    pub fn message_count(&self) -> usize {
        self.messages.len()
    }

    /// Get the first (oldest) message, if any.
    pub fn first_message(&self) -> Option<&BufferedMessage> {
        self.messages.first()
    }

    /// Get the last (most recent) message, if any.
    pub fn latest_message(&self) -> Option<&BufferedMessage> {
        self.messages.last()
    }

    /// Get the data from the first message (convenience for common pattern).
    ///
    /// Returns `None` if there are no messages or if the first message is a system message.
    pub fn first_data(&self) -> Option<&[u8]> {
        self.messages.first().and_then(|m| m.as_data())
    }

    /// Get the data from the latest message (convenience for common pattern).
    ///
    /// Returns `None` if there are no messages or if the latest message is a system message.
    pub fn latest_data(&self) -> Option<&[u8]> {
        self.messages.last().and_then(|m| m.as_data())
    }

    /// Iterate over only the data messages (filtering out system messages).
    pub fn data_messages(&self) -> impl Iterator<Item = &BufferedMessage> {
        self.messages.iter().filter(|m| !m.is_system_message())
    }

    /// Iterate over only the data bytes (convenience for processing).
    ///
    /// This filters out system messages and returns only the raw data bytes.
    pub fn iter_data(&self) -> impl Iterator<Item = &[u8]> {
        self.messages.iter().filter_map(|m| m.as_data())
    }

    /// Returns the number of data messages (excluding system messages).
    pub fn data_message_count(&self) -> usize {
        self.messages
            .iter()
            .filter(|m| !m.is_system_message())
            .count()
    }

    /// Returns true if any message is a system message (like OversizedWarning).
    pub fn has_system_messages(&self) -> bool {
        self.messages.iter().any(|m| m.is_system_message())
    }
}

/// WebSocket operation types for the WebSocket REX.
///
/// This enum defines the different operations that can be performed on WebSocket connections.
/// All WebSocket operations are handled through a single `TargetRexProgram::WebSocket` variant.
///
/// # Variants
/// * `Connect` - Establish a new WebSocket connection to an external server
/// * `Read` - Read the latest data from an existing WebSocket connection
///
/// # Future Extensions
/// Additional operations like `Send` and `Close` may be added in the future.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
#[cfg_attr(feature = "non-pdk", derive(Subcommand))]
pub enum WebSocketOperation {
    /// Establish a new WebSocket connection.
    ///
    /// When this operation is executed:
    /// - A validator is randomly selected to handle the connection
    /// - The TEE establishes the WebSocket connection to the specified URL
    /// - On success, the connection is registered in the WebSocket Registry
    /// - The assigned validator is recorded for future routing
    Connect {
        /// WebSocket URL (typically wss://...)
        #[cfg_attr(feature = "non-pdk", clap(
            long = "target-url",
            value_parser = clap::value_parser!(RexUrl)
        ))]
        url: RexUrl,
        /// The RexId that represents this connection request
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_rex_id))]
        rex_id: RexId,
    },
    /// Read data from an existing WebSocket connection.
    ///
    /// When this operation is executed:
    /// - The operation is automatically routed to the TEE holding the connection
    /// - Returns messages from the connection's buffer based on the read mode
    /// - The referenced connection must exist and be active
    ///
    /// # Read Modes
    ///
    /// - `Latest` (default): Returns only the most recent message
    /// - `All`: Returns all accumulated messages in the buffer
    /// - `FromIndex(n)`: Returns all messages with index > n, with gap detection
    ///
    /// # Validation
    ///
    /// The handler implementation must verify that:
    /// - `connection_rex_id` references an existing Connect operation in the WebSocket Registry
    /// - The connection is still active and hasn't been closed
    /// - The requesting validator has permission to read from this connection
    ///
    /// # Error Handling
    ///
    /// If validation fails, the REX should return an appropriate error indicating:
    /// - `ConnectionNotFound` - if the referenced REX ID doesn't exist
    /// - `ConnectionClosed` - if the connection has been terminated
    /// - `InvalidConnectionType` - if the referenced REX is not a WebSocket Connect operation
    Read {
        /// References the RexId of the Connect operation that established the connection
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_rex_id))]
        connection_rex_id: RexId,
        /// The read mode determining which messages to return (defaults to Latest for backward compatibility)
        #[serde(default)]
        #[cfg_attr(feature = "non-pdk", clap(skip))]
        mode: WebSocketReadMode,
    },
    /// Send messages to an existing WebSocket connection.
    ///
    /// When this operation is executed:
    /// - The operation is automatically routed to the TEE holding the connection
    /// - Each `RexValue` in `messages` is sent as a separate WebSocket message
    /// - The inner value is extracted: `Plain(String)` → sent as-is, `Encrypted(String)` → decrypted then sent
    /// - Returns the number of messages successfully sent
    /// - The referenced connection must exist and be active
    ///
    /// # Validation
    ///
    /// The handler implementation must verify that:
    /// - `connection_rex_id` references an existing Connect operation in the WebSocket Registry
    /// - The connection is still active and hasn't been closed
    /// - The requesting validator has permission to send to this connection
    ///
    /// # Error Handling
    ///
    /// If validation fails, the REX should return an appropriate error indicating:
    /// - `ConnectionNotFound` - if the referenced REX ID doesn't exist
    /// - `ConnectionClosed` - if the connection has been terminated
    /// - `InvalidConnectionType` - if the referenced REX is not a WebSocket Connect operation
    /// - `SecretDecryptionFailed` - if an encrypted message cannot be decrypted
    Send {
        /// References the RexId of the Connect operation that established the connection
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_rex_id))]
        connection_rex_id: RexId,
        /// Messages to send to the WebSocket connection
        /// Each RexValue will be sent as a separate WebSocket message
        #[cfg_attr(feature = "non-pdk", clap(skip))]
        messages: Vec<RexValue>,
    },
    /// Close an existing WebSocket connection.
    ///
    /// When this operation is executed:
    /// - The operation is automatically routed to the TEE holding the connection
    /// - The WebSocket connection is gracefully closed
    /// - The connection status in the WebSocket Registry is updated to `Closed`
    /// - The referenced connection must exist and be active
    ///
    /// # Validation
    ///
    /// The handler implementation must verify that:
    /// - `connection_rex_id` references an existing Connect operation in the WebSocket Registry
    /// - The requesting validator has permission to close this connection
    ///
    /// # Error Handling
    ///
    /// If validation fails, the REX should return an appropriate error indicating:
    /// - `ConnectionNotFound` - if the referenced REX ID doesn't exist
    /// - `ConnectionAlreadyClosed` - if the connection has already been terminated
    /// - `InvalidConnectionType` - if the referenced REX is not a WebSocket Connect operation
    Close {
        /// References the RexId of the Connect operation to close
        #[cfg_attr(feature = "non-pdk", clap(long, value_parser = parse_rex_id))]
        connection_rex_id: RexId,
    },
}

/// Parse a RexId from a string (for clap CLI parsing)
#[cfg(feature = "non-pdk")]
fn parse_rex_id(s: &str) -> Result<RexId, String> {
    RexId::from_str(s)
}

impl WebSocketOperation {
    /// Create a new Connect operation.
    pub fn connect(url: impl Into<RexUrl>, rex_id: RexId) -> Self {
        WebSocketOperation::Connect {
            url: url.into(),
            rex_id,
        }
    }

    /// Create a new Read operation with default Latest mode.
    pub fn read(connection_rex_id: RexId) -> Self {
        WebSocketOperation::Read {
            connection_rex_id,
            mode: WebSocketReadMode::default(),
        }
    }

    /// Create a new Read operation with a specific mode.
    pub fn read_with_mode(connection_rex_id: RexId, mode: WebSocketReadMode) -> Self {
        WebSocketOperation::Read {
            connection_rex_id,
            mode,
        }
    }

    /// Create a new Send operation.
    pub fn send(connection_rex_id: RexId, messages: Vec<RexValue>) -> Self {
        WebSocketOperation::Send {
            connection_rex_id,
            messages,
        }
    }

    /// Create a new Close operation.
    pub fn close(connection_rex_id: RexId) -> Self {
        WebSocketOperation::Close { connection_rex_id }
    }

    /// Returns true if this is a Connect operation.
    pub fn is_connect(&self) -> bool {
        matches!(self, WebSocketOperation::Connect { .. })
    }

    /// Returns true if this is a Read operation.
    pub fn is_read(&self) -> bool {
        matches!(self, WebSocketOperation::Read { .. })
    }

    /// Returns true if this is a Send operation.
    pub fn is_send(&self) -> bool {
        matches!(self, WebSocketOperation::Send { .. })
    }

    /// Returns true if this is a Close operation.
    pub fn is_close(&self) -> bool {
        matches!(self, WebSocketOperation::Close { .. })
    }

    /// Returns the URL if this is a Connect operation, None otherwise.
    pub fn url(&self) -> Option<&RexUrl> {
        match self {
            WebSocketOperation::Connect { url, .. } => Some(url),
            _ => None,
        }
    }

    /// Returns the rex_id if this is a Connect operation, None otherwise.
    pub fn rex_id(&self) -> Option<&RexId> {
        match self {
            WebSocketOperation::Connect { rex_id, .. } => Some(rex_id),
            _ => None,
        }
    }

    /// Returns the connection_rex_id if this is a Read, Send, or Close operation, None otherwise.
    pub fn connection_rex_id(&self) -> Option<&RexId> {
        match self {
            WebSocketOperation::Read {
                connection_rex_id, ..
            }
            | WebSocketOperation::Send {
                connection_rex_id, ..
            }
            | WebSocketOperation::Close { connection_rex_id } => Some(connection_rex_id),
            _ => None,
        }
    }

    /// Returns the read mode if this is a Read operation, None otherwise.
    pub fn read_mode(&self) -> Option<&WebSocketReadMode> {
        match self {
            WebSocketOperation::Read { mode, .. } => Some(mode),
            _ => None,
        }
    }

    /// Returns the messages if this is a Send operation, None otherwise.
    pub fn messages(&self) -> Option<&[RexValue]> {
        match self {
            WebSocketOperation::Send { messages, .. } => Some(messages),
            _ => None,
        }
    }
}

impl fmt::Display for WebSocketOperation {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            WebSocketOperation::Connect { url, rex_id } => {
                write!(f, "Connect(url={url}, rex_id={rex_id})")
            }
            WebSocketOperation::Read {
                connection_rex_id,
                mode,
            } => {
                write!(
                    f,
                    "Read(connection_rex_id={}, mode={:?})",
                    connection_rex_id, mode
                )
            }
            WebSocketOperation::Send {
                connection_rex_id,
                messages,
            } => {
                write!(
                    f,
                    "Send(connection_rex_id={}, messages_count={})",
                    connection_rex_id,
                    messages.len()
                )
            }
            WebSocketOperation::Close { connection_rex_id } => {
                write!(f, "Close(connection_rex_id={})", connection_rex_id)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use rialo_s_pubkey::Pubkey;

    use super::*;
    use crate::rex_info::{RexInfo, StartingTimestamp, TargetRexProgram, UpdateFrequency};

    fn base_valid_rex_info() -> RexInfo {
        RexInfo {
            description: "test".to_string(),
            target_rex_programs: vec![TargetRexProgram::Time],
            // Default is OneShot, which is allowed for both Asap and Timestamp
            update_frequency: UpdateFrequency::OneShot,
            // Start at timestamp 0 by default
            starting_timestamp: StartingTimestamp::Timestamp(0),
            // Keep other fields as default
            ..RexInfo::default()
        }
    }

    // ========================================================================
    // WebSocket Operation Tests
    // ========================================================================

    #[test]
    fn test_websocket_connect_operation_creation() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::connect("wss://example.com/stream", rex_id);
        assert!(op.is_connect());
        assert!(!op.is_read());

        if let WebSocketOperation::Connect {
            url,
            rex_id: op_rex_id,
        } = op
        {
            assert_eq!(url.to_string(), "wss://example.com/stream");
            assert_eq!(op_rex_id, rex_id);
        } else {
            panic!("Expected Connect variant");
        }
    }

    #[test]
    fn test_websocket_read_operation_creation() {
        let rex_id = RexId::new(Pubkey::default(), 42u64);
        let op = WebSocketOperation::read(rex_id);
        assert!(op.is_read());
        assert!(!op.is_connect());

        if let WebSocketOperation::Read {
            connection_rex_id,
            mode,
        } = op
        {
            assert_eq!(connection_rex_id, rex_id);
            assert_eq!(mode, WebSocketReadMode::Latest);
        } else {
            panic!("Expected Read variant");
        }
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_connect() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            rex_id,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_read() {
        let rex_id = RexId::new(Pubkey::default(), 123u64);
        let op = WebSocketOperation::Read {
            connection_rex_id: rex_id,
            mode: WebSocketReadMode::default(),
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_target_rex_program_websocket_serde_roundtrip() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let target = TargetRexProgram::WebSocket(WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            rex_id,
        });

        // Serialize to JSON
        let json = serde_json::to_string(&target).expect("Failed to serialize");
        // Deserialize back
        let deserialized: TargetRexProgram =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(target, deserialized);
    }

    #[test]
    fn test_websocket_operation_display() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com".into(),
            rex_id,
        };
        let display = format!("{}", connect_op);
        assert!(display.contains("Connect"));
        assert!(display.contains("wss://example.com"));

        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let read_op = WebSocketOperation::Read {
            connection_rex_id: rex_id,
            mode: WebSocketReadMode::default(),
        };
        let display = format!("{}", read_op);
        assert!(display.contains("Read"));
        assert!(display.contains("connection_rex_id"));
    }

    #[test]
    fn test_websocket_connect_with_encrypted_url() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "enc://encrypted_websocket_url".into(),
            rex_id,
        };

        if let WebSocketOperation::Connect { url, .. } = &op {
            assert_eq!(url.to_string(), "enc://encrypted_websocket_url");
        }

        // Verify serde roundtrip preserves encrypted URL
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_rex_info_with_websocket_target() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let mut info = base_valid_rex_info();
        info.target_rex_programs = vec![TargetRexProgram::WebSocket(WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            rex_id,
        })];

        // Should validate successfully
        assert!(info.validate().is_ok());
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_connect() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            rex_id,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_read() {
        let rex_id = RexId::new(Pubkey::default(), 123u64);
        let op = WebSocketOperation::Read {
            connection_rex_id: rex_id,
            mode: WebSocketReadMode::default(),
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_encrypted_url() {
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let op = WebSocketOperation::Connect {
            url: "enc://encrypted_websocket_url".into(),
            rex_id,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_url_getter() {
        // Connect operation should return Some(url)
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            rex_id,
        };
        assert!(connect_op.url().is_some());
        assert_eq!(
            connect_op.url().unwrap().to_string(),
            "wss://example.com/stream"
        );

        // Read operation should return None
        let rex_id = RexId::new(Pubkey::default(), 42u64);
        let read_op = WebSocketOperation::Read {
            connection_rex_id: rex_id,
            mode: WebSocketReadMode::default(),
        };
        assert!(read_op.url().is_none());
    }

    #[test]
    fn test_websocket_operation_connection_rex_id_getter() {
        // Read operation should return Some(connection_rex_id)
        let rex_id = RexId::new(Pubkey::default(), 42u64);
        let read_op = WebSocketOperation::Read {
            connection_rex_id: rex_id,
            mode: WebSocketReadMode::default(),
        };
        assert!(read_op.connection_rex_id().is_some());
        assert_eq!(*read_op.connection_rex_id().unwrap(), rex_id);

        // Connect operation should return None
        let rex_id = RexId::new(Pubkey::default(), 1u64);
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com/stream".into(),
            rex_id,
        };
        assert!(connect_op.connection_rex_id().is_none());

        // Send operation should return Some(connection_rex_id)
        let rex_id = RexId::new(Pubkey::default(), 99u64);
        let send_op = WebSocketOperation::Send {
            connection_rex_id: rex_id,
            messages: vec![],
        };
        assert!(send_op.connection_rex_id().is_some());
        assert_eq!(*send_op.connection_rex_id().unwrap(), rex_id);
    }

    // ========================================================================
    // WebSocket Send Operation Tests
    // ========================================================================

    #[test]
    fn test_websocket_send_operation_creation() {
        let rex_id = RexId::new(Pubkey::default(), 5u64);
        let messages = vec![
            RexValue::plain_string("message1"),
            RexValue::plain_string("message2"),
        ];
        let op = WebSocketOperation::send(rex_id, messages.clone());
        assert!(op.is_send());
        assert!(!op.is_connect());
        assert!(!op.is_read());

        if let WebSocketOperation::Send {
            connection_rex_id,
            messages: op_messages,
        } = op
        {
            assert_eq!(connection_rex_id, rex_id);
            assert_eq!(op_messages, messages);
        } else {
            panic!("Expected Send variant");
        }
    }

    #[test]
    fn test_websocket_send_messages_getter() {
        let rex_id = RexId::new(Pubkey::default(), 10u64);
        let messages = vec![RexValue::plain_string("test")];
        let send_op = WebSocketOperation::Send {
            connection_rex_id: rex_id,
            messages: messages.clone(),
        };

        assert!(send_op.messages().is_some());
        assert_eq!(send_op.messages().unwrap(), &messages[..]);

        // Connect and Read should return None
        let connect_op = WebSocketOperation::Connect {
            url: "wss://example.com".into(),
            rex_id,
        };
        assert!(connect_op.messages().is_none());

        let read_op = WebSocketOperation::Read {
            connection_rex_id: rex_id,
            mode: WebSocketReadMode::default(),
        };
        assert!(read_op.messages().is_none());
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_send() {
        let rex_id = RexId::new(Pubkey::default(), 7u64);
        let messages = vec![
            RexValue::plain_string("plain_message"),
            RexValue::encrypted(b"encrypted_data".to_vec()),
        ];
        let op = WebSocketOperation::Send {
            connection_rex_id: rex_id,
            messages,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_send() {
        let rex_id = RexId::new(Pubkey::default(), 8u64);
        let messages = vec![
            RexValue::plain_string("hello"),
            RexValue::plain_string("world"),
        ];
        let op = WebSocketOperation::Send {
            connection_rex_id: rex_id,
            messages,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_send_display() {
        let rex_id = RexId::new(Pubkey::default(), 20u64);
        let messages = vec![
            RexValue::plain_string("msg1"),
            RexValue::plain_string("msg2"),
            RexValue::plain_string("msg3"),
        ];
        let send_op = WebSocketOperation::Send {
            connection_rex_id: rex_id,
            messages,
        };

        let display = format!("{}", send_op);
        assert!(display.contains("Send"));
        assert!(display.contains("connection_rex_id"));
        assert!(display.contains("messages_count=3"));
    }

    #[test]
    fn test_websocket_send_empty_messages() {
        let rex_id = RexId::new(Pubkey::default(), 15u64);
        let op = WebSocketOperation::send(rex_id, vec![]);

        assert!(op.is_send());
        assert_eq!(op.messages().unwrap().len(), 0);

        let display = format!("{}", op);
        assert!(display.contains("messages_count=0"));
    }

    #[test]
    fn test_websocket_send_with_encrypted_messages() {
        let rex_id = RexId::new(Pubkey::default(), 25u64);
        let messages = vec![
            RexValue::encrypted(b"encrypted1".to_vec()),
            RexValue::encrypted(b"encrypted2".to_vec()),
        ];
        let op = WebSocketOperation::Send {
            connection_rex_id: rex_id,
            messages: messages.clone(),
        };

        // Verify serde roundtrip preserves encrypted messages
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");
        assert_eq!(op, deserialized);

        // Verify borsh roundtrip preserves encrypted messages
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");
        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_target_rex_program_websocket_send_serde_roundtrip() {
        let rex_id = RexId::new(Pubkey::default(), 30u64);
        let messages = vec![RexValue::plain_string("test_message")];
        let target = TargetRexProgram::WebSocket(WebSocketOperation::Send {
            connection_rex_id: rex_id,
            messages,
        });

        // Serialize to JSON
        let json = serde_json::to_string(&target).expect("Failed to serialize");
        // Deserialize back
        let deserialized: TargetRexProgram =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(target, deserialized);
    }

    // ========================================================================
    // WebSocket Close Operation Tests
    // ========================================================================

    #[test]
    fn test_websocket_close_operation_creation() {
        let rex_id = RexId::new(Pubkey::default(), 50u64);
        let op = WebSocketOperation::close(rex_id);
        assert!(op.is_close());
        assert!(!op.is_connect());
        assert!(!op.is_read());
        assert!(!op.is_send());

        if let WebSocketOperation::Close { connection_rex_id } = op {
            assert_eq!(connection_rex_id, rex_id);
        } else {
            panic!("Expected Close variant");
        }
    }

    #[test]
    fn test_websocket_close_connection_rex_id_getter() {
        let rex_id = RexId::new(Pubkey::default(), 55u64);
        let close_op = WebSocketOperation::Close {
            connection_rex_id: rex_id,
        };
        assert!(close_op.connection_rex_id().is_some());
        assert_eq!(*close_op.connection_rex_id().unwrap(), rex_id);
    }

    #[test]
    fn test_websocket_operation_serde_roundtrip_close() {
        let rex_id = RexId::new(Pubkey::default(), 60u64);
        let op = WebSocketOperation::Close {
            connection_rex_id: rex_id,
        };

        // Serialize to JSON
        let json = serde_json::to_string(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_operation_borsh_roundtrip_close() {
        let rex_id = RexId::new(Pubkey::default(), 65u64);
        let op = WebSocketOperation::Close {
            connection_rex_id: rex_id,
        };

        // Serialize to Borsh
        let bytes = borsh::to_vec(&op).expect("Failed to serialize");
        // Deserialize back
        let deserialized: WebSocketOperation =
            borsh::from_slice(&bytes).expect("Failed to deserialize");

        assert_eq!(op, deserialized);
    }

    #[test]
    fn test_websocket_close_display() {
        let rex_id = RexId::new(Pubkey::default(), 70u64);
        let close_op = WebSocketOperation::Close {
            connection_rex_id: rex_id,
        };

        let display = format!("{}", close_op);
        assert!(display.contains("Close"));
        assert!(display.contains("connection_rex_id"));
        assert!(display.contains(&rex_id.to_string()));
    }

    #[test]
    fn test_target_rex_program_websocket_close_serde_roundtrip() {
        let rex_id = RexId::new(Pubkey::default(), 75u64);
        let target = TargetRexProgram::WebSocket(WebSocketOperation::Close {
            connection_rex_id: rex_id,
        });

        // Serialize to JSON
        let json = serde_json::to_string(&target).expect("Failed to serialize");
        // Deserialize back
        let deserialized: TargetRexProgram =
            serde_json::from_str(&json).expect("Failed to deserialize");

        assert_eq!(target, deserialized);
    }
}