返回 CodeWhale
handle.rs
根目录 / crates / tui / src / core / engine / handle.rs
1 //! Public `EngineHandle` methods.
2 //!
3 //! The struct itself lives next door in `engine.rs` because two
4 //! construction sites (`Engine::new` and the test-only
5 //! `mock_engine_handle`) need access to its private mpsc channels.
6 //! The method surface — `send`, `cancel*`, `is_cancelled`,
7 //! `approve_tool_call` / `deny_tool_call` / `retry_tool_with_policy`,
8 //! `submit_user_input` / `cancel_user_input`, and `steer` — moves here
9 //! so the agent loop's mailbox API is reviewable on its own.
10
11 use anyhow::Result;
12 use std::collections::VecDeque;
13 use std::sync::{Arc, Mutex as StdMutex};
14 use tokio::sync::{mpsc, oneshot};
15 use tokio_util::sync::CancellationToken;
16
17 use codewhale_config::AppMode;
18 use codewhale_execpolicy::ApprovalMode;
19
20 use super::approval::{ApprovalDecision, UserInputDecision};
21 use super::{
22 CancelReason, EngineHandle, LiveRuntimeAuthority, Op, RuntimePermissionAuthority,
23 UserInputResponse,
24 };
25
26 #[derive(Clone)]
27 pub(super) struct TurnControl {
28 pub id: u64,
29 pub cancel: CancellationToken,
30 pub reason: Arc<StdMutex<Option<CancelReason>>>,
31 }
32
33 #[derive(Default)]
34 pub(super) struct TurnControls {
35 next_id: u64,
36 pub active: Option<TurnControl>,
37 pub pending: VecDeque<TurnControl>,
38 }
39
40 impl TurnControls {
41 pub fn fresh(&mut self) -> TurnControl {
42 self.next_id = self
43 .next_id
44 .checked_add(1)
45 .expect("turn control id exhausted");
46 TurnControl {
47 id: self.next_id,
48 cancel: CancellationToken::new(),
49 reason: Arc::new(StdMutex::new(None)),
50 }
51 }
52
53 fn target(&self) -> Option<&TurnControl> {
54 self.active.as_ref().or_else(|| self.pending.front())
55 }
56 }
57
58 pub(super) struct TurnControlGuard {
59 pub controls: Arc<StdMutex<TurnControls>>,
60 pub id: u64,
61 }
62
63 impl Drop for TurnControlGuard {
64 fn drop(&mut self) {
65 let mut controls = self
66 .controls
67 .lock()
68 .unwrap_or_else(std::sync::PoisonError::into_inner);
69 if controls
70 .active
71 .as_ref()
72 .is_some_and(|active| active.id == self.id)
73 {
74 controls.active = None;
75 }
76 }
77 }
78
79 #[derive(Debug)]
80 pub(crate) struct SteerInput {
81 pub(super) turn_id: Option<u64>,
82 pub(crate) content: String,
83 pub(super) outcome: Option<oneshot::Sender<SteerOutcome>>,
84 }
85
86 /// The engine's verdict on one steer. A steer whose turn had already moved
87 /// on is discarded by `next_turn_steer`; the verdict tells the sender which
88 /// happened, because "the channel accepted the text" is not "the model saw
89 /// it" (#6276).
90 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
91 pub(crate) enum SteerOutcome {
92 /// The steer's text was committed into the session record inside the
93 /// turn it was sent to; the model received it.
94 Accepted,
95 /// The turn had already moved on (or ended) before the steer reached a
96 /// commit boundary. The model never saw the text.
97 Dropped,
98 }
99
100 /// A steer the engine has taken ownership of. Committing it reports
101 /// [`SteerOutcome::Accepted`]; any other exit — interrupt, failure, early
102 /// return, silent drop of the pending queue — reports `Dropped` from `Drop`,
103 /// so no path can lose a verdict.
104 pub(crate) struct PendingSteer {
105 pub(crate) content: String,
106 outcome: Option<oneshot::Sender<SteerOutcome>>,
107 }
108
109 impl PendingSteer {
110 pub(crate) fn new(content: String, outcome: Option<oneshot::Sender<SteerOutcome>>) -> Self {
111 Self { content, outcome }
112 }
113
114 /// Commit the steer into the turn's record: report `Accepted`, then hand
115 /// back the text. Consuming `self` without calling this reports
116 /// `Dropped` via `Drop`.
117 pub(crate) fn commit(mut self) -> String {
118 if let Some(outcome) = self.outcome.take() {
119 let _ = outcome.send(SteerOutcome::Accepted);
120 }
121 // `Drop` runs after this returns and finds `outcome` already taken,
122 // so the verdict stays exactly one `Accepted`.
123 std::mem::take(&mut self.content)
124 }
125 }
126
127 impl Drop for PendingSteer {
128 fn drop(&mut self) {
129 if let Some(outcome) = self.outcome.take() {
130 let _ = outcome.send(SteerOutcome::Dropped);
131 }
132 }
133 }
134
135 impl SteerInput {
136 /// Take ownership of this steer as an unsettled [`PendingSteer`].
137 ///
138 /// This is the only way to claim a steer off the channel. Whatever the
139 /// claimant then does — `commit()` or drop — settles it exactly once, so
140 /// there is one settlement mechanism rather than two (#6276).
141 pub(crate) fn into_pending(mut self) -> PendingSteer {
142 // Both fields are taken, so the `Drop` below finds nothing left to
143 // settle and the verdict travels with the `PendingSteer`.
144 PendingSteer::new(std::mem::take(&mut self.content), self.outcome.take())
145 }
146 }
147
148 impl Drop for SteerInput {
149 fn drop(&mut self) {
150 if let Some(outcome) = self.outcome.take() {
151 let _ = outcome.send(SteerOutcome::Dropped);
152 }
153 }
154 }
155
156 impl std::ops::Deref for SteerInput {
157 type Target = str;
158 fn deref(&self) -> &str {
159 &self.content
160 }
161 }
162
163 impl std::fmt::Display for SteerInput {
164 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
165 self.content.fmt(f)
166 }
167 }
168
169 pub(crate) struct SteerPermit {
170 permit: mpsc::OwnedPermit<SteerInput>,
171 turn_id: Option<u64>,
172 }
173
174 impl SteerPermit {
175 pub(crate) fn send(self, content: String) {
176 self.permit.send(SteerInput {
177 turn_id: self.turn_id,
178 content,
179 outcome: None,
180 });
181 }
182
183 /// Send a steer and receive the engine's verdict on it. The receiver
184 /// resolves to [`SteerOutcome::Accepted`] when the turn commits the text
185 /// into its record, [`SteerOutcome::Dropped`] when the turn moved on
186 /// first, and closes without a verdict only if the engine itself is gone
187 /// (#6276).
188 pub(crate) fn send_with_outcome(self, content: String) -> oneshot::Receiver<SteerOutcome> {
189 let (outcome_tx, outcome_rx) = oneshot::channel();
190 self.permit.send(SteerInput {
191 turn_id: self.turn_id,
192 content,
193 outcome: Some(outcome_tx),
194 });
195 outcome_rx
196 }
197 }
198
199 impl EngineHandle {
200 /// Called only while Runtime holds the idle turn admission claim. The
201 /// following SendMessage refreshes the existing prompt/config projection.
202 pub(crate) fn restore_runtime_goal(
203 &self,
204 goal: Option<&codewhale_protocol::ThreadGoal>,
205 ) -> Result<()> {
206 let mut state = self
207 .goal_state
208 .lock()
209 .map_err(|_| anyhow::anyhow!("goal state lock poisoned"))?;
210 let current = state.snapshot();
211 if current.goal_id.as_deref() != goal.map(|goal| goal.goal_id.as_str()) {
212 *state = goal.map_or_else(crate::tools::goal::GoalState::default, |goal| {
213 crate::tools::goal::GoalState::from_snapshot(
214 &crate::tools::goal::GoalSnapshot::from_thread_goal(goal),
215 )
216 });
217 }
218 Ok(())
219 }
220
221 /// True when the caller must preflight a concrete provider client before
222 /// committing UI/runtime turn state. Test and embedding handles with an
223 /// injected model client return false because that client owns model I/O.
224 #[must_use]
225 pub(crate) fn client_preflight_required(&self) -> bool {
226 self.client_preflight_required
227 }
228
229 /// Send an operation to the engine
230 ///
231 /// This awaits channel capacity, and the engine drains `rx_op` only
232 /// between turns — so on the UI event loop an awaited send into a
233 /// saturated mailbox freezes input for the rest of the turn (#6150).
234 /// Input-path callers instead either `try_send` a droppable op (report
235 /// the rejection) or `try_reserve_owned` before committing UI state and
236 /// hand off with `send_reserved_op`. An awaited `send` remains correct
237 /// only where the operation is part of a committed, ordered transition
238 /// (session/provider reload) whose drop would desync engine and UI.
239 pub async fn send(&self, op: Op) -> Result<()> {
240 let authority = Self::change_mode_authority(&op);
241 let permit = self.tx_op.clone().reserve_owned().await?;
242 if let Some(authority) = authority {
243 self.publish_runtime_authority(authority);
244 }
245 self.send_reserved_op(permit, op);
246 Ok(())
247 }
248
249 /// Try to send an operation without blocking.
250 ///
251 /// Returns `Err` if the channel is full or closed. Use this for
252 /// non-critical, refresh-type ops (e.g. `Op::ListSubAgents`) that can
253 /// safely be dropped and re-requested on the next drain cycle.
254 pub fn try_send(&self, op: Op) -> Result<()> {
255 let authority = Self::change_mode_authority(&op);
256 let result = self.tx_op.clone().try_reserve_owned();
257 // A full channel already guarantees that the engine will wake and
258 // drain an operation. Publish the typed authority anyway: the drain
259 // applies pending authority before handling that queued operation, so
260 // a posture edit never blocks behind refresh traffic. A closed
261 // channel has no engine left to observe the update.
262 if !matches!(&result, Err(mpsc::error::TrySendError::Closed(_)))
263 && let Some(authority) = authority
264 {
265 self.publish_runtime_authority(authority);
266 }
267 // Keep the public error bound to the rejected operation. Callers use
268 // TrySendError<Op> to distinguish a retryable full mailbox from a
269 // stopped engine; reservation errors otherwise carry a Sender<Op>.
270 match result {
271 Ok(permit) => {
272 self.send_reserved_op(permit, op);
273 Ok(())
274 }
275 Err(mpsc::error::TrySendError::Full(_)) => {
276 Err(mpsc::error::TrySendError::Full(op).into())
277 }
278 Err(mpsc::error::TrySendError::Closed(_)) => {
279 Err(mpsc::error::TrySendError::Closed(op).into())
280 }
281 }
282 }
283
284 /// Bind controls and enqueue under one lock, preserving the same FIFO as
285 /// the operation mailbox even when several senders hold reserved slots.
286 pub(crate) fn send_reserved_op(&self, permit: mpsc::OwnedPermit<Op>, op: Op) {
287 let mut controls = self
288 .turn_controls
289 .lock()
290 .unwrap_or_else(std::sync::PoisonError::into_inner);
291 if matches!(&op, Op::SendMessage(_)) {
292 let control = controls.fresh();
293 controls.pending.push_back(control);
294 }
295 permit.send(op);
296 }
297
298 fn change_mode_authority(op: &Op) -> Option<LiveRuntimeAuthority> {
299 let Op::ChangeMode {
300 mode,
301 allow_shell,
302 trust_mode,
303 auto_approve,
304 approval_mode,
305 configured_sandbox_mode,
306 } = op
307 else {
308 return None;
309 };
310 Some(LiveRuntimeAuthority::from_fields(
311 *mode,
312 *allow_shell,
313 *trust_mode,
314 *auto_approve,
315 *approval_mode,
316 configured_sandbox_mode.clone(),
317 ))
318 }
319
320 fn publish_runtime_authority(&self, authority: LiveRuntimeAuthority) {
321 let mut state = self
322 .live_runtime_authority
323 .lock()
324 .unwrap_or_else(std::sync::PoisonError::into_inner);
325 state.revision = state.revision.wrapping_add(1).max(1);
326 state.authority = authority;
327 }
328
329 pub(crate) fn publish_turn_authority(
330 &self,
331 mode: AppMode,
332 allow_shell: bool,
333 trust_mode: bool,
334 auto_approve: bool,
335 approval_mode: ApprovalMode,
336 configured_sandbox_mode: Option<String>,
337 ) {
338 self.publish_runtime_authority(LiveRuntimeAuthority::from_fields(
339 mode,
340 allow_shell,
341 trust_mode,
342 auto_approve,
343 approval_mode,
344 configured_sandbox_mode,
345 ));
346 }
347
348 /// Exact live permission authority for runtime approval and elevation
349 /// gates. This is the same typed state the active engine turn drains.
350 #[must_use]
351 pub(crate) fn runtime_permission_authority(&self) -> RuntimePermissionAuthority {
352 self.live_runtime_authority
353 .lock()
354 .unwrap_or_else(std::sync::PoisonError::into_inner)
355 .authority
356 .permission_snapshot()
357 }
358
359 /// Reserve capacity for a runtime steer before it mutates durable state.
360 /// The owned permit lets the caller persist and dispatch synchronously,
361 /// without a cancellation point between those two operations.
362 pub(crate) async fn reserve_steer(&self) -> Result<SteerPermit> {
363 let permit = self.tx_steer.clone().reserve_owned().await?;
364 let turn_id = self
365 .turn_controls
366 .lock()
367 .unwrap_or_else(std::sync::PoisonError::into_inner)
368 .target()
369 .map(|control| control.id);
370 Ok(SteerPermit { permit, turn_id })
371 }
372
373 /// Cancel the current request (user-initiated path — keeps the
374 /// public `cancel()` signature stable). Equivalent to
375 /// `cancel_with_reason(CancelReason::User)`.
376 pub fn cancel(&self) {
377 self.cancel_with_reason(CancelReason::User);
378 }
379
380 /// Cancel the current request and latch the reason so downstream
381 /// "request cancelled" error messages can name a cause.
382 pub fn cancel_with_reason(&self, reason: CancelReason) {
383 // Keep turn activation excluded until both the admitted control and
384 // the legacy shared token have been canceled.
385 let controls = self
386 .turn_controls
387 .lock()
388 .unwrap_or_else(std::sync::PoisonError::into_inner);
389 if let Some(control) = controls.target() {
390 *control
391 .reason
392 .lock()
393 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(reason);
394 control.cancel.cancel();
395 }
396 match self.cancel_reason.lock() {
397 Ok(mut slot) => *slot = Some(reason),
398 Err(poisoned) => *poisoned.into_inner() = Some(reason),
399 }
400 match self.cancel_token.lock() {
401 Ok(token) => token.cancel(),
402 Err(poisoned) => poisoned.into_inner().cancel(),
403 }
404 crate::retry_status::clear();
405 }
406
407 /// Check if a request is currently cancelled
408 #[must_use]
409 pub fn is_cancelled(&self) -> bool {
410 if let Some(control) = self
411 .turn_controls
412 .lock()
413 .unwrap_or_else(std::sync::PoisonError::into_inner)
414 .target()
415 {
416 return control.cancel.is_cancelled();
417 }
418 match self.cancel_token.lock() {
419 Ok(token) => token.is_cancelled(),
420 Err(poisoned) => poisoned.into_inner().is_cancelled(),
421 }
422 }
423
424 /// Pause or resume the current pausable command.
425 pub fn set_paused(&self, paused: bool) {
426 match self.shared_paused.lock() {
427 Ok(mut slot) => *slot = paused,
428 Err(poisoned) => *poisoned.into_inner() = paused,
429 }
430 }
431
432 /// Check whether the engine pause gate is set.
433 #[cfg(test)]
434 #[must_use]
435 pub fn is_paused(&self) -> bool {
436 match self.shared_paused.lock() {
437 Ok(slot) => *slot,
438 Err(poisoned) => *poisoned.into_inner(),
439 }
440 }
441
442 /// Approve a pending tool call
443 pub async fn approve_tool_call(&self, id: impl Into<String>) -> Result<()> {
444 self.tx_approval
445 .send(ApprovalDecision::Approved { id: id.into() })
446 .await?;
447 Ok(())
448 }
449
450 /// Deny a pending tool call
451 pub async fn deny_tool_call(&self, id: impl Into<String>) -> Result<()> {
452 self.tx_approval
453 .send(ApprovalDecision::Denied { id: id.into() })
454 .await?;
455 Ok(())
456 }
457
458 /// Deny a pending tool call because its interactive approval card
459 /// expired (#6101). Kept distinct from [`Self::deny_tool_call`] so the
460 /// receipt records a timeout instead of an operator denial.
461 pub async fn deny_tool_call_timed_out(&self, id: impl Into<String>) -> Result<()> {
462 self.tx_approval
463 .send(ApprovalDecision::TimedOut { id: id.into() })
464 .await?;
465 Ok(())
466 }
467
468 /// Retry a tool call with an elevated sandbox policy.
469 pub async fn retry_tool_with_policy(
470 &self,
471 id: impl Into<String>,
472 policy: crate::sandbox::SandboxPolicy,
473 ) -> Result<()> {
474 self.tx_approval
475 .send(ApprovalDecision::RetryWithPolicy {
476 id: id.into(),
477 policy,
478 })
479 .await?;
480 Ok(())
481 }
482
483 /// Submit a response for request_user_input.
484 pub async fn submit_user_input(
485 &self,
486 id: impl Into<String>,
487 response: UserInputResponse,
488 ) -> Result<()> {
489 self.tx_user_input
490 .send(UserInputDecision::Submitted {
491 id: id.into(),
492 response,
493 })
494 .await?;
495 Ok(())
496 }
497
498 /// Cancel a request_user_input prompt.
499 pub async fn cancel_user_input(&self, id: impl Into<String>) -> Result<()> {
500 self.tx_user_input
501 .send(UserInputDecision::Cancelled { id: id.into() })
502 .await?;
503 Ok(())
504 }
505
506 /// Steer an in-flight turn with additional user input.
507 pub async fn steer(&self, content: impl Into<String>) -> Result<()> {
508 self.reserve_steer().await?.send(content.into());
509 Ok(())
510 }
511
512 /// Request the live context-window budget for this session's route.
513 /// `None` means the route cannot express a bounded window (e.g. an
514 /// unknown model with no catalog or configured limits) — callers should
515 /// surface "unavailable" rather than inventing a number.
516 pub async fn get_context_budget(
517 &self,
518 ) -> Result<Option<crate::core::ops::SessionContextBudget>> {
519 let (tx, rx) = tokio::sync::oneshot::channel();
520 let tx = std::sync::Arc::new(std::sync::Mutex::new(Some(tx)));
521 self.send(Op::GetContextBudget { tx }).await?;
522 rx.await
523 .map_err(|_| anyhow::anyhow!("Engine dropped context budget oneshot"))
524 }
525
526 /// Request a snapshot of the current session state.
527 /// Returns the snapshot directly via a oneshot channel, avoiding
528 /// competition with the SSE event stream on the mpsc receiver.
529 pub async fn get_session_snapshot(&self) -> Result<crate::core::ops::SessionSnapshot> {
530 let (tx, rx) = tokio::sync::oneshot::channel();
531 let tx = std::sync::Arc::new(std::sync::Mutex::new(Some(tx)));
532 self.send(Op::GetSessionSnapshot { tx }).await?;
533 rx.await
534 .map_err(|_| anyhow::anyhow!("Engine dropped session snapshot oneshot"))
535 }
536
537 /// Query after the active turn settles, without competing with events.
538 /// The caller must keep draining events and bound this future: an active
539 /// turn can be awaiting provider/tool work or a full event channel.
540 pub(crate) async fn get_subagent_settlement(
541 &self,
542 ) -> Result<crate::core::ops::SubAgentSettlement> {
543 let (tx, rx) = tokio::sync::oneshot::channel();
544 let tx = Arc::new(StdMutex::new(Some(tx)));
545 self.send(Op::GetSubAgentSettlement { tx }).await?;
546 rx.await
547 .map_err(|_| anyhow::anyhow!("Engine dropped child settlement receipt"))
548 }
549
550 /// Request active provider request concurrency state.
551 pub async fn get_provider_runtime_status(
552 &self,
553 ) -> Result<crate::core::ops::ProviderRuntimeStatus> {
554 let (tx, rx) = tokio::sync::oneshot::channel();
555 let tx = std::sync::Arc::new(std::sync::Mutex::new(Some(tx)));
556 self.send(Op::GetProviderRuntimeStatus { tx }).await?;
557 rx.await
558 .map_err(|_| anyhow::anyhow!("Engine dropped provider runtime status oneshot"))
559 }
560
561 /// Run the bounded initial connection pass on the engine-owned MCP pool.
562 ///
563 /// The returned manager snapshot and every later tool call therefore see
564 /// the same connections and catalog generation. Unlike `reload_mcp`, this
565 /// does not force a config re-read or drop ready transports. Optional
566 /// servers are connected in the background at engine spawn; this waits
567 /// only if the caller explicitly asked for the settled receipt.
568 pub async fn bootstrap_mcp(&self) -> Result<crate::core::ops::McpManagerUpdate> {
569 let (tx, rx) = tokio::sync::oneshot::channel();
570 let tx = std::sync::Arc::new(std::sync::Mutex::new(Some(tx)));
571 self.send(Op::BootstrapMcp { tx }).await?;
572 rx.await
573 .map_err(|_| anyhow::anyhow!("Engine dropped MCP bootstrap oneshot"))?
574 .map_err(anyhow::Error::msg)
575 }
576
577 /// Retry one failed server through the existing engine-owned pool.
578 pub async fn retry_mcp_server(
579 &self,
580 name: impl Into<String>,
581 ) -> Result<crate::core::ops::McpManagerUpdate> {
582 let (tx, rx) = tokio::sync::oneshot::channel();
583 let tx = std::sync::Arc::new(std::sync::Mutex::new(Some(tx)));
584 self.send(Op::RetryMcpServer {
585 name: name.into(),
586 tx,
587 })
588 .await?;
589 rx.await
590 .map_err(|_| anyhow::anyhow!("Engine dropped MCP retry oneshot"))?
591 .map_err(anyhow::Error::msg)
592 }
593
594 /// Force the engine-owned MCP pool to reload and reconnect, returning a
595 /// snapshot from the exact live pool that supplies the next model turn.
596 pub async fn reload_mcp(
597 &self,
598 config_path: std::path::PathBuf,
599 ) -> Result<crate::core::ops::McpManagerUpdate> {
600 let (tx, rx) = tokio::sync::oneshot::channel();
601 let tx = std::sync::Arc::new(std::sync::Mutex::new(Some(tx)));
602 self.send(Op::ReloadMcp { config_path, tx }).await?;
603 rx.await
604 .map_err(|_| anyhow::anyhow!("Engine dropped MCP reload oneshot"))?
605 .map_err(anyhow::Error::msg)
606 }
607 }
608
608 lines RUST