codewhale-hq/Codewhale
1
0 Fork 0
mirror of https://github.com/codewhale-hq/Codewhale.git synced 2026-11-08 08:04:25 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

fix(repl): keep rounds in step with the interpreter and bound buffered output

Browse Source 
F02-07 (Refs #6559), F02-09 (Refs #6559, partial).

F02-07: a timed-out round returned an error but left the interpreter
running the old code, and a round accepted any `__RLM_DONE_<sid>__` line
without checking its round id, so the interrupted round's late DONE ended
the next round early and its output was attributed to the wrong round.
DONE now has to carry this round's id (another round's, or user output
imitating one, is kept as output). A round that fails (timeout, I/O) kills
the interpreter and marks the runtime unusable; a round dropped mid-flight
by a caller's own timeout leaves it marked too, so later rounds refuse
instead of desynchronizing. Callers already discard a failed kernel.
Known limit: code running in the REPL can read `_SID` and forge its own
round's sentinels; it already has full Python and `_rpc`, so this grants
nothing new.

F02-09 (partial): stdout lines were read with an unbounded `read_until`,
so one RPC line could make Rust buffer and deserialize an arbitrarily
large request before the bridge's batch cap applied, and `full_stdout`
grew without limit. Lines are now held to 32 MiB (the rest is consumed and
noted; an over-long RPC fails to parse and is answered with an error) and
a round retains at most 32 MiB of stdout, counting the remainder. Not done
here: `RlmSessionConfig::sub_query_timeout_secs` is still ignored by
rlm/bridge.rs (hardcoded 120 s); that wiring is outside this slice.

Validation: new tests
`timed_out_round_kills_the_interpreter_and_refuses_later_rounds` and
`another_rounds_done_sentinel_does_not_end_this_round`; all 30
repl::runtime tests pass. Focused run: `403 passed; 1 failed` (failure
unrelated: T07-01 fixture).

Co-Authored-By: Claude Opus 5.5 <[email protected]>
Signed-off-by: CodeWhale Bot <[email protected]>
This commit is contained in:
CodeWhale BotandClaude Opus 5.5 committed 2026-09-30 01:08:51 -07:00
1 parent d0c640a11f
commit 88cdc1ca5c
1 file changed
+155 -24
Download Patch File Download Diff File Expand all files Collapse all files
crates/tui/src/repl/runtime.rs
+155 -24
View File
@@ -141,6 +141,14 @@ pub trait RpcDispatcher: Send + Sync {
// ---------------------------------------------------------------------------
const DEFAULT_STDOUT_LIMIT: usize = 8_192;
/// Longest single stdout line (protocol or plain output) held in memory. An
/// RPC request is one line, so this also bounds what a request can make Rust
/// deserialize before the bridge's own caps apply; an over-long request is
/// truncated, fails to parse, and is answered with an error.
const MAX_STDOUT_LINE_BYTES: usize = 32 * 1024 * 1024;
/// Most stdout one round retains (`full_stdout`); the rest is counted and
/// reported, never buffered.
const MAX_ROUND_STDOUT_BYTES: usize = 32 * 1024 * 1024;
const ROUND_TIMEOUT: Duration = Duration::from_secs(180);
#[cfg(not(windows))]
const SPAWN_READY_TIMEOUT: Duration = Duration::from_secs(10);
@@ -165,6 +173,13 @@ pub struct PythonRuntime {
round_count: u64,
started: Instant,
round_timeout: Option<Duration>,
/// Set while a round has not yet read its own DONE sentinel, and kept
/// when a round ends without it (timeout, I/O failure, or a caller that
/// dropped the round mid-flight). The protocol stream is then out of step
/// — the interrupted round may still print — so a failed round kills the
/// interpreter and every later round refuses instead of reading another
/// round's output as its own.
broken: Option<String>,
}
impl PythonRuntime {
@@ -250,6 +265,7 @@ impl PythonRuntime {
round_count: 0,
started: Instant::now(),
round_timeout,
broken: None,
};
// Wait for `__RLM_READY_<sid>__` before handing control back. If
@@ -289,18 +305,43 @@ impl PythonRuntime {
}
}
/// One stdout line, holding at most [`MAX_STDOUT_LINE_BYTES`] of it: the
/// rest of an over-long line is consumed and replaced by an omission note.
async fn read_stdout_line_lossy(&mut self) -> Result<Option<String>, String> {
let mut buf = Vec::new();
let n = self
.stdout
.read_until(b'\n', &mut buf)
.await
.map_err(|e| format!("stdout read: {e}"))?;
if n == 0 {
Ok(None)
} else {
Ok(Some(String::from_utf8_lossy(&buf).into_owned()))
let mut omitted = 0usize;
loop {
let available = self
.stdout
.fill_buf()
.await
.map_err(|e| format!("stdout read: {e}"))?;
if available.is_empty() {
break;
}
let (take, ends_line) = match available.iter().position(|&byte| byte == b'\n') {
Some(index) => (index + 1, true),
None => (available.len(), false),
};
let keep = take.min(MAX_STDOUT_LINE_BYTES.saturating_sub(buf.len()));
buf.extend_from_slice(&available[..keep]);
omitted += take - keep;
self.stdout.consume(take);
if ends_line {
break;
}
}
if buf.is_empty() && omitted == 0 {
return Ok(None);
}
let mut line = String::from_utf8_lossy(&buf).into_owned();
if omitted > 0 {
line.truncate(line.trim_end_matches(['\n', '\r']).len());
line.push_str(&format!(
"[... {omitted} bytes of this line omitted by the REPL line limit ...]\n"
));
}
Ok(Some(line))
}
/// Execute a Python code block with no RPC dispatcher. Used for inline
@@ -359,9 +400,18 @@ impl PythonRuntime {
where
D: RpcDispatcher + ?Sized,
{
if let Some(reason) = &self.broken {
return Err(format!(
"Python REPL is unusable after an earlier round failed ({reason}); start a new REPL"
));
}
// Pessimistic until this round reads its own DONE: a caller that
// drops this future mid-round leaves the stream out of step too.
self.broken = Some("an earlier round was cancelled before it finished".to_string());
let started = Instant::now();
self.round_count += 1;
let round_id = self.round_count;
let round_tag = round_id.to_string();
// Send the code header + body + end marker in one write.
let header = format!("__RLM_RUN_{}__::{round_id}\n", self.session_id);
@@ -383,6 +433,7 @@ impl PythonRuntime {
let done_prefix = format!("__RLM_DONE_{}__::", self.session_id);
let mut stdout_buf = String::new();
let mut stdout_omitted = 0usize;
let mut final_value: Option<String> = None;
let mut final_json: Option<Value> = None;
let mut final_confidence: Option<Value> = None;
@@ -401,8 +452,13 @@ impl PythonRuntime {
let trimmed = line.trim_end_matches(['\n', '\r']);
if let Some(rest) = trimmed.strip_prefix(&done_prefix) {
let _ = rest;
break;
if rest == round_tag {
break;
}
// Another round's completion (or user output imitating
// one) never ends this round.
push_bounded(&mut stdout_buf, &line, &mut stdout_omitted);
continue;
}
if let Some(rest) = trimmed.strip_prefix(&final_prefix) {
// New sessions emit an object with value/confidence;
@@ -442,7 +498,11 @@ impl PythonRuntime {
let traceback =
serde_json::from_str::<String>(rest).unwrap_or_else(|_| rest.to_string());
had_error = true;
stdout_buf.push_str(&format!("[traceback]\n{traceback}\n"));
push_bounded(
&mut stdout_buf,
&format!("[traceback]\n{traceback}\n"),
&mut stdout_omitted,
);
continue;
}
if let Some(rest) = trimmed.strip_prefix(&req_prefix) {
@@ -470,24 +530,35 @@ impl PythonRuntime {
continue;
}
stdout_buf.push_str(&line);
push_bounded(&mut stdout_buf, &line, &mut stdout_omitted);
}
Ok::<_, String>(())
};
if let Some(round_timeout) = round_timeout {
match tokio::time::timeout(round_timeout, read_loop).await {
Ok(Ok(())) => {}
Ok(Err(e)) => return Err(e),
Err(_) => {
return Err(format!(
let outcome = match round_timeout {
Some(round_timeout) => tokio::time::timeout(round_timeout, read_loop)
.await
.unwrap_or_else(|_| {
Err(format!(
"REPL round timed out after {}s",
round_timeout.as_secs()
));
}
}
} else {
read_loop.await?;
))
}),
None => read_loop.await,
};
if let Err(error) = outcome {
// The interpreter may still be executing this round. Kill it now
// (the handle is reaped on drop) rather than leave it running and
// let its late output land in a later round.
let _ = self.child.start_kill();
self.broken = Some(error.clone());
return Err(error);
}
self.broken = None;
if stdout_omitted > 0 {
stdout_buf.push_str(&format!(
"\n[... REPL round output capped: {stdout_omitted} bytes omitted ...]\n"
));
}
let stderr = self.drain_stderr().await;
@@ -1038,6 +1109,22 @@ _main_loop()
// Helpers
// ---------------------------------------------------------------------------
/// Append to a round's retained stdout without exceeding
/// [`MAX_ROUND_STDOUT_BYTES`]; whatever does not fit is only counted.
fn push_bounded(buf: &mut String, text: &str, omitted: &mut usize) {
let room = MAX_ROUND_STDOUT_BYTES.saturating_sub(buf.len());
if text.len() <= room {
buf.push_str(text);
return;
}
let cut = (0..=room)
.rev()
.find(|&index| text.is_char_boundary(index))
.unwrap_or(0);
buf.push_str(&text[..cut]);
*omitted += text.len() - cut;
}
fn truncate_stdout(stdout: &str, limit: usize) -> String {
if stdout.len() <= limit {
return stdout.to_string();
@@ -1352,6 +1439,50 @@ mod tests {
rt.shutdown().await;
}
#[tokio::test]
async fn timed_out_round_kills_the_interpreter_and_refuses_later_rounds() {
let mut rt = PythonRuntime::new().await.expect("spawn");
rt.round_timeout = Some(Duration::from_millis(300));
let error = rt
.execute("import time\ntime.sleep(30)\nprint('late output')")
.await
.expect_err("the round must time out");
assert!(error.contains("timed out"), "{error}");
let refused = rt
.execute("print('next round')")
.await
.expect_err("a desynchronized REPL must not run another round");
assert!(refused.contains("unusable"), "{refused}");
let deadline = Instant::now() + Duration::from_secs(5);
while rt.child.try_wait().expect("child status").is_none() {
assert!(
Instant::now() < deadline,
"the timed-out interpreter must be killed"
);
tokio::time::sleep(Duration::from_millis(25)).await;
}
}
#[tokio::test]
async fn another_rounds_done_sentinel_does_not_end_this_round() {
let mut rt = PythonRuntime::new().await.expect("spawn");
let round = rt
.execute("print(_DONE + '999')\nprint('still this round')")
.await
.expect("execute");
assert!(
round.stdout.contains("still this round"),
"{}",
round.stdout
);
let next = rt.execute("print('second')").await.expect("next round");
assert!(next.stdout.contains("second"), "{}", next.stdout);
assert!(!next.stdout.contains("still this round"), "{}", next.stdout);
rt.shutdown().await;
}
#[tokio::test]
async fn inline_runtime_keeps_bounded_round_timeout() {
let rt = PythonRuntime::new().await.expect("spawn");