Centralize stdin control semantics

Author: t-kalinowski

src/main.rs

@@ -23,6 +23,7 @@ mod r_session;
 mod sandbox;
 mod sandbox_cli;
 mod server;
+mod stdin_payload;
 #[cfg(target_os = "windows")]
 mod windows_sandbox;
 mod worker;

src/server/response.rs

@@ -11,6 +11,7 @@ use rmcp::model::{
 use serde_json::Value;
 use tempfile::Builder;
 
+pub(crate) use crate::stdin_payload::{TimeoutBundleReuse, timeout_bundle_reuse_for_input};
 use crate::worker_process::WorkerError;
 use crate::worker_protocol::{
     ContentOrigin, TextStream, WorkerContent, WorkerErrorCode, WorkerReply,
@@ -170,40 +171,6 @@ struct TimeoutReplyView<'a> {
     protected_bundle_id: Option<u64>,
 }
 
-#[derive(Clone, Copy)]
-pub(crate) enum TimeoutBundleReuse {
-    None,
-    FullReply,
-    FollowUpInput,
-}
-
-pub(crate) fn timeout_bundle_reuse_for_input(input: &str) -> TimeoutBundleReuse {
-    if input.is_empty() {
-        return TimeoutBundleReuse::FullReply;
-    }
-
-    let Some(first) = input.chars().next() else {
-        return TimeoutBundleReuse::FullReply;
-    };
-    let tail = &input[first.len_utf8()..];
-    let tail = if let Some(rest) = tail.strip_prefix("\r\n") {
-        rest
-    } else if let Some(rest) = tail.strip_prefix('\n') {
-        rest
-    } else if let Some(rest) = tail.strip_prefix('\r') {
-        rest
-    } else {
-        tail
-    };
-
-    match first {
-        '\u{3}' if tail.is_empty() => TimeoutBundleReuse::FullReply,
-        '\u{3}' => TimeoutBundleReuse::FollowUpInput,
-        '\u{4}' => TimeoutBundleReuse::None,
-        _ => TimeoutBundleReuse::FollowUpInput,
-    }
-}
-
 impl ResponseState {
     pub(crate) fn new() -> Result<Self, WorkerError> {
         Ok(Self {

src/server/tests.rs

@@ -195,6 +195,22 @@ fn timeout_bundle_reuse_treats_blank_lines_as_fresh_input() {
         super::response::timeout_bundle_reuse_for_input("\r\n"),
         super::response::TimeoutBundleReuse::FollowUpInput
     ));
+    assert!(matches!(
+        super::response::timeout_bundle_reuse_for_input("\u{3}"),
+        super::response::TimeoutBundleReuse::FullReply
+    ));
+    assert!(matches!(
+        super::response::timeout_bundle_reuse_for_input("\u{3}\n"),
+        super::response::TimeoutBundleReuse::FollowUpInput
+    ));
+    assert!(matches!(
+        super::response::timeout_bundle_reuse_for_input("\u{3}\r\n"),
+        super::response::TimeoutBundleReuse::FollowUpInput
+    ));
+    assert!(matches!(
+        super::response::timeout_bundle_reuse_for_input("\u{4}"),
+        super::response::TimeoutBundleReuse::None
+    ));
 }
 
 fn split_lines(text: &str) -> Vec<String> {

src/stdin_payload.rs

@@ -0,0 +1,45 @@
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum WriteStdinControlAction {
+    Interrupt,
+    Restart,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum TimeoutBundleReuse {
+    None,
+    FullReply,
+    FollowUpInput,
+}
+
+pub(crate) fn prepare_worker_stdin_payload(input: &str) -> Vec<u8> {
+    let mut payload = input.as_bytes().to_vec();
+    if !payload.is_empty() && !payload.ends_with(b"\n") {
+        payload.push(b'\n');
+    }
+    payload
+}
+
+pub(crate) fn split_write_stdin_control_prefix(
+    input: &str,
+) -> Option<(WriteStdinControlAction, &str)> {
+    let first = input.chars().next()?;
+    let action = match first {
+        '\u{3}' => WriteStdinControlAction::Interrupt,
+        '\u{4}' => WriteStdinControlAction::Restart,
+        _ => return None,
+    };
+    Some((action, &input[first.len_utf8()..]))
+}
+
+pub(crate) fn timeout_bundle_reuse_for_input(input: &str) -> TimeoutBundleReuse {
+    if input.is_empty() {
+        return TimeoutBundleReuse::FullReply;
+    }
+
+    match split_write_stdin_control_prefix(input) {
+        Some((WriteStdinControlAction::Interrupt, "")) => TimeoutBundleReuse::FullReply,
+        Some((WriteStdinControlAction::Interrupt, _)) => TimeoutBundleReuse::FollowUpInput,
+        Some((WriteStdinControlAction::Restart, _)) => TimeoutBundleReuse::None,
+        None => TimeoutBundleReuse::FollowUpInput,
+    }
+}

src/worker_process.rs

@@ -51,6 +51,8 @@ use crate::sandbox_cli::{
     resolve_effective_sandbox_state_with_defaults, sandbox_plan_requests_inherited_state,
     validate_sandbox_plan_with_defaults,
 };
+use crate::stdin_payload::prepare_worker_stdin_payload;
+pub(crate) use crate::stdin_payload::{WriteStdinControlAction, split_write_stdin_control_prefix};
 use crate::worker_protocol::{
     ContentOrigin, TextStream, WORKER_MODE_ARG, WorkerContent, WorkerErrorCode, WorkerReply,
 };
@@ -275,7 +277,14 @@ fn prechecked_follow_up_requires_meta_error() -> WorkerError {
 }
 
 trait BackendDriver: Send {
-    fn prepare_input_payload(&self, text: &str) -> Vec<u8>;
+    fn prepare_input_text(&self, text: String) -> String {
+        text
+    }
+
+    fn prepare_input_payload(&self, text: &str) -> Vec<u8> {
+        prepare_worker_stdin_payload(text)
+    }
+
     fn on_input_start(
         &mut self,
         text: &str,
@@ -473,12 +482,8 @@ fn driver_refresh_worker_ready(
 }
 
 impl BackendDriver for RBackendDriver {
-    fn prepare_input_payload(&self, text: &str) -> Vec<u8> {
-        let mut payload = text.as_bytes().to_vec();
-        if !payload.is_empty() && !payload.ends_with(b"\n") {
-            payload.push(b'\n');
-        }
-        payload
+    fn prepare_input_text(&self, text: String) -> String {
+        normalize_input_newlines(&text)
     }
 
     fn on_input_start(
@@ -850,14 +855,6 @@ fn strip_one_line_ending(text: &str) -> Option<&str> {
 
 #[cfg(not(target_family = "unix"))]
 impl BackendDriver for PythonBackendDriver {
-    fn prepare_input_payload(&self, text: &str) -> Vec<u8> {
-        let mut payload = text.as_bytes().to_vec();
-        if !payload.is_empty() && !payload.ends_with(b"\n") {
-            payload.push(b'\n');
-        }
-        payload
-    }
-
     fn on_input_start(
         &mut self,
         text: &str,
@@ -934,14 +931,6 @@ impl ProtocolBackendDriver {
 }
 
 impl BackendDriver for ProtocolBackendDriver {
-    fn prepare_input_payload(&self, text: &str) -> Vec<u8> {
-        let mut payload = text.as_bytes().to_vec();
-        if !payload.is_empty() && !payload.ends_with(b"\n") {
-            payload.push(b'\n');
-        }
-        payload
-    }
-
     fn on_input_start(
         &mut self,
         _text: &str,
@@ -1196,12 +1185,6 @@ fn completion_info_from_ipc(
 
 const DEFERRED_SANDBOX_UPDATE_TIMEOUT: Duration = Duration::from_secs(5);
 
-#[derive(Clone, Copy)]
-pub(crate) enum WriteStdinControlAction {
-    Interrupt,
-    Restart,
-}
-
 #[derive(Debug, Clone, Default)]
 pub(crate) struct WriteStdinOptions {
     pub page_bytes_override: Option<u64>,
@@ -1223,29 +1206,6 @@ impl WriteStdinOptions {
     }
 }
 
-pub(crate) fn split_write_stdin_control_prefix(
-    input: &str,
-) -> Option<(WriteStdinControlAction, &str)> {
-    let first = input.chars().next()?;
-    let action = match first {
-        '\u{3}' => WriteStdinControlAction::Interrupt,
-        '\u{4}' => WriteStdinControlAction::Restart,
-        _ => return None,
-    };
-
-    let tail = &input[first.len_utf8()..];
-    let tail = if let Some(rest) = tail.strip_prefix("\r\n") {
-        rest
-    } else if let Some(rest) = tail.strip_prefix('\n') {
-        rest
-    } else if let Some(rest) = tail.strip_prefix('\r') {
-        rest
-    } else {
-        tail
-    };
-    Some((action, tail))
-}
-
 fn worker_context_event_payload(
     worker_launch: &WorkerLaunch,
     backend: Backend,
@@ -2574,7 +2534,7 @@ impl WorkerManager {
         worker_timeout: Duration,
         server_timeout: Duration,
     ) -> Result<RequestState, WorkerError> {
-        let text = normalize_input_newlines(&text);
+        let text = self.driver.prepare_input_text(text);
         let started_at = std::time::Instant::now();
         let prompt = self.current_prompt_hint();
         self.remember_prompt(prompt);
@@ -7605,11 +7565,11 @@ mod tests {
     }
 
     #[test]
-    fn control_prefix_strips_single_separator_newline() {
+    fn control_prefix_preserves_immediate_newline_tail() {
         let (action, remaining) =
             split_write_stdin_control_prefix("\u{4}\nprint(1)").expect("expected control prefix");
         assert!(matches!(action, WriteStdinControlAction::Restart));
-        assert_eq!(remaining, "print(1)");
+        assert_eq!(remaining, "\nprint(1)");
     }
 
     #[test]

tests/common/mod.rs

@@ -669,7 +669,7 @@ impl McpTestSession {
     #[allow(dead_code)]
     pub async fn write_stdin_with(&mut self, input: impl Into<String>, timeout: Option<f64>) {
         let mut input = input.into();
-        if !input.ends_with('\n') {
+        if test_input_needs_trailing_newline(&input) {
             input.push('\n');
         }
         let timeout = normalized_test_timeout(timeout);
@@ -797,7 +797,7 @@ impl McpTestSession {
         meta: Option<Value>,
     ) -> Result<rmcp::model::CallToolResult, ServiceError> {
         let mut input = input.into();
-        if !input.is_empty() && !input.ends_with('\n') {
+        if !input.is_empty() && test_input_needs_trailing_newline(&input) {
             input.push('\n');
         }
         let timeout = normalized_test_timeout(timeout);
@@ -872,6 +872,10 @@ impl McpTestSession {
     }
 }
 
+fn test_input_needs_trailing_newline(input: &str) -> bool {
+    !input.ends_with('\n') && !matches!(input, "\u{3}" | "\u{4}")
+}
+
 pub struct McpSnapshot {
     sessions: Vec<(String, Vec<SnapshotStep>)>,
 }

tests/fixtures/zod-worker.rs

@@ -63,13 +63,15 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut line = String::new();
     let mut next_prompt = "zod> ".to_string();
     let mut timeline = Timeline::default();
+    let mut line_number = 0_u64;
     loop {
         line.clear();
         let bytes = reader.read_line(&mut line)?;
         if bytes == 0 {
             send_session_end(&writer, &mut timeline, "shutdown")?;
             return Ok(());
         }
+        line_number += 1;
 
         let command = line.trim_end_matches(['\r', '\n']);
         let reported_input = if let Some(text) = command.strip_prefix("misreport-input ") {
@@ -96,6 +98,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             &interrupted,
             command,
             &line,
+            line_number,
             &mut next_prompt,
             &mut timeline,
         )?;
@@ -113,6 +116,7 @@ fn run_command(
     interrupted: &AtomicBool,
     command: &str,
     raw_line: &str,
+    line_number: u64,
     next_prompt: &mut String,
     timeline: &mut Timeline,
 ) -> io::Result<()> {
@@ -141,6 +145,15 @@ fn run_command(
         return Ok(());
     }
 
+    if command.starts_with("raw-line-escape") {
+        let escaped = escape_bytes(raw_line.as_bytes());
+        writer.output_text(
+            "stdout",
+            format!("raw-line[{line_number}]={escaped}\n").as_bytes(),
+        )?;
+        return Ok(());
+    }
+
     if let Some(millis) = command.strip_prefix("prompt-then-sleep ") {
         writer.send(&WorkerToServer::ReadlineStart {
             prompt: "buffered> ".to_string(),
@@ -199,6 +212,21 @@ fn run_command(
     writer.output_text("stdout", raw_line.as_bytes())
 }
 
+fn escape_bytes(bytes: &[u8]) -> String {
+    let mut escaped = String::new();
+    for byte in bytes {
+        match byte {
+            b'\n' => escaped.push_str("\\n"),
+            b'\r' => escaped.push_str("\\r"),
+            b'\t' => escaped.push_str("\\t"),
+            b'\\' => escaped.push_str("\\\\"),
+            b' '..=b'~' => escaped.push(char::from(*byte)),
+            _ => escaped.push_str(&format!("\\x{byte:02x}")),
+        }
+    }
+    escaped
+}
+
 fn send_readline_start(
     writer: &IpcWriter,
     timeline: &mut Timeline,

tests/snapshots/mcp_transcripts__snapshots_interrupt_handler_output.snap

@@ -28,7 +28,7 @@ call:
 {
   "tool": "r_repl",
   "arguments": {
-    "input": "\u0003\n",
+    "input": "\u0003",
     "timeout_ms": 5000
   }
 }

tests/snapshots/mcp_transcripts__snapshots_support_multiple_calls_and_sessions.snap

@@ -52,7 +52,7 @@ call:
 {
   "tool": "r_repl",
   "arguments": {
-    "input": "\u0004\n"
+    "input": "\u0004"
   }
 }
 response:

tests/snapshots/mcp_transcripts__snapshots_tempdir_session_restart.snap

@@ -32,7 +32,7 @@ call:
 {
   "tool": "r_repl",
   "arguments": {
-    "input": "\u0004\n"
+    "input": "\u0004"
   }
 }
 response:
@@ -75,7 +75,7 @@ call:
 {
   "tool": "r_repl",
   "arguments": {
-    "input": "\u0004\n"
+    "input": "\u0004"
   }
 }
 response: