Skip to main content

vv_agent/interactive/
runtime_support.rs

1use super::*;
2
3impl InteractiveSessionSubscription {
4    pub(super) fn new(
5        receiver: broadcast::Receiver<InteractiveSessionEvent>,
6        closed: Arc<AtomicBool>,
7    ) -> Self {
8        Self { receiver, closed }
9    }
10
11    pub async fn recv(&mut self) -> Result<InteractiveSessionEvent, InteractiveSessionError> {
12        if self.closed.load(Ordering::SeqCst) && self.receiver.is_empty() {
13            return Err(InteractiveSessionError::EventStreamClosed);
14        }
15        self.receiver.recv().await.map_err(|error| match error {
16            broadcast::error::RecvError::Closed => InteractiveSessionError::EventStreamClosed,
17            broadcast::error::RecvError::Lagged(missed) => {
18                InteractiveSessionError::EventGap { missed }
19            }
20        })
21    }
22
23    pub fn try_recv(&mut self) -> Result<InteractiveSessionEvent, InteractiveSessionError> {
24        self.receiver.try_recv().map_err(|error| match error {
25            broadcast::error::TryRecvError::Closed => InteractiveSessionError::EventStreamClosed,
26            broadcast::error::TryRecvError::Empty if self.closed.load(Ordering::SeqCst) => {
27                InteractiveSessionError::EventStreamClosed
28            }
29            broadcast::error::TryRecvError::Empty => InteractiveSessionError::EventStreamEmpty,
30            broadcast::error::TryRecvError::Lagged(missed) => {
31                InteractiveSessionError::EventGap { missed }
32            }
33        })
34    }
35}
36
37pub(super) struct InteractiveRunHandleController {
38    pub(super) inner: std::sync::Weak<InteractiveSessionInner>,
39    pub(super) session_id: String,
40}
41
42impl InteractiveRunHandleController {
43    fn session(&self) -> Result<InteractiveSession, String> {
44        self.inner
45            .upgrade()
46            .map(|inner| InteractiveSession { inner })
47            .ok_or_else(|| format!("interactive session `{}` is closed", self.session_id))
48    }
49}
50
51impl RunHandleController for InteractiveRunHandleController {
52    fn steer(&self, message: String) -> Result<(), String> {
53        self.session()
54            .and_then(|session| session.steer(message).map_err(|error| error.to_string()))
55    }
56
57    fn follow_up(&self, message: String) -> Result<(), String> {
58        self.session().and_then(|session| {
59            session
60                .follow_up(message)
61                .map_err(|error| error.to_string())
62        })
63    }
64}
65
66pub(super) struct RunEventForwarder {
67    task: Option<tokio::task::JoinHandle<()>>,
68}
69
70impl RunEventForwarder {
71    pub(super) fn new(task: tokio::task::JoinHandle<()>) -> Self {
72        Self { task: Some(task) }
73    }
74
75    pub(super) async fn finish(&mut self) -> Result<(), tokio::task::JoinError> {
76        let result = match self.task.as_mut() {
77            Some(task) => task.await,
78            None => Ok(()),
79        };
80        self.task.take();
81        result
82    }
83}
84
85impl Drop for RunEventForwarder {
86    fn drop(&mut self) {
87        if let Some(task) = self.task.take() {
88            task.abort();
89        }
90    }
91}
92
93pub(super) struct RunLifecycleGuard {
94    inner: Arc<InteractiveSessionInner>,
95    armed: bool,
96}
97
98impl RunLifecycleGuard {
99    pub(super) fn begin(
100        inner: Arc<InteractiveSessionInner>,
101    ) -> Result<Self, InteractiveSessionError> {
102        {
103            let mut state = lock_unpoisoned(&inner.state);
104            if state.closed {
105                return Err(InteractiveSessionError::Closed {
106                    session_id: inner.session_id.clone(),
107                });
108            }
109            if state.running {
110                return Err(InteractiveSessionError::AlreadyRunning {
111                    session_id: inner.session_id.clone(),
112                });
113            }
114            state.running = true;
115            state.active_cancellation_token = Some(
116                inner
117                    .run_config
118                    .cancellation_token
119                    .as_ref()
120                    .map(CancellationToken::child)
121                    .unwrap_or_default(),
122            );
123        }
124        Ok(Self { inner, armed: true })
125    }
126
127    pub(super) fn finish(&mut self) {
128        self.reset(false);
129        self.armed = false;
130    }
131
132    fn reset(&self, aborted: bool) {
133        let (active, controller) = {
134            let mut state = lock_unpoisoned(&self.inner.state);
135            state.running = false;
136            state.active_cancellation_token = None;
137            (
138                state.active_handle.take(),
139                state.active_handle_controller.take(),
140            )
141        };
142        if let Some(handle) = active {
143            if let Some(controller) = controller {
144                handle.detach_controller(controller);
145            }
146            handle.cancel_with_reason("interactive session run aborted");
147            let _ = self
148                .inner
149                .events
150                .send(InteractiveSessionEvent::ActiveHandleChanged {
151                    session_id: self.inner.session_id.clone(),
152                    active: false,
153                });
154        }
155        if aborted && !self.inner.closed.load(Ordering::SeqCst) {
156            let _ = self.inner.events.send(InteractiveSessionEvent::RunAborted {
157                session_id: self.inner.session_id.clone(),
158            });
159        }
160    }
161}
162
163impl Drop for RunLifecycleGuard {
164    fn drop(&mut self) {
165        if self.armed {
166            self.reset(true);
167        }
168    }
169}
170
171pub(super) struct SteeringRuntimeHook {
172    pub(super) queue: SteeringQueue,
173    pub(super) session_id: String,
174    pub(super) events: broadcast::Sender<InteractiveSessionEvent>,
175    pub(super) inner: std::sync::Weak<InteractiveSessionInner>,
176}
177
178impl RuntimeHook for SteeringRuntimeHook {
179    fn before_llm(&self, event: BeforeLlmEvent<'_>) -> Option<BeforeLlmPatch> {
180        let queued = {
181            let mut queue = lock_queue(&self.queue);
182            queue.drain(..).collect::<Vec<_>>()
183        };
184        if queued.is_empty() {
185            return None;
186        }
187        let mut messages = event.messages.to_vec();
188        for prompt in queued {
189            messages.push(Message::user(prompt.clone()));
190            let _ = self.events.send(InteractiveSessionEvent::SteerDequeued {
191                session_id: self.session_id.clone(),
192                prompt,
193                cycle_index: Some(event.cycle_index),
194            });
195        }
196        Some(BeforeLlmPatch {
197            messages: Some(messages),
198            tool_schemas: None,
199        })
200    }
201
202    fn before_tool_call(&self, event: BeforeToolCallEvent<'_>) -> Option<BeforeToolCallPatch> {
203        if lock_queue(&self.queue).is_empty() {
204            return None;
205        }
206        let mut result = ToolExecutionResult::error(
207            event.call.id.clone(),
208            serde_json::json!({
209                "ok": false,
210                "error": "Tool skipped due to queued steering message.",
211                "error_code": "skipped_due_to_steering",
212            })
213            .to_string(),
214        );
215        result.error_code = Some("skipped_due_to_steering".to_string());
216        Some(result.into())
217    }
218
219    fn after_tool_call(&self, event: AfterToolCallEvent<'_>) -> Option<ToolExecutionResult> {
220        if let Some(inner) = self.inner.upgrade() {
221            InteractiveSession { inner }
222                .sync_background_command_result(&event.call.name, event.result);
223        }
224        None
225    }
226}
227
228pub(super) fn build_background_command_notification(payload: &Value) -> String {
229    let status = payload
230        .get("status")
231        .and_then(Value::as_str)
232        .unwrap_or_default()
233        .trim()
234        .to_ascii_lowercase();
235    let status_text = match status.as_str() {
236        "completed" => "completed",
237        "failed" => "failed",
238        "timeout" => "timed out",
239        "" => "updated",
240        other => other,
241    };
242    let background_session_id = payload
243        .get("session_id")
244        .and_then(Value::as_str)
245        .unwrap_or_default()
246        .trim();
247    let command = payload
248        .get("command")
249        .and_then(Value::as_str)
250        .unwrap_or_default()
251        .trim();
252    let output = payload
253        .get("output")
254        .and_then(Value::as_str)
255        .unwrap_or_default()
256        .trim();
257    let mut summary = if output.is_empty() {
258        payload
259            .get("exit_code")
260            .map(|exit_code| format!("exit_code={exit_code}"))
261            .unwrap_or_default()
262    } else {
263        output.to_string()
264    };
265    if summary.chars().count() > 500 {
266        summary = format!(
267            "{}...",
268            summary.chars().take(497).collect::<String>().trim_end()
269        );
270    }
271
272    let mut lines = vec![format!(
273        "System notification: background command {background_session_id} {status_text}."
274    )];
275    if !command.is_empty() {
276        lines.push(format!("Command: {command}"));
277    }
278    if !summary.is_empty() {
279        lines.push(format!("Summary: {summary}"));
280    }
281    lines.join("\n")
282}
283
284pub(super) fn resolve_session(
285    requested_id: Option<String>,
286    session: Option<Arc<dyn Session>>,
287) -> Result<(String, Arc<dyn Session>), InteractiveSessionError> {
288    let requested_id = requested_id
289        .map(|value| value.trim().to_string())
290        .filter(|value| !value.is_empty());
291    match session {
292        Some(session) => {
293            let actual = session.session_id().trim().to_string();
294            if actual.is_empty() {
295                return Err(InteractiveSessionError::EmptySessionId);
296            }
297            if let Some(requested) = requested_id.as_ref() {
298                if requested != &actual {
299                    return Err(InteractiveSessionError::SessionIdMismatch {
300                        requested: requested.clone(),
301                        actual,
302                    });
303                }
304            }
305            Ok((requested_id.unwrap_or(actual), session))
306        }
307        None => {
308            let session_id = requested_id.unwrap_or_else(new_session_id);
309            let session: Arc<dyn Session> = Arc::new(MemorySession::new(session_id.clone()));
310            Ok((session_id, session))
311        }
312    }
313}
314
315pub(super) fn select_session_source(
316    option_session: Option<Arc<dyn Session>>,
317    run_config_session: Option<Arc<dyn Session>>,
318) -> Result<Option<Arc<dyn Session>>, InteractiveSessionError> {
319    match (option_session, run_config_session) {
320        (Some(option), Some(configured)) => {
321            let option_id = option.session_id().trim();
322            let configured_id = configured.session_id().trim();
323            if option_id != configured_id {
324                return Err(InteractiveSessionError::SessionIdMismatch {
325                    requested: option_id.to_string(),
326                    actual: configured_id.to_string(),
327                });
328            }
329            Ok(Some(option))
330        }
331        (Some(session), None) | (None, Some(session)) => Ok(Some(session)),
332        (None, None) => Ok(None),
333    }
334}
335
336pub(super) fn normalized_prompt(
337    prompt: impl Into<String>,
338    operation: &'static str,
339) -> Result<String, InteractiveSessionError> {
340    let prompt = prompt.into().trim().to_string();
341    if prompt.is_empty() {
342        return Err(InteractiveSessionError::EmptyPrompt { operation });
343    }
344    Ok(prompt)
345}
346
347pub(super) fn new_session_id() -> String {
348    let id = uuid::Uuid::new_v4().simple().to_string();
349    format!("session_{}", &id[..12])
350}
351
352pub(super) fn lock_queue(queue: &SteeringQueue) -> MutexGuard<'_, VecDeque<String>> {
353    lock_unpoisoned(queue)
354}
355
356pub(super) fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
357    mutex.lock().unwrap_or_else(|error| error.into_inner())
358}