Skip to main content

oino_tools/bash/
detach.rs

1//! Detaching running `bash` calls into background tasks (Oino addition, runtime ยง8 C11).
2//!
3//! Every running shell call registers under its tool-call id. Core's background-task service
4//! calls [`BashDetachRegistry::detach`] with the task id it allocated; the call then returns at
5//! once with the output so far plus `\n\n[Detached to background task {taskId}: {reason}]`
6//! (`isError: false`) and the still-running process is handed to the caller.
7
8use std::collections::HashMap;
9use std::sync::{Arc, Mutex, PoisonError};
10
11use tokio::sync::{mpsc, oneshot};
12
13use crate::env::SpawnedProcess;
14
15/// A request sent to a running call.
16#[derive(Debug)]
17pub(crate) struct DetachRequest {
18    pub(crate) task_id: String,
19    pub(crate) reason: String,
20    pub(crate) reply: oneshot::Sender<DetachedProcess>,
21}
22
23/// What a running call is asked for.
24#[derive(Debug)]
25pub(crate) enum CallMessage {
26    /// Hand the process over.
27    Detach(DetachRequest),
28    /// Report the process id.
29    Pid(oneshot::Sender<Option<u32>>),
30}
31
32/// A process taken over from a detached `bash` call.
33#[derive(Debug)]
34pub struct DetachedProcess {
35    /// The running process; further output arrives on `process.output`.
36    pub process: SpawnedProcess,
37    /// The text the tool call returned before the detach marker.
38    pub output_so_far: String,
39    /// Temp file holding the complete output so far (when output was truncated).
40    pub full_output_path: Option<String>,
41}
42
43/// Running shell calls by tool-call id. Cheap to clone; clones share the registry.
44#[derive(Clone, Debug, Default)]
45pub struct BashDetachRegistry {
46    calls: Arc<Mutex<HashMap<String, mpsc::UnboundedSender<CallMessage>>>>,
47}
48
49/// Removes the registration when the call ends.
50pub(crate) struct Registration {
51    registry: BashDetachRegistry,
52    id: String,
53}
54
55impl Drop for Registration {
56    fn drop(&mut self) {
57        self.registry.lock().remove(&self.id);
58    }
59}
60
61impl BashDetachRegistry {
62    fn lock(
63        &self,
64    ) -> std::sync::MutexGuard<'_, HashMap<String, mpsc::UnboundedSender<CallMessage>>> {
65        self.calls.lock().unwrap_or_else(PoisonError::into_inner)
66    }
67
68    pub(crate) fn register(
69        &self,
70        tool_call_id: &str,
71    ) -> (Registration, mpsc::UnboundedReceiver<CallMessage>) {
72        let (tx, rx) = mpsc::unbounded_channel();
73        self.lock().insert(tool_call_id.to_owned(), tx);
74        (
75            Registration {
76                registry: self.clone(),
77                id: tool_call_id.to_owned(),
78            },
79            rx,
80        )
81    }
82
83    /// Whether a shell call with this id is running.
84    #[must_use]
85    pub fn is_running(&self, tool_call_id: &str) -> bool {
86        self.lock().contains_key(tool_call_id)
87    }
88
89    /// Detaches the running call `tool_call_id` into background task `task_id`. `None` when no
90    /// such call runs or it finished before the request was seen.
91    pub async fn detach(
92        &self,
93        tool_call_id: &str,
94        task_id: impl Into<String>,
95        reason: impl Into<String>,
96    ) -> Option<DetachedProcess> {
97        let sender = self.lock().get(tool_call_id).cloned()?;
98        let (reply, response) = oneshot::channel();
99        sender
100            .send(CallMessage::Detach(DetachRequest {
101                task_id: task_id.into(),
102                reason: reason.into(),
103                reply,
104            }))
105            .ok()?;
106        response.await.ok()
107    }
108
109    /// Process id of the running call `tool_call_id` (`None` when it is not running or the
110    /// process id is unknown).
111    pub async fn pid(&self, tool_call_id: &str) -> Option<u32> {
112        let sender = self.lock().get(tool_call_id).cloned()?;
113        let (reply, response) = oneshot::channel();
114        sender.send(CallMessage::Pid(reply)).ok()?;
115        response.await.ok().flatten()
116    }
117
118    /// Tool-call ids of the running calls.
119    #[must_use]
120    pub fn running(&self) -> Vec<String> {
121        self.lock().keys().cloned().collect()
122    }
123}