返回 CodeWhale
turn.rs
根目录 / crates / tui / src / rlm / turn.rs
1 //! RLM turn loop — paper Algorithm 1 driven over a long-lived Python
2 //! subprocess + stdin/stdout RPC bridge (no HTTP sidecar).
3
4 use std::path::PathBuf;
5 use std::sync::Arc;
6 use std::time::{Duration, Instant};
7
8 use tokio::sync::mpsc;
9 use uuid::Uuid;
10
11 use crate::client::CodewhaleClient;
12 use crate::core::events::Event;
13 use crate::repl::PythonRuntime;
14 use codewhale_models::{
15 ContentBlock, Message, MessageRequest, SystemPrompt, Usage, is_incomplete_stop_reason,
16 stop_reason_detail,
17 };
18
19 use super::bridge::{RlmBridge, RlmLlmClient, RlmUsageAccumulator};
20 use super::prompt::rlm_system_prompt;
21 use codewhale_models::Role;
22
23 // ---------------------------------------------------------------------------
24 // Constants
25 // ---------------------------------------------------------------------------
26
27 /// Maximum number of RLM iterations before the loop gives up.
28 const MAX_RLM_ITERATIONS: u32 = 25;
29 /// Max consecutive rounds where the model returns no `repl` fence before we
30 /// hard-fail. The paper requires `code → REPL → Final`; anything else is
31 /// not the RLM contract.
32 const MAX_CONSECUTIVE_NO_CODE: u32 = 3;
33 /// Max output tokens for the root LLM — it just needs to generate code.
34 /// Max chars of stdout shown as metadata to the root LLM in next iteration.
35 const STDOUT_METADATA_PREVIEW_LEN: usize = 800;
36 /// Max chars of `context` shown as a preview in the metadata.
37 const PROMPT_PREVIEW_LEN: usize = 500;
38 /// Temperature for root LLM calls.
39 /// Bound on conversation history we keep across iterations.
40 const MAX_HISTORY_MESSAGES: usize = 20;
41
42 // ---------------------------------------------------------------------------
43 // Public API
44 // ---------------------------------------------------------------------------
45
46 /// How an RLM turn ended.
47 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
48 pub enum RlmTermination {
49 /// `FINAL(value)` was called inside the REPL or `FINAL(...)` appeared
50 /// at the top of the model's response on its own line.
51 Final,
52 /// The model failed to emit a `repl` block for too many rounds in a
53 /// row. The accumulated last response text is surfaced as the answer
54 /// rather than being thrown away.
55 NoCode,
56 /// Iteration cap reached without `FINAL`.
57 Exhausted,
58 /// Hard error — LLM call failed, REPL crashed, timeout.
59 Error,
60 }
61
62 /// Per-round trace entry. Surfaced in the tool result so the user can see
63 /// exactly what the sub-agent did.
64 #[derive(Debug, Clone)]
65 pub struct RlmRoundTrace {
66 pub round: u32,
67 pub code_summary: String,
68 pub stdout_preview: String,
69 pub had_error: bool,
70 pub rpc_count: u32,
71 pub elapsed_ms: u64,
72 }
73
74 /// Result of an RLM turn.
75 #[derive(Debug, Clone)]
76 pub struct RlmTurnResult {
77 pub answer: String,
78 pub iterations: u32,
79 pub duration: Duration,
80 pub error: Option<String>,
81 pub usage: Usage,
82 /// One exact frozen route/quote receipt per admitted provider request.
83 /// Distinct calls are never coalesced, even when they share a route.
84 pub routed_usage: Vec<crate::cost_status::RuntimeUsageRecord>,
85 /// Exact routes for provider-success responses that omitted authoritative
86 /// usage metadata.
87 pub routed_usage_drop_records: Vec<crate::cost_status::RuntimeUsageDropRecord>,
88 pub routed_usage_dropped_records: u64,
89 pub termination: RlmTermination,
90 /// Per-round trace. Empty when the loop never reached the REPL.
91 pub trace: Vec<RlmRoundTrace>,
92 /// Total sub-LLM RPCs made by the sub-agent (sum of `rpc_count` across
93 /// rounds). Useful for verifying that the model engaged with `context`
94 /// rather than answering directly.
95 pub total_rpcs: u32,
96 }
97
98 /// Run a full RLM turn. `prompt` is loaded into the REPL as `context`; it
99 /// never enters the root LLM's window.
100 pub async fn run_rlm_turn(
101 client: &CodewhaleClient,
102 model: String,
103 prompt: String,
104 child_model: String,
105 tx_event: mpsc::Sender<Event>,
106 max_depth: u32,
107 ) -> RlmTurnResult {
108 run_rlm_turn_inner(
109 Arc::new(client.clone()),
110 model,
111 prompt,
112 None,
113 child_model,
114 tx_event,
115 max_depth,
116 )
117 .await
118 }
119
120 /// Variant that also passes a small `root_prompt` (the user-facing task)
121 /// shown to the root LLM each iteration so it remembers its objective.
122 pub async fn run_rlm_turn_with_root(
123 client: &CodewhaleClient,
124 model: String,
125 prompt: String,
126 root_prompt: Option<String>,
127 child_model: String,
128 tx_event: mpsc::Sender<Event>,
129 max_depth: u32,
130 ) -> RlmTurnResult {
131 run_rlm_turn_inner(
132 Arc::new(client.clone()),
133 model,
134 prompt,
135 root_prompt,
136 child_model,
137 tx_event,
138 max_depth,
139 )
140 .await
141 }
142
143 /// Inner entry point — also used by the bridge when it recurses. Returns
144 /// a boxed future to break the recursive opaque-future-type cycle:
145 /// `run_rlm_turn_inner` → `RlmBridge::dispatch` → `run_rlm_turn_inner`.
146 pub(crate) fn run_rlm_turn_inner(
147 client: Arc<dyn RlmLlmClient>,
148 model: String,
149 prompt: String,
150 root_prompt: Option<String>,
151 child_model: String,
152 tx_event: mpsc::Sender<Event>,
153 max_depth: u32,
154 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = RlmTurnResult> + Send>> {
155 run_rlm_turn_inner_with_usage(
156 client,
157 model,
158 prompt,
159 root_prompt,
160 child_model,
161 tx_event,
162 max_depth,
163 RlmUsageAccumulator::new(),
164 )
165 }
166
167 /// Recursive entry point that keeps one pre-dispatch receipt bound across the
168 /// entire nested/batched RLM tree.
169 pub(crate) fn run_rlm_turn_inner_with_usage(
170 client: Arc<dyn RlmLlmClient>,
171 model: String,
172 prompt: String,
173 root_prompt: Option<String>,
174 child_model: String,
175 tx_event: mpsc::Sender<Event>,
176 max_depth: u32,
177 usage: RlmUsageAccumulator,
178 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = RlmTurnResult> + Send>> {
179 Box::pin(async move {
180 let mut result = run_rlm_turn_impl(
181 client,
182 model,
183 prompt,
184 root_prompt,
185 child_model,
186 tx_event,
187 max_depth,
188 usage.clone(),
189 )
190 .await;
191 let snapshot = usage.snapshot().await;
192 result.usage = snapshot.usage;
193 result.routed_usage = snapshot.records;
194 result.routed_usage_drop_records = snapshot.drop_records;
195 result.routed_usage_dropped_records = snapshot.dropped_records;
196 result
197 })
198 }
199
200 /// RLM turns are long-running background-style work. Do not kill the whole
201 /// turn with the old fixed 180s wall-clock cap; per-request cancellation still
202 /// comes from the parent turn token and the user can cancel from the TUI.
203 fn turn_timeout() -> Option<Duration> {
204 None
205 }
206
207 // ---------------------------------------------------------------------------
208 // Implementation
209 // ---------------------------------------------------------------------------
210
211 async fn run_rlm_turn_impl(
212 client: Arc<dyn RlmLlmClient>,
213 model: String,
214 prompt: String,
215 root_prompt: Option<String>,
216 child_model: String,
217 tx_event: mpsc::Sender<Event>,
218 max_depth: u32,
219 routed_usage: RlmUsageAccumulator,
220 ) -> RlmTurnResult {
221 let start = Instant::now();
222 let mut total_usage = Usage::default();
223 let mut trace: Vec<RlmRoundTrace> = Vec::new();
224 let mut total_rpcs: u32 = 0;
225
226 // 1. Stage `context` to a temp file. The REPL reads it on bootstrap so
227 // the big string never enters the process command line and doesn't
228 // show up in `ps`.
229 let ctx_path = match write_context_file(&prompt) {
230 Ok(p) => p,
231 Err(e) => {
232 return RlmTurnResult {
233 answer: String::new(),
234 iterations: 0,
235 duration: start.elapsed(),
236 error: Some(format!("rlm: failed to stage context: {e}")),
237 usage: total_usage,
238 routed_usage: Vec::new(),
239 routed_usage_drop_records: Vec::new(),
240 routed_usage_dropped_records: 0,
241 termination: RlmTermination::Error,
242 trace,
243 total_rpcs,
244 };
245 }
246 };
247
248 // 2. Spawn the long-lived REPL.
249 let mut repl = match PythonRuntime::spawn_with_context(&ctx_path).await {
250 Ok(rt) => rt,
251 Err(e) => {
252 let _ = tokio::fs::remove_file(&ctx_path).await;
253 return RlmTurnResult {
254 answer: String::new(),
255 iterations: 0,
256 duration: start.elapsed(),
257 error: Some(format!("rlm: failed to spawn REPL: {e}")),
258 usage: total_usage,
259 routed_usage: Vec::new(),
260 routed_usage_drop_records: Vec::new(),
261 routed_usage_dropped_records: 0,
262 termination: RlmTermination::Error,
263 trace,
264 total_rpcs,
265 };
266 }
267 };
268
269 // 3. Build the bridge that services llm_query / rlm_query RPCs.
270 let bridge = RlmBridge::with_usage_accumulator(
271 Arc::clone(&client),
272 child_model.clone(),
273 max_depth,
274 routed_usage.clone(),
275 );
276
277 let _ = tx_event
278 .send(Event::status(format!(
279 "RLM: spawned Python REPL (root={model}, child={child_model}, max_depth={max_depth}, ctx={} chars)",
280 prompt.chars().count()
281 )))
282 .await;
283
284 // 4. Build initial metadata-only history.
285 let system = rlm_system_prompt();
286 let mut messages: Vec<Message> = vec![build_metadata_message(
287 &prompt,
288 root_prompt.as_deref(),
289 0,
290 None,
291 None,
292 )];
293
294 let mut consecutive_no_code: u32 = 0;
295 let mut consecutive_empty_rounds: u32 = 0;
296 let mut last_response_text = String::new();
297
298 let result = 'turn: {
299 for iteration in 0..MAX_RLM_ITERATIONS {
300 if let Some(timeout) = turn_timeout()
301 && start.elapsed() > timeout
302 {
303 break 'turn RlmTurnResult {
304 answer: String::new(),
305 iterations: iteration,
306 duration: start.elapsed(),
307 error: Some(format!("RLM turn timed out after {}s", timeout.as_secs())),
308 usage: total_usage,
309 routed_usage: Vec::new(),
310 routed_usage_drop_records: Vec::new(),
311 routed_usage_dropped_records: 0,
312 termination: RlmTermination::Error,
313 trace: trace.clone(),
314 total_rpcs,
315 };
316 }
317
318 let _ = tx_event
319 .send(Event::status(format!(
320 "RLM iteration {}/{}",
321 iteration + 1,
322 MAX_RLM_ITERATIONS
323 )))
324 .await;
325
326 // 4a. Root LLM generates code from metadata-only context.
327 let request_route = client.effective_route_envelope(&model, chrono::Utc::now());
328 let reservation = match routed_usage.reserve(request_route.clone()).await {
329 Ok(reservation) => reservation,
330 Err(error) => {
331 break 'turn RlmTurnResult {
332 answer: String::new(),
333 iterations: iteration,
334 duration: start.elapsed(),
335 error: Some(error),
336 usage: total_usage,
337 routed_usage: Vec::new(),
338 routed_usage_drop_records: Vec::new(),
339 routed_usage_dropped_records: 0,
340 termination: RlmTermination::Error,
341 trace: trace.clone(),
342 total_rpcs,
343 };
344 }
345 };
346 let request = build_root_request(
347 &model,
348 &messages,
349 &system,
350 client.effective_max_output_tokens(&request_route.model),
351 );
352
353 let response = match client.create_message_boxed(request).await {
354 Ok(r) => r,
355 Err(e) => {
356 routed_usage.cancel(reservation, false).await;
357 break 'turn RlmTurnResult {
358 answer: String::new(),
359 iterations: iteration + 1,
360 duration: start.elapsed(),
361 error: Some(format!("Root LLM call failed: {e}")),
362 usage: total_usage,
363 routed_usage: Vec::new(),
364 routed_usage_drop_records: Vec::new(),
365 routed_usage_dropped_records: 0,
366 termination: RlmTermination::Error,
367 trace: trace.clone(),
368 total_rpcs,
369 };
370 }
371 };
372
373 // Preserve billed usage even when the response is incomplete and
374 // its partial FINAL/REPL output is rejected below.
375 routed_usage
376 .settle_provider_success(reservation, &response.usage)
377 .await;
378 super::add_usage_with_prompt_cache(&mut total_usage, &response.usage);
379
380 if is_incomplete_stop_reason(response.stop_reason.as_deref()) {
381 let reason = stop_reason_detail(response.stop_reason.as_deref());
382 break 'turn RlmTurnResult {
383 answer: String::new(),
384 iterations: iteration + 1,
385 duration: start.elapsed(),
386 error: Some(format!(
387 "RLM root model response incomplete: provider stop reason `{reason}`; partial FINAL/REPL output was not accepted."
388 )),
389 usage: total_usage,
390 routed_usage: Vec::new(),
391 routed_usage_drop_records: Vec::new(),
392 routed_usage_dropped_records: 0,
393 termination: RlmTermination::Error,
394 trace: trace.clone(),
395 total_rpcs,
396 };
397 }
398
399 let response_text = extract_text_blocks(&response.content);
400 last_response_text = response_text.clone();
401
402 // 4b. Top-level FINAL(...) lets the model close out without
403 // touching the REPL — but only if it has done some work
404 // (non-zero rpc_count) on a prior round. Otherwise it's a
405 // shortcut and we reject it.
406 if let Some(final_val) = parse_text_final(&response_text) {
407 if total_rpcs == 0 {
408 // Discard the top-level FINAL — the model is bypassing
409 // the loop. Force it to use the REPL by appending a
410 // strict reminder.
411 consecutive_no_code = consecutive_no_code.saturating_add(1);
412 if consecutive_no_code >= MAX_CONSECUTIVE_NO_CODE {
413 break 'turn RlmTurnResult {
414 answer: final_val,
415 iterations: iteration + 1,
416 duration: start.elapsed(),
417 error: None,
418 usage: total_usage,
419 routed_usage: Vec::new(),
420 routed_usage_drop_records: Vec::new(),
421 routed_usage_dropped_records: 0,
422 termination: RlmTermination::NoCode,
423 trace: trace.clone(),
424 total_rpcs,
425 };
426 }
427 messages.push(Message {
428 role: Role::Assistant,
429 content: vec![ContentBlock::Text {
430 text: response_text.clone(),
431 cache_control: None,
432 }],
433 });
434 messages.push(Message {
435 role: Role::User,
436 content: vec![ContentBlock::Text {
437 text: "You called FINAL(...) without ever running a ```repl block. \
438 That defeats the recursive language model — you're guessing \
439 from the preview alone. Emit a ```repl block now that uses \
440 `llm_query`, `sub_query_sequence`, or an explicitly independent \
441 `llm_query_batched(..., dependency_mode=\"independent\")` against \
442 `context` to actually compute the answer."
443 .to_string(),
444 cache_control: None,
445 }],
446 });
447 continue;
448 }
449 let _ = tx_event
450 .send(Event::status(
451 "RLM: FINAL detected in response text".to_string(),
452 ))
453 .await;
454 break 'turn RlmTurnResult {
455 answer: final_val,
456 iterations: iteration + 1,
457 duration: start.elapsed(),
458 error: None,
459 usage: total_usage,
460 routed_usage: Vec::new(),
461 routed_usage_drop_records: Vec::new(),
462 routed_usage_dropped_records: 0,
463 termination: RlmTermination::Final,
464 trace: trace.clone(),
465 total_rpcs,
466 };
467 }
468
469 // 4c. Extract a ```repl block.
470 let code = extract_repl_code(&response_text);
471 let code_to_run = match code {
472 Some(c) => {
473 consecutive_no_code = 0;
474 c
475 }
476 None => {
477 consecutive_no_code = consecutive_no_code.saturating_add(1);
478 if consecutive_no_code >= MAX_CONSECUTIVE_NO_CODE {
479 break 'turn RlmTurnResult {
480 answer: response_text,
481 iterations: iteration + 1,
482 duration: start.elapsed(),
483 error: Some(format!(
484 "RLM: model failed to emit ```repl after {MAX_CONSECUTIVE_NO_CODE} consecutive rounds"
485 )),
486 usage: total_usage,
487 routed_usage: Vec::new(),
488 routed_usage_drop_records: Vec::new(),
489 routed_usage_dropped_records: 0,
490 termination: RlmTermination::NoCode,
491 trace: trace.clone(),
492 total_rpcs,
493 };
494 }
495 messages.push(Message {
496 role: Role::Assistant,
497 content: vec![ContentBlock::Text {
498 text: response_text.clone(),
499 cache_control: None,
500 }],
501 });
502 messages.push(Message {
503 role: Role::User,
504 content: vec![ContentBlock::Text {
505 text: "Reminder: emit Python inside a ```repl … ``` fence. \
506 Use `llm_query`, `sub_query_sequence`, or \
507 `llm_query_batched(..., dependency_mode=\"independent\")` to \
508 process `context` and call `FINAL(value)` when done."
509 .to_string(),
510 cache_control: None,
511 }],
512 });
513 continue;
514 }
515 };
516
517 let _ = tx_event
518 .send(Event::MessageDelta {
519 index: iteration as usize,
520 content: format!(
521 "\n[RLM round {} — code]\n```repl\n{code_to_run}\n```\n",
522 iteration + 1
523 ),
524 })
525 .await;
526
527 // 4d. Execute the code in the REPL with the bridge servicing
528 // llm_query / rlm_query callbacks.
529 let round = match repl.run(&code_to_run, Some(&bridge)).await {
530 Ok(r) => r,
531 Err(e) => {
532 break 'turn RlmTurnResult {
533 answer: String::new(),
534 iterations: iteration + 1,
535 duration: start.elapsed(),
536 error: Some(format!("REPL execution failed: {e}")),
537 usage: total_usage,
538 routed_usage: Vec::new(),
539 routed_usage_drop_records: Vec::new(),
540 routed_usage_dropped_records: 0,
541 termination: RlmTermination::Error,
542 trace: trace.clone(),
543 total_rpcs,
544 };
545 }
546 };
547
548 total_rpcs = total_rpcs.saturating_add(round.rpc_count);
549
550 // Trace this round.
551 let stdout_preview = truncate_text(round.stdout.trim(), STDOUT_METADATA_PREVIEW_LEN);
552 trace.push(RlmRoundTrace {
553 round: iteration + 1,
554 code_summary: summarize_code(&code_to_run),
555 stdout_preview: stdout_preview.clone(),
556 had_error: round.has_error,
557 rpc_count: round.rpc_count,
558 elapsed_ms: round.elapsed.as_millis() as u64,
559 });
560
561 let _ = tx_event
562 .send(Event::status(format!(
563 "RLM round {}: {} bytes stdout, {} sub-LLM call(s){}",
564 iteration + 1,
565 round.full_stdout.len(),
566 round.rpc_count,
567 if round.has_error { " (error)" } else { "" },
568 )))
569 .await;
570
571 // 4e. FINAL detection.
572 if let Some(final_val) = round.final_value.clone() {
573 let _ = tx_event
574 .send(Event::status(
575 "RLM: FINAL detected in REPL, ending loop".to_string(),
576 ))
577 .await;
578 break 'turn RlmTurnResult {
579 answer: final_val,
580 iterations: iteration + 1,
581 duration: start.elapsed(),
582 error: None,
583 usage: total_usage,
584 routed_usage: Vec::new(),
585 routed_usage_drop_records: Vec::new(),
586 routed_usage_dropped_records: 0,
587 termination: RlmTermination::Final,
588 trace: trace.clone(),
589 total_rpcs,
590 };
591 }
592
593 // 4e+. Empty/no-op guard — same contract as the normal Agent REPL path.
594 // If the block produced no stdout, no RPC, and no finalize(), tell the
595 // model plainly and count consecutive empties to avoid an infinite loop.
596 let is_empty_round = !round.has_error
597 && round.stdout.trim().is_empty()
598 && round.stderr.trim().is_empty()
599 && round.rpc_count == 0
600 && round.final_value.is_none();
601 if is_empty_round {
602 consecutive_empty_rounds = consecutive_empty_rounds.saturating_add(1);
603 let empty_feedback = if consecutive_empty_rounds >= MAX_CONSECUTIVE_NO_CODE {
604 format!(
605 "Your emitted ```repl block (round {}) result: no observable output — print something, call a helper, or stop emitting REPL blocks and answer. No output for {consecutive_empty_rounds} consecutive rounds; stopping empty loop.",
606 iteration + 1
607 )
608 } else {
609 format!(
610 "Your emitted ```repl block (round {}) result: no observable output — print something, call a helper, or stop emitting REPL blocks and answer",
611 iteration + 1
612 )
613 };
614 messages.push(Message {
615 role: Role::Assistant,
616 content: vec![ContentBlock::Text {
617 text: format!("```repl\n{code_to_run}\n```"),
618 cache_control: None,
619 }],
620 });
621 messages.push(build_metadata_message(
622 &prompt,
623 root_prompt.as_deref(),
624 iteration + 1,
625 Some(&code_to_run),
626 Some(&empty_feedback),
627 ));
628 if consecutive_empty_rounds >= MAX_CONSECUTIVE_NO_CODE {
629 break 'turn RlmTurnResult {
630 answer: last_response_text.clone(),
631 iterations: iteration + 1,
632 duration: start.elapsed(),
633 error: Some(format!(
634 "RLM: {MAX_CONSECUTIVE_NO_CODE} consecutive empty REPL rounds"
635 )),
636 usage: total_usage,
637 routed_usage: Vec::new(),
638 routed_usage_drop_records: Vec::new(),
639 routed_usage_dropped_records: 0,
640 termination: RlmTermination::NoCode,
641 trace: trace.clone(),
642 total_rpcs,
643 };
644 }
645 if messages.len() > MAX_HISTORY_MESSAGES {
646 let drop_from = messages.len() - MAX_HISTORY_MESSAGES + 1;
647 let mut kept = vec![messages[0].clone()];
648 kept.extend(messages.drain(drop_from..));
649 messages = kept;
650 }
651 continue;
652 } else {
653 consecutive_empty_rounds = 0;
654 }
655
656 // Provenance: make round feedback unambiguous — it is always the
657 // assistant's own emitted block, never the user's.
658 let provenance_prefix = format!(
659 "Your emitted ```repl block (round {}) result:",
660 iteration + 1
661 );
662 let stdout_for_feedback = if round.has_error {
663 format!(
664 "{provenance_prefix} error\nstdout:\n{}\nstderr:\n{}",
665 round.stdout, round.stderr
666 )
667 } else if round.stdout.trim().is_empty() && round.rpc_count == 0 {
668 format!(
669 "{provenance_prefix} no output — block produced no observable output — print something, call a helper, or stop emitting REPL blocks and answer\nstdout:\n{}\n[{} child query RPC(s)]",
670 round.stdout, round.rpc_count
671 )
672 } else {
673 format!(
674 "{provenance_prefix}\n[{} child query RPC(s)]\n{}",
675 round.rpc_count, round.stdout
676 )
677 };
678 let stdout_preview_for_next =
679 truncate_text(stdout_for_feedback.trim(), STDOUT_METADATA_PREVIEW_LEN);
680
681 // 4f. Build metadata for next iteration.
682 messages.push(Message {
683 role: Role::Assistant,
684 content: vec![ContentBlock::Text {
685 text: format!("```repl\n{code_to_run}\n```"),
686 cache_control: None,
687 }],
688 });
689 messages.push(build_metadata_message(
690 &prompt,
691 root_prompt.as_deref(),
692 iteration + 1,
693 Some(&code_to_run),
694 Some(&stdout_preview_for_next),
695 ));
696
697 if messages.len() > MAX_HISTORY_MESSAGES {
698 let drop_from = messages.len() - MAX_HISTORY_MESSAGES + 1;
699 let mut kept = vec![messages[0].clone()];
700 kept.extend(messages.drain(drop_from..));
701 messages = kept;
702 }
703 }
704
705 let _ = last_response_text;
706 RlmTurnResult {
707 answer: String::new(),
708 iterations: MAX_RLM_ITERATIONS,
709 duration: start.elapsed(),
710 error: Some(format!(
711 "RLM loop exhausted after {MAX_RLM_ITERATIONS} iterations without FINAL"
712 )),
713 usage: total_usage,
714 routed_usage: Vec::new(),
715 routed_usage_drop_records: Vec::new(),
716 routed_usage_dropped_records: 0,
717 termination: RlmTermination::Exhausted,
718 trace: trace.clone(),
719 total_rpcs,
720 }
721 };
722
723 repl.shutdown().await;
724 result
725 }
726
727 // ---------------------------------------------------------------------------
728 // Helpers
729 // ---------------------------------------------------------------------------
730
731 fn write_context_file(prompt: &str) -> std::io::Result<PathBuf> {
732 let dir = std::env::temp_dir().join("deepseek_rlm_ctx");
733 std::fs::create_dir_all(&dir)?;
734 let path = dir.join(format!(
735 "ctx_{}_{}.txt",
736 std::process::id(),
737 Uuid::new_v4().simple()
738 ));
739 std::fs::write(&path, prompt)?;
740 Ok(path)
741 }
742
743 fn build_root_request(
744 model: &str,
745 messages: &[Message],
746 system: &SystemPrompt,
747 max_tokens: u32,
748 ) -> MessageRequest {
749 MessageRequest {
750 model: model.to_string(),
751 messages: messages.to_vec(),
752 max_tokens,
753 system: Some(system.clone()),
754 tools: None,
755 tool_choice: None,
756 metadata: None,
757 thinking: None,
758 reasoning_effort: None,
759 stream: Some(false),
760 temperature: None,
761 top_p: None,
762 }
763 }
764
765 /// Build `Metadata(state)` from the paper. Surfaces:
766 /// - the small `root_prompt` (if any) — repeated each iteration
767 /// - `context` length + preview
768 /// - the REPL helpers
769 /// - the previous round's code summary + stdout preview
770 fn build_metadata_message(
771 prompt: &str,
772 root_prompt: Option<&str>,
773 iteration: u32,
774 previous_code: Option<&str>,
775 previous_stdout: Option<&str>,
776 ) -> Message {
777 let prompt_len = prompt.chars().count();
778 let prompt_preview = truncate_text(prompt, PROMPT_PREVIEW_LEN);
779
780 let mut parts = Vec::new();
781 parts.push(format!("## REPL state (round {iteration})"));
782 parts.push(String::new());
783 if let Some(rp) = root_prompt
784 && !rp.trim().is_empty()
785 {
786 parts.push("**Original task** (re-shown every round)".to_string());
787 parts.push(format!("> {}", truncate_text(rp.trim(), 600)));
788 parts.push(String::new());
789 }
790 parts.push("**`context`** — the long input lives in the REPL only".to_string());
791 parts.push(format!("- Length: {prompt_len} chars"));
792 parts.push(format!("- Preview: \"{prompt_preview}\""));
793 parts.push(String::new());
794
795 parts.push("**REPL helpers** (use inside ```repl blocks)".to_string());
796 parts.push("- `context` / `ctx` — the full input string".to_string());
797 parts.push("- `len(context)` / `context[a:b]` / `context.splitlines()` — slice it".to_string());
798 parts.push(
799 "- `chunk_context(max_chars=20000, overlap=0)` — full-coverage chunks with index/start/end/text"
800 .to_string(),
801 );
802 parts.push(
803 "- `chunk_coverage(chunks)` — coverage report for chunk_context output"
804 .to_string(),
805 );
806 parts.push(
807 "- `llm_query(prompt, model=None)` — one-shot child LLM; `model` is ignored and child calls stay pinned to Flash"
808 .to_string(),
809 );
810 parts.push(
811 "- `llm_query_batched([p1, p2, ...], dependency_mode=\"independent\")` — concurrent fan-out for independent prompts only; `model` is ignored"
812 .to_string(),
813 );
814 parts.push(
815 "- `rlm_query(prompt, model=None)` — recursive sub-RLM; `model` is ignored"
816 .to_string(),
817 );
818 parts.push(
819 "- `rlm_query_batched([p1, p2, ...], dependency_mode=\"independent\")` — concurrent recursive sub-RLMs for independent prompts only; `model` is ignored"
820 .to_string(),
821 );
822 parts.push(
823 "- `sub_query_sequence(prompt, slices)` — sequential child calls for A->B dependencies and rollback-sensitive work"
824 .to_string(),
825 );
826 parts.push(
827 "- Batch safety: never batch dependent steps, global-state refactors, schema migrations, or rollback-sensitive tasks"
828 .to_string(),
829 );
830 parts.push("- `SHOW_VARS()` — list user variables".to_string());
831 parts.push("- `repl_set(name, value)` / `repl_get(name)` — explicit store".to_string());
832 parts.push(
833 "- `FINAL(value)` — end the loop with this answer".to_string(),
834 );
835 parts.push(
836 "- `FINAL_VAR(name)` — end the loop with a variable's value"
837 .to_string(),
838 );
839 parts.push(String::new());
840
841 if iteration > 0 {
842 parts.push("**Previous round**".to_string());
843 if let Some(code) = previous_code {
844 parts.push(format!("- Code: {}", summarize_code(code)));
845 }
846 if let Some(stdout) = previous_stdout {
847 let stdout_clean = stdout.trim();
848 if !stdout_clean.is_empty() {
849 parts.push(format!("- Stdout preview: \"{stdout_clean}\""));
850 } else {
851 parts.push("- Stdout: (empty)".to_string());
852 }
853 }
854 }
855
856 let text = parts.join("\n");
857
858 Message {
859 role: Role::User,
860 content: vec![ContentBlock::Text {
861 text,
862 cache_control: None,
863 }],
864 }
865 }
866
867 fn summarize_code(code: &str) -> String {
868 let lines: Vec<&str> = code.lines().collect();
869 if lines.len() <= 8 {
870 return code.to_string();
871 }
872 let head = lines[..4].join("\n");
873 let tail = lines[lines.len() - 4..].join("\n");
874 format!("{} lines:\n{head}\n…\n{tail}", lines.len())
875 }
876
877 fn extract_text_blocks(blocks: &[ContentBlock]) -> String {
878 blocks
879 .iter()
880 .filter_map(|b| match b {
881 ContentBlock::Text { text, .. } => Some(text.as_str()),
882 _ => None,
883 })
884 .collect::<Vec<_>>()
885 .join("\n")
886 }
887
888 /// Extract the first ` ```repl ` block from `text`. Falls back to
889 /// ` ```python `/`` ```py `` for compatibility with prompts that learned
890 /// the older fence style.
891 fn extract_repl_code(text: &str) -> Option<String> {
892 let start_markers = [
893 "```repl\n",
894 "```repl\r\n",
895 "```python\n",
896 "```py\n",
897 "```python\r\n",
898 "```py\r\n",
899 ];
900 let mut best_start: Option<(usize, &str)> = None;
901
902 for marker in &start_markers {
903 if let Some(idx) = text.find(marker) {
904 let end_pos = idx + marker.len();
905 match best_start {
906 Some((best_idx, _)) if idx < best_idx => {
907 best_start = Some((idx, &text[end_pos..]));
908 }
909 None => {
910 best_start = Some((idx, &text[end_pos..]));
911 }
912 _ => {}
913 }
914 }
915 }
916
917 let after_fence = best_start.map(|(_, rest)| rest)?;
918
919 let end_idx = after_fence
920 .find("\n```")
921 .or_else(|| after_fence.find("```"))?;
922
923 let code = after_fence[..end_idx].trim().to_string();
924 if code.is_empty() {
925 return None;
926 }
927 Some(code)
928 }
929
930 /// Parse a top-level `FINAL(...)` directive from the model's raw text.
931 /// Mirrors the reference RLM's `find_final_answer`: directive must appear
932 /// at the start of a line, *outside* any code fence.
933 fn parse_text_final(text: &str) -> Option<String> {
934 let outside_fence = strip_code_fences(text);
935
936 for line in outside_fence.lines() {
937 let trimmed = line.trim_start();
938 if trimmed.starts_with("FINAL_VAR(") {
939 // FINAL_VAR can't be resolved from text alone — defer to REPL.
940 continue;
941 }
942 if let Some(rest) = trimmed.strip_prefix("FINAL(") {
943 let inner = rest.trim_end();
944 if let Some(end) = inner.rfind(')') {
945 let value = inner[..end].trim();
946 if !value.is_empty() {
947 return Some(strip_quotes(value));
948 }
949 }
950 }
951 }
952 None
953 }
954
955 fn strip_code_fences(text: &str) -> String {
956 let mut out = String::with_capacity(text.len());
957 let mut in_fence = false;
958 for line in text.lines() {
959 if line.trim_start().starts_with("```") {
960 in_fence = !in_fence;
961 continue;
962 }
963 if !in_fence {
964 out.push_str(line);
965 out.push('\n');
966 }
967 }
968 out
969 }
970
971 fn strip_quotes(s: &str) -> String {
972 let bytes = s.as_bytes();
973 if bytes.len() >= 2
974 && ((bytes[0] == b'"' && bytes[bytes.len() - 1] == b'"')
975 || (bytes[0] == b'\'' && bytes[bytes.len() - 1] == b'\''))
976 {
977 return s[1..s.len() - 1].to_string();
978 }
979 s.to_string()
980 }
981
982 fn truncate_text(text: &str, max_chars: usize) -> String {
983 let count = text.chars().count();
984 if count <= max_chars {
985 return text.to_string();
986 }
987 let take = max_chars.saturating_sub(3);
988 let mut result: String = text.chars().take(take).collect();
989 result.push_str("...");
990 result
991 }
992
993 // ---------------------------------------------------------------------------
994 // Tests
995 // ---------------------------------------------------------------------------
996
997 #[cfg(test)]
998 mod tests {
999 use super::*;
1000 use crate::llm_client::mock::MockLlmClient;
1001 use codewhale_models::MessageResponse;
1002
1003 #[tokio::test]
1004 async fn max_tokens_complete_repl_is_not_executed_or_accepted() {
1005 let workspace = tempfile::tempdir().expect("tempdir");
1006 let marker = workspace.path().join("truncated-repl-executed.txt");
1007 let marker_literal = serde_json::to_string(&marker.to_string_lossy())
1008 .expect("marker path should serialize as a Python string literal");
1009 let partial = format!(
1010 "```repl\nfrom pathlib import Path\nPath({marker_literal}).write_text('executed')\nFINAL('partial answer')\n```"
1011 );
1012 let usage = Usage {
1013 input_tokens: 17,
1014 output_tokens: 4096,
1015 ..Usage::default()
1016 };
1017 let mock = Arc::new(MockLlmClient::new(Vec::new()));
1018 mock.push_message_response(MessageResponse {
1019 id: "mock_truncated_rlm".to_string(),
1020 r#type: "message".to_string(),
1021 role: "assistant".to_string(),
1022 content: vec![ContentBlock::Text {
1023 text: partial,
1024 cache_control: None,
1025 }],
1026 model: "mock-model".to_string(),
1027 stop_reason: Some("max_tokens".to_string()),
1028 stop_sequence: None,
1029 container: None,
1030 usage: usage.clone(),
1031 });
1032 let client: Arc<dyn RlmLlmClient> = mock.clone();
1033 let (tx, _rx) = mpsc::channel(8);
1034
1035 let result = run_rlm_turn_inner(
1036 client,
1037 "root-model".to_string(),
1038 "long context".to_string(),
1039 None,
1040 "child-model".to_string(),
1041 tx,
1042 0,
1043 )
1044 .await;
1045
1046 assert_eq!(result.termination, RlmTermination::Error);
1047 assert!(
1048 result.answer.is_empty(),
1049 "partial FINAL must not be accepted"
1050 );
1051 let error = result.error.expect("truncation must fail the RLM turn");
1052 assert!(error.contains("incomplete"), "{error}");
1053 assert!(error.contains("max_tokens"), "{error}");
1054 assert_eq!(result.usage, usage, "billed usage must still be charged");
1055 assert_eq!(result.routed_usage.len(), 1);
1056 assert_eq!(result.routed_usage[0].usage.usage, usage);
1057 assert_eq!(result.routed_usage_dropped_records, 0);
1058 assert_eq!(mock.call_count(), 1, "truncation must not retry");
1059 assert!(
1060 !marker.exists(),
1061 "complete-looking code from a truncated response must not execute"
1062 );
1063 }
1064
1065 #[tokio::test]
1066 async fn root_provider_success_without_usage_retains_exact_missing_receipt() {
1067 let mock = Arc::new(MockLlmClient::new(Vec::new()));
1068 mock.push_message_response(MessageResponse {
1069 id: "mock_missing_rlm_usage".to_string(),
1070 r#type: "message".to_string(),
1071 role: "assistant".to_string(),
1072 content: vec![ContentBlock::Text {
1073 text: "partial".to_string(),
1074 cache_control: None,
1075 }],
1076 model: "mock-model".to_string(),
1077 stop_reason: Some("max_tokens".to_string()),
1078 stop_sequence: None,
1079 container: None,
1080 usage: Usage::default(),
1081 });
1082 let client: Arc<dyn RlmLlmClient> = mock;
1083 let (tx, _rx) = mpsc::channel(8);
1084
1085 let result = run_rlm_turn_inner(
1086 client,
1087 "root-model".to_string(),
1088 "long context".to_string(),
1089 None,
1090 "child-model".to_string(),
1091 tx,
1092 0,
1093 )
1094 .await;
1095
1096 assert_eq!(result.termination, RlmTermination::Error);
1097 assert_eq!(result.usage, Usage::default());
1098 assert!(result.routed_usage.is_empty());
1099 assert_eq!(result.routed_usage_drop_records.len(), 1);
1100 assert_eq!(result.routed_usage_dropped_records, 1);
1101 assert_eq!(
1102 result.routed_usage_drop_records[0].route.model,
1103 "root-model"
1104 );
1105 }
1106
1107 #[test]
1108 fn extract_repl_code_finds_simple_block() {
1109 let text = "Here:\n```repl\nprint('hi')\n```\nEnd.";
1110 let code = extract_repl_code(text).unwrap();
1111 assert_eq!(code, "print('hi')");
1112 }
1113
1114 #[test]
1115 fn extract_repl_code_falls_back_to_python_marker() {
1116 let text = "Code:\n```python\nx = 1 + 2\n```";
1117 let code = extract_repl_code(text).unwrap();
1118 assert_eq!(code, "x = 1 + 2");
1119 }
1120
1121 #[test]
1122 fn extract_repl_code_returns_none_when_missing() {
1123 assert!(extract_repl_code("Just text.").is_none());
1124 }
1125
1126 #[test]
1127 fn extract_repl_code_returns_none_on_empty_block() {
1128 assert!(extract_repl_code("```repl\n\n```").is_none());
1129 }
1130
1131 #[test]
1132 fn extract_repl_code_handles_multiple_blocks() {
1133 let text = "```repl\na=1\n```\n```repl\nb=2\n```";
1134 let code = extract_repl_code(text).unwrap();
1135 assert_eq!(code, "a=1");
1136 }
1137
1138 #[test]
1139 fn extract_repl_code_ignores_other_fences() {
1140 let text = "```\nfoo\n```\n```repl\nreal_code()\n```";
1141 let code = extract_repl_code(text).unwrap();
1142 assert_eq!(code, "real_code()");
1143 }
1144
1145 #[test]
1146 fn parse_text_final_extracts_simple_value() {
1147 let text = "OK.\nFINAL(42)\nThanks.";
1148 assert_eq!(parse_text_final(text).as_deref(), Some("42"));
1149 }
1150
1151 #[test]
1152 fn parse_text_final_strips_quotes() {
1153 let text = "FINAL(\"the answer is yes\")";
1154 assert_eq!(parse_text_final(text).as_deref(), Some("the answer is yes"));
1155 }
1156
1157 #[test]
1158 fn parse_text_final_ignores_inside_code_fence() {
1159 let text =
1160 "Some prose.\n```repl\n# Note: when ready, call FINAL(value)\nx = 1\n```\nMore prose.";
1161 assert!(parse_text_final(text).is_none());
1162 }
1163
1164 #[test]
1165 fn parse_text_final_returns_none_when_absent() {
1166 assert!(parse_text_final("just talking, no final.").is_none());
1167 }
1168
1169 #[test]
1170 fn build_metadata_contains_key_information() {
1171 let msg = build_metadata_message("Hello, world!", None, 0, None, None);
1172 let text = extract_text_blocks(&msg.content);
1173 assert!(text.contains("context"));
1174 assert!(text.contains("Hello, world!"));
1175 assert!(text.contains("round 0"));
1176 assert!(text.contains("llm_query"));
1177 assert!(text.contains("rlm_query"));
1178 assert!(text.contains("FINAL"));
1179 }
1180
1181 #[test]
1182 fn build_metadata_truncates_long_context_without_leaking_tail() {
1183 let secret_tail = "DO_NOT_LEAK_CONTEXT_TAIL";
1184 let prompt = format!("{}{}", "a".repeat(PROMPT_PREVIEW_LEN + 100), secret_tail);
1185 let msg = build_metadata_message(&prompt, None, 0, None, None);
1186 let text = extract_text_blocks(&msg.content);
1187
1188 assert!(text.contains(&format!("- Length: {} chars", prompt.chars().count())));
1189 assert!(text.contains("- Preview: \""));
1190 assert!(text.contains("..."));
1191 assert!(
1192 !text.contains(secret_tail),
1193 "metadata leaked the non-preview tail of context"
1194 );
1195 }
1196
1197 #[test]
1198 fn build_root_request_keeps_context_tail_out_of_root_payload() {
1199 let secret_tail = "DO_NOT_LEAK_ROOT_REQUEST";
1200 let prompt = format!("{}{}", "a".repeat(PROMPT_PREVIEW_LEN + 100), secret_tail);
1201 let messages = vec![build_metadata_message(
1202 &prompt,
1203 Some("answer from the long context"),
1204 0,
1205 None,
1206 None,
1207 )];
1208
1209 let request = build_root_request("root-model", &messages, &rlm_system_prompt(), 8192);
1210 let payload = serde_json::to_string(&request).expect("request should serialize");
1211
1212 assert_eq!(request.max_tokens, 8192);
1213 assert_eq!(request.temperature, None);
1214 assert_eq!(request.top_p, None);
1215 assert!(payload.contains(&format!("- Length: {} chars", prompt.chars().count())));
1216 assert!(
1217 !payload.contains(secret_tail),
1218 "root LLM request leaked the non-preview tail of context"
1219 );
1220 }
1221
1222 #[test]
1223 fn build_metadata_with_iteration_shows_previous_code() {
1224 let msg = build_metadata_message("Test prompt", None, 3, Some("print('hi')"), Some("hi"));
1225 let text = extract_text_blocks(&msg.content);
1226 assert!(text.contains("round 3"));
1227 assert!(text.contains("print('hi')"));
1228 assert!(text.contains("hi"));
1229 }
1230
1231 #[test]
1232 fn build_metadata_includes_root_prompt() {
1233 let msg = build_metadata_message(
1234 "long context",
1235 Some("Summarize the security model"),
1236 1,
1237 Some("# noop"),
1238 Some("ok"),
1239 );
1240 let text = extract_text_blocks(&msg.content);
1241 assert!(text.contains("Original task"));
1242 assert!(text.contains("Summarize the security model"));
1243 }
1244
1245 #[test]
1246 fn truncate_text_leaves_short_alone() {
1247 assert_eq!(truncate_text("hello", 100), "hello");
1248 }
1249
1250 #[test]
1251 fn truncate_text_shortens_long_text() {
1252 let long = "a".repeat(1000);
1253 let truncated = truncate_text(&long, 10);
1254 assert_eq!(truncated.chars().count(), 10);
1255 assert!(truncated.ends_with("..."));
1256 }
1257
1258 #[test]
1259 fn truncate_text_is_unicode_safe() {
1260 let s = "日本語テスト";
1261 let out = truncate_text(s, 4);
1262 assert_eq!(out.chars().count(), 4);
1263 assert!(out.ends_with("..."));
1264 assert!(std::str::from_utf8(out.as_bytes()).is_ok());
1265 }
1266
1267 #[test]
1268 fn extract_text_blocks_joins_text() {
1269 let blocks = vec![
1270 ContentBlock::Text {
1271 text: "first".to_string(),
1272 cache_control: None,
1273 },
1274 ContentBlock::Thinking {
1275 signature: None,
1276 state: None,
1277 thinking: "skip".to_string(),
1278 },
1279 ContentBlock::Text {
1280 text: "second".to_string(),
1281 cache_control: None,
1282 },
1283 ];
1284 assert_eq!(extract_text_blocks(&blocks), "first\nsecond");
1285 }
1286
1287 #[test]
1288 fn metadata_msg_role_is_user() {
1289 let msg = build_metadata_message("test", None, 0, None, None);
1290 assert_eq!(msg.role, "user");
1291 }
1292
1293 #[test]
1294 fn summarize_code_keeps_short() {
1295 assert_eq!(summarize_code("a\nb\nc"), "a\nb\nc");
1296 }
1297
1298 #[test]
1299 fn summarize_code_compresses_long() {
1300 let lines: Vec<String> = (0..20).map(|i| format!("line{i}")).collect();
1301 let code = lines.join("\n");
1302 let s = summarize_code(&code);
1303 assert!(s.starts_with("20 lines:"));
1304 assert!(s.contains("line0"));
1305 assert!(s.contains("line19"));
1306 assert!(s.contains("…"));
1307 }
1308
1309 #[test]
1310 fn rlm_turn_has_no_fixed_wall_clock_timeout() {
1311 assert!(
1312 turn_timeout().is_none(),
1313 "RLM turns should not be killed by the old fixed 180s wall-clock cap"
1314 );
1315 }
1316 }
1317
1317 lines RUST