| 1 | //! Long-lived Python REPL runtime. |
| 2 | //! |
| 3 | //! One Python subprocess lives for the duration of an RLM turn (or an |
| 4 | //! inline `repl` block sequence in the agent loop). Code blocks are sent |
| 5 | //! over stdin framed by `__RLM_RUN__`/`__RLM_END__` sentinels; the bootstrap |
| 6 | //! `exec()`s them into the same global namespace so variables, imports, |
| 7 | //! and even open file handles persist naturally across rounds. |
| 8 | //! |
| 9 | //! Sub-LLM helpers (`sub_query`, `sub_query_batch`, `sub_rlm`, plus legacy |
| 10 | //! `llm_query`, `llm_query_batched`, `rlm_query`, `rlm_query_batched`) are |
| 11 | //! wired through a stdin/stdout RPC protocol: |
| 12 | //! Python emits `__RLM_REQ_<sid>__::{json}` on stdout, Rust dispatches the |
| 13 | //! request and writes `__RLM_RESP_<sid>__::{json}` back on stdin. No HTTP |
| 14 | //! sidecar, no temp ports — the same pipes carry both control and data. |
| 15 | //! |
| 16 | //! The session id (`<sid>`) is a UUID generated per spawn, so user output |
| 17 | //! that happens to contain "REQ" or "FINAL" can't be confused with control |
| 18 | //! messages. |
| 19 | |
| 20 | use std::ffi::OsString; |
| 21 | use std::path::{Path, PathBuf}; |
| 22 | use std::process::Stdio; |
| 23 | use std::time::{Duration, Instant}; |
| 24 | |
| 25 | use serde::{Deserialize, Serialize}; |
| 26 | use serde_json::Value; |
| 27 | use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; |
| 28 | use tokio::process::{Child, ChildStdin, ChildStdout}; |
| 29 | use uuid::Uuid; |
| 30 | |
| 31 | use crate::child_env; |
| 32 | use crate::dependencies::ExternalTool; |
| 33 | |
| 34 | // --------------------------------------------------------------------------- |
| 35 | // Public types |
| 36 | // --------------------------------------------------------------------------- |
| 37 | |
| 38 | /// Result of executing one code block. |
| 39 | #[derive(Debug, Clone)] |
| 40 | pub struct ReplRound { |
| 41 | /// Stdout shown to the model as metadata next round. |
| 42 | pub stdout: String, |
| 43 | /// Full stdout (with sentinels stripped, but otherwise raw). |
| 44 | pub full_stdout: String, |
| 45 | /// Stderr from this round (if any). |
| 46 | pub stderr: String, |
| 47 | /// `True` if the user code raised an unhandled Python exception. |
| 48 | pub has_error: bool, |
| 49 | /// Captured `finalize(value, confidence=...)` payload, if any. |
| 50 | pub final_value: Option<String>, |
| 51 | /// Captured final value before string fallback. Structured `finalize` |
| 52 | /// payloads use this so `handle_read` can expose JSON instead of a Python |
| 53 | /// repr string. |
| 54 | pub final_json: Option<Value>, |
| 55 | /// Optional confidence supplied to `finalize(...)`. |
| 56 | pub final_confidence: Option<Value>, |
| 57 | /// Number of `sub_query`/`sub_rlm` RPCs the round issued. |
| 58 | pub rpc_count: u32, |
| 59 | /// Wall-clock duration of the round. |
| 60 | pub elapsed: Duration, |
| 61 | } |
| 62 | |
| 63 | /// One RPC request emitted by Python during a round. |
| 64 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 65 | #[serde(tag = "type", rename_all = "snake_case")] |
| 66 | pub enum RpcRequest { |
| 67 | /// `llm_query(prompt, model=None, max_tokens=None, system=None)` |
| 68 | Llm { |
| 69 | prompt: String, |
| 70 | #[serde(default)] |
| 71 | model: Option<String>, |
| 72 | #[serde(default)] |
| 73 | max_tokens: Option<u32>, |
| 74 | #[serde(default)] |
| 75 | system: Option<String>, |
| 76 | }, |
| 77 | /// `llm_query_batched(prompts, model=None, dependency_mode="independent")` |
| 78 | LlmBatch { |
| 79 | prompts: Vec<String>, |
| 80 | #[serde(default)] |
| 81 | model: Option<String>, |
| 82 | #[serde(default)] |
| 83 | dependency_mode: Option<String>, |
| 84 | #[serde(default)] |
| 85 | safety_note: Option<String>, |
| 86 | }, |
| 87 | /// `rlm_query(prompt, model=None)` — recursive sub-RLM (paper's `sub_RLM`). |
| 88 | Rlm { |
| 89 | prompt: String, |
| 90 | #[serde(default)] |
| 91 | model: Option<String>, |
| 92 | }, |
| 93 | /// `rlm_query_batched(prompts, model=None, dependency_mode="independent")` |
| 94 | RlmBatch { |
| 95 | prompts: Vec<String>, |
| 96 | #[serde(default)] |
| 97 | model: Option<String>, |
| 98 | #[serde(default)] |
| 99 | dependency_mode: Option<String>, |
| 100 | #[serde(default)] |
| 101 | safety_note: Option<String>, |
| 102 | }, |
| 103 | } |
| 104 | |
| 105 | /// Response for one RPC request. |
| 106 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 107 | #[serde(untagged)] |
| 108 | pub enum RpcResponse { |
| 109 | /// Single-text reply (Llm / Rlm). |
| 110 | Single(SingleResp), |
| 111 | /// Batch reply (LlmBatch / RlmBatch). |
| 112 | Batch(BatchResp), |
| 113 | } |
| 114 | |
| 115 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 116 | pub struct SingleResp { |
| 117 | #[serde(default)] |
| 118 | pub text: String, |
| 119 | #[serde(default, skip_serializing_if = "Option::is_none")] |
| 120 | pub error: Option<String>, |
| 121 | } |
| 122 | |
| 123 | #[derive(Debug, Clone, Serialize, Deserialize)] |
| 124 | pub struct BatchResp { |
| 125 | pub results: Vec<SingleResp>, |
| 126 | } |
| 127 | |
| 128 | /// Trait-object handle for dispatching Python RPCs back into Rust. |
| 129 | /// |
| 130 | /// Each RLM turn supplies one. Implementations forward to the LLM client |
| 131 | /// (and recursively into `run_rlm_turn_inner` for `Rlm` / `RlmBatch`). |
| 132 | pub trait RpcDispatcher: Send + Sync { |
| 133 | fn dispatch<'a>( |
| 134 | &'a self, |
| 135 | req: RpcRequest, |
| 136 | ) -> std::pin::Pin<Box<dyn std::future::Future<Output = RpcResponse> + Send + 'a>>; |
| 137 | } |
| 138 | |
| 139 | // --------------------------------------------------------------------------- |
| 140 | // Constants |
| 141 | // --------------------------------------------------------------------------- |
| 142 | |
| 143 | const DEFAULT_STDOUT_LIMIT: usize = 8_192; |
| 144 | const ROUND_TIMEOUT: Duration = Duration::from_secs(180); |
| 145 | #[cfg(not(windows))] |
| 146 | const SPAWN_READY_TIMEOUT: Duration = Duration::from_secs(10); |
| 147 | #[cfg(windows)] |
| 148 | const SPAWN_READY_TIMEOUT: Duration = Duration::from_secs(30); |
| 149 | |
| 150 | // --------------------------------------------------------------------------- |
| 151 | // PythonRuntime |
| 152 | // --------------------------------------------------------------------------- |
| 153 | |
| 154 | /// Long-lived Python REPL. |
| 155 | #[derive(Debug)] |
| 156 | pub struct PythonRuntime { |
| 157 | child: Child, |
| 158 | stdin: ChildStdin, |
| 159 | stdout: BufReader<ChildStdout>, |
| 160 | /// Per-spawn session id used in protocol sentinels. |
| 161 | session_id: String, |
| 162 | /// Path to the file holding `context` (kept around for cleanup). |
| 163 | context_path: Option<PathBuf>, |
| 164 | stdout_limit: usize, |
| 165 | round_count: u64, |
| 166 | started: Instant, |
| 167 | round_timeout: Option<Duration>, |
| 168 | } |
| 169 | |
| 170 | impl PythonRuntime { |
| 171 | /// Spawn a REPL with no `context` variable and no LLM helpers wired up. |
| 172 | /// Used by the agent loop for inline `repl` blocks the model emits in |
| 173 | /// regular conversation. |
| 174 | pub async fn new() -> Result<Self, String> { |
| 175 | Self::spawn_inner(None, Some(ROUND_TIMEOUT)).await |
| 176 | } |
| 177 | |
| 178 | /// Compatibility shim — older RLM code path used to pass a state file. |
| 179 | /// The state file is no longer used, but the path doubles as an extra |
| 180 | /// scratch location callers can rely on for cleanup symmetry. |
| 181 | pub fn with_state_path(_path: PathBuf) -> Self { |
| 182 | // Synchronous constructor is no longer meaningful: spawning Python |
| 183 | // is async. Callers in turn.rs already use `spawn_with_context` — |
| 184 | // this stub is kept only so the public surface compiles for any |
| 185 | // out-of-tree user. It returns a deliberately broken runtime that |
| 186 | // panics on first use, which is preferable to silently lying. |
| 187 | unreachable!( |
| 188 | "PythonRuntime::with_state_path is deprecated — \ |
| 189 | use PythonRuntime::new() or PythonRuntime::spawn_with_context()" |
| 190 | ) |
| 191 | } |
| 192 | |
| 193 | /// Spawn a REPL with the long input preloaded from a file. Used by the |
| 194 | /// RLM turn loop. |
| 195 | pub async fn spawn_with_context(context_path: &Path) -> Result<Self, String> { |
| 196 | Self::spawn_inner(Some(context_path), None).await |
| 197 | } |
| 198 | |
| 199 | async fn spawn_inner( |
| 200 | context_path: Option<&Path>, |
| 201 | round_timeout: Option<Duration>, |
| 202 | ) -> Result<Self, String> { |
| 203 | let session_id = Uuid::new_v4().simple().to_string(); |
| 204 | let bootstrap = render_bootstrap(&session_id); |
| 205 | |
| 206 | let mut cmd = crate::dependencies::Python::tokio_command().ok_or_else(|| { |
| 207 | "no Python interpreter found on PATH (tried python3, python, py -3). \ |
| 208 | Install Python 3 and restart codewhale." |
| 209 | .to_string() |
| 210 | })?; |
| 211 | cmd.arg("-u") |
| 212 | .arg("-c") |
| 213 | .arg(&bootstrap) |
| 214 | .stdin(Stdio::piped()) |
| 215 | .stdout(Stdio::piped()) |
| 216 | .stderr(Stdio::piped()) |
| 217 | .kill_on_drop(true); |
| 218 | |
| 219 | let context_env = context_path |
| 220 | .map(|path| { |
| 221 | vec![( |
| 222 | OsString::from("RLM_CONTEXT_FILE"), |
| 223 | path.as_os_str().to_os_string(), |
| 224 | )] |
| 225 | }) |
| 226 | .unwrap_or_default(); |
| 227 | child_env::apply_to_tokio_command(&mut cmd, context_env); |
| 228 | |
| 229 | let mut child = cmd |
| 230 | .spawn() |
| 231 | .map_err(|e| format!("failed to spawn Python interpreter: {e}"))?; |
| 232 | |
| 233 | let stdin = child |
| 234 | .stdin |
| 235 | .take() |
| 236 | .ok_or_else(|| "Python interpreter stdin pipe missing".to_string())?; |
| 237 | let raw_stdout = child |
| 238 | .stdout |
| 239 | .take() |
| 240 | .ok_or_else(|| "Python interpreter stdout pipe missing".to_string())?; |
| 241 | let stdout = BufReader::new(raw_stdout); |
| 242 | |
| 243 | let mut rt = Self { |
| 244 | child, |
| 245 | stdin, |
| 246 | stdout, |
| 247 | session_id: session_id.clone(), |
| 248 | context_path: context_path.map(Path::to_path_buf), |
| 249 | stdout_limit: DEFAULT_STDOUT_LIMIT, |
| 250 | round_count: 0, |
| 251 | started: Instant::now(), |
| 252 | round_timeout, |
| 253 | }; |
| 254 | |
| 255 | // Wait for `__RLM_READY_<sid>__` before handing control back. If |
| 256 | // Python failed to start (missing module, syntax error in the |
| 257 | // bootstrap, etc.), this is where we'll find out. |
| 258 | let ready_sentinel = format!("__RLM_READY_{session_id}__"); |
| 259 | match tokio::time::timeout(SPAWN_READY_TIMEOUT, rt.read_until_ready(&ready_sentinel)).await |
| 260 | { |
| 261 | Ok(Ok(())) => Ok(rt), |
| 262 | Ok(Err(e)) => { |
| 263 | let _ = rt.child.kill().await; |
| 264 | Err(format!("Python interpreter bootstrap failed: {e}")) |
| 265 | } |
| 266 | Err(_) => { |
| 267 | let _ = rt.child.kill().await; |
| 268 | Err(format!( |
| 269 | "Python interpreter bootstrap did not signal ready within {}s", |
| 270 | SPAWN_READY_TIMEOUT.as_secs() |
| 271 | )) |
| 272 | } |
| 273 | } |
| 274 | } |
| 275 | |
| 276 | async fn read_until_ready(&mut self, ready_sentinel: &str) -> Result<(), String> { |
| 277 | loop { |
| 278 | let line = match self.read_stdout_line_lossy().await? { |
| 279 | Some(line) => line, |
| 280 | None => { |
| 281 | return Err("Python interpreter closed stdout before ready signal".to_string()); |
| 282 | } |
| 283 | }; |
| 284 | let trimmed = line.trim_end_matches(['\n', '\r']); |
| 285 | if trimmed == ready_sentinel { |
| 286 | return Ok(()); |
| 287 | } |
| 288 | // Pre-ready output is rare; ignore it. |
| 289 | } |
| 290 | } |
| 291 | |
| 292 | async fn read_stdout_line_lossy(&mut self) -> Result<Option<String>, String> { |
| 293 | let mut buf = Vec::new(); |
| 294 | let n = self |
| 295 | .stdout |
| 296 | .read_until(b'\n', &mut buf) |
| 297 | .await |
| 298 | .map_err(|e| format!("stdout read: {e}"))?; |
| 299 | if n == 0 { |
| 300 | Ok(None) |
| 301 | } else { |
| 302 | Ok(Some(String::from_utf8_lossy(&buf).into_owned())) |
| 303 | } |
| 304 | } |
| 305 | |
| 306 | /// Execute a Python code block with no RPC dispatcher. Used for inline |
| 307 | /// `repl` blocks where `llm_query()` should fall back to a sentinel. |
| 308 | pub async fn execute(&mut self, code: &str) -> Result<ReplRound, String> { |
| 309 | self.run(code, None::<&dyn RpcDispatcher>).await |
| 310 | } |
| 311 | |
| 312 | /// Replace the long context visible to bounded helpers without restarting |
| 313 | /// the Python process. User-created variables, imports, and handles stay |
| 314 | /// alive across the refresh, which is what lets the normal agent loop use |
| 315 | /// one working kernel instead of rebuilding a throwaway REPL each turn. |
| 316 | /// |
| 317 | /// The payload travels through a temporary file rather than a generated |
| 318 | /// Python string so large transcripts neither bloat the command stream nor |
| 319 | /// acquire quoting semantics. The old owned file is released only after |
| 320 | /// the kernel has successfully loaded the replacement. |
| 321 | pub async fn replace_context(&mut self, body: &str) -> Result<(), String> { |
| 322 | let path = crate::rlm::session::write_context_file(body) |
| 323 | .map_err(|e| format!("write refreshed REPL context: {e}"))?; |
| 324 | let path_literal = serde_json::to_string(&path.to_string_lossy()) |
| 325 | .map_err(|e| format!("encode refreshed REPL context path: {e}"))?; |
| 326 | let code = format!("_replace_context_file({path_literal})"); |
| 327 | |
| 328 | match self.execute(&code).await { |
| 329 | Ok(round) if !round.has_error => { |
| 330 | if let Some(previous) = self.context_path.replace(path) { |
| 331 | let _ = tokio::fs::remove_file(previous).await; |
| 332 | } |
| 333 | Ok(()) |
| 334 | } |
| 335 | Ok(round) => { |
| 336 | let _ = tokio::fs::remove_file(&path).await; |
| 337 | Err(format!( |
| 338 | "refresh REPL context failed: {}{}", |
| 339 | round.stdout, |
| 340 | if round.stderr.is_empty() { |
| 341 | String::new() |
| 342 | } else { |
| 343 | format!("\nstderr: {}", round.stderr) |
| 344 | } |
| 345 | )) |
| 346 | } |
| 347 | Err(error) => { |
| 348 | let _ = tokio::fs::remove_file(&path).await; |
| 349 | Err(error) |
| 350 | } |
| 351 | } |
| 352 | } |
| 353 | |
| 354 | /// Execute a code block, dispatching any sub-LLM RPCs through `bridge`. |
| 355 | /// |
| 356 | /// Returns once Python emits `__RLM_DONE_<sid>__` or the round timeout |
| 357 | /// elapses (whichever happens first). |
| 358 | pub async fn run<D>(&mut self, code: &str, bridge: Option<&D>) -> Result<ReplRound, String> |
| 359 | where |
| 360 | D: RpcDispatcher + ?Sized, |
| 361 | { |
| 362 | let started = Instant::now(); |
| 363 | self.round_count += 1; |
| 364 | let round_id = self.round_count; |
| 365 | |
| 366 | // Send the code header + body + end marker in one write. |
| 367 | let header = format!("__RLM_RUN_{}__::{round_id}\n", self.session_id); |
| 368 | let footer = format!("__RLM_END_{}__\n", self.session_id); |
| 369 | let payload = format!("{header}{code}\n{footer}"); |
| 370 | self.stdin |
| 371 | .write_all(payload.as_bytes()) |
| 372 | .await |
| 373 | .map_err(|e| format!("stdin write: {e}"))?; |
| 374 | self.stdin |
| 375 | .flush() |
| 376 | .await |
| 377 | .map_err(|e| format!("stdin flush: {e}"))?; |
| 378 | |
| 379 | // Sentinels for this session. |
| 380 | let req_prefix = format!("__RLM_REQ_{}__::", self.session_id); |
| 381 | let final_prefix = format!("__RLM_FINAL_{}__::", self.session_id); |
| 382 | let err_prefix = format!("__RLM_ERR_{}__::", self.session_id); |
| 383 | let done_prefix = format!("__RLM_DONE_{}__::", self.session_id); |
| 384 | |
| 385 | let mut stdout_buf = String::new(); |
| 386 | let mut final_value: Option<String> = None; |
| 387 | let mut final_json: Option<Value> = None; |
| 388 | let mut final_confidence: Option<Value> = None; |
| 389 | let mut had_error = false; |
| 390 | let mut rpc_count: u32 = 0; |
| 391 | let round_timeout = self.round_timeout; |
| 392 | |
| 393 | let read_loop = async { |
| 394 | loop { |
| 395 | let line = match self.read_stdout_line_lossy().await? { |
| 396 | Some(line) => line, |
| 397 | None => { |
| 398 | return Err("Python interpreter closed stdout mid-round".to_string()); |
| 399 | } |
| 400 | }; |
| 401 | let trimmed = line.trim_end_matches(['\n', '\r']); |
| 402 | |
| 403 | if let Some(rest) = trimmed.strip_prefix(&done_prefix) { |
| 404 | let _ = rest; |
| 405 | break; |
| 406 | } |
| 407 | if let Some(rest) = trimmed.strip_prefix(&final_prefix) { |
| 408 | // New sessions emit an object with value/confidence; |
| 409 | // legacy helpers emitted a JSON string. |
| 410 | match serde_json::from_str::<Value>(rest) { |
| 411 | Ok(Value::Object(map)) => { |
| 412 | let value_json = map |
| 413 | .get("value") |
| 414 | .cloned() |
| 415 | .unwrap_or(Value::String(rest.to_string())); |
| 416 | let value = value_json |
| 417 | .as_str() |
| 418 | .map(str::to_string) |
| 419 | .unwrap_or_else(|| value_json.to_string()); |
| 420 | final_json = Some(value_json); |
| 421 | final_value = Some(value); |
| 422 | final_confidence = map.get("confidence").cloned(); |
| 423 | } |
| 424 | Ok(Value::String(value)) => { |
| 425 | final_json = Some(Value::String(value.clone())); |
| 426 | final_value = Some(value); |
| 427 | final_confidence = None; |
| 428 | } |
| 429 | Ok(other) => { |
| 430 | final_json = Some(other.clone()); |
| 431 | final_value = Some(other.to_string()); |
| 432 | final_confidence = None; |
| 433 | } |
| 434 | Err(_) => { |
| 435 | final_value = Some(rest.to_string()); |
| 436 | final_confidence = None; |
| 437 | } |
| 438 | } |
| 439 | continue; |
| 440 | } |
| 441 | if let Some(rest) = trimmed.strip_prefix(&err_prefix) { |
| 442 | let traceback = |
| 443 | serde_json::from_str::<String>(rest).unwrap_or_else(|_| rest.to_string()); |
| 444 | had_error = true; |
| 445 | stdout_buf.push_str(&format!("[traceback]\n{traceback}\n")); |
| 446 | continue; |
| 447 | } |
| 448 | if let Some(rest) = trimmed.strip_prefix(&req_prefix) { |
| 449 | rpc_count = rpc_count.saturating_add(1); |
| 450 | let req: RpcRequest = match serde_json::from_str(rest) { |
| 451 | Ok(r) => r, |
| 452 | Err(e) => { |
| 453 | // Send an error response so Python isn't blocked. |
| 454 | self.send_resp(&RpcResponse::Single(SingleResp { |
| 455 | text: String::new(), |
| 456 | error: Some(format!("malformed RPC: {e}")), |
| 457 | })) |
| 458 | .await?; |
| 459 | continue; |
| 460 | } |
| 461 | }; |
| 462 | let resp = match bridge { |
| 463 | Some(b) => b.dispatch(req).await, |
| 464 | None => RpcResponse::Single(SingleResp { |
| 465 | text: String::new(), |
| 466 | error: Some("no LLM bridge bound to this REPL".to_string()), |
| 467 | }), |
| 468 | }; |
| 469 | self.send_resp(&resp).await?; |
| 470 | continue; |
| 471 | } |
| 472 | |
| 473 | stdout_buf.push_str(&line); |
| 474 | } |
| 475 | Ok::<_, String>(()) |
| 476 | }; |
| 477 | |
| 478 | if let Some(round_timeout) = round_timeout { |
| 479 | match tokio::time::timeout(round_timeout, read_loop).await { |
| 480 | Ok(Ok(())) => {} |
| 481 | Ok(Err(e)) => return Err(e), |
| 482 | Err(_) => { |
| 483 | return Err(format!( |
| 484 | "REPL round timed out after {}s", |
| 485 | round_timeout.as_secs() |
| 486 | )); |
| 487 | } |
| 488 | } |
| 489 | } else { |
| 490 | read_loop.await?; |
| 491 | } |
| 492 | |
| 493 | let stderr = self.drain_stderr().await; |
| 494 | let display = truncate_stdout(stdout_buf.trim_end_matches('\n'), self.stdout_limit); |
| 495 | |
| 496 | Ok(ReplRound { |
| 497 | stdout: display, |
| 498 | full_stdout: stdout_buf, |
| 499 | stderr, |
| 500 | has_error: had_error, |
| 501 | final_value, |
| 502 | final_json, |
| 503 | final_confidence, |
| 504 | rpc_count, |
| 505 | elapsed: started.elapsed(), |
| 506 | }) |
| 507 | } |
| 508 | |
| 509 | async fn send_resp(&mut self, resp: &RpcResponse) -> Result<(), String> { |
| 510 | let body = serde_json::to_string(resp).map_err(|e| format!("encode rpc resp: {e}"))?; |
| 511 | let line = format!("__RLM_RESP_{}__::{body}\n", self.session_id); |
| 512 | self.stdin |
| 513 | .write_all(line.as_bytes()) |
| 514 | .await |
| 515 | .map_err(|e| format!("stdin write resp: {e}"))?; |
| 516 | self.stdin |
| 517 | .flush() |
| 518 | .await |
| 519 | .map_err(|e| format!("stdin flush resp: {e}"))?; |
| 520 | Ok(()) |
| 521 | } |
| 522 | |
| 523 | async fn drain_stderr(&mut self) -> String { |
| 524 | // We don't continuously read stderr — drain whatever's pending after |
| 525 | // a round so it can show up in error reports without deadlocking |
| 526 | // anything during normal operation. |
| 527 | let Some(stderr) = self.child.stderr.as_mut() else { |
| 528 | return String::new(); |
| 529 | }; |
| 530 | use tokio::io::AsyncReadExt; |
| 531 | let mut buf = Vec::new(); |
| 532 | // Best-effort read with a tight deadline; we don't want to block. |
| 533 | let fut = async { |
| 534 | let mut chunk = [0u8; 4096]; |
| 535 | loop { |
| 536 | match tokio::time::timeout(Duration::from_millis(20), stderr.read(&mut chunk)).await |
| 537 | { |
| 538 | Ok(Ok(0)) => break, |
| 539 | Ok(Ok(n)) => buf.extend_from_slice(&chunk[..n]), |
| 540 | _ => break, |
| 541 | } |
| 542 | } |
| 543 | }; |
| 544 | let _ = fut.await; |
| 545 | String::from_utf8_lossy(&buf).to_string() |
| 546 | } |
| 547 | |
| 548 | /// Total rounds executed. |
| 549 | pub fn round_count(&self) -> u64 { |
| 550 | self.round_count |
| 551 | } |
| 552 | |
| 553 | /// Current per-round timeout policy. RLM context runs intentionally return |
| 554 | /// `None` so long map-reduce jobs are not killed by the old 180s cap. |
| 555 | pub fn round_timeout(&self) -> Option<Duration> { |
| 556 | self.round_timeout |
| 557 | } |
| 558 | |
| 559 | /// Wall-clock uptime since spawn. |
| 560 | pub fn uptime(&self) -> Duration { |
| 561 | self.started.elapsed() |
| 562 | } |
| 563 | |
| 564 | /// Cleanly tear down the subprocess. |
| 565 | pub async fn shutdown(mut self) { |
| 566 | let _ = self.stdin.shutdown().await; |
| 567 | let _ = self.child.kill().await; |
| 568 | if let Some(path) = self.context_path.take() { |
| 569 | let _ = tokio::fs::remove_file(path).await; |
| 570 | } |
| 571 | } |
| 572 | } |
| 573 | |
| 574 | impl Drop for PythonRuntime { |
| 575 | fn drop(&mut self) { |
| 576 | // tokio sets `kill_on_drop(true)` on the child; the context file |
| 577 | // (if any) is removed on `shutdown()` — drop is best-effort. |
| 578 | if let Some(path) = self.context_path.take() { |
| 579 | let _ = std::fs::remove_file(path); |
| 580 | } |
| 581 | } |
| 582 | } |
| 583 | |
| 584 | // --------------------------------------------------------------------------- |
| 585 | // Bootstrap script |
| 586 | // --------------------------------------------------------------------------- |
| 587 | |
| 588 | /// Render the Python bootstrap with session-specific sentinels baked in. |
| 589 | /// The sentinels include a UUID to prevent user prints from being mistaken |
| 590 | /// for control messages. |
| 591 | fn render_bootstrap(session_id: &str) -> String { |
| 592 | BOOTSTRAP_TEMPLATE.replace("__SID__", session_id) |
| 593 | } |
| 594 | |
| 595 | const BOOTSTRAP_TEMPLATE: &str = r#" |
| 596 | import json as _json |
| 597 | import os as _os |
| 598 | import re as _re |
| 599 | import sys as _sys |
| 600 | import traceback as _traceback |
| 601 | |
| 602 | _SID = "__SID__" |
| 603 | _REQ = f"__RLM_REQ_{_SID}__::" |
| 604 | _RESP = f"__RLM_RESP_{_SID}__::" |
| 605 | _FINAL = f"__RLM_FINAL_{_SID}__::" |
| 606 | _ERR = f"__RLM_ERR_{_SID}__::" |
| 607 | _RUN = f"__RLM_RUN_{_SID}__::" |
| 608 | _END = f"__RLM_END_{_SID}__" |
| 609 | _DONE = f"__RLM_DONE_{_SID}__::" |
| 610 | _READY = f"__RLM_READY_{_SID}__" |
| 611 | |
| 612 | def _rpc(req): |
| 613 | _sys.stdout.write(_REQ + _json.dumps(req) + "\n") |
| 614 | _sys.stdout.flush() |
| 615 | line = _sys.stdin.readline() |
| 616 | if not line: |
| 617 | return {"error": "rust driver closed stdin"} |
| 618 | if line.startswith(_RESP): |
| 619 | try: |
| 620 | return _json.loads(line[len(_RESP):]) |
| 621 | except Exception as e: |
| 622 | return {"error": f"malformed rpc resp: {e}"} |
| 623 | return {"error": f"unexpected protocol line: {line[:120]!r}"} |
| 624 | |
| 625 | def llm_query(prompt, model=None, max_tokens=None, system=None): |
| 626 | """One-shot sub-LLM call. The model arg is accepted for compatibility but ignored by Rust.""" |
| 627 | resp = _rpc({"type":"llm","prompt":str(prompt),"model":model, |
| 628 | "max_tokens":max_tokens,"system":system}) |
| 629 | if isinstance(resp, dict) and resp.get("error"): |
| 630 | return f"[llm_query error: {resp['error']}]" |
| 631 | if isinstance(resp, dict): |
| 632 | return resp.get("text","") |
| 633 | return str(resp) |
| 634 | |
| 635 | def _normalize_dependency_mode(mode): |
| 636 | if mode is None: |
| 637 | return "" |
| 638 | return str(mode).strip().lower().replace("-", "_").replace(" ", "_") |
| 639 | |
| 640 | def _batch_dependency_error(helper, prompts, dependency_mode): |
| 641 | mode = _normalize_dependency_mode(dependency_mode) |
| 642 | if mode in ("independent", "parallel_safe", "map_reduce"): |
| 643 | return None |
| 644 | if mode in ("sequential", "dependent", "ordered", "chain", "serial"): |
| 645 | return ( |
| 646 | f"[{helper}: refused parallel batch because dependency_mode={dependency_mode!r}. " |
| 647 | "Use sub_query_sequence(...) or an explicit for-loop with sub_query(...) so each step can consume the previous result.]" |
| 648 | ) |
| 649 | return ( |
| 650 | f"[{helper}: batch helpers require dependency_mode='independent'. " |
| 651 | "Use only for independent slices/items; for A->B dependencies, global-state refactors, migrations, or rollback-sensitive work, use sub_query_sequence(...).]" |
| 652 | ) |
| 653 | |
| 654 | def llm_query_batched(prompts, model=None, dependency_mode=None, safety_note=None): |
| 655 | """Run independent sub-LLM calls concurrently. Declare dependency_mode='independent'.""" |
| 656 | if not isinstance(prompts, (list, tuple)): |
| 657 | return ["[llm_query_batched: prompts must be a list]"] |
| 658 | err = _batch_dependency_error("llm_query_batched", prompts, dependency_mode) |
| 659 | if err is not None: |
| 660 | return [err for _ in prompts] |
| 661 | resp = _rpc({ |
| 662 | "type":"llm_batch", |
| 663 | "prompts":[str(p) for p in prompts], |
| 664 | "model":model, |
| 665 | "dependency_mode":dependency_mode, |
| 666 | "safety_note":safety_note, |
| 667 | }) |
| 668 | if isinstance(resp, dict) and resp.get("error"): |
| 669 | return [f"[llm_query_batched: {resp['error']}]" for _ in prompts] |
| 670 | results = (resp or {}).get("results", []) if isinstance(resp, dict) else [] |
| 671 | if len(results) != len(prompts): |
| 672 | return [f"[llm_query_batched: size mismatch ({len(results)}/{len(prompts)})]" for _ in prompts] |
| 673 | out = [] |
| 674 | for r in results: |
| 675 | if r.get("error"): |
| 676 | out.append(f"[child err: {r['error']}]") |
| 677 | else: |
| 678 | out.append(r.get("text","")) |
| 679 | return out |
| 680 | |
| 681 | def rlm_query(prompt, model=None): |
| 682 | """Recursive sub-RLM. The model arg is accepted for compatibility but ignored by Rust.""" |
| 683 | resp = _rpc({"type":"rlm","prompt":str(prompt),"model":model}) |
| 684 | if isinstance(resp, dict) and resp.get("error"): |
| 685 | return f"[rlm_query error: {resp['error']}]" |
| 686 | if isinstance(resp, dict): |
| 687 | return resp.get("text","") |
| 688 | return str(resp) |
| 689 | |
| 690 | def rlm_query_batched(prompts, model=None, dependency_mode=None, safety_note=None): |
| 691 | """Run independent recursive sub-RLMs in parallel. Declare dependency_mode='independent'.""" |
| 692 | if not isinstance(prompts, (list, tuple)): |
| 693 | return ["[rlm_query_batched: prompts must be a list]"] |
| 694 | err = _batch_dependency_error("rlm_query_batched", prompts, dependency_mode) |
| 695 | if err is not None: |
| 696 | return [err for _ in prompts] |
| 697 | resp = _rpc({ |
| 698 | "type":"rlm_batch", |
| 699 | "prompts":[str(p) for p in prompts], |
| 700 | "model":model, |
| 701 | "dependency_mode":dependency_mode, |
| 702 | "safety_note":safety_note, |
| 703 | }) |
| 704 | if isinstance(resp, dict) and resp.get("error"): |
| 705 | return [f"[rlm_query_batched: {resp['error']}]" for _ in prompts] |
| 706 | results = (resp or {}).get("results", []) if isinstance(resp, dict) else [] |
| 707 | if len(results) != len(prompts): |
| 708 | return [f"[rlm_query_batched: size mismatch ({len(results)}/{len(prompts)})]" for _ in prompts] |
| 709 | out = [] |
| 710 | for r in results: |
| 711 | if r.get("error"): |
| 712 | out.append(f"[child err: {r['error']}]") |
| 713 | else: |
| 714 | out.append(r.get("text","")) |
| 715 | return out |
| 716 | |
| 717 | def _slice_text(slice_value): |
| 718 | if slice_value is None: |
| 719 | return "" |
| 720 | if isinstance(slice_value, dict): |
| 721 | if "text" in slice_value: |
| 722 | return str(slice_value["text"]) |
| 723 | return _json.dumps(slice_value, ensure_ascii=False) |
| 724 | return str(slice_value) |
| 725 | |
| 726 | def _prompt_with_slice(prompt, slice_value): |
| 727 | text = _slice_text(slice_value) |
| 728 | if not text: |
| 729 | return str(prompt) |
| 730 | if isinstance(slice_value, dict) and ("index" in slice_value or ("start" in slice_value and "end" in slice_value)): |
| 731 | label = f"slice index={slice_value.get('index', '?')} range={slice_value.get('start', '?')}:{slice_value.get('end', '?')}" |
| 732 | else: |
| 733 | label = "slice" |
| 734 | return f"{prompt}\n\n--- {label} ---\n{text}" |
| 735 | |
| 736 | def sub_query(prompt, slice=None, timeout_secs=None, **kwargs): |
| 737 | """One child LLM call, optionally scoped to a bounded slice.""" |
| 738 | return llm_query(_prompt_with_slice(prompt, slice)) |
| 739 | |
| 740 | def sub_query_batch(prompt, slices, timeout_secs=None, dependency_mode=None, safety_note=None, **kwargs): |
| 741 | """Apply one prompt to many independent bounded slices concurrently.""" |
| 742 | if not isinstance(slices, (list, tuple)): |
| 743 | return ["[sub_query_batch: slices must be a list]"] |
| 744 | return llm_query_batched( |
| 745 | [_prompt_with_slice(prompt, s) for s in slices], |
| 746 | dependency_mode=dependency_mode, |
| 747 | safety_note=safety_note, |
| 748 | ) |
| 749 | |
| 750 | def sub_query_map(prompts, slices=None, timeout_secs=None, dependency_mode=None, safety_note=None, **kwargs): |
| 751 | """Run N distinct independent prompts, optionally paired with N bounded slices.""" |
| 752 | if not isinstance(prompts, (list, tuple)): |
| 753 | return ["[sub_query_map: prompts must be a list]"] |
| 754 | if slices is None: |
| 755 | return llm_query_batched( |
| 756 | [str(p) for p in prompts], |
| 757 | dependency_mode=dependency_mode, |
| 758 | safety_note=safety_note, |
| 759 | ) |
| 760 | if not isinstance(slices, (list, tuple)): |
| 761 | return ["[sub_query_map: slices must be a list]"] |
| 762 | if len(prompts) != len(slices): |
| 763 | return [f"[sub_query_map: size mismatch ({len(prompts)}/{len(slices)})]" for _ in prompts] |
| 764 | return llm_query_batched( |
| 765 | [_prompt_with_slice(p, s) for p, s in zip(prompts, slices)], |
| 766 | dependency_mode=dependency_mode, |
| 767 | safety_note=safety_note, |
| 768 | ) |
| 769 | |
| 770 | def sub_query_sequence(prompt, slices, carry_prompt=None, timeout_secs=None, **kwargs): |
| 771 | """Apply one prompt to slices sequentially, feeding each result into the next step.""" |
| 772 | if not isinstance(slices, (list, tuple)): |
| 773 | return ["[sub_query_sequence: slices must be a list]"] |
| 774 | out = [] |
| 775 | previous = "" |
| 776 | carry = str(carry_prompt or "Previous step result; treat it as required input for this step:") |
| 777 | total = len(slices) |
| 778 | for i, s in enumerate(slices): |
| 779 | step_prompt = _prompt_with_slice(prompt, s) |
| 780 | if previous: |
| 781 | step_prompt = ( |
| 782 | f"{step_prompt}\n\n--- dependency_state step {i}/{total} ---\n" |
| 783 | f"{carry}\n{previous}" |
| 784 | ) |
| 785 | result = llm_query(step_prompt) |
| 786 | out.append(result) |
| 787 | previous = result |
| 788 | return out |
| 789 | |
| 790 | def sub_rlm(prompt, source=None, timeout_secs=None, **kwargs): |
| 791 | """Recursive sub-RLM call for tasks that need their own decomposition.""" |
| 792 | return rlm_query(_prompt_with_slice(prompt, source)) |
| 793 | |
| 794 | def _json_safe(value): |
| 795 | try: |
| 796 | _json.dumps(value, ensure_ascii=False) |
| 797 | return value |
| 798 | except Exception: |
| 799 | return str(value) |
| 800 | |
| 801 | def _emit_final(value, confidence=None): |
| 802 | safe_value = _json_safe(value) |
| 803 | _sys.stdout.write(_FINAL + _json.dumps({ |
| 804 | "value": safe_value, |
| 805 | "confidence": confidence, |
| 806 | }, ensure_ascii=False) + "\n") |
| 807 | _sys.stdout.flush() |
| 808 | |
| 809 | def FINAL(value): |
| 810 | """Legacy compatibility alias for finalize(value).""" |
| 811 | _emit_final(value) |
| 812 | |
| 813 | def FINAL_VAR(name): |
| 814 | """Legacy compatibility alias for finalize(repl_get(name)).""" |
| 815 | name_str = str(name).strip().strip("'\"") |
| 816 | if name_str in globals(): |
| 817 | _emit_final(globals()[name_str]) |
| 818 | else: |
| 819 | print(f"FINAL_VAR error: variable '{name_str}' not found. " |
| 820 | f"Use SHOW_VARS() to list available variables.", flush=True) |
| 821 | |
| 822 | def SHOW_VARS(): |
| 823 | """Return a dict of {name: type-name} for all user variables in the REPL.""" |
| 824 | out = {} |
| 825 | for k, v in list(globals().items()): |
| 826 | if k.startswith('_') or k in _BOOTSTRAP_NAMES: |
| 827 | continue |
| 828 | out[k] = type(v).__name__ |
| 829 | return out |
| 830 | |
| 831 | def repl_get(name, default=None): |
| 832 | return globals().get(str(name), default) |
| 833 | |
| 834 | def repl_set(name, value): |
| 835 | globals()[str(name)] = value |
| 836 | |
| 837 | def context_meta(): |
| 838 | """Return bounded metadata about the loaded input; never includes the full text.""" |
| 839 | text = _context |
| 840 | line_count = 0 if text == "" else text.count("\n") + (0 if text.endswith("\n") else 1) |
| 841 | return { |
| 842 | "chars": len(text), |
| 843 | "lines": line_count, |
| 844 | "preview": text[:500], |
| 845 | "tail_preview": text[-500:] if len(text) > 500 else text, |
| 846 | } |
| 847 | |
| 848 | def _slice_chars(start, end): |
| 849 | total = len(_context) |
| 850 | s = max(0, int(start)) |
| 851 | e = max(s, min(total, int(end))) |
| 852 | return _context[s:e] |
| 853 | |
| 854 | def _slice_lines(start, end): |
| 855 | lines = _context.splitlines() |
| 856 | s = max(0, int(start)) |
| 857 | e = max(s, min(len(lines), int(end))) |
| 858 | return "\n".join(lines[s:e]) |
| 859 | |
| 860 | def peek(start, end, unit="chars"): |
| 861 | """Return a bounded slice of the input by char offsets or line numbers.""" |
| 862 | if str(unit).lower() in ("line", "lines"): |
| 863 | return _slice_lines(start, end) |
| 864 | if str(unit).lower() not in ("char", "chars"): |
| 865 | raise ValueError("unit must be 'chars' or 'lines'") |
| 866 | return _slice_chars(start, end) |
| 867 | |
| 868 | def search(pattern, max_hits=100): |
| 869 | """Regex-search the input and return bounded hit records with snippets.""" |
| 870 | max_hits = max(0, int(max_hits)) |
| 871 | hits = [] |
| 872 | if max_hits == 0: |
| 873 | return hits |
| 874 | rx = _re.compile(str(pattern), _re.MULTILINE) |
| 875 | for i, m in enumerate(rx.finditer(_context)): |
| 876 | if i >= max_hits: |
| 877 | break |
| 878 | start, end = m.span() |
| 879 | snippet_start = max(0, start - 120) |
| 880 | snippet_end = min(len(_context), end + 120) |
| 881 | hits.append({ |
| 882 | "index": i, |
| 883 | "start": start, |
| 884 | "end": end, |
| 885 | "match": m.group(0), |
| 886 | "snippet": _context[snippet_start:snippet_end], |
| 887 | }) |
| 888 | return hits |
| 889 | |
| 890 | def chunk(max_chars=20000, overlap=0): |
| 891 | """Return full-coverage input chunks with index/start/end/text fields.""" |
| 892 | max_chars = int(max_chars) |
| 893 | overlap = max(0, int(overlap)) |
| 894 | if max_chars <= 0: |
| 895 | raise ValueError("max_chars must be > 0") |
| 896 | if overlap >= max_chars: |
| 897 | raise ValueError("overlap must be smaller than max_chars") |
| 898 | chunks = [] |
| 899 | start = 0 |
| 900 | idx = 0 |
| 901 | total = len(_context) |
| 902 | while start < total: |
| 903 | end = min(total, start + max_chars) |
| 904 | chunks.append({"index": idx, "start": start, "end": end, "text": _context[start:end]}) |
| 905 | idx += 1 |
| 906 | if end >= total: |
| 907 | break |
| 908 | start = end - overlap |
| 909 | return chunks |
| 910 | |
| 911 | def chunk_context(max_chars=20000, overlap=0): |
| 912 | """Compatibility alias for chunk().""" |
| 913 | return chunk(max_chars=max_chars, overlap=overlap) |
| 914 | |
| 915 | def chunk_coverage(chunks): |
| 916 | """Summarize coverage for chunks produced by chunk().""" |
| 917 | spans = [] |
| 918 | for c in chunks: |
| 919 | try: |
| 920 | spans.append((int(c["start"]), int(c["end"]))) |
| 921 | except Exception: |
| 922 | continue |
| 923 | spans.sort() |
| 924 | covered = 0 |
| 925 | cursor = 0 |
| 926 | gaps = [] |
| 927 | for start, end in spans: |
| 928 | if start > cursor: |
| 929 | gaps.append((cursor, start)) |
| 930 | if end > cursor: |
| 931 | covered += end - max(start, cursor) |
| 932 | cursor = end |
| 933 | if cursor < len(_context): |
| 934 | gaps.append((cursor, len(_context))) |
| 935 | return { |
| 936 | "chunks": len(chunks), |
| 937 | "context_chars": len(_context), |
| 938 | "input_chars": len(_context), |
| 939 | "covered_chars": covered, |
| 940 | "gaps": gaps, |
| 941 | "complete": covered >= len(_context) and not gaps, |
| 942 | } |
| 943 | |
| 944 | def finalize(value, confidence=None): |
| 945 | """Signal the session's final answer and persist confidence metadata.""" |
| 946 | global final_answer, final_confidence, final_result |
| 947 | final_answer = _json_safe(value) |
| 948 | final_confidence = confidence |
| 949 | final_result = { |
| 950 | "value": final_answer, |
| 951 | "confidence": confidence, |
| 952 | } |
| 953 | _emit_final(final_answer, confidence=confidence) |
| 954 | return final_answer |
| 955 | |
| 956 | def evaluate_progress(): |
| 957 | """Return lightweight state useful before deciding the next REPL step.""" |
| 958 | vars_now = SHOW_VARS() |
| 959 | return { |
| 960 | "has_final_answer": "final_answer" in globals(), |
| 961 | "final_confidence": globals().get("final_confidence", None), |
| 962 | "user_variables": vars_now, |
| 963 | } |
| 964 | |
| 965 | # Load the long input from a file. This keeps the big string out of the |
| 966 | # process command-line and out of the LLM's window. |
| 967 | _ctx_file = _os.environ.get("RLM_CONTEXT_FILE","") |
| 968 | _context = "" |
| 969 | if _ctx_file: |
| 970 | try: |
| 971 | with open(_ctx_file, "r", encoding="utf-8", errors="replace") as f: |
| 972 | _context = f.read() |
| 973 | except Exception as e: |
| 974 | _sys.stderr.write(f"[bootstrap] failed to load context: {e}\n") |
| 975 | content = _context |
| 976 | |
| 977 | def _replace_context_file(path): |
| 978 | """Atomically switch bounded helpers to a freshly written context file.""" |
| 979 | global _context, content |
| 980 | with open(path, "r", encoding="utf-8", errors="replace") as f: |
| 981 | _context = f.read() |
| 982 | content = _context |
| 983 | return context_meta() |
| 984 | |
| 985 | _BOOTSTRAP_NAMES = { |
| 986 | "_SID","_REQ","_RESP","_FINAL","_ERR","_RUN","_END","_DONE","_READY", |
| 987 | "_rpc","_ctx_file","_context","_slice_chars","_slice_lines","_replace_context_file","_BOOTSTRAP_NAMES","_main_loop", |
| 988 | "_emit_final","_json_safe","_slice_text","_prompt_with_slice", |
| 989 | "_normalize_dependency_mode","_batch_dependency_error", |
| 990 | "llm_query","llm_query_batched","rlm_query","rlm_query_batched", |
| 991 | "sub_query","sub_query_batch","sub_query_map","sub_query_sequence","sub_rlm", |
| 992 | "FINAL","FINAL_VAR","SHOW_VARS","repl_get","repl_set", |
| 993 | "context_meta","peek","search","chunk","chunk_context","chunk_coverage", |
| 994 | "finalize","evaluate_progress","content", |
| 995 | "_json","_os","_re","_sys","_traceback", |
| 996 | } |
| 997 | |
| 998 | def _main_loop(): |
| 999 | _sys.stdout.write(_READY + "\n") |
| 1000 | _sys.stdout.flush() |
| 1001 | while True: |
| 1002 | header = _sys.stdin.readline() |
| 1003 | if not header: |
| 1004 | return |
| 1005 | if not header.startswith(_RUN): |
| 1006 | continue |
| 1007 | round_id = header.rstrip("\n")[len(_RUN):] |
| 1008 | code_lines = [] |
| 1009 | while True: |
| 1010 | line = _sys.stdin.readline() |
| 1011 | if not line: |
| 1012 | return |
| 1013 | if line.rstrip("\n") == _END: |
| 1014 | break |
| 1015 | code_lines.append(line) |
| 1016 | code = "".join(code_lines) |
| 1017 | try: |
| 1018 | exec(compile(code, f"<repl-{round_id}>", "exec"), globals()) |
| 1019 | except SystemExit: |
| 1020 | _sys.stdout.write(_DONE + round_id + "\n") |
| 1021 | _sys.stdout.flush() |
| 1022 | return |
| 1023 | except BaseException: |
| 1024 | tb = _traceback.format_exc() |
| 1025 | _sys.stdout.write(_ERR + _json.dumps(tb) + "\n") |
| 1026 | _sys.stdout.flush() |
| 1027 | _sys.stdout.write(_DONE + round_id + "\n") |
| 1028 | _sys.stdout.flush() |
| 1029 | |
| 1030 | _main_loop() |
| 1031 | "#; |
| 1032 | |
| 1033 | // --------------------------------------------------------------------------- |
| 1034 | // Helpers |
| 1035 | // --------------------------------------------------------------------------- |
| 1036 | |
| 1037 | fn truncate_stdout(stdout: &str, limit: usize) -> String { |
| 1038 | if stdout.len() <= limit { |
| 1039 | return stdout.to_string(); |
| 1040 | } |
| 1041 | let take = limit.saturating_sub(80); |
| 1042 | let mut out: String = stdout.chars().take(take).collect(); |
| 1043 | let omitted = stdout.len().saturating_sub(out.len()); |
| 1044 | out.push_str(&format!( |
| 1045 | "\n\n[... REPL output truncated: {omitted} bytes omitted ...]\n" |
| 1046 | )); |
| 1047 | out |
| 1048 | } |
| 1049 | |
| 1050 | // --------------------------------------------------------------------------- |
| 1051 | // Tests |
| 1052 | // --------------------------------------------------------------------------- |
| 1053 | |
| 1054 | #[cfg(test)] |
| 1055 | mod tests { |
| 1056 | use super::*; |
| 1057 | use std::sync::Arc; |
| 1058 | use std::sync::atomic::{AtomicU32, Ordering}; |
| 1059 | use tokio::sync::Mutex; |
| 1060 | |
| 1061 | /// In-process dispatcher that records what was asked and replies with |
| 1062 | /// canned text. Lets tests verify the round-trip without real network. |
| 1063 | struct StubBridge { |
| 1064 | calls: Arc<Mutex<Vec<RpcRequest>>>, |
| 1065 | canned: Arc<AtomicU32>, |
| 1066 | } |
| 1067 | |
| 1068 | impl StubBridge { |
| 1069 | fn new() -> Self { |
| 1070 | Self { |
| 1071 | calls: Arc::new(Mutex::new(Vec::new())), |
| 1072 | canned: Arc::new(AtomicU32::new(0)), |
| 1073 | } |
| 1074 | } |
| 1075 | } |
| 1076 | |
| 1077 | impl RpcDispatcher for StubBridge { |
| 1078 | fn dispatch<'a>( |
| 1079 | &'a self, |
| 1080 | req: RpcRequest, |
| 1081 | ) -> std::pin::Pin<Box<dyn std::future::Future<Output = RpcResponse> + Send + 'a>> { |
| 1082 | Box::pin(async move { |
| 1083 | self.calls.lock().await.push(req.clone()); |
| 1084 | let n = self.canned.fetch_add(1, Ordering::Relaxed); |
| 1085 | match req { |
| 1086 | RpcRequest::Llm { prompt, .. } | RpcRequest::Rlm { prompt, .. } => { |
| 1087 | RpcResponse::Single(SingleResp { |
| 1088 | text: format!("stub#{n}: {prompt}"), |
| 1089 | error: None, |
| 1090 | }) |
| 1091 | } |
| 1092 | RpcRequest::LlmBatch { prompts, .. } | RpcRequest::RlmBatch { prompts, .. } => { |
| 1093 | let results = prompts |
| 1094 | .into_iter() |
| 1095 | .enumerate() |
| 1096 | .map(|(i, p)| SingleResp { |
| 1097 | text: format!("stub#{n}.{i}: {p}"), |
| 1098 | error: None, |
| 1099 | }) |
| 1100 | .collect(); |
| 1101 | RpcResponse::Batch(BatchResp { results }) |
| 1102 | } |
| 1103 | } |
| 1104 | }) |
| 1105 | } |
| 1106 | } |
| 1107 | |
| 1108 | fn write_temp_context(body: &str) -> std::path::PathBuf { |
| 1109 | let dir = std::env::temp_dir().join("deepseek_repl_runtime_tests"); |
| 1110 | std::fs::create_dir_all(&dir).unwrap(); |
| 1111 | let path = dir.join(format!("ctx_{}_{}.txt", std::process::id(), Uuid::new_v4())); |
| 1112 | std::fs::write(&path, body).unwrap(); |
| 1113 | path |
| 1114 | } |
| 1115 | |
| 1116 | #[tokio::test] |
| 1117 | async fn spawns_and_executes_simple_print() { |
| 1118 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1119 | let round = rt.execute("print('hello world')").await.expect("execute"); |
| 1120 | assert!(round.stdout.contains("hello world")); |
| 1121 | assert!(!round.has_error); |
| 1122 | assert!(round.final_value.is_none()); |
| 1123 | assert_eq!(round.rpc_count, 0); |
| 1124 | rt.shutdown().await; |
| 1125 | } |
| 1126 | |
| 1127 | #[tokio::test] |
| 1128 | async fn non_utf8_stdout_decodes_lossy_and_runtime_survives() { |
| 1129 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1130 | let round = rt |
| 1131 | .execute( |
| 1132 | "import sys\n\ |
| 1133 | sys.stdout.buffer.write(b'bad:\\xff\\n')\n\ |
| 1134 | sys.stdout.buffer.flush()\n\ |
| 1135 | print('after invalid')", |
| 1136 | ) |
| 1137 | .await |
| 1138 | .expect("execute"); |
| 1139 | |
| 1140 | assert!(round.stdout.contains("bad:\u{fffd}"), "{}", round.stdout); |
| 1141 | assert!(round.stdout.contains("after invalid"), "{}", round.stdout); |
| 1142 | rt.shutdown().await; |
| 1143 | } |
| 1144 | |
| 1145 | #[tokio::test] |
| 1146 | async fn variables_persist_across_rounds() { |
| 1147 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1148 | rt.execute("x = [1, 2, 3]").await.expect("r1"); |
| 1149 | rt.execute("x.append(99)").await.expect("r2"); |
| 1150 | let round = rt.execute("print(x)").await.expect("r3"); |
| 1151 | assert!(round.stdout.contains("[1, 2, 3, 99]")); |
| 1152 | rt.shutdown().await; |
| 1153 | } |
| 1154 | |
| 1155 | #[tokio::test] |
| 1156 | async fn imports_persist_across_rounds() { |
| 1157 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1158 | rt.execute("import math").await.expect("r1"); |
| 1159 | let round = rt.execute("print(math.pi)").await.expect("r2"); |
| 1160 | assert!(round.stdout.contains("3.14")); |
| 1161 | rt.shutdown().await; |
| 1162 | } |
| 1163 | |
| 1164 | #[tokio::test] |
| 1165 | async fn context_loads_from_file() { |
| 1166 | let path = write_temp_context("the quick brown fox"); |
| 1167 | let mut rt = PythonRuntime::spawn_with_context(&path) |
| 1168 | .await |
| 1169 | .expect("spawn"); |
| 1170 | let round = rt |
| 1171 | .execute("print(context_meta()['chars'], peek(0, 5))") |
| 1172 | .await |
| 1173 | .expect("execute"); |
| 1174 | assert!(round.stdout.contains("19")); |
| 1175 | assert!(round.stdout.contains("the q")); |
| 1176 | rt.shutdown().await; |
| 1177 | } |
| 1178 | |
| 1179 | #[tokio::test] |
| 1180 | async fn context_aliases_keep_common_content_name_bounded() { |
| 1181 | let path = write_temp_context("aleph-style"); |
| 1182 | let mut rt = PythonRuntime::spawn_with_context(&path) |
| 1183 | .await |
| 1184 | .expect("spawn"); |
| 1185 | let round = rt |
| 1186 | .execute("print(content == _context, 'context' in globals(), 'ctx' in globals())") |
| 1187 | .await |
| 1188 | .expect("execute"); |
| 1189 | assert!(round.stdout.contains("True False False")); |
| 1190 | rt.shutdown().await; |
| 1191 | } |
| 1192 | |
| 1193 | #[tokio::test] |
| 1194 | async fn replacing_context_keeps_kernel_variables_and_refreshes_helpers() { |
| 1195 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1196 | rt.execute("remembered = {'answer': 42}") |
| 1197 | .await |
| 1198 | .expect("seed persistent variable"); |
| 1199 | rt.replace_context("fresh transcript\nwith a needle") |
| 1200 | .await |
| 1201 | .expect("refresh context"); |
| 1202 | |
| 1203 | let round = rt |
| 1204 | .execute( |
| 1205 | "print(remembered['answer'])\n\ |
| 1206 | print(context_meta()['chars'])\n\ |
| 1207 | print(search('needle')[0]['match'])", |
| 1208 | ) |
| 1209 | .await |
| 1210 | .expect("inspect refreshed context"); |
| 1211 | |
| 1212 | assert!(round.stdout.contains("42"), "{}", round.stdout); |
| 1213 | assert!(round.stdout.contains("30"), "{}", round.stdout); |
| 1214 | assert!(round.stdout.contains("needle"), "{}", round.stdout); |
| 1215 | rt.shutdown().await; |
| 1216 | } |
| 1217 | |
| 1218 | #[tokio::test] |
| 1219 | async fn context_chunk_helpers_report_full_coverage() { |
| 1220 | let path = write_temp_context("abcdefghijklmnopqrstuvwxyz"); |
| 1221 | let mut rt = PythonRuntime::spawn_with_context(&path) |
| 1222 | .await |
| 1223 | .expect("spawn"); |
| 1224 | let round = rt |
| 1225 | .execute( |
| 1226 | "chunks = chunk_context(max_chars=10)\n\ |
| 1227 | coverage = chunk_coverage(chunks)\n\ |
| 1228 | print(len(chunks), coverage['covered_chars'], coverage['complete'])", |
| 1229 | ) |
| 1230 | .await |
| 1231 | .expect("execute"); |
| 1232 | assert!(round.stdout.contains("3 26 True"), "{}", round.stdout); |
| 1233 | rt.shutdown().await; |
| 1234 | } |
| 1235 | |
| 1236 | #[tokio::test] |
| 1237 | async fn bounded_input_helpers_work() { |
| 1238 | let path = write_temp_context("alpha\nbeta needle\ngamma needle\nomega"); |
| 1239 | let mut rt = PythonRuntime::spawn_with_context(&path) |
| 1240 | .await |
| 1241 | .expect("spawn"); |
| 1242 | let round = rt |
| 1243 | .execute( |
| 1244 | "meta = context_meta()\n\ |
| 1245 | hits = search('needle', max_hits=1)\n\ |
| 1246 | print(meta['chars'], meta['lines'])\n\ |
| 1247 | print(peek(6, 17))\n\ |
| 1248 | print(peek(1, 3, unit='lines'))\n\ |
| 1249 | print(len(hits), hits[0]['match'], hits[0]['start'])", |
| 1250 | ) |
| 1251 | .await |
| 1252 | .expect("execute"); |
| 1253 | let stdout = round.stdout.replace("\r\n", "\n"); |
| 1254 | assert!(stdout.contains("36 4"), "{stdout}"); |
| 1255 | assert!(stdout.contains("beta needle"), "{stdout}"); |
| 1256 | assert!(stdout.contains("beta needle\ngamma needle"), "{stdout}"); |
| 1257 | assert!(stdout.contains("1 needle 11"), "{stdout}"); |
| 1258 | rt.shutdown().await; |
| 1259 | } |
| 1260 | |
| 1261 | #[tokio::test] |
| 1262 | async fn new_chunk_helper_reports_full_coverage() { |
| 1263 | let path = write_temp_context("abcdefghijklmnopqrstuvwxyz"); |
| 1264 | let mut rt = PythonRuntime::spawn_with_context(&path) |
| 1265 | .await |
| 1266 | .expect("spawn"); |
| 1267 | let round = rt |
| 1268 | .execute( |
| 1269 | "chunks = chunk(max_chars=10)\n\ |
| 1270 | coverage = chunk_coverage(chunks)\n\ |
| 1271 | print(len(chunks), coverage['input_chars'], coverage['covered_chars'], coverage['complete'])", |
| 1272 | ) |
| 1273 | .await |
| 1274 | .expect("execute"); |
| 1275 | assert!(round.stdout.contains("3 26 26 True"), "{}", round.stdout); |
| 1276 | rt.shutdown().await; |
| 1277 | } |
| 1278 | |
| 1279 | #[tokio::test] |
| 1280 | async fn finalize_helper_is_captured_directly() { |
| 1281 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1282 | let round = rt |
| 1283 | .execute("finalize('computed answer', confidence='high')") |
| 1284 | .await |
| 1285 | .expect("execute"); |
| 1286 | assert_eq!(round.final_value.as_deref(), Some("computed answer")); |
| 1287 | assert_eq!( |
| 1288 | round.final_json.as_ref().and_then(Value::as_str), |
| 1289 | Some("computed answer") |
| 1290 | ); |
| 1291 | assert_eq!( |
| 1292 | round.final_confidence.as_ref().and_then(Value::as_str), |
| 1293 | Some("high") |
| 1294 | ); |
| 1295 | rt.shutdown().await; |
| 1296 | } |
| 1297 | |
| 1298 | #[tokio::test] |
| 1299 | async fn finalize_preserves_json_values_for_handles() { |
| 1300 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1301 | let round = rt |
| 1302 | .execute("finalize({'answer': 42, 'items': ['a', 'b']})") |
| 1303 | .await |
| 1304 | .expect("execute"); |
| 1305 | |
| 1306 | assert_eq!( |
| 1307 | round.final_value.as_deref(), |
| 1308 | Some(r#"{"answer":42,"items":["a","b"]}"#) |
| 1309 | ); |
| 1310 | assert_eq!( |
| 1311 | round.final_json, |
| 1312 | Some(serde_json::json!({"answer": 42, "items": ["a", "b"]})) |
| 1313 | ); |
| 1314 | rt.shutdown().await; |
| 1315 | } |
| 1316 | |
| 1317 | #[tokio::test] |
| 1318 | async fn sub_query_accepts_timeout_keyword_for_agent_guesses() { |
| 1319 | let bridge = StubBridge::new(); |
| 1320 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1321 | let round = rt |
| 1322 | .run( |
| 1323 | "answer = sub_query('summarize', timeout_secs=2)\nprint(answer)", |
| 1324 | Some(&bridge), |
| 1325 | ) |
| 1326 | .await |
| 1327 | .expect("execute"); |
| 1328 | |
| 1329 | assert!(!round.has_error, "{}", round.stdout); |
| 1330 | assert!( |
| 1331 | round.stdout.contains("stub#0: summarize"), |
| 1332 | "{}", |
| 1333 | round.stdout |
| 1334 | ); |
| 1335 | rt.shutdown().await; |
| 1336 | } |
| 1337 | |
| 1338 | #[tokio::test] |
| 1339 | async fn rlm_context_runtime_has_no_fixed_round_timeout() { |
| 1340 | let path = write_temp_context("long input"); |
| 1341 | let rt = PythonRuntime::spawn_with_context(&path) |
| 1342 | .await |
| 1343 | .expect("spawn"); |
| 1344 | assert!( |
| 1345 | rt.round_timeout().is_none(), |
| 1346 | "RLM context runs must not inherit the old 180s REPL round timeout" |
| 1347 | ); |
| 1348 | rt.shutdown().await; |
| 1349 | } |
| 1350 | |
| 1351 | #[tokio::test] |
| 1352 | async fn inline_runtime_keeps_bounded_round_timeout() { |
| 1353 | let rt = PythonRuntime::new().await.expect("spawn"); |
| 1354 | assert_eq!(rt.round_timeout(), Some(ROUND_TIMEOUT)); |
| 1355 | rt.shutdown().await; |
| 1356 | } |
| 1357 | |
| 1358 | #[tokio::test] |
| 1359 | async fn legacy_final_is_captured() { |
| 1360 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1361 | let round = rt |
| 1362 | .execute("FINAL('the answer is 42')") |
| 1363 | .await |
| 1364 | .expect("execute"); |
| 1365 | assert_eq!(round.final_value.as_deref(), Some("the answer is 42")); |
| 1366 | rt.shutdown().await; |
| 1367 | } |
| 1368 | |
| 1369 | #[tokio::test] |
| 1370 | async fn legacy_final_var_is_captured() { |
| 1371 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1372 | rt.execute("answer = 'computed'").await.expect("r1"); |
| 1373 | let round = rt.execute("FINAL_VAR('answer')").await.expect("r2"); |
| 1374 | assert_eq!(round.final_value.as_deref(), Some("computed")); |
| 1375 | rt.shutdown().await; |
| 1376 | } |
| 1377 | |
| 1378 | #[tokio::test] |
| 1379 | async fn errors_are_reported_without_killing_runtime() { |
| 1380 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1381 | let r1 = rt.execute("raise ValueError('boom')").await.expect("r1"); |
| 1382 | assert!(r1.has_error); |
| 1383 | assert!(r1.full_stdout.contains("boom") || r1.stdout.contains("boom")); |
| 1384 | // The runtime is still alive — next round should work. |
| 1385 | let r2 = rt.execute("print('still here')").await.expect("r2"); |
| 1386 | assert!(r2.stdout.contains("still here")); |
| 1387 | rt.shutdown().await; |
| 1388 | } |
| 1389 | |
| 1390 | #[tokio::test] |
| 1391 | async fn rpc_dispatcher_round_trips_llm_query() { |
| 1392 | let bridge = StubBridge::new(); |
| 1393 | let calls = Arc::clone(&bridge.calls); |
| 1394 | |
| 1395 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1396 | let round = rt |
| 1397 | .run("print(llm_query('hello'))", Some(&bridge)) |
| 1398 | .await |
| 1399 | .expect("execute"); |
| 1400 | assert!( |
| 1401 | round.stdout.contains("stub#0: hello"), |
| 1402 | "stdout: {:?}", |
| 1403 | round.stdout |
| 1404 | ); |
| 1405 | assert_eq!(round.rpc_count, 1); |
| 1406 | |
| 1407 | let recorded = calls.lock().await; |
| 1408 | assert_eq!(recorded.len(), 1); |
| 1409 | match &recorded[0] { |
| 1410 | RpcRequest::Llm { prompt, .. } => assert_eq!(prompt, "hello"), |
| 1411 | other => panic!("expected Llm request, got {other:?}"), |
| 1412 | } |
| 1413 | drop(recorded); |
| 1414 | rt.shutdown().await; |
| 1415 | } |
| 1416 | |
| 1417 | #[tokio::test] |
| 1418 | async fn rpc_dispatcher_round_trips_sub_query_alias() { |
| 1419 | let bridge = StubBridge::new(); |
| 1420 | let calls = Arc::clone(&bridge.calls); |
| 1421 | |
| 1422 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1423 | let round = rt |
| 1424 | .run("print(sub_query('hello from sub'))", Some(&bridge)) |
| 1425 | .await |
| 1426 | .expect("execute"); |
| 1427 | assert!( |
| 1428 | round.stdout.contains("stub#0: hello from sub"), |
| 1429 | "stdout: {:?}", |
| 1430 | round.stdout |
| 1431 | ); |
| 1432 | assert_eq!(round.rpc_count, 1); |
| 1433 | |
| 1434 | let recorded = calls.lock().await; |
| 1435 | assert_eq!(recorded.len(), 1); |
| 1436 | match &recorded[0] { |
| 1437 | RpcRequest::Llm { prompt, .. } => assert_eq!(prompt, "hello from sub"), |
| 1438 | other => panic!("expected Llm request, got {other:?}"), |
| 1439 | } |
| 1440 | drop(recorded); |
| 1441 | rt.shutdown().await; |
| 1442 | } |
| 1443 | |
| 1444 | #[tokio::test] |
| 1445 | async fn rpc_dispatcher_round_trips_batch() { |
| 1446 | let bridge = StubBridge::new(); |
| 1447 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1448 | let round = rt |
| 1449 | .run( |
| 1450 | "outs = llm_query_batched(['a','b','c'], dependency_mode='independent', safety_note='same independent classification')\n\ |
| 1451 | print('|'.join(outs))", |
| 1452 | Some(&bridge), |
| 1453 | ) |
| 1454 | .await |
| 1455 | .expect("execute"); |
| 1456 | assert!(round.stdout.contains("stub#0.0: a")); |
| 1457 | assert!(round.stdout.contains("stub#0.1: b")); |
| 1458 | assert!(round.stdout.contains("stub#0.2: c")); |
| 1459 | assert_eq!(round.rpc_count, 1); |
| 1460 | rt.shutdown().await; |
| 1461 | } |
| 1462 | |
| 1463 | #[tokio::test] |
| 1464 | async fn batched_helpers_require_independence_declaration() { |
| 1465 | let bridge = StubBridge::new(); |
| 1466 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1467 | let round = rt |
| 1468 | .run( |
| 1469 | "outs = sub_query_batch('summarize', [{'text': 'a'}, {'text': 'b'}])\n\ |
| 1470 | print(outs[0])", |
| 1471 | Some(&bridge), |
| 1472 | ) |
| 1473 | .await |
| 1474 | .expect("execute"); |
| 1475 | |
| 1476 | assert!( |
| 1477 | round.stdout.contains("dependency_mode='independent'"), |
| 1478 | "{}", |
| 1479 | round.stdout |
| 1480 | ); |
| 1481 | assert_eq!(round.rpc_count, 0); |
| 1482 | rt.shutdown().await; |
| 1483 | } |
| 1484 | |
| 1485 | #[tokio::test] |
| 1486 | async fn dependent_batch_mode_points_to_sequence_helper() { |
| 1487 | let bridge = StubBridge::new(); |
| 1488 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1489 | let round = rt |
| 1490 | .run( |
| 1491 | "outs = llm_query_batched(['migrate A', 'migrate B'], dependency_mode='sequential')\n\ |
| 1492 | print(outs[0])", |
| 1493 | Some(&bridge), |
| 1494 | ) |
| 1495 | .await |
| 1496 | .expect("execute"); |
| 1497 | |
| 1498 | assert!( |
| 1499 | round.stdout.contains("sub_query_sequence"), |
| 1500 | "{}", |
| 1501 | round.stdout |
| 1502 | ); |
| 1503 | assert_eq!(round.rpc_count, 0); |
| 1504 | rt.shutdown().await; |
| 1505 | } |
| 1506 | |
| 1507 | #[tokio::test] |
| 1508 | async fn sub_query_sequence_feeds_prior_result_into_next_prompt() { |
| 1509 | let bridge = StubBridge::new(); |
| 1510 | let calls = Arc::clone(&bridge.calls); |
| 1511 | |
| 1512 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1513 | let round = rt |
| 1514 | .run( |
| 1515 | "outs = sub_query_sequence('process this step', [{'text': 'A'}, {'text': 'B'}])\n\ |
| 1516 | print(len(outs))", |
| 1517 | Some(&bridge), |
| 1518 | ) |
| 1519 | .await |
| 1520 | .expect("execute"); |
| 1521 | |
| 1522 | assert!(round.stdout.contains("2"), "{}", round.stdout); |
| 1523 | assert_eq!(round.rpc_count, 2); |
| 1524 | |
| 1525 | let recorded = calls.lock().await; |
| 1526 | assert_eq!(recorded.len(), 2); |
| 1527 | let second_prompt = match &recorded[1] { |
| 1528 | RpcRequest::Llm { prompt, .. } => prompt, |
| 1529 | other => panic!("expected second Llm request, got {other:?}"), |
| 1530 | }; |
| 1531 | assert!(second_prompt.contains("--- dependency_state step 1/2 ---")); |
| 1532 | assert!(second_prompt.contains("stub#0: process this step")); |
| 1533 | drop(recorded); |
| 1534 | rt.shutdown().await; |
| 1535 | } |
| 1536 | |
| 1537 | #[tokio::test] |
| 1538 | async fn no_dispatcher_returns_unavailable_sentinel() { |
| 1539 | let mut rt = PythonRuntime::new().await.expect("spawn"); |
| 1540 | let round = rt.execute("print(llm_query('hi'))").await.expect("execute"); |
| 1541 | assert!( |
| 1542 | round.stdout.contains("[llm_query error:") || round.stdout.contains("no LLM bridge"), |
| 1543 | "stdout: {:?}", |
| 1544 | round.stdout |
| 1545 | ); |
| 1546 | rt.shutdown().await; |
| 1547 | } |
| 1548 | |
| 1549 | #[test] |
| 1550 | fn truncate_keeps_short_unchanged() { |
| 1551 | assert_eq!(truncate_stdout("hello", 100), "hello"); |
| 1552 | } |
| 1553 | |
| 1554 | #[test] |
| 1555 | fn truncate_clips_long() { |
| 1556 | let long = "a".repeat(10_000); |
| 1557 | let out = truncate_stdout(&long, 1024); |
| 1558 | assert!(out.len() < 1500); |
| 1559 | assert!(out.contains("truncated")); |
| 1560 | } |
| 1561 | } |
| 1562 |