fix(ai): route o2-ai requests through NATS in HA mode (#15036)

src/api/management/src/request/ai/chat.rs

@@ -264,7 +264,7 @@ pub async fn chat(Path(org_id): Path<String>, 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");
         }
 

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")

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.

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;
     }

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<Option<(Instant, Vec<Replica>)>> = RwLock::new(None);
 
+static DIRECTORY: LazyLock<nats::NatsDb> = 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<Vec<Replica>> {
-    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<Vec<Replica>> {
     )
 }
 
+async fn list_replica_values() -> Result<Vec<bytes::Bytes>> {
+    directory().await?.list_values(&replicas_prefix()).await
+}
+
+async fn get_session_claim(session_id: &str) -> Result<Option<bytes::Bytes>> {
+    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<String> {
     // 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<String> {
     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<String> {
+    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!(

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,