From 65ba4aff9bdd198dd46cae168102ea375fcaff79 Mon Sep 17 00:00:00 2001 From: Harsh Mahajan <115500013+007harshmahajan@users.noreply.github.com> Date: Thu, 1 Oct 2026 12:51:55 +0000 Subject: [PATCH] fix(ai): route o2-ai requests through NATS in HA mode (#15036) --- src/api/management/src/request/ai/chat.rs | 2 +- .../src/request/alerts/incidents.rs | 2 +- src/core/src/alerts/incidents.rs | 4 +- src/core/src/alerts/scheduler/handlers.rs | 2 +- src/infra/src/cluster/ai_sessions.rs | 69 +++++++++++++------ src/infra/src/db/nats.rs | 5 ++ 6 files changed, 58 insertions(+), 26 deletions(-) diff --git a/src/api/management/src/request/ai/chat.rs b/src/api/management/src/request/ai/chat.rs index a0bc7c9f5a..cc5c8e5467 100644 --- a/src/api/management/src/request/ai/chat.rs +++ b/src/api/management/src/request/ai/chat.rs @@ -264,7 +264,7 @@ pub async fn chat(Path(org_id): Path, in_req: axum::extract::Request) -> return MetaHttpResponse::bad_request("AI is not enabled"); } - if o2_cfg.ai.agent_url.is_empty() { + if !o2_cfg.ai.has_agent_target() { return MetaHttpResponse::bad_request("AI agent URL is not set"); } diff --git a/src/api/management/src/request/alerts/incidents.rs b/src/api/management/src/request/alerts/incidents.rs index bf665aa736..93eb0e4098 100644 --- a/src/api/management/src/request/alerts/incidents.rs +++ b/src/api/management/src/request/alerts/incidents.rs @@ -470,7 +470,7 @@ pub async fn trigger_incident_rca( return MetaHttpResponse::bad_request("RCA is not enabled"); } - if o2_cfg.ai.agent_url.is_empty() { + if !o2_cfg.ai.has_agent_target() { return axum::response::Response::builder() .status(axum::http::StatusCode::SERVICE_UNAVAILABLE) .header(axum::http::header::CONTENT_TYPE, "application/json") diff --git a/src/core/src/alerts/incidents.rs b/src/core/src/alerts/incidents.rs index 0de5bb31b8..113a007d73 100644 --- a/src/core/src/alerts/incidents.rs +++ b/src/core/src/alerts/incidents.rs @@ -2381,7 +2381,7 @@ pub async fn rca_will_run(org_id: &str, incident_id: &str) -> bool { let cfg = get_config(); let will_run = cfg.incidents.enabled && cfg.incidents.rca_enabled - && !cfg.ai.agent_url.is_empty() + && cfg.ai.has_agent_target() && !rca_service::l0_off_for_incident(org_id, incident_id).await; if !will_run { rca_service::skip_analysis_for_incident(org_id, incident_id).await; @@ -2431,7 +2431,7 @@ pub async fn trigger_rca_for_incident( return Ok(()); // Not an error - just not configured } - if config.ai.agent_url.is_empty() { + if !config.ai.has_agent_target() { log::debug!("[INCIDENTS::RCA] RCA agent URL not set, skipping immediate trigger"); // §6: no verdict is coming, so nothing may go on holding a page for // one. Every guard below reaches this same state. diff --git a/src/core/src/alerts/scheduler/handlers.rs b/src/core/src/alerts/scheduler/handlers.rs index d2af80acc2..519c0701a1 100644 --- a/src/core/src/alerts/scheduler/handlers.rs +++ b/src/core/src/alerts/scheduler/handlers.rs @@ -1821,7 +1821,7 @@ async fn trigger_rca_for_alert_firing( let cfg = get_o2_config(); // Must agree with the condition `analysis_at_start` used to mark these records `Pending`. - if !cfg.incidents.rca_enabled || !cfg.ai.enabled || cfg.ai.agent_url.is_empty() { + if !cfg.incidents.rca_enabled || !cfg.ai.enabled || !cfg.ai.has_agent_target() { release_alert_firing_hold(&alert.org_id, &opened).await; return; } diff --git a/src/infra/src/cluster/ai_sessions.rs b/src/infra/src/cluster/ai_sessions.rs index adfdb4abd5..fd34efe62d 100644 --- a/src/infra/src/cluster/ai_sessions.rs +++ b/src/infra/src/cluster/ai_sessions.rs @@ -15,20 +15,25 @@ //! Read side of the o2-ai session->owner directory. o2-ai keeps each //! conversation node-local, claims it and heartbeats its liveness; both are read -//! here via [`crate::db::get_coordinator`] (see o2-ai's `cluster/directory.py`). +//! here straight from NATS (see o2-ai's `cluster/directory.py`). //! -//! A routing *hint* only: every failure degrades to "unknown owner" and the -//! caller falls back to the configured agent URL. +//! Every failure degrades to "unknown owner"; the caller decides what that means. use std::{ - sync::{Once, RwLock}, + sync::{ + LazyLock, Once, RwLock, + atomic::{AtomicUsize, Ordering}, + }, time::{Duration, Instant}, }; use config::utils::hash::Sum64; use serde::Deserialize; -use crate::db::get_coordinator; +use crate::{ + db::{Db, nats}, + errors::Result, +}; /// Bucket holding session->owner claims. Claims outlive a replica restart so a /// session stays pinned to its owner. @@ -83,6 +88,16 @@ const EMPTY_REPLICAS_TTL: Duration = Duration::from_secs(30); /// serialising every caller. The critical section never awaits. static LIVE_REPLICAS_CACHE: RwLock)>> = RwLock::new(None); +static DIRECTORY: LazyLock = LazyLock::new(nats::NatsDb::default); + +static NEXT_REPLICA: AtomicUsize = AtomicUsize::new(0); + +/// NATS directly: the coordinator is SQLite in local mode, which o2-ai never writes to. +async fn directory() -> Result<&'static nats::NatsDb> { + nats::try_get_nats_client().await?; + Ok(&DIRECTORY) +} + /// The o2-ai replicas currently heartbeating, or `None` if the registry could not /// be read. Not the same answer: an empty registry proves a claimed owner is /// gone, a failed read proves nothing. @@ -116,11 +131,7 @@ fn ttl_for(replicas: &[Replica]) -> Duration { /// `None` when the registry could not be read; `Some(vec![])` when it was read /// and nobody is registered. async fn fetch_live_replicas() -> Option> { - let values = match get_coordinator() - .await - .list_values(&replicas_prefix()) - .await - { + let values = match list_replica_values().await { Ok(v) => v, Err(e) => { log::debug!("[AI_SESSIONS] cannot list replicas: {e}"); @@ -138,12 +149,22 @@ async fn fetch_live_replicas() -> Option> { ) } +async fn list_replica_values() -> Result> { + directory().await?.list_values(&replicas_prefix()).await +} + +async fn get_session_claim(session_id: &str) -> Result> { + directory() + .await? + .get_if_exists(&session_key(session_id)) + .await +} + /// Pick a live o2-ai replica to host a NEW session — never an existing one, /// which would land on a replica that has never seen it. Hashing the sorted live /// set makes every openobserve node place a session identically, uncoordinated. /// -/// `None` when no replica is heartbeating; the caller falls back to the -/// configured agent URL. +/// `None` when no replica is heartbeating. pub async fn pick_replica_for_new_session(session_id: &str) -> Option { // Unreadable and empty are the same answer here: nowhere to place it. let mut replicas = live_replicas().await.unwrap_or_default(); @@ -167,6 +188,18 @@ pub async fn pick_replica_for_new_session(session_id: &str) -> Option { Some(addr.clone()) } +/// Pick a live o2-ai replica, in rotation, for a request that belongs to no session. +pub async fn pick_any_replica() -> Option { + let mut replicas = live_replicas().await.unwrap_or_default(); + if replicas.is_empty() { + log::debug!("[AI_SESSIONS] no live o2-ai replicas registered"); + return None; + } + replicas.sort_by(|a, b| a.0.cmp(&b.0)); + let idx = NEXT_REPLICA.fetch_add(1, Ordering::Relaxed) % replicas.len(); + Some(replicas.swap_remove(idx).1) +} + /// Where a session should be routed. #[derive(Debug, Clone, PartialEq)] pub enum SessionRoute { @@ -197,11 +230,7 @@ pub async fn get_session_route(session_id: &str) -> SessionRoute { return SessionRoute::Unknown; } - let entry = match get_coordinator() - .await - .get_if_exists(&session_key(session_id)) - .await - { + let entry = match get_session_claim(session_id).await { Ok(Some(v)) => v, // The only "definitely no claim" answer, so the only placeable one. Ok(None) => return SessionRoute::Unclaimed, @@ -217,8 +246,7 @@ pub async fn get_session_route(session_id: &str) -> SessionRoute { // Claimed but not dialable (no O2_AI_ADVERTISE_URL). NOT `Unclaimed`: // the session belongs to someone we just can't reach. log::warn!( - "[AI_SESSIONS] session {session_id} is owned by {} but has no advertised address; \ - falling back to the configured agent URL", + "[AI_SESSIONS] session {session_id} is owned by {} but has no advertised address", v.owner ); return SessionRoute::Unknown; @@ -243,8 +271,7 @@ pub async fn get_session_route(session_id: &str) -> SessionRoute { }; if live.iter().any(|(name, _)| name == &owner.owner) { - // Routing has no config flag, so this once-per-process line is the only - // signal an operator gets that session affinity is in effect. + // Logged once per process so an operator can see affinity took effect. static ROUTING_ENGAGED: Once = Once::new(); ROUTING_ENGAGED.call_once(|| { log::info!( diff --git a/src/infra/src/db/nats.rs b/src/infra/src/db/nats.rs index 3ecb65889f..0d9129d18a 100644 --- a/src/infra/src/db/nats.rs +++ b/src/infra/src/db/nats.rs @@ -57,6 +57,11 @@ pub async fn get_nats_client() -> &'static Client { .unwrap_or_else(|e| panic!("{e}")) } +/// Like [`get_nats_client`], but reports a failed connect instead of panicking. +pub async fn try_get_nats_client() -> Result<&'static Client> { + NATS_CLIENT.get_or_try_init(connect).await +} + async fn get_bucket_by_key<'a>( prefix: &'a str, key: &'a str,