feat: ws v2 (#6404)

# WebSocket v2 Architecture

## NOTE: Open Issues:
- [ ] `session_drain` from server is experimental and not ready for
prod. Instead, the client maintains ongoing `trace_id`'s and when cookie
expires it will retry search for all the ongoing searches.

## Overview
Implements websocket v2 with the main goal to establish `Client-Router`
connections `per tab` and `Router-Querier` `persistent` connections.

## New Env (Default):
```sh
ZO_WEBSOCKET_CHECK_COOKIE_EXPIRY=true
ZO_WEBSOCKET_HEALTH_CHECK_INTERVAL_SECS=60
ZO_WEBSOCKET_IS_SESSION_DRAIN_ENABLED=false # This is experimental, not meant for prod
```

## High-Level Flow

1. Client establishes WebSocket connection with the router.  
2. Router creates a session.  
3. Router maintains persistent connections to all queriers.  
4. Client requests are routed to the appropriate querier using
consistent hashing.
5. Responses are routed back using `trace_id` mappings.  
6. Health checks maintain connection validity.  
7. Failed connections trigger automatic cleanup and reconnection.  
8. **Internal Communication Validation Using:**  
   - **8a.** Internal gRPC tokens  
   - **8b.** IP validation  
   - **8c.** WebSocket upgrade headers  

## Connection Lifecycle
```mermaid
sequenceDiagram
    participant Client
    participant Router
    participant ConnectionPool
    participant QuerierConn
    participant Querier
    participant Cluster

    Client->>Router: WS Connection Request
    Router->>SessionManager: register_client(client_id)
    
    Note over Router,SessionManager: Creates channels:<br/>response_tx/rx<br/>disconnect_tx/rx

    par Handle Incoming Messages
        Router->>Router: spawn handle_incoming task
    and Handle Outgoing Messages
        Router->>Router: spawn handle_outgoing task
    end

    Client->>Router: WS Client Event
    Router->>SessionManager: get_querier_for_trace(trace_id)
    
    alt No Existing Mapping
        Router->>Cluster: get_node_from_consistent_hash(trace_id)
        Router->>ConnectionPool: get_or_create_connection(querier)
        
        alt No Existing Connection
            ConnectionPool->>QuerierConn: create_connection
            QuerierConn->>Querier: WS Connection Request
            
            par Connection Tasks
                QuerierConn->>QuerierConn: spawn listen_to_querier_response
                QuerierConn->>QuerierConn: spawn health_check
            end
            
            Querier-->>QuerierConn: Connection Established
        end
        
        Router->>SessionManager: set_querier_for_trace(trace_id, querier)
    end

    Router->>QuerierConn: register_request(trace_id)
    Router->>QuerierConn: send_message(event)
    Querier-->>QuerierConn: Response
    QuerierConn->>Router: Route Response via response_tx
    Router->>Client: Forward Response

    Note over Router,Client: On Disconnect/Error:<br/>1. Send disconnect signal<br/>2. Clean up sessions<br/>3. Drain if needed
```

## Connections
```mermaid
graph TB
    subgraph "Client Layer"
        UI1[UI Client 1]
        UI2[UI Client 2]
        UI3[UI Client 3]
    end

    subgraph "Router Layer"
        R1[Router 1]
        R2[Router 2]
        
        subgraph "Router 1 Connection Pool"
            CP1[Connection Pool]
            QC11[Querier1 Conn]
            QC12[Querier2 Conn]
            QC13[Querier3 Conn]
            CP1 --- QC11
            CP1 --- QC12
            CP1 --- QC13
        end
        
        subgraph "Router 2 Connection Pool"
            CP2[Connection Pool]
            QC21[Querier1 Conn]
            QC22[Querier2 Conn]
            QC23[Querier3 Conn]
            CP2 --- QC21
            CP2 --- QC22
            CP2 --- QC23
        end
        
        R1 --- CP1
        R2 --- CP2
    end

    subgraph "Querier Layer"
        Q1[Querier 1]
        Q2[Querier 2]
        Q3[Querier 3]
    end

    UI1 -.->|WS| R1
    UI2 -.->|WS| R1
    UI3 -.->|WS| R2

    QC11 ===|Persistent WS| Q1
    QC12 ===|Persistent WS| Q2
    QC13 ===|Persistent WS| Q3
    QC21 ===|Persistent WS| Q1
    QC22 ===|Persistent WS| Q2
    QC23 ===|Persistent WS| Q3

    classDef client fill:#f9f,stroke:#333,stroke-width:2px
    classDef router fill:#bbf,stroke:#333,stroke-width:2px
    classDef querier fill:#bfb,stroke:#333,stroke-width:2px
    classDef pool fill:#ffd,stroke:#333,stroke-width:2px
    
    class UI1,UI2,UI3 client
    class R1,R2 router
    class Q1,Q2,Q3 querier
    class CP1,CP2 pool
```

## Components
```mermaid
classDiagram
    class WsHandler {
        +session_manager: Arc<SessionManager>
        +connection_pool: Arc<QuerierConnectionPool>
        +handle_connection()
        +remove_querier_connection()
        +_shutdown()
    }

    class SessionManager {
        +sessions: RwAHashMap<ClientId, SessionInfo>
        +mapped_queriers: RwAHashMap<QuerierName, HashSet<TraceId>>
        +register_client()
        +update_session_activity()
        +remove_trace_id()
        +unregister_client()
        +set_querier_for_trace()
        +get_querier_for_trace()
    }

    class QuerierConnectionPool {
        +connections: RwAHashMap<QuerierName, Arc<QuerierConnection>>
        +get_or_create_connection()
        +remove_querier_connection()
        +maintain_connections()
        +get_active_connection()
    }

    class QuerierConnection {
        +querier_name: QuerierName
        +write: Arc<Mutex<Option<SplitSink>>>
        +response_router: Arc<ResponseRouter>
        +is_connected: Arc<AtomicBool>
        +connect()
        +disconnect()
        +send_message()
        +register_request()
        +unregister_request()
    }

    class ResponseRouter {
        +routes: RwAHashMap<TraceId, Sender<WsServerEvents>>
        +route_response()
        +flush()
        +spawn_cleanup_task()
    }

    WsHandler --> SessionManager
    WsHandler --> QuerierConnectionPool
    QuerierConnectionPool --> QuerierConnection
    QuerierConnection --> ResponseRouter
```


Related Issues: #6396

---------

Co-authored-by: Taiming Liu <liutaiming3@gmail.com>
Co-authored-by: Hengfei Yang <hengfei.yang@gmail.com>
Co-authored-by: Abhay Padamani <141122205+ktx-abhay@users.noreply.github.com>
Co-authored-by: Ashish Kolhe <35160958+oasisk@users.noreply.github.com>
Co-authored-by: udev <uddhav@openobserve.ai>
Co-authored-by: omkark06 <omkar@zinclabs.io>
Co-authored-by: ktx-abhay <abhay.padamani@kiara.tech>
Co-authored-by: Omkar Kesarkhane <127078473+omkarK06@users.noreply.github.com>
Co-authored-by: udev <uddhavdave1@gmail.com>
Co-authored-by: Subhra264 <chakrabortysubhradeep556@gmail.com>
Co-authored-by: ktx-vaidehi <134508096+ktx-vaidehi@users.noreply.github.com>

src/common/infra/cluster/mod.rs

@@ -337,6 +337,9 @@ async fn watch_node_list() -> Result<()> {
                         )
                         .await;
                     }
+                    if item_value.is_querier() && LOCAL_NODE.is_router() {
+                        crate::router::http::remove_querier_from_handler(&item_value.name).await;
+                    }
                     if item_value.is_compactor() {
                         remove_node_from_consistent_hash(&item_value, &Role::Compactor, None).await;
                     }
@@ -408,6 +411,9 @@ async fn watch_node_list() -> Result<()> {
                     )
                     .await;
                 }
+                if item_value.is_querier() && LOCAL_NODE.is_router() {
+                    crate::router::http::remove_querier_from_handler(&item_value.name).await;
+                }
                 if item_value.is_compactor() {
                     remove_node_from_consistent_hash(&item_value, &Role::Compactor, None).await;
                 }
@@ -543,6 +549,12 @@ pub async fn get_cached_online_ingester_nodes() -> Option<Vec<Node>> {
     .await
 }
 
+#[inline]
+pub async fn get_cached_online_router_nodes() -> Option<Vec<Node>> {
+    get_cached_nodes(|node| node.status == NodeStatus::Online && node.scheduled && node.is_router())
+        .await
+}
+
 #[inline]
 pub async fn get_cached_online_querier_nodes(group: Option<RoleGroup>) -> Option<Vec<Node>> {
     let nodes = get_cached_nodes(|node| {

src/common/infra/config.rs

@@ -38,10 +38,11 @@ use crate::{
     common::meta::{
         maxmind::MaxmindClient, organization::OrganizationSetting, syslog::SyslogRoute, user::User,
     },
-    handler::http::request::websocket::session::WsSession,
+    handler::http::request::ws_v2::session::WsSession,
     service::{
         db::scheduler as db_scheduler, enrichment::StreamTable, enrichment_table::geoip::Geoip,
         pipeline::batch_execution::ExecutablePipeline,
+        websocket_events::search_registry_utils::SearchState,
     },
 };
 
@@ -95,3 +96,5 @@ pub static USER_SESSIONS: Lazy<RwHashMap<String, String>> = Lazy::new(Default::d
 pub static SHORT_URLS: Lazy<RwHashMap<String, ShortUrlRecord>> = Lazy::new(DashMap::default);
 // TODO: Implement rate limiting for maximum number of sessions
 pub static WS_SESSIONS: Lazy<RwHashMap<String, WsSession>> = Lazy::new(DashMap::default);
+// Global registry for search requests by `trace_id`
+pub static WS_SEARCH_REGISTRY: Lazy<DashMap<String, SearchState>> = Lazy::new(DashMap::new);

src/common/utils/auth.rs

@@ -20,15 +20,22 @@ use argon2::{Algorithm, Argon2, Params, PasswordHasher, Version, password_hash::
 use base64::Engine;
 use config::utils::json;
 use futures::future::{Ready, ready};
-#[cfg(feature = "enterprise")]
-use o2_openfga::config::get_config as get_openfga_config;
-#[cfg(feature = "enterprise")]
-use o2_openfga::meta::mapping::OFGA_MODELS;
 use once_cell::sync::Lazy;
 use regex::Regex;
-
 #[cfg(feature = "enterprise")]
-use crate::common::infra::config::USER_SESSIONS;
+use {
+    crate::common::{
+        infra::config::USER_SESSIONS,
+        {meta, meta::ingestion::INGESTION_EP},
+    },
+    jsonwebtoken::TokenData,
+    o2_dex::service::auth::get_dex_jwks,
+    o2_openfga::config::get_config as get_openfga_config,
+    o2_openfga::meta::mapping::OFGA_MODELS,
+    serde_json::Value,
+    std::{collections::HashMap, str::FromStr},
+};
+
 use crate::common::{
     infra::config::{PASSWORD_HASH, USERS},
     meta::{
@@ -37,8 +44,6 @@ use crate::common::{
         user::{AuthTokens, UserRole},
     },
 };
-#[cfg(feature = "enterprise")]
-use crate::common::{meta, meta::ingestion::INGESTION_EP};
 
 pub const V2_API_PREFIX: &str = "v2";
 
@@ -202,7 +207,7 @@ impl FromRequest for AuthExtractor {
     #[cfg(feature = "enterprise")]
     fn from_request(req: &HttpRequest, _: &mut Payload) -> Self::Future {
         use actix_web::web;
-        use config::meta::stream::StreamType;
+        use config::{get_config, meta::stream::StreamType};
         use hashbrown::HashMap;
         use o2_openfga::meta::mapping::OFGA_MODELS;
 
@@ -573,6 +578,26 @@ impl FromRequest for AuthExtractor {
             )
         };
 
+        // Check if the ws request is using internal grpc token
+        if method.eq("GET") && path.contains("/ws") {
+            if let Some(auth_header) = req.headers().get("Authorization") {
+                if auth_header
+                    .to_str()
+                    .unwrap()
+                    .eq(&get_config().grpc.internal_grpc_token)
+                {
+                    return ready(Ok(AuthExtractor {
+                        auth: auth_header.to_str().unwrap().to_string(),
+                        method,
+                        o2_type: format!("stream:{org_id}"),
+                        org_id,
+                        bypass_check: true,
+                        parent_id: folder,
+                    }));
+                }
+            }
+        }
+
         let auth_str = extract_auth_str(req);
 
         // if let Some(auth_header) = req.headers().get("Authorization") {
@@ -853,6 +878,93 @@ pub async fn check_permissions(
     true
 }
 
+#[cfg(feature = "enterprise")]
+pub async fn extract_auth_expiry_and_user_id(
+    req: &HttpRequest,
+) -> (Option<chrono::DateTime<chrono::Utc>>, Option<String>) {
+    use crate::handler::http::auth::validator::get_user_email_from_auth_str;
+
+    let decode = async |token: &str| match decode_expiry(token).await {
+        Ok(token_data) => token_data
+            .claims
+            .get("exp")
+            .and_then(|exp| exp.as_i64())
+            .and_then(|exp_ts| chrono::DateTime::from_timestamp(exp_ts, 0)),
+        Err(e) => {
+            log::error!("Error verifying token: {}", e);
+            None
+        }
+    };
+
+    let auth_str = extract_auth_str(req);
+    if auth_str.is_empty() {
+        return (None, None);
+    } else if auth_str.starts_with("Basic") {
+        let user_id = get_user_email_from_auth_str(&auth_str).await;
+        return (None, user_id);
+    } else if auth_str.starts_with("Bearer") {
+        let user_id = get_user_email_from_auth_str(&auth_str).await;
+        let stripped_bearer_token = auth_str.strip_prefix("Bearer ").unwrap();
+        let exp = decode(stripped_bearer_token).await;
+        return (exp, user_id);
+    } else if auth_str.starts_with("session ") {
+        let session_key = auth_str.strip_prefix("session ").unwrap();
+        let stripped_bearer_token = match crate::service::db::session::get(session_key).await {
+            Ok(bearer_token) => bearer_token,
+            Err(e) => {
+                log::error!("Error getting session: {}", e);
+                return (None, None);
+            }
+        };
+        let exp = decode(&stripped_bearer_token).await;
+        let bearer_full_token = format!("Bearer {}", stripped_bearer_token);
+        let user_id = get_user_email_from_auth_str(&bearer_full_token).await;
+        return (exp, user_id);
+    }
+    (None, None)
+}
+
+#[cfg(feature = "enterprise")]
+async fn decode_expiry(token: &str) -> Result<TokenData<HashMap<String, Value>>, anyhow::Error> {
+    use infra::errors::JwtError;
+    use jsonwebtoken::{
+        Algorithm, DecodingKey, Validation, decode, decode_header,
+        jwk::{self, AlgorithmParameters},
+    };
+
+    let header = decode_header(token)?;
+    let kid = match header.kid {
+        Some(k) => k,
+        None => return Err(JwtError::MissingAttribute("`kid` header".to_owned()).into()),
+    };
+    let dex_jwks = get_dex_jwks().await;
+    let jwks: jwk::JwkSet = serde_json::from_str(&dex_jwks).unwrap();
+
+    if let Some(j) = jwks.find(&kid) {
+        match &j.algorithm {
+            AlgorithmParameters::RSA(rsa) => {
+                let decoding_key = DecodingKey::from_rsa_components(&rsa.n, &rsa.e).unwrap();
+
+                let mut validation = Validation::new(
+                    Algorithm::from_str(j.common.key_algorithm.unwrap().to_string().as_str())
+                        .unwrap(),
+                );
+                validation.validate_exp = true;
+                let aud = &o2_dex::config::get_config().client_id;
+                validation.set_audience(&[aud]);
+                Ok(decode::<HashMap<String, serde_json::Value>>(
+                    token,
+                    &decoding_key,
+                    &validation,
+                )?)
+            }
+            _ => Err(JwtError::ValidationFailed().into()),
+        }
+    } else {
+        Err(JwtError::KeyNotExists().into())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use infra::db as infra_db;

src/common/utils/auth_tests.rs

@@ -3005,6 +3005,7 @@ mod tests {
                 session_max_lifetime_secs: i64::default(),
                 session_gc_interval_secs: i64::default(),
                 ping_interval_secs: i64::default(),
+                check_cookie_expiry: bool::default(),
             },
             route: config::Route {
                 timeout: u64::default(),

src/config/src/config.rs

@@ -431,6 +431,12 @@ pub struct WebSocket {
     pub session_gc_interval_secs: i64,
     #[env_config(name = "ZO_WEBSOCKET_PING_INTERVAL_SECS", default = 15)]
     pub ping_interval_secs: i64,
+    #[env_config(name = "ZO_WEBSOCKET_CHECK_COOKIE_EXPIRY", default = true)]
+    pub check_cookie_expiry: bool,
+    #[env_config(name = "ZO_WEBSOCKET_HEALTH_CHECK_INTERVAL_SECS", default = 60)]
+    pub health_check_interval: i64,
+    #[env_config(name = "ZO_WEBSOCKET_IS_SESSION_DRAIN_ENABLED", default = false)]
+    pub is_session_drain_enabled: bool,
 }
 
 #[derive(EnvConfig)]

src/config/src/meta/websocket.rs

@@ -36,4 +36,21 @@ pub struct SearchEventReq {
     pub search_event_context: Option<SearchEventContext>,
     #[serde(default)]
     pub fallback_order_by_col: Option<String>,
+    // TODO: modify this once v1 is deprecated
+    pub org_id: String,
+    // TODO: is it PII safe?
+    pub user_id: Option<String>,
+}
+
+impl SearchEventReq {
+    pub fn is_valid(&self) -> bool {
+        // TODO: add event payload validation in the future
+        if self.trace_id.is_empty() {
+            return false;
+        }
+        if self.org_id.is_empty() {
+            return false;
+        }
+        true
+    }
 }

src/handler/http/auth/validator.rs

@@ -13,6 +13,8 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+use std::net::IpAddr;
+
 use actix_web::{
     Error,
     dev::ServiceRequest,
@@ -25,9 +27,11 @@ use config::{get_config, utils::base64};
 use o2_dex::config::get_config as get_dex_config;
 #[cfg(feature = "enterprise")]
 use o2_openfga::config::get_config as get_openfga_config;
+use url::Url;
 
 use crate::{
     common::{
+        infra::cluster,
         meta::{
             ingestion::INGESTION_EP,
             user::{
@@ -599,6 +603,11 @@ async fn oo_validator_internal(
     auth_info: AuthExtractor,
     path_prefix: &str,
 ) -> Result<ServiceRequest, (Error, ServiceRequest)> {
+    // Check if the ws request is using internal grpc token
+    if get_config().websocket.enabled && auth_info.auth.eq(&get_config().grpc.internal_grpc_token) {
+        return validate_http_internal(req).await;
+    }
+
     if auth_info.auth.starts_with("Basic") {
         let decoded = match base64::decode(auth_info.auth.strip_prefix("Basic").unwrap().trim()) {
             Ok(val) => val,
@@ -760,6 +769,62 @@ pub async fn validator_proxy_url(
     oo_validator_internal(req, auth_info, path_prefix).await
 }
 
+pub async fn validate_http_internal(
+    req: ServiceRequest,
+) -> Result<ServiceRequest, (Error, ServiceRequest)> {
+    let router_nodes = cluster::get_cached_online_router_nodes()
+        .await
+        .unwrap_or_default();
+
+    // Get the peer address early and own the string
+    let peer = req
+        .connection_info()
+        .peer_addr()
+        .unwrap_or("unknown")
+        .to_string();
+
+    let router_node = router_nodes.iter().find(|node| {
+        let node_url = match Url::parse(&node.http_addr) {
+            Ok(node_url) => node_url,
+            Err(e) => {
+                log::error!("Failed to parse node URL: {}", e);
+                return false;
+            }
+        };
+
+        // Get IP from peer address (strips port if present)
+        let peer_ip = match peer
+            .split(':')
+            .next()
+            .and_then(|addr| addr.parse::<IpAddr>().ok())
+        {
+            Some(ip) => ip,
+            None => {
+                log::debug!("Failed to parse peer IP from: {}", peer);
+                return false;
+            }
+        };
+
+        let node_ip = match node_url
+            .host()
+            .and_then(|h| h.to_string().parse::<IpAddr>().ok())
+        {
+            Some(ip) => ip,
+            None => {
+                log::debug!("Failed to parse node IP");
+                return false;
+            }
+        };
+
+        peer_ip == node_ip
+    });
+
+    if router_node.is_none() {
+        return Err((ErrorUnauthorized("Unauthorized Access"), req));
+    }
+    Ok(req)
+}
+
 #[cfg(feature = "enterprise")]
 pub(crate) async fn check_permissions(
     user_id: &str,

src/handler/http/request/mod.rs

@@ -43,6 +43,7 @@ pub mod syslog;
 pub mod traces;
 pub mod users;
 pub mod websocket;
+pub mod ws_v2;
 
 pub const CONTENT_TYPE_JSON: &str = "application/json";
 pub const CONTENT_TYPE_PROTO: &str = "application/x-protobuf";

src/handler/http/request/websocket/mod.rs

@@ -13,15 +13,15 @@
 // You should have received a copy of the GNU Affero General Public License
 // along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-pub mod search;
 pub mod session;
-pub mod sort;
-pub mod utils;
 
 use actix_web::{Error, HttpRequest, HttpResponse, get, web};
 use config::{cluster::LOCAL_NODE, get_config};
-use session::WsSession;
-use utils::sessions_cache_utils;
+
+use crate::{
+    handler::http::request::ws_v2::session::WsSession,
+    service::websocket_events::sessions_cache_utils,
+};
 
 #[get("{org_id}/ws/{request_id}")]
 pub async fn websocket(
@@ -54,7 +54,7 @@ pub async fn websocket(
 
     let (res, session, msg_stream) = actix_ws::handle(&req, stream)?;
 
-    let ws_session = WsSession::new(session);
+    let ws_session = WsSession::new(session, None);
     sessions_cache_utils::insert_session(&request_id, ws_session);
     log::info!(
         "[WS_HANDLER]: Node Role: {} Got websocket request for request_id: {}",
@@ -71,7 +71,7 @@ pub async fn websocket(
 /// Initialize the job init for websocket
 pub async fn init() -> Result<(), anyhow::Error> {
     // Run the garbage collector for websocket sessions
-    if !LOCAL_NODE.is_compactor() && !LOCAL_NODE.is_alert_manager() {
+    if LOCAL_NODE.is_querier() {
         sessions_cache_utils::run_gc_ws_sessions().await;
     }
     Ok(())

src/handler/http/request/websocket/session.rs

@@ -16,12 +16,11 @@
 use std::time::Duration;
 
 use actix_http::ws::{CloseCode, CloseReason};
-use actix_ws::{MessageStream, Session};
+use actix_ws::MessageStream;
 use config::{
     get_config,
     meta::websocket::{SearchEventReq, SearchResultType},
 };
-use dashmap::DashMap;
 use futures::StreamExt;
 use infra::errors::{self, Error};
 #[cfg(feature = "enterprise")]
@@ -29,98 +28,21 @@ use o2_enterprise::enterprise::common::{
     auditor::{AuditMessage, Protocol, WsMeta},
     infra::config::get_config as get_o2_config,
 };
-use once_cell::sync::Lazy;
 use rand::prelude::SliceRandom;
 use tokio::sync::mpsc;
 
-use super::utils::search_registry_utils::SearchState;
 #[cfg(feature = "enterprise")]
-use crate::handler::http::request::websocket::utils::search_registry_utils;
-use crate::handler::http::request::websocket::{
-    search,
-    utils::{WsClientEvents, WsServerEvents, sessions_cache_utils},
+use crate::service::{
+    self_reporting::audit,
+    websocket_events::{handle_cancel, search_registry_utils},
+};
+use crate::{
+    common::infra::config::WS_SEARCH_REGISTRY,
+    service::websocket_events::{
+        WsClientEvents, WsServerEvents, handle_search_request, search_registry_utils::SearchState,
+        sessions_cache_utils,
+    },
 };
-#[cfg(feature = "enterprise")]
-use crate::service::self_reporting::audit;
-
-// Global registry for search requests by `trace_id`
-pub static SEARCH_REGISTRY: Lazy<DashMap<String, SearchState>> = Lazy::new(DashMap::new);
-
-// Do not clone the session, instead use a reference to the session
-pub struct WsSession {
-    inner: Option<Session>,
-    // Utc timestamp in microseconds
-    last_activity_ts: i64,
-    // Utc timestamp in microseconds
-    created_ts: i64,
-}
-
-impl WsSession {
-    pub fn new(inner: Session) -> Self {
-        let now = chrono::Utc::now().timestamp_micros();
-        Self {
-            inner: Some(inner),
-            last_activity_ts: now,
-            created_ts: now,
-        }
-    }
-
-    pub fn update_activity(&mut self) {
-        self.last_activity_ts = chrono::Utc::now().timestamp_micros();
-    }
-
-    pub fn is_expired(&self) -> bool {
-        let cfg = get_config();
-        let now = chrono::Utc::now().timestamp_micros();
-        let idle_timeout_micros = cfg.websocket.session_idle_timeout_secs * 1_000_000;
-        let max_lifetime_micros = cfg.websocket.session_max_lifetime_secs * 1_000_000;
-
-        // 1. if the session has been idle for too long
-        // 2. if the session has exceeded the max lifetime
-        (now - self.last_activity_ts) > idle_timeout_micros
-            || (now - self.created_ts) > max_lifetime_micros
-    }
-
-    /// Send a text message to the client
-    pub async fn text(&mut self, msg: String) -> Result<(), actix_ws::Closed> {
-        self.update_activity();
-        if let Some(ref mut session) = self.inner {
-            session.text(msg).await
-        } else {
-            Err(actix_ws::Closed)
-        }
-    }
-
-    /// Close the session with a reason
-    pub async fn close(&mut self, reason: Option<CloseReason>) -> Result<(), actix_ws::Closed> {
-        self.update_activity();
-        if let Some(session) = self.inner.take() {
-            session.close(reason).await
-        } else {
-            Err(actix_ws::Closed)
-        }
-    }
-
-    /// Send a pong response
-    pub async fn pong(&mut self, payload: &[u8]) -> Result<(), actix_ws::Closed> {
-        self.update_activity();
-        if let Some(ref mut session) = self.inner {
-            session.pong(payload).await
-        } else {
-            Err(actix_ws::Closed)
-        }
-    }
-
-    /// Send a ping request
-    pub async fn ping(&mut self, payload: &[u8]) -> Result<(), actix_ws::Closed> {
-        self.update_activity();
-        if let Some(ref mut session) = self.inner {
-            session.ping(payload).await
-        } else {
-            Err(actix_ws::Closed)
-        }
-    }
-}
 
 pub async fn run(
     mut msg_stream: MessageStream,
@@ -218,7 +140,7 @@ pub async fn handle_text_message(
                     handle_search_event(search_req, org_id, user_id, req_id, path.clone()).await;
                 }
                 #[cfg(feature = "enterprise")]
-                WsClientEvents::Cancel { trace_id } => {
+                WsClientEvents::Cancel { trace_id, .. } => {
                     // First handle the cancel event
                     // send a cancel flag to the search task
                     if let Err(e) = handle_cancel_event(&trace_id).await {
@@ -232,7 +154,7 @@ pub async fn handle_text_message(
                         search_registry_utils::is_cancelled(&trace_id)
                     );
 
-                    let res = search::handle_cancel(&trace_id, org_id).await;
+                    let res = handle_cancel(&trace_id, org_id).await;
                     let _ = send_message(req_id, res.to_json().to_string()).await;
                     let close_reason = Some(CloseReason {
                         code: CloseCode::Normal,
@@ -240,7 +162,11 @@ pub async fn handle_text_message(
                     });
 
                     #[cfg(feature = "enterprise")]
-                    let client_msg = WsClientEvents::Cancel { trace_id };
+                    let client_msg = WsClientEvents::Cancel {
+                        trace_id,
+                        org_id: org_id.to_string(),
+                        user_id: None,
+                    };
 
                     // Add audit before closing
                     #[cfg(feature = "enterprise")]
@@ -287,6 +213,18 @@ pub async fn handle_text_message(
                     });
                     cleanup_and_close_session(req_id, close_reason).await;
                 }
+                WsClientEvents::TestAbnormalClose { req_id } => {
+                    log::info!(
+                        "[WS_HANDLER]: Request Id: {} Node Role: {} Test Abnormal Close",
+                        req_id,
+                        get_config().common.node_role
+                    );
+                    let close_reason = Some(CloseReason {
+                        code: CloseCode::Error,
+                        description: None,
+                    });
+                    cleanup_and_close_session(&req_id, close_reason).await;
+                }
             }
         }
         Err(e) => {
@@ -296,7 +234,12 @@ pub async fn handle_text_message(
                 msg,
                 e
             );
-            let err_res = WsServerEvents::error_response(e.into(), Some(req_id.to_string()), None);
+            let err_res = WsServerEvents::error_response(
+                e.into(),
+                Some(req_id.to_string()),
+                None,
+                Default::default(),
+            );
             let _ = send_message(req_id, err_res.to_json().to_string()).await;
             let close_reason = Some(CloseReason {
                 code: CloseCode::Error,
@@ -385,7 +328,7 @@ async fn handle_search_event(
     let client_msg = WsClientEvents::Search(Box::new(search_req.clone()));
 
     // Register running search BEFORE spawning the search task
-    SEARCH_REGISTRY.insert(
+    WS_SEARCH_REGISTRY.insert(
         trace_id.clone(),
         SearchState::Running {
             cancel_tx,
@@ -402,7 +345,7 @@ async fn handle_search_event(
         // The task will also close the session
         // The task will also cleanup the search resources
         tokio::select! {
-            search_result = search::handle_search_request(
+            search_result = handle_search_request(
                 &req_id,
                 &mut accumulated_results,
                 &org_id,
@@ -411,7 +354,7 @@ async fn handle_search_event(
             ) => {
                 match search_result {
                     Ok(_) => {
-                        if let Some(mut state) = SEARCH_REGISTRY.get_mut(&trace_id_for_task) {
+                        if let Some(mut state) = WS_SEARCH_REGISTRY.get_mut(&trace_id_for_task) {
                             *state = SearchState::Completed {
                                 req_id: req_id.to_string(),
                             };
@@ -483,7 +426,7 @@ async fn handle_search_event(
 // Cancel handler
 #[cfg(feature = "enterprise")]
 async fn handle_cancel_event(trace_id: &str) -> Result<(), anyhow::Error> {
-    if let Some(mut entry) = SEARCH_REGISTRY.get_mut(trace_id) {
+    if let Some(mut entry) = WS_SEARCH_REGISTRY.get_mut(trace_id) {
         let state = entry.value_mut();
         let cancel_tx = match state {
             SearchState::Running { cancel_tx, .. } => cancel_tx.clone(),
@@ -516,7 +459,7 @@ async fn handle_search_error(e: Error, req_id: &str, trace_id: &str) -> Option<C
             trace_id
         );
         // Update state to cancelled before returning
-        if let Some(mut state) = SEARCH_REGISTRY.get_mut(trace_id) {
+        if let Some(mut state) = WS_SEARCH_REGISTRY.get_mut(trace_id) {
             *state = SearchState::Cancelled {
                 req_id: trace_id.to_string(),
             };
@@ -526,8 +469,12 @@ async fn handle_search_error(e: Error, req_id: &str, trace_id: &str) -> Option<C
 
     log::error!("[WS_HANDLER]: trace_id: {} Search error: {}", trace_id, e);
     // Send error response
-    let err_res =
-        WsServerEvents::error_response(e, Some(trace_id.to_string()), Some(req_id.to_string()));
+    let err_res = WsServerEvents::error_response(
+        e,
+        Some(trace_id.to_string()),
+        Some(req_id.to_string()),
+        Default::default(),
+    );
     let _ = send_message(req_id, err_res.to_json().to_string()).await;
 
     // Close with error
@@ -537,7 +484,7 @@ async fn handle_search_error(e: Error, req_id: &str, trace_id: &str) -> Option<C
     };
 
     // Update registry state
-    if let Some(mut state) = SEARCH_REGISTRY.get_mut(trace_id) {
+    if let Some(mut state) = WS_SEARCH_REGISTRY.get_mut(trace_id) {
         *state = SearchState::Completed {
             req_id: trace_id.to_string(),
         };
@@ -548,6 +495,6 @@ async fn handle_search_error(e: Error, req_id: &str, trace_id: &str) -> Option<C
 
 // Add cleanup function
 async fn cleanup_search_resources(trace_id: &str) {
-    SEARCH_REGISTRY.remove(trace_id);
+    WS_SEARCH_REGISTRY.remove(trace_id);
     log::debug!("[WS_HANDLER]: trace_id: {}, Resources cleaned up", trace_id);
 }

src/handler/http/request/ws_v2/mod.rs

@@ -0,0 +1,91 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+pub mod session;
+
+use actix_web::{Error, HttpRequest, HttpResponse, get, web};
+use config::{cluster::LOCAL_NODE, get_config};
+
+#[cfg(feature = "enterprise")]
+use crate::common::utils::auth::extract_auth_expiry_and_user_id;
+use crate::{
+    handler::http::request::ws_v2::session::WsSession,
+    service::websocket_events::sessions_cache_utils,
+};
+
+#[get("{org_id}/ws/v2/{router_id}")]
+pub async fn websocket(
+    path_params: web::Path<(String, String)>,
+    req: HttpRequest,
+    stream: web::Payload,
+    in_req: HttpRequest,
+) -> Result<HttpResponse, Error> {
+    let cfg = get_config();
+
+    if !cfg.websocket.enabled {
+        log::info!(
+            "[WS_HANDLER]: Node Role: {} Websocket is disabled",
+            cfg.common.node_role
+        );
+        return Ok(HttpResponse::NotFound().body("WebSocket is disabled"));
+    }
+
+    let (_, router_id) = path_params.into_inner();
+
+    let prefix = format!("{}/api/", get_config().common.base_uri);
+    let path = req.path().strip_prefix(&prefix).unwrap().to_string();
+
+    let user_id = in_req
+        .headers()
+        .get("user_id")
+        .and_then(|v| v.to_str().ok())
+        .unwrap_or("")
+        .to_string();
+
+    #[cfg(feature = "enterprise")]
+    let cookie_expiry = if cfg.websocket.check_cookie_expiry && cfg.common.local_mode {
+        let (cookie_expiry, _) = extract_auth_expiry_and_user_id(&req).await;
+        cookie_expiry.map(|expiry| expiry.timestamp_micros())
+    } else {
+        None
+    };
+
+    #[cfg(not(feature = "enterprise"))]
+    let cookie_expiry = None;
+
+    let (res, session, msg_stream) = actix_ws::handle(&req, stream)?;
+
+    let ws_session = WsSession::new(session, cookie_expiry);
+    sessions_cache_utils::insert_session(&router_id, ws_session);
+    log::info!(
+        "[WS_HANDLER]: Node Role: {} Got websocket request for router_id: {}",
+        cfg.common.node_role,
+        router_id,
+    );
+
+    // Spawn the handler
+    actix_web::rt::spawn(session::run(msg_stream, user_id, router_id, path));
+
+    Ok(res)
+}
+
+/// Initialize the job init for websocket
+pub async fn init() -> Result<(), anyhow::Error> {
+    // Run the garbage collector for websocket sessions
+    if LOCAL_NODE.is_querier() {
+        sessions_cache_utils::run_gc_ws_sessions().await;
+    }
+    Ok(())
+}

src/handler/http/request/ws_v2/session.rs

@@ -0,0 +1,622 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::time::Duration;
+
+use actix_http::ws::{CloseCode, CloseReason};
+use actix_ws::{MessageStream, Session};
+use config::{
+    get_config,
+    meta::websocket::{SearchEventReq, SearchResultType},
+};
+use futures::StreamExt;
+use infra::errors::{self, Error};
+#[cfg(feature = "enterprise")]
+use o2_enterprise::enterprise::common::{
+    auditor::{AuditMessage, Protocol, WsMeta},
+    infra::config::get_config as get_o2_config,
+};
+use rand::prelude::SliceRandom;
+use tokio::sync::mpsc;
+
+#[cfg(feature = "enterprise")]
+use crate::service::{
+    self_reporting::audit,
+    websocket_events::{handle_cancel, search_registry_utils},
+};
+use crate::{
+    common::infra::config::WS_SEARCH_REGISTRY,
+    // router::http::ws_v2::types::StreamMessage,
+    service::websocket_events::{
+        WsClientEvents, WsServerEvents, handle_search_request, search_registry_utils::SearchState,
+        sessions_cache_utils,
+    },
+};
+
+// Do not clone the session, instead use a reference to the session
+pub struct WsSession {
+    inner: Option<Session>,
+    // Utc timestamp in microseconds
+    last_activity_ts: i64,
+    // Utc timestamp in microseconds
+    created_ts: i64,
+    // Utc timestamp in microseconds
+    cookie_expiry: Option<i64>,
+}
+
+impl WsSession {
+    pub fn new(inner: Session, cookie_expiry: Option<i64>) -> Self {
+        let now = chrono::Utc::now().timestamp_micros();
+        Self {
+            inner: Some(inner),
+            last_activity_ts: now,
+            created_ts: now,
+            cookie_expiry,
+        }
+    }
+
+    pub fn update_activity(&mut self) {
+        self.last_activity_ts = chrono::Utc::now().timestamp_micros();
+    }
+
+    pub fn is_expired(&self) -> bool {
+        let cfg = get_config();
+        let now = chrono::Utc::now().timestamp_micros();
+        let idle_timeout_micros = cfg.websocket.session_idle_timeout_secs * 1_000_000;
+        let max_lifetime_micros = cfg.websocket.session_max_lifetime_secs * 1_000_000;
+
+        // 1. if the session has been idle for too long
+        // 2. if the session has exceeded the max lifetime
+        (now - self.last_activity_ts) > idle_timeout_micros
+            || (now - self.created_ts) > max_lifetime_micros
+    }
+
+    pub fn is_client_cookie_valid(&self) -> bool {
+        let now = chrono::Utc::now().timestamp_micros();
+        self.cookie_expiry.is_none_or(|expiry| expiry > now)
+    }
+
+    /// Send a text message to the client
+    pub async fn text(&mut self, msg: String) -> Result<(), actix_ws::Closed> {
+        self.update_activity();
+        if let Some(ref mut session) = self.inner {
+            session.text(msg).await
+        } else {
+            Err(actix_ws::Closed)
+        }
+    }
+
+    /// Close the session with a reason
+    pub async fn close(&mut self, reason: Option<CloseReason>) -> Result<(), actix_ws::Closed> {
+        self.update_activity();
+        if let Some(session) = self.inner.take() {
+            session.close(reason).await
+        } else {
+            Err(actix_ws::Closed)
+        }
+    }
+
+    /// Send a pong response
+    pub async fn pong(&mut self, payload: &[u8]) -> Result<(), actix_ws::Closed> {
+        self.update_activity();
+        if let Some(ref mut session) = self.inner {
+            session.pong(payload).await
+        } else {
+            Err(actix_ws::Closed)
+        }
+    }
+
+    /// Send a ping request
+    pub async fn ping(&mut self, payload: &[u8]) -> Result<(), actix_ws::Closed> {
+        self.update_activity();
+        if let Some(ref mut session) = self.inner {
+            session.ping(payload).await
+        } else {
+            Err(actix_ws::Closed)
+        }
+    }
+}
+
+pub async fn run(mut msg_stream: MessageStream, user_id: String, req_id: String, path: String) {
+    let cfg = get_config();
+    let mut ping_interval =
+        tokio::time::interval(Duration::from_secs(cfg.websocket.ping_interval_secs as u64));
+    let mut close_reason: Option<CloseReason> = None;
+
+    loop {
+        tokio::select! {
+            Some(msg) = msg_stream.next() => {
+                // Update activity and check cookie_expiry on any message
+                if let Some(mut session) = sessions_cache_utils::get_mut_session(&req_id) {
+                    session.update_activity();
+                    if !session.is_client_cookie_valid() {
+                        log::error!("[WS_HANDLER]: Session cookie expired");
+                        break;
+                    }
+                }
+
+                match msg {
+                    Ok(actix_ws::Message::Ping(bytes)) => {
+                        if let Some(mut session) = sessions_cache_utils::get_mut_session(&req_id) {
+                            if let Err(e) = session.pong(&bytes).await {
+                                log::error!("[WS_HANDLER]: Failed to send pong: {}", e);
+                                break;
+                            }
+                        } else {
+                            log::error!("[WS_HANDLER]: Session not found for ping response");
+                            break;
+                        }
+                    }
+                    Ok(actix_ws::Message::Pong(_)) => {
+                        log::debug!("[WS_HANDLER] Received pong from {}", req_id);
+                    }
+                    Ok(actix_ws::Message::Text(msg)) => {
+                        log::info!("[WS_HANDLER]: Request Id: {} Node Role: {} Received message: {}",
+                            req_id,
+                            get_config().common.node_role,
+                            msg
+                        );
+                        handle_text_message(&user_id, &req_id, msg.to_string(), path.clone()).await;
+                    }
+                    Ok(actix_ws::Message::Close(reason)) => {
+                        if let Some(reason) = reason.as_ref() {
+                            match reason.code {
+                                CloseCode::Normal | CloseCode::Error => {
+                                    log::info!("[WS_HANDLER]: Request Id: {} Node Role: {} Closing connection with reason: {:?}",
+                                        req_id,
+                                        get_config().common.node_role,
+                                        reason
+                                    );
+                                },
+                                _ => {
+                                    log::error!("[WS_HANDLER]: Request Id: {} Node Role: {} Abnormal closure with reason: {:?}",req_id,get_config().common.node_role,reason);
+                                },
+                            }
+                        }
+                        close_reason = reason;
+                        break;
+                    }
+                    _ => ()
+                }
+            }
+            // Heartbeat to keep the connection alive
+            _ = ping_interval.tick() => {
+                if let Some(mut session) = sessions_cache_utils::get_mut_session(&req_id) {
+                    if let Err(e) = session.ping(&[]).await {
+                        log::error!("[WS_HANDLER] Failed to send ping: {}", e);
+                        break;
+                    }
+                }
+            }
+        }
+    }
+    cleanup_and_close_session(&req_id, close_reason).await;
+}
+
+/// Handle the incoming text message
+/// Text message is parsed into `WsClientEvents` and processed accordingly
+/// Depending on each event type, audit must be done
+/// Currently audit is done only for the search event
+pub async fn handle_text_message(user_id: &str, req_id: &str, msg: String, path: String) {
+    match serde_json::from_str::<WsClientEvents>(&msg) {
+        Ok(client_msg) => {
+            // Validate the events
+            if !client_msg.is_valid() {
+                log::error!("[WS_HANDLER]: Invalid event: {:?}", client_msg);
+                let err_res = WsServerEvents::error_response(
+                    errors::Error::Message("Invalid event".to_string()),
+                    Some(req_id.to_string()),
+                    None,
+                    Default::default(),
+                );
+                let _ = send_message(req_id, err_res.to_json()).await;
+                return;
+            }
+
+            match client_msg {
+                WsClientEvents::Search(ref search_req) => {
+                    #[allow(unused_mut)]
+                    let mut user_id = user_id;
+                    // verify user_id for handling stream permissions
+                    #[cfg(feature = "enterprise")]
+                    {
+                        user_id = if get_config().common.local_mode {
+                            // single node mode, use user_id from the http request header
+                            user_id
+                        } else {
+                            // cluster mode, use user_id from the ws event added by the router
+                            if let Some(user_id) = &search_req.user_id {
+                                user_id
+                            } else {
+                                log::error!("[WS_HANDLER]: User id not found in search request");
+                                let err_res = WsServerEvents::error_response(
+                                    errors::Error::Message(
+                                        "User id not found in search request".to_string(),
+                                    ),
+                                    Some(req_id.to_string()),
+                                    Some(search_req.trace_id.clone()),
+                                    Default::default(),
+                                );
+                                let _ = send_message(req_id, err_res.to_json()).await;
+                                return;
+                            }
+                        };
+                    }
+                    handle_search_event(
+                        search_req,
+                        &search_req.org_id,
+                        user_id,
+                        req_id,
+                        path.clone(),
+                    )
+                    .await;
+                }
+                #[cfg(feature = "enterprise")]
+                WsClientEvents::Cancel {
+                    trace_id, org_id, ..
+                } => {
+                    if org_id.is_empty() {
+                        log::error!(
+                            "[WS_HANDLER]: Request Id: {} Node Role: {} Org id not found",
+                            req_id,
+                            get_config().common.node_role
+                        );
+                        return;
+                    };
+
+                    // First handle the cancel event
+                    // send a cancel flag to the search task
+                    if let Err(e) = handle_cancel_event(&trace_id).await {
+                        log::warn!("[WS_HANDLER]: Error in cancelling : {}", e);
+                        return;
+                    }
+
+                    log::info!(
+                        "[WS_HANDLER]: trace_id: {}, Cancellation flag set to: {:?}",
+                        trace_id,
+                        search_registry_utils::is_cancelled(&trace_id)
+                    );
+
+                    let res = handle_cancel(&trace_id, &org_id).await;
+                    let _ = send_message(req_id, res.to_json()).await;
+
+                    // Only used for audit
+                    #[cfg(feature = "enterprise")]
+                    let client_msg = WsClientEvents::Cancel {
+                        trace_id,
+                        org_id: org_id.to_string(),
+                        // setting user_id to None to handle PII
+                        user_id: None,
+                    };
+
+                    // Add audit before closing
+                    #[cfg(feature = "enterprise")]
+                    let is_audit_enabled = get_o2_config().common.audit_enabled;
+
+                    #[cfg(feature = "enterprise")]
+                    if is_audit_enabled {
+                        audit(AuditMessage {
+                            user_email: user_id.to_string(),
+                            org_id: org_id.to_string(),
+                            _timestamp: chrono::Utc::now().timestamp(),
+                            protocol: Protocol::Ws(WsMeta {
+                                path: path.clone(),
+                                message_type: client_msg.get_type(),
+                                content: client_msg.to_json(),
+                                close_reason: "".to_string(),
+                            }),
+                        })
+                        .await;
+                    }
+                }
+                WsClientEvents::Benchmark { id } => {
+                    // simulate random delay for benchmarking by sleep for 10/20/30/60/90
+                    // seconds
+                    let delay: Vec<u64> = vec![10, 20, 30, 60, 90];
+                    let delay = delay.choose(&mut rand::thread_rng()).unwrap();
+                    log::info!(
+                        "[WS_HANDLER]: Sleeping for benchmark, id: {}, delay: {}",
+                        id,
+                        delay
+                    );
+                    tokio::time::sleep(tokio::time::Duration::from_secs(*delay)).await;
+
+                    let response = serde_json::json!({
+                        "id": id,
+                        "took": delay,
+                    });
+                    let _ = send_message(req_id, response.to_string()).await;
+                    let close_reason = Some(CloseReason {
+                        code: CloseCode::Normal,
+                        description: None,
+                    });
+                    cleanup_and_close_session(req_id, close_reason).await;
+                }
+                WsClientEvents::TestAbnormalClose { req_id } => {
+                    log::info!(
+                        "[WS_HANDLER]: Request Id: {} Node Role: {} Test Abnormal Close",
+                        req_id,
+                        get_config().common.node_role
+                    );
+                    let close_reason = Some(CloseReason {
+                        code: CloseCode::Error,
+                        description: None,
+                    });
+                    cleanup_and_close_session(&req_id, close_reason).await;
+                }
+            }
+        }
+        Err(e) => {
+            log::error!(
+                "[WS_HANDLER]: Request id: {} Failed to parse message: {:?}, error: {:?}",
+                req_id,
+                msg,
+                e
+            );
+            let err_res = WsServerEvents::error_response(
+                e.into(),
+                Some(req_id.to_string()),
+                None,
+                Default::default(),
+            );
+            let _ = send_message(req_id, err_res.to_json()).await;
+            let close_reason = Some(CloseReason {
+                code: CloseCode::Error,
+                description: None,
+            });
+            let mut session = if let Some(session) = sessions_cache_utils::get_mut_session(req_id) {
+                session
+            } else {
+                log::error!("[WS_HANDLER]: req_id: {} session not found", req_id);
+                return;
+            };
+            let _ = session.close(close_reason).await;
+        }
+    }
+}
+
+pub async fn send_message(req_id: &str, msg: String) -> Result<(), Error> {
+    let mut session = if let Some(session) = sessions_cache_utils::get_mut_session(req_id) {
+        session
+    } else {
+        return Err(Error::Message(format!(
+            "[req_id {}] session not found",
+            req_id
+        )));
+    };
+
+    log::debug!("[WS_HANDLER]: req_id: {} sending message: {}", req_id, msg);
+
+    session.text(msg).await.map_err(|e| {
+        log::error!("[WS_HANDLER]: Failed to send message: {:?}", e);
+        Error::Message(e.to_string())
+    })
+}
+
+async fn cleanup_and_close_session(req_id: &str, close_reason: Option<CloseReason>) {
+    if let Some(mut session) = sessions_cache_utils::get_mut_session(req_id) {
+        if let Some(reason) = close_reason.as_ref() {
+            log::info!(
+                "[WS_HANDLER]: req_id: {} Closing session with reason: {:?}",
+                req_id,
+                reason
+            );
+        }
+
+        // Attempt to close the session
+        if let Err(e) = session.close(close_reason).await {
+            log::error!(
+                "[WS_HANDLER]: req_id: {} Failed to close session gracefully: {:?}",
+                req_id,
+                e
+            );
+        }
+    }
+
+    // Remove the session from the cache
+    sessions_cache_utils::remove_session(req_id);
+    log::info!(
+        "[WS_HANDLER]: req_id: {} Session removed from cache. Remaining sessions: {}",
+        req_id,
+        sessions_cache_utils::len_sessions()
+    );
+}
+
+// Main search handler
+async fn handle_search_event(
+    search_req: &SearchEventReq,
+    org_id: &str,
+    user_id: &str,
+    req_id: &str,
+    #[allow(unused_variables)] path: String,
+) {
+    let (cancel_tx, mut cancel_rx) = mpsc::channel(1);
+    let mut accumulated_results: Vec<SearchResultType> = Vec::new();
+
+    let org_id = org_id.to_string();
+    let user_id = user_id.to_string();
+    let req_id = req_id.to_string();
+    let trace_id = search_req.trace_id.clone();
+    let trace_id_for_task = trace_id.clone();
+    let search_req = search_req.clone();
+
+    #[cfg(feature = "enterprise")]
+    let is_audit_enabled = get_o2_config().common.audit_enabled;
+
+    #[cfg(feature = "enterprise")]
+    let client_msg = WsClientEvents::Search(Box::new(search_req.clone()));
+
+    // Register running search BEFORE spawning the search task
+    WS_SEARCH_REGISTRY.insert(
+        trace_id.clone(),
+        SearchState::Running {
+            cancel_tx,
+            req_id: req_id.to_string(),
+        },
+    );
+
+    // Spawn the search task
+    tokio::spawn(async move {
+        // Handle the search request
+        // If search is cancelled, the task will exit
+        // Otherwise, the task will complete and the results will be sent to the client
+        // The task will also update the search state to completed
+        // The task will also close the session
+        // The task will also cleanup the search resources
+        tokio::select! {
+            search_result = handle_search_request(
+                &req_id,
+                &mut accumulated_results,
+                &org_id,
+                &user_id,
+                search_req.clone(),
+            ) => {
+                match search_result {
+                    Ok(_) => {
+                        if let Some(mut state) = WS_SEARCH_REGISTRY.get_mut(&trace_id_for_task) {
+                            *state = SearchState::Completed {
+                                req_id: req_id.to_string(),
+                            };
+                        }
+
+                        // Add audit before closing
+                        #[cfg(feature = "enterprise")]
+                        if is_audit_enabled {
+                            audit(AuditMessage {
+                                user_email: user_id,
+                                org_id,
+                                _timestamp: chrono::Utc::now().timestamp(),
+                                protocol: Protocol::Ws(WsMeta {
+                                    path: path.clone(),
+                                    message_type: client_msg.get_type(),
+                                    content: client_msg.to_json(),
+                                    close_reason: "".to_string(),
+                                }),
+                            })
+                            .await;
+                        }
+
+                        cleanup_search_resources(&trace_id_for_task).await;
+                    }
+                    Err(e) => {
+                        let _ = handle_search_error(e, &req_id, &trace_id_for_task).await;
+                        // Add audit before closing
+                        #[cfg(feature = "enterprise")]
+                        if is_audit_enabled {
+                          audit(AuditMessage {
+                                  user_email: user_id,
+                                  org_id,
+                                  _timestamp: chrono::Utc::now().timestamp(),
+                                  protocol: Protocol::Ws(WsMeta {
+                                      path: path.clone(),
+                                      message_type: client_msg.get_type(),
+                                      content: client_msg.to_json(),
+                                      close_reason: "".to_string(),
+                                  }),
+                              })
+                              .await;
+                        }
+
+                        // Even if the search is cancelled, we need to cleanup the resources
+                        cleanup_search_resources(&trace_id_for_task).await;
+                    }
+                }
+            }
+            _ = cancel_rx.recv() => {
+                // if search is cancelled, update the state
+                // the cancel handler will close the session
+
+                // Just cleanup resources when cancelled
+                cleanup_search_resources(&trace_id_for_task).await;
+            }
+        }
+    });
+}
+
+// Cancel handler
+#[cfg(feature = "enterprise")]
+async fn handle_cancel_event(trace_id: &str) -> Result<(), anyhow::Error> {
+    if let Some(mut entry) = WS_SEARCH_REGISTRY.get_mut(trace_id) {
+        let state = entry.value_mut();
+        let cancel_tx = match state {
+            SearchState::Running { cancel_tx, .. } => cancel_tx.clone(),
+            state => {
+                let err_msg = format!("Cannot cancel search in state: {:?}", state);
+                log::warn!("[WS_HANDLER]: {}", err_msg);
+                return Err(anyhow::anyhow!(err_msg));
+            }
+        };
+
+        *entry.value_mut() = SearchState::Cancelled {
+            req_id: trace_id.to_string(),
+        };
+
+        if let Err(e) = cancel_tx.send(()).await {
+            log::error!("[WS_HANDLER]: Failed to send cancel signal: {}", e);
+        }
+
+        log::info!("[WS_HANDLER]: Search cancelled for trace_id: {}", trace_id);
+    }
+    Ok(())
+}
+
+async fn handle_search_error(e: Error, req_id: &str, trace_id: &str) -> Option<CloseReason> {
+    // if the error is due to search cancellation, return.
+    // the cancel handler will close the session
+    if let errors::Error::ErrorCode(errors::ErrorCodes::SearchCancelQuery(_)) = e {
+        log::info!(
+            "[WS_HANDLER]: trace_id: {}, Return from search handler, search canceled",
+            trace_id
+        );
+        // Update state to cancelled before returning
+        if let Some(mut state) = WS_SEARCH_REGISTRY.get_mut(trace_id) {
+            *state = SearchState::Cancelled {
+                req_id: trace_id.to_string(),
+            };
+        }
+        return None;
+    }
+
+    log::error!("[WS_HANDLER]: trace_id: {} Search error: {}", trace_id, e);
+    // Send error response
+    let err_res = WsServerEvents::error_response(
+        e,
+        Some(req_id.to_string()),
+        Some(trace_id.to_string()),
+        Default::default(),
+    );
+    let _ = send_message(req_id, err_res.to_json()).await;
+
+    // Close with error
+    let close_reason = CloseReason {
+        code: CloseCode::Normal,
+        description: None,
+    };
+
+    // Update registry state
+    if let Some(mut state) = WS_SEARCH_REGISTRY.get_mut(trace_id) {
+        *state = SearchState::Completed {
+            req_id: trace_id.to_string(),
+        };
+    }
+
+    Some(close_reason)
+}
+
+// Add cleanup function
+async fn cleanup_search_resources(trace_id: &str) {
+    WS_SEARCH_REGISTRY.remove(trace_id);
+    log::debug!("[WS_HANDLER]: trace_id: {}, Resources cleaned up", trace_id);
+}

src/handler/http/router/mod.rs

@@ -527,7 +527,8 @@ pub fn get_service_routes(svc: &mut web::ServiceConfig) {
         .service(service_accounts::delete)
         .service(service_accounts::update)
         .service(service_accounts::get_api_token)
-        .service(websocket::websocket);
+        .service(websocket::websocket)
+        .service(ws_v2::websocket);
 
     #[cfg(feature = "enterprise")]
     let service = service

src/router/http/mod.rs

@@ -27,10 +27,12 @@ use actix_web::{
     route, web,
 };
 use hashbrown::HashMap;
+pub use ws_v2::remove_querier_from_handler;
 
 use crate::common::{infra::cluster, utils::http::get_search_type_from_request};
 
 mod ws;
+pub(crate) mod ws_v2;
 
 const QUERIER_ROUTES: [&str; 20] = [
     "/config",
@@ -194,17 +196,17 @@ async fn dispatch(
             .body(new_url.error.unwrap_or("internal server error".to_string())));
     }
 
+    // check if the request is a websocket request
+    let path_columns: Vec<&str> = path.split('/').collect();
+    if path_columns.get(3).unwrap_or(&"").starts_with("ws") {
+        return proxy_ws(req, payload, new_url, start).await;
+    }
+
     // check if the request need to be proxied by body
     if cfg.common.metrics_cache_enabled && is_querier_route_by_body(&path) {
         return proxy_querier_by_body(req, payload, client, new_url, start).await;
     }
 
-    // check if the request is a websocket request
-    let path_columns: Vec<&str> = path.split('/').collect();
-    if *path_columns.get(3).unwrap_or(&"") == "ws" {
-        return proxy_ws(req, payload, new_url, start).await;
-    }
-
     // send query
     default_proxy(req, payload, client, new_url, start).await
 }
@@ -428,31 +430,58 @@ async fn proxy_ws(
 ) -> actix_web::Result<HttpResponse, Error> {
     let cfg = get_config();
     if cfg.websocket.enabled {
-        // Convert the HTTP/HTTPS URL to a WebSocket URL (WS/WSS)
-        let ws_url = match ws::convert_to_websocket_url(&new_url.full_url) {
-            Ok(url) => url,
-            Err(e) => {
-                log::error!("Error converting URL to WebSocket: {:?}", e);
-                return Ok(HttpResponse::BadRequest()
-                    .force_close()
-                    .body("Invalid WebSocket URL"));
-            }
-        };
+        // Check if this is a WebSocket v2 request (e.g., contains a specific path segment or
+        // header)
+        let path = req.uri().path();
+        // Extract client ID from the path or query parameters
+        // Path format example: /api/{org_id}/ws/v2/{client_id}
+        if path.contains("/ws/v2/") {
+            let path_parts: Vec<&str> = path.split('/').collect();
+            let client_id = path_parts[path_parts.len() - 1].to_string();
 
-        match ws::ws_proxy(req, payload, &ws_url).await {
-            Ok(res) => {
-                log::info!(
-                    "[WS_ROUTER] Successfully proxied WebSocket connection to backend: {}, took: {} ms",
-                    ws_url,
-                    start.elapsed().as_millis()
-                );
-                Ok(res)
+            log::info!(
+                "[WS_V2_ROUTER] Handling WS v2 connection for client: {}, took: {} ms",
+                client_id,
+                start.elapsed().as_millis()
+            );
+
+            // Use the WebSocket v2 handler
+            let ws_handler = ws_v2::get_ws_handler().await;
+            match ws_handler.handle_connection(req, payload, client_id).await {
+                Ok(response) => Ok(response),
+                Err(e) => {
+                    log::error!("[WS_V2_ROUTER] failed: {}", e);
+                    Ok(HttpResponse::InternalServerError().body("WebSocket v2 error"))
+                }
             }
-            Err(e) => {
-                log::error!("[WS_ROUTER] failed: {:?}", e);
-                Ok(HttpResponse::InternalServerError()
-                    .force_close()
-                    .body("WebSocket proxy error"))
+        } else {
+            // Use the legacy WebSocket proxy implementation
+            // Convert the HTTP/HTTPS URL to a WebSocket URL (WS/WSS)
+            let ws_url = match ws::convert_to_websocket_url(&new_url.full_url) {
+                Ok(url) => url,
+                Err(e) => {
+                    log::error!("Error converting URL to WebSocket: {:?}", e);
+                    return Ok(HttpResponse::BadRequest()
+                        .force_close()
+                        .body("Invalid WebSocket URL"));
+                }
+            };
+
+            match ws::ws_proxy(req, payload, &ws_url).await {
+                Ok(res) => {
+                    log::info!(
+                        "[WS_ROUTER] Successfully proxied WebSocket connection to backend: {}, took: {} ms",
+                        ws_url,
+                        start.elapsed().as_millis()
+                    );
+                    Ok(res)
+                }
+                Err(e) => {
+                    log::error!("[WS_ROUTER] failed: {:?}", e);
+                    Ok(HttpResponse::InternalServerError()
+                        .force_close()
+                        .body("WebSocket proxy error"))
+                }
             }
         }
     } else {

src/router/http/ws.rs

@@ -298,7 +298,7 @@ impl Drop for WsProxySession {
 }
 
 /// Convert actix-web WebSocket message to tungstenite message format
-fn from_actix_message(msg: Message) -> tungstenite::protocol::Message {
+pub fn from_actix_message(msg: Message) -> tungstenite::protocol::Message {
     match msg {
         Message::Text(text) => tungstenite::protocol::Message::Text(text.to_string()),
         Message::Binary(bin) => tungstenite::protocol::Message::Binary(bin.to_vec()),
@@ -324,7 +324,7 @@ fn from_actix_message(msg: Message) -> tungstenite::protocol::Message {
 }
 
 /// Convert tungstenite WebSocket message to actix-web message format
-fn from_tungstenite_msg_to_actix_msg(msg: tungstenite::protocol::Message) -> Message {
+pub fn from_tungstenite_msg_to_actix_msg(msg: tungstenite::protocol::Message) -> Message {
     match msg {
         tungstenite::protocol::Message::Text(text) => {
             log_frame_details("Converting Text ->", &text, false);

src/router/http/ws_v2/connection.rs

@@ -0,0 +1,480 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::sync::{
+    Arc,
+    atomic::{AtomicBool, Ordering},
+};
+
+use actix_http::StatusCode;
+use async_trait::async_trait;
+use config::{RwAHashMap, get_config, utils::json};
+use futures_util::{
+    SinkExt, StreamExt,
+    stream::{SplitSink, SplitStream},
+};
+use reqwest::header::{HeaderName, HeaderValue};
+use tokio::{
+    net::TcpStream,
+    sync::{
+        Mutex,
+        mpsc::{Sender, channel},
+    },
+};
+use tokio_tungstenite::{
+    MaybeTlsStream, WebSocketStream, connect_async,
+    tungstenite::{self, client::IntoClientRequest, protocol::Message as WsMessage},
+};
+
+use crate::{
+    common::infra::cluster,
+    router::http::ws_v2::{
+        error::*,
+        handler::{QuerierName, TraceId},
+    },
+    service::websocket_events::{WsClientEvents, WsServerEvents},
+};
+
+type WsStreamType = WebSocketStream<MaybeTlsStream<TcpStream>>;
+
+#[async_trait]
+pub trait Connection: Send + Sync {
+    async fn connect(node_name: &str, http_url: &str) -> WsResult<Arc<Self>>;
+    async fn disconnect(&self);
+    async fn send_message(&self, message: WsClientEvents) -> WsResult<()>;
+    async fn is_connected(&self) -> bool;
+}
+
+#[derive(Debug)]
+pub struct QuerierConnection {
+    querier_name: QuerierName,
+    write: Arc<Mutex<Option<SplitSink<WsStreamType, WsMessage>>>>,
+    shutdown_tx: Sender<()>,
+    response_router: Arc<ResponseRouter>,
+    is_connected: Arc<AtomicBool>,
+}
+
+impl Drop for QuerierConnection {
+    fn drop(&mut self) {
+        log::info!(
+            "[WS::QuerierConnection] connection to querier {} closed",
+            self.querier_name
+        );
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct ResponseRouter {
+    routes: RwAHashMap<TraceId, Sender<WsServerEvents>>,
+}
+
+impl Drop for ResponseRouter {
+    fn drop(&mut self) {
+        log::info!("[WS::ResponseRouter] shutting down. Stop routing responses from querier");
+    }
+}
+
+impl ResponseRouter {
+    pub async fn flush(&self) {
+        let mut count = 0;
+        for (trace_id, sender) in self.routes.write().await.drain() {
+            // Send error response to client
+            let _ = sender
+                .send(WsServerEvents::Error {
+                    code: StatusCode::SERVICE_UNAVAILABLE.into(),
+                    message: "Querier connection closed".to_string(),
+                    error_detail: None,
+                    trace_id: Some(trace_id),
+                    request_id: None,
+                    should_client_retry: true, /* Indicate client should retry since the querier
+                                                * connection is closed */
+                })
+                .await;
+            count += 1;
+        }
+
+        log::info!("[WS::QuerierConnection::ResponseRouter] flushed {count} routes",);
+    }
+}
+
+pub async fn create_connection(querier_name: &QuerierName) -> WsResult<Arc<QuerierConnection>> {
+    // Get querier info from cluster
+    let node = cluster::get_cached_node_by_name(querier_name)
+        .await
+        .ok_or_else(|| WsError::QuerierNotAvailable(querier_name.to_string()))?;
+
+    let conn = QuerierConnection::connect(&node.name, &node.http_addr).await?;
+    Ok(conn)
+}
+
+impl QuerierConnection {
+    pub async fn register_request(&self, trace_id: TraceId, response_tx: Sender<WsServerEvents>) {
+        self.response_router
+            .routes
+            .write()
+            .await
+            .insert(trace_id, response_tx);
+    }
+
+    pub async fn unregister_request(&self, trace_id: &TraceId) {
+        let mut write_guard = self.response_router.routes.write().await;
+        write_guard.remove(trace_id);
+        write_guard.shrink_to_fit();
+        log::debug!("[WS::Connection] removed trace_id {trace_id} from response_router-routes");
+    }
+
+    async fn listen_to_querier_response(
+        &self,
+        mut read: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
+        mut shutdown_rx: tokio::sync::mpsc::Receiver<()>,
+    ) {
+        loop {
+            // Handle incoming messages from querier to router
+            tokio::select! {
+                Some(msg) = read.next() => {
+                    match msg {
+                        Ok(msg) => {
+                            match msg {
+                                tungstenite::protocol::Message::Close(reason) => {
+                                    log::debug!("[WS::Router::QuerierConnection] received close message: {:?}", reason);
+                                    break;
+                                }
+                                tungstenite::protocol::Message::Ping(ping) => {
+                                    log::debug!("[WS::Router::QuerierConnection] received ping message: {:?}", ping);
+                                    let pong = tungstenite::protocol::Message::Pong(ping);
+                                    let write_clone = self.write.clone();
+                                    let mut write_guard = write_clone.lock().await;
+                                    if let Some(write) = write_guard.as_mut() {
+                                        let _ = write.send(pong).await;
+                                    }
+                                }
+                                tungstenite::protocol::Message::Pong(pong) => {
+                                    log::debug!("[WS::Router::QuerierConnection] received pong message: {:?}", pong);
+                                }
+                                tungstenite::protocol::Message::Binary(binary) => {
+                                    log::debug!("[WS::Router::QuerierConnection] received binary message: {:?}", binary);
+                                }
+                                tungstenite::protocol::Message::Frame(frame) => {
+                                    log::debug!("[WS::Router::QuerierConnection] received raw frame from querier: {:?}", frame);
+                                }
+                                tungstenite::protocol::Message::Text(text) => {
+                                    let svr_event = match json::from_str::<WsServerEvents>(&text) {
+                                        Ok(event) => event,
+                                        Err(e) => {
+                                            log::error!(
+                                                "[WS::Router::QuerierConnection] Error parsing message received from querier: {}. Trying to get trace_id",
+                                                e
+                                            );
+                                            match json::from_str::<json::Value>(&text).map(|val| val.get("trace_id").map(ToString::to_string)) {
+                                                Err(e) => {
+                                                    log::error!(
+                                                        "[WS::Router::QuerierConnection] Error parsing trace_id from message received from querier: {}. Drop the message",
+                                                        e
+                                                    );
+                                                    // scenario 1 where the trace_id & sender are not cleaned up -> left for clean job
+                                                    continue;
+                                                }
+                                                Ok(trace_id) => {
+                                                    WsServerEvents::Error {
+                                                        code: StatusCode::INTERNAL_SERVER_ERROR.into(),
+                                                        message: "Failed to parse event received from querier".to_string(),
+                                                        error_detail: None,
+                                                        trace_id,
+                                                        request_id: None,
+                                                        should_client_retry: true,
+                                                    }
+                                                }
+                                            }
+                                        }
+                                    };
+                                    let remove_trace_id = svr_event.should_clean_trace_id();
+                                    if let Err(e) = self.response_router.route_response(svr_event.clone()).await {
+                                        // scenario 2 where the trace_id & sender are not cleaned up -> left for clean job
+                                        log::error!(
+                                            "[WS::Router::QuerierConnection] Error routing response from querier back to client socket: {}, message: {}",
+                                            e,
+                                            svr_event.to_json()
+                                        );
+                                    }
+                                    if let Some(trace_id) = remove_trace_id {
+                                        log::info!("[WS::Router::QuerierConnection] Unregistering trace_id: {}, svr_event: {:?}", trace_id, svr_event);
+                                        self.unregister_request(&trace_id).await;
+                                    }
+                                }
+                            }
+                        }
+                        Err(e) => {
+                            log::error!("[WS::Router::QuerierConnection] Read error: {}, Querier: {}", e, self.querier_name);
+                            break;
+                        }
+                    }
+                }
+                _ = shutdown_rx.recv() => {
+                    log::info!("[WS::Router::QuerierConnection] shutting down. Stop listening from the querier {}", self.querier_name);
+                    break;
+                }
+            }
+        }
+
+        // reaches here when connection is closed/error from the querier side
+        // mark the connection disconnected and exist
+        self.is_connected.store(false, Ordering::SeqCst);
+        // flush in case of any remaining trace_ids
+        self.response_router.flush().await;
+    }
+
+    async fn health_check(&self) {
+        let cfg = get_config();
+        let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(
+            cfg.websocket.health_check_interval as _,
+        ));
+
+        loop {
+            interval.tick().await;
+            let mut write_guard = self.write.lock().await;
+            match write_guard.as_mut() {
+                None => {
+                    break;
+                }
+                Some(w) => {
+                    if w.send(WsMessage::Ping(vec![])).await.is_err() {
+                        break;
+                    }
+                }
+            }
+        }
+
+        // mark the connection disconnected and exist
+        self.is_connected.store(false, Ordering::SeqCst);
+        // flush in case of any remaining trace_ids
+        self.response_router.flush().await;
+    }
+
+    pub async fn is_active_trace_id(&self, trace_id: &TraceId) -> bool {
+        self.response_router
+            .routes
+            .read()
+            .await
+            .contains_key(trace_id)
+    }
+}
+
+impl ResponseRouter {
+    pub fn new() -> Arc<Self> {
+        let response_router = Arc::new(Self {
+            routes: Default::default(),
+        });
+
+        // Spawn cleanup task
+        let response_router_cp = response_router.clone();
+        tokio::spawn(async move {
+            response_router_cp.spawn_cleanup_task().await;
+        });
+
+        response_router
+    }
+
+    /// Ideally all the registered request should be removed from the routes after done
+    /// The only
+    async fn spawn_cleanup_task(&self) {
+        let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(60));
+        loop {
+            interval.tick().await;
+            let mut write_guard = self.routes.write().await;
+            write_guard.retain(|trace_id, response_tx| {
+                if response_tx.is_closed() {
+                    log::debug!("[WS::QuerierConnection] channel closed for trace_id {trace_id}. Removed from routes");
+                    false
+                } else {
+                    true
+                }
+            });
+            write_guard.shrink_to_fit();
+        }
+    }
+
+    pub async fn route_response(&self, message: WsServerEvents) -> WsResult<()> {
+        let trace_id = message.get_trace_id();
+
+        match self.routes.read().await.get(&trace_id) {
+            None => Err(WsError::ResponseChannelNotFound(trace_id.clone())),
+            Some(resp_sender) => {
+                resp_sender
+                    .send(message)
+                    .await
+                    .map_err(|_| WsError::ResponseChannelClosed(trace_id.clone()))?;
+                Ok(())
+            }
+        }
+    }
+}
+
+#[async_trait]
+impl Connection for QuerierConnection {
+    async fn connect(node_name: &str, http_url: &str) -> WsResult<Arc<Self>> {
+        let ws_req = get_default_querier_request(http_url)?;
+
+        // Router -> Querier
+        let (ws_stream, _) = connect_async(ws_req).await.map_err(|e| {
+            log::error!(
+                "[WS::QuerierConnection] error connecting to querier {}: {}",
+                http_url,
+                e
+            );
+            WsError::ConnectionError(e.to_string())
+        })?;
+        let (write, read) = ws_stream.split();
+        let write = Arc::new(Mutex::new(Some(write)));
+
+        // Signal to shutdown
+        let (shutdown_tx, shutdown_rx) = channel::<()>(1);
+
+        // Setting up components needed for the two tasks
+        let response_router = ResponseRouter::new();
+        let is_connected = Arc::new(AtomicBool::new(true));
+
+        let conn: Arc<QuerierConnection> = Arc::new(Self {
+            querier_name: node_name.to_string(),
+            write,
+            shutdown_tx,
+            response_router,
+            is_connected,
+        });
+
+        // Spawn task to listen to querier responses
+        let conn_t1 = conn.clone();
+        tokio::spawn(async move {
+            let _ = conn_t1.listen_to_querier_response(read, shutdown_rx).await;
+            if !conn_t1.is_connected().await {
+                log::info!(
+                    "[WS::QuerierConnection] connection to querier {} closed",
+                    conn_t1.querier_name
+                );
+            }
+        });
+
+        // Spawn health check task
+        let conn_t2 = conn.clone();
+        tokio::spawn(async move {
+            let _ = conn_t2.health_check().await;
+            if !conn_t2.is_connected().await {
+                log::info!(
+                    "[WS::QuerierConnection] connection to querier {} closed",
+                    conn_t2.querier_name
+                );
+            }
+        });
+
+        Ok(conn)
+    }
+
+    async fn disconnect(&self) {
+        let mut write_guard = self.write.lock().await;
+        if let Some(write) = write_guard.as_mut() {
+            let close_frame = tungstenite::protocol::Message::Close(None);
+            let _ = write.send(close_frame).await;
+            _ = write.close().await;
+            *write_guard = None;
+        }
+        _ = self.shutdown_tx.send(()).await;
+    }
+
+    async fn send_message(&self, message: WsClientEvents) -> WsResult<()> {
+        let mut write_guard = self.write.lock().await;
+        if let Some(write) = write_guard.as_mut() {
+            let trace_id = message.get_trace_id();
+            // Convert `WsClientEvents` to `tungstenite::protocol::Message`
+            let message = tungstenite::protocol::Message::from(message);
+            match write.send(message).await {
+                Ok(_) => {
+                    log::info!(
+                        "[WS::QuerierConnection] request w/ trace_id {} successfully forwarded to querier {}",
+                        trace_id,
+                        self.querier_name
+                    );
+                    Ok(())
+                }
+                Err(e) => {
+                    log::error!(
+                        "[WS::QuerierConnection] error sending messages via connection {e}. Mark the connection disconnected"
+                    );
+                    self.is_connected.store(false, Ordering::SeqCst);
+                    Err(WsError::ConnectionError(format!(
+                        "[WS::QuerierConnection] error sending messages via connection {e}"
+                    )))
+                }
+            }
+        } else {
+            Err(WsError::ConnectionError("Not connected".into()))
+        }
+    }
+
+    async fn is_connected(&self) -> bool {
+        if self.write.lock().await.is_some() {
+            self.is_connected.load(Ordering::SeqCst)
+        } else {
+            false
+        }
+    }
+}
+
+/// Helper function to create a mock request used to establish websocket connection with querier
+fn get_default_querier_request(http_url: &str) -> WsResult<tungstenite::http::Request<()>> {
+    let mut parsed_url = match url::Url::parse(http_url) {
+        Ok(url) => url,
+        Err(e) => return Err(WsError::QuerierUrlInvalid(e.to_string())),
+    };
+
+    // Check the scheme and update it accordingly
+    match parsed_url.scheme() {
+        "http" => {
+            parsed_url
+                .set_scheme("ws")
+                .map_err(|_| WsError::QuerierWSUrlError("Failed to set scheme to ws".into()))?;
+        }
+        "https" => {
+            parsed_url
+                .set_scheme("wss")
+                .map_err(|_| WsError::QuerierWSUrlError("Failed to set scheme to wss".into()))?;
+        }
+        other_schema => {
+            return Err(WsError::QuerierWSUrlError(format!(
+                "Unsupported URL scheme: {}",
+                other_schema
+            )));
+        }
+    }
+
+    // TODO: v2 ws endpoint should be agnostic from `org` and `12345678`
+    let parsed_url = parsed_url.to_string() + "api/default/ws/v2/12345678";
+
+    let mut ws_req = parsed_url
+        .into_client_request()
+        .map_err(|e| WsError::ConnectionError(e.to_string()))?;
+
+    // additional headers to the req
+    let token = get_config().grpc.internal_grpc_token.clone();
+    ws_req.headers_mut().insert(
+        HeaderName::from_static("authorization"),
+        HeaderValue::from_str(&token).map_err(|e| WsError::ConnectionError(e.to_string()))?,
+    );
+    // Add hearder upgrade websocket
+    ws_req.headers_mut().insert(
+        HeaderName::from_static("upgrade"),
+        HeaderValue::from_str("websocket").map_err(|e| WsError::ConnectionError(e.to_string()))?,
+    );
+
+    Ok(ws_req)
+}

src/router/http/ws_v2/error.rs

@@ -0,0 +1,156 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+use actix_http::ws::ProtocolError;
+use actix_web::{error::ResponseError, http::StatusCode};
+use actix_ws::CloseReason;
+use serde_json::error::Error as SerdeError;
+use thiserror::Error;
+
+use crate::service::websocket_events::WsServerEvents;
+
+#[derive(Error, Debug)]
+pub enum WsError {
+    #[error("ProtocolError# {0}")]
+    ProtocolError(#[from] ProtocolError),
+
+    #[error("SerdeError# {0}")]
+    SerdeError(#[from] SerdeError),
+
+    #[error("Connection error: {0}")]
+    ConnectionError(String),
+
+    #[error("Pooled connection disconnected")]
+    ConnectionDisconnected,
+
+    #[error("Session not found: {0}")]
+    SessionNotFound(String),
+
+    #[error("Response channel not found: {0}")]
+    ResponseChannelNotFound(String),
+
+    #[error("Response channel registered for trace_id {0} not found")]
+    ResponseChannelClosed(String),
+
+    #[error("Querier not available: {0}")]
+    QuerierNotAvailable(String),
+
+    #[error("Querier http url {0} not valid")]
+    QuerierUrlInvalid(String),
+
+    #[error("Querier WS url error: {0}")]
+    QuerierWSUrlError(String),
+
+    #[error(transparent)]
+    Other(#[from] anyhow::Error),
+}
+
+impl ResponseError for WsError {
+    fn status_code(&self) -> StatusCode {
+        match self {
+            WsError::ProtocolError(_) => StatusCode::BAD_REQUEST,
+            WsError::SerdeError(_) => StatusCode::BAD_REQUEST,
+            WsError::ConnectionError(_) => StatusCode::SERVICE_UNAVAILABLE,
+            WsError::ConnectionDisconnected => StatusCode::INTERNAL_SERVER_ERROR,
+            WsError::SessionNotFound(_) => StatusCode::BAD_REQUEST,
+            WsError::QuerierUrlInvalid(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            WsError::QuerierWSUrlError(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            WsError::QuerierNotAvailable(_) => StatusCode::SERVICE_UNAVAILABLE,
+            WsError::Other(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            WsError::ResponseChannelNotFound(_) => StatusCode::INTERNAL_SERVER_ERROR,
+            WsError::ResponseChannelClosed(_) => StatusCode::INTERNAL_SERVER_ERROR,
+        }
+    }
+}
+
+pub type WsResult<T> = Result<T, WsError>;
+
+#[derive(Debug)]
+pub enum DisconnectMessage {
+    Error(ErrorMessage),
+    Close(Option<CloseReason>),
+}
+
+#[derive(Debug)]
+pub struct ErrorMessage {
+    pub ws_server_events: WsServerEvents,
+    pub should_disconnect: bool,
+}
+
+impl ErrorMessage {
+    pub fn new(ws_error: WsError, trace_id: Option<String>, request_id: Option<String>) -> Self {
+        let should_client_retry = ws_error.should_client_retry();
+        let should_disconnect = ws_error.should_disconnect();
+        Self {
+            ws_server_events: ws_error.into_ws_server_events(
+                trace_id,
+                request_id,
+                should_client_retry,
+            ),
+            should_disconnect,
+        }
+    }
+
+    pub fn new_unauthorized(trace_id: Option<String>) -> Self {
+        let ws_server_events = WsServerEvents::Error {
+            code: StatusCode::UNAUTHORIZED.into(),
+            message: "Cookie expired. Please retry".to_string(),
+            error_detail: None,
+            trace_id,
+            request_id: None,
+            // client should not retry on the same ws session, only when cookie is expired
+            should_client_retry: false,
+        };
+        let should_disconnect = !config::get_config().websocket.is_session_drain_enabled;
+        Self {
+            ws_server_events,
+            should_disconnect,
+        }
+    }
+}
+
+impl WsError {
+    /// Disconnect the ws conn to client from router
+    pub fn should_disconnect(&self) -> bool {
+        matches!(
+            self,
+            WsError::ProtocolError(_) | WsError::QuerierNotAvailable(_)
+        )
+    }
+
+    pub fn should_client_retry(&self) -> bool {
+        matches!(
+            self,
+            WsError::ConnectionError(_) | WsError::ConnectionDisconnected
+        )
+    }
+
+    pub fn into_ws_server_events(
+        self,
+        trace_id: Option<String>,
+        request_id: Option<String>,
+        should_client_retry: bool,
+    ) -> WsServerEvents {
+        let code = self.status_code();
+        WsServerEvents::Error {
+            code: code.as_u16(),
+            message: code.to_string(),
+            error_detail: Some(self.to_string()),
+            trace_id,
+            request_id,
+            should_client_retry,
+        }
+    }
+}

src/router/http/ws_v2/handler.rs

@@ -0,0 +1,557 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::sync::{
+    Arc,
+    atomic::{AtomicBool, Ordering},
+};
+
+use actix_http::StatusCode;
+use actix_web::{Error, HttpRequest, HttpResponse, web};
+use actix_ws::{CloseCode, CloseReason};
+use config::{get_config, meta::cluster::RoleGroup, utils::json};
+use futures_util::StreamExt;
+
+use super::{
+    connection::{Connection, QuerierConnection},
+    error::*,
+    pool::QuerierConnectionPool,
+    session::SessionManager,
+};
+#[cfg(feature = "enterprise")]
+use crate::common::utils::auth::extract_auth_expiry_and_user_id;
+use crate::service::websocket_events::{WsClientEvents, WsServerEvents};
+
+pub type ClientId = String;
+pub type QuerierName = String;
+pub type TraceId = String;
+
+#[derive(Debug)]
+pub struct WsHandler {
+    pub session_manager: Arc<SessionManager>,
+    pub connection_pool: Arc<QuerierConnectionPool>,
+}
+
+impl WsHandler {
+    pub fn new(
+        session_manager: Arc<SessionManager>,
+        connection_pool: Arc<QuerierConnectionPool>,
+    ) -> Self {
+        Self {
+            session_manager,
+            connection_pool,
+        }
+    }
+
+    pub async fn handle_connection(
+        &self,
+        req: HttpRequest,
+        stream: web::Payload,
+        client_id: ClientId,
+    ) -> Result<HttpResponse, Error> {
+        let cfg = get_config();
+
+        // Client -> Router connection
+        let (response, mut ws_session, mut msg_stream) = actix_ws::handle(&req, stream)?;
+
+        // Create session by registering the client & extract the user_id from the auth
+        #[cfg(feature = "enterprise")]
+        let (mut cookie_expiry, user_id) = extract_auth_expiry_and_user_id(&req).await;
+
+        #[cfg(feature = "enterprise")]
+        {
+            if !cfg.websocket.check_cookie_expiry {
+                cookie_expiry = None;
+            }
+        }
+
+        #[cfg(not(feature = "enterprise"))]
+        let cookie_expiry = None;
+
+        self.session_manager
+            .register_client(&client_id, cookie_expiry)
+            .await;
+
+        // Setup message channels
+        // TODO: add env variable for this for channel size
+        let (response_tx, mut response_rx) = tokio::sync::mpsc::channel::<WsServerEvents>(100);
+        let (disconnect_tx, mut disconnect_rx) =
+            tokio::sync::mpsc::channel::<Option<DisconnectMessage>>(10);
+        let session_manager = self.session_manager.clone();
+        let connection_pool = self.connection_pool.clone();
+
+        // Spawn message handling tasks between client and router
+        actix_web::rt::spawn(async move {
+            let mut ping_interval = tokio::time::interval(tokio::time::Duration::from_secs(
+                cfg.websocket.ping_interval_secs as _,
+            ));
+            ping_interval.tick().await; // first tick is immediate
+
+            // Handle incoming messages from client
+            let handle_incoming = async {
+                loop {
+                    tokio::select! {
+                        _ = tokio::time::sleep(tokio::time::Duration::from_secs(cfg.websocket.session_idle_timeout_secs as _ )) => {
+                            // check if last_active time is updated by the outgoing thread
+                            if session_manager.reached_max_idle_time(&client_id).await {
+                                log::info!("[WS::Router::Handler]: MAX_IDLE_TIME reached. Normal shutdown");
+                                if let Err(e) = disconnect_tx.send(None).await {
+                                    log::error!(
+                                        "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                    );
+                                };
+                                break;
+                            }
+                        }
+                        _ = ping_interval.tick() => {
+                            if let Err(e) = response_tx.send(WsServerEvents::Ping(vec![])).await {
+                                log::error!("[WS::Router::Handler] error sending ping to outgoing thread via response channel: {}", e);
+                                _ = disconnect_tx.send(None).await;
+                                break;
+                            }
+                        }
+                        Some(msg) = msg_stream.next() => {
+                            match msg {
+                                Err(e) => {
+                                    log::error!(
+                                        "[WS::Router::Handler] error receiving websocket message from client {e}"
+                                    );
+                                    let err_msg = ErrorMessage::new(e.into(), None, None);
+                                    let should_disconnect = err_msg.should_disconnect;
+                                    if let Err(e) = disconnect_tx.send(Some(DisconnectMessage::Error(err_msg))).await {
+                                        log::error!(
+                                            "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                        );
+                                    };
+                                    if should_disconnect {
+                                        break;
+                                    }
+                                }
+                                Ok(actix_ws::Message::Text(text)) => {
+                                    match json::from_str::<WsClientEvents>(&text) {
+                                        Err(e) => {
+                                            log::error!(
+                                                "[WS::Router::Handler] received invalid request message: {:?}",
+                                                e
+                                            );
+                                            let trace_id = json::from_str::<json::Value>(&text)
+                                                .ok()
+                                                .and_then(|val| val.get("trace_id").cloned())
+                                                .map(|v| v.to_string());
+                                            let err_msg =
+                                                ErrorMessage::new(e.into(), trace_id, None);
+                                            let should_disconnect = err_msg.should_disconnect;
+                                            if let Err(e) = disconnect_tx.send(Some(DisconnectMessage::Error(err_msg))).await {
+                                                log::error!(
+                                                    "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                                );
+                                            }
+                                            if should_disconnect {
+                                                break;
+                                            }
+                                        }
+                                        #[allow(unused_mut)]
+                                        Ok(mut message) => {
+                                            // check if cookie is valid for each client event only
+                                            // for enterprise
+                                            if cfg.websocket.check_cookie_expiry
+                                                && !session_manager.is_client_cookie_valid(&client_id).await
+                                            {
+                                                log::info!(
+                                                    "[WS::Router::Handler] Client cookie expired. Disconnect..."
+                                                );
+                                                let err_msg = ErrorMessage::new_unauthorized(
+                                                    Some(message.get_trace_id()),
+                                                );
+                                                if let Err(e) = disconnect_tx.send(Some(DisconnectMessage::Error(err_msg))).await
+                                                {
+                                                    log::error!(
+                                                        "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                                    );
+                                                };
+                                                log::debug!(
+                                                    "[WS::Router::Handler] Stop handle_incoming"
+                                                );
+                                                break;
+                                            }
+                                            // end of cookie check
+
+                                            // Append `user_id` to the ws client events when run in
+                                            // cluster mode to handle stream permissions
+                                            #[cfg(feature = "enterprise")]
+                                            {
+                                                message.append_user_id(user_id.clone());
+                                            }
+
+                                            log::debug!(
+                                                "[WS::Router::Handler] received message: {:?}",
+                                                message
+                                            );
+                                            let trace_id = message.get_trace_id();
+                                            let role_group =
+                                                if let WsClientEvents::Search(req) = &message {
+                                                    Some(RoleGroup::from(req.search_type))
+                                                } else {
+                                                    None
+                                                };
+                                            let querier_conn = match get_querier_connection(
+                                                session_manager.clone(),
+                                                connection_pool.clone(),
+                                                &client_id,
+                                                &trace_id,
+                                                role_group,
+                                            )
+                                            .await
+                                            {
+                                                Err(e) => {
+                                                    log::error!(
+                                                        "[WS::Router::Handler] error getting querier_conn: {e}"
+                                                    );
+                                                    let err_msg = ErrorMessage::new(e, Some(trace_id), None);
+                                                    let should_disconnect =
+                                                        err_msg.should_disconnect;
+                                                    if let Err(e) = disconnect_tx.send(Some(DisconnectMessage::Error(err_msg))).await
+                                                    {
+                                                        log::error!(
+                                                            "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                                        );
+                                                    }
+                                                    if should_disconnect {
+                                                        break;
+                                                    } else {
+                                                        continue;
+                                                    }
+                                                }
+                                                Ok(querier_conn) => querier_conn,
+                                            };
+                                            querier_conn
+                                                .register_request(
+                                                    trace_id.clone(),
+                                                    response_tx.clone(),
+                                                )
+                                                .await;
+
+                                            if let Err(e) = querier_conn.send_message(message).await
+                                            {
+                                                log::error!(
+                                                    "[WS::Router::Handler] error forwarding client message via selected querier connection: {e}"
+                                                );
+                                                let err_msg = ErrorMessage::new(e, Some(trace_id), None);
+                                                let should_disconnect = err_msg.should_disconnect;
+                                                if let Err(e) = disconnect_tx.send(Some(DisconnectMessage::Error(err_msg))).await {
+                                                    log::error!(
+                                                        "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                                    );
+                                                }
+                                                if should_disconnect {
+                                                    break;
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
+                                Ok(actix_ws::Message::Close(close_reason)) => {
+                                    log::info!(
+                                        "[WS::Router::Handler] disconnect signal received from client."
+                                    );
+                                    if let Err(e) = disconnect_tx.send(Some(DisconnectMessage::Close(close_reason))).await {
+                                        log::error!(
+                                            "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                        );
+                                    };
+                                    log::debug!("[WS::Router::Handler] Stop handle_incoming");
+                                    break;
+                                }
+                                Ok(actix_ws::Message::Ping(ping)) => {
+                                    let _ = response_tx.send(WsServerEvents::Ping(ping.to_vec())).await;
+                                }
+                                Ok(actix_ws::Message::Pong(pong)) => {
+                                    log::debug!("[WS::Router::Handler]: Pong received from client : {:?}", pong);
+                                }
+                                Ok(_) => {}
+                            }
+                        }
+                    }
+                }
+                Ok::<_, Error>(())
+            };
+
+            // Handle outgoing messages from router to client
+            let handle_outgoing = async {
+                loop {
+                    tokio::select! {
+                        // response from querier
+                        Some(message) = response_rx.recv() => {
+                            match message {
+                                WsServerEvents::Ping(ping) => {
+                                    log::debug!("[WS::Router::Handler]: pinging client");
+                                    if let Err(e) = ws_session.pong(&ping).await {
+                                        log::error!("[WS::Router::Handler]: Error sending pong: {}", e);
+                                    }
+                                }
+                                WsServerEvents::Pong(pong) => {
+                                    log::debug!("[WS::Router::Handler]: Pong received from client : {:?}", pong);
+                                }
+                                _ => {
+                                    let Ok(message_str) = serde_json::to_string(&message) else {
+                                        log::error!(
+                                            "[WS::Router::Handler]: error convert WsServerEvents to string before sending back to client "
+                                        );
+                                        continue;
+                                    };
+                                    if let Err(e) = ws_session.text(message_str).await {
+                                        log::error!("Error sending message to client: {}", e);
+                                        break;
+                                    }
+                                    if let Some(trace_id) = message.should_clean_trace_id() {
+                                        log::info!("[WS::Router::Handler] Unregistering trace_id: {}, message: {:?}", trace_id, message);
+                                        session_manager.update_session_activity(&client_id).await;
+                                        session_manager.remove_trace_id(&client_id, &trace_id).await;
+                                    }
+                                }
+                            }
+                        }
+                        // interruption from handling_incoming thread
+                        Some(msg) = disconnect_rx.recv() => {
+                            match msg {
+                                None => {
+                                    // proper disconnecting
+                                    log::debug!(
+                                        "[WS::Handler]: disconnect signal received from client. handle_outgoing stopped"
+                                    );
+                                    break;
+                                }
+                                Some(DisconnectMessage::Error(err_msg)) => {
+                                    // send error message to client first
+                                    _ = ws_session.text(err_msg.ws_server_events.to_json()).await;
+                                    // Experimental feature: `is_session_drain_enabled`
+                                    if !cfg.websocket.is_session_drain_enabled && err_msg.should_disconnect {
+                                        break;
+                                    }
+
+                                    // then drain the session
+                                    if cfg.websocket.is_session_drain_enabled {
+                                        match err_msg.ws_server_events {
+                                            WsServerEvents::Error { code, .. } if code == <actix_http::StatusCode as Into<u16>>::into(StatusCode::UNAUTHORIZED) => {
+                                                let is_session_drain_complete = Arc::new(AtomicBool::new(false));
+                                                let is_session_drain_complete_clone = is_session_drain_complete.clone();
+
+                                                // Clone values before moving into spawn
+                                                let session_manager = session_manager.clone();
+                                                let session_manager_clone = session_manager.clone();
+                                                let connection_pool = connection_pool.clone();
+                                                let client_id = client_id.clone();
+                                                let client_id_clone = client_id.clone();
+                                                let cfg_clone = cfg.clone();
+
+                                                // Set a flag to indicate we're in drain mode
+                                                let is_draining = Arc::new(AtomicBool::new(true));
+                                                let is_draining_clone = is_draining.clone();
+
+                                                let _session_drain_task = tokio::spawn(async move {
+                                                    let querier_connections = session_manager_clone.get_querier_connections(&client_id_clone).await;
+                                                    let max_wait_time = cfg_clone.websocket.session_idle_timeout_secs as u64;
+                                                    let start_time = std::time::Instant::now();
+
+                                                    'outer: loop {
+                                                        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
+
+                                                        if start_time.elapsed().as_secs() > max_wait_time {
+                                                            log::warn!("[WS::Router::Handler]: Session drain timeout reached");
+                                                            break;
+                                                        }
+
+                                                        let current_trace_ids = session_manager_clone.get_trace_ids(&client_id_clone).await;
+
+                                                        // If all trace_ids are processed, we're done
+                                                        if current_trace_ids.is_empty() {
+                                                            is_session_drain_complete_clone.store(true, Ordering::Relaxed);
+                                                            log::info!("[WS::Router::Handler]: All trace IDs drained successfully");
+                                                            break;
+                                                        }
+
+                                                        // Log remaining trace IDs for debugging
+                                                        log::debug!(
+                                                            "[WS::Router::Handler]: Waiting for trace IDs to complete: {:?}",
+                                                            current_trace_ids
+                                                        );
+
+                                                        // Check each querier connection for active trace IDs
+                                                        for querier_name in &querier_connections {
+                                                            let Some(active_connection) = connection_pool.get_active_connection(querier_name).await else {
+                                                                continue;
+                                                            };
+
+                                                            for trace_id in &current_trace_ids {
+                                                                if active_connection.is_active_trace_id(trace_id).await {
+                                                                    continue 'outer;
+                                                                }
+                                                            }
+                                                        }
+                                                    }
+
+                                                    // Clear drain mode flag when complete
+                                                    is_draining_clone.store(false, Ordering::Relaxed);
+                                                });
+
+                                                // Continue processing messages while draining
+                                                while !is_session_drain_complete.load(Ordering::Relaxed) {
+                                                    if let Some(message) = response_rx.recv().await {
+                                                        // Forward all messages to client during drain
+                                                        if let Ok(message_str) = serde_json::to_string(&message) {
+                                                            if let Err(e) = ws_session.text(message_str).await {
+                                                                log::error!("[WS::Router::Handler]: Error sending message during drain: {}", e);
+                                                                break;
+                                                            }
+
+                                                            // Important: Process End messages and clean up trace_ids
+                                                            if let Some(trace_id) = message.should_clean_trace_id() {
+                                                                log::info!(
+                                                                    "[WS::Router::Handler] Cleaning up trace_id during drain: {}, message: {:?}",
+                                                                    trace_id,
+                                                                    message
+                                                                );
+                                                                session_manager.remove_trace_id(&client_id, &trace_id).await;
+
+                                                                // Check if this was the last trace_id
+                                                                let remaining_trace_ids = session_manager.get_trace_ids(&client_id).await;
+                                                                if remaining_trace_ids.is_empty() {
+                                                                    is_session_drain_complete.store(true, Ordering::Relaxed);
+                                                                    // close the session
+                                                                    if let Err(e) = disconnect_tx.send(Some(DisconnectMessage::Close(Some(CloseReason::from(CloseCode::Normal))))).await {
+                                                                        log::error!(
+                                                                            "[WS::Router::Handler] Error informing handle_outgoing to stop: {e}"
+                                                                        );
+                                                                    }
+                                                                    log::info!("[WS::Router::Handler]: All trace IDs drained successfully");
+                                                                    break;
+                                                                }
+                                                            }
+                                                        }
+                                                    } else {
+                                                        // Channel closed, no more messages coming
+                                                        log::warn!("[WS::Router::Handler]: Response channel closed during drain");
+                                                        break;
+                                                    }
+                                                }
+                                            }
+                                            _ => {}
+                                        }
+                                    }
+                                }
+                                Some(DisconnectMessage::Close(close_reason)) => {
+                                    if let Err(e) = ws_session.close(close_reason).await {
+                                        log::error!("[WS::Handler]: Error closing websocket session: {}", e);
+                                    };
+                                    return Ok(());
+                                }
+                            }
+                        }
+                    }
+                }
+
+                _ = ws_session
+                    .close(Some(CloseReason::from(CloseCode::Normal)))
+                    .await;
+                log::info!("[WS::Router::Handler]: client ws closed: {}", client_id);
+                Ok::<_, Error>(())
+            };
+
+            // Run both tasks concurrently
+            let _ = tokio::join!(handle_incoming, handle_outgoing);
+
+            // clean up resource related to this client
+            session_manager.unregister_client(&client_id).await;
+        });
+
+        Ok(response)
+    }
+
+    pub async fn remove_querier_connection(&self, querier_name: &str) {
+        self.connection_pool
+            .remove_querier_connection(querier_name)
+            .await;
+        self.session_manager
+            .remove_querier_connection(querier_name)
+            .await;
+    }
+
+    pub async fn _shutdown(&self) {
+        // just need to disconnect all the ws_conns in connection_pool
+        self.connection_pool._shutdown().await;
+    }
+}
+
+pub async fn get_querier_connection(
+    session_manager: Arc<SessionManager>,
+    connection_pool: Arc<QuerierConnectionPool>,
+    client_id: &ClientId,
+    trace_id: &TraceId,
+    role_group: Option<RoleGroup>,
+) -> WsResult<Arc<QuerierConnection>> {
+    let interval = tokio::time::Duration::from_secs(1);
+    // Retry max of 3 times to get a valid querier connection
+    for try_num in 0..3 {
+        // Get or assign querier for this trace_id included in message
+        let querier_name = match session_manager
+            .get_querier_for_trace(client_id, trace_id)
+            .await?
+        {
+            Some(querier_name) => querier_name,
+            None => {
+                let querier_name = select_querier(trace_id, role_group).await?;
+                session_manager
+                    .set_querier_for_trace(client_id, trace_id, &querier_name)
+                    .await?;
+                querier_name
+            }
+        };
+
+        // Get connection for selected querier
+        match connection_pool
+            .get_or_create_connection(&querier_name)
+            .await
+        {
+            Ok(conn) => return Ok(conn),
+            Err(e) => {
+                log::error!(
+                    "[WS::Router::Handler] error getting or creating querier connection: {e} try number: {}",
+                    try_num
+                );
+                session_manager
+                    .remove_querier_connection(&querier_name)
+                    .await;
+            }
+        }
+        tokio::time::sleep(interval).await;
+    }
+
+    Err(WsError::QuerierNotAvailable(
+        "Reached maximum retries of 3".to_string(),
+    ))
+}
+
+async fn select_querier(trace_id: &str, role_group: Option<RoleGroup>) -> WsResult<QuerierName> {
+    use crate::common::infra::cluster;
+
+    let node = cluster::get_node_from_consistent_hash(
+        trace_id,
+        &config::meta::cluster::Role::Querier,
+        role_group,
+    )
+    .await
+    .ok_or_else(|| WsError::QuerierNotAvailable(format!("for trace_id {}", trace_id)))?;
+    Ok(node)
+}

src/router/http/ws_v2/mod.rs

@@ -0,0 +1,111 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+//! WebSocket v2 Implementation for Router-Querier Communication
+//!
+//! This module implements a WebSocket-based communication system between clients, routers, and
+//! queriers. It provides connection pooling, session management, and reliable message routing.
+//!
+//! # Architecture
+//!
+//! The system consists of several key components:
+//!
+//! * `WsHandler`: Main handler for client WebSocket connections
+//! * `QuerierConnectionPool`: Manages persistent WebSocket connections to querier nodes
+//! * `SessionManager`: Tracks client sessions and their querier mappings
+//! * `QuerierConnection`: Handles individual WebSocket connections to querier nodes
+//!
+//! # Flow
+//!
+//! 1. Client establishes WebSocket connection with Router
+//! 2. Router creates session and validates authentication
+//! 3. Client sends requests with trace_ids
+//! 4. Router maps trace_ids to querier nodes using consistent hashing
+//! 5. Router maintains persistent connections to queriers
+//! 6. Responses are routed back to clients using trace_ids
+//!
+//! # Features
+//!
+//! * Connection pooling for querier connections
+//! * Session management with idle timeout
+//! * Health checking of querier connections
+//! * Graceful session draining on disconnects
+//! * Request/response routing using trace_ids
+//! * Support for enterprise authentication
+//!
+//! # Components
+//!
+//! ## Handler
+//! Handles the main WebSocket lifecycle and message routing
+//!
+//! ## Connection Pool
+//! Maintains and reuses WebSocket connections to querier nodes
+//!
+//! ## Session Manager
+//! Tracks client sessions, querier mappings, and handles cleanup
+//!
+//! ## Connection
+//! Implements the low-level WebSocket connection handling
+
+use std::sync::Arc;
+
+use handler::WsHandler;
+use once_cell::sync::OnceCell;
+use pool::QuerierConnectionPool;
+use session::SessionManager;
+
+mod connection;
+mod error;
+mod handler;
+mod pool;
+mod session;
+
+// Initialize WsHandler global instance
+static WS_HANDLER: OnceCell<Arc<WsHandler>> = OnceCell::new();
+
+// Helper function to get or initialize the WsHandler
+pub async fn get_ws_handler() -> Arc<WsHandler> {
+    if let Some(handler) = WS_HANDLER.get() {
+        return handler.clone();
+    }
+
+    // Initialize if not already done
+    let handler = init().await;
+
+    // This may fail if another thread initialized it first, which is fine
+    let _ = WS_HANDLER.set(handler.clone());
+
+    handler
+}
+
+pub async fn remove_querier_from_handler(querier_name: &str) {
+    get_ws_handler()
+        .await
+        .remove_querier_connection(querier_name)
+        .await;
+}
+
+async fn init() -> Arc<WsHandler> {
+    let session_manager = Arc::new(SessionManager::default());
+    let connection_pool = Arc::new(QuerierConnectionPool::new());
+    let handler = Arc::new(WsHandler::new(session_manager, connection_pool.clone()));
+
+    // Start connection maintenance task
+    tokio::spawn(async move {
+        connection_pool.maintain_connections().await;
+    });
+
+    handler
+}

src/router/http/ws_v2/pool.rs

@@ -0,0 +1,114 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::sync::Arc;
+
+use config::{RwAHashMap, get_config};
+
+use super::{
+    connection::{Connection, QuerierConnection},
+    error::*,
+    handler::QuerierName,
+};
+
+#[derive(Debug)]
+pub struct QuerierConnectionPool {
+    connections: RwAHashMap<QuerierName, Arc<QuerierConnection>>,
+}
+
+impl QuerierConnectionPool {
+    pub fn new() -> Self {
+        Self {
+            connections: RwAHashMap::default(),
+        }
+    }
+
+    pub async fn get_or_create_connection(
+        &self,
+        querier_name: &QuerierName,
+    ) -> WsResult<Arc<QuerierConnection>> {
+        if let Some(conn) = self.connections.read().await.get(querier_name) {
+            // double check if the connection is still connected
+            return if conn.is_connected().await {
+                Ok(conn.clone())
+            } else {
+                Err(WsError::ConnectionDisconnected)
+            };
+        }
+
+        // Create new connection
+        let conn = super::connection::create_connection(querier_name).await?;
+        self.connections
+            .write()
+            .await
+            .insert(querier_name.to_string(), conn.clone());
+        Ok(conn)
+    }
+
+    pub async fn remove_querier_connection(&self, querier_name: &str) {
+        if let Some(conn) = self.connections.write().await.remove(querier_name) {
+            log::warn!("[WS::ConnectionPool] removing connection to querier {querier_name}");
+            conn.disconnect().await;
+        }
+    }
+
+    pub async fn maintain_connections(&self) {
+        let cfg = get_config();
+        loop {
+            let mut to_remove = Vec::new();
+
+            let read_guard = self.connections.read().await;
+            for (querier_name, conn) in read_guard.iter() {
+                if !conn.is_connected().await {
+                    // Just drop it. A new connection will be made to the querier when chosen again
+                    to_remove.push(querier_name.clone());
+                }
+            }
+            drop(read_guard);
+
+            // Remove connections that failed to reconnect
+            for querier in to_remove {
+                log::warn!(
+                    "[WS::ConnectionPool] Removing disconnected connection to querier: {}",
+                    querier
+                );
+                self.remove_querier_connection(&querier).await;
+            }
+
+            tokio::time::sleep(tokio::time::Duration::from_secs(
+                cfg.websocket.health_check_interval as _,
+            ))
+            .await;
+        }
+    }
+
+    pub async fn _shutdown(&self) {
+        let mut writer_guard = self.connections.write().await;
+        for (querier_name, conn) in writer_guard.drain() {
+            log::warn!(
+                "[WS::ConnectionPool] error disconnect connection to querier {querier_name} as it's still in use. Force remove."
+            );
+            conn.disconnect().await;
+        }
+    }
+
+    pub async fn get_active_connection(
+        &self,
+        querier_name: &QuerierName,
+    ) -> Option<Arc<QuerierConnection>> {
+        let read_guard = self.connections.read().await;
+        read_guard.get(querier_name).cloned()
+    }
+}

src/router/http/ws_v2/session.rs

@@ -0,0 +1,222 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+use std::collections::{HashMap, HashSet};
+
+use chrono::{DateTime, Utc};
+use config::RwAHashMap;
+
+use super::{
+    error::*,
+    handler::{ClientId, QuerierName, TraceId},
+};
+
+#[derive(Debug, Default)]
+pub struct SessionManager {
+    sessions: RwAHashMap<ClientId, SessionInfo>,
+    mapped_queriers: RwAHashMap<QuerierName, HashSet<TraceId>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct SessionInfo {
+    pub querier_mappings: HashMap<TraceId, QuerierName>,
+    pub cookie_expiry: Option<DateTime<Utc>>,
+    pub last_active: DateTime<Utc>,
+}
+
+impl SessionManager {
+    pub async fn register_client(
+        &self,
+        client_id: &ClientId,
+        cookie_expiry: Option<DateTime<Utc>>,
+    ) {
+        if self.sessions.read().await.get(client_id).is_some() {
+            return;
+        }
+
+        let session_info = SessionInfo {
+            querier_mappings: HashMap::default(),
+            cookie_expiry,
+            last_active: Utc::now(),
+        };
+
+        let mut write_guard = self.sessions.write().await;
+        if !write_guard.contains_key(client_id) {
+            write_guard.insert(client_id.clone(), session_info.clone());
+        }
+    }
+
+    pub async fn update_session_activity(&self, client_id: &ClientId) {
+        let mut write_guard = self.sessions.write().await;
+        if let Some(session_info) = write_guard.get_mut(client_id) {
+            session_info.last_active = chrono::Utc::now();
+        }
+    }
+
+    pub async fn remove_trace_id(&self, client_id: &str, trace_id: &str) {
+        let querier_name = {
+            let mut session_write = self.sessions.write().await;
+            session_write
+                .get_mut(client_id)
+                .and_then(|session_info| session_info.querier_mappings.remove(trace_id))
+        };
+
+        if let Some(querier_name) = querier_name {
+            log::debug!(
+                "[WS::Session] removed querier {querier_name} from sessions-session_info-querier_mappings"
+            );
+            let mut mapping_write = self.mapped_queriers.write().await;
+            if let Some(trace_ids) = mapping_write.get_mut(&querier_name) {
+                trace_ids.remove(trace_id);
+                log::debug!(
+                    "[WS::Session] removed trace_id {trace_id} from mapped_queriers-trace_ids"
+                );
+                if trace_ids.is_empty() {
+                    log::debug!(
+                        "[WS::Session] no more trace_id's mapped to querier {querier_name}. removing querier from mapped_queriers"
+                    );
+                    mapping_write.remove(&querier_name);
+                }
+            }
+        }
+    }
+
+    pub async fn reached_max_idle_time(&self, client_id: &ClientId) -> bool {
+        self.sessions
+            .read()
+            .await
+            .get(client_id)
+            .is_none_or(|session_info| {
+                Utc::now()
+                    .signed_duration_since(session_info.last_active)
+                    .num_seconds()
+                    > config::get_config().websocket.session_idle_timeout_secs
+            })
+    }
+
+    pub async fn unregister_client(&self, client_id: &ClientId) {
+        if let Some(session_info) = self.sessions.write().await.remove(client_id) {
+            let mut mapped_querier_write = self.mapped_queriers.write().await;
+
+            for (trace_id, querier_name) in session_info.querier_mappings {
+                if let Some(mut trace_ids) = mapped_querier_write.remove(&querier_name) {
+                    trace_ids.retain(|tid| tid != &trace_id);
+                    if !trace_ids.is_empty() {
+                        trace_ids.shrink_to_fit();
+                        mapped_querier_write.insert(querier_name.to_string(), trace_ids);
+                    }
+                }
+            }
+        }
+    }
+
+    pub async fn is_client_cookie_valid(&self, client_id: &ClientId) -> bool {
+        match self.sessions.read().await.get(client_id) {
+            Some(session_info) => session_info
+                .cookie_expiry
+                .is_none_or(|expiry| expiry > Utc::now()),
+            None => false, // not set is treated as unauthenticated
+        }
+    }
+
+    pub async fn remove_querier_connection(&self, querier_name: &str) {
+        let client_ids = {
+            let (mapped_read, sessions_read) =
+                tokio::join!(self.mapped_queriers.read(), self.sessions.read());
+
+            match mapped_read.get(querier_name) {
+                Some(_) => sessions_read.keys().cloned().collect::<Vec<_>>(),
+                None => return,
+            }
+        };
+
+        // Remove from mapped_querier
+        let trace_ids = self
+            .mapped_queriers
+            .write()
+            .await
+            .remove(querier_name)
+            .unwrap(); // existence validated
+
+        // Batch update sessions
+        let mut session_write = self.sessions.write().await;
+        for client_id in client_ids {
+            if let Some(session_info) = session_write.get_mut(&client_id) {
+                session_info
+                    .querier_mappings
+                    .retain(|tid, _| !trace_ids.contains(tid));
+                session_info.querier_mappings.shrink_to_fit();
+            }
+        }
+    }
+
+    pub async fn set_querier_for_trace(
+        &self,
+        client_id: &ClientId,
+        trace_id: &TraceId,
+        querier_name: &QuerierName,
+    ) -> WsResult<()> {
+        self.sessions
+            .write()
+            .await
+            .get_mut(client_id)
+            .ok_or(WsError::SessionNotFound(format!("client_id {}", client_id)))?
+            .querier_mappings
+            .insert(trace_id.clone(), querier_name.clone());
+
+        // mapped_queriers
+        self.mapped_queriers
+            .write()
+            .await
+            .entry(querier_name.clone())
+            .or_insert_with(HashSet::new)
+            .insert(trace_id.clone());
+        Ok(())
+    }
+
+    pub async fn get_querier_for_trace(
+        &self,
+        client_id: &ClientId,
+        trace_id: &TraceId,
+    ) -> WsResult<Option<QuerierName>> {
+        Ok(self
+            .sessions
+            .read()
+            .await
+            .get(client_id)
+            .ok_or(WsError::SessionNotFound(format!("client_id {}", client_id)))?
+            .querier_mappings
+            .get(trace_id)
+            .cloned())
+    }
+
+    pub async fn get_querier_connections(&self, client_id: &ClientId) -> Vec<QuerierName> {
+        self.sessions
+            .read()
+            .await
+            .get(client_id)
+            .map(|session_info| session_info.querier_mappings.values().cloned().collect())
+            .unwrap_or_default()
+    }
+
+    pub async fn get_trace_ids(&self, client_id: &ClientId) -> Vec<TraceId> {
+        self.sessions
+            .read()
+            .await
+            .get(client_id)
+            .map(|session_info| session_info.querier_mappings.keys().cloned().collect())
+            .unwrap_or_default()
+    }
+}

src/service/alerts/alert.rs

@@ -1612,7 +1612,9 @@ pub fn get_alert_start_end_time(
 }
 
 fn format_variable_value(val: String) -> String {
-    val.chars().flat_map(|c| c.escape_default()).collect::<String>()
+    val.chars()
+        .flat_map(|c| c.escape_default())
+        .collect::<String>()
 }
 
 pub(super) fn to_float(val: &Value) -> f64 {
@@ -1779,7 +1781,10 @@ mod tests {
         assert_eq!(format_variable_value("".to_string()), "");
 
         // Test string with no special characters
-        assert_eq!(format_variable_value("Hello World".to_string()), "Hello World");
+        assert_eq!(
+            format_variable_value("Hello World".to_string()),
+            "Hello World"
+        );
     }
 
     #[tokio::test]

src/service/db/session.rs

@@ -30,8 +30,12 @@ pub async fn get(session_id: &str) -> Result<String, anyhow::Error> {
         Some(val) => Ok(val.to_string()),
         None => {
             let val = db::get(&format!("{USER_SESSION_KEY}{session_id}")).await?;
-
-            Ok(json::to_string(&val).unwrap())
+            // base64 format: convert bytes to string and trim any quotes
+            let val = String::from_utf8(val.to_vec())
+                .unwrap()
+                .trim_matches('"')
+                .to_string();
+            Ok(val)
         }
     }
 }

src/service/mod.rs

@@ -38,6 +38,8 @@ pub mod pipeline;
 pub mod promql;
 pub mod schema;
 pub mod search;
+pub mod websocket_events;
+
 #[cfg(feature = "enterprise")]
 pub mod search_jobs;
 pub mod self_reporting;

src/service/websocket_events/mod.rs

@@ -0,0 +1,21 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+pub mod search;
+pub mod sort;
+pub mod utils;
+
+pub use search::*;
+pub use utils::*;

src/service/websocket_events/search.rs

@@ -29,7 +29,7 @@ use tracing::Instrument;
 
 use super::sort::order_search_results;
 #[allow(unused_imports)]
-use crate::handler::http::request::websocket::utils::enterprise_utils;
+use crate::service::websocket_events::enterprise_utils;
 use crate::{
     common::{
         meta::search::{CachedQueryResponse, MultiCachedQueryResponse, QueryDelta},
@@ -40,12 +40,13 @@ use crate::{
             },
         },
     },
-    handler::http::request::websocket::{
-        session::send_message,
-        utils::{TimeOffset, WsServerEvents, search_registry_utils},
-    },
-    service::search::{
-        self as SearchService, cache, datafusion::distributed_plan::streaming_aggs_exec, sql::Sql,
+    handler::http::request::websocket::session::send_message,
+    service::{
+        search::{
+            self as SearchService, cache, datafusion::distributed_plan::streaming_aggs_exec,
+            sql::Sql,
+        },
+        websocket_events::{TimeOffset, WsServerEvents, search_registry_utils},
     },
 };
 
@@ -110,8 +111,9 @@ pub async fn handle_search_request(
                 Error::Message(e.to_string()),
                 Some(req_id.to_string()),
                 Some(trace_id),
+                Default::default(),
             );
-            send_message(req_id, err_res.to_json().to_string()).await?;
+            send_message(req_id, err_res.to_json()).await?;
             return Ok(());
         }
     };
@@ -126,8 +128,9 @@ pub async fn handle_search_request(
                 Error::Message(e),
                 Some(req_id.to_string()),
                 Some(trace_id),
+                Default::default(),
             );
-            send_message(req_id, err_res.to_json().to_string()).await?;
+            send_message(req_id, err_res.to_json()).await?;
             return Ok(());
         }
     }
@@ -231,7 +234,6 @@ pub async fn handle_search_request(
                 cached_resp,
                 deltas,
                 accumulated_results,
-                org_id,
                 user_id,
                 max_query_range,
                 remaining_query_range,
@@ -253,7 +255,6 @@ pub async fn handle_search_request(
                 &mut req,
                 &trace_id,
                 req_size,
-                org_id,
                 user_id,
                 accumulated_results,
                 max_query_range,
@@ -284,7 +285,6 @@ pub async fn handle_search_request(
             &mut req,
             &trace_id,
             req_size,
-            org_id,
             user_id,
             accumulated_results,
             max_query_range,
@@ -297,28 +297,27 @@ pub async fn handle_search_request(
     let end_res = WsServerEvents::End {
         trace_id: Some(trace_id.clone()),
     };
-    send_message(req_id, end_res.to_json().to_string()).await?;
+    send_message(req_id, end_res.to_json()).await?;
 
     Ok(())
 }
 
 async fn do_search(
     req: &SearchEventReq,
-    org_id: &str,
     user_id: &str,
     use_cache: bool,
 ) -> Result<Response, Error> {
     let span = tracing::info_span!(
         "src::handler::http::request::websocket::search::do_search",
         trace_id = %req.trace_id,
-        org_id = %org_id,
+        org_id = %req.org_id,
     );
 
     let mut req = req.clone();
     req.payload.use_cache = Some(use_cache);
     let res = SearchService::cache::search(
         &req.trace_id,
-        org_id,
+        &req.org_id,
         req.stream_type,
         Some(user_id.to_string()),
         &req.payload,
@@ -352,7 +351,6 @@ async fn handle_cache_responses_and_deltas(
     mut cached_resp: Vec<CachedQueryResponse>,
     mut deltas: Vec<QueryDelta>,
     accumulated_results: &mut Vec<SearchResultType>,
-    org_id: &str,
     user_id: &str,
     max_query_range: i64,
     remaining_query_range: i64,
@@ -435,7 +433,6 @@ async fn handle_cache_responses_and_deltas(
                     req_size,
                     accumulated_results,
                     &mut curr_res_size,
-                    org_id,
                     user_id,
                     &mut remaining_query_range,
                     cached_search_duration,
@@ -470,7 +467,6 @@ async fn handle_cache_responses_and_deltas(
                 req_size,
                 accumulated_results,
                 &mut curr_res_size,
-                org_id,
                 user_id,
                 &mut remaining_query_range,
                 cached_search_duration,
@@ -515,7 +511,6 @@ async fn process_delta(
     req_size: i64,
     accumulated_results: &mut Vec<SearchResultType>,
     curr_res_size: &mut i64,
-    org_id: &str,
     user_id: &str,
     remaining_query_range: &mut f64,
     cache_req_duration: i64,
@@ -531,7 +526,7 @@ async fn process_delta(
     req.payload.query.start_time = delta.delta_start_time;
     req.payload.query.end_time = delta.delta_end_time;
 
-    let partition_resp = get_partitions(&req, org_id, user_id).await?;
+    let partition_resp = get_partitions(&req, user_id).await?;
     let mut partitions = partition_resp.partitions;
 
     if partitions.is_empty() {
@@ -578,7 +573,7 @@ async fn process_delta(
         }
 
         // use cache for delta search
-        let mut search_res = do_search(&req, org_id, user_id, true).await?;
+        let mut search_res = do_search(&req, user_id, true).await?;
         *curr_res_size += search_res.hits.len() as i64;
 
         log::info!(
@@ -643,7 +638,7 @@ async fn process_delta(
                 result_cache_ratio,
                 accumulated_results.len()
             );
-            send_message(req_id, ws_search_res.to_json().to_string()).await?;
+            send_message(req_id, ws_search_res.to_json()).await?;
         }
 
         // Stop if `remaining_query_range` is less than 0
@@ -690,7 +685,6 @@ async fn process_delta(
 
 async fn get_partitions(
     req: &SearchEventReq,
-    org_id: &str,
     user_id: &str,
 ) -> Result<SearchPartitionResponse, Error> {
     let search_payload = req.payload.clone();
@@ -708,7 +702,7 @@ async fn get_partitions(
 
     let res = SearchService::search_partition(
         &req.trace_id,
-        org_id,
+        &req.org_id,
         Some(user_id),
         req.stream_type,
         &search_partition_req,
@@ -798,7 +792,7 @@ async fn send_cached_responses(
         cached.cached_response.result_cache_ratio,
         accumulated_results.len()
     );
-    send_message(req_id, ws_search_res.to_json().to_string()).await?;
+    send_message(req_id, ws_search_res.to_json()).await?;
 
     Ok(())
 }
@@ -810,7 +804,6 @@ async fn do_partitioned_search(
     req: &mut SearchEventReq,
     trace_id: &str,
     req_size: i64,
-    org_id: &str,
     user_id: &str,
     accumulated_results: &mut Vec<SearchResultType>,
     max_query_range: i64, // hours
@@ -837,7 +830,7 @@ async fn do_partitioned_search(
     let modified_start_time = req.payload.query.start_time;
     let modified_end_time = req.payload.query.end_time;
 
-    let partition_resp = get_partitions(req, org_id, user_id).await?;
+    let partition_resp = get_partitions(req, user_id).await?;
     let mut partitions = partition_resp.partitions;
 
     if partitions.is_empty() {
@@ -889,7 +882,7 @@ async fn do_partitioned_search(
         }
 
         // do not use cache for partitioned search without cache
-        let mut search_res = do_search(&req, org_id, user_id, false).await?;
+        let mut search_res = do_search(&req, user_id, false).await?;
         curr_res_size += search_res.hits.len() as i64;
 
         if !search_res.hits.is_empty() {
@@ -928,7 +921,7 @@ async fn do_partitioned_search(
                 },
                 streaming_aggs: is_streaming_aggs,
             };
-            send_message(req_id, ws_search_res.to_json().to_string()).await?;
+            send_message(req_id, ws_search_res.to_json()).await?;
         }
 
         // Stop if reached the requested result size
@@ -957,7 +950,7 @@ async fn do_partitioned_search(
             },
             streaming_aggs: is_streaming_aggs,
         };
-        send_message(req_id, ws_search_res.to_json().to_string()).await?;
+        send_message(req_id, ws_search_res.to_json()).await?;
     }
 
     // Remove the streaming_aggs cache
@@ -1005,7 +998,7 @@ async fn send_partial_search_resp(
         trace_id
     );
 
-    send_message(req_id, ws_search_res.to_json().to_string()).await?;
+    send_message(req_id, ws_search_res.to_json()).await?;
 
     Ok(())
 }

src/service/websocket_events/sort.rs

@@ -1,3 +1,18 @@
+// Copyright 2025 OpenObserve Inc.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
 use config::meta::{search::Response, sql::OrderBy};
 use serde_json::Value;
 

src/service/websocket_events/utils.rs

@@ -15,8 +15,9 @@
 
 use actix_web::http::StatusCode;
 use config::meta::websocket::SearchEventReq;
-use infra::{errors, errors::Error};
+use infra::errors;
 use serde::{Deserialize, Serialize};
+use tokio_tungstenite::tungstenite;
 
 pub mod enterprise_utils {
     #[allow(unused_imports)]
@@ -92,8 +93,8 @@ pub mod sessions_cache_utils {
 
     use super::search_registry_utils::SearchState;
     use crate::{
-        common::infra::config::WS_SESSIONS,
-        handler::http::request::websocket::session::{SEARCH_REGISTRY, WsSession},
+        common::infra::config::{WS_SEARCH_REGISTRY, WS_SESSIONS},
+        handler::http::request::ws_v2::session::WsSession,
     };
 
     pub async fn run_gc_ws_sessions() {
@@ -162,7 +163,7 @@ pub mod sessions_cache_utils {
     }
 
     fn cleanup_searches_for_session(session_id: &str) {
-        let searches_to_remove: Vec<String> = SEARCH_REGISTRY
+        let searches_to_remove: Vec<String> = WS_SEARCH_REGISTRY
             .iter()
             .filter_map(|entry| {
                 if entry.value().get_req_id() == session_id {
@@ -174,7 +175,7 @@ pub mod sessions_cache_utils {
             .collect();
 
         for trace_id in searches_to_remove {
-            if let Some((_, state)) = SEARCH_REGISTRY.remove(&trace_id) {
+            if let Some((_, state)) = WS_SEARCH_REGISTRY.remove(&trace_id) {
                 match state {
                     SearchState::Running { cancel_tx, req_id } => {
                         let _ = cancel_tx.try_send(());
@@ -227,7 +228,7 @@ pub mod sessions_cache_utils {
 pub mod search_registry_utils {
     use tokio::sync::mpsc;
 
-    use crate::handler::http::request::websocket::session::SEARCH_REGISTRY;
+    use crate::common::infra::config::WS_SEARCH_REGISTRY;
 
     #[derive(Debug)]
     pub enum SearchState {
@@ -255,7 +256,7 @@ pub mod search_registry_utils {
 
     // Add this function to check if a search is cancelled
     pub fn is_cancelled(trace_id: &str) -> Option<bool> {
-        SEARCH_REGISTRY
+        WS_SEARCH_REGISTRY
             .get(trace_id)
             .map(|state| matches!(*state, SearchState::Cancelled { .. }))
     }
@@ -286,10 +287,17 @@ pub enum WsClientEvents {
     #[cfg(feature = "enterprise")]
     Cancel {
         trace_id: String,
+        // TODO: remove this once v1 is deprecated
+        org_id: String,
+        // TODO: is it PII safe?
+        user_id: Option<String>,
     },
     Benchmark {
         id: String,
     },
+    TestAbnormalClose {
+        req_id: String,
+    },
 }
 
 impl WsClientEvents {
@@ -299,6 +307,7 @@ impl WsClientEvents {
             #[cfg(feature = "enterprise")]
             WsClientEvents::Cancel { .. } => "cancel",
             WsClientEvents::Benchmark { .. } => "benchmark",
+            WsClientEvents::TestAbnormalClose { .. } => "test_abnormal_close",
         }
         .to_string()
     }
@@ -306,6 +315,39 @@ impl WsClientEvents {
     pub fn to_json(&self) -> String {
         serde_json::to_string(self).expect("Failed to serialize WsClientEvents")
     }
+
+    pub fn get_trace_id(&self) -> String {
+        match &self {
+            Self::Search(req) => req.trace_id.clone(),
+            #[cfg(feature = "enterprise")]
+            Self::Cancel { trace_id, .. } => trace_id.clone(),
+            Self::Benchmark { id } => id.clone(),
+            Self::TestAbnormalClose { req_id } => req_id.clone(),
+        }
+    }
+
+    pub fn is_valid(&self) -> bool {
+        match self {
+            Self::Search(req) => req.is_valid(),
+            #[cfg(feature = "enterprise")]
+            Self::Cancel {
+                trace_id, org_id, ..
+            } => !trace_id.is_empty() && !org_id.is_empty(),
+            Self::Benchmark { id } => !id.is_empty(),
+            Self::TestAbnormalClose { req_id } => !req_id.is_empty(),
+        }
+    }
+
+    // Append `user_id` to the ws client events when run in cluster mode to handle stream
+    // permissions
+    #[cfg(feature = "enterprise")]
+    pub fn append_user_id(&mut self, user_id: Option<String>) {
+        match self {
+            Self::Search(req) => req.user_id = user_id,
+            Self::Cancel { user_id: uid, .. } => *uid = user_id,
+            _ => {}
+        }
+    }
 }
 
 /// To represent the query start and end time based of partition or cache
@@ -339,10 +381,13 @@ pub enum WsServerEvents {
         error_detail: Option<String>,
         trace_id: Option<String>,
         request_id: Option<String>,
+        should_client_retry: bool,
     },
     End {
         trace_id: Option<String>,
     },
+    Ping(Vec<u8>),
+    Pong(Vec<u8>),
 }
 
 impl WsServerEvents {
@@ -351,25 +396,83 @@ impl WsServerEvents {
     }
 
     pub fn error_response(
-        err: Error,
+        err: errors::Error,
         request_id: Option<String>,
         trace_id: Option<String>,
+        should_client_retry: bool,
     ) -> Self {
         match err {
-            errors::Error::ErrorCode(code) => Self::Error {
+            errors::Error::ErrorCode(code) => WsServerEvents::Error {
                 code: code.get_code(),
                 message: code.get_message(),
                 error_detail: Some(code.get_error_detail()),
                 trace_id: trace_id.clone(),
                 request_id: request_id.clone(),
+                should_client_retry,
             },
-            _ => Self::Error {
+            _ => WsServerEvents::Error {
                 code: StatusCode::INTERNAL_SERVER_ERROR.into(),
                 message: err.to_string(),
                 error_detail: None,
                 trace_id,
                 request_id,
+                should_client_retry,
             },
         }
     }
+
+    pub fn get_trace_id(&self) -> String {
+        match self {
+            Self::SearchResponse { trace_id, .. } => trace_id.to_string(),
+            #[cfg(feature = "enterprise")]
+            Self::CancelResponse { trace_id, .. } => trace_id.to_string(),
+            Self::Error { trace_id, .. } => trace_id.clone().unwrap_or_default(),
+            Self::End { trace_id } => trace_id.clone().unwrap_or_default(),
+            Self::Ping(_) => "".to_string(),
+            Self::Pong(_) => "".to_string(),
+        }
+    }
+
+    pub fn should_clean_trace_id(&self) -> Option<String> {
+        match self {
+            #[cfg(feature = "enterprise")]
+            Self::CancelResponse {
+                trace_id,
+                is_success,
+            } if *is_success => Some(trace_id.to_owned()),
+            Self::Error { trace_id, .. } => trace_id.to_owned(),
+            Self::End { trace_id } => trace_id.to_owned(),
+            _ => None,
+        }
+    }
+}
+
+impl From<WsClientEvents> for tungstenite::protocol::Message {
+    fn from(event: WsClientEvents) -> Self {
+        let payload = serde_json::to_value(&event).unwrap_or_default();
+        tungstenite::protocol::Message::Text(payload.to_string())
+    }
+}
+
+impl From<WsServerEvents> for tungstenite::protocol::Message {
+    fn from(event: WsServerEvents) -> Self {
+        let payload = serde_json::to_value(&event).unwrap_or_default();
+        tungstenite::protocol::Message::Text(payload.to_string())
+    }
+}
+
+impl TryFrom<serde_json::Value> for WsServerEvents {
+    type Error = String;
+
+    fn try_from(
+        value: serde_json::Value,
+    ) -> Result<Self, <Self as TryFrom<serde_json::Value>>::Error> {
+        match value["type"].as_str() {
+            Some("search_response") => serde_json::from_value(value).map_err(|e| e.to_string()),
+            Some("cancel_response") => serde_json::from_value(value).map_err(|e| e.to_string()),
+            Some("error") => serde_json::from_value(value).map_err(|e| e.to_string()),
+            Some("end") => serde_json::from_value(value).map_err(|e| e.to_string()),
+            _ => Err("Unknown message type".to_string()),
+        }
+    }
 }

web/src/composables/dashboard/usePanelDataLoader.ts

@@ -76,23 +76,13 @@ export const usePanelDataLoader = (
   };
   let runCount = 0;
 
-  const searchRetriesCount = ref<{ [key: string]: number }>({});
-
   const store = useStore();
 
-  // Add cleanup function
-  const cleanupSearchRetries = (traceId: string) => {
-    removeTraceId(traceId);
-    if (searchRetriesCount.value[traceId]) {
-      delete searchRetriesCount.value[traceId];
-    }
-  };
-
   const {
     fetchQueryDataWithWebSocket,
     sendSearchMessageBasedOnRequestId,
     cancelSearchQueryBasedOnRequestId,
-    closeSocketBasedOnRequestId,
+    cleanUpListeners,
   } = useSearchWebSocket();
 
   const { refreshAnnotations } = useAnnotations(
@@ -145,9 +135,6 @@ export const usePanelDataLoader = (
     lastTriggeredAt: null as any,
     isCachedDataDifferWithCurrentTimeRange: false,
     searchRequestTraceIds: <string[]>[],
-    searchWebSocketRequestIdsAndTraceIds: <
-      { requestId: string; traceId: string }[]
-    >[],
     isOperationCancelled: false,
   });
 
@@ -322,18 +309,19 @@ export const usePanelDataLoader = (
 
     state.isOperationCancelled = true;
 
-    if (isWebSocketEnabled() && state.searchWebSocketRequestIdsAndTraceIds) {
+    if (isWebSocketEnabled() && state.searchRequestTraceIds) {
       try {
-        // loop on state.searchWebSocketRequestIdsAndTraceIds
-        state.searchWebSocketRequestIdsAndTraceIds.forEach((it) => {
-          if (it?.requestId && it?.traceId) {
-            cancelSearchQueryBasedOnRequestId(it?.requestId, it?.traceId);
-          }
+        // loop on state.searchRequestTraceIds
+        state.searchRequestTraceIds.forEach((traceId) => {
+          cancelSearchQueryBasedOnRequestId({
+            trace_id: traceId,
+            org_id: store?.state?.selectedOrganization?.identifier,
+          });
         });
       } catch (error) {
         console.error("Error during WebSocket cleanup:", error);
       } finally {
-        state.searchWebSocketRequestIdsAndTraceIds = [];
+        state.searchRequestTraceIds = [];
       }
     }
     if (abortController) {
@@ -608,11 +596,7 @@ export const usePanelDataLoader = (
   };
 
   // Limit, aggregation, vrl function, pagination, function error and query error
-  const handleSearchResponse = (
-    requestId: string,
-    payload: any,
-    response: any,
-  ) => {
+  const handleSearchResponse = (payload: any, response: any) => {
     try {
       if (response.type === "search_response") {
         handleHistogramResponse(payload, response);
@@ -639,15 +623,18 @@ export const usePanelDataLoader = (
     }
   };
 
-  const sendSearchMessage = async (requestId: string, payload: any) => {
+  const sendSearchMessage = async (payload: any) => {
     // check if query is already canceled, if it is, close the socket
     if (state.isOperationCancelled) {
-      closeSocketBasedOnRequestId(requestId);
       state.isOperationCancelled = false;
+
+      // clean up the listeners
+      cleanUpListeners(payload.traceId);
+
       return;
     }
 
-    sendSearchMessageBasedOnRequestId(requestId, {
+    sendSearchMessageBasedOnRequestId({
       type: "search",
       content: {
         trace_id: payload.traceId,
@@ -662,6 +649,7 @@ export const usePanelDataLoader = (
         },
         stream_type: payload.pageType,
         search_type: searchType.value ?? "dashboards",
+        org_id: store?.state?.selectedOrganization?.identifier,
         use_cache: (window as any).use_cache ?? true,
         dashboard_id: dashboardId?.value,
         folder_id: folderId?.value,
@@ -670,70 +658,27 @@ export const usePanelDataLoader = (
     });
   };
 
-  const handleSearchClose = (
-    requestId: string,
-    payload: any,
-    response: any,
-  ) => {
-    const MAX_RETRIES = 2;
-    const RECONNECT_DELAY = 1000; // 1 second
-
-    removeRequestId(requestId);
-
-    if (response.code === 1001 || response.code === 1006) {
-      const retryCount = searchRetriesCount.value[payload.traceId] || 0;
-      searchRetriesCount.value[payload.traceId] = retryCount + 1;
-
-      if (retryCount < MAX_RETRIES) {
-        state.loading = true;
-        state.isOperationCancelled = false;
-
-        setTimeout(() => {
-          try {
-            const newRequestId = fetchQueryDataWithWebSocket(payload, {
-              open: sendSearchMessage,
-              close: handleSearchClose,
-              error: handleSearchError,
-              message: handleSearchResponse,
-            }) as string;
-
-            addRequestId(newRequestId, payload.traceId);
-          } catch (error: any) {
-            console.error("Error reconnecting WebSocket:", error);
-            cleanupSearchRetries(payload?.traceId);
-            handleSearchError(requestId, payload, {
-              content: {
-                message: "Failed to reconnect WebSocket",
-                trace_id: payload.traceId,
-                code: response.code,
-                error_detail: error?.message,
-              },
-            });
-          }
-        }, RECONNECT_DELAY);
-
-        return;
-      } else {
-        // remove current traceId
-        cleanupSearchRetries(payload?.traceId);
-        handleSearchError(requestId, payload, {
-          content: {
-            message:
-              "WebSocket connection terminated unexpectedly. Please check your network and try again",
-            trace_id: payload.traceId,
-            code: response.code,
-            error_detail: "",
-          },
-        });
-      }
-    }
+  const handleSearchClose = (payload: any, response: any) => {
+    removeTraceId(payload?.traceId);
 
     if (response.type === "error") {
       processApiError(response?.content, "sql");
     }
 
-    // remove current traceId
-    cleanupSearchRetries(payload?.traceId);
+
+    const errorCodes = [1001, 1006, 1010, 1011, 1012, 1013];
+
+    if (errorCodes.includes(response.code)) {
+      handleSearchError(payload, {
+        content: {
+          message:
+            "WebSocket connection terminated unexpectedly. Please check your network and try again",
+          trace_id: payload.traceId,
+          code: response.code,
+          error_detail: "",
+        }
+      });
+    }
 
     // set loading to false
     state.loading = false;
@@ -743,14 +688,23 @@ export const usePanelDataLoader = (
     saveCurrentStateToCache();
   };
 
-  const handleSearchError = (
-    requestId: string,
-    payload: any,
-    response: any,
-  ) => {
-    removeRequestId(requestId);
+  const handleSearchReset = (payload: any) => {
+    // reset old state data
+    state.data = [];
+    state.resultMetaData = [];
 
-    cleanupSearchRetries(payload?.traceId);
+    getDataThroughWebSocket(
+      payload.queryReq.query,
+      payload.queryReq.it,
+      payload.queryReq.startISOTimestamp,
+      payload.queryReq.endISOTimestamp,
+      payload.pageType,
+      payload.currentQueryIndex,
+    );
+  }
+
+  const handleSearchError = (payload: any, response: any) => {
+    removeTraceId(payload.traceId);
 
     // set loading to false
     state.loading = false;
@@ -765,7 +719,6 @@ export const usePanelDataLoader = (
     startISOTimestamp: string,
     endISOTimestamp: string,
     pageType: string,
-    abortControllerRef: AbortController,
     currentQueryIndex: number,
   ) => {
     try {
@@ -785,7 +738,6 @@ export const usePanelDataLoader = (
           it,
           startISOTimestamp,
           endISOTimestamp,
-          abortControllerRef,
           currentQueryIndex,
         },
         type: "histogram",
@@ -794,14 +746,16 @@ export const usePanelDataLoader = (
         org_id: store?.state?.selectedOrganization?.identifier,
         pageType,
       };
-      const requestId = fetchQueryDataWithWebSocket(payload, {
+
+      fetchQueryDataWithWebSocket(payload, {
         open: sendSearchMessage,
         close: handleSearchClose,
         error: handleSearchError,
         message: handleSearchResponse,
-      }) as string;
+        reset: handleSearchReset,
+      });
 
-      addRequestId(requestId, traceId);
+      addTraceId(traceId);
     } catch (e: any) {
       state.errorDetail = e?.message || e;
       state.loading = false;
@@ -969,7 +923,8 @@ export const usePanelDataLoader = (
         const abortControllerRef = abortController;
 
         try {
-          // reset old state data
+
+          // Call search API
           state.data = [];
           state.metadata = {
             queries: [],
@@ -977,9 +932,7 @@ export const usePanelDataLoader = (
           state.resultMetaData = [];
           state.annotations = [];
           state.isOperationCancelled = false;
-
-          // Call search API
-
+          
           // Get the page type from the first query in the panel schema
           const pageType = panelSchema.value.queries[0]?.fields?.stream_type;
 
@@ -1225,7 +1178,6 @@ export const usePanelDataLoader = (
                   startISOTimestamp,
                   endISOTimestamp,
                   pageType,
-                  abortControllerRef,
                   panelQueryIndex,
                 );
               } else {
@@ -1525,20 +1477,6 @@ export const usePanelDataLoader = (
     );
   };
 
-  const addRequestId = (requestId: string, traceId: string) => {
-    state.searchWebSocketRequestIdsAndTraceIds = [
-      ...state.searchWebSocketRequestIdsAndTraceIds,
-      { requestId, traceId },
-    ];
-  };
-
-  const removeRequestId = (requestId: string) => {
-    state.searchWebSocketRequestIdsAndTraceIds =
-      state.searchWebSocketRequestIdsAndTraceIds.filter(
-        (id: any) => id.requestId !== requestId,
-      );
-  };
-
   const hasAtLeastOneQuery = () =>
     panelSchema.value.queries?.some((q: any) => q?.query);
 
@@ -1927,6 +1865,23 @@ export const usePanelDataLoader = (
       observer.disconnect();
     }
 
+    // for websocket
+    if (isWebSocketEnabled() && state.searchRequestTraceIds) {
+      try {
+        // loop on state.searchRequestTraceIds
+        state.searchRequestTraceIds.forEach((traceId) => {
+          cancelSearchQueryBasedOnRequestId({
+            trace_id: traceId,
+            org_id: store?.state?.selectedOrganization?.identifier,
+          });
+        });
+      } catch (error) {
+        console.error("Error during WebSocket cleanup:", error);
+      } finally {
+        state.searchRequestTraceIds = [];
+      }
+    }
+
     // remove cancelquery event
     window.removeEventListener("cancelQuery", cancelQueryAbort);
   });

web/src/composables/useLogs.ts

@@ -242,9 +242,7 @@ const defaultObject = {
     customDownloadQueryObj: <any>{},
     functionError: "",
     searchRequestTraceIds: <string[]>[],
-    searchWebSocketRequestIdsAndTraceIds: <
-      { requestId: string; traceId: string }[]
-    >[],
+    searchWebSocketRequestIdsAndTraceIds: <{ traceId: string }[]>[],
     isOperationCancelled: false,
     searchRetriesCount: <{ [key: string]: number }>{},
     actionId: null,
@@ -635,8 +633,9 @@ const useLogs = () => {
   const updateUrlQueryParams = () => {
     const query = generateURLQuery(false);
     if (
-      Object.hasOwn(query, "type") &&
-      query.type == "search_history_re_apply" || query.type == "search_scheduler"
+      (Object.hasOwn(query, "type") &&
+        query.type == "search_history_re_apply") ||
+      query.type == "search_scheduler"
     ) {
       delete query.type;
     }
@@ -1658,7 +1657,6 @@ const useLogs = () => {
           JSON.stringify(queryReq),
         );
 
-
         // get the current page detail and set it into query request
         queryReq.query.start_time =
           searchObj.data.queryResults.partitionDetail.paginations[
@@ -1729,7 +1727,6 @@ const useLogs = () => {
           searchObj.meta.refreshHistogram = false;
           if (searchObj.data.queryResults.hits.length > 0) {
             if (searchObj.data.stream.selectedStream.length > 1) {
-
               searchObj.data.histogram = {
                 xData: [],
                 yData: [],
@@ -1839,7 +1836,6 @@ const useLogs = () => {
           }
 
           if (searchObj.data.stream.selectedStream.length > 1) {
-
             searchObj.data.histogram = {
               xData: [],
               yData: [],
@@ -1978,7 +1974,6 @@ const useLogs = () => {
   }
 
   function resetHistogramWithError(errorMsg: string, errorCode: number = 0) {
-
     searchObj.data.histogram = {
       xData: [],
       yData: [],
@@ -2010,7 +2005,6 @@ const useLogs = () => {
   }
 
   function generateHistogramSkeleton() {
-
     if (
       searchObj.data.queryResults.hasOwnProperty("aggs") &&
       searchObj.data.queryResults.aggs
@@ -2123,12 +2117,12 @@ const useLogs = () => {
         if(queryReq.query.action_id) delete queryReq.query.action_id;
 
         queryReq.query["sql_mode"] = "full";
-  
+
         queryReq.query.track_total_hits = true;
-  
+
         const { traceparent, traceId } = generateTraceContext();
         addTraceId(traceId);
-  
+
         searchService
           .search(
             {
@@ -2165,8 +2159,7 @@ const useLogs = () => {
                 }
               }
             }
-  
-  
+
             let regeratePaginationFlag = false;
             if (res.data.hits.length != searchObj.meta.resultGrid.rowsPerPage) {
               regeratePaginationFlag = true;
@@ -2182,11 +2175,11 @@ const useLogs = () => {
           .catch((err) => {
             searchObj.loading = false;
             searchObj.loadingCounter = false;
-  
-          // Reset cancel query on search error
-          searchObj.data.isOperationCancelled = false;
 
-          let trace_id = "";
+            // Reset cancel query on search error
+            searchObj.data.isOperationCancelled = false;
+
+            let trace_id = "";
             searchObj.data.countErrorMsg =
               "Error while retrieving total events: ";
             if (err.response != undefined) {
@@ -2198,17 +2191,17 @@ const useLogs = () => {
                 trace_id = err?.trace_id;
               }
             }
-  
+
             const customMessage = logsErrorMessage(err?.response?.data.code);
             searchObj.data.errorCode = err?.response?.data.code;
-  
+
             notificationMsg.value = searchObj.data.countErrorMsg;
-  
+
             if (err?.request?.status >= 429) {
               notificationMsg.value = err?.response?.data?.message;
               searchObj.data.countErrorMsg += err?.response?.data?.message;
             }
-  
+
             if (trace_id) {
               searchObj.data.countErrorMsg += " TraceID:" + trace_id;
               notificationMsg.value += " TraceID:" + trace_id;
@@ -2223,7 +2216,6 @@ const useLogs = () => {
         searchObj.loadingCounter = false;
         reject(false);
       }
-
     });
   };
 
@@ -2237,8 +2229,7 @@ const useLogs = () => {
     updateGridColumns();
 
     filterHitsColumns();
-      searchObj.data.histogram.chartParams.title = getHistogramTitle();
-
+    searchObj.data.histogram.chartParams.title = getHistogramTitle();
   };
 
   const getPaginatedData = async (
@@ -2266,7 +2257,7 @@ const useLogs = () => {
         showCancelSearchNotification();
         return;
       }
-      if (searchObj.meta.jobId != ""){
+      if (searchObj.meta.jobId != "") {
         searchObj.meta.resultGrid.rowsPerPage = queryReq.query.size;
       }
       const parsedSQL: any = fnParsedSQL();
@@ -2350,7 +2341,7 @@ const useLogs = () => {
           // check for total records update for the partition and update pagination accordingly
           // searchObj.data.queryResults.partitionDetail.partitions.forEach(
           //   (item: any, index: number) => {
-          if (searchObj.meta.jobId == ""){
+          if (searchObj.meta.jobId == "") {
             for (const [
               index,
               item,
@@ -2425,9 +2416,8 @@ const useLogs = () => {
               );
               searchObj.data.queryResults.hits = res.data.hits;
             } else {
-              if (searchObj.meta.jobId != ""){
+              if (searchObj.meta.jobId != "") {
                 searchObj.data.queryResults.total = res.data.total;
-                
               }
               if (!queryReq.query.hasOwnProperty("track_total_hits")) {
                 delete res.data.total;
@@ -2470,7 +2460,7 @@ const useLogs = () => {
             searchObj.data.queryResults.subpage++;
 
             setTimeout(async () => {
-                processPostPaginationData();
+              processPostPaginationData();
 
               // searchObj.data.functionError = "";
               // if (
@@ -2482,7 +2472,7 @@ const useLogs = () => {
             }, 0);
             await getPaginatedData(queryReq, true);
           }
-          if (searchObj.meta.jobId != ""){
+          if (searchObj.meta.jobId != "") {
             searchObj.meta.resultGrid.rowsPerPage = queryReq.query.size;
             searchObj.data.queryResults.pagination = [];
             refreshJobPagination(true);
@@ -2592,13 +2582,13 @@ const useLogs = () => {
 
   const routeToSearchSchedule = () => {
     router.push({
-      query:{
+      query: {
         action: "search_scheduler",
         org_identifier: store.state.selectedOrganization.identifier,
-        type: "search_scheduler_list"
-      }
+        type: "search_scheduler_list",
+      },
     });
-  }
+  };
 
   const getHistogramQueryData = (queryReq: any) => {
     return new Promise((resolve, reject) => {
@@ -3949,18 +3939,14 @@ const useLogs = () => {
       await getFunctions();
       if (isActionsEnabled.value) await getActions();
       await extractFields();
-      if (searchObj.meta.jobId == ""){
+      if (searchObj.meta.jobId == "") {
         await getQueryData();
-      }
-      else{
-
+      } else {
         await getJobData();
-
       }
       refreshData();
     } catch (e: any) {
       searchObj.loading = false;
-
     }
   };
   const loadJobData = async () => {
@@ -4008,7 +3994,7 @@ const useLogs = () => {
         }
       }, 120000);
       await getQueryData();
-      clearTimeout(queryTimeout); 
+      clearTimeout(queryTimeout);
     } catch (e: any) {
       console.log("Error while loading logs data");
     }
@@ -4381,17 +4367,16 @@ const useLogs = () => {
       );
   };
 
-  const addRequestId = (requestId: string, traceId: string) => {
+  const addRequestId = (traceId: string) => {
     searchObj.data.searchWebSocketRequestIdsAndTraceIds.push({
-      requestId,
       traceId,
     });
   };
 
-  const removeRequestId = (requestId: string) => {
+  const removeRequestId = (traceId: string) => {
     searchObj.data.searchWebSocketRequestIdsAndTraceIds =
       searchObj.data.searchWebSocketRequestIdsAndTraceIds.filter(
-        (id) => id.requestId !== requestId,
+        (id) => id.traceId !== traceId,
       );
   };
 
@@ -4626,10 +4611,10 @@ const useLogs = () => {
 
   const extractValueQuery = () => {
     try {
-      if(searchObj.meta.sqlMode == false || searchObj.data.query == "") {
+      if (searchObj.meta.sqlMode == false || searchObj.data.query == "") {
         return {};
       }
-      
+
       const orgQuery: string = searchObj.data.query
         .split("\n")
         .filter((line: string) => !line.trim().startsWith("--"))
@@ -4895,7 +4880,7 @@ const useLogs = () => {
         throw new Error("Failed to initialize WebSocket connection");
       }
 
-      addRequestId(requestId, payload.traceId);
+      addRequestId(payload.traceId);
     } catch (e: any) {
       console.error("Error while getting data through web socket", e);
       searchObj.loading = false;
@@ -4940,13 +4925,93 @@ const useLogs = () => {
       close: handleSearchClose,
       error: handleSearchError,
       message: handleSearchResponse,
+      reset: handleSearchReset,
     }) as string;
   };
 
-  const sendSearchMessage = (requestId: string, queryReq: any) => {
+
+  const handleSearchReset = async (data: any) => {
+    // reset query data
+    try {
+      if(data.type === "search") {
+        if(!data.isPagination) {
+          resetQueryData();
+          searchObj.data.queryResults = {};
+        }
+
+        // reset searchAggData
+        searchAggData.total = 0;
+        searchAggData.hasAggregation = false;
+        searchObj.meta.showDetailTab = false;
+        searchObj.meta.searchApplied = true;
+        searchObj.data.functionError = "";
+      
+        searchObj.data.errorCode = 0;
+      
+        if(!data.isPagination) {
+          // Histogram reset
+          searchObj.data.histogram = {
+            xData: [],
+            yData: [],
+            chartParams: {
+              title: "",
+              unparsed_x_data: [],
+              timezone: "",
+            },
+            errorCode: 0,
+            errorMsg: "",
+            errorDetail: "",
+          };
+          resetHistogramError();
+        }
+
+        const payload = buildWebSocketPayload(data.queryReq, data.isPagination, "search");
+
+        initializeWebSocketConnection(payload);
+        addRequestId(payload.traceId);
+      }
+
+      if(data.type === "histogram" || data.type === "pageCount") {
+        searchObj.data.queryResults.aggs = [];
+        searchObj.data.histogram = {
+          xData: [],
+          yData: [],
+          chartParams: {
+            title: "",
+            unparsed_x_data: [],
+            timezone: "",
+          },
+          errorCode: 0,
+          errorMsg: "",
+          errorDetail: "",
+        };
+        resetHistogramError();
+
+        searchObj.loadingHistogram = true;
+
+        if(data.type === "histogram") await generateHistogramSkeleton();
+
+        histogramResults = [];
+
+        const payload = buildWebSocketPayload(
+          data.queryReq,
+          false,
+          "histogram",
+        );
+
+        initializeWebSocketConnection(payload);
+
+        addRequestId(payload.traceId);
+      }
+    } catch(e: any){
+      console.error("Error while resetting search", e);
+    }
+  };
+
+  const sendSearchMessage = (queryReq: any) => {
     try {
       if (searchObj.data.isOperationCancelled) {
-        closeSocketBasedOnRequestId(requestId);
+        closeSocketBasedOnRequestId(queryReq.traceId);
         return;
       }
 
@@ -4960,6 +5025,7 @@ const useLogs = () => {
           stream_type: searchObj.data.stream.streamType,
           search_type: "ui",
           use_cache: (window as any).use_cache ?? true,
+          org_id: searchObj.organizationIdentifier,
         },
       };
 
@@ -4971,7 +5037,7 @@ const useLogs = () => {
         payload.content.payload["clusters"] = queryReq.queryReq.clusters;
       }
 
-      sendSearchMessageBasedOnRequestId(requestId, payload);
+      sendSearchMessageBasedOnRequestId(payload);
     } catch (e: any) {
       searchObj.loading = false;
       showErrorNotification(
@@ -4989,7 +5055,6 @@ const useLogs = () => {
 
   // Limit, aggregation, vrl function, pagination, function error and query error
   const handleSearchResponse = (
-    requestId: string,
     payload: WebSocketSearchPayload,
     response: WebSocketSearchResponse | WebSocketErrorResponse,
   ) => {
@@ -5025,7 +5090,7 @@ const useLogs = () => {
     }
 
     if (response.type === "error") {
-      handleSearchError(requestId, payload.traceId, response);
+      handleSearchError(payload.traceId, response);
     }
 
     if (response.type === "cancel_response") {
@@ -5363,7 +5428,7 @@ const useLogs = () => {
         );
         const requestId = initializeWebSocketConnection(payload);
 
-        addRequestId(requestId, payload.traceId);
+        addRequestId(payload.traceId);
       }
     } else if (searchObj.meta.sqlMode && isLimitQuery(parsedSQL)) {
       resetHistogramWithError("Histogram is not available for limit queries.");
@@ -5428,12 +5493,8 @@ const useLogs = () => {
     return searchObj.data.queryResults.aggs === undefined;
   }
 
-  const handleSearchClose = (
-    requestId: string,
-    payload: any,
-    response: any,
-  ) => {
-    if (requestId) removeRequestId(requestId);
+  const handleSearchClose = (payload: any, response: any) => {
+    if (payload.traceId) removeRequestId(payload.traceId);
 
     // Any case where below logic may end in recursion
     if (payload.traceId) delete searchPartitionMap[payload.traceId];
@@ -5448,44 +5509,25 @@ const useLogs = () => {
       return;
     }
 
-    if (response.code === 1001 || response.code === 1006) {
-      if (!searchObj.data.searchRetriesCount[payload.traceId]) {
-        searchObj.data.searchRetriesCount[payload.traceId] = 1;
-      } else {
-        searchObj.data.searchRetriesCount[payload.traceId] += 1;
-      }
+    //TODO Omkar: Remove the duplicate error codes, are present same in useSearchWebSocket.ts
+    const errorCodes = [1001, 1006, 1010, 1011, 1012, 1013];
 
-      if (
-        searchObj.data.searchRetriesCount[payload.traceId] <=
-        searchReconnectDelay
-      ) {
-        if (payload.type === "search") searchObj.loading = true;
-        if (payload.type === "histogram") searchObj.loadingHistogram = true;
-
-        setTimeout(() => {
-          const requestId = initializeWebSocketConnection(payload);
-
-          addRequestId(requestId, payload.traceId);
-        }, maxSearchRetries);
-
-        return;
-      } else {
-        handleSearchError(requestId, payload, {
-          content: {
-            message:
-              "WebSocket connection terminated unexpectedly. Please check your network and try again",
-            trace_id: payload.traceId,
-            code: response.code,
-            error_detail: "",
-          },
-          type: "error",
-        });
-      }
+    if (errorCodes.includes(response.code)) {
+      handleSearchError(payload, {
+        content: {
+          message:
+            "WebSocket connection terminated unexpectedly. Please check your network and try again",
+          trace_id: payload.traceId,
+          code: response.code,
+          error_detail: "",
+        },
+        type: "error",
+      });
     }
 
-    if (searchObj.data.searchRetriesCount[payload.traceId]) {
-      delete searchObj.data.searchRetriesCount[payload.traceId];
-    }
+    // if (searchObj.data.searchRetriesCount[payload.traceId]) {
+    //   delete searchObj.data.searchRetriesCount[payload.traceId];
+    // }
 
     if (payload.traceId) removeTraceId(payload.traceId);
 
@@ -5504,11 +5546,7 @@ const useLogs = () => {
     searchObj.data.isOperationCancelled = false;
   };
 
-  const handleSearchError = (
-    requestId: string,
-    request: any,
-    err: WebSocketErrorResponse,
-  ) => {
+  const handleSearchError = (request: any, err: WebSocketErrorResponse) => {
     searchObj.loading = false;
     searchObj.loadingHistogram = false;
 
@@ -5564,7 +5602,8 @@ const useLogs = () => {
       const { rowsPerPage } = searchObj.meta.resultGrid;
       const { currentPage } = searchObj.data.resultGrid;
 
-      if(searchObj.meta.jobId != "") searchObj.meta.resultGrid.rowsPerPage = 100;
+      if (searchObj.meta.jobId != "")
+        searchObj.meta.resultGrid.rowsPerPage = 100;
 
       let total = 0;
       let totalPages = 0;
@@ -5667,7 +5706,7 @@ const useLogs = () => {
 
     const requestId = initializeWebSocketConnection(payload);
 
-    addRequestId(requestId, payload.traceId);
+    addRequestId(payload.traceId);
   };
 
   const sendCancelSearchMessage = (searchRequests: any[]) => {
@@ -5680,8 +5719,11 @@ const useLogs = () => {
       searchObj.data.isOperationCancelled = true;
 
       // loop on all requestIds
-      searchRequests.forEach(({ requestId, traceId }) => {
-        cancelSearchQueryBasedOnRequestId(requestId, traceId);
+      searchRequests.forEach(({ traceId }) => {
+        cancelSearchQueryBasedOnRequestId({
+          trace_id: traceId,
+          org_id: store?.state?.selectedOrganization?.identifier,
+        });
       });
     } catch (error: any) {
       console.error("Failed to cancel WebSocket searches:", error);
@@ -5807,7 +5849,8 @@ const useLogs = () => {
     initializeWebSocketConnection,
     addRequestId,
     routeToSearchSchedule,
-    isActionsEnabled
+    isActionsEnabled,
+    sendCancelSearchMessage,
   };
 };
 

web/src/composables/useSearchWebSocket.ts

@@ -1,9 +1,157 @@
 import { getUUID, getWebSocketUrl } from "@/utils/zincutils";
 import useWebSocket from "@/composables/useWebSocket";
 import type { SearchRequestPayload } from "@/ts/interfaces";
+import { useStore } from "vuex";
+import { ref } from "vue";
+import authService from "@/services/auth";
+
+type MessageHandler = (event: MessageEvent) => void;
+type OpenHandler = (event: Event) => void;
+type CloseHandler = (event: CloseEvent) => void;
+type ErrorHandler = (event: Event) => void;
+
+type TraceRecord = {
+  open: OpenHandler[];
+  message: MessageHandler[];
+  close: CloseHandler[];
+  error: ErrorHandler[];
+  reset: ((data: any) => void)[];
+  isInitiated: boolean;
+  isActive: boolean;
+  socketId: string | null;
+  data: any;
+};
+
+const webSocket = useWebSocket();
+
+const traces: Record<string, TraceRecord> = {};
+// const openHandlers: OpenHandler[] = [];
+const socketId = ref<string | null>(null);
+
+const inactiveSocketId = ref<string | null>(null);
+
+const isCreatingSocket = ref(false);
+
+const socketRetryCodes = [1001, 1006, 1010, 1011, 1012, 1013];
+
+const socketFailureCount = ref(0);
+
+const maxSearchRetries = 3;
+
+const isInDrainMode = ref(false);
 
 const useSearchWebSocket = () => {
-  const webSocket = useWebSocket();
+  const store = useStore();
+
+  const onOpen = (response: any) => {
+    isCreatingSocket.value = false;
+    Object.keys(traces).forEach((traceId) => {
+      if(traces[traceId].isActive) {
+        traces[traceId].isInitiated = true;
+        traces[traceId].open.forEach((handler: any) => handler(response));
+      }
+    });
+  };
+
+  const onMessage = (response: any) => {
+    if (response.type === "end") {
+      socketFailureCount.value = 0;
+      traces[response.content.trace_id]?.close?.forEach((handler: any) =>
+        handler(response),
+      );
+      cleanUpListeners(response.content.trace_id);
+      return;
+    }
+
+    traces[response.content.trace_id]?.message?.forEach((handler: any) =>
+      handler(response),
+    );
+  };
+
+  const onClose = async (response: any, _socketId: string) => {    
+    isCreatingSocket.value = false;
+    socketId.value = null;
+
+    const shouldRetry = socketRetryCodes.includes(response.code) || isInDrainMode.value;
+    
+    if(shouldRetry) socketFailureCount.value++;
+
+
+    // if(inactiveSocketId.value && _socketId === inactiveSocketId.value) {
+    //   inactiveSocketId.value = null;
+    // }
+
+    if (shouldRetry && socketFailureCount.value < maxSearchRetries) {
+      if(isInDrainMode.value) {
+        await resetAuthToken();
+        return;
+      }
+
+      setTimeout(() => {
+        Object.keys(traces).forEach((traceId) => {
+          if(traces[traceId].isInitiated) {
+            response.code = 1000;
+            traces[traceId]?.close.forEach((handler: any) => handler(response));
+            traces[traceId]?.reset.forEach((handler: any) => handler(traces[traceId].data));
+            cleanUpListeners(traceId);        
+          }
+        });
+      }, 1000);
+    } else {
+      Object.keys(traces).forEach((traceId) => {
+        if(traces[traceId].isInitiated) {
+          traces[traceId]?.close.forEach((handler: any) => handler(response));
+          cleanUpListeners(traceId);
+        }
+      });
+    }
+  };
+
+  const onError = async (response: any) => {
+    if(response.type === 'error'){
+      if(response.content.should_client_retry && response.content.trace_id){
+        setTimeout(() => {
+          retryActiveTrace(response.content.trace_id, response);
+        }, 300)
+
+        return;
+      }
+
+      if(response.content.code === 401) {
+        // Store the current socketId as inactive and clear it
+        isInDrainMode.value = true;
+      }
+    }
+
+    traces[response.content.trace_id]?.error?.forEach((handler: any) =>
+      handler(response),
+    );
+    // cleanUpListeners(response.traceId)
+  };
+
+  const createSocketConnection = (org_id: string) => {
+    isCreatingSocket.value = true;
+
+    socketId.value = getUUID();
+
+    Object.keys(traces).forEach((traceId) => {
+      if(traces[traceId].isActive) traces[traceId].socketId = socketId.value;
+    });
+
+    const url = getWebSocketUrl(socketId.value, org_id);
+    // If needed we can store the socketID in global state
+    webSocket.connect(socketId.value, url);
+
+    webSocket.addOpenHandler(socketId.value, onOpen);
+
+    // When we receive message from BE/server
+    webSocket.addMessageHandler(socketId.value, onMessage);
+
+    // On closing of ws, when search is completed Server closes the WS
+    webSocket.addCloseHandler(socketId.value, onClose);
+
+    webSocket.addErrorHandler(socketId.value, onError);
+  };
 
   const fetchQueryDataWithWebSocket = (
     data: {
@@ -14,42 +162,39 @@ const useSearchWebSocket = () => {
       org_id: string;
     },
     handlers: {
-      open: (requestId: string, data: any, response: any) => void;
-      message: (requestId: string, data: any, response: any) => void;
-      close: (requestId: string, data: any, response: any) => void;
-      error: (requestId: string, data: any, response: any) => void;
-    },
+      open: (data: any, response: any) => void;
+      message: (data: any, response: any) => void;
+      close: (data: any, response: any) => void;
+      error: (data: any, response: any) => void;
+      reset: (data: any, response: any) => void;
+    }  
   ) => {
     try {
-      const requestId = getUUID();
-      const url = getWebSocketUrl(requestId, data.org_id);
+      traces[data.traceId] = {
+        open: [],
+        message: [],
+        close: [],
+        error: [],
+        reset: [],
+        data: data,
+        isActive: true, //  True if the trace id is on current active socket. If false, If we receive 401, we will initiate a new socket and mark the old socket traas inactive
+        socketId: socketId.value, // Track which socket this search was initiated on
+        isInitiated: false, // True if the search was initiated on the current socket
+      };
 
-      webSocket.connect(requestId, url);
+      traces[data.traceId].message.push(handlers.message.bind(null, data));
+      traces[data.traceId].close.push(handlers.close.bind(null, data));
+      traces[data.traceId].error.push(handlers.error.bind(null, data));
+      traces[data.traceId].open.push(handlers.open.bind(null, data));
+      traces[data.traceId].reset.push(handlers.reset.bind(null, data));
 
-      // Gets called when socket connect is established
-      webSocket.addOpenHandler(
-        requestId,
-        handlers.open.bind(null, requestId, data),
-      );
+      if (isInDrainMode.value) {
+        return data.traceId;
+      }
 
-      // When we receive message from BE/server
-      webSocket.addMessageHandler(
-        requestId,
-        handlers.message.bind(null, requestId, data),
-      );
+      initiateSocketConnection(data, handlers);
 
-      // On closing of ws, when search is completed Server closes the WS
-      webSocket.addCloseHandler(
-        requestId,
-        handlers.close.bind(null, requestId, data),
-      );
-
-      webSocket.addErrorHandler(
-        requestId,
-        handlers.error.bind(null, requestId, data),
-      );
-      
-      return requestId;
+      return data.traceId;
     } catch (error: any) {
       console.error(
         `Error in fetching search data: ${error instanceof Error ? error.message : String(error)}`,
@@ -58,9 +203,36 @@ const useSearchWebSocket = () => {
     }
   };
 
-  const sendSearchMessageBasedOnRequestId = (requestId: string, data: any) => {
+  const initiateSocketConnection = (data: {
+    queryReq: SearchRequestPayload;
+    type: "search" | "histogram" | "pageCount";
+    isPagination: boolean;
+    traceId: string;
+    org_id: string;
+  }, handlers: {
+    open: (data: any, response: any) => void;
+    message: (data: any, response: any) => void;
+    close: (data: any, response: any) => void;
+    error: (data: any, response: any) => void;
+    reset: (data: any, response: any) => void;
+  }) => {
+    if (!socketId.value) {
+      createSocketConnection(data.org_id);
+    } else if (!isCreatingSocket.value) {
+      traces[data.traceId].isInitiated = true;
+      handlers.open(data, null);
+    }
+  }
+
+  const sendSearchMessageBasedOnRequestId = (data: any) => {
     try {
-      webSocket.sendMessage(requestId, JSON.stringify(data));
+      
+      if(!traces[data.content.traceId]?.isActive && inactiveSocketId.value) {
+        webSocket.sendMessage(inactiveSocketId.value as string, JSON.stringify(data));
+        return;
+      }
+
+      webSocket.sendMessage(socketId.value as string, JSON.stringify(data));
     } catch (error: any) {
       console.error(
         `Failed to send WebSocket message: ${error instanceof Error ? error.message : String(error)}`,
@@ -68,38 +240,104 @@ const useSearchWebSocket = () => {
     }
   };
 
-  const cancelSearchQueryBasedOnRequestId = (
-    requestId: string,
-    trace_id: string,
-  ) => {
-    const socket = webSocket.getWebSocketBasedOnSocketId(requestId);
+  const cancelSearchQueryBasedOnRequestId = (payload: {
+    trace_id: string;
+    org_id: string;
+  }) => {
+    const socket = webSocket.getWebSocketBasedOnSocketId(
+      socketId.value as string,
+    );
     // check state of socket
     if (socket && socket.readyState === WebSocket.OPEN) {
       webSocket.sendMessage(
-        requestId,
+        socketId.value as string,
         JSON.stringify({
           type: "cancel",
-          content: {
-            trace_id: trace_id,
-          },
+          content: payload,
         }),
       );
     }
   };
 
-  const closeSocketBasedOnRequestId = (requestId: string) => {
+  const closeSocketBasedOnRequestId = (traceId: string) => {
     try {
-      webSocket.closeSocket(requestId);
+      cleanUpListeners(traceId);
     } catch (error: any) {
-      console.error(`Failed to close socket ${requestId}:`, error);
+      console.error(`Failed to clean search trace ${traceId}:`, error);
     }
   };
 
+  const cleanUpListeners = (traceId: string) => {
+    // Remove all event listeners
+    if (traces[traceId]) traces[traceId].close.length = 0;
+    if (traces[traceId]) traces[traceId].error.length = 0;
+    if (traces[traceId]) traces[traceId].message.length = 0;
+    if (traces[traceId]) traces[traceId].reset.length = 0;
+    if (traces[traceId]) traces[traceId].open.length = 0;
+
+    delete traces[traceId];
+  };
+
+  const closeSocketWithError = () => {
+      webSocket.sendMessage(
+        socketId.value as string,
+        JSON.stringify({
+          type: "test_abnormal_close",
+          content: {
+            req_id: socketId.value,
+          },
+        }),
+      );
+  };
+
+  const closeSocket = () => {
+    webSocket.cleanupSocket(socketId.value as string);
+  }
+
+  const retryActiveTrace = (traceId: string, response: any) => {
+    traces[traceId]?.close.forEach((handler: any) => handler(response));
+    traces[traceId]?.reset.forEach((handler: any) => handler(traces[traceId].data));
+    cleanUpListeners(traceId);   
+  }
+
+  const resetAuthToken = async () => {
+    isInDrainMode.value = true;
+
+    return new Promise(async (resolve, reject) => {
+      authService.refresh_token().then((res: any) => {
+        isInDrainMode.value = false;
+        // Retry the request
+        Object.keys(traces).forEach((traceId) => {
+          if(!traces[traceId].isInitiated) {
+            initiateSocketConnection(traces[traceId].data, {
+              open: traces[traceId].open[0],
+              message: traces[traceId].message[0],
+              close: traces[traceId].close[0],
+              error: traces[traceId].error[0],
+              reset: traces[traceId].reset[0],
+            });
+          }else{
+            retryActiveTrace(traceId, { code: 1000 });
+          } 
+        });
+        resolve(res);
+      }).catch((err: any) => {
+        console.error("Error in refreshing auth token", err);
+        reject(err);
+      }).finally(() => {
+        isInDrainMode.value = false;
+      });
+    })
+  }
+
   return {
     fetchQueryDataWithWebSocket,
     sendSearchMessageBasedOnRequestId,
     cancelSearchQueryBasedOnRequestId,
     closeSocketBasedOnRequestId,
+    cleanUpListeners,
+    closeSocketWithError,
+    closeSocket
   };
 };
 

web/src/composables/useWebSocket.ts

@@ -4,7 +4,7 @@ import { onBeforeUnmount } from "vue";
 
 type MessageHandler = (event: MessageEvent) => void;
 type OpenHandler = (event: Event) => void;
-type CloseHandler = (event: CloseEvent) => void;
+type CloseHandler = (event: CloseEvent, socketId: string) => void;
 type ErrorHandler = (event: Event) => void;
 
 type WebSocketHandler =
@@ -71,7 +71,7 @@ const onClose = (socketId: string, event: CloseEvent) => {
   delete pingIntervals[socketId];
   delete sockets[socketId];
 
-  closeHandlers[socketId]?.forEach((handler) => handler(event));
+  closeHandlers[socketId]?.forEach((handler) => handler(event, socketId));
 };
 
 const onError = (socketId: string, event: Event) => {
@@ -132,14 +132,17 @@ const closeSocket = (socketId: string) => {
 const useWebSocket = () => {
   onBeforeUnmount(() => {
     for (const socketId in sockets) {
-      const socket = sockets[socketId];
-      if (socket) {
-        cleanupSocket(socketId, socket);
-      }
+      cleanupSocket(socketId);
     }
   });
 
-  const cleanupSocket = (socketId: string, socket: WebSocket) => {
+  const cleanupSocket = (socketId: string) => {
+    const socket = sockets[socketId];
+
+    if(!socket) {
+      console.error("Cleanup socket failed, socket not found", socketId);
+      return;
+    }
     // Remove all event listeners
     socket.onopen = null;
     socket.onmessage = null;
@@ -169,6 +172,7 @@ const useWebSocket = () => {
     connect,
     sendMessage,
     closeSocket,
+    cleanupSocket,
     getWebSocketBasedOnSocketId,
     addMessageHandler: (socketId: string, handler: MessageHandler) =>
       addHandler(messageHandlers, socketId, handler),

web/src/layouts/MainLayout.vue

@@ -514,6 +514,7 @@ import ManagementIcon from "@/components/icons/ManagementIcon.vue";
 import organizations from "@/services/organizations";
 import useStreams from "@/composables/useStreams";
 import { openobserveRum } from "@openobserve/browser-rum";
+import useSearchWebSocket from "@/composables/useSearchWebSocket";
 
 let mainLayoutMixin: any = null;
 if (config.isCloud == "true") {
@@ -564,6 +565,8 @@ export default defineComponent({
       window.open(zoBackendUrl + "/swagger/index.html", "_blank");
     },
     signout() {
+      this.closeSocket();
+
       if (config.isEnterprise == "true") {
         invlidateLoginData();
       }
@@ -596,6 +599,7 @@ export default defineComponent({
     const zoBackendUrl = store.state.API_ENDPOINT;
     const isLoading = ref(false);
     const { getStreams, resetStreams } = useStreams();
+    const { closeSocket } = useSearchWebSocket();
 
     const isMonacoEditorLoaded = ref(false);
 
@@ -1237,6 +1241,7 @@ export default defineComponent({
       slackIcon: markRaw(SlackIcon),
       openSlack,
       outlinedSettings,
+      closeSocket,
     };
   },
   computed: {

web/src/plugins/logs/Index.vue

@@ -18,7 +18,11 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 <!-- eslint-disable vue/v-on-event-hyphenation -->
 <template>
   <q-page class="logPage q-my-xs" id="logPage">
-    <div v-show="!showSearchHistory && !showSearchScheduler" id="secondLevel" class="full-height">
+    <div
+      v-show="!showSearchHistory && !showSearchScheduler"
+      id="secondLevel"
+      class="full-height"
+    >
       <q-splitter
         class="logs-horizontal-splitter full-height"
         v-model="splitterModel"
@@ -302,7 +306,10 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
         @closeSearchHistory="closeSearchHistoryfn"
         :isClicked="showSearchHistory"
       />
-      <div v-else-if="showSearchHistory && !store.state.zoConfig.usage_enabled " style="height: 200px">
+      <div
+        v-else-if="showSearchHistory && !store.state.zoConfig.usage_enabled"
+        style="height: 200px"
+      >
         <div style="height: 80vh" class="text-center q-pa-md flex flex-center">
           <div>
             <div>
@@ -344,16 +351,14 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
           </div>
         </div>
       </div>
-     
     </div>
     <div v-show="showSearchScheduler">
-        <SearchSchedulersList
+      <SearchSchedulersList
         ref="searchSchedulerRef"
         @closeSearchHistory="closeSearchSchedulerFn"
         :isClicked="showSearchScheduler"
-         />
-      </div>
-
+      />
+    </div>
   </q-page>
 </template>
 
@@ -371,6 +376,7 @@ import {
   defineAsyncComponent,
   provide,
   onMounted,
+  onBeforeUnmount,
 } from "vue";
 import { useQuasar } from "quasar";
 import { useStore } from "vuex";
@@ -460,12 +466,10 @@ export default defineComponent({
         if (this.searchObj.meta.jobId == "") {
           await this.getQueryData(false);
           this.refreshHistogramChart();
-        }
-        else{
+        } else {
           await this.getJobData(false);
         }
 
-
         if (config.isCloud == "true") {
           segment.track("Button Click", {
             button: "Get More Data",
@@ -487,12 +491,10 @@ export default defineComponent({
         // this.searchObj.data.resultGrid.currentPage =
         //   this.searchObj.data.resultGrid.currentPage + 1;
         this.searchObj.loading = true;
-        if(this.searchObj.meta.jobId == ""){
+        if (this.searchObj.meta.jobId == "") {
           await this.getQueryData(true);
           this.refreshHistogramChart();
-
-        }
-        else{
+        } else {
           await this.getJobData(false);
         }
 
@@ -581,6 +583,7 @@ export default defineComponent({
       buildWebSocketPayload,
       initializeWebSocketConnection,
       addRequestId,
+      sendCancelSearchMessage,
     } = useLogs();
     const searchResultRef = ref(null);
     const searchBarRef = ref(null);
@@ -660,15 +663,19 @@ export default defineComponent({
         router.currentRoute.value.query.hasOwnProperty("action") &&
         router.currentRoute.value.query.action == "search_scheduler"
       ) {
-        if(config.isEnterprise == 'true'){
+        if (config.isEnterprise == "true") {
           showSearchScheduler.value = true;
-        }
-        else{
+        } else {
           router.back();
         }
       }
     });
 
+    onBeforeUnmount(() => {
+      // Cancel all the search queries
+      cancelOnGoingSearchQueries();
+    });
+
     onActivated(() => {
       if (isLogsMounted.value) handleActivation();
     });
@@ -711,12 +718,11 @@ export default defineComponent({
           router.currentRoute.value.query.hasOwnProperty("action") &&
           router.currentRoute.value.query.action == "search_scheduler"
         ) {
-          if(config.isEnterprise == 'true'){
+          if (config.isEnterprise == "true") {
             showSearchScheduler.value = true;
+          } else {
+            router.back();
           }
-          else{
-          router.back();
-        }
         }
       },
       // (action) => {
@@ -799,7 +805,6 @@ export default defineComponent({
         await getQueryData();
         refreshHistogramChart();
         showJobScheduler.value = true;
-
       } catch (e) {
         console.log(e);
       }
@@ -1544,6 +1549,12 @@ export default defineComponent({
 
     // [END] cancel running queries
 
+    const cancelOnGoingSearchQueries = () => {
+      sendCancelSearchMessage(
+        searchObj.data.searchWebSocketRequestIdsAndTraceIds,
+      );
+    };
+
     return {
       t,
       store,

web/src/plugins/logs/SearchBar.vue

@@ -37,10 +37,10 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
                 size="sm"
                 icon="search"
                 class="button button-right tw-flex tw-justify-center tw-items-center no-border no-outline !tw-rounded-r-none q-px-sm"
-                style="height: 32px;"
+                style="height: 32px"
               >
                 <q-tooltip>
-                  {{ t('common.search') }}
+                  {{ t("common.search") }}
                 </q-tooltip>
               </q-btn>
             </div>
@@ -52,51 +52,67 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
                     ? 'selected'
                     : ''
                 "
-                class="button button-right tw-flex tw-justify-center tw-items-center  no-border no-outline !tw-rounded-l-none q-px-sm"
+                class="button button-right tw-flex tw-justify-center tw-items-center no-border no-outline !tw-rounded-l-none q-px-sm"
                 @click="onLogsVisualizeToggleUpdate('visualize')"
                 :disabled="isVisualizeToggleDisabled"
                 no-caps
                 size="sm"
-                style="height: 32px;"
+                style="height: 32px"
               >
-              <q-tooltip>
-                {{ 
-                  isVisualizeToggleDisabled
-                    ? t('search.visualizeDisabledForMultiStream')
-                    : t('search.visualize')
-                }}
-              </q-tooltip>
-                <img :src="visualizeIcon" alt="Visualize" style="width: 20px; height: 20px; ">
+                <q-tooltip>
+                  {{
+                    isVisualizeToggleDisabled
+                      ? t("search.visualizeDisabledForMultiStream")
+                      : t("search.visualize")
+                  }}
+                </q-tooltip>
+                <img
+                  :src="visualizeIcon"
+                  alt="Visualize"
+                  style="width: 20px; height: 20px"
+                />
               </q-btn>
             </div>
           </div>
         </div>
 
-        <div style="border: 1px solid #c4c4c4; border-radius: 5px ;" class="q-pr-xs q-ml-xs">
-          <q-toggle
-          data-test="logs-search-bar-show-histogram-toggle-btn"
-          v-model="searchObj.meta.showHistogram"
+        <div
+          style="border: 1px solid #c4c4c4; border-radius: 5px"
+          class="q-pr-xs q-ml-xs"
         >
-
-        <img :src="histogramIcon" alt="Histogram" style="width: 20px; height: 20px; ">
-        <q-tooltip> 
-          {{ t('search.showHistogramLabel') }}
-        </q-tooltip>
-        </q-toggle>
-        </div>
-        <div style="border: 1px solid #c4c4c4; border-radius: 5px ;" class="q-pr-xs q-ml-xs">
-
-        <q-toggle
-          data-test="logs-search-bar-sql-mode-toggle-btn"
-          v-model="searchObj.meta.sqlMode"
+          <q-toggle
+            data-test="logs-search-bar-show-histogram-toggle-btn"
+            v-model="searchObj.meta.showHistogram"
           >
-          <img :src="sqlIcon" alt="SQL Mode" style="width: 20px; height: 20px;">
-          <q-tooltip>
-            {{ t('search.sqlModeLabel') }}
-          </q-tooltip>
-        </q-toggle>
+            <img
+              :src="histogramIcon"
+              alt="Histogram"
+              style="width: 20px; height: 20px"
+            />
+            <q-tooltip>
+              {{ t("search.showHistogramLabel") }}
+            </q-tooltip>
+          </q-toggle>
         </div>
-        
+        <div
+          style="border: 1px solid #c4c4c4; border-radius: 5px"
+          class="q-pr-xs q-ml-xs"
+        >
+          <q-toggle
+            data-test="logs-search-bar-sql-mode-toggle-btn"
+            v-model="searchObj.meta.sqlMode"
+          >
+            <img
+              :src="sqlIcon"
+              alt="SQL Mode"
+              style="width: 20px; height: 20px"
+            />
+            <q-tooltip>
+              {{ t("search.sqlModeLabel") }}
+            </q-tooltip>
+          </q-toggle>
+        </div>
+
         <q-btn
           data-test="logs-search-bar-reset-filters-btn"
           no-caps
@@ -106,14 +122,14 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
           @click="resetFilters"
         >
           <q-tooltip>
-            {{ t('search.resetFilters') }}
+            {{ t("search.resetFilters") }}
           </q-tooltip>
         </q-btn>
         <syntax-guide
           data-test="logs-search-bar-sql-mode-toggle-btn"
           :sqlmode="searchObj.meta.sqlMode"
         >
-      </syntax-guide>
+        </syntax-guide>
         <q-btn-group class="q-ml-xs no-outline q-pa-none no-border">
           <q-btn-dropdown
             data-test="logs-search-saved-views-btn"
@@ -347,26 +363,31 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
                 </q-item-section>
               </q-item>
             </q-list>
-
           </q-btn-dropdown>
           <q-tooltip>
-              {{ t('search.savedViewsLabel') }}
-            </q-tooltip>
-        </q-btn-group>
-        <div style="border: 1px solid #c4c4c4; border-radius: 5px ;" class="q-pr-xs q-ml-xs">
-          <q-toggle
-          data-test="logs-search-bar-quick-mode-toggle-btn"
-          v-model="searchObj.meta.quickMode"
-          @click="handleQuickMode"
-          >
-          <img :src="quickModeIcon" alt="Quick Mode" style="width: 20px; height: 20px;">
-          <q-tooltip>
-            {{ t('search.quickModeLabel') }}
+            {{ t("search.savedViewsLabel") }}
           </q-tooltip>
-        </q-toggle>
+        </q-btn-group>
+        <div
+          style="border: 1px solid #c4c4c4; border-radius: 5px"
+          class="q-pr-xs q-ml-xs"
+        >
+          <q-toggle
+            data-test="logs-search-bar-quick-mode-toggle-btn"
+            v-model="searchObj.meta.quickMode"
+            @click="handleQuickMode"
+          >
+            <img
+              :src="quickModeIcon"
+              alt="Quick Mode"
+              style="width: 20px; height: 20px"
+            />
+            <q-tooltip>
+              {{ t("search.quickModeLabel") }}
+            </q-tooltip>
+          </q-toggle>
         </div>
       </div>
-      
 
       <div class="float-right col-auto q-mb-xs">
         <q-toggle
@@ -377,7 +398,7 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
           size="32px"
         >
           <q-tooltip>
-            {{ t('search.messageWrapContent') }}
+            {{ t("search.messageWrapContent") }}
           </q-tooltip>
         </q-toggle>
 
@@ -402,32 +423,39 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
           :loading="shareLink.isLoading.value"
         >
           <q-tooltip>
-            {{ t('search.shareLink') }}
+            {{ t("search.shareLink") }}
           </q-tooltip>
         </q-btn>
 
         <q-btn
-              data-test="logs-search-bar-more-options-btn"
-              class="q-mr-xs download-logs-btn q-px-sm"
-              size="sm"
-              icon="menu"
-            >
-            <q-menu>
- 
-              <q-list>
-                <q-item data-test="search-history-item-btn" class="q-pa-sm saved-view-item" clickable v-close-popup>
-                <q-item-section
-                  @click.stop="showSearchHistoryfn"
-                  >
+          data-test="logs-search-bar-more-options-btn"
+          class="q-mr-xs download-logs-btn q-px-sm"
+          size="sm"
+          icon="menu"
+        >
+          <q-menu>
+            <q-list>
+              <q-item
+                data-test="search-history-item-btn"
+                class="q-pa-sm saved-view-item"
+                clickable
+                v-close-popup
+              >
+                <q-item-section @click.stop="showSearchHistoryfn">
                   <q-item-label class="tw-flex tw-items-center tw-gap-2">
-                    <img :src="searchHistoryIcon" alt="Search History"  style="width: 20px; height: 20px;" />
+                    <img
+                      :src="searchHistoryIcon"
+                      alt="Search History"
+                      style="width: 20px; height: 20px"
+                    />
 
-                    Search History</q-item-label>
+                    Search History</q-item-label
+                  >
                 </q-item-section>
               </q-item>
               <q-separator />
               <q-item
-              style="min-width: 150px;"
+                style="min-width: 150px"
                 class="q-pa-sm saved-view-item"
                 clickable
                 v-close-popup
@@ -442,24 +470,39 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
                   v-close-popup
                 >
                   <q-item-label class="tw-flex tw-items-center tw-gap-2">
-                   <img :src="downloadTableIcon" alt="Download Table" style="width: 20px; height: 20px;" />
-                    {{ t("search.downloadTable") }}</q-item-label>
+                    <img
+                      :src="downloadTableIcon"
+                      alt="Download Table"
+                      style="width: 20px; height: 20px"
+                    />
+                    {{ t("search.downloadTable") }}</q-item-label
+                  >
                 </q-item-section>
               </q-item>
-              <q-item class="q-pa-sm saved-view-item" style="min-width: 150px;" clickable v-close-popup>
+              <q-item
+                class="q-pa-sm saved-view-item"
+                style="min-width: 150px"
+                clickable
+                v-close-popup
+              >
                 <q-item-section
                   @click.stop="toggleCustomDownloadDialog"
                   v-close-popup
                 >
                   <q-item-label class="tw-flex tw-items-center tw-gap-2">
-                    <img :src="customRangeIcon" alt="Custom Range" style="width: 20px; height: 20px;" />
+                    <img
+                      :src="customRangeIcon"
+                      alt="Custom Range"
+                      style="width: 20px; height: 20px"
+                    />
 
-                    {{ t("search.customRange") }}</q-item-label>
+                    {{ t("search.customRange") }}</q-item-label
+                  >
                 </q-item-section>
               </q-item>
               <q-separator />
               <q-item
-               v-if="config.isEnterprise == 'true'"
+                v-if="config.isEnterprise == 'true'"
                 data-test="search-scheduler-create-new-btn"
                 class="q-pa-sm saved-view-item"
                 clickable
@@ -467,8 +510,15 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
                 @click="createScheduleJob"
               >
                 <q-item-section v-close-popup>
-                  <q-item-label class="tw-flex tw-items-center tw-gap-2" data-test="search-scheduler-create-new-label">
-                    <img :src="createScheduledSearchIcon" alt="Create Scheduled Search" style="width: 20px; height: 20px;" />
+                  <q-item-label
+                    class="tw-flex tw-items-center tw-gap-2"
+                    data-test="search-scheduler-create-new-label"
+                  >
+                    <img
+                      :src="createScheduledSearchIcon"
+                      alt="Create Scheduled Search"
+                      style="width: 20px; height: 20px"
+                    />
                     Create Scheduled Search</q-item-label
                   >
                 </q-item-section>
@@ -482,20 +532,25 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
                 @click="routeToSearchSchedule"
               >
                 <q-item-section v-close-popup>
-                  <q-item-label class="tw-flex tw-items-center tw-gap-2" data-test="search-scheduler-list-label">
-                    <img :src="listScheduledSearchIcon" alt="List Scheduled Search" style=" width: 20px; height: 20px;" />
+                  <q-item-label
+                    class="tw-flex tw-items-center tw-gap-2"
+                    data-test="search-scheduler-list-label"
+                  >
+                    <img
+                      :src="listScheduledSearchIcon"
+                      alt="List Scheduled Search"
+                      style="width: 20px; height: 20px"
+                    />
 
                     List Scheduled Search</q-item-label
                   >
                 </q-item-section>
               </q-item>
-
             </q-list>
-            </q-menu>
-            <q-tooltip style="width: 80px;" >
-              {{ t('search.moreActions') }}
-            </q-tooltip>
-
+          </q-menu>
+          <q-tooltip style="width: 80px">
+            {{ t("search.moreActions") }}
+          </q-tooltip>
         </q-btn>
         <div class="float-left">
           <date-time
@@ -1167,9 +1222,9 @@ import { computed } from "vue";
 import { useLoading } from "@/composables/useLoading";
 import TransformSelector from "./TransformSelector.vue";
 import FunctionSelector from "./FunctionSelector.vue";
+import useSearchWebSocket from "@/composables/useSearchWebSocket";
 import histogram_svg from "../../assets/images/common/histogram_image.svg";
 
-
 const defaultValue: any = () => {
   return {
     name: "",
@@ -1345,6 +1400,8 @@ export default defineComponent({
     const formData: any = ref(defaultValue());
     const functionOptions = ref(searchObj.data.transforms);
 
+    const { closeSocketWithError } = useSearchWebSocket();
+
     const transformsExpandState = ref({
       actions: false,
       functions: false,
@@ -3241,33 +3298,50 @@ export default defineComponent({
       }
     };
     const visualizeIcon = computed(() => {
-      return searchObj.meta.logsVisualizeToggle === 'visualize' ? getImageURL('images/common/visualize_icon_light.svg') : getImageURL('images/common/visualize_icon_dark.svg')
+      return searchObj.meta.logsVisualizeToggle === "visualize"
+        ? getImageURL("images/common/visualize_icon_light.svg")
+        : getImageURL("images/common/visualize_icon_dark.svg");
     });
     const histogramIcon = computed(() => {
-      return store.state.theme === 'dark' ? getImageURL('images/common/bar_chart_histogram_light.svg') : getImageURL('images/common/bar_chart_histogram.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/bar_chart_histogram_light.svg")
+        : getImageURL("images/common/bar_chart_histogram.svg");
     });
     const sqlIcon = computed(() => {
-      return store.state.theme === 'dark' ? getImageURL('images/common/hugeicons_sql_light.svg') : getImageURL('images/common/hugeicons_sql.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/hugeicons_sql_light.svg")
+        : getImageURL("images/common/hugeicons_sql.svg");
     });
     const quickModeIcon = computed(() => {
-      return store.state.theme === 'dark' ? getImageURL('images/common/quick_mode_light.svg') : getImageURL('images/common/quick_mode.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/quick_mode_light.svg")
+        : getImageURL("images/common/quick_mode.svg");
     });
     const searchHistoryIcon = computed(() => {
-      return store.state.theme === 'dark' ? getImageURL('images/common/search_history_light.svg') : getImageURL('images/common/search_history.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/search_history_light.svg")
+        : getImageURL("images/common/search_history.svg");
     });
     const downloadTableIcon = computed(() => {
-      return store.state.theme === 'dark' ? getImageURL('images/common/download_table_light.svg') : getImageURL('images/common/download_table.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/download_table_light.svg")
+        : getImageURL("images/common/download_table.svg");
     });
     const customRangeIcon = computed(() => {
-      return store.state.theme === 'dark' ? getImageURL('images/common/custom_range_light.svg') : getImageURL('images/common/custom_range.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/custom_range_light.svg")
+        : getImageURL("images/common/custom_range.svg");
     });
     const createScheduledSearchIcon = computed(() => {
-      return store.state.theme === 'dark'   ? getImageURL('images/common/create_scheduled_search_light.svg') : getImageURL('images/common/create_scheduled_search.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/create_scheduled_search_light.svg")
+        : getImageURL("images/common/create_scheduled_search.svg");
     });
     const listScheduledSearchIcon = computed(() => {
-      return store.state.theme === 'dark' ? getImageURL('images/common/list_scheduled_search_light.svg') : getImageURL('images/common/list_scheduled_search.svg')
+      return store.state.theme === "dark"
+        ? getImageURL("images/common/list_scheduled_search_light.svg")
+        : getImageURL("images/common/list_scheduled_search.svg");
     });
-    
 
     // [END] cancel running queries
 
@@ -3385,6 +3459,7 @@ export default defineComponent({
       actionEditorQuery,
       isActionsEnabled,
       showFunctionEditor,
+      closeSocketWithError,
       histogram_svg,
       visualizeIcon,
       histogramIcon,
@@ -3394,7 +3469,7 @@ export default defineComponent({
       downloadTableIcon,
       customRangeIcon,
       createScheduledSearchIcon,
-      listScheduledSearchIcon
+      listScheduledSearchIcon,
     };
   },
   computed: {

web/src/services/auth.ts

@@ -25,6 +25,10 @@ const auth = {
     const res = await http().get("/config/dex_login");
     return res.data;
   },
+  refresh_token: async () => {
+    const res = await http().get("/config/dex_refresh");
+    return res.data;
+  }
 };
 
 export default auth;

web/src/utils/zincutils.ts

@@ -979,7 +979,7 @@ export const deepCopy = (value: any) => {
 
 export const getWebSocketUrl = (request_id: string, org_identifier: string) => {
   const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
-  return `${protocol}//${store.state.API_ENDPOINT.split("//")[1]}/api/${org_identifier}/ws/${request_id}`;
+  return `${protocol}//${store.state.API_ENDPOINT.split("//")[1]}/api/${org_identifier}/ws/v2/${request_id}`;
 };
 
 export const isWebSocketEnabled = () => {