pub mod actions;
pub mod wire;
use crate::llm::action_helper::call_llm;
use crate::llm::actions::executor::ExecutionResult;
use crate::llm::actions::protocol_trait::ActionResult;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::{Event, EventType};
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use crate::state::client_handles::{ClientCommand, ClientSendOutcome};
use actions::{Answer, Delivery};
use anyhow::Result;
use serde_json::json;
use std::collections::{BTreeMap, VecDeque};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, WriteHalf};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
pub const FIRST_BYTE_TIMEOUT: Duration = Duration::from_secs(300);
pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30);
pub const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
pub const MAX_PENDING_MESSAGES: usize = 1000;
pub const MAX_PENDING_BYTES: usize = 16 * 1024 * 1024;
const EVENT_MESSAGES: usize = 20;
const EVENT_BODY_BYTES: usize = 1000;
fn failed_text(command: &str, overloaded: bool) -> &'static str {
match (command, overloaded) {
("MPUB", false) => "MPUB failed: broker backend unavailable",
("MPUB", true) => "MPUB failed: broker backend at capacity",
("DPUB", false) => "DPUB failed: broker backend unavailable",
("DPUB", true) => "DPUB failed: broker backend at capacity",
("SUB", false) => "SUB failed: broker backend unavailable",
("SUB", true) => "SUB failed: broker backend at capacity",
(_, false) => "PUB failed: broker backend unavailable",
(_, true) => "PUB failed: broker backend at capacity",
}
}
#[derive(Debug, Clone, Copy)]
pub struct NsqConfig {
pub first_byte_timeout: Duration,
pub heartbeat_interval: Duration,
pub idle_timeout: Duration,
}
pub struct NsqServer;
impl NsqServer {
pub async fn spawn_with_llm_actions(
listen_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: crate::state::ServerId,
config: NsqConfig,
) -> Result<SocketAddr> {
let listener = TcpListener::bind(listen_addr).await?;
let local_addr = listener.local_addr()?;
Log::new(Some(&status_tx)).info(format!("NSQ broker listening on {}", local_addr));
let protocol = Arc::new(actions::NsqProtocol::new());
let ids = Arc::new(AtomicU64::new(0));
let task_registrar = app_state.clone();
let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
let accept_handle = tokio::spawn(async move {
loop {
match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
b"",
"NSQ",
Some(&status_tx),
)
.await
{
Ok((socket, peer_addr, permit)) => {
let connection_id =
ConnectionId::new(app_state.get_next_unified_id().await);
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus,
ProtocolConnectionInfo,
};
let now = crate::utils::clock::Instant::now();
app_state
.add_connection_to_server(
server_id,
ServerConnectionState {
id: connection_id,
remote_addr: peer_addr,
local_addr,
bytes_sent: 0,
bytes_received: 0,
packets_sent: 0,
packets_received: 0,
last_activity: now,
status: ConnectionStatus::Active,
status_changed_at: now,
protocol_info: ProtocolConnectionInfo::empty(),
},
)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
Log::new(Some(&status_tx))
.info(format!("NSQ client connected from {}", peer_addr));
let session = Session {
peer_addr,
llm_client: llm_client.clone(),
app_state: app_state.clone(),
status_tx: status_tx.clone(),
server_id,
protocol: protocol.clone(),
connection_id,
config,
ids: ids.clone(),
};
let task_owner = app_state.clone();
task_owner
.spawn_server_task(server_id, async move {
let _permit = permit;
session.run(socket).await
})
.await;
}
Err(e) => {
Log::new(Some(&status_tx)).error(format!("NSQ accept error: {}", e));
break;
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
type Writer = WriteHalf<TcpStream>;
struct Session {
peer_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: crate::state::ServerId,
protocol: Arc<actions::NsqProtocol>,
connection_id: ConnectionId,
config: NsqConfig,
ids: Arc<AtomicU64>,
}
struct Flow {
heartbeat: Option<Duration>,
ticker: Option<tokio::time::Interval>,
identified: bool,
client_id: String,
user_agent: String,
sub: Option<(String, String)>,
closing: bool,
rdy: u64,
in_flight: BTreeMap<String, (String, u16)>,
pending: VecDeque<Delivery>,
pending_bytes: usize,
}
fn ticker(every: Duration) -> tokio::time::Interval {
let mut t = tokio::time::interval_at(tokio::time::Instant::now() + every, every);
t.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
t
}
impl Flow {
fn new(heartbeat: Duration) -> Self {
Self {
heartbeat: Some(heartbeat),
ticker: Some(ticker(heartbeat)),
identified: false,
client_id: String::new(),
user_agent: String::new(),
sub: None,
closing: false,
rdy: 0,
in_flight: BTreeMap::new(),
pending: VecDeque::new(),
pending_bytes: 0,
}
}
fn silence_window(&self, idle: Duration) -> Duration {
match self.heartbeat {
Some(hb) => hb * 2,
None => idle,
}
}
fn can_deliver(&self) -> u64 {
if self.sub.is_none() || self.closing {
return 0;
}
self.rdy.saturating_sub(self.in_flight.len() as u64)
}
fn topic(&self) -> (String, String) {
self.sub.clone().unwrap_or_default()
}
}
async fn tick(ticker: &mut Option<tokio::time::Interval>) {
match ticker {
Some(t) => {
t.tick().await;
}
None => std::future::pending::<()>().await,
}
}
async fn next_injected(rx: &mut Option<mpsc::Receiver<ClientCommand>>) -> Option<ClientCommand> {
match rx {
Some(r) => r.recv().await,
None => std::future::pending().await,
}
}
enum ReadError {
Eof,
TimedOut,
Wire(wire::WireError),
Io(std::io::Error),
}
enum Step {
Continue,
Close,
}
enum Wake {
Command(Result<(wire::Command, usize), ReadError>),
Heartbeat,
Injected(Option<ClientCommand>),
}
impl Session {
async fn run(self, socket: TcpStream) {
let (mut reader, mut writer) = tokio::io::split(socket);
let peer_rx = crate::server::peer_support::register_peer_channel(
&self.app_state,
self.server_id,
self.connection_id.as_u32(),
)
.await;
{
let mut framer = Framer::new(&mut reader);
self.session(&mut framer, &mut writer, peer_rx).await;
}
self.app_state
.remove_peer_handle(self.server_id, self.connection_id.as_u32())
.await;
let _ = writer.shutdown().await;
linger(&mut reader).await;
self.app_state
.update_connection_status(
self.server_id,
self.connection_id,
crate::state::server::ConnectionStatus::Closed,
)
.await;
let _ = self.status_tx.send("__UPDATE_UI__".to_string());
}
async fn write(&self, writer: &mut Writer, data: &[u8]) -> bool {
let ok = writer.write_all(data).await.is_ok() && writer.flush().await.is_ok();
if ok {
self.app_state
.update_connection_stats(
self.server_id,
self.connection_id,
None,
Some(data.len() as u64),
None,
Some(1),
)
.await;
}
ok
}
async fn refuse(&self, writer: &mut Writer, code: &'static str, message: &str) -> Step {
let written = self.write(writer, &wire::error_frame(code, message)).await;
if written && !wire::is_fatal(code) {
Step::Continue
} else {
Step::Close
}
}
async fn session<R: AsyncRead + Unpin>(
&self,
framer: &mut Framer<'_, R>,
writer: &mut Writer,
peer_rx: mpsc::Receiver<ClientCommand>,
) {
let log = Log::new(Some(&self.status_tx));
let magic_deadline = tokio::time::Instant::now() + self.config.first_byte_timeout;
match framer.magic(magic_deadline).await {
Ok(true) => {}
Ok(false) => {
log.warn(format!(
"NSQ peer {} did not send the V2 magic decision=fail_closed_bad_magic",
self.peer_addr
));
let _ = self
.write(writer, &wire::error_frame("E_BAD_PROTOCOL", ""))
.await;
return;
}
Err(ReadError::TimedOut) => {
log.info(format!(
"NSQ client {} sent nothing within {}s; closing",
self.peer_addr,
self.config.first_byte_timeout.as_secs()
));
return;
}
Err(_) => return,
}
let mut flow = Flow::new(self.config.heartbeat_interval);
let mut peer_rx = Some(peer_rx);
let mut answered_at = tokio::time::Instant::now();
let heartbeat = wire::response_frame(wire::HEARTBEAT);
loop {
let window = flow.silence_window(self.config.idle_timeout);
let deadline = framer.last_read.max(answered_at) + window;
let wake = tokio::select! {
r = framer.next_command(deadline) => Wake::Command(r),
_ = tick(&mut flow.ticker) => Wake::Heartbeat,
c = next_injected(&mut peer_rx) => Wake::Injected(c),
};
let step = match wake {
Wake::Heartbeat => {
if self.write(writer, &heartbeat).await {
Step::Continue
} else {
Step::Close
}
}
Wake::Injected(None) => {
peer_rx = None;
Step::Continue
}
Wake::Injected(Some(command)) => self.inject(command, &mut flow, writer).await,
Wake::Command(Err(ReadError::Eof)) => {
log.info(format!("NSQ client {} disconnected", self.peer_addr));
Step::Close
}
Wake::Command(Err(ReadError::TimedOut)) => {
log.info(format!(
"NSQ client {} silent for {}s decision=fail_closed_idle; closing",
self.peer_addr,
window.as_secs()
));
Step::Close
}
Wake::Command(Err(ReadError::Io(e))) => {
log.error(format!("NSQ read error from {}: {}", self.peer_addr, e));
Step::Close
}
Wake::Command(Err(ReadError::Wire(e))) => {
log.warn(format!(
"NSQ command from {} refused: {} {} decision={}",
self.peer_addr,
e.code,
e.message,
e.decision()
));
let _ = self.write(writer, &e.to_frame()).await;
Step::Close
}
Wake::Command(Ok((command, used))) => {
self.app_state
.update_connection_stats(
self.server_id,
self.connection_id,
Some(used as u64),
None,
Some(1),
None,
)
.await;
let step = self.command(command, &mut flow, writer).await;
answered_at = tokio::time::Instant::now();
step
}
};
if let Step::Close = step {
return;
}
}
}
async fn command(&self, command: wire::Command, flow: &mut Flow, writer: &mut Writer) -> Step {
use wire::Command as C;
let log = Log::new(Some(&self.status_tx));
let name = command.name();
let subscribed_or_closing = flow.sub.is_some() || flow.closing;
match command {
C::Nop => Step::Continue,
C::Identify(body) => {
if flow.identified {
return self
.refuse(writer, "E_INVALID", "cannot IDENTIFY in current state")
.await;
}
let identify = match wire::parse_identify(&body) {
Ok(i) => i,
Err(e) => {
log.warn(format!(
"NSQ IDENTIFY from {} refused: {} decision=fail_closed_bad_identify",
self.peer_addr, e.message
));
let _ = self.write(writer, &e.to_frame()).await;
return Step::Close;
}
};
flow.identified = true;
match identify.heartbeat_ms {
None => {
flow.heartbeat = None;
flow.ticker = None;
}
Some(0) => {}
Some(ms) => {
let every = Duration::from_millis(ms as u64);
flow.heartbeat = Some(every);
flow.ticker = Some(ticker(every));
}
}
flow.client_id = identify.client_id;
flow.user_agent = identify.user_agent;
let reply = if identify.feature_negotiation {
wire::identify_response(identify.msg_timeout_ms)
} else {
b"OK".to_vec()
};
log.debug(format!(
"NSQ IDENTIFY from {} (user_agent {:?}, heartbeat {:?})",
self.peer_addr, flow.user_agent, flow.heartbeat
));
self.step(self.write(writer, &wire::response_frame(&reply)).await)
}
C::Auth(_) => {
log.warn(format!(
"NSQ AUTH from {} refused decision=fail_closed_auth_disabled",
self.peer_addr
));
self.refuse(writer, "E_AUTH_DISABLED", "AUTH disabled")
.await
}
C::Cls => {
if flow.sub.is_none() {
return self
.refuse(writer, "E_INVALID", "cannot CLS in current state")
.await;
}
flow.closing = true;
flow.pending.clear();
flow.pending_bytes = 0;
self.step(
self.write(writer, &wire::response_frame(b"CLOSE_WAIT"))
.await,
)
}
C::Touch(id) => {
if !subscribed_or_closing {
return self
.refuse(writer, "E_INVALID", "cannot TOUCH in current state")
.await;
}
if flow.in_flight.contains_key(&id) {
Step::Continue
} else {
self.refuse(
writer,
"E_TOUCH_FAILED",
&format!("TOUCH {id} failed ID not in flight"),
)
.await
}
}
C::Sub { topic, channel } => {
if flow.sub.is_some() || flow.closing {
return self
.refuse(writer, "E_INVALID", "cannot SUB in current state")
.await;
}
self.subscribe(topic, channel, flow, writer).await
}
C::Rdy(count) => {
if flow.closing {
return Step::Continue;
}
if flow.sub.is_none() {
return self
.refuse(writer, "E_INVALID", "cannot RDY in current state")
.await;
}
flow.rdy = count;
if !self.flush(flow, writer).await {
return Step::Close;
}
if flow.can_deliver() == 0 {
return Step::Continue;
}
let (topic, _) = flow.topic();
let mut data = self.flow_data(flow);
data["count"] = json!(count);
data["answer_with"] = json!(actions::deliver_answer_with(
&topic,
flow.can_deliver(),
flow.pending.len()
));
self.flow_event(&actions::NSQ_READY_EVENT, data, name, flow, writer)
.await
}
C::Fin(id) => {
if !subscribed_or_closing {
return self
.refuse(writer, "E_INVALID", "cannot FIN in current state")
.await;
}
if flow.in_flight.remove(&id).is_none() {
return self
.refuse(
writer,
"E_FIN_FAILED",
&format!("FIN {id} failed ID not in flight"),
)
.await;
}
if !self.flush(flow, writer).await {
return Step::Close;
}
if flow.closing {
return Step::Continue;
}
let (topic, _) = flow.topic();
let mut data = self.flow_data(flow);
data["message_id"] = json!(id);
data["answer_with"] = json!(actions::deliver_answer_with(
&topic,
flow.can_deliver(),
flow.pending.len()
));
self.flow_event(&actions::NSQ_FINISH_EVENT, data, name, flow, writer)
.await
}
C::Req { id, timeout_ms } => {
if !subscribed_or_closing {
return self
.refuse(writer, "E_INVALID", "cannot REQ in current state")
.await;
}
let Some((body, attempts)) = flow.in_flight.remove(&id) else {
return self
.refuse(
writer,
"E_REQ_FAILED",
&format!("REQ {id} failed ID not in flight"),
)
.await;
};
if !self.flush(flow, writer).await {
return Step::Close;
}
if flow.closing {
return Step::Continue;
}
let mut data = self.flow_data(flow);
data["message_id"] = json!(id);
data["timeout_ms"] = json!(timeout_ms);
data["attempts"] = json!(attempts);
data["body"] = json!(crate::utils::truncate_for_llm(&body, EVENT_BODY_BYTES));
data["answer_with"] = json!(actions::requeue_answer_with(attempts));
self.flow_event(&actions::NSQ_REQUEUE_EVENT, data, name, flow, writer)
.await
}
C::Pub { topic, body } => {
self.publish(name, topic, vec![body], None, flow, writer)
.await
}
C::Dpub {
topic,
defer_ms,
body,
} => {
self.publish(name, topic, vec![body], Some(defer_ms), flow, writer)
.await
}
C::Mpub { topic, messages } => {
self.publish(name, topic, messages, None, flow, writer)
.await
}
}
}
fn step(&self, written: bool) -> Step {
if written {
Step::Continue
} else {
Step::Close
}
}
fn flow_data(&self, flow: &Flow) -> serde_json::Value {
let (topic, channel) = flow.topic();
json!({
"topic": topic,
"channel": channel,
"ready": flow.rdy,
"in_flight": flow.in_flight.len(),
"pending": flow.pending.len(),
"can_deliver": flow.can_deliver(),
})
}
async fn ask(
&self,
event: &Event,
flow: &mut Flow,
writer: &mut Writer,
) -> Result<ExecutionResult> {
let heartbeat = wire::response_frame(wire::HEARTBEAT);
let call = call_llm(
&self.llm_client,
&self.app_state,
self.server_id,
Some(self.connection_id),
event,
self.protocol.as_ref(),
);
tokio::pin!(call);
loop {
tokio::select! {
result = &mut call => return result,
_ = tick(&mut flow.ticker) => {
let _ = self.write(writer, &heartbeat).await;
}
}
}
}
fn answers(&self, result: ExecutionResult) -> Vec<Answer> {
let log = Log::new(Some(&self.status_tx));
for message in &result.messages {
log.info(message);
}
let mut out = Vec::new();
let mut stack = result.protocol_results;
stack.reverse();
while let Some(item) = stack.pop() {
match item {
ActionResult::Multiple(inner) => stack.extend(inner.into_iter().rev()),
other => match Answer::from_result(&other) {
Some(answer) => out.push(answer),
None => log.debug(format!("NSQ: ignoring non-NSQ result {other:?}")),
},
}
}
out
}
async fn subscribe(
&self,
topic: String,
channel: String,
flow: &mut Flow,
writer: &mut Writer,
) -> Step {
let log = Log::new(Some(&self.status_tx));
let event = Event::new(
&actions::NSQ_SUBSCRIBE_EVENT,
json!({
"topic": topic,
"channel": channel,
"ephemeral": wire::is_ephemeral(&channel),
"client_id": flow.client_id,
"user_agent": flow.user_agent,
"answer_with": actions::subscribe_answer_with(&topic, &channel),
}),
);
let answers = match self.ask(&event, flow, writer).await {
Ok(result) => self.answers(result),
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
log.warn(format!(
"NSQ SUB {topic} {channel} from {} decision=fail_closed_llm_error category={:?}",
self.peer_addr, failure
));
log.debug(format!("NSQ LLM call failed: {}", e));
let text = failed_text("SUB", failure.is_overloaded());
return self.refuse(writer, "E_INVALID", text).await;
}
};
let step = match reply_of(&answers) {
Some(Answer::Ok) => {
log.info(format!(
"NSQ SUB {topic} {channel} from {} decision=model_answer",
self.peer_addr
));
flow.sub = Some((topic, channel));
self.step(self.write(writer, &wire::response_frame(b"OK")).await)
}
Some(Answer::Error { code, message }) => {
log.info(format!(
"NSQ SUB {topic} {channel} from {} decision=model_reject code={code}",
self.peer_addr
));
self.refuse(writer, static_code(code), message).await
}
_ => {
log.warn(format!(
"NSQ SUB {topic} {channel} from {} decision=model_silent; refusing",
self.peer_addr
));
return self.refuse(writer, "E_INVALID", "SUB failed").await;
}
};
self.finish_answers(step, &answers, flow, writer).await
}
async fn publish(
&self,
command: &'static str,
topic: String,
bodies: Vec<Vec<u8>>,
defer_ms: Option<i64>,
flow: &mut Flow,
writer: &mut Writer,
) -> Step {
let log = Log::new(Some(&self.status_tx));
let failed_code: &'static str = match command {
"MPUB" => "E_MPUB_FAILED",
"DPUB" => "E_DPUB_FAILED",
_ => "E_PUB_FAILED",
};
let total_bytes: usize = bodies.iter().map(Vec::len).sum();
let shown: Vec<String> = bodies
.iter()
.take(EVENT_MESSAGES)
.map(|b| crate::utils::truncate_for_llm(&String::from_utf8_lossy(b), EVENT_BODY_BYTES))
.collect();
let mut data = json!({
"command": command,
"topic": topic,
"messages": shown,
"message_count": bodies.len(),
"total_bytes": total_bytes,
"answer_with": actions::publish_answer_with(command, &topic, bodies.len()),
});
if let Some(defer) = defer_ms {
data["defer_ms"] = json!(defer);
}
let event = Event::new(&actions::NSQ_PUBLISH_EVENT, data);
let answers = match self.ask(&event, flow, writer).await {
Ok(result) => self.answers(result),
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
log.warn(format!(
"NSQ {command} {topic} from {} decision=fail_closed_llm_error category={:?}",
self.peer_addr, failure
));
log.debug(format!("NSQ LLM call failed: {}", e));
let text = failed_text(command, failure.is_overloaded());
return self.refuse(writer, failed_code, text).await;
}
};
let step = match reply_of(&answers) {
Some(Answer::Ok) => {
log.info(format!(
"NSQ {command} {topic} ({} message(s)) from {} decision=model_answer",
bodies.len(),
self.peer_addr
));
self.step(self.write(writer, &wire::response_frame(b"OK")).await)
}
Some(Answer::Error { code, message }) => {
log.info(format!(
"NSQ {command} {topic} from {} decision=model_reject code={code}",
self.peer_addr
));
self.refuse(writer, static_code(code), message).await
}
_ => {
log.warn(format!(
"NSQ {command} {topic} from {} decision=model_silent; refusing",
self.peer_addr
));
return self
.refuse(writer, failed_code, &format!("{command} failed"))
.await;
}
};
self.finish_answers(step, &answers, flow, writer).await
}
async fn finish_answers(
&self,
step: Step,
answers: &[Answer],
flow: &mut Flow,
writer: &mut Writer,
) -> Step {
if let Step::Close = step {
return Step::Close;
}
let deliveries: Vec<Delivery> = answers
.iter()
.filter_map(|a| match a {
Answer::Deliver(d) => Some(d.clone()),
_ => None,
})
.flatten()
.collect();
if !deliveries.is_empty() {
self.enqueue(flow, deliveries);
if !self.flush(flow, writer).await {
return Step::Close;
}
}
if answers.iter().any(|a| matches!(a, Answer::Close)) {
return Step::Close;
}
Step::Continue
}
async fn flow_event(
&self,
event_type: &'static EventType,
data: serde_json::Value,
command: &'static str,
flow: &mut Flow,
writer: &mut Writer,
) -> Step {
let log = Log::new(Some(&self.status_tx));
let event = Event::new(event_type, data);
let answers = match self.ask(&event, flow, writer).await {
Ok(result) => self.answers(result),
Err(e) => {
log.warn(format!(
"NSQ {command} from {} decision=fail_closed_llm_error category={:?}; \
delivering nothing",
self.peer_addr,
crate::utils::WireFailure::classify(&e)
));
log.debug(format!("NSQ LLM call failed: {}", e));
return Step::Continue;
}
};
if let Some(Answer::Error { code, message }) =
answers.iter().find(|a| matches!(a, Answer::Error { .. }))
{
log.info(format!(
"NSQ {command} from {} decision=model_reject code={code}",
self.peer_addr
));
if let Step::Close = self.refuse(writer, static_code(code), message).await {
return Step::Close;
}
}
let offered: usize = answers
.iter()
.map(|a| match a {
Answer::Deliver(d) => d.len(),
_ => 0,
})
.sum();
let step = self
.finish_answers(Step::Continue, &answers, flow, writer)
.await;
if offered > 0 {
log.info(format!(
"NSQ {command} from {} decision=model_answer offered={offered} in_flight={} \
pending={}",
self.peer_addr,
flow.in_flight.len(),
flow.pending.len()
));
} else if !answers.iter().any(|a| matches!(a, Answer::Error { .. })) {
log.info(format!(
"NSQ {command} from {} decision=model_silent (nothing delivered)",
self.peer_addr
));
}
step
}
fn enqueue(&self, flow: &mut Flow, deliveries: Vec<Delivery>) {
let log = Log::new(Some(&self.status_tx));
if flow.sub.is_none() || flow.closing {
log.warn(format!(
"NSQ: {} message(s) for {} dropped - the connection is not subscribed{} \
decision=fail_closed_not_subscribed",
deliveries.len(),
self.peer_addr,
if flow.closing { " (closing)" } else { "" }
));
return;
}
let mut dropped = 0usize;
for delivery in deliveries {
if flow.pending.len() >= MAX_PENDING_MESSAGES
|| flow.pending_bytes + delivery.body.len() > MAX_PENDING_BYTES
{
dropped += 1;
continue;
}
flow.pending_bytes += delivery.body.len();
flow.pending.push_back(delivery);
}
if dropped > 0 {
log.warn(format!(
"NSQ: {dropped} message(s) for {} dropped over the pending bound ({} messages, \
{} bytes) decision=fail_closed_pending_full",
self.peer_addr, MAX_PENDING_MESSAGES, MAX_PENDING_BYTES
));
}
}
async fn flush(&self, flow: &mut Flow, writer: &mut Writer) -> bool {
self.flush_counting(flow, writer).await.is_some()
}
async fn flush_counting(&self, flow: &mut Flow, writer: &mut Writer) -> Option<usize> {
let mut written = 0usize;
while flow.can_deliver() > 0 {
let Some(delivery) = flow.pending.pop_front() else {
break;
};
flow.pending_bytes -= delivery.body.len();
let n = self.ids.fetch_add(1, Ordering::SeqCst) + 1;
let id = wire::message_id_for(n);
let timestamp = crate::utils::clock::SystemTime::now()
.duration_since(crate::utils::clock::UNIX_EPOCH)
.map(|d| d.as_nanos() as i64)
.unwrap_or(0);
let frame =
wire::message_frame(timestamp, delivery.attempts, &id, delivery.body.as_bytes());
if !self.write(writer, &frame).await {
return None;
}
written += frame.len();
flow.in_flight.insert(
String::from_utf8_lossy(&id).into_owned(),
(delivery.body, delivery.attempts),
);
}
Some(written)
}
async fn inject(&self, command: ClientCommand, flow: &mut Flow, writer: &mut Writer) -> Step {
let action = command.action.clone();
let result = crate::llm::actions::executor::execute_actions(
vec![action.clone()],
&self.app_state,
Some(self.protocol.as_ref()),
Some(self.server_id),
None,
)
.await;
let (outcome, step) = match result {
Err(e) => (Err(e), Step::Continue),
Ok(result) => {
let failures = result.failure_summary();
let answers = self.answers(result);
if answers.is_empty() {
let error = failures.unwrap_or_else(|| "no NSQ action".to_string());
(Ok(ClientSendOutcome::Rejected { error }), Step::Continue)
} else {
let mut bytes = 0usize;
let mut step = Step::Continue;
for answer in &answers {
let frame = match answer {
Answer::Ok => wire::response_frame(b"OK"),
Answer::Error { code, message } => {
let code = static_code(code);
if wire::is_fatal(code) {
step = Step::Close;
}
wire::error_frame(code, message)
}
Answer::Close => {
step = Step::Close;
continue;
}
Answer::Deliver(d) => {
self.enqueue(flow, d.clone());
match self.flush_counting(flow, writer).await {
Some(n) => bytes += n,
None => step = Step::Close,
}
continue;
}
};
if self.write(writer, &frame).await {
bytes += frame.len();
} else {
step = Step::Close;
}
}
let outcome = match step {
Step::Close => ClientSendOutcome::Disconnected,
Step::Continue if bytes > 0 => {
ClientSendOutcome::Sent { bytes_sent: bytes }
}
Step::Continue => ClientSendOutcome::Executed {
detail: format!("queued; {} pending for RDY", flow.pending.len()),
},
};
(Ok(outcome), step)
}
}
};
let outcome_json = match &outcome {
Ok(o) => serde_json::to_value(o).unwrap_or(serde_json::Value::Null),
Err(e) => json!({"error": e.to_string()}),
};
self.app_state
.record_access_log(
crate::state::AccessLogOwner::Server(self.server_id.as_u32()),
"NSQ",
Some(self.connection_id.as_u32()),
"injected_action",
action,
vec![outcome_json],
)
.await;
let _ = self.status_tx.send("__UPDATE_UI__".to_string());
let _ = command.reply_tx.send(outcome);
step
}
}
fn reply_of(answers: &[Answer]) -> Option<&Answer> {
answers
.iter()
.find(|a| matches!(a, Answer::Ok | Answer::Error { .. }))
}
fn static_code(code: &str) -> &'static str {
wire::ERROR_CODES
.iter()
.copied()
.find(|c| *c == code)
.unwrap_or("E_INVALID")
}
const LINGER_TIME: Duration = Duration::from_secs(2);
const LINGER_BYTES: usize = 64 * 1024;
async fn linger<R: AsyncRead + Unpin>(reader: &mut R) {
let deadline = tokio::time::Instant::now() + LINGER_TIME;
let mut sink = [0u8; 4096];
let mut drained = 0usize;
while drained < LINGER_BYTES {
match tokio::time::timeout_at(deadline, reader.read(&mut sink)).await {
Ok(Ok(n)) if n > 0 => drained += n,
_ => break,
}
}
}
struct Framer<'a, R> {
reader: &'a mut R,
pending: Vec<u8>,
chunk: Vec<u8>,
last_read: tokio::time::Instant,
}
impl<'a, R: AsyncRead + Unpin> Framer<'a, R> {
fn new(reader: &'a mut R) -> Self {
Self {
reader,
pending: Vec::new(),
chunk: vec![0u8; 16 * 1024],
last_read: tokio::time::Instant::now(),
}
}
async fn fill(&mut self, deadline: tokio::time::Instant) -> Result<(), ReadError> {
match tokio::time::timeout_at(deadline, self.reader.read(&mut self.chunk)).await {
Err(_) => Err(ReadError::TimedOut),
Ok(Ok(0)) => Err(ReadError::Eof),
Ok(Ok(n)) => {
self.pending.extend_from_slice(&self.chunk[..n]);
self.last_read = tokio::time::Instant::now();
Ok(())
}
Ok(Err(e)) => Err(ReadError::Io(e)),
}
}
async fn magic(&mut self, deadline: tokio::time::Instant) -> Result<bool, ReadError> {
while self.pending.len() < wire::MAGIC_V2.len() {
self.fill(deadline).await?;
}
let ok = &self.pending[..4] == wire::MAGIC_V2;
self.pending.drain(..4);
Ok(ok)
}
async fn next_command(
&mut self,
deadline: tokio::time::Instant,
) -> Result<(wire::Command, usize), ReadError> {
loop {
match wire::parse_command(&self.pending) {
Ok(Some((command, used))) => {
self.pending.drain(..used);
return Ok((command, used));
}
Ok(None) => self.fill(deadline).await?,
Err(e) => return Err(ReadError::Wire(e)),
}
}
}
}