pub mod actions;
pub mod frame;
use crate::llm::action_helper::call_llm;
use crate::llm::actions::protocol_trait::ActionResult;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::Event;
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use actions::{
StompProtocol, STOMP_ACK_EVENT, STOMP_CONNECT_EVENT, STOMP_DISCONNECT_EVENT, STOMP_NACK_EVENT,
STOMP_SEND_EVENT, STOMP_SUBSCRIBE_EVENT, STOMP_UNSUBSCRIBE_EVENT, STOMP_VERSION,
};
use anyhow::Result;
use frame::{error_frame, receipt_frame, ParseOutcome, StompFrame};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::{mpsc, Mutex};
const FIRST_FRAME_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const IDLE_BETWEEN_FRAMES_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(1800);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL_MESSAGE: &str = "too many connections";
const CONNECTION_CAP_REFUSAL_BODY: &str =
"The server has reached its connection limit. Retry later.";
pub struct StompServer;
impl StompServer {
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,
) -> Result<SocketAddr> {
let listener = crate::server::socket_helpers::create_reusable_tcp_listener(listen_addr)
.await
.map_err(|e| anyhow::anyhow!("STOMP could not bind {listen_addr}: {e}"))?;
let local_addr = listener.local_addr()?;
Log::new(Some(&status_tx)).info(format!("STOMP server listening on {}", local_addr));
let protocol = Arc::new(StompProtocol::new());
let task_registrar = app_state.clone();
let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
let cap_refusal = error_frame(CONNECTION_CAP_REFUSAL_MESSAGE, CONNECTION_CAP_REFUSAL_BODY);
let accept_handle = tokio::spawn(async move {
loop {
match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
&cap_refusal,
"STOMP",
Some(&status_tx),
)
.await
{
Ok((socket, peer_addr, permit)) => {
let connection_id =
ConnectionId::new(app_state.get_next_unified_id().await);
let local_addr_conn = socket.local_addr().unwrap_or(local_addr);
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: local_addr_conn,
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::new(
serde_json::json!({"session": "pre-connect"}),
),
},
)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
Log::new(Some(&status_tx))
.info(format!("STOMP client connected from {}", peer_addr));
let llm_clone = llm_client.clone();
let state_clone = app_state.clone();
let status_clone = status_tx.clone();
let protocol_clone = protocol.clone();
let conn_handle = tokio::spawn(async move {
let _permit = permit;
handle_stomp_connection(
socket,
peer_addr,
llm_clone,
state_clone,
status_clone,
server_id,
protocol_clone,
connection_id,
)
.await
});
app_state.register_server_task(server_id, conn_handle).await;
}
Err(e) => {
Log::new(Some(&status_tx)).error(format!("STOMP accept error: {}", e));
break;
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
async fn write_counted<W>(
write_half: &Arc<Mutex<W>>,
data: &[u8],
app_state: &AppState,
server_id: crate::state::ServerId,
connection_id: ConnectionId,
) -> std::io::Result<()>
where
W: tokio::io::AsyncWrite + Unpin,
{
{
let mut writer = write_half.lock().await;
writer.write_all(data).await?;
writer.flush().await?;
}
app_state
.update_connection_stats(
server_id,
connection_id,
None,
Some(data.len() as u64),
None,
Some(1),
)
.await;
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn handle_stomp_connection(
socket: tokio::net::TcpStream,
peer_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: crate::state::ServerId,
protocol: Arc<StompProtocol>,
connection_id: ConnectionId,
) {
let (reader, write_half) = tokio::io::split(socket);
let write_half = Arc::new(Mutex::new(write_half));
let peer_rx = crate::server::peer_support::register_peer_channel(
&app_state,
server_id,
connection_id.as_u32(),
)
.await;
crate::server::peer_support::spawn_peer_command_task(
peer_rx,
protocol.clone(),
app_state.clone(),
server_id,
connection_id.as_u32(),
write_half.clone(),
status_tx.clone(),
);
run_stomp_session(
reader,
&write_half,
peer_addr,
&llm_client,
&app_state,
&status_tx,
server_id,
&protocol,
connection_id,
)
.await;
app_state
.remove_peer_handle(server_id, connection_id.as_u32())
.await;
let _ = write_half.lock().await.shutdown().await;
use crate::state::server::ConnectionStatus;
app_state
.update_connection_status(server_id, connection_id, ConnectionStatus::Closed)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
}
enum Flow {
Continue,
Close,
}
#[allow(clippy::too_many_arguments)]
async fn run_stomp_session<R, W>(
mut reader: R,
write_half: &Arc<Mutex<W>>,
peer_addr: SocketAddr,
llm_client: &OllamaClient,
app_state: &Arc<AppState>,
status_tx: &mpsc::UnboundedSender<String>,
server_id: crate::state::ServerId,
protocol: &Arc<StompProtocol>,
connection_id: ConnectionId,
) where
R: tokio::io::AsyncRead + Unpin,
W: tokio::io::AsyncWrite + Unpin,
{
let log = Log::new(Some(status_tx));
let mut pending: Vec<u8> = Vec::new();
let mut read_buf = vec![0u8; 8192];
let mut connected = false;
loop {
loop {
match frame::parse_frame(&pending) {
Ok(ParseOutcome::Incomplete) => break,
Ok(ParseOutcome::Heartbeat { consumed }) => {
log.trace(format!(
"STOMP drained {consumed} inter-frame EOL byte(s) from {peer_addr}"
));
pending.drain(..consumed);
}
Ok(ParseOutcome::Frame { frame, consumed }) => {
pending.drain(..consumed);
let flow = handle_frame(
frame,
&mut connected,
write_half,
peer_addr,
llm_client,
app_state,
status_tx,
server_id,
protocol,
connection_id,
)
.await;
if matches!(flow, Flow::Close) {
return;
}
}
Err(e) => {
log.warn(format!("STOMP framing error from {peer_addr}: {e}"));
let _ = write_counted(
write_half,
&error_frame("malformed frame", &e.to_string()),
app_state,
server_id,
connection_id,
)
.await;
return;
}
}
}
let read_deadline = if connected {
IDLE_BETWEEN_FRAMES_TIMEOUT
} else {
FIRST_FRAME_READ_TIMEOUT
};
let read_result =
match tokio::time::timeout(read_deadline, reader.read(&mut read_buf)).await {
Ok(v) => v,
Err(_) => {
log.info(format!(
"STOMP client {} sent nothing for {}s; closing \
decision=fail_closed_idle_timeout",
peer_addr,
read_deadline.as_secs()
));
let _ = write_counted(
write_half,
&error_frame("idle timeout", "No frame received within the idle timeout."),
app_state,
server_id,
connection_id,
)
.await;
return;
}
};
match read_result {
Ok(0) => {
log.info(format!("STOMP client {} disconnected", peer_addr));
return;
}
Ok(n) => {
app_state
.update_connection_stats(
server_id,
connection_id,
Some(n as u64),
None,
Some(1),
None,
)
.await;
log.debug(format!("STOMP received {} bytes from {}", n, peer_addr));
pending.extend_from_slice(&read_buf[..n]);
}
Err(e) => {
log.error(format!("STOMP read error from {}: {}", peer_addr, e));
return;
}
}
}
}
async fn refuse_frame<W>(
write_half: &Arc<Mutex<W>>,
app_state: &AppState,
server_id: crate::state::ServerId,
connection_id: ConnectionId,
message: &str,
body: &str,
) -> Flow
where
W: tokio::io::AsyncWrite + Unpin,
{
let _ = write_counted(
write_half,
&error_frame(message, body),
app_state,
server_id,
connection_id,
)
.await;
Flow::Close
}
#[allow(clippy::too_many_arguments)]
async fn handle_frame<W>(
frame: StompFrame,
connected: &mut bool,
write_half: &Arc<Mutex<W>>,
peer_addr: SocketAddr,
llm_client: &OllamaClient,
app_state: &Arc<AppState>,
status_tx: &mpsc::UnboundedSender<String>,
server_id: crate::state::ServerId,
protocol: &Arc<StompProtocol>,
connection_id: ConnectionId,
) -> Flow
where
W: tokio::io::AsyncWrite + Unpin,
{
let log = Log::new(Some(status_tx));
let command = frame.command.clone();
log.debug(format!("STOMP {} frame from {}", command, peer_addr));
if !*connected && !matches!(command.as_str(), "CONNECT" | "STOMP") {
log.warn(format!(
"STOMP {command} from {peer_addr} before CONNECT: decision=protocol_error"
));
return refuse_frame(
write_half,
app_state,
server_id,
connection_id,
"expected CONNECT",
&format!("The first frame of a STOMP session must be CONNECT or STOMP, not {command}."),
)
.await;
}
let receipt_id = frame.header("receipt").map(str::to_string);
let event = match command.as_str() {
"CONNECT" | "STOMP" => {
let accept_version = frame.header("accept-version").unwrap_or("");
if !accept_version.split(',').any(|v| v.trim() == STOMP_VERSION) {
log.warn(format!(
"STOMP CONNECT from {peer_addr} wants accept-version={accept_version:?}: \
decision=version_refused"
));
let body = format!(
"This server implements STOMP {STOMP_VERSION} only. Send \
accept-version:{STOMP_VERSION}."
);
let refusal = StompFrame::new(
"ERROR",
vec![
("version".to_string(), STOMP_VERSION.to_string()),
("message".to_string(), "unsupported version".to_string()),
("content-type".to_string(), "text/plain".to_string()),
],
body.into_bytes(),
)
.encode();
let _ =
write_counted(write_half, &refusal, app_state, server_id, connection_id).await;
return Flow::Close;
}
Event::new(
&STOMP_CONNECT_EVENT,
serde_json::json!({
"accept_version": accept_version,
"host": frame.header("host").unwrap_or(""),
"login": frame.header("login").unwrap_or(""),
"passcode": frame.header("passcode").unwrap_or(""),
"heart_beat": frame.header("heart-beat").unwrap_or("0,0"),
}),
)
}
"SEND" => {
let Some(destination) = frame.header("destination") else {
return refuse_frame(
write_half,
app_state,
server_id,
connection_id,
"missing destination",
"A SEND frame must carry a destination header.",
)
.await;
};
let (body, body_encoding) = body_for_event(&frame.body);
Event::new(
&STOMP_SEND_EVENT,
serde_json::json!({
"destination": destination,
"body": body,
"body_encoding": body_encoding,
"headers": frame.headers_except(&["destination", "content-length"]),
}),
)
}
"SUBSCRIBE" => {
let (Some(destination), Some(id)) = (frame.header("destination"), frame.header("id"))
else {
return refuse_frame(
write_half,
app_state,
server_id,
connection_id,
"missing subscription headers",
"A SUBSCRIBE frame must carry both destination and id headers.",
)
.await;
};
Event::new(
&STOMP_SUBSCRIBE_EVENT,
serde_json::json!({
"destination": destination,
"id": id,
"ack_mode": frame.header("ack").unwrap_or("auto"),
}),
)
}
"UNSUBSCRIBE" => {
let Some(id) = frame.header("id") else {
return refuse_frame(
write_half,
app_state,
server_id,
connection_id,
"missing id",
"An UNSUBSCRIBE frame must carry an id header.",
)
.await;
};
Event::new(&STOMP_UNSUBSCRIBE_EVENT, serde_json::json!({ "id": id }))
}
"ACK" | "NACK" => {
let Some(id) = frame.header("id") else {
return refuse_frame(
write_half,
app_state,
server_id,
connection_id,
"missing id",
&format!("A {command} frame must carry an id header."),
)
.await;
};
let event_type = if command == "ACK" {
&*STOMP_ACK_EVENT
} else {
&*STOMP_NACK_EVENT
};
Event::new(event_type, serde_json::json!({ "id": id }))
}
"DISCONNECT" => Event::new(
&STOMP_DISCONNECT_EVENT,
serde_json::json!({ "receipt": receipt_id.clone().unwrap_or_default() }),
),
"BEGIN" | "COMMIT" | "ABORT" => {
log.info(format!(
"STOMP {command} from {peer_addr} acknowledged; transactions are not implemented"
));
if let Some(id) = &receipt_id {
if write_counted(
write_half,
&receipt_frame(id),
app_state,
server_id,
connection_id,
)
.await
.is_err()
{
return Flow::Close;
}
}
return Flow::Continue;
}
other => {
log.warn(format!(
"STOMP unknown command {other:?} from {peer_addr}: decision=protocol_error"
));
return refuse_frame(
write_half,
app_state,
server_id,
connection_id,
"unknown command",
&format!("{other} is not a STOMP 1.2 client command."),
)
.await;
}
};
let is_connect = matches!(command.as_str(), "CONNECT" | "STOMP");
let is_disconnect = command == "DISCONNECT";
let execution = match call_llm(
llm_client,
app_state,
server_id,
Some(connection_id),
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
let class = if failure.is_overloaded() {
"overloaded"
} else {
"unavailable"
};
log.warn(format!(
"STOMP {command} on {connection_id} could not be answered: \
decision=fail_closed_llm_error class={class} error={e}"
));
let _ = write_counted(
write_half,
&error_frame(failure.text(), failure.prefixed_text()),
app_state,
server_id,
connection_id,
)
.await;
return Flow::Close;
}
};
for message in &execution.messages {
log.info(message);
}
let mut wrote_output = false;
let mut wrote_receipt = false;
let mut wrote_error = false;
let mut model_closed = false;
for result in execution.protocol_results {
match result {
ActionResult::Output(bytes) => {
if write_counted(write_half, &bytes, app_state, server_id, connection_id)
.await
.is_err()
{
return Flow::Close;
}
wrote_output = true;
wrote_receipt |= receipt_id
.as_deref()
.is_some_and(|id| is_receipt_for(&bytes, id));
wrote_error |= bytes.starts_with(b"ERROR\n");
log.debug(format!("STOMP sent {} bytes to {}", bytes.len(), peer_addr));
log.trace(format!(
"STOMP sent: {}",
String::from_utf8_lossy(&bytes).replace('\0', "\\0")
));
}
ActionResult::CloseConnection => model_closed = true,
_ => {}
}
}
if is_connect {
if wrote_error {
log.info(format!(
"STOMP CONNECT from {peer_addr} refused: decision=model_reject"
));
return Flow::Close;
}
if !wrote_output {
log.warn(format!(
"STOMP CONNECT from {peer_addr} produced no frame ({} failed action(s)): \
decision=fail_closed_no_answer",
execution.failures.len()
));
let _ = write_counted(
write_half,
&error_frame(
"connection refused",
"The server did not accept this connection.",
),
app_state,
server_id,
connection_id,
)
.await;
return Flow::Close;
}
*connected = true;
log.info(format!("STOMP session opened for {peer_addr}"));
} else if wrote_error {
log.info(format!(
"STOMP closing {connection_id} after ERROR: decision=model_reject"
));
return Flow::Close;
}
let owes_receipt = receipt_id.is_some() && !is_connect && !wrote_receipt;
if let (true, Some(id)) = (owes_receipt, &receipt_id) {
if write_counted(
write_half,
&receipt_frame(id),
app_state,
server_id,
connection_id,
)
.await
.is_err()
{
return Flow::Close;
}
}
if model_closed {
log.info(format!(
"STOMP closing {connection_id}: decision=model_close"
));
return Flow::Close;
}
if is_disconnect {
log.info(format!("STOMP session closed by {peer_addr}"));
return Flow::Close;
}
if !wrote_output {
log.debug(format!(
"STOMP {command} on {connection_id} answered with no frame: \
decision=model_no_actions"
));
}
Flow::Continue
}
fn is_receipt_for(bytes: &[u8], receipt_id: &str) -> bool {
bytes.starts_with(b"RECEIPT\n") && bytes == frame::receipt_frame(receipt_id)
}
fn body_for_event(body: &[u8]) -> (String, &'static str) {
if body
.iter()
.all(|b| b.is_ascii_graphic() || b.is_ascii_whitespace())
{
(String::from_utf8_lossy(body).to_string(), "utf8")
} else {
(hex::encode(body), "hex")
}
}