1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
use std::sync::{Arc, Mutex};
use serde_json::Value as JsonValue;
use tokio::sync::mpsc;
use super::types::{ConnectionState, McpOrchestratorService};
use harn_serve::transport::{
read_jsonrpc_stdio_frame, write_jsonrpc_stdio_message, JsonRpcStdioFrameStyle,
};
pub(super) async fn run_stdio(service: Arc<McpOrchestratorService>) -> Result<(), String> {
let mut stdin = tokio::io::BufReader::new(tokio::io::stdin());
let mut session = ConnectionState::default();
let (in_tx, mut in_rx) = mpsc::unbounded_channel();
let input_reader = tokio::spawn(async move {
loop {
match read_jsonrpc_stdio_frame(&mut stdin).await {
Ok(Some(frame)) => {
if in_tx.send(Ok(frame)).is_err() {
break;
}
}
Ok(None) => break,
Err(error) => {
let _ = in_tx.send(Err(error));
break;
}
}
}
});
// Single mpsc fan-in for everything we write to stdout: per-request
// responses, broadcast notifications, and progress updates emitted
// by tool handlers via `harn_vm::mcp_progress`. Funnelling all
// outbound JSON through one writer task means progress lines and
// their final response can never interleave mid-line, and the
// ProgressBus can hand its sender to handle_tools_call without a
// separate channel per request.
let (out_tx, mut out_rx) = mpsc::unbounded_channel::<JsonValue>();
let output_style = Arc::new(Mutex::new(JsonRpcStdioFrameStyle::default()));
let writer_style = output_style.clone();
// Explicit shutdown signal for the writer. We must not couple the
// writer's lifetime to "every `out_tx` clone has been dropped": a
// background task (notably an MCP task-mode tool spawned by
// `create_tool_task`) can hold a progress sender past connection
// shutdown, which would leave `out_rx.recv()` pending forever and wedge
// the awaited `writer` join at EOF — an intermittent, load-dependent
// hang to the nextest slow-test cap (harn#5397). On shutdown we instead
// tell the writer to flush what is already queued and exit.
let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
let writer = tokio::spawn(async move {
let mut stdout = tokio::io::stdout();
loop {
tokio::select! {
message = out_rx.recv() => match message {
Some(message) => {
let style = *writer_style.lock().expect("stdio frame style poisoned");
if write_jsonrpc_stdio_message(&mut stdout, &message, style)
.await
.is_err()
{
break;
}
}
None => break,
},
_ = &mut shutdown_rx => {
// Drain anything already queued so a final response is
// never truncated, then exit without waiting on stray
// sender clones.
while let Ok(message) = out_rx.try_recv() {
let style = *writer_style.lock().expect("stdio frame style poisoned");
if write_jsonrpc_stdio_message(&mut stdout, &message, style)
.await
.is_err()
{
break;
}
}
break;
}
}
}
});
let progress_bus = harn_vm::mcp_progress::ProgressBus::from_mpsc(out_tx.clone());
let _bus_guard = harn_vm::mcp_progress::ActiveBusGuard::install(Some(progress_bus));
eprintln!("[harn] MCP stdio server ready");
loop {
tokio::select! {
frame = in_rx.recv() => {
let Some(frame) = frame else {
break;
};
let frame = frame?;
*output_style.lock().expect("stdio frame style poisoned") = frame.style;
let request: JsonValue = match frame.parse_json() {
Ok(value) => value,
Err(_) => continue,
};
let response = service.handle_request(&mut session, request).await;
if !response.is_null() {
let _ = out_tx.send(response);
}
}
}
}
// Signal the writer to flush and exit. Dropping the two senders we
// hold here (the loop's clone and the ProgressBus's clone in the
// install guard) is not sufficient on its own, because a detached
// background task can still hold a progress sender; the explicit
// shutdown makes teardown independent of that. See the writer's
// construction above for the full rationale (harn#5397).
drop(_bus_guard);
let _ = shutdown_tx.send(());
drop(out_tx);
let _ = writer.await;
input_reader.abort();
let _ = input_reader.await;
Ok(())
}