use std::collections::HashMap;
use std::io::{BufRead, BufReader, Read, Write};
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use hh_core::blob::BlobStore;
use hh_core::event::{AdapterStatus, AgentKind, Event, EventKind, NewSession, SessionStatus};
use hh_core::store::{EventWriter, Store};
use crate::git::GitMeta;
const SPILL_BYTES: usize = 256 * 1024;
#[derive(Debug, Clone)]
pub struct McpProxyOptions {
pub command: Vec<String>,
pub cwd: PathBuf,
pub session_hint: Option<String>,
pub hh_version: String,
pub redactor: Option<Arc<hh_core::redact::Detectors>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct McpProxyOutcome {
pub session_id: String,
pub short_id: String,
pub attached: bool,
pub exit_code: Option<i32>,
}
#[allow(clippy::too_many_lines)] pub fn run_mcp_proxy(store: &Store, opts: &McpProxyOptions) -> crate::Result<McpProxyOutcome> {
let now = now_unix_ms();
let attached_id = std::env::var_os("HH_SESSION_ID").and_then(|v| v.into_string().ok());
let (session_id, short_id, base_started_at, attached) = if let Some(id) = attached_id.as_deref()
{
let resolved = store.resolve_session(id).map_err(|e| {
crate::RecordError::McpSession(format!(
"session `{id}` not found ({e}); run `hh mcp-proxy -- …` standalone, \
or re-run inside `hh run` so the proxy can attach to a live session"
))
})?;
let parent_started = store.session_started_at(&resolved).map_err(|e| {
crate::RecordError::McpSession(format!(
"could not read start time of session `{resolved}`: {e}"
))
})?;
let short = store.session_short_id(&resolved).unwrap_or_default();
(resolved, short, parent_started, true)
} else {
let session_uuid = hh_core::event::now_v7();
let git = GitMeta::capture(&opts.cwd);
let mut command = vec!["mcp-proxy".to_string()];
command.extend_from_slice(&opts.command);
let new_session = NewSession {
id: session_uuid,
started_at: now,
agent_kind: AgentKind::McpOnly,
adapter_status: AdapterStatus::None,
command,
cwd: opts.cwd.clone(),
hostname: hostname(),
hh_version: opts.hh_version.clone(),
model: None,
git_branch: git.branch,
git_sha: git.sha,
git_dirty: git.dirty,
};
let created = store.create_session(&new_session)?;
(created.id, created.short_id, now, false)
};
let writer = Arc::new(Mutex::new(
store
.event_writer_with_redactor(opts.redactor.clone())
.map_err(|e| crate::RecordError::Mcp(format!("open event writer: {e}")))?,
));
let blobs = Arc::new(crate::make_blob_store(store, opts.redactor.clone()));
let pending: Arc<Mutex<HashMap<String, (i64, i64)>>> = Arc::new(Mutex::new(HashMap::new()));
let stop = Arc::new(AtomicBool::new(false));
let mut cmd = Command::new(&opts.command[0]);
cmd.args(&opts.command[1..])
.current_dir(&opts.cwd)
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::inherit());
let mut child = cmd
.spawn()
.map_err(|e| crate::RecordError::Mcp(format!("spawn server `{}`: {e}", opts.command[0])))?;
let server_stdin = child
.stdin
.take()
.ok_or_else(|| crate::RecordError::Mcp("server stdin not piped".into()))?;
let server_stdout = child
.stdout
.take()
.ok_or_else(|| crate::RecordError::Mcp("server stdout not piped".into()))?;
let ctx = ProxyCtx {
writer: Arc::clone(&writer),
blobs: Arc::clone(&blobs),
pending: Arc::clone(&pending),
session_id: session_id.clone(),
base_started_at,
};
let up_stop = Arc::clone(&stop);
let upstream = std::thread::Builder::new()
.name("hh-mcp-upstream".into())
.spawn({
let ctx = ctx.clone();
move || run_upstream(server_stdin, ctx, up_stop)
})
.map_err(|e| crate::RecordError::Mcp(format!("spawn upstream thread: {e}")))?;
let down_stop = Arc::clone(&stop);
let mut stdout = std::io::stdout();
let downstream = std::thread::Builder::new()
.name("hh-mcp-downstream".into())
.spawn({
let ctx = ctx.clone();
move || run_downstream(server_stdout, &mut stdout, ctx, down_stop)
})
.map_err(|e| crate::RecordError::Mcp(format!("spawn downstream thread: {e}")))?;
let _ = upstream.join();
let _ = downstream.join();
stop.store(true, Ordering::Release);
let exit_code = wait_for_child(&mut child);
{
let mut w = writer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
w.flush()
.map_err(|e| crate::RecordError::Mcp(format!("flush writer: {e}")))?;
w.close()
.map_err(|e| crate::RecordError::Mcp(format!("close writer: {e}")))?;
}
if attached {
let _ = store.assign_steps(&session_id);
} else {
let ended_at = now_unix_ms();
let status = match exit_code {
Some(0) => SessionStatus::Ok,
Some(_) => SessionStatus::Error,
None => SessionStatus::Interrupted,
};
store
.assign_steps(&session_id)
.map_err(|e| crate::RecordError::Mcp(format!("assign steps: {e}")))?;
store
.finalize_session(&session_id, ended_at, exit_code, status)
.map_err(|e| crate::RecordError::Mcp(format!("finalize session: {e}")))?;
}
Ok(McpProxyOutcome {
session_id,
short_id,
attached,
exit_code,
})
}
#[derive(Clone)]
struct ProxyCtx {
writer: Arc<Mutex<EventWriter>>,
blobs: Arc<BlobStore>,
pending: Arc<Mutex<HashMap<String, (i64, i64)>>>,
session_id: String,
base_started_at: i64,
}
impl ProxyCtx {
fn append(&self, event: Event) -> Option<i64> {
let w = self
.writer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
w.append_event(event).ok()
}
fn take_pending(&self, key: &str) -> Option<(i64, i64)> {
self.pending.lock().ok()?.get(key).copied()
}
fn register_pending(&self, key: String, event_id: i64, ts_abs: i64) {
if let Ok(mut m) = self.pending.lock() {
m.insert(key, (event_id, ts_abs));
}
}
fn event(
&self,
kind: EventKind,
ts_ms: i64,
summary: String,
body: serde_json::Value,
) -> Event {
let (body_json, blob_hash, blob_size) = maybe_spill(body, &self.blobs);
Event {
session_id: self.session_id.clone(),
ts_ms,
kind,
step: None,
summary,
body_json,
blob_hash,
blob_size,
correlates: None,
}
}
}
#[allow(clippy::needless_pass_by_value)] fn run_upstream(mut server_stdin: impl Write + Send, ctx: ProxyCtx, stop: Arc<AtomicBool>) {
let stdin = std::io::stdin();
let mut reader = BufReader::new(stdin.lock());
let mut line = Vec::with_capacity(8 * 1024);
while !stop.load(Ordering::Acquire) {
line.clear();
let n = match reader.read_until(b'\n', &mut line) {
Ok(0) => break, Ok(n) => n,
Err(e) => {
if e.kind() != std::io::ErrorKind::Interrupted {
break;
}
continue;
}
};
record_upstream_line(line.as_slice(), &ctx);
if server_stdin.write_all(&line[..n]).is_err() {
break; }
let _ = server_stdin.flush();
}
let _ = server_stdin.flush();
}
#[allow(clippy::needless_pass_by_value)] fn run_downstream(
server_stdout: impl Read + Send,
stdout: &mut (impl Write + Send),
ctx: ProxyCtx,
stop: Arc<AtomicBool>,
) {
let mut reader = BufReader::new(server_stdout);
let mut line = Vec::with_capacity(8 * 1024);
while !stop.load(Ordering::Acquire) {
line.clear();
let n = match reader.read_until(b'\n', &mut line) {
Ok(0) => break, Ok(n) => n,
Err(e) => {
if e.kind() != std::io::ErrorKind::Interrupted {
break;
}
continue;
}
};
record_downstream_line(line.as_slice(), &ctx);
if stdout.write_all(&line[..n]).is_err() {
break; }
let _ = stdout.flush();
}
}
fn record_upstream_line(line: &[u8], ctx: &ProxyCtx) {
let Some(parsed) = parse_json_line(line) else {
return;
};
let now_abs = now_unix_ms();
let ts_ms = now_abs - ctx.base_started_at;
match &parsed {
serde_json::Value::Object(obj) => {
let id = obj.get("id").cloned();
let method = obj
.get("method")
.and_then(|v| v.as_str())
.map(str::to_string);
match (id.as_ref(), method.as_deref()) {
(Some(id_val), Some(method)) => {
let key = id_key(id_val);
let summary = truncate_summary(&format!("mcp request: {method}"));
let body = serde_json::to_value(&parsed).unwrap_or_else(|_| parsed.clone());
let mut event = ctx.event(EventKind::McpRequest, ts_ms, summary, body);
if let Some(k) = &key {
if let Some(b) = event.body_json.as_mut() {
if let Some(o) = b.as_object_mut() {
o.insert("correlate_key".into(), serde_json::json!(k));
}
}
}
let event_id = ctx.append(event);
if let (Some(k), Some(eid)) = (key, event_id) {
ctx.register_pending(k, eid, now_abs);
}
}
(None, Some(method)) => record_notification(ctx, ts_ms, method, parsed),
_ => {}
}
}
serde_json::Value::Array(arr) => record_batch(ctx, ts_ms, arr.len()),
_ => {}
}
}
fn record_downstream_line(line: &[u8], ctx: &ProxyCtx) {
let Some(parsed) = parse_json_line(line) else {
return;
};
let ts_ms = now_unix_ms() - ctx.base_started_at;
match &parsed {
serde_json::Value::Object(obj) => {
let id = obj.get("id").cloned();
let method = obj
.get("method")
.and_then(|v| v.as_str())
.map(str::to_string);
match (id.as_ref(), method.as_deref()) {
(Some(id_val), None) => record_response(ctx, ts_ms, id_val, &parsed),
(Some(_), Some(method)) => record_server_request(ctx, ts_ms, method, parsed),
(None, Some(method)) => record_notification(ctx, ts_ms, method, parsed),
_ => {}
}
}
serde_json::Value::Array(arr) => record_batch(ctx, ts_ms, arr.len()),
_ => {}
}
}
fn record_response(
ctx: &ProxyCtx,
ts_ms: i64,
id_val: &serde_json::Value,
parsed: &serde_json::Value,
) {
let key = id_key(id_val);
let (correlates, request_ts) = match &key {
Some(k) => ctx
.take_pending(k)
.map_or((None, None), |(eid, ts)| (Some(eid), Some(ts))),
None => (None, None),
};
let latency_ms = request_ts.map(|t| now_unix_ms() - t);
let id_disp = id_key(id_val).unwrap_or_else(|| "?".into());
let summary = truncate_summary(&format!("mcp response: id {id_disp}"));
let mut body = serde_json::to_value(parsed).unwrap_or_else(|_| parsed.clone());
if let Some(lat) = latency_ms {
if let Some(b) = body.as_object_mut() {
b.insert("latency_ms".into(), serde_json::json!(lat));
}
}
let mut event = ctx.event(EventKind::McpResponse, ts_ms, summary, body);
event.correlates = correlates;
ctx.append(event);
}
fn record_server_request(ctx: &ProxyCtx, ts_ms: i64, method: &str, parsed: serde_json::Value) {
let summary = truncate_summary(&format!("mcp request: {method}"));
let body = serde_json::to_value(&parsed).unwrap_or(parsed);
ctx.append(ctx.event(EventKind::McpRequest, ts_ms, summary, body));
}
fn record_notification(ctx: &ProxyCtx, ts_ms: i64, method: &str, parsed: serde_json::Value) {
let summary = truncate_summary(&format!("mcp notification: {method}"));
let body = serde_json::to_value(&parsed).unwrap_or(parsed);
ctx.append(ctx.event(EventKind::McpNotification, ts_ms, summary, body));
}
fn record_batch(ctx: &ProxyCtx, ts_ms: i64, len: usize) {
let summary = truncate_summary(&format!("mcp batch: {len} messages"));
let body = serde_json::json!({ "batch": true, "count": len });
ctx.append(ctx.event(EventKind::McpNotification, ts_ms, summary, body));
}
fn parse_json_line(line: &[u8]) -> Option<serde_json::Value> {
let trimmed = line.strip_suffix(b"\n").unwrap_or(line);
let trimmed = trimmed.strip_suffix(b"\r").unwrap_or(trimmed);
serde_json::from_slice(trimmed).ok()
}
fn id_key(v: &serde_json::Value) -> Option<String> {
match v {
serde_json::Value::Number(n) => Some(n.to_string()),
serde_json::Value::String(s) => Some(s.clone()),
_ => None,
}
}
fn maybe_spill(
body: serde_json::Value,
blobs: &BlobStore,
) -> (Option<serde_json::Value>, Option<String>, Option<u64>) {
let serialized = serde_json::to_vec(&body).unwrap_or_default();
if serialized.len() >= SPILL_BYTES {
if let Ok(outcome) = blobs.put(&serialized) {
let envelope = serde_json::json!({
"overflow": true,
"size": outcome.size,
"blob_hash": outcome.hash,
"encoding": "blob",
});
return (Some(envelope), Some(outcome.hash), Some(outcome.size));
}
}
(Some(body), None, None)
}
fn truncate_summary(s: &str) -> String {
const LIMIT: usize = 120;
if s.chars().count() <= LIMIT {
return s.to_string();
}
let truncated: String = s.chars().take(LIMIT - 1).collect();
format!("{truncated}…")
}
fn now_unix_ms() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX))
}
fn hostname() -> Option<String> {
std::process::Command::new("hostname")
.output()
.ok()
.filter(|o| o.status.success())
.and_then(|o| String::from_utf8(o.stdout).ok())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
}
fn wait_for_child(child: &mut Child) -> Option<i32> {
match child.wait() {
Ok(status) => status.code(),
Err(_) => None,
}
}
#[cfg(feature = "fuzzing")]
pub mod fuzzing {
use super::{record_downstream_line, record_upstream_line, ProxyCtx};
use hh_core::blob::BlobStore;
use hh_core::event::{AdapterStatus, AgentKind, NewSession};
use hh_core::store::Store;
use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};
fn ctx() -> &'static ProxyCtx {
static CTX: OnceLock<ProxyCtx> = OnceLock::new();
CTX.get_or_init(|| {
let dir = std::env::temp_dir().join(format!("hh-fuzz-mcp-{}", std::process::id()));
let store =
Store::open(&dir.join("hh.db"), &dir.join("blobs")).expect("open fuzz store");
let new_session = NewSession {
id: hh_core::event::now_v7(),
started_at: 0,
agent_kind: AgentKind::McpOnly,
adapter_status: AdapterStatus::None,
command: vec!["fuzz".to_string()],
cwd: std::env::temp_dir(),
hostname: None,
hh_version: "fuzz".to_string(),
model: None,
git_branch: None,
git_sha: None,
git_dirty: None,
};
let created = store
.create_session(&new_session)
.expect("create fuzz session");
let writer = Arc::new(Mutex::new(store.event_writer().expect("fuzz writer")));
let blobs = Arc::new(BlobStore::new(store.blobs().root().to_path_buf()));
ProxyCtx {
writer,
blobs,
pending: Arc::new(Mutex::new(HashMap::new())),
session_id: created.id,
base_started_at: 0,
}
})
}
pub fn fuzz_upstream_line(line: &[u8]) {
record_upstream_line(line, ctx());
}
pub fn fuzz_downstream_line(line: &[u8]) {
record_downstream_line(line, ctx());
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn id_key_handles_string_number_null() {
assert_eq!(id_key(&serde_json::json!("abc")), Some("abc".into()));
assert_eq!(id_key(&serde_json::json!(42)), Some("42".into()));
assert_eq!(id_key(&serde_json::json!(null)), None);
assert_eq!(id_key(&serde_json::json!(true)), None);
}
#[test]
fn parse_json_line_handles_trailing_newline_and_cr() {
assert!(parse_json_line(b"{\"a\":1}\n").is_some());
assert!(parse_json_line(b"{\"a\":1}\r\n").is_some());
assert!(parse_json_line(b"not json\n").is_none());
assert!(parse_json_line(b"").is_none());
}
#[test]
fn maybe_spill_inline_below_threshold_and_spills_above() {
let dir = std::env::temp_dir().join(format!("hh-mcp-spill-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let blobs = BlobStore::new(dir.clone());
let small = serde_json::json!({"method": "tools/call"});
let (body, hash, size) = maybe_spill(small.clone(), &blobs);
assert_eq!(body.as_ref().unwrap(), &small);
assert!(hash.is_none() && size.is_none());
let big = serde_json::json!({"data": "x".repeat(SPILL_BYTES)});
let (body, hash, size) = maybe_spill(big, &blobs);
let envelope = body.expect("large body must spill");
assert_eq!(envelope["overflow"], serde_json::json!(true));
assert_eq!(envelope["encoding"], serde_json::json!("blob"));
let hash = hash.expect("blob hash set");
assert!(size.unwrap() >= SPILL_BYTES as u64);
let stored = blobs.get(&hash).unwrap();
assert!(stored.len() >= SPILL_BYTES);
std::fs::remove_dir_all(&dir).ok();
}
}