pub mod actions;
use std::convert::Infallible;
use std::net::SocketAddr;
use std::sync::Arc;
use base64::Engine;
use bytes::Bytes;
use http_body_util::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;
use tokio::sync::mpsc;
use tracing::{error, trace};
use crate::llm::action_helper::call_llm;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::EventType;
use crate::server::accept_bounded::{
accept_bounded, watch_idle, ConnectionActivity, ConnectionLimiter,
};
use crate::server::connection::ConnectionId;
use crate::server::npm::actions::NpmProtocol;
use crate::state::app_state::AppState;
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(300);
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 NpmServer;
impl NpmServer {
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,
) -> 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!("NPM registry server listening on {}", local_addr));
let protocol = Arc::new(NpmProtocol::new());
let task_registrar = app_state.clone();
let limiter = ConnectionLimiter::new(MAX_CONNECTIONS);
let accept_handle = tokio::spawn(async move {
loop {
match accept_bounded(
&listener,
&limiter,
CONNECTION_CAP_REFUSAL,
"NPM",
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!(
"NPM 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 {
let io = TokioIo::new(stream);
let status_for_service = status_tx_clone.clone();
let app_state_for_service = app_state_clone.clone();
let activity = Arc::new(ConnectionActivity::new());
let activity_for_service = 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 = Arc::clone(&activity_for_service);
async move {
let _busy = activity.busy();
handle_npm_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 NPM connection: {:?}", err);
}
}
_ = watch_idle(
Arc::clone(&activity),
IDLE_BETWEEN_REQUESTS_TIMEOUT,
) => {
Log::new(Some(&status_tx_clone)).debug(format!(
"NPM connection {} idle for {}s; closing",
connection_id,
IDLE_BETWEEN_REQUESTS_TIMEOUT.as_secs()
));
}
}
} else {
Log::new(Some(&status_tx_clone)).debug(format!(
"NPM peer {} sent no request within {}s; closing \
decision=fail_closed_first_byte_timeout",
remote_addr,
FIRST_BYTE_READ_TIMEOUT.as_secs()
));
}
app_state_clone
.close_connection_on_server(server_id, connection_id)
.await;
Log::new(Some(&status_tx_clone))
.info(format!("NPM 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 NPM connection: {}", e));
break;
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
fn approximate_request_bytes(req: &Request<Incoming>) -> u64 {
use hyper::body::Body;
let head: usize = req.method().as_str().len()
+ req.uri().to_string().len()
+ 12 + req
.headers()
.iter()
.map(|(name, value)| name.as_str().len() + value.len() + 4)
.sum::<usize>();
head as u64 + req.body().size_hint().lower()
}
#[allow(clippy::too_many_arguments)]
async fn handle_npm_request(
req: Request<Incoming>,
connection_id: ConnectionId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
protocol: Arc<NpmProtocol>,
server_id: crate::state::ServerId,
) -> Result<Response<Full<Bytes>>, Infallible> {
let bytes_received = approximate_request_bytes(&req);
let response = handle_npm_request_inner(
req,
connection_id,
llm_client,
app_state.clone(),
status_tx,
protocol,
server_id,
)
.await;
let bytes_sent = response
.as_ref()
.ok()
.and_then(|resp| {
use hyper::body::Body;
resp.body().size_hint().exact()
})
.unwrap_or(0);
app_state
.update_connection_stats(
server_id,
connection_id,
Some(bytes_received),
Some(bytes_sent),
Some(1),
Some(1),
)
.await;
response
}
async fn handle_npm_request_inner(
req: Request<Incoming>,
_connection_id: ConnectionId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
protocol: Arc<NpmProtocol>,
server_id: crate::state::ServerId,
) -> Result<Response<Full<Bytes>>, Infallible> {
let method = req.method().clone();
let uri = req.uri().clone();
let path = uri.path();
let query = uri.query().unwrap_or("");
Log::new(Some(&status_tx)).debug(format!("NPM request: {} {}", method, path));
if method != Method::GET {
Log::new(Some(&status_tx)).warn(format!(
"NPM {} {} decision=protocol_error (only GET is served)",
method, path
));
let response = json!({
"error": "Method not allowed"
});
return Ok(Response::builder()
.status(StatusCode::METHOD_NOT_ALLOWED)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(response.to_string())))
.unwrap());
}
let (event_type, description) = if path == "/-/all" {
("NPM_LIST_REQUEST", "NPM package list request".to_string())
} else if path.starts_with("/-/v1/search") {
(
"NPM_SEARCH_REQUEST",
format!("NPM package search: {}", query),
)
} else if path.contains("/-/") {
let parts: Vec<&str> = path.split("/-/").collect();
let package_name = parts.get(0).unwrap_or(&"").trim_start_matches('/');
let tarball_name = parts.get(1).unwrap_or(&"");
(
"NPM_TARBALL_REQUEST",
format!(
"NPM tarball request: package={}, tarball={}",
package_name, tarball_name
),
)
} else {
let package_name = path.trim_start_matches('/');
(
"NPM_PACKAGE_REQUEST",
format!("NPM package metadata request: {}", package_name),
)
};
trace!("NPM event: {}: {}", event_type, &description);
if app_state.get_instruction(server_id).await.is_none() {
Log::new(Some(&status_tx)).error(format!(
"NPM {} {} decision=fail_closed_server_missing (server {} not found)",
method, path, server_id
));
let response = json!({
"error": "Server not found"
});
return Ok(Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(response.to_string())))
.unwrap());
}
let event_type_static: &'static EventType = match &event_type[..] {
"NPM_PACKAGE_REQUEST" => &actions::NPM_PACKAGE_REQUEST,
"NPM_TARBALL_REQUEST" => &actions::NPM_TARBALL_REQUEST,
"NPM_LIST_REQUEST" => &actions::NPM_LIST_REQUEST,
"NPM_SEARCH_REQUEST" => &actions::NPM_SEARCH_REQUEST,
_ => {
error!(
"Unknown NPM event type: {} decision=protocol_error",
event_type
);
let error_response = json!({
"error": format!("Internal error: unknown event type '{}'", event_type)
});
return Ok(Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(error_response.to_string())))
.unwrap());
}
};
let event = crate::protocol::Event::new(
event_type_static,
json!({
"method": method.as_str(),
"path": path,
"query": query,
"description": description,
}),
);
Log::new(Some(&status_tx)).debug(format!("Calling LLM for NPM request: {} {}", method, path));
let llm_result = call_llm(
&llm_client,
&app_state,
server_id,
None,
&event,
protocol.as_ref(),
)
.await;
match llm_result {
Ok(execution_result) => {
let failures = execution_result.failures.len();
for result in execution_result.protocol_results {
if let Some(response) = process_npm_action_result(result, &status_tx).await {
let decision = if response.status().is_client_error()
|| response.status().is_server_error()
{
"model_reject"
} else {
"model_answer"
};
Log::new(Some(&status_tx)).info(format!(
"NPM {} {} -> {} decision={}",
method,
path,
response.status().as_u16(),
decision
));
return Ok(response);
}
}
let msg = format!(
"NPM {} {} -> 500 decision={} ({} failed action(s))",
method,
path,
if failures == 0 {
"model_silent"
} else {
"fail_closed_bad_action"
},
failures
);
if failures == 0 {
Log::new(Some(&status_tx)).warn(msg);
} else {
Log::new(Some(&status_tx)).error(msg);
}
let error_response = json!({
"error": "No NPM action returned"
});
Ok(Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(error_response.to_string())))
.unwrap())
}
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
Log::new(Some(&status_tx)).error(format!(
"NPM {} {} -> 500 decision={}: {}",
method,
path,
if failure.is_overloaded() {
"fail_closed_llm_overloaded"
} else {
"fail_closed_llm_error"
},
e
));
let error_response = json!({
"error": failure.text()
});
Ok(Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(error_response.to_string())))
.unwrap())
}
}
}
fn npm_server_error(
status_tx: &mpsc::UnboundedSender<String>,
reason: &str,
) -> Response<Full<Bytes>> {
Log::new(Some(status_tx)).error(format!(
"NPM -> 500 decision=fail_closed_bad_action: could not build a response — {}",
reason
));
Response::builder()
.status(StatusCode::INTERNAL_SERVER_ERROR)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(
json!({ "error": crate::utils::WireFailure::Unavailable.prefixed_text() }).to_string(),
)))
.unwrap()
}
async fn process_npm_action_result(
action_result: crate::llm::ActionResult,
status_tx: &mpsc::UnboundedSender<String>,
) -> Option<Response<Full<Bytes>>> {
use crate::llm::ActionResult;
match action_result {
ActionResult::Custom { name, data } => {
match name.as_str() {
"npm_package_metadata" => {
let Some(metadata) = data.get("metadata") else {
return Some(npm_server_error(
status_tx,
"npm_package_metadata carried no 'metadata' field",
));
};
Log::new(Some(status_tx)).debug("NPM package metadata response");
Some(
Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(metadata.to_string())))
.unwrap(),
)
}
"npm_package_tarball" => {
let Some(tarball_data) = data.get("tarball_data").and_then(|v| v.as_str())
else {
return Some(npm_server_error(
status_tx,
"npm_package_tarball carried no 'tarball_data' string",
));
};
let decoded =
match base64::engine::general_purpose::STANDARD.decode(tarball_data) {
Ok(bytes) => bytes,
Err(e) => {
return Some(npm_server_error(
status_tx,
&format!(
"tarball_data is not valid base64 ({}). A tarball is \
binary, so base64 is the only faithful form; send \
the whole encoded string, never an abbreviation \
ending in \"...\"",
e
),
));
}
};
if decoded.is_empty() {
return Some(npm_server_error(
status_tx,
"tarball_data decoded to zero bytes; an empty .tgz is not a \
package npm can install",
));
}
Log::new(Some(status_tx)).debug(format!(
"NPM package tarball response: {} bytes",
decoded.len()
));
Some(
Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "application/octet-stream")
.body(Full::new(Bytes::from(decoded)))
.unwrap(),
)
}
"npm_package_list" => {
let Some(packages) = data.get("packages") else {
return Some(npm_server_error(
status_tx,
"npm_package_list carried no 'packages' field",
));
};
Log::new(Some(status_tx)).debug("NPM package list response");
Some(
Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(packages.to_string())))
.unwrap(),
)
}
"npm_package_search" => {
let Some(results) = data.get("results") else {
return Some(npm_server_error(
status_tx,
"npm_package_search carried no 'results' field",
));
};
Log::new(Some(status_tx)).debug("NPM package search response");
Some(
Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(results.to_string())))
.unwrap(),
)
}
"npm_error" => {
let error_message = data
.get("error")
.and_then(|v| v.as_str())
.unwrap_or("Unknown error");
let status_code = data
.get("status_code")
.and_then(|v| v.as_u64())
.unwrap_or(500) as u16;
Log::new(Some(status_tx))
.debug(format!("NPM error: {} ({})", error_message, status_code));
let error_response = json!({
"error": error_message
});
Some(
Response::builder()
.status(
StatusCode::from_u16(status_code)
.unwrap_or(StatusCode::INTERNAL_SERVER_ERROR),
)
.header("Content-Type", "application/json")
.body(Full::new(Bytes::from(error_response.to_string())))
.unwrap(),
)
}
_ => {
error!("Unknown NPM action: {}", name);
None
}
}
}
_ => {
error!("Unexpected action result type for NPM request");
None
}
}
}