1use std::collections::HashMap;
23use std::sync::Arc;
24
25use parking_lot::RwLock;
26use tokio::sync::{Mutex, mpsc, oneshot};
27use tokio::task::JoinHandle;
28
29use crate::Error;
30use crate::client::ControlDispatchHandle;
31use crate::transport::codec::{DecodedLine, decode_dispatch};
32use crate::transport::process::Subprocess;
33use forge_primitives::Message;
34
35pub(crate) type ControlOutcome = Result<serde_json::Value, Error>;
38
39pub(crate) type PendingControls = Arc<Mutex<HashMap<String, oneshot::Sender<ControlOutcome>>>>;
41
42pub(crate) type InflightDispatches = Arc<Mutex<HashMap<String, JoinHandle<()>>>>;
47
48#[derive(Default)]
53pub(crate) struct SkipCounters {
54 consecutive: u64,
55 total: u64,
56}
57
58const UNPARSABLE_ESCALATE_AFTER: u64 = 50;
61
62#[derive(Debug, PartialEq, Eq)]
64pub(crate) enum LineOutcome {
65 Continue,
67 Skipped,
69 Stop,
71}
72
73pub(crate) fn spawn_reader_task(
80 mut subprocess: Subprocess,
81 dispatch: ControlDispatchHandle,
82 pending_controls: PendingControls,
83 events_tx: mpsc::UnboundedSender<Result<Message, Error>>,
84 pre_init_messages: Vec<Message>,
85 mut shutdown_rx: oneshot::Receiver<()>,
86) -> JoinHandle<()> {
87 use tracing::Instrument;
88 let inflight: InflightDispatches = Arc::new(Mutex::new(HashMap::new()));
89 let span = tracing::info_span!("forge_sdk::control_reader");
90 tokio::spawn(
91 async move {
92 for msg in pre_init_messages {
94 dispatch.capture_session_id_from(&msg);
95 if events_tx.send(Ok(msg)).is_err() {
96 close_subprocess(&mut subprocess).await;
97 drain_pending(&pending_controls).await;
98 return;
99 }
100 }
101
102 let mut line_number: u64 = 0;
103 let mut counters = SkipCounters::default();
104 loop {
105 tokio::select! {
106 biased;
107 _ = &mut shutdown_rx => break,
108 line = subprocess.read_line() => {
109 match line {
110 Ok(Some(line)) => {
111 line_number += 1;
112 if handle_line(
113 &dispatch,
114 &pending_controls,
115 &inflight,
116 &events_tx,
117 line_number,
118 &line,
119 &mut counters,
120 )
121 .await
122 == LineOutcome::Stop
123 {
124 break;
125 }
126 }
127 Ok(None) => break,
128 Err(e) => {
129 let err_text = e.to_string();
130 if events_tx.send(Err(e)).is_err() {
131 tracing::warn!(
132 target: crate::logging::targets::SDK_READER,
133 error = %err_text,
134 "events channel closed; transport error dropped",
135 );
136 }
137 break;
138 }
139 }
140 }
141 }
142 }
143
144 if counters.total > 0 {
145 tracing::warn!(
146 target: crate::logging::targets::SDK_READER,
147 total = counters.total,
148 "stream ended with unparsable lines skipped",
149 );
150 }
151
152 close_subprocess(&mut subprocess).await;
153 drain_pending(&pending_controls).await;
154 }
155 .instrument(span),
156 )
157}
158
159async fn handle_line(
163 dispatch: &ControlDispatchHandle,
164 pending_controls: &PendingControls,
165 inflight: &InflightDispatches,
166 events_tx: &mpsc::UnboundedSender<Result<Message, Error>>,
167 line_number: u64,
168 line: &str,
169 counters: &mut SkipCounters,
170) -> LineOutcome {
171 let outcome = match decode_dispatch(line, line_number) {
172 DecodedLine::Message(msg) => {
173 dispatch.capture_session_id_from(&msg);
174 if events_tx.send(Ok(*msg)).is_ok() { LineOutcome::Continue } else { LineOutcome::Stop }
175 }
176 DecodedLine::Malformed { line: line_no, reason } => {
177 counters.consecutive += 1;
178 counters.total += 1;
179 if let Some(request_id) = control_request_id(line) {
184 let dispatch = dispatch.clone();
185 let reason = reason.clone();
186 tokio::spawn(async move {
187 dispatch.write_error_response(&request_id, &reason).await;
188 });
189 }
190 let raw = excerpt(line, 160);
191 tracing::warn!(
192 target: crate::logging::targets::SDK_READER,
193 line = line_no,
194 reason = %reason,
195 raw = %raw,
196 "unparsable stream-json line - skipping, session continues",
197 );
198 if counters.consecutive == UNPARSABLE_ESCALATE_AFTER {
199 tracing::error!(
200 target: crate::logging::targets::BRIDGE_LIFECYCLE,
201 consecutive = counters.consecutive,
202 "subprocess is not speaking stream-json - surfacing a session failure",
203 );
204 let _ = events_tx.send(Err(Error::Connection {
205 reason: format!(
206 "{} consecutive unparsable stream-json lines; \
207 the subprocess is not speaking stream-json",
208 counters.consecutive
209 ),
210 }));
211 }
212 LineOutcome::Skipped
213 }
214 DecodedLine::Control(req) => {
215 let dispatch_clone = dispatch.clone();
216 let inflight_clone = Arc::clone(inflight);
217 let req = *req;
218 let request_id = req.request_id.clone();
219 let request_id_for_task = request_id.clone();
220 let (gate_tx, gate_rx) = oneshot::channel::<()>();
227 let handle = tokio::spawn(async move {
228 let _ = gate_rx.await;
229 let _ = dispatch_clone.dispatch(req).await;
233 inflight_clone.lock().await.remove(&request_id_for_task);
234 });
235 inflight.lock().await.insert(request_id, handle);
236 let _ = gate_tx.send(());
238 LineOutcome::Continue
239 }
240 DecodedLine::ControlCancel { request_id } => {
241 if let Some(handle) = inflight.lock().await.remove(&request_id) {
242 handle.abort();
243 tracing::debug!(
244 target: crate::logging::targets::SDK_READER,
245 %request_id,
246 "control_cancel_request: dispatch task aborted",
247 );
248 } else {
249 tracing::debug!(
250 target: crate::logging::targets::SDK_READER,
251 %request_id,
252 "control_cancel_request: no in-flight dispatch (already completed)",
253 );
254 }
255 LineOutcome::Continue
256 }
257 DecodedLine::ControlResponse { request_id, raw: value } => {
258 let resp_subtype =
259 value.pointer("/response/subtype").and_then(serde_json::Value::as_str);
260 let outcome = if resp_subtype == Some("success") {
261 Ok(value.pointer("/response/response").cloned().unwrap_or(serde_json::Value::Null))
262 } else {
263 let err = value
264 .pointer("/response/error")
265 .and_then(serde_json::Value::as_str)
266 .map_or_else(
267 || {
268 format!(
269 "no `error` string field; full response: {}",
270 value
271 .pointer("/response")
272 .map_or_else(|| "<missing>".to_string(), ToString::to_string,)
273 )
274 },
275 ToString::to_string,
276 );
277 Err(Error::message_parse(format!("control failed: {err}")))
278 };
279 let mut pending = pending_controls.lock().await;
280 if let Some(tx) = pending.remove(&request_id) {
281 let _ = tx.send(outcome);
282 } else {
283 tracing::warn!(
284 target: crate::logging::targets::SDK_READER,
285 %request_id,
286 "unexpected control_response - dropping",
287 );
288 }
289 LineOutcome::Continue
290 }
291 DecodedLine::Unknown { type_str, raw } => {
292 tracing::warn!(
293 target: crate::logging::targets::SDK_READER,
294 type = %type_str,
295 raw = %raw,
296 line = line_number,
297 "unknown top-level stream-json type - surfacing as Message::Unknown",
298 );
299 if events_tx.send(Ok(Message::Unknown { type_str, raw })).is_ok() {
300 LineOutcome::Continue
301 } else {
302 LineOutcome::Stop
303 }
304 }
305 };
306 if outcome != LineOutcome::Skipped {
308 counters.consecutive = 0;
309 }
310 outcome
311}
312
313pub(crate) fn excerpt(line: &str, max_chars: usize) -> String {
316 line.chars().take(max_chars).collect()
317}
318
319fn control_request_id(line: &str) -> Option<String> {
321 let value: serde_json::Value = serde_json::from_str(line).ok()?;
322 if value.get("type").and_then(serde_json::Value::as_str) != Some("control_request") {
323 return None;
324 }
325 value.pointer("/request_id").and_then(serde_json::Value::as_str).map(str::to_string)
326}
327
328async fn close_subprocess(subprocess: &mut Subprocess) {
329 if let Err(e) = subprocess.close().await {
330 tracing::warn!(
334 target: crate::logging::targets::SDK_READER,
335 error = %e,
336 "reader task: subprocess close error",
337 );
338 }
339}
340
341async fn drain_pending(pending_controls: &PendingControls) {
342 let mut pending = pending_controls.lock().await;
343 for (_, tx) in pending.drain() {
344 let _ = tx.send(Err(Error::Connection {
345 reason: "subprocess closed before control_response".into(),
346 }));
347 }
348}
349
350pub(crate) type SharedSessionId = Arc<RwLock<String>>;
354
355#[cfg(test)]
361pub(crate) fn new_shared_session_id() -> SharedSessionId {
362 Arc::new(RwLock::new(String::new()))
363}
364
365#[cfg(test)]
366mod tests {
367 use super::*;
368 use crate::transport::process::SharedWriter;
369
370 #[tokio::test]
377 async fn a_tool_progress_heartbeat_is_delivered_as_a_message() {
378 let (writer, _lines) = SharedWriter::test_stub();
379 let dispatch = ControlDispatchHandle::new(
380 Arc::new(writer),
381 None,
382 None,
383 crate::mcp::orchestration::McpHosts::new(Vec::new(), HashMap::new()),
384 HashMap::new(),
385 new_shared_session_id(),
386 );
387 let pending: PendingControls = Arc::new(Mutex::new(HashMap::new()));
388 let inflight: InflightDispatches = Arc::new(Mutex::new(HashMap::new()));
389 let (events_tx, mut events_rx) = mpsc::unbounded_channel();
390
391 let line = r#"{"type":"tool_progress","tool_use_id":"toolu_01QhFqNDEgKeskhhiYpzeHnL-heartbeat-0","tool_name":"Bash","parent_tool_use_id":"toolu_01QhFqNDEgKeskhhiYpzeHnL","elapsed_time_seconds":30,"heartbeat":true,"session_id":"s","uuid":"u"}"#;
392 let mut counters = SkipCounters::default();
393 let outcome =
394 handle_line(&dispatch, &pending, &inflight, &events_tx, 40, line, &mut counters).await;
395
396 assert_eq!(outcome, LineOutcome::Continue, "a heartbeat must not end the read loop");
397 let delivered = events_rx
398 .try_recv()
399 .expect("a heartbeat must surface an event")
400 .expect("a heartbeat decodes");
401 let forge_primitives::Message::ToolProgress {
402 tool_use_id,
403 tool_name,
404 parent_tool_use_id,
405 elapsed_time_seconds,
406 heartbeat,
407 ..
408 } = delivered
409 else {
410 panic!("expected ToolProgress, got {delivered:?}");
411 };
412 assert_eq!(tool_use_id, "toolu_01QhFqNDEgKeskhhiYpzeHnL-heartbeat-0");
413 assert_eq!(tool_name, "Bash");
414 assert_eq!(parent_tool_use_id.as_deref(), Some("toolu_01QhFqNDEgKeskhhiYpzeHnL"));
415 assert_eq!(elapsed_time_seconds, 30.0);
416 assert!(heartbeat);
417 }
418
419 #[tokio::test]
424 async fn a_malformed_line_is_skipped_and_the_stream_continues() {
425 let (writer, mut lines) = SharedWriter::test_stub();
426 let dispatch = ControlDispatchHandle::new(
427 Arc::new(writer),
428 None,
429 None,
430 crate::mcp::orchestration::McpHosts::new(Vec::new(), HashMap::new()),
431 HashMap::new(),
432 new_shared_session_id(),
433 );
434 let pending: PendingControls = Arc::new(Mutex::new(HashMap::new()));
435 let inflight: InflightDispatches = Arc::new(Mutex::new(HashMap::new()));
436 let (events_tx, mut events_rx) = mpsc::unbounded_channel();
437
438 let corrupt = r#"{"type":"stream_event""#;
439 let mut counters = SkipCounters::default();
440 let outcome =
441 handle_line(&dispatch, &pending, &inflight, &events_tx, 7, corrupt, &mut counters)
442 .await;
443
444 assert_eq!(outcome, LineOutcome::Skipped, "a decode error must not end the read loop");
445 assert!(events_rx.try_recv().is_err(), "a skipped line must not surface an event");
446 assert!(
447 lines.try_recv().is_err(),
448 "a corrupt non-control line must not be answered on stdin"
449 );
450
451 let valid = r#"{"type":"stream_event","uuid":"evt-1","session_id":"sess-1","event":{"type":"message_start"}}"#;
452 let not_a_request = r#"{"type":"rate_limit_event","request_id":"req_7"}"#;
453 let outcome = handle_line(
454 &dispatch,
455 &pending,
456 &inflight,
457 &events_tx,
458 8,
459 not_a_request,
460 &mut counters,
461 )
462 .await;
463 assert_eq!(outcome, LineOutcome::Skipped);
464 assert!(
466 tokio::time::timeout(std::time::Duration::from_millis(100), lines.recv())
467 .await
468 .is_err(),
469 "only a corrupt control_request may be answered on stdin"
470 );
471
472 let outcome =
473 handle_line(&dispatch, &pending, &inflight, &events_tx, 9, valid, &mut counters).await;
474
475 assert_eq!(
476 outcome,
477 LineOutcome::Continue,
478 "the read loop must still be alive after a skip"
479 );
480 match events_rx.try_recv() {
481 Ok(Ok(Message::StreamEvent { uuid, .. })) => assert_eq!(uuid, "evt-1"),
482 other => panic!("expected the next valid frame on the event stream, got {other:?}"),
483 }
484 }
485
486 #[tokio::test]
491 async fn a_malformed_control_request_is_answered_then_skipped() {
492 let (writer, mut lines) = SharedWriter::test_stub();
493 let dispatch = ControlDispatchHandle::new(
494 Arc::new(writer),
495 None,
496 None,
497 crate::mcp::orchestration::McpHosts::new(Vec::new(), HashMap::new()),
498 HashMap::new(),
499 new_shared_session_id(),
500 );
501 let pending: PendingControls = Arc::new(Mutex::new(HashMap::new()));
502 let inflight: InflightDispatches = Arc::new(Mutex::new(HashMap::new()));
503 let (events_tx, mut events_rx) = mpsc::unbounded_channel();
504
505 let line = r#"{"type":"control_request","request_id":"req_9","request":{"subtype":"can_use_tool","tool_name":"Edit"}}"#;
508 let mut counters = SkipCounters::default();
509 let outcome =
510 handle_line(&dispatch, &pending, &inflight, &events_tx, 7, line, &mut counters).await;
511
512 assert_eq!(
513 outcome,
514 LineOutcome::Skipped,
515 "a malformed control_request must not end the read loop"
516 );
517 assert!(events_rx.try_recv().is_err(), "nothing surfaces on the event stream");
518 let written = tokio::time::timeout(std::time::Duration::from_secs(5), lines.recv())
519 .await
520 .expect("error control_response written to stdin within 5s")
521 .expect("writer channel open");
522 assert!(
523 written.contains(r#""request_id":"req_9""#),
524 "the response must name the request: {written}"
525 );
526 assert!(
527 written.contains(r#""subtype":"error""#),
528 "the response must be an error: {written}"
529 );
530 assert!(
531 written.contains("tool_use_id"),
532 "the error body must carry the real decode reason: {written}"
533 );
534 }
535
536 #[tokio::test]
541 async fn a_degenerate_malformed_stream_escalates_once() {
542 let (writer, _lines) = SharedWriter::test_stub();
543 let dispatch = ControlDispatchHandle::new(
544 Arc::new(writer),
545 None,
546 None,
547 crate::mcp::orchestration::McpHosts::new(Vec::new(), HashMap::new()),
548 HashMap::new(),
549 new_shared_session_id(),
550 );
551 let pending: PendingControls = Arc::new(Mutex::new(HashMap::new()));
552 let inflight: InflightDispatches = Arc::new(Mutex::new(HashMap::new()));
553 let (events_tx, mut events_rx) = mpsc::unbounded_channel();
554 let mut counters = SkipCounters::default();
555 let corrupt = r#"{"type":"stream_event""#;
556
557 for line_no in 1..UNPARSABLE_ESCALATE_AFTER {
558 let outcome = handle_line(
559 &dispatch,
560 &pending,
561 &inflight,
562 &events_tx,
563 line_no,
564 corrupt,
565 &mut counters,
566 )
567 .await;
568 assert_eq!(outcome, LineOutcome::Skipped);
569 assert!(
570 events_rx.try_recv().is_err(),
571 "no escalation below the threshold (line {line_no})"
572 );
573 }
574
575 let outcome = handle_line(
576 &dispatch,
577 &pending,
578 &inflight,
579 &events_tx,
580 UNPARSABLE_ESCALATE_AFTER,
581 corrupt,
582 &mut counters,
583 )
584 .await;
585 assert_eq!(outcome, LineOutcome::Skipped, "the escalation itself must not stop the loop");
586 match events_rx.try_recv() {
587 Ok(Err(_)) => {}
588 other => panic!("expected the session-visible escalation error, got {other:?}"),
589 }
590 assert!(events_rx.try_recv().is_err(), "the escalation must fire exactly once");
591 }
592
593 #[tokio::test]
594 async fn a_good_line_resets_the_consecutive_skip_count() {
595 let (writer, _lines) = SharedWriter::test_stub();
596 let dispatch = ControlDispatchHandle::new(
597 Arc::new(writer),
598 None,
599 None,
600 crate::mcp::orchestration::McpHosts::new(Vec::new(), HashMap::new()),
601 HashMap::new(),
602 new_shared_session_id(),
603 );
604 let pending: PendingControls = Arc::new(Mutex::new(HashMap::new()));
605 let inflight: InflightDispatches = Arc::new(Mutex::new(HashMap::new()));
606 let (events_tx, mut events_rx) = mpsc::unbounded_channel();
607 let mut counters = SkipCounters::default();
608 let corrupt = r#"{"type":"stream_event""#;
609 let valid = r#"{"type":"stream_event","uuid":"evt-1","session_id":"sess-1","event":{"type":"message_start"}}"#;
610
611 let mut line_no = 0;
612 for _ in 0..(UNPARSABLE_ESCALATE_AFTER - 1) {
613 line_no += 1;
614 handle_line(
615 &dispatch,
616 &pending,
617 &inflight,
618 &events_tx,
619 line_no,
620 corrupt,
621 &mut counters,
622 )
623 .await;
624 }
625 line_no += 1;
626 handle_line(&dispatch, &pending, &inflight, &events_tx, line_no, valid, &mut counters)
627 .await;
628 for _ in 0..(UNPARSABLE_ESCALATE_AFTER - 1) {
629 line_no += 1;
630 handle_line(
631 &dispatch,
632 &pending,
633 &inflight,
634 &events_tx,
635 line_no,
636 corrupt,
637 &mut counters,
638 )
639 .await;
640 }
641
642 match events_rx.try_recv() {
643 Ok(Ok(Message::StreamEvent { .. })) => {}
644 other => panic!("expected the valid frame's event, got {other:?}"),
645 }
646 assert!(
647 events_rx.try_recv().is_err(),
648 "a good line between runs of corrupt lines must reset the escalation count"
649 );
650 }
651}