pub mod actions;
use std::collections::HashMap;
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::{Request, Response};
use hyper_util::rt::TokioIo;
use tokio::sync::mpsc;
use crate::llm::action_helper::call_llm;
use crate::llm::ollama_client::OllamaClient;
use crate::llm::ActionResult;
use crate::logging::emit::Log;
use crate::protocol::Event;
use crate::server::accept_bounded::{
accept_bounded, watch_idle, ConnectionActivity, ConnectionLimiter,
};
use crate::server::connection::ConnectionId;
use crate::server::MavenProtocol;
use crate::state::app_state::AppState;
use actions::MAVEN_ARTIFACT_REQUEST_EVENT;
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 MavenServer;
impl MavenServer {
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!(
"Maven repository server listening on {}",
local_addr
));
let protocol = Arc::new(MavenProtocol::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,
"MAVEN",
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!(
"Maven 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_maven_request_with_llm(
req,
connection_id,
server_id,
llm_clone,
state_clone,
status_clone,
protocol_clone,
)
.await
}
});
let conn = http1::Builder::new().serve_connection(io, service);
tokio::pin!(conn);
tokio::select! {
result = &mut conn => {
if let Err(err) = result {
Log::new(Some(&status_tx_clone))
.warn(format!("Error serving Maven connection: {:?}", err));
}
}
_ = watch_idle(
Arc::clone(&activity),
IDLE_BETWEEN_REQUESTS_TIMEOUT,
) => {
Log::new(Some(&status_tx_clone)).debug(format!(
"Maven connection {} idle for {}s; closing",
connection_id,
IDLE_BETWEEN_REQUESTS_TIMEOUT.as_secs()
));
}
}
} else {
Log::new(Some(&status_tx_clone)).debug(format!(
"Maven 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!("Maven 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 Maven connection: {}", e));
break;
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
pub const MAX_REQUEST_BODY_BYTES: usize = 8 * 1024 * 1024;
#[derive(Debug)]
struct MavenArtifact {
group_id: String,
artifact_id: String,
version: Option<String>,
classifier: Option<String>,
extension: String,
is_metadata: bool,
is_checksum: bool,
checksum_type: Option<String>, }
impl MavenArtifact {
fn parse(uri: &str) -> Option<Self> {
let path = uri.trim_start_matches('/');
let parts: Vec<&str> = path.split('/').collect();
if parts.is_empty() {
return None;
}
if parts.last() == Some(&"maven-metadata.xml")
|| parts.last() == Some(&"maven-metadata.xml.sha1")
|| parts.last() == Some(&"maven-metadata.xml.md5")
{
let is_checksum = parts.last()?.contains(".sha1") || parts.last()?.contains(".md5");
let checksum_type = if parts.last()?.ends_with(".sha1") {
Some("sha1".to_string())
} else if parts.last()?.ends_with(".md5") {
Some("md5".to_string())
} else {
None
};
if parts.len() < 2 {
return None;
}
let group_id = parts[..parts.len() - 2].join(".");
let artifact_id = parts[parts.len() - 2].to_string();
return Some(Self {
group_id,
artifact_id,
version: None,
classifier: None,
extension: "xml".to_string(),
is_metadata: true,
is_checksum,
checksum_type,
});
}
if parts.len() < 4 {
return None;
}
let filename = parts.last()?;
let version = parts[parts.len() - 2];
let artifact_id = parts[parts.len() - 3];
let group_parts = &parts[..parts.len() - 3];
let group_id = group_parts.join(".");
let (filename, is_checksum, checksum_type) = if filename.ends_with(".sha1") {
(
filename.trim_end_matches(".sha1"),
true,
Some("sha1".to_string()),
)
} else if filename.ends_with(".md5") {
(
filename.trim_end_matches(".md5"),
true,
Some("md5".to_string()),
)
} else if filename.ends_with(".sha256") {
(
filename.trim_end_matches(".sha256"),
true,
Some("sha256".to_string()),
)
} else if filename.ends_with(".sha512") {
(
filename.trim_end_matches(".sha512"),
true,
Some("sha512".to_string()),
)
} else {
(*filename, false, None)
};
let extension = filename.split('.').last()?.to_string();
let base = filename.trim_end_matches(&format!(".{}", extension));
let expected_prefix = format!("{}-{}", artifact_id, version);
if !base.starts_with(&expected_prefix) {
return None;
}
let classifier = if base.len() > expected_prefix.len() {
let remainder = &base[expected_prefix.len()..];
if remainder.starts_with('-') {
Some(remainder[1..].to_string())
} else {
return None; }
} else {
None
};
Some(Self {
group_id,
artifact_id: artifact_id.to_string(),
version: Some(version.to_string()),
classifier,
extension,
is_metadata: false,
is_checksum,
checksum_type,
})
}
}
fn bad_gateway() -> Response<Full<Bytes>> {
Response::builder()
.status(502)
.body(Full::new(Bytes::from(
"Bad Gateway: invalid Maven response",
)))
.expect("502 response with a literal body is always valid")
}
fn payload_too_large() -> Response<Full<Bytes>> {
Response::builder()
.status(413)
.header("Content-Type", "text/plain; charset=utf-8")
.body(Full::new(Bytes::from(format!(
"Payload Too Large: request bodies are limited to {} bytes\n",
MAX_REQUEST_BODY_BYTES
))))
.expect("413 response with a literal body is always valid")
}
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_maven_request_with_llm(
req: Request<Incoming>,
connection_id: ConnectionId,
server_id: crate::state::ServerId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
protocol: Arc<MavenProtocol>,
) -> Result<Response<Full<Bytes>>, Infallible> {
let bytes_received = approximate_request_bytes(&req);
let response = handle_maven_request_with_llm_inner(
req,
connection_id,
server_id,
llm_client,
app_state.clone(),
status_tx,
protocol,
)
.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_maven_request_with_llm_inner(
req: Request<Incoming>,
connection_id: ConnectionId,
server_id: crate::state::ServerId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
protocol: Arc<MavenProtocol>,
) -> Result<Response<Full<Bytes>>, Infallible> {
let method = req.method().to_string();
let uri = req.uri().to_string();
let log = Log::new(Some(&status_tx));
let mut headers = HashMap::new();
for (name, value) in req.headers() {
if let Ok(value_str) = value.to_str() {
headers.insert(name.to_string(), value_str.to_string());
}
}
if let Err(e) = http_body_util::Limited::new(req.into_body(), MAX_REQUEST_BODY_BYTES)
.collect()
.await
{
log.warn(format!(
"Maven {} {} decision=refused_body_too_large (limit {} bytes) -> 413: {}",
method, uri, MAX_REQUEST_BODY_BYTES, e
));
return Ok(payload_too_large());
}
log.debug(format!(
"Maven request: {} {} from {:?}",
method, uri, connection_id
));
let artifact = MavenArtifact::parse(&uri);
if let Some(ref art) = artifact {
log.trace(format!("Parsed Maven artifact: {:?}", art));
}
let event_params = if let Some(art) = artifact {
serde_json::json!({
"method": method,
"uri": uri,
"group_id": art.group_id,
"artifact_id": art.artifact_id,
"version": art.version,
"classifier": art.classifier,
"extension": art.extension,
"is_metadata": art.is_metadata,
"is_checksum": art.is_checksum,
"checksum_type": art.checksum_type,
"headers": headers,
})
} else {
log.debug(format!("Invalid Maven artifact path: {}", uri));
return Ok(Response::builder()
.status(404)
.body(Full::new(Bytes::from("Invalid Maven artifact path")))
.unwrap());
};
let event = Event::new(&MAVEN_ARTIFACT_REQUEST_EVENT, event_params);
match call_llm(
&llm_client,
&app_state,
server_id,
Some(connection_id),
&event,
protocol.as_ref(),
)
.await
{
Ok(execution_result) => {
log.debug("LLM Maven response received");
for msg in execution_result.messages {
let _ = status_tx.send(msg);
}
let mut status_code = 404;
let mut response_headers = HashMap::new();
let mut response_body = Vec::new();
let mut produced_response = false;
for protocol_result in execution_result.protocol_results {
if let ActionResult::Output(output_data) = protocol_result {
if let Ok(json_value) =
serde_json::from_slice::<serde_json::Value>(&output_data)
{
produced_response = true;
if let Some(status) = json_value.get("status").and_then(|v| v.as_u64()) {
status_code = status as u16;
}
if let Some(headers_obj) =
json_value.get("headers").and_then(|v| v.as_object())
{
for (k, v) in headers_obj {
if let Some(v_str) = v.as_str() {
response_headers.insert(k.clone(), v_str.to_string());
}
}
}
if let Some(encoded) =
json_value.get("body_base64").and_then(|v| v.as_str())
{
use base64::Engine;
match base64::engine::general_purpose::STANDARD.decode(encoded) {
Ok(decoded) => response_body = decoded,
Err(e) => {
tracing::error!(
"Maven {} {} decision=fail_closed_bad_body error={}",
method,
uri,
e
);
log.warn(format!(
"Maven {} {} decision=fail_closed_bad_body: response \
body_base64 is not valid base64: {}",
method, uri, e
));
return Ok(bad_gateway());
}
}
} else if let Some(body_str) =
json_value.get("body").and_then(|v| v.as_str())
{
response_body = body_str.as_bytes().to_vec();
}
}
}
}
let decision = if produced_response {
"model_answer"
} else {
"model_no_action"
};
log.debug(format!(
"Maven {} {} decision={} -> {} ({} bytes)",
method,
uri,
decision,
status_code,
response_body.len()
));
let mut response = Response::builder().status(status_code);
for (name, value) in response_headers {
response = response.header(name, value);
}
match response.body(Full::new(Bytes::from(response_body))) {
Ok(resp) => Ok(resp),
Err(e) => {
log.warn(format!("Invalid Maven response ({}), sending 502", e));
Ok(bad_gateway())
}
}
}
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
let decision = if failure.is_overloaded() {
"fail_closed_llm_overloaded"
} else {
"fail_closed_llm_error"
};
tracing::error!("Maven {} {} decision={} error={}", method, uri, decision, e);
log.warn(format!(
"Maven {} {} decision={}: {}",
method, uri, decision, e
));
if failure.is_overloaded() {
return Ok(Response::builder()
.status(503)
.header(hyper::header::RETRY_AFTER, "1")
.body(Full::new(Bytes::from("Service Unavailable")))
.expect("503 response with a literal body is always valid"));
}
Ok(Response::builder()
.status(500)
.body(Full::new(Bytes::from("Internal Server Error")))
.expect("500 response with a literal body is always valid"))
}
}
}