laburnum 1.17.1

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

mod laburnum_commands;
mod router;

pub use router::RequestRouter;
use {
  crate::{
    Ident,
    Partitions,
    TRACER,
    connect::ipc::Connection,
    protocol::{
      jsonrpc::{
        self,
        Id,
        Message,
        Response,
      },
      lsp::LanguageServer,
      otel::TraceContext,
    },
    scheduler::{
      Scheduler,
      lanes::RPC_LANE_3,
      task::{
        LaburnumTask,
        TaskContext,
      },
    },
    connect::lsp::ClientId,
  },
  async_channel::Receiver,
  std::{
    collections::HashSet,
    ops::ControlFlow,
    sync::{
      Arc,
      atomic::{
        AtomicBool,
        Ordering,
      },
    },
  },
};

struct ShutdownFlag(Arc<AtomicBool>);

impl ShutdownFlag {
  fn set(&self) {
    self.0.store(true, Ordering::SeqCst);
  }

  pub(crate) fn is_shutdown(&self) -> bool {
    self.0.load(Ordering::SeqCst)
  }
}

/// The main RPC message processing task for the Language Server Protocol.
///
/// # Purpose
///
/// `RpcTask` manages the lifecycle of LSP communication by:
/// - Receiving and dispatching incoming JSON-RPC messages from the client
/// - Managing server state transitions (uninitialized → initializing →
///   initialized → shutdown → exited)
/// - Spawning child tasks on appropriate lanes to handle requests
///   asynchronously
/// - Coordinating response delivery back to the client
///
/// # Why It Exists
///
/// LSP servers must continuously process incoming messages while remaining
/// responsive. The RpcTask provides a central event loop that:
/// 1. Polls the connection for incoming messages without blocking
/// 2. Enforces LSP protocol state machine requirements
/// 3. Delegates actual work to prioritized worker tasks
/// 4. Aggregates responses from child tasks and sends them to the client
///
/// # Scheduling
///
/// The RpcTask runs on RPC lanes (LANE08-LANE11), which use a priority queue
/// with aging semantics. Workers check RPC lanes periodically with rate
/// limiting. This prevents the RpcTask from starving user-facing work while
/// ensuring messages are processed regularly.
///
/// When the RpcTask receives an LSP request, it spawns a child task on
/// [`SYNC_LANE`] to handle the request, allowing the RpcTask to immediately
/// return to polling for new messages.
///
/// # Lifecycle
///
/// Created during [`Scheduler`] initialization and runs until receiving an LSP
/// `exit` notification. The task yields after each `tick()` to allow other work
/// to progress.
///
/// [`SYNC_LANE`]: crate::scheduler::lanes::DEFAULT_LANE
/// [`Scheduler`]: crate::scheduler::Scheduler
pub struct RpcTask<P: Partitions, T: LanguageServer<P>> {
  state:            ServerState,
  ctx:              TaskContext<P, T>,
  conn:             Connection,
  router:           RequestRouter<P, T>,
  pending_requests: Arc<parking_lot::Mutex<HashSet<jsonrpc::Id>>>,
  response_tx:      async_channel::Sender<Response>,
  response_rx:      Receiver<Response>,

  init_request_id: Option<Id>,

  shutdown_flag: ShutdownFlag,
}

impl<P: Partitions, T: LanguageServer<P>> RpcTask<P, T> {
  pub(crate) fn create(
    scheduler: Arc<Scheduler<P, T>>,
    conn: Connection,
    client_id: ClientId,
    server: Arc<T>,
    shutdown_flag: Arc<AtomicBool>,
    response_capacity: usize,
  ) -> Arc<LaburnumTask<P, T>> {
    let shutdown_flag = ShutdownFlag(shutdown_flag);

    LaburnumTask::new_with_parent(
      scheduler,
      move |ctx| {
        Box::pin(async move {
          let (response_tx, response_rx) =
            async_channel::bounded(response_capacity);

          let mut task = RpcTask {
            state: ServerState::new(),
            ctx,
            conn,
            router: RequestRouter::new(server),
            pending_requests: Arc::new(parking_lot::Mutex::new(HashSet::new())),
            response_tx,
            response_rx,
            init_request_id: None,
            shutdown_flag,
          };

          loop {
            if task.shutdown_flag.is_shutdown() {
              break;
            }

            // Wait for either a message from the connection, a response from
            // handlers, or a 100ms timeout to check shutdown flag
            // periodically
            enum Event {
              Message(Message),
              Response(Response),
              Timeout,
              Closed,
            }

            let event = futures_lite::future::or(
              futures_lite::future::or(
                async {
                  match task.conn.receiver.recv().await {
                    | Ok(msg) => Event::Message(msg),
                    | Err(_) => Event::Closed,
                  }
                },
                async {
                  match task.response_rx.recv().await {
                    | Ok(resp) => Event::Response(resp),
                    | Err(_) => Event::Closed,
                  }
                },
              ),
              async {
                smol::Timer::after(std::time::Duration::from_millis(100)).await;
                Event::Timeout
              },
            )
            .await;

            match event {
              | Event::Message(message) => {
                if task.handle_message(message).is_break() {
                  break;
                }
              },
              | Event::Response(response) => {
                task.handle_response(response);
              },
              | Event::Timeout => {},
              | Event::Closed => break,
            }
          }

          None
        })
      },
      RPC_LANE_3,
      None,
      client_id,
    )
  }
}

impl<P: Partitions, T: LanguageServer<P>> RpcTask<P, T> {
  fn is_initialize_request(method: &str) -> bool {
    method == "initialize"
  }

  fn is_shutdown_request(method: &str) -> bool {
    method == "shutdown"
  }

  fn is_initialized_notification(method: &str) -> bool {
    method == "initialized"
  }

  fn is_exit_notification(method: &str) -> bool {
    method == "exit"
  }

  fn is_execute_command_request(method: &str) -> bool {
    method == "workspace/executeCommand"
  }

  fn send_error_response(&self, id: jsonrpc::Id, error: jsonrpc::Error) {
    use opentelemetry::trace::SpanKind;
    otel::span!("rpc.send_error_response", kind = SpanKind::Producer);

    let response = Response::from_error(id.clone(), error);
    if let Err(e) = self.conn.sender.try_send(Message::Response(response)) {
      otel::error!(
        "error_response_send_failed",
        format!("Failed to send error response for request {}: {:?}", id, e)
      );
    }
  }

  fn handle_message(&mut self, message: Message) -> ControlFlow<()> {
    match message {
      | Message::Request(request) => {
        self.handle_request(request);
        ControlFlow::Continue(())
      },
      | Message::Notification(notification) => {
        self.handle_notification(notification)
      },
      | Message::Response(_) => ControlFlow::Continue(()),
    }
  }

  fn handle_request(&mut self, request: jsonrpc::Request) {
    let trace_ctx = request.trace_context();
    let _guard = trace_ctx.attach();

    use opentelemetry::trace::SpanKind;
    otel::span!("rpc.receive_request", kind = SpanKind::Consumer);

    // Extract client_id from request before into_parts() consumes it
    let msg_client_id = request.client_id();

    // Validate client_id if present - must match this connection's assigned ID
    if let Some(id) = msg_client_id
      && id != self.ctx.client_id()
    {
      let request_id = request.id().clone();
      self.send_error_response(request_id, jsonrpc::Error {
        code:    jsonrpc::ErrorCode::InvalidRequest,
        message: "client_id mismatch".into(),
        data:    None,
      });
      return;
    }

    // Use the validated client_id (or fallback to connection's ID if not provided)
    let task_client_id = msg_client_id.unwrap_or_else(|| self.ctx.client_id());

    let (method, request_id, params) = request.into_parts();
    eprintln!("Received request: method={}, id={}", method, request_id);
    let current_state = self.state.get();

    if Self::is_initialize_request(&method) {
      if current_state != State::Uninitialized {
        self.send_error_response(request_id, jsonrpc::Error {
          code:    jsonrpc::ErrorCode::InvalidRequest,
          message: "Server already initialized".into(),
          data:    None,
        });
        return;
      }

      self.state.set(State::Initializing);
      self.init_request_id = Some(request_id.clone());
    }

    if Self::is_shutdown_request(&method) {
      self.handle_shutdown(request_id);
      return;
    }

    if current_state != State::Initialized
      && !Self::is_initialize_request(&method)
      && !Self::is_execute_command_request(&method)
    {
      self.send_error_response(
        request_id,
        crate::protocol::jsonrpc::error::not_initialized_error(),
      );
      return;
    }

    self.pending_requests.lock().insert(request_id.clone());

    #[cfg(feature = "testing-commands")]
    if let Some(params_value) = &params
      && self.try_handle_laburnum_command(
        &method,
        params_value,
        request_id.clone(),
      )
    {
      return;
    }

    let (handler, method_lane) = self.router.get_request_handler(&method);
    let tx = self.response_tx.clone();

    let request_trace_context = TraceContext::from_current_span();

    self.ctx.spawn_task_for_client(
      move |ctx| {
        use opentelemetry::trace::FutureExt as _;

        let parent_cx = {
          use opentelemetry::global;
          let mut map = std::collections::HashMap::new();
          if let Some(tp) = &request_trace_context.traceparent {
            map.insert("traceparent".to_string(), tp.clone());
          }
          if let Some(ts) = &request_trace_context.tracestate {
            map.insert("tracestate".to_string(), ts.clone());
          }
          global::get_text_map_propagator(|propagator| {
            propagator
              .extract(&crate::protocol::otel::JsonRpcExtractor::new(&map))
          })
        };

        async move {
          let writer = ctx
            .new_record_writer(Ident::new(&format!("request({})", request_id)));

          let (response, writer) = handler(
            method.to_string(),
            request_id.clone(),
            params,
            ctx,
            writer,
          )
          .await;

          if let Some(response) = response
            && let Err(e) = tx.send(response).await
          {
            otel::error!(
              "response_channel_send_failed",
              format!(
                "Failed to send response to RPC task for request {}: {:?}",
                request_id, e
              )
            );
          }

          Some(writer)
        }
        .with_context(parent_cx)
      },
      method_lane,
      task_client_id,
    );
  }

  fn handle_response(&mut self, response: jsonrpc::Response) {
    let request_id = response.id().clone();

    if let Some(init_id) = &self.init_request_id
      && *init_id == request_id
    {
      self.init_request_id = None;

      if response.error().is_some() {
        self.state.set(State::Uninitialized);
      } else {
        self.state.set(State::Initialized);
      }
    }

    self.pending_requests.lock().remove(&request_id);

    use opentelemetry::trace::SpanKind;
    otel::span!("rpc.send_response", kind = SpanKind::Producer);

    if let Err(e) = self.conn.sender.try_send(Message::Response(response)) {
      otel::error!(
        "response_send_failed",
        format!(
          "Failed to send response for request {}: {:?}",
          request_id, e
        )
      );
    }
  }

  fn handle_notification(
    &mut self,
    notification: jsonrpc::Notification,
  ) -> ControlFlow<()> {
    let trace_ctx = notification.trace_context();
    let _guard = trace_ctx.attach();

    use opentelemetry::trace::SpanKind;
    otel::span!("rpc.receive_notification", kind = SpanKind::Consumer);

    // Extract client_id from notification before into_parts() consumes it
    let msg_client_id = notification.client_id();

    // Validate client_id if present - must match this connection's assigned ID
    if let Some(id) = msg_client_id
      && id != self.ctx.client_id()
    {
      // Silently ignore notifications with mismatched client_id
      // (notifications don't have responses, so we can't send an error)
      otel::error!(
        "notification_client_id_mismatch",
        format!(
          "Notification has client_id {:?} but connection has {:?}",
          id,
          self.ctx.client_id()
        )
      );
      return ControlFlow::Continue(());
    }

    // Use the validated client_id (or fallback to connection's ID if not provided)
    let task_client_id = msg_client_id.unwrap_or_else(|| self.ctx.client_id());

    let (method, params) = notification.into_parts();

    if Self::is_initialized_notification(&method) {
      return ControlFlow::Continue(());
    }

    if Self::is_exit_notification(&method) {
      self.state.set(State::Exited);
      self.shutdown_flag.set();
      return ControlFlow::Break(());
    }

    let (handler, method_lane) = self.router.get_notification_handler(&method);

    let notification_trace_context = TraceContext::from_current_span();

    self.ctx.spawn_task_for_client(
      move |ctx| {
        use opentelemetry::trace::FutureExt as _;

        let parent_cx = {
          use opentelemetry::global;
          let mut map = std::collections::HashMap::new();
          if let Some(tp) = &notification_trace_context.traceparent {
            map.insert("traceparent".to_string(), tp.clone());
          }
          if let Some(ts) = &notification_trace_context.tracestate {
            map.insert("tracestate".to_string(), ts.clone());
          }
          global::get_text_map_propagator(|propagator| {
            propagator
              .extract(&crate::protocol::otel::JsonRpcExtractor::new(&map))
          })
        };

        async move {
          let writer = ctx.new_record_writer(Ident::new(&format!(
            "notification({})",
            method
          )));

          let writer = handler(method.to_string(), params, ctx, writer).await;

          Some(writer)
        }
        .with_context(parent_cx)
      },
      method_lane,
      task_client_id,
    );

    ControlFlow::Continue(())
  }

  fn handle_shutdown(&mut self, request_id: jsonrpc::Id) {
    let current_state = self.state.get();

    if current_state != State::Initialized {
      self.send_error_response(request_id, jsonrpc::Error {
        code:    jsonrpc::ErrorCode::InvalidRequest,
        message: "Server is shutting down".into(),
        data:    None,
      });
      return;
    }

    self.state.set(State::ShutDown);
    self.shutdown_flag.set();

    let response =
      Response::from_ok(request_id.clone(), serde_json::Value::Null);

    if let Err(e) = self.conn.sender.try_send(Message::Response(response)) {
      otel::error!(
        "shutdown_response_send_failed",
        format!("Failed to send shutdown response: {:?}", e)
      );
    }
  }
}

use std::{
  fmt::{
    self,
    Debug,
    Formatter,
  },
  sync::atomic::AtomicU8,
};

/// A list of possible states the language server can be in.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum State {
  /// Server has not received an `initialize` request.
  Uninitialized = 0,
  /// Server received an `initialize` request, but has not yet responded.
  Initializing  = 1,
  /// Server received and responded success to an `initialize` request.
  Initialized   = 2,
  /// Server received a `shutdown` request.
  ShutDown      = 3,
  /// Server received an `exit` notification.
  Exited        = 4,
}

/// Atomic value which represents the current state of the server.
#[derive(Clone)]
pub struct ServerState(Arc<AtomicU8>);

impl ServerState {
  pub fn new() -> Self {
    ServerState(Arc::new(AtomicU8::new(State::Uninitialized as u8)))
  }

  pub fn set(&self, state: State) {
    self.0.store(state as u8, Ordering::SeqCst);
  }

  pub fn get(&self) -> State {
    match self.0.load(Ordering::SeqCst) {
      | 0 => State::Uninitialized,
      | 1 => State::Initializing,
      | 2 => State::Initialized,
      | 3 => State::ShutDown,
      | 4 => State::Exited,
      | _ => State::Uninitialized,
    }
  }
}

impl Debug for ServerState {
  fn fmt(&self, f: &mut Formatter) -> fmt::Result {
    self.get().fmt(f)
  }
}