use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::thread;
use std::time::{Duration, Instant};
use anyhow::{Result, anyhow};
use tiny_http::{Header, Method, Request, Response, Server};
use crate::assets::{self, Delivery};
use crate::graph::CodeGraph;
pub struct ServeOptions {
pub port: u16,
pub open: bool,
pub project: PathBuf,
}
const WORKERS: usize = 4;
const ZERO_GRACE: Duration = Duration::from_millis(2000);
const STARTUP_BACKSTOP: Duration = Duration::from_secs(90);
const KEEPALIVE: Duration = Duration::from_millis(2000);
pub fn serve(graph: &CodeGraph, opts: ServeOptions) -> Result<()> {
let data = serde_json::to_string(graph)?;
let stat = format!(
"{} nodes \u{00b7} {} edges",
graph.nodes.len(),
graph.edges.len()
);
let page = Arc::new(assets::graph_page(&data, &stat, Delivery::Server));
let server = Server::http(("127.0.0.1", opts.port))
.map_err(|e| anyhow!("cannot bind 127.0.0.1:{}: {e}", opts.port))?;
let port = server
.server_addr()
.to_ip()
.map(|a| a.port())
.unwrap_or(opts.port);
let url = format!("http://127.0.0.1:{port}/");
let server = Arc::new(server);
let shutdown = Arc::new(AtomicBool::new(false));
let live = Arc::new(AtomicUsize::new(0));
let project = Arc::new(opts.project);
{
let shutdown = shutdown.clone();
let _ = ctrlc::set_handler(move || shutdown.store(true, Ordering::SeqCst));
}
spawn_reaper(shutdown.clone(), live.clone());
println!("code-rcl graph -> {url}");
println!(
" serving {} nodes / {} edges; the server exits when you close the tab (or press Ctrl-C)",
graph.nodes.len(),
graph.edges.len()
);
if opts.open && webbrowser::open(&url).is_err() {
eprintln!(" (couldn't open a browser automatically — open the URL above)");
}
let mut workers = Vec::with_capacity(WORKERS);
for _ in 0..WORKERS {
let server = server.clone();
let shutdown = shutdown.clone();
let live = live.clone();
let page = page.clone();
let project = project.clone();
workers.push(thread::spawn(move || {
worker_loop(&server, &shutdown, &live, &page, &project);
}));
}
for w in workers {
let _ = w.join();
}
println!("code-rcl serve: stopped.");
Ok(())
}
fn spawn_reaper(shutdown: Arc<AtomicBool>, live: Arc<AtomicUsize>) {
thread::spawn(move || {
let started = Instant::now();
let mut ever_connected = false;
let mut zero_since: Option<Instant> = None;
loop {
thread::sleep(Duration::from_millis(300));
if shutdown.load(Ordering::SeqCst) {
return;
}
let n = live.load(Ordering::SeqCst);
if n > 0 {
ever_connected = true;
zero_since = None;
continue;
}
let expired = if ever_connected {
zero_since.get_or_insert_with(Instant::now).elapsed() > ZERO_GRACE
} else {
started.elapsed() > STARTUP_BACKSTOP
};
if expired {
shutdown.store(true, Ordering::SeqCst);
return;
}
}
});
}
fn worker_loop(
server: &Arc<Server>,
shutdown: &Arc<AtomicBool>,
live: &Arc<AtomicUsize>,
page: &str,
project: &Path,
) {
while !shutdown.load(Ordering::SeqCst) {
let req = match server.recv_timeout(Duration::from_millis(200)) {
Ok(Some(r)) => r,
Ok(None) => continue, Err(_) => break, };
let is_get = *req.method() == Method::Get;
let url = req.url().to_owned();
let path = url.split('?').next().unwrap_or("/").to_owned();
match (is_get, path.as_str()) {
(_, "/dump") => {
let body = run_dump(project, &url);
let _ = req.respond(with_type(
text(
if body.contains("\"ok\":true") {
200
} else {
500
},
&body,
),
"application/json; charset=utf-8",
));
}
(_, "/quit") => {
let _ = req.respond(text(204, ""));
shutdown.store(true, Ordering::SeqCst);
break;
}
(true, "/live") => {
let live = live.clone();
let shutdown = shutdown.clone();
thread::spawn(move || live_stream(req, live, shutdown));
}
(true, "/") | (true, "/index.html") => {
let _ = req.respond(with_type(text(200, page), "text/html; charset=utf-8"));
}
(true, "/assets/d3.min.js") => {
let _ = req.respond(js(assets::D3_JS));
}
(true, "/assets/graph-view.js") => {
let _ = req.respond(js(assets::GRAPH_VIEW_JS.as_str()));
}
(true, "/assets/live.js") => {
let _ = req.respond(js(assets::LIVE_JS));
}
(true, "/assets/graph.css") => {
let _ = req.respond(with_type(
text(200, assets::GRAPH_CSS),
"text/css; charset=utf-8",
));
}
_ => {
let _ = req.respond(text(404, "not found"));
}
}
}
}
fn live_stream(req: Request, live: Arc<AtomicUsize>, shutdown: Arc<AtomicBool>) {
live.fetch_add(1, Ordering::SeqCst);
let _guard = Decrement(&live);
let mut sock = req.into_writer();
let head: &[u8] = b"HTTP/1.1 200 OK\r\n\
Content-Type: text/event-stream\r\n\
Cache-Control: no-store\r\n\
Connection: close\r\n\
\r\n\
: connected\n\n";
if sock.write_all(head).and_then(|_| sock.flush()).is_err() {
return;
}
while !shutdown.load(Ordering::SeqCst) {
let mut waited = Duration::ZERO;
while waited < KEEPALIVE {
if shutdown.load(Ordering::SeqCst) {
return;
}
thread::sleep(Duration::from_millis(100));
waited += Duration::from_millis(100);
}
if sock
.write_all(b": keep-alive\n\n")
.and_then(|_| sock.flush())
.is_err()
{
return; }
}
}
struct Decrement<'a>(&'a Arc<AtomicUsize>);
impl Drop for Decrement<'_> {
fn drop(&mut self) {
self.0.fetch_sub(1, Ordering::SeqCst);
}
}
fn run_dump(project: &Path, url: &str) -> String {
let Some(target) = query_param(url, "target").filter(|t| !t.is_empty()) else {
return r#"{"ok":false,"error":"missing target"}"#.to_string();
};
let depth = query_param(url, "depth")
.and_then(|d| d.parse::<u32>().ok())
.unwrap_or(2)
.clamp(1, 5);
let name = query_param(url, "name")
.map(|s| s.trim().to_string())
.and_then(|s| {
let base = s.rsplit(['/', '\\']).next().unwrap_or("").to_string();
(!base.is_empty() && base != "." && base != ".." && !base.contains(':')).then_some(base)
})
.unwrap_or_else(|| "codebase-context.md".to_string());
let output = project.join(&name);
match crate::commands::dump::relation_bundle(project, &target, depth, 50, false, &output) {
Ok((path, n)) => format!(
r#"{{"ok":true,"path":"{}","files":{}}}"#,
json_escape(&path.display().to_string()),
n
),
Err(e) => format!(
r#"{{"ok":false,"error":"{}"}}"#,
json_escape(&e.to_string())
),
}
}
fn query_param(url: &str, key: &str) -> Option<String> {
let q = url.split_once('?')?.1;
for pair in q.split('&') {
let (k, v) = pair.split_once('=').unwrap_or((pair, ""));
if k == key {
return Some(percent_decode(v));
}
}
None
}
fn percent_decode(s: &str) -> String {
let b = s.as_bytes();
let mut out = Vec::with_capacity(b.len());
let mut i = 0;
while i < b.len() {
match b[i] {
b'+' => out.push(b' '),
b'%' if i + 2 < b.len() => {
let hex = |c: u8| (c as char).to_digit(16);
match (hex(b[i + 1]), hex(b[i + 2])) {
(Some(h), Some(l)) => {
out.push((h * 16 + l) as u8);
i += 2;
}
_ => out.push(b'%'),
}
}
c => out.push(c),
}
i += 1;
}
String::from_utf8_lossy(&out).into_owned()
}
fn json_escape(s: &str) -> String {
let mut o = String::with_capacity(s.len());
for c in s.chars() {
match c {
'"' => o.push_str("\\\""),
'\\' => o.push_str("\\\\"),
'\n' => o.push_str("\\n"),
'\r' => o.push_str("\\r"),
'\t' => o.push_str("\\t"),
c if (c as u32) < 0x20 => o.push_str(&format!("\\u{:04x}", c as u32)),
c => o.push(c),
}
}
o
}
fn header(name: &str, value: &str) -> Header {
Header::from_bytes(name.as_bytes(), value.as_bytes()).expect("static header")
}
fn text(status: u16, body: &str) -> Response<io::Cursor<Vec<u8>>> {
Response::from_string(body).with_status_code(status)
}
fn with_type(mut r: Response<io::Cursor<Vec<u8>>>, ct: &str) -> Response<io::Cursor<Vec<u8>>> {
r.add_header(header("Content-Type", ct));
r
}
fn js(src: &str) -> Response<io::Cursor<Vec<u8>>> {
with_type(text(200, src), "application/javascript; charset=utf-8")
}