pub mod actions;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Method, Request, Response, StatusCode};
use hyper_util::rt::TokioIo;
use serde_json::{json, Value};
use tokio::sync::mpsc;
use tracing::error;
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::server::jsonrpc::actions::{JsonRpcProtocol, JSONRPC_METHOD_CALL_EVENT};
use crate::state::app_state::AppState;
const PARSE_ERROR: i32 = -32700;
const INVALID_REQUEST: i32 = -32600;
const INTERNAL_ERROR: i32 = -32603;
const SERVER_BUSY: i32 = -32000;
pub const MAX_REQUEST_BODY_BYTES: usize = 4 * 1024 * 1024;
const MAX_BATCH_LEN: usize = 128;
const FIRST_BYTE_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const IDLE_BETWEEN_REQUESTS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(75);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL: &[u8] = b"HTTP/1.1 503 Service Unavailable\r\n\
Content-Length: 0\r\nRetry-After: 5\r\nConnection: close\r\n\r\n";
pub struct JsonRpcServer;
impl JsonRpcServer {
pub async fn spawn_with_llm_actions(
listen_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
_send_first: bool,
server_id: crate::state::ServerId,
) -> anyhow::Result<SocketAddr> {
let listener =
crate::server::socket_helpers::create_reusable_tcp_listener(listen_addr).await?;
let local_addr = listener.local_addr()?;
Log::new(Some(&status_tx)).info(format!("JSON-RPC server listening on {}", local_addr));
let protocol = Arc::new(JsonRpcProtocol::new());
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,
CONNECTION_CAP_REFUSAL,
"JSON-RPC",
Some(&status_tx),
)
.await
{
Ok((stream, remote_addr, permit)) => {
let connection_id =
ConnectionId::new(app_state.get_next_unified_id().await);
let local_addr_conn = stream.local_addr().unwrap_or(local_addr);
Log::new(Some(&status_tx)).info(format!(
"JSON-RPC connection {} from {}",
connection_id, remote_addr
));
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus,
ProtocolConnectionInfo,
};
let now = crate::utils::clock::Instant::now();
let conn_state = ServerConnectionState {
id: connection_id,
remote_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::empty(),
};
app_state
.add_connection_to_server(server_id, conn_state)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
let llm_client_clone = llm_client.clone();
let app_state_clone = app_state.clone();
let status_tx_clone = status_tx.clone();
let protocol_clone = protocol.clone();
let task_owner = app_state.clone();
task_owner
.spawn_server_task(server_id, async move {
let _permit = permit;
let spoke = matches!(
tokio::time::timeout(
FIRST_BYTE_READ_TIMEOUT,
stream.peek(&mut [0u8; 1]),
)
.await,
Ok(Ok(n)) if n > 0
);
if !spoke {
Log::new(Some(&status_tx_clone)).debug(format!(
"JSON-RPC peer {} sent nothing for {}s; closing before \
any request",
remote_addr,
FIRST_BYTE_READ_TIMEOUT.as_secs()
));
} else {
let io = TokioIo::new(stream);
let status_for_service = status_tx_clone.clone();
let app_state_for_service = app_state_clone.clone();
let activity = std::sync::Arc::new(
crate::server::accept_bounded::ConnectionActivity::new(),
);
let activity_for_service = std::sync::Arc::clone(&activity);
let service = service_fn(move |req: Request<Incoming>| {
let llm_clone = llm_client_clone.clone();
let state_clone = app_state_for_service.clone();
let status_clone = status_for_service.clone();
let protocol_clone = protocol_clone.clone();
let activity = std::sync::Arc::clone(&activity_for_service);
async move {
let _busy = activity.busy();
handle_jsonrpc_request(
req,
connection_id,
llm_clone,
state_clone,
status_clone,
protocol_clone,
server_id,
)
.await
}
});
let conn = http1::Builder::new().serve_connection(io, service);
tokio::pin!(conn);
tokio::select! {
result = &mut conn => {
if let Err(err) = result {
error!(
"Error serving JSON-RPC connection: {:?}",
err
);
}
}
_ = crate::server::accept_bounded::watch_idle(
std::sync::Arc::clone(&activity),
IDLE_BETWEEN_REQUESTS_TIMEOUT,
) => {
Log::new(Some(&status_tx_clone)).debug(format!(
"JSON-RPC connection {} idle for {}s; closing",
connection_id,
IDLE_BETWEEN_REQUESTS_TIMEOUT.as_secs()
));
}
}
}
app_state_clone
.close_connection_on_server(server_id, connection_id)
.await;
Log::new(Some(&status_tx_clone))
.info(format!("JSON-RPC connection {} closed", connection_id));
let _ = status_tx_clone.send("__UPDATE_UI__".to_string());
})
.await;
}
Err(e) => {
Log::new(Some(&status_tx))
.error(format!("Failed to accept JSON-RPC connection: {}", e));
break;
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
async fn handle_jsonrpc_request(
req: Request<Incoming>,
connection_id: ConnectionId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
protocol: Arc<JsonRpcProtocol>,
server_id: crate::state::ServerId,
) -> Result<Response<Full<Bytes>>, Infallible> {
let method = req.method().clone();
let uri = req.uri().clone();
let log = Log::new(Some(&status_tx));
log.debug(format!("JSON-RPC request: {} {}", method, uri.path()));
if method != Method::POST {
return Ok(build_error_response(
INVALID_REQUEST,
"JSON-RPC requires POST method",
None,
None,
));
}
let body_bytes = match http_body_util::Limited::new(req.into_body(), MAX_REQUEST_BODY_BYTES)
.collect()
.await
{
Ok(collected) => collected.to_bytes(),
Err(e) => {
tracing::warn!("JSON-RPC request body rejected (decision=reject_oversized): {e}");
log.warn(format!("Failed to read request body: {}", e));
return Ok(build_error_response(
INVALID_REQUEST,
&format!(
"Request body rejected (limit {} bytes)",
MAX_REQUEST_BODY_BYTES
),
None,
None,
));
}
};
let request_value: Value = match serde_json::from_slice(&body_bytes) {
Ok(json) => json,
Err(e) => {
log.warn(format!("Failed to parse JSON: {}", e));
return Ok(build_error_response(PARSE_ERROR, "Parse error", None, None));
}
};
log.trace(format!(
"JSON-RPC request body: {}",
serde_json::to_string_pretty(&request_value).unwrap_or_default()
));
match request_value {
Value::Array(requests) if !requests.is_empty() => {
if requests.len() > MAX_BATCH_LEN {
tracing::warn!(
"JSON-RPC batch of {} rejected (decision=reject_oversized_batch, limit={})",
requests.len(),
MAX_BATCH_LEN
);
log.warn(format!(
"JSON-RPC batch of {} rejected (limit {})",
requests.len(),
MAX_BATCH_LEN
));
return Ok(build_error_response(
INVALID_REQUEST,
&format!("Batch too large (limit {} requests)", MAX_BATCH_LEN),
None,
None,
));
}
log.debug(format!(
"Processing batch JSON-RPC request with {} items",
requests.len()
));
let mut responses = Vec::new();
for request in requests {
if let Some(response) = process_single_request(
request,
connection_id,
&llm_client,
&app_state,
&status_tx,
&protocol,
server_id,
)
.await
{
responses.push(response);
}
}
if responses.is_empty() {
return Ok(Response::builder()
.status(StatusCode::NO_CONTENT)
.body(Full::new(Bytes::new()))
.unwrap());
}
let response_json = json!(responses);
Ok(Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(response_json.to_string())))
.unwrap())
}
Value::Object(_) => {
if let Some(response) = process_single_request(
request_value,
connection_id,
&llm_client,
&app_state,
&status_tx,
&protocol,
server_id,
)
.await
{
let response_json = response;
Ok(Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(response_json.to_string())))
.unwrap())
} else {
Ok(Response::builder()
.status(StatusCode::NO_CONTENT)
.body(Full::new(Bytes::new()))
.unwrap())
}
}
Value::Array(_) => {
Ok(build_error_response(
INVALID_REQUEST,
"Empty batch request",
None,
None,
))
}
_ => {
Ok(build_error_response(
INVALID_REQUEST,
"Request must be an object or array",
None,
None,
))
}
}
}
async fn process_single_request(
request: Value,
connection_id: ConnectionId,
llm_client: &OllamaClient,
app_state: &Arc<AppState>,
status_tx: &mpsc::UnboundedSender<String>,
protocol: &Arc<JsonRpcProtocol>,
server_id: crate::state::ServerId,
) -> Option<Value> {
let log = Log::new(Some(status_tx));
if !request.is_object() {
return Some(json!({
"jsonrpc": "2.0",
"error": {
"code": INVALID_REQUEST,
"message": "Request must be a JSON object"
},
"id": Value::Null
}));
}
let jsonrpc_version = request.get("jsonrpc").and_then(|v| v.as_str());
let method = request.get("method").and_then(|v| v.as_str());
let params = request.get("params").cloned();
let id = request.get("id").cloned();
let is_notification = id.is_none();
if jsonrpc_version != Some("2.0") {
if !is_notification {
return Some(json!({
"jsonrpc": "2.0",
"error": {
"code": INVALID_REQUEST,
"message": "Invalid JSON-RPC version, must be '2.0'"
},
"id": id
}));
} else {
return None; }
}
let method = match method {
Some(m) if !m.is_empty() => m,
_ => {
if !is_notification {
return Some(json!({
"jsonrpc": "2.0",
"error": {
"code": INVALID_REQUEST,
"message": "Missing or invalid 'method' field"
},
"id": id
}));
} else {
return None; }
}
};
log.debug(format!(
"JSON-RPC method call: method={}, is_notification={}",
method, is_notification
));
let response_value = match call_llm_for_method(
method,
params.as_ref(),
id.clone(),
llm_client,
app_state,
status_tx,
protocol,
connection_id,
server_id,
)
.await
{
Ok(result) => result,
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
let (code, decision) = if failure.is_overloaded() {
(SERVER_BUSY, "fail_closed_llm_overloaded")
} else {
(INTERNAL_ERROR, "fail_closed_llm_error")
};
tracing::error!(
"JSON-RPC {} failed (decision={decision}, code={code}, notification={is_notification}): {e:#}",
method
);
log.warn(format!(
"JSON-RPC {} failed (decision={}): {}",
method, decision, e
));
if !is_notification {
return Some(json!({
"jsonrpc": "2.0",
"error": {
"code": code,
"message": failure.text(),
"data": {"retryable": failure.is_overloaded()}
},
"id": id
}));
} else {
return None; }
}
};
if !is_notification {
Some(response_value)
} else {
log.trace("Notification processed, no response sent");
None
}
}
async fn call_llm_for_method(
method: &str,
params: Option<&Value>,
request_id: Option<Value>,
llm_client: &OllamaClient,
app_state: &Arc<AppState>,
status_tx: &mpsc::UnboundedSender<String>,
protocol: &Arc<JsonRpcProtocol>,
connection_id: ConnectionId,
server_id: crate::state::ServerId,
) -> anyhow::Result<Value> {
let event_data = if let Some(params_val) = params {
json!({
"method": method,
"params": params_val,
"id": request_id,
"is_notification": request_id.is_none(),
})
} else {
json!({
"method": method,
"id": request_id,
"is_notification": request_id.is_none(),
})
};
let event = Event::new(&JSONRPC_METHOD_CALL_EVENT, event_data);
let log = Log::new(Some(status_tx));
log.debug(format!("Calling LLM for JSON-RPC method: {}", method));
let llm_result = call_llm(
llm_client,
app_state,
server_id,
Some(connection_id),
&event,
protocol.as_ref(),
)
.await?;
log.trace(format!(
"LLM actions for JSON-RPC: {:?}",
llm_result.raw_actions.len()
));
let response = llm_result
.protocol_results
.iter()
.find_map(|result| collect_jsonrpc_response(result));
let Some(mut response) = response else {
tracing::warn!(
"JSON-RPC {} answered -32603 (decision=model_no_answer): handler produced neither \
jsonrpc_success nor jsonrpc_error",
method
);
log.warn(format!(
"JSON-RPC {} (decision=model_no_answer): no jsonrpc_success or jsonrpc_error",
method
));
return Ok(json!({
"jsonrpc": "2.0",
"error": {
"code": INTERNAL_ERROR,
"message": "Handler did not produce a jsonrpc_success or jsonrpc_error action"
},
"id": request_id.unwrap_or(Value::Null)
}));
};
if let Some(obj) = response.as_object_mut() {
obj.insert("id".to_string(), request_id.clone().unwrap_or(Value::Null));
obj.insert("jsonrpc".to_string(), Value::String("2.0".to_string()));
}
Ok(response)
}
fn collect_jsonrpc_response(result: &ActionResult) -> Option<Value> {
match result {
ActionResult::Custom { name, data } if name == "jsonrpc_response" => Some(data.clone()),
ActionResult::Multiple(inner) => inner.iter().find_map(collect_jsonrpc_response),
_ => None,
}
}
fn build_error_response(
code: i32,
message: &str,
data: Option<Value>,
id: Option<Value>,
) -> Response<Full<Bytes>> {
let mut error = json!({
"code": code,
"message": message
});
if let Some(data_val) = data {
error
.as_object_mut()
.unwrap()
.insert("data".to_string(), data_val);
}
let response = json!({
"jsonrpc": "2.0",
"error": error,
"id": id.unwrap_or(Value::Null)
});
Response::builder()
.status(StatusCode::OK) .header("Content-Type", "application/json")
.body(Full::new(Bytes::from(response.to_string())))
.unwrap()
}