Skip to main content

newt_core/agentic/
driver.rs

1//! A **non-blocking turn driver** around [`chat_complete`](super::chat_complete)
2//! (issue #308 — the cowork foundation).
3//!
4//! [`chat_complete`] is `async` and, inside `run_chat`, is pumped by a blocking
5//! `rustyline` REPL (`block_in_place` + `block_on`). A downstream `ratatui`
6//! app cannot block its event loop like that: it must poll input, redraw, and
7//! poll the turn's progress every frame. This module is that seam — a struct an
8//! external `crossterm` event loop drives without ever calling `run_chat`:
9//!
10//! ```text
11//! loop {                                 // the consumer's crossterm loop
12//!     if event::poll(..) { /* read key, maybe driver.submit(line) */ }
13//!     match driver.poll() {              // non-blocking
14//!         TurnStatus::Completed { .. } => { /* redraw transcript */ }
15//!         TurnStatus::Running          => { /* show a spinner */ }
16//!         _ => {}
17//!     }
18//!     // turn driver.transcript() into rows via transcript_lines(.., width)
19//!     terminal.draw(|f| draw(transcript_lines(driver.transcript(), f.area().width as usize)));
20//! }
21//! ```
22//!
23//! ## How it spans the async boundary
24//!
25//! [`ChatCtx`](super::ChatCtx) is borrow-heavy — most of its fields are
26//! references with the caller's lifetime, several are `&mut dyn` session
27//! handles (note sink, recall source, permission gate, …). Those trait objects
28//! aren't `Send`, so `chat_complete`'s future isn't `Send` either and cannot go
29//! through `tokio::spawn`. So the driver owns a [`TurnDriverConfig`] of plain
30//! owned values and, on [`submit`](TurnDriver::submit), launches a **dedicated
31//! OS thread** running its own current-thread tokio runtime. That thread builds
32//! a **headless** `ChatCtx` from a clone of the config plus a snapshot of the
33//! transcript, `block_on`s one turn against [`NoMcp`] (raced against a cancel
34//! signal), and sends the result back over a [`oneshot`] channel. The
35//! session-bound `Option` fields (`note_sink`, `recall_source`, `summarizer`,
36//! `compress_state`, …) are all `None` here: a cowork turn is a clean drive,
37//! exactly like the ACP worker / `newt-eval` headless callers. Compression and
38//! budget logic are untouched — the driver wraps `chat_complete`, it does not
39//! reach inside it.
40//!
41//! Running on its own thread also means the driver does **not** require the
42//! consumer to call from inside a tokio runtime — a plain `crossterm` loop can
43//! own a `TurnDriver` with no async scaffolding of its own.
44//!
45//! ## Minimal by design
46//!
47//! `submit` → `poll` → `cancel`, plus `submit_observation` to fold a
48//! [`ShellObservation`] into the next turn's context. That is the whole
49//! surface; richer wiring (live MCP tools, a session note store) stays with the
50//! downstream consumer, which can build a full `ChatCtx` itself if it needs
51//! one.
52
53use std::thread::JoinHandle;
54
55use tokio::sync::oneshot;
56
57use super::observation::ShellObservation;
58use super::{chat_complete, openai_chat_complete, ChatCtx, NoMcp};
59use crate::{BackendKind, MemMessage, Role, TokenUsage};
60
61/// Owned, `'static` configuration for one [`TurnDriver`] — the cloneable
62/// counterpart of the borrow-heavy [`ChatCtx`]. The driver clones it into each
63/// spawned turn task so nothing borrows across the async boundary.
64///
65/// The fields mirror the inference-relevant subset of [`ChatCtx`]; the
66/// session-bound `&mut` handles (note sink, recall, permission gate, summarizer,
67/// compress state) are intentionally absent — a driven cowork turn is headless.
68#[derive(Debug, Clone)]
69pub struct TurnDriverConfig {
70    /// Inference endpoint base URL.
71    pub url: String,
72    /// Model name.
73    pub model: String,
74    /// Wire protocol of the backend.
75    pub kind: BackendKind,
76    /// Bearer token for authenticated OpenAI-compatible endpoints.
77    pub api_key: Option<String>,
78    /// Absolute workspace path the tool loop runs against.
79    pub workspace: String,
80    /// Permission caveats enforced for this turn's tool calls.
81    pub caveats: crate::caveats::Caveats,
82    /// Maximum tool-call rounds before a forced final completion.
83    pub max_tool_rounds: usize,
84    /// Max narrate-then-stop rescue nudges per turn (see
85    /// `[tui] narration_nudge_cap`); cowork default 1 keeps the historical
86    /// one-shot rescue.
87    pub narration_nudge_cap: usize,
88    /// Additional progress-aware rounds after `max_tool_rounds`; `0` makes the
89    /// normal cap hard.
90    pub workflow_grace_rounds: usize,
91    /// Max lines of tool output shown inline.
92    pub tool_output_lines: usize,
93    /// Ollama `options.num_ctx` (ignored on the OpenAI path).
94    pub num_ctx: Option<u32>,
95    /// TCP connect timeout.
96    pub connect_timeout_secs: u64,
97    /// Total inference timeout.
98    pub inference_timeout_secs: u64,
99    /// Message count at which the in-flight conversation is trimmed mid-turn.
100    pub mid_loop_trim_threshold: usize,
101    /// Token threshold that also triggers a mid-loop trim. `None` disables it.
102    pub mid_loop_trim_tokens: Option<usize>,
103    /// Highest input-token count the model has accepted (pre-send budget gate).
104    pub max_ok_input: Option<u32>,
105    /// Shell command run after every successful write for ground-truth feedback.
106    pub build_check_cmd: Option<String>,
107    /// Empirically safe context size used to detect likely overflow.
108    pub safe_context: Option<u32>,
109    /// #727: percent of `num_ctx` used as the input-token ceiling (default 80).
110    pub input_ceiling_pct: u32,
111    /// #727: remaining-budget percent below which the loop nudges (default 15).
112    pub low_budget_pct: usize,
113}
114
115impl TurnDriverConfig {
116    /// Construct a config with sensible cowork defaults; only the endpoint, the
117    /// model, the backend kind, and the workspace are required. Caveats default
118    /// to [`Caveats::top`](crate::caveats::Caveats::top) — the downstream
119    /// consumer is expected to narrow them.
120    pub fn new(
121        url: impl Into<String>,
122        model: impl Into<String>,
123        kind: BackendKind,
124        workspace: impl Into<String>,
125    ) -> Self {
126        Self {
127            url: url.into(),
128            model: model.into(),
129            kind,
130            api_key: None,
131            workspace: workspace.into(),
132            caveats: crate::caveats::Caveats::top(),
133            max_tool_rounds: 40,
134            narration_nudge_cap: 1,
135            workflow_grace_rounds: 5,
136            tool_output_lines: 20,
137            num_ctx: None,
138            connect_timeout_secs: 5,
139            inference_timeout_secs: 120,
140            mid_loop_trim_threshold: 40,
141            mid_loop_trim_tokens: None,
142            max_ok_input: None,
143            build_check_cmd: None,
144            safe_context: None,
145            input_ceiling_pct: 80,
146            low_budget_pct: 15,
147        }
148    }
149}
150
151/// The outcome of one completed turn.
152#[derive(Debug, Clone)]
153pub struct TurnOutcome {
154    /// The model's reply text.
155    pub reply: String,
156    /// Whether the reply was streamed token-by-token by the loop (it prints to
157    /// stdout when `color` is on; the driver always runs headless, so this is
158    /// informational).
159    pub was_streamed: bool,
160    /// Token usage for the turn, when the backend reported it.
161    pub usage: Option<TokenUsage>,
162    /// Count of hallucinated (non-existent) tool calls the loop saw.
163    pub hallucinations: u32,
164}
165
166/// Non-blocking snapshot of the driver's state, returned by
167/// [`TurnDriver::poll`].
168#[derive(Debug)]
169pub enum TurnStatus {
170    /// No turn in flight (nothing submitted, or the last result was drained).
171    Idle,
172    /// A turn is running; the result is not ready yet.
173    Running,
174    /// A turn finished successfully. Returned **once**; the reply has already
175    /// been appended to the transcript as an assistant message.
176    Completed(TurnOutcome),
177    /// A turn failed (transport error, dispatch error, …). Returned **once**.
178    Failed(String),
179}
180
181/// Internal in-flight turn handle: the worker thread, the channel it reports
182/// on, and the cancel signal that races its `block_on`.
183struct InFlight {
184    handle: JoinHandle<()>,
185    rx: oneshot::Receiver<Result<TurnOutcome, String>>,
186    cancel_tx: Option<oneshot::Sender<()>>,
187}
188
189/// A pumpable, non-blocking driver for one agentic turn at a time.
190///
191/// Owns the running transcript and at most one in-flight turn. The external
192/// event loop [`submit`](Self::submit)s input (or
193/// [`submit_observation`](Self::submit_observation)s shell activity),
194/// [`poll`](Self::poll)s for completion every frame, and may
195/// [`cancel`](Self::cancel) the in-flight turn.
196pub struct TurnDriver {
197    config: TurnDriverConfig,
198    transcript: Vec<MemMessage>,
199    in_flight: Option<InFlight>,
200}
201
202impl TurnDriver {
203    /// Build a driver with an empty transcript.
204    pub fn new(config: TurnDriverConfig) -> Self {
205        Self {
206            config,
207            transcript: Vec::new(),
208            in_flight: None,
209        }
210    }
211
212    /// Build a driver seeded with an existing transcript (e.g. a system prompt
213    /// plus prior turns the consumer already assembled).
214    pub fn with_transcript(config: TurnDriverConfig, transcript: Vec<MemMessage>) -> Self {
215        Self {
216            config,
217            transcript,
218            in_flight: None,
219        }
220    }
221
222    /// The current transcript — every message the model has seen plus the
223    /// replies it produced. Feed it to
224    /// [`transcript_lines`](super::transcript_lines) for renderer-agnostic,
225    /// width-wrapped rows, or render it however the consumer likes.
226    pub fn transcript(&self) -> &[MemMessage] {
227        &self.transcript
228    }
229
230    /// Whether a turn is currently in flight.
231    pub fn is_running(&self) -> bool {
232        self.in_flight.is_some()
233    }
234
235    /// Append a [`ShellObservation`] to the transcript so it becomes part of the
236    /// **next** turn's context. The observation is redacted and framed by
237    /// [`ShellObservation::into_mem_message`] — credentials in shell output are
238    /// scrubbed before they enter the transcript. This does **not** start a
239    /// turn on its own; a real human message ([`submit`](Self::submit)) drives
240    /// the loop, with the accumulated observations already in context.
241    pub fn submit_observation(&mut self, obs: ShellObservation) {
242        self.transcript.push(obs.into_mem_message());
243    }
244
245    /// Submit a human message and start a turn.
246    ///
247    /// Appends the message as a `User` turn, snapshots the transcript, and
248    /// spawns a task that runs one [`chat_complete`] against the configured
249    /// backend. Returns an error (without starting anything) if a turn is
250    /// already in flight — the driver runs one turn at a time. Poll
251    /// [`poll`](Self::poll) for the result.
252    pub fn submit(&mut self, input: impl Into<String>) -> Result<(), TurnDriverError> {
253        if self.in_flight.is_some() {
254            return Err(TurnDriverError::Busy);
255        }
256        let task = input.into();
257        self.transcript.push(MemMessage::user(task.clone()));
258        self.spawn_turn(task);
259        Ok(())
260    }
261
262    /// Launch the turn on a dedicated thread over a clone of the config + a
263    /// snapshot of the transcript. The `task` string is the current user
264    /// message, threaded into `ChatCtx.task` (the loop's nudges key on it).
265    ///
266    /// `chat_complete`'s future is `!Send` (its `ChatCtx` holds `&mut dyn`
267    /// session handles), so it can't ride `tokio::spawn`. A current-thread
268    /// runtime on its own OS thread keeps the non-`Send` future local; a
269    /// cancel oneshot races the turn so [`cancel`](Self::cancel) returns the
270    /// thread promptly instead of waiting out the inference timeout.
271    fn spawn_turn(&mut self, task: String) {
272        let config = self.config.clone();
273        let messages = self.transcript.clone();
274        let (tx, rx) = oneshot::channel();
275        let (cancel_tx, cancel_rx) = oneshot::channel();
276        let handle = std::thread::spawn(move || {
277            // A current-thread runtime: no work-stealing, nothing to require
278            // `Send`, and self-contained so the consumer needn't be in a
279            // runtime itself.
280            let rt = match tokio::runtime::Builder::new_current_thread()
281                .enable_all()
282                .build()
283            {
284                Ok(rt) => rt,
285                Err(e) => {
286                    let _ = tx.send(Err(format!("failed to start turn runtime: {e}")));
287                    return;
288                }
289            };
290            let result = rt.block_on(async move {
291                tokio::select! {
292                    biased;
293                    // Cancellation wins the race when signalled.
294                    _ = cancel_rx => Err("turn cancelled".to_string()),
295                    out = run_one_turn(&config, &messages, &task) => out,
296                }
297            });
298            // The receiver may have been dropped (driver went away); fine.
299            let _ = tx.send(result);
300        });
301        self.in_flight = Some(InFlight {
302            handle,
303            rx,
304            cancel_tx: Some(cancel_tx),
305        });
306    }
307
308    /// Non-blocking poll for the in-flight turn's status.
309    ///
310    /// - [`TurnStatus::Idle`] — nothing in flight.
311    /// - [`TurnStatus::Running`] — in flight, not done.
312    /// - [`TurnStatus::Completed`] — done; the reply has been appended to the
313    ///   transcript as an assistant message. Returned exactly once.
314    /// - [`TurnStatus::Failed`] — errored; nothing appended. Returned once.
315    pub fn poll(&mut self) -> TurnStatus {
316        let Some(in_flight) = self.in_flight.as_mut() else {
317            return TurnStatus::Idle;
318        };
319        match in_flight.rx.try_recv() {
320            Ok(Ok(outcome)) => {
321                self.transcript
322                    .push(MemMessage::assistant(outcome.reply.clone()));
323                self.in_flight = None;
324                TurnStatus::Completed(outcome)
325            }
326            Ok(Err(err)) => {
327                self.in_flight = None;
328                TurnStatus::Failed(err)
329            }
330            Err(oneshot::error::TryRecvError::Empty) => TurnStatus::Running,
331            Err(oneshot::error::TryRecvError::Closed) => {
332                // The task ended without sending (it was aborted, or panicked).
333                self.in_flight = None;
334                TurnStatus::Failed("turn task ended without a result".to_string())
335            }
336        }
337    }
338
339    /// Cancel the in-flight turn, if any. Signals the worker thread to abandon
340    /// the turn (winning the `select!` race against `chat_complete`) and joins
341    /// it so the runtime tears down before returning. The user message that
342    /// started the turn stays in the transcript — the consumer can pop it for a
343    /// clean retry. No-op when idle.
344    ///
345    /// A turn already blocked deep inside a single HTTP dispatch returns when
346    /// that dispatch resolves (or times out); the cancel signal is honored at
347    /// the next `select!` poll point, which in practice is between dispatches.
348    pub fn cancel(&mut self) {
349        if let Some(mut in_flight) = self.in_flight.take() {
350            if let Some(cancel_tx) = in_flight.cancel_tx.take() {
351                let _ = cancel_tx.send(());
352            }
353            let _ = in_flight.handle.join();
354        }
355    }
356}
357
358/// Run exactly one turn: assemble a headless [`ChatCtx`] from owned config +
359/// messages and dispatch on the right wire protocol. Factored out of the spawn
360/// closure so it is directly unit-testable against a wiremock backend.
361async fn run_one_turn(
362    config: &TurnDriverConfig,
363    messages: &[MemMessage],
364    task: &str,
365) -> Result<TurnOutcome, String> {
366    let ctx = ChatCtx {
367        url: &config.url,
368        model: &config.model,
369        kind: config.kind,
370        api_key: config.api_key.as_deref(),
371        messages,
372        task,
373        workspace: &config.workspace,
374        // The driver is headless: the loop's inline progress prints are
375        // suppressed so nothing fights the consumer's ratatui frame.
376        color: false,
377        // Cowork renders into the consumer's frame, not the scroller — no
378        // markdown ANSI here (it would fight the host UI).
379        markdown: false,
380        // Headless: no tool-offload (26.3) / scratchpad (26.4) — bit-for-bit.
381        tool_offload: false,
382        spill_store: None,
383        compaction_store: None,
384        scratchpad: false,
385        scratchpad_store: None,
386        code_search: None,
387        experience_store: None,
388        step_ledger: None,
389        caveats: &config.caveats,
390        // Headless cowork driver carries no persona surface (FR-1 part 2, #997).
391        persona_tools: None,
392        max_tool_rounds: config.max_tool_rounds,
393        narration_nudge_cap: config.narration_nudge_cap,
394        workflow_grace_rounds: config.workflow_grace_rounds,
395        tool_output_lines: config.tool_output_lines,
396        debug: false,
397        trace: false,
398        num_ctx: config.num_ctx,
399        // #727 tunables: headless driver keeps the historical defaults
400        // (80% input ceiling, 15% low-budget nudge threshold).
401        input_ceiling_pct: config.input_ceiling_pct,
402        low_budget_pct: config.low_budget_pct,
403        connect_timeout_secs: config.connect_timeout_secs,
404        inference_timeout_secs: config.inference_timeout_secs,
405        mid_loop_trim_threshold: config.mid_loop_trim_threshold,
406        mid_loop_trim_tokens: config.mid_loop_trim_tokens,
407        max_ok_input: config.max_ok_input,
408        build_check_cmd: config.build_check_cmd.clone(),
409        safe_context: config.safe_context,
410        // Session-bound seams a driven cowork turn does not carry (headless,
411        // exactly like the ACP worker / newt-eval callers).
412        recover_cw_400: None,
413        note_sink: None,
414        note_nudge: None,
415        recall_source: None,
416        memory_source: None,
417        summarizer: None,
418        compress_state: None,
419        tool_events: None,
420        phantom_reaches: None,
421        end_reason: None,
422        permission_gate: None,
423        // Phase 20 (spec §5): headless surfaces neither read nor write the
424        // capability cache — the hook stays absent and no calibration is
425        // applied, preserving today's behavior exactly.
426        on_round_usage: None,
427        estimate_ratio: None,
428        estimation: crate::tokens::TokenEstimation::default(),
429        summary_input_cap_floor_chars: 8_192,
430        // #307: the headless driver carries no preset clamp — exec authority is
431        // whatever `config.caveats` already grants, exactly like the ACP worker
432        // / newt-eval callers. A consumer enforcing a mode clamps `caveats` itself.
433        exec_floor: None,
434        write_ledger: None,
435        // The headless driver has no interactive keyboard to interrupt from;
436        // a consumer that wants cancellation drives its own ChatCtx.
437        cancel: None,
438        git_tool: None,
439        crew_runner: None,
440    };
441    // NoMcp: the cowork driver advertises only the built-in tools. A consumer
442    // that wants live MCP tools assembles its own ChatCtx.
443    let mut mcp = NoMcp;
444    let dispatch = if config.kind == BackendKind::Openai {
445        openai_chat_complete(ctx, &mut mcp).await
446    } else {
447        chat_complete(ctx, &mut mcp).await
448    };
449    match dispatch {
450        Ok((reply, was_streamed, usage, hallucinations)) => Ok(TurnOutcome {
451            reply,
452            was_streamed,
453            usage,
454            hallucinations,
455        }),
456        Err(e) => Err(e.to_string()),
457    }
458}
459
460/// Errors from driving a turn.
461#[derive(Debug, thiserror::Error)]
462pub enum TurnDriverError {
463    /// A turn is already in flight; the driver runs one at a time.
464    #[error("a turn is already in flight — poll() for it or cancel() before submitting another")]
465    Busy,
466}
467
468/// Roles whose messages a transcript renderer typically shows to the human. The
469/// driver exposes it so a consumer's render can filter on the same set the
470/// newt-tui render uses (system + tool turns are usually hidden).
471pub const VISIBLE_TRANSCRIPT_ROLES: [Role; 2] = [Role::User, Role::Assistant];
472
473#[cfg(test)]
474mod tests {
475    use super::*;
476    use crate::caveats::Caveats;
477    use std::sync::atomic::{AtomicUsize, Ordering};
478    use std::sync::Arc;
479    use std::time::Duration;
480    use wiremock::matchers::{method, path};
481    use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate};
482
483    /// Ollama-shaped responder: a plain text answer (no tools), echoing back a
484    /// fixed reply. Counts how many chat requests it served.
485    struct PlainOllama {
486        served: Arc<AtomicUsize>,
487        reply: String,
488    }
489
490    impl Respond for PlainOllama {
491        fn respond(&self, _req: &Request) -> ResponseTemplate {
492            self.served.fetch_add(1, Ordering::SeqCst);
493            ResponseTemplate::new(200).set_body_json(serde_json::json!({
494                "message": { "content": self.reply }
495            }))
496        }
497    }
498
499    fn cfg(url: &str) -> TurnDriverConfig {
500        let mut c = TurnDriverConfig::new(url, "test-model", BackendKind::Ollama, ".");
501        c.caveats = Caveats::top();
502        c
503    }
504
505    /// Pump the driver to completion the way a crossterm loop would: poll on an
506    /// interval until the turn resolves, never blocking the "frame".
507    async fn pump_to_done(driver: &mut TurnDriver) -> TurnStatus {
508        for _ in 0..600 {
509            match driver.poll() {
510                TurnStatus::Running => tokio::time::sleep(Duration::from_millis(10)).await,
511                other => return other,
512            }
513        }
514        panic!("turn did not complete within the pump budget");
515    }
516
517    /// THE acceptance test: a consumer drives one turn through the driver and
518    /// gets the answer back — no `run_chat`, no blocking REPL, only the
519    /// submit → poll surface.
520    #[tokio::test]
521    async fn driver_pumps_one_turn_and_yields_the_answer() {
522        let server = MockServer::start().await;
523        let served = Arc::new(AtomicUsize::new(0));
524        Mock::given(method("POST"))
525            .and(path("/api/chat"))
526            .respond_with(PlainOllama {
527                served: served.clone(),
528                reply: "the driver answered".into(),
529            })
530            .mount(&server)
531            .await;
532
533        let mut driver = TurnDriver::new(cfg(&server.uri()));
534        // Idle before anything is submitted.
535        assert!(matches!(driver.poll(), TurnStatus::Idle));
536        assert!(!driver.is_running());
537
538        driver
539            .submit("what is 2 + 2?")
540            .expect("submit starts a turn");
541        assert!(driver.is_running());
542
543        let status = pump_to_done(&mut driver).await;
544        match status {
545            TurnStatus::Completed(outcome) => {
546                assert_eq!(outcome.reply, "the driver answered");
547            }
548            other => panic!("expected Completed, got {other:?}"),
549        }
550
551        // The transcript now holds the user turn AND the assistant reply, with
552        // no run_chat involved.
553        let t = driver.transcript();
554        assert_eq!(t.len(), 2);
555        assert_eq!(t[0].role, Role::User);
556        assert_eq!(t[0].content, "what is 2 + 2?");
557        assert_eq!(t[1].role, Role::Assistant);
558        assert_eq!(t[1].content, "the driver answered");
559
560        // The backend served exactly the streaming re-issue + probe (the
561        // tools-present probe then the stream); at least one request landed.
562        assert!(served.load(Ordering::SeqCst) >= 1);
563        // Drained — back to idle.
564        assert!(matches!(driver.poll(), TurnStatus::Idle));
565        assert!(!driver.is_running());
566    }
567
568    /// A shell observation feeds the NEXT turn's context, redacted. We prove the
569    /// observation reaches the backend's request body (so the model sees it) and
570    /// that a secret in it never does.
571    #[tokio::test]
572    async fn observation_is_in_context_for_the_next_turn_and_redacted() {
573        let server = MockServer::start().await;
574        // Capture the request body the model would see.
575        let seen = Arc::new(std::sync::Mutex::new(String::new()));
576        struct Capture {
577            seen: Arc<std::sync::Mutex<String>>,
578        }
579        impl Respond for Capture {
580            fn respond(&self, req: &Request) -> ResponseTemplate {
581                *self.seen.lock().unwrap() = String::from_utf8_lossy(&req.body).into_owned();
582                ResponseTemplate::new(200).set_body_json(serde_json::json!({
583                    "message": { "content": "ack" }
584                }))
585            }
586        }
587        Mock::given(method("POST"))
588            .and(path("/api/chat"))
589            .respond_with(Capture { seen: seen.clone() })
590            .mount(&server)
591            .await;
592
593        let mut driver = TurnDriver::new(cfg(&server.uri()));
594        driver.submit_observation(ShellObservation::new(
595            "zsh",
596            "$ cat creds\nsecret_key=wJalrXUtnFEMIabcdEFGHIJKLMNOPbPxRfiCY",
597        ));
598        // Observation alone does not start a turn.
599        assert!(!driver.is_running());
600        assert_eq!(driver.transcript().len(), 1);
601
602        driver.submit("did you see my shell?").expect("submit");
603        let _ = pump_to_done(&mut driver).await;
604
605        let body = seen.lock().unwrap().clone();
606        // The observation framing reached the model …
607        assert!(
608            body.contains("shell observation"),
609            "observation missing from request body: {body}"
610        );
611        // … but the secret value did not.
612        assert!(
613            !body.contains("wJalrXUtnFEMIabcdEFGHIJKLMNOPbPxRfiCY"),
614            "secret leaked into the request body: {body}"
615        );
616    }
617
618    /// Submitting while a turn is in flight is rejected — one turn at a time.
619    #[tokio::test]
620    async fn second_submit_while_running_is_busy() {
621        let server = MockServer::start().await;
622        Mock::given(method("POST"))
623            .and(path("/api/chat"))
624            // Delay so the first turn is reliably still in flight.
625            .respond_with(
626                ResponseTemplate::new(200)
627                    .set_delay(Duration::from_millis(200))
628                    .set_body_json(serde_json::json!({ "message": { "content": "slow" } })),
629            )
630            .mount(&server)
631            .await;
632
633        let mut driver = TurnDriver::new(cfg(&server.uri()));
634        driver.submit("first").expect("first submit");
635        let err = driver
636            .submit("second")
637            .expect_err("second must be rejected");
638        assert!(matches!(err, TurnDriverError::Busy));
639        // Let it finish so the test doesn't leak the task.
640        let _ = pump_to_done(&mut driver).await;
641    }
642
643    /// Cancel aborts the in-flight turn and returns the driver to idle.
644    #[tokio::test]
645    async fn cancel_aborts_the_in_flight_turn() {
646        let server = MockServer::start().await;
647        Mock::given(method("POST"))
648            .and(path("/api/chat"))
649            .respond_with(
650                ResponseTemplate::new(200)
651                    .set_delay(Duration::from_secs(30))
652                    .set_body_json(serde_json::json!({ "message": { "content": "never" } })),
653            )
654            .mount(&server)
655            .await;
656
657        let mut driver = TurnDriver::new(cfg(&server.uri()));
658        driver.submit("start a slow turn").expect("submit");
659        assert!(driver.is_running());
660        driver.cancel();
661        assert!(!driver.is_running());
662        assert!(matches!(driver.poll(), TurnStatus::Idle));
663        // The user message that started it survives for a clean retry.
664        assert_eq!(driver.transcript().len(), 1);
665        assert_eq!(driver.transcript()[0].content, "start a slow turn");
666    }
667}