deepseek-tui 0.8.36

Terminal UI for DeepSeek
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
//! Long-lived Python REPL runtime.
//!
//! One Python subprocess lives for the duration of an RLM turn (or an
//! inline `repl` block sequence in the agent loop). Code blocks are sent
//! over stdin framed by `__RLM_RUN__`/`__RLM_END__` sentinels; the bootstrap
//! `exec()`s them into the same global namespace so variables, imports,
//! and even open file handles persist naturally across rounds.
//!
//! Sub-LLM helpers (`sub_query`, `sub_query_batch`, `sub_rlm`, plus legacy
//! `llm_query`, `llm_query_batched`, `rlm_query`, `rlm_query_batched`) are
//! wired through a stdin/stdout RPC protocol:
//! Python emits `__RLM_REQ_<sid>__::{json}` on stdout, Rust dispatches the
//! request and writes `__RLM_RESP_<sid>__::{json}` back on stdin. No HTTP
//! sidecar, no temp ports — the same pipes carry both control and data.
//!
//! The session id (`<sid>`) is a UUID generated per spawn, so user output
//! that happens to contain "REQ" or "FINAL" can't be confused with control
//! messages.

use std::ffi::OsString;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::{Duration, Instant};

use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use uuid::Uuid;

use crate::child_env;
use crate::dependencies::{PYTHON_CANDIDATES, resolve_python_interpreter, split_interpreter_spec};

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Result of executing one code block.
#[derive(Debug, Clone)]
pub struct ReplRound {
    /// Stdout shown to the model as metadata next round.
    pub stdout: String,
    /// Full stdout (with sentinels stripped, but otherwise raw).
    pub full_stdout: String,
    /// Stderr from this round (if any).
    pub stderr: String,
    /// `True` if the user code raised an unhandled Python exception.
    pub has_error: bool,
    /// Captured `finalize(value, confidence=...)` payload, if any.
    pub final_value: Option<String>,
    /// Captured final value before string fallback. Structured `finalize`
    /// payloads use this so `handle_read` can expose JSON instead of a Python
    /// repr string.
    pub final_json: Option<Value>,
    /// Optional confidence supplied to `finalize(...)`.
    pub final_confidence: Option<Value>,
    /// Number of `sub_query`/`sub_rlm` RPCs the round issued.
    pub rpc_count: u32,
    /// Wall-clock duration of the round.
    pub elapsed: Duration,
}

/// One RPC request emitted by Python during a round.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum RpcRequest {
    /// `llm_query(prompt, model=None, max_tokens=None, system=None)`
    Llm {
        prompt: String,
        #[serde(default)]
        model: Option<String>,
        #[serde(default)]
        max_tokens: Option<u32>,
        #[serde(default)]
        system: Option<String>,
    },
    /// `llm_query_batched(prompts, model=None, dependency_mode="independent")`
    LlmBatch {
        prompts: Vec<String>,
        #[serde(default)]
        model: Option<String>,
        #[serde(default)]
        dependency_mode: Option<String>,
        #[serde(default)]
        safety_note: Option<String>,
    },
    /// `rlm_query(prompt, model=None)` — recursive sub-RLM (paper's `sub_RLM`).
    Rlm {
        prompt: String,
        #[serde(default)]
        model: Option<String>,
    },
    /// `rlm_query_batched(prompts, model=None, dependency_mode="independent")`
    RlmBatch {
        prompts: Vec<String>,
        #[serde(default)]
        model: Option<String>,
        #[serde(default)]
        dependency_mode: Option<String>,
        #[serde(default)]
        safety_note: Option<String>,
    },
}

/// Response for one RPC request.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum RpcResponse {
    /// Single-text reply (Llm / Rlm).
    Single(SingleResp),
    /// Batch reply (LlmBatch / RlmBatch).
    Batch(BatchResp),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SingleResp {
    #[serde(default)]
    pub text: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BatchResp {
    pub results: Vec<SingleResp>,
}

/// Trait-object handle for dispatching Python RPCs back into Rust.
///
/// Each RLM turn supplies one. Implementations forward to the LLM client
/// (and recursively into `run_rlm_turn_inner` for `Rlm` / `RlmBatch`).
pub trait RpcDispatcher: Send + Sync {
    fn dispatch<'a>(
        &'a self,
        req: RpcRequest,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = RpcResponse> + Send + 'a>>;
}

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

const DEFAULT_STDOUT_LIMIT: usize = 8_192;
const ROUND_TIMEOUT: Duration = Duration::from_secs(180);
#[cfg(not(windows))]
const SPAWN_READY_TIMEOUT: Duration = Duration::from_secs(10);
#[cfg(windows)]
const SPAWN_READY_TIMEOUT: Duration = Duration::from_secs(30);

// ---------------------------------------------------------------------------
// PythonRuntime
// ---------------------------------------------------------------------------

/// Long-lived Python REPL.
#[derive(Debug)]
pub struct PythonRuntime {
    child: Child,
    stdin: ChildStdin,
    stdout: BufReader<ChildStdout>,
    /// Per-spawn session id used in protocol sentinels.
    session_id: String,
    /// Path to the file holding `context` (kept around for cleanup).
    context_path: Option<PathBuf>,
    stdout_limit: usize,
    round_count: u64,
    started: Instant,
    round_timeout: Option<Duration>,
}

impl PythonRuntime {
    /// Spawn a REPL with no `context` variable and no LLM helpers wired up.
    /// Used by the agent loop for inline `repl` blocks the model emits in
    /// regular conversation.
    pub async fn new() -> Result<Self, String> {
        Self::spawn_inner(None, Some(ROUND_TIMEOUT)).await
    }

    /// Compatibility shim — older RLM code path used to pass a state file.
    /// The state file is no longer used, but the path doubles as an extra
    /// scratch location callers can rely on for cleanup symmetry.
    pub fn with_state_path(_path: PathBuf) -> Self {
        // Synchronous constructor is no longer meaningful: spawning Python
        // is async. Callers in turn.rs already use `spawn_with_context` —
        // this stub is kept only so the public surface compiles for any
        // out-of-tree user. It returns a deliberately broken runtime that
        // panics on first use, which is preferable to silently lying.
        unreachable!(
            "PythonRuntime::with_state_path is deprecated — \
             use PythonRuntime::new() or PythonRuntime::spawn_with_context()"
        )
    }

    /// Spawn a REPL with the long input preloaded from a file. Used by the
    /// RLM turn loop.
    pub async fn spawn_with_context(context_path: &Path) -> Result<Self, String> {
        Self::spawn_inner(Some(context_path), None).await
    }

    async fn spawn_inner(
        context_path: Option<&Path>,
        round_timeout: Option<Duration>,
    ) -> Result<Self, String> {
        let session_id = Uuid::new_v4().simple().to_string();
        let bootstrap = render_bootstrap(&session_id);

        let interpreter = resolve_python_interpreter().ok_or_else(|| {
            format!(
                "no Python interpreter found on PATH (tried {:?}). \
                 Install Python 3 and ensure one of these commands works, then restart deepseek-tui.",
                PYTHON_CANDIDATES,
            )
        })?;
        let (program, interpreter_args) = split_interpreter_spec(&interpreter);
        if program.is_empty() {
            return Err(format!(
                "resolved Python interpreter is empty: {interpreter:?}"
            ));
        }

        let mut cmd = Command::new(&program);
        cmd.args(&interpreter_args)
            .arg("-u")
            .arg("-c")
            .arg(&bootstrap)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .kill_on_drop(true);

        let context_env = context_path
            .map(|path| {
                vec![(
                    OsString::from("RLM_CONTEXT_FILE"),
                    path.as_os_str().to_os_string(),
                )]
            })
            .unwrap_or_default();
        child_env::apply_to_tokio_command(&mut cmd, context_env);

        let mut child = cmd
            .spawn()
            .map_err(|e| format!("failed to spawn Python interpreter `{interpreter}`: {e}"))?;

        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| format!("Python interpreter `{interpreter}` stdin pipe missing"))?;
        let raw_stdout = child
            .stdout
            .take()
            .ok_or_else(|| format!("Python interpreter `{interpreter}` stdout pipe missing"))?;
        let stdout = BufReader::new(raw_stdout);

        let mut rt = Self {
            child,
            stdin,
            stdout,
            session_id: session_id.clone(),
            context_path: context_path.map(Path::to_path_buf),
            stdout_limit: DEFAULT_STDOUT_LIMIT,
            round_count: 0,
            started: Instant::now(),
            round_timeout,
        };

        // Wait for `__RLM_READY_<sid>__` before handing control back. If
        // Python failed to start (missing module, syntax error in the
        // bootstrap, etc.), this is where we'll find out.
        let ready_sentinel = format!("__RLM_READY_{session_id}__");
        match tokio::time::timeout(SPAWN_READY_TIMEOUT, rt.read_until_ready(&ready_sentinel)).await
        {
            Ok(Ok(())) => Ok(rt),
            Ok(Err(e)) => {
                let _ = rt.child.kill().await;
                Err(format!(
                    "Python interpreter `{interpreter}` bootstrap failed: {e}"
                ))
            }
            Err(_) => {
                let _ = rt.child.kill().await;
                Err(format!(
                    "Python interpreter `{interpreter}` bootstrap did not signal ready within {}s",
                    SPAWN_READY_TIMEOUT.as_secs()
                ))
            }
        }
    }

    async fn read_until_ready(&mut self, ready_sentinel: &str) -> Result<(), String> {
        loop {
            let mut line = String::new();
            let n = self
                .stdout
                .read_line(&mut line)
                .await
                .map_err(|e| format!("stdout read: {e}"))?;
            if n == 0 {
                return Err("Python interpreter closed stdout before ready signal".to_string());
            }
            let trimmed = line.trim_end_matches(['\n', '\r']);
            if trimmed == ready_sentinel {
                return Ok(());
            }
            // Pre-ready output is rare; ignore it.
        }
    }

    /// Execute a Python code block with no RPC dispatcher. Used for inline
    /// `repl` blocks where `llm_query()` should fall back to a sentinel.
    pub async fn execute(&mut self, code: &str) -> Result<ReplRound, String> {
        self.run(code, None::<&dyn RpcDispatcher>).await
    }

    /// Execute a code block, dispatching any sub-LLM RPCs through `bridge`.
    ///
    /// Returns once Python emits `__RLM_DONE_<sid>__` or the round timeout
    /// elapses (whichever happens first).
    pub async fn run<D>(&mut self, code: &str, bridge: Option<&D>) -> Result<ReplRound, String>
    where
        D: RpcDispatcher + ?Sized,
    {
        let started = Instant::now();
        self.round_count += 1;
        let round_id = self.round_count;

        // Send the code header + body + end marker in one write.
        let header = format!("__RLM_RUN_{}__::{round_id}\n", self.session_id);
        let footer = format!("__RLM_END_{}__\n", self.session_id);
        let payload = format!("{header}{code}\n{footer}");
        self.stdin
            .write_all(payload.as_bytes())
            .await
            .map_err(|e| format!("stdin write: {e}"))?;
        self.stdin
            .flush()
            .await
            .map_err(|e| format!("stdin flush: {e}"))?;

        // Sentinels for this session.
        let req_prefix = format!("__RLM_REQ_{}__::", self.session_id);
        let final_prefix = format!("__RLM_FINAL_{}__::", self.session_id);
        let err_prefix = format!("__RLM_ERR_{}__::", self.session_id);
        let done_prefix = format!("__RLM_DONE_{}__::", self.session_id);

        let mut stdout_buf = String::new();
        let mut final_value: Option<String> = None;
        let mut final_json: Option<Value> = None;
        let mut final_confidence: Option<Value> = None;
        let mut had_error = false;
        let mut rpc_count: u32 = 0;
        let round_timeout = self.round_timeout;

        let read_loop = async {
            loop {
                let mut line = String::new();
                let n = self
                    .stdout
                    .read_line(&mut line)
                    .await
                    .map_err(|e| format!("stdout read: {e}"))?;
                if n == 0 {
                    return Err("Python interpreter closed stdout mid-round".to_string());
                }
                let trimmed = line.trim_end_matches(['\n', '\r']);

                if let Some(rest) = trimmed.strip_prefix(&done_prefix) {
                    let _ = rest;
                    break;
                }
                if let Some(rest) = trimmed.strip_prefix(&final_prefix) {
                    // New sessions emit an object with value/confidence;
                    // legacy helpers emitted a JSON string.
                    match serde_json::from_str::<Value>(rest) {
                        Ok(Value::Object(map)) => {
                            let value_json = map
                                .get("value")
                                .cloned()
                                .unwrap_or(Value::String(rest.to_string()));
                            let value = value_json
                                .as_str()
                                .map(str::to_string)
                                .unwrap_or_else(|| value_json.to_string());
                            final_json = Some(value_json);
                            final_value = Some(value);
                            final_confidence = map.get("confidence").cloned();
                        }
                        Ok(Value::String(value)) => {
                            final_json = Some(Value::String(value.clone()));
                            final_value = Some(value);
                            final_confidence = None;
                        }
                        Ok(other) => {
                            final_json = Some(other.clone());
                            final_value = Some(other.to_string());
                            final_confidence = None;
                        }
                        Err(_) => {
                            final_value = Some(rest.to_string());
                            final_confidence = None;
                        }
                    }
                    continue;
                }
                if let Some(rest) = trimmed.strip_prefix(&err_prefix) {
                    let traceback =
                        serde_json::from_str::<String>(rest).unwrap_or_else(|_| rest.to_string());
                    had_error = true;
                    stdout_buf.push_str(&format!("[traceback]\n{traceback}\n"));
                    continue;
                }
                if let Some(rest) = trimmed.strip_prefix(&req_prefix) {
                    rpc_count = rpc_count.saturating_add(1);
                    let req: RpcRequest = match serde_json::from_str(rest) {
                        Ok(r) => r,
                        Err(e) => {
                            // Send an error response so Python isn't blocked.
                            self.send_resp(&RpcResponse::Single(SingleResp {
                                text: String::new(),
                                error: Some(format!("malformed RPC: {e}")),
                            }))
                            .await?;
                            continue;
                        }
                    };
                    let resp = match bridge {
                        Some(b) => b.dispatch(req).await,
                        None => RpcResponse::Single(SingleResp {
                            text: String::new(),
                            error: Some("no LLM bridge bound to this REPL".to_string()),
                        }),
                    };
                    self.send_resp(&resp).await?;
                    continue;
                }

                stdout_buf.push_str(&line);
            }
            Ok::<_, String>(())
        };

        if let Some(round_timeout) = round_timeout {
            match tokio::time::timeout(round_timeout, read_loop).await {
                Ok(Ok(())) => {}
                Ok(Err(e)) => return Err(e),
                Err(_) => {
                    return Err(format!(
                        "REPL round timed out after {}s",
                        round_timeout.as_secs()
                    ));
                }
            }
        } else {
            read_loop.await?;
        }

        let stderr = self.drain_stderr().await;
        let display = truncate_stdout(stdout_buf.trim_end_matches('\n'), self.stdout_limit);

        Ok(ReplRound {
            stdout: display,
            full_stdout: stdout_buf,
            stderr,
            has_error: had_error,
            final_value,
            final_json,
            final_confidence,
            rpc_count,
            elapsed: started.elapsed(),
        })
    }

    async fn send_resp(&mut self, resp: &RpcResponse) -> Result<(), String> {
        let body = serde_json::to_string(resp).map_err(|e| format!("encode rpc resp: {e}"))?;
        let line = format!("__RLM_RESP_{}__::{body}\n", self.session_id);
        self.stdin
            .write_all(line.as_bytes())
            .await
            .map_err(|e| format!("stdin write resp: {e}"))?;
        self.stdin
            .flush()
            .await
            .map_err(|e| format!("stdin flush resp: {e}"))?;
        Ok(())
    }

    async fn drain_stderr(&mut self) -> String {
        // We don't continuously read stderr — drain whatever's pending after
        // a round so it can show up in error reports without deadlocking
        // anything during normal operation.
        let Some(stderr) = self.child.stderr.as_mut() else {
            return String::new();
        };
        use tokio::io::AsyncReadExt;
        let mut buf = Vec::new();
        // Best-effort read with a tight deadline; we don't want to block.
        let fut = async {
            let mut chunk = [0u8; 4096];
            loop {
                match tokio::time::timeout(Duration::from_millis(20), stderr.read(&mut chunk)).await
                {
                    Ok(Ok(0)) => break,
                    Ok(Ok(n)) => buf.extend_from_slice(&chunk[..n]),
                    _ => break,
                }
            }
        };
        let _ = fut.await;
        String::from_utf8_lossy(&buf).to_string()
    }

    /// Total rounds executed.
    pub fn round_count(&self) -> u64 {
        self.round_count
    }

    /// Current per-round timeout policy. RLM context runs intentionally return
    /// `None` so long map-reduce jobs are not killed by the old 180s cap.
    pub fn round_timeout(&self) -> Option<Duration> {
        self.round_timeout
    }

    /// Wall-clock uptime since spawn.
    pub fn uptime(&self) -> Duration {
        self.started.elapsed()
    }

    /// Cleanly tear down the subprocess.
    pub async fn shutdown(mut self) {
        let _ = self.stdin.shutdown().await;
        let _ = self.child.kill().await;
        if let Some(path) = self.context_path.take() {
            let _ = tokio::fs::remove_file(path).await;
        }
    }
}

impl Drop for PythonRuntime {
    fn drop(&mut self) {
        // tokio sets `kill_on_drop(true)` on the child; the context file
        // (if any) is removed on `shutdown()` — drop is best-effort.
        if let Some(path) = self.context_path.take() {
            let _ = std::fs::remove_file(path);
        }
    }
}

// ---------------------------------------------------------------------------
// Bootstrap script
// ---------------------------------------------------------------------------

/// Render the Python bootstrap with session-specific sentinels baked in.
/// The sentinels include a UUID to prevent user prints from being mistaken
/// for control messages.
fn render_bootstrap(session_id: &str) -> String {
    BOOTSTRAP_TEMPLATE.replace("__SID__", session_id)
}

const BOOTSTRAP_TEMPLATE: &str = r#"
import json as _json
import os as _os
import re as _re
import sys as _sys
import traceback as _traceback

_SID = "__SID__"
_REQ = f"__RLM_REQ_{_SID}__::"
_RESP = f"__RLM_RESP_{_SID}__::"
_FINAL = f"__RLM_FINAL_{_SID}__::"
_ERR = f"__RLM_ERR_{_SID}__::"
_RUN = f"__RLM_RUN_{_SID}__::"
_END = f"__RLM_END_{_SID}__"
_DONE = f"__RLM_DONE_{_SID}__::"
_READY = f"__RLM_READY_{_SID}__"

def _rpc(req):
    _sys.stdout.write(_REQ + _json.dumps(req) + "\n")
    _sys.stdout.flush()
    line = _sys.stdin.readline()
    if not line:
        return {"error": "rust driver closed stdin"}
    if line.startswith(_RESP):
        try:
            return _json.loads(line[len(_RESP):])
        except Exception as e:
            return {"error": f"malformed rpc resp: {e}"}
    return {"error": f"unexpected protocol line: {line[:120]!r}"}

def llm_query(prompt, model=None, max_tokens=None, system=None):
    """One-shot sub-LLM call. The model arg is accepted for compatibility but ignored by Rust."""
    resp = _rpc({"type":"llm","prompt":str(prompt),"model":model,
                 "max_tokens":max_tokens,"system":system})
    if isinstance(resp, dict) and resp.get("error"):
        return f"[llm_query error: {resp['error']}]"
    if isinstance(resp, dict):
        return resp.get("text","")
    return str(resp)

def _normalize_dependency_mode(mode):
    if mode is None:
        return ""
    return str(mode).strip().lower().replace("-", "_").replace(" ", "_")

def _batch_dependency_error(helper, prompts, dependency_mode):
    mode = _normalize_dependency_mode(dependency_mode)
    if mode in ("independent", "parallel_safe", "map_reduce"):
        return None
    if mode in ("sequential", "dependent", "ordered", "chain", "serial"):
        return (
            f"[{helper}: refused parallel batch because dependency_mode={dependency_mode!r}. "
            "Use sub_query_sequence(...) or an explicit for-loop with sub_query(...) so each step can consume the previous result.]"
        )
    return (
        f"[{helper}: batch helpers require dependency_mode='independent'. "
        "Use only for independent slices/items; for A->B dependencies, global-state refactors, migrations, or rollback-sensitive work, use sub_query_sequence(...).]"
    )

def llm_query_batched(prompts, model=None, dependency_mode=None, safety_note=None):
    """Run independent sub-LLM calls concurrently. Declare dependency_mode='independent'."""
    if not isinstance(prompts, (list, tuple)):
        return ["[llm_query_batched: prompts must be a list]"]
    err = _batch_dependency_error("llm_query_batched", prompts, dependency_mode)
    if err is not None:
        return [err for _ in prompts]
    resp = _rpc({
        "type":"llm_batch",
        "prompts":[str(p) for p in prompts],
        "model":model,
        "dependency_mode":dependency_mode,
        "safety_note":safety_note,
    })
    if isinstance(resp, dict) and resp.get("error"):
        return [f"[llm_query_batched: {resp['error']}]" for _ in prompts]
    results = (resp or {}).get("results", []) if isinstance(resp, dict) else []
    if len(results) != len(prompts):
        return [f"[llm_query_batched: size mismatch ({len(results)}/{len(prompts)})]" for _ in prompts]
    out = []
    for r in results:
        if r.get("error"):
            out.append(f"[child err: {r['error']}]")
        else:
            out.append(r.get("text",""))
    return out

def rlm_query(prompt, model=None):
    """Recursive sub-RLM. The model arg is accepted for compatibility but ignored by Rust."""
    resp = _rpc({"type":"rlm","prompt":str(prompt),"model":model})
    if isinstance(resp, dict) and resp.get("error"):
        return f"[rlm_query error: {resp['error']}]"
    if isinstance(resp, dict):
        return resp.get("text","")
    return str(resp)

def rlm_query_batched(prompts, model=None, dependency_mode=None, safety_note=None):
    """Run independent recursive sub-RLMs in parallel. Declare dependency_mode='independent'."""
    if not isinstance(prompts, (list, tuple)):
        return ["[rlm_query_batched: prompts must be a list]"]
    err = _batch_dependency_error("rlm_query_batched", prompts, dependency_mode)
    if err is not None:
        return [err for _ in prompts]
    resp = _rpc({
        "type":"rlm_batch",
        "prompts":[str(p) for p in prompts],
        "model":model,
        "dependency_mode":dependency_mode,
        "safety_note":safety_note,
    })
    if isinstance(resp, dict) and resp.get("error"):
        return [f"[rlm_query_batched: {resp['error']}]" for _ in prompts]
    results = (resp or {}).get("results", []) if isinstance(resp, dict) else []
    if len(results) != len(prompts):
        return [f"[rlm_query_batched: size mismatch ({len(results)}/{len(prompts)})]" for _ in prompts]
    out = []
    for r in results:
        if r.get("error"):
            out.append(f"[child err: {r['error']}]")
        else:
            out.append(r.get("text",""))
    return out

def _slice_text(slice_value):
    if slice_value is None:
        return ""
    if isinstance(slice_value, dict):
        if "text" in slice_value:
            return str(slice_value["text"])
        return _json.dumps(slice_value, ensure_ascii=False)
    return str(slice_value)

def _prompt_with_slice(prompt, slice_value):
    text = _slice_text(slice_value)
    if not text:
        return str(prompt)
    if isinstance(slice_value, dict) and ("index" in slice_value or ("start" in slice_value and "end" in slice_value)):
        label = f"slice index={slice_value.get('index', '?')} range={slice_value.get('start', '?')}:{slice_value.get('end', '?')}"
    else:
        label = "slice"
    return f"{prompt}\n\n--- {label} ---\n{text}"

def sub_query(prompt, slice=None, timeout_secs=None, **kwargs):
    """One child LLM call, optionally scoped to a bounded slice."""
    return llm_query(_prompt_with_slice(prompt, slice))

def sub_query_batch(prompt, slices, timeout_secs=None, dependency_mode=None, safety_note=None, **kwargs):
    """Apply one prompt to many independent bounded slices concurrently."""
    if not isinstance(slices, (list, tuple)):
        return ["[sub_query_batch: slices must be a list]"]
    return llm_query_batched(
        [_prompt_with_slice(prompt, s) for s in slices],
        dependency_mode=dependency_mode,
        safety_note=safety_note,
    )

def sub_query_map(prompts, slices=None, timeout_secs=None, dependency_mode=None, safety_note=None, **kwargs):
    """Run N distinct independent prompts, optionally paired with N bounded slices."""
    if not isinstance(prompts, (list, tuple)):
        return ["[sub_query_map: prompts must be a list]"]
    if slices is None:
        return llm_query_batched(
            [str(p) for p in prompts],
            dependency_mode=dependency_mode,
            safety_note=safety_note,
        )
    if not isinstance(slices, (list, tuple)):
        return ["[sub_query_map: slices must be a list]"]
    if len(prompts) != len(slices):
        return [f"[sub_query_map: size mismatch ({len(prompts)}/{len(slices)})]" for _ in prompts]
    return llm_query_batched(
        [_prompt_with_slice(p, s) for p, s in zip(prompts, slices)],
        dependency_mode=dependency_mode,
        safety_note=safety_note,
    )

def sub_query_sequence(prompt, slices, carry_prompt=None, timeout_secs=None, **kwargs):
    """Apply one prompt to slices sequentially, feeding each result into the next step."""
    if not isinstance(slices, (list, tuple)):
        return ["[sub_query_sequence: slices must be a list]"]
    out = []
    previous = ""
    carry = str(carry_prompt or "Previous step result; treat it as required input for this step:")
    total = len(slices)
    for i, s in enumerate(slices):
        step_prompt = _prompt_with_slice(prompt, s)
        if previous:
            step_prompt = (
                f"{step_prompt}\n\n--- dependency_state step {i}/{total} ---\n"
                f"{carry}\n{previous}"
            )
        result = llm_query(step_prompt)
        out.append(result)
        previous = result
    return out

def sub_rlm(prompt, source=None, timeout_secs=None, **kwargs):
    """Recursive sub-RLM call for tasks that need their own decomposition."""
    return rlm_query(_prompt_with_slice(prompt, source))

def _json_safe(value):
    try:
        _json.dumps(value, ensure_ascii=False)
        return value
    except Exception:
        return str(value)

def _emit_final(value, confidence=None):
    safe_value = _json_safe(value)
    _sys.stdout.write(_FINAL + _json.dumps({
        "value": safe_value,
        "confidence": confidence,
    }, ensure_ascii=False) + "\n")
    _sys.stdout.flush()

def FINAL(value):
    """Legacy compatibility alias for finalize(value)."""
    _emit_final(value)

def FINAL_VAR(name):
    """Legacy compatibility alias for finalize(repl_get(name))."""
    name_str = str(name).strip().strip("'\"")
    if name_str in globals():
        _emit_final(globals()[name_str])
    else:
        print(f"FINAL_VAR error: variable '{name_str}' not found. "
              f"Use SHOW_VARS() to list available variables.", flush=True)

def SHOW_VARS():
    """Return a dict of {name: type-name} for all user variables in the REPL."""
    out = {}
    for k, v in list(globals().items()):
        if k.startswith('_') or k in _BOOTSTRAP_NAMES:
            continue
        out[k] = type(v).__name__
    return out

def repl_get(name, default=None):
    return globals().get(str(name), default)

def repl_set(name, value):
    globals()[str(name)] = value

def context_meta():
    """Return bounded metadata about the loaded input; never includes the full text."""
    text = _context
    line_count = 0 if text == "" else text.count("\n") + (0 if text.endswith("\n") else 1)
    return {
        "chars": len(text),
        "lines": line_count,
        "preview": text[:500],
        "tail_preview": text[-500:] if len(text) > 500 else text,
    }

def _slice_chars(start, end):
    total = len(_context)
    s = max(0, int(start))
    e = max(s, min(total, int(end)))
    return _context[s:e]

def _slice_lines(start, end):
    lines = _context.splitlines()
    s = max(0, int(start))
    e = max(s, min(len(lines), int(end)))
    return "\n".join(lines[s:e])

def peek(start, end, unit="chars"):
    """Return a bounded slice of the input by char offsets or line numbers."""
    if str(unit).lower() in ("line", "lines"):
        return _slice_lines(start, end)
    if str(unit).lower() not in ("char", "chars"):
        raise ValueError("unit must be 'chars' or 'lines'")
    return _slice_chars(start, end)

def search(pattern, max_hits=100):
    """Regex-search the input and return bounded hit records with snippets."""
    max_hits = max(0, int(max_hits))
    hits = []
    if max_hits == 0:
        return hits
    rx = _re.compile(str(pattern), _re.MULTILINE)
    for i, m in enumerate(rx.finditer(_context)):
        if i >= max_hits:
            break
        start, end = m.span()
        snippet_start = max(0, start - 120)
        snippet_end = min(len(_context), end + 120)
        hits.append({
            "index": i,
            "start": start,
            "end": end,
            "match": m.group(0),
            "snippet": _context[snippet_start:snippet_end],
        })
    return hits

def chunk(max_chars=20000, overlap=0):
    """Return full-coverage input chunks with index/start/end/text fields."""
    max_chars = int(max_chars)
    overlap = max(0, int(overlap))
    if max_chars <= 0:
        raise ValueError("max_chars must be > 0")
    if overlap >= max_chars:
        raise ValueError("overlap must be smaller than max_chars")
    chunks = []
    start = 0
    idx = 0
    total = len(_context)
    while start < total:
        end = min(total, start + max_chars)
        chunks.append({"index": idx, "start": start, "end": end, "text": _context[start:end]})
        idx += 1
        if end >= total:
            break
        start = end - overlap
    return chunks

def chunk_context(max_chars=20000, overlap=0):
    """Compatibility alias for chunk()."""
    return chunk(max_chars=max_chars, overlap=overlap)

def chunk_coverage(chunks):
    """Summarize coverage for chunks produced by chunk()."""
    spans = []
    for c in chunks:
        try:
            spans.append((int(c["start"]), int(c["end"])))
        except Exception:
            continue
    spans.sort()
    covered = 0
    cursor = 0
    gaps = []
    for start, end in spans:
        if start > cursor:
            gaps.append((cursor, start))
        if end > cursor:
            covered += end - max(start, cursor)
            cursor = end
    if cursor < len(_context):
        gaps.append((cursor, len(_context)))
    return {
        "chunks": len(chunks),
        "context_chars": len(_context),
        "input_chars": len(_context),
        "covered_chars": covered,
        "gaps": gaps,
        "complete": covered >= len(_context) and not gaps,
    }

def finalize(value, confidence=None):
    """Signal the session's final answer and persist confidence metadata."""
    global final_answer, final_confidence, final_result
    final_answer = _json_safe(value)
    final_confidence = confidence
    final_result = {
        "value": final_answer,
        "confidence": confidence,
    }
    _emit_final(final_answer, confidence=confidence)
    return final_answer

def evaluate_progress():
    """Return lightweight state useful before deciding the next REPL step."""
    vars_now = SHOW_VARS()
    return {
        "has_final_answer": "final_answer" in globals(),
        "final_confidence": globals().get("final_confidence", None),
        "user_variables": vars_now,
    }

# Load the long input from a file. This keeps the big string out of the
# process command-line and out of the LLM's window.
_ctx_file = _os.environ.get("RLM_CONTEXT_FILE","")
_context = ""
if _ctx_file:
    try:
        with open(_ctx_file, "r", encoding="utf-8", errors="replace") as f:
            _context = f.read()
    except Exception as e:
        _sys.stderr.write(f"[bootstrap] failed to load context: {e}\n")
content = _context

_BOOTSTRAP_NAMES = {
    "_SID","_REQ","_RESP","_FINAL","_ERR","_RUN","_END","_DONE","_READY",
    "_rpc","_ctx_file","_context","_slice_chars","_slice_lines","_BOOTSTRAP_NAMES","_main_loop",
    "_emit_final","_json_safe","_slice_text","_prompt_with_slice",
    "_normalize_dependency_mode","_batch_dependency_error",
    "llm_query","llm_query_batched","rlm_query","rlm_query_batched",
    "sub_query","sub_query_batch","sub_query_map","sub_query_sequence","sub_rlm",
    "FINAL","FINAL_VAR","SHOW_VARS","repl_get","repl_set",
    "context_meta","peek","search","chunk","chunk_context","chunk_coverage",
    "finalize","evaluate_progress","content",
    "_json","_os","_re","_sys","_traceback",
}

def _main_loop():
    _sys.stdout.write(_READY + "\n")
    _sys.stdout.flush()
    while True:
        header = _sys.stdin.readline()
        if not header:
            return
        if not header.startswith(_RUN):
            continue
        round_id = header.rstrip("\n")[len(_RUN):]
        code_lines = []
        while True:
            line = _sys.stdin.readline()
            if not line:
                return
            if line.rstrip("\n") == _END:
                break
            code_lines.append(line)
        code = "".join(code_lines)
        try:
            exec(compile(code, f"<repl-{round_id}>", "exec"), globals())
        except SystemExit:
            _sys.stdout.write(_DONE + round_id + "\n")
            _sys.stdout.flush()
            return
        except BaseException:
            tb = _traceback.format_exc()
            _sys.stdout.write(_ERR + _json.dumps(tb) + "\n")
            _sys.stdout.flush()
        _sys.stdout.write(_DONE + round_id + "\n")
        _sys.stdout.flush()

_main_loop()
"#;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn truncate_stdout(stdout: &str, limit: usize) -> String {
    if stdout.len() <= limit {
        return stdout.to_string();
    }
    let take = limit.saturating_sub(80);
    let mut out: String = stdout.chars().take(take).collect();
    let omitted = stdout.len().saturating_sub(out.len());
    out.push_str(&format!(
        "\n\n[... REPL output truncated: {omitted} bytes omitted ...]\n"
    ));
    out
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU32, Ordering};
    use tokio::sync::Mutex;

    /// In-process dispatcher that records what was asked and replies with
    /// canned text. Lets tests verify the round-trip without real network.
    struct StubBridge {
        calls: Arc<Mutex<Vec<RpcRequest>>>,
        canned: Arc<AtomicU32>,
    }

    impl StubBridge {
        fn new() -> Self {
            Self {
                calls: Arc::new(Mutex::new(Vec::new())),
                canned: Arc::new(AtomicU32::new(0)),
            }
        }
    }

    impl RpcDispatcher for StubBridge {
        fn dispatch<'a>(
            &'a self,
            req: RpcRequest,
        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = RpcResponse> + Send + 'a>> {
            Box::pin(async move {
                self.calls.lock().await.push(req.clone());
                let n = self.canned.fetch_add(1, Ordering::Relaxed);
                match req {
                    RpcRequest::Llm { prompt, .. } | RpcRequest::Rlm { prompt, .. } => {
                        RpcResponse::Single(SingleResp {
                            text: format!("stub#{n}: {prompt}"),
                            error: None,
                        })
                    }
                    RpcRequest::LlmBatch { prompts, .. } | RpcRequest::RlmBatch { prompts, .. } => {
                        let results = prompts
                            .into_iter()
                            .enumerate()
                            .map(|(i, p)| SingleResp {
                                text: format!("stub#{n}.{i}: {p}"),
                                error: None,
                            })
                            .collect();
                        RpcResponse::Batch(BatchResp { results })
                    }
                }
            })
        }
    }

    fn write_temp_context(body: &str) -> std::path::PathBuf {
        let dir = std::env::temp_dir().join("deepseek_repl_runtime_tests");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join(format!("ctx_{}_{}.txt", std::process::id(), Uuid::new_v4()));
        std::fs::write(&path, body).unwrap();
        path
    }

    #[tokio::test]
    async fn spawns_and_executes_simple_print() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt.execute("print('hello world')").await.expect("execute");
        assert!(round.stdout.contains("hello world"));
        assert!(!round.has_error);
        assert!(round.final_value.is_none());
        assert_eq!(round.rpc_count, 0);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn variables_persist_across_rounds() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        rt.execute("x = [1, 2, 3]").await.expect("r1");
        rt.execute("x.append(99)").await.expect("r2");
        let round = rt.execute("print(x)").await.expect("r3");
        assert!(round.stdout.contains("[1, 2, 3, 99]"));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn imports_persist_across_rounds() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        rt.execute("import math").await.expect("r1");
        let round = rt.execute("print(math.pi)").await.expect("r2");
        assert!(round.stdout.contains("3.14"));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn context_loads_from_file() {
        let path = write_temp_context("the quick brown fox");
        let mut rt = PythonRuntime::spawn_with_context(&path)
            .await
            .expect("spawn");
        let round = rt
            .execute("print(context_meta()['chars'], peek(0, 5))")
            .await
            .expect("execute");
        assert!(round.stdout.contains("19"));
        assert!(round.stdout.contains("the q"));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn context_aliases_keep_common_content_name_bounded() {
        let path = write_temp_context("aleph-style");
        let mut rt = PythonRuntime::spawn_with_context(&path)
            .await
            .expect("spawn");
        let round = rt
            .execute("print(content == _context, 'context' in globals(), 'ctx' in globals())")
            .await
            .expect("execute");
        assert!(round.stdout.contains("True False False"));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn context_chunk_helpers_report_full_coverage() {
        let path = write_temp_context("abcdefghijklmnopqrstuvwxyz");
        let mut rt = PythonRuntime::spawn_with_context(&path)
            .await
            .expect("spawn");
        let round = rt
            .execute(
                "chunks = chunk_context(max_chars=10)\n\
                 coverage = chunk_coverage(chunks)\n\
                 print(len(chunks), coverage['covered_chars'], coverage['complete'])",
            )
            .await
            .expect("execute");
        assert!(round.stdout.contains("3 26 True"), "{}", round.stdout);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn bounded_input_helpers_work() {
        let path = write_temp_context("alpha\nbeta needle\ngamma needle\nomega");
        let mut rt = PythonRuntime::spawn_with_context(&path)
            .await
            .expect("spawn");
        let round = rt
            .execute(
                "meta = context_meta()\n\
                 hits = search('needle', max_hits=1)\n\
                 print(meta['chars'], meta['lines'])\n\
                 print(peek(6, 17))\n\
                 print(peek(1, 3, unit='lines'))\n\
                 print(len(hits), hits[0]['match'], hits[0]['start'])",
            )
            .await
            .expect("execute");
        let stdout = round.stdout.replace("\r\n", "\n");
        assert!(stdout.contains("36 4"), "{stdout}");
        assert!(stdout.contains("beta needle"), "{stdout}");
        assert!(stdout.contains("beta needle\ngamma needle"), "{stdout}");
        assert!(stdout.contains("1 needle 11"), "{stdout}");
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn new_chunk_helper_reports_full_coverage() {
        let path = write_temp_context("abcdefghijklmnopqrstuvwxyz");
        let mut rt = PythonRuntime::spawn_with_context(&path)
            .await
            .expect("spawn");
        let round = rt
            .execute(
                "chunks = chunk(max_chars=10)\n\
                 coverage = chunk_coverage(chunks)\n\
                 print(len(chunks), coverage['input_chars'], coverage['covered_chars'], coverage['complete'])",
            )
            .await
            .expect("execute");
        assert!(round.stdout.contains("3 26 26 True"), "{}", round.stdout);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn finalize_helper_is_captured_directly() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .execute("finalize('computed answer', confidence='high')")
            .await
            .expect("execute");
        assert_eq!(round.final_value.as_deref(), Some("computed answer"));
        assert_eq!(
            round.final_json.as_ref().and_then(Value::as_str),
            Some("computed answer")
        );
        assert_eq!(
            round.final_confidence.as_ref().and_then(Value::as_str),
            Some("high")
        );
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn finalize_preserves_json_values_for_handles() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .execute("finalize({'answer': 42, 'items': ['a', 'b']})")
            .await
            .expect("execute");

        assert_eq!(
            round.final_value.as_deref(),
            Some(r#"{"answer":42,"items":["a","b"]}"#)
        );
        assert_eq!(
            round.final_json,
            Some(serde_json::json!({"answer": 42, "items": ["a", "b"]}))
        );
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn sub_query_accepts_timeout_keyword_for_agent_guesses() {
        let bridge = StubBridge::new();
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .run(
                "answer = sub_query('summarize', timeout_secs=2)\nprint(answer)",
                Some(&bridge),
            )
            .await
            .expect("execute");

        assert!(!round.has_error, "{}", round.stdout);
        assert!(
            round.stdout.contains("stub#0: summarize"),
            "{}",
            round.stdout
        );
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn rlm_context_runtime_has_no_fixed_round_timeout() {
        let path = write_temp_context("long input");
        let rt = PythonRuntime::spawn_with_context(&path)
            .await
            .expect("spawn");
        assert!(
            rt.round_timeout().is_none(),
            "RLM context runs must not inherit the old 180s REPL round timeout"
        );
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn inline_runtime_keeps_bounded_round_timeout() {
        let rt = PythonRuntime::new().await.expect("spawn");
        assert_eq!(rt.round_timeout(), Some(ROUND_TIMEOUT));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn legacy_final_is_captured() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .execute("FINAL('the answer is 42')")
            .await
            .expect("execute");
        assert_eq!(round.final_value.as_deref(), Some("the answer is 42"));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn legacy_final_var_is_captured() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        rt.execute("answer = 'computed'").await.expect("r1");
        let round = rt.execute("FINAL_VAR('answer')").await.expect("r2");
        assert_eq!(round.final_value.as_deref(), Some("computed"));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn errors_are_reported_without_killing_runtime() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let r1 = rt.execute("raise ValueError('boom')").await.expect("r1");
        assert!(r1.has_error);
        assert!(r1.full_stdout.contains("boom") || r1.stdout.contains("boom"));
        // The runtime is still alive — next round should work.
        let r2 = rt.execute("print('still here')").await.expect("r2");
        assert!(r2.stdout.contains("still here"));
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn rpc_dispatcher_round_trips_llm_query() {
        let bridge = StubBridge::new();
        let calls = Arc::clone(&bridge.calls);

        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .run("print(llm_query('hello'))", Some(&bridge))
            .await
            .expect("execute");
        assert!(
            round.stdout.contains("stub#0: hello"),
            "stdout: {:?}",
            round.stdout
        );
        assert_eq!(round.rpc_count, 1);

        let recorded = calls.lock().await;
        assert_eq!(recorded.len(), 1);
        match &recorded[0] {
            RpcRequest::Llm { prompt, .. } => assert_eq!(prompt, "hello"),
            other => panic!("expected Llm request, got {other:?}"),
        }
        drop(recorded);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn rpc_dispatcher_round_trips_sub_query_alias() {
        let bridge = StubBridge::new();
        let calls = Arc::clone(&bridge.calls);

        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .run("print(sub_query('hello from sub'))", Some(&bridge))
            .await
            .expect("execute");
        assert!(
            round.stdout.contains("stub#0: hello from sub"),
            "stdout: {:?}",
            round.stdout
        );
        assert_eq!(round.rpc_count, 1);

        let recorded = calls.lock().await;
        assert_eq!(recorded.len(), 1);
        match &recorded[0] {
            RpcRequest::Llm { prompt, .. } => assert_eq!(prompt, "hello from sub"),
            other => panic!("expected Llm request, got {other:?}"),
        }
        drop(recorded);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn rpc_dispatcher_round_trips_batch() {
        let bridge = StubBridge::new();
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .run(
                "outs = llm_query_batched(['a','b','c'], dependency_mode='independent', safety_note='same independent classification')\n\
                 print('|'.join(outs))",
                Some(&bridge),
            )
            .await
            .expect("execute");
        assert!(round.stdout.contains("stub#0.0: a"));
        assert!(round.stdout.contains("stub#0.1: b"));
        assert!(round.stdout.contains("stub#0.2: c"));
        assert_eq!(round.rpc_count, 1);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn batched_helpers_require_independence_declaration() {
        let bridge = StubBridge::new();
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .run(
                "outs = sub_query_batch('summarize', [{'text': 'a'}, {'text': 'b'}])\n\
                 print(outs[0])",
                Some(&bridge),
            )
            .await
            .expect("execute");

        assert!(
            round.stdout.contains("dependency_mode='independent'"),
            "{}",
            round.stdout
        );
        assert_eq!(round.rpc_count, 0);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn dependent_batch_mode_points_to_sequence_helper() {
        let bridge = StubBridge::new();
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .run(
                "outs = llm_query_batched(['migrate A', 'migrate B'], dependency_mode='sequential')\n\
                 print(outs[0])",
                Some(&bridge),
            )
            .await
            .expect("execute");

        assert!(
            round.stdout.contains("sub_query_sequence"),
            "{}",
            round.stdout
        );
        assert_eq!(round.rpc_count, 0);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn sub_query_sequence_feeds_prior_result_into_next_prompt() {
        let bridge = StubBridge::new();
        let calls = Arc::clone(&bridge.calls);

        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt
            .run(
                "outs = sub_query_sequence('process this step', [{'text': 'A'}, {'text': 'B'}])\n\
                 print(len(outs))",
                Some(&bridge),
            )
            .await
            .expect("execute");

        assert!(round.stdout.contains("2"), "{}", round.stdout);
        assert_eq!(round.rpc_count, 2);

        let recorded = calls.lock().await;
        assert_eq!(recorded.len(), 2);
        let second_prompt = match &recorded[1] {
            RpcRequest::Llm { prompt, .. } => prompt,
            other => panic!("expected second Llm request, got {other:?}"),
        };
        assert!(second_prompt.contains("--- dependency_state step 1/2 ---"));
        assert!(second_prompt.contains("stub#0: process this step"));
        drop(recorded);
        rt.shutdown().await;
    }

    #[tokio::test]
    async fn no_dispatcher_returns_unavailable_sentinel() {
        let mut rt = PythonRuntime::new().await.expect("spawn");
        let round = rt.execute("print(llm_query('hi'))").await.expect("execute");
        assert!(
            round.stdout.contains("[llm_query error:") || round.stdout.contains("no LLM bridge"),
            "stdout: {:?}",
            round.stdout
        );
        rt.shutdown().await;
    }

    #[test]
    fn truncate_keeps_short_unchanged() {
        assert_eq!(truncate_stdout("hello", 100), "hello");
    }

    #[test]
    fn truncate_clips_long() {
        let long = "a".repeat(10_000);
        let out = truncate_stdout(&long, 1024);
        assert!(out.len() < 1500);
        assert!(out.contains("truncated"));
    }
}