vv_agent/app_server/
thread_state.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3
4use tokio::sync::Mutex;
5
6use crate::app_server::protocol::AppTurn;
7use crate::app_server::transport::ConnectionId;
8use crate::RunHandle;
9
10#[derive(Clone, Default)]
11pub struct ThreadStateManager {
12    inner: Arc<Mutex<ThreadStateInner>>,
13}
14
15#[derive(Default)]
16struct ThreadStateInner {
17    subscribers: HashMap<String, HashSet<ConnectionId>>,
18    active_turns: HashMap<String, ActiveTurn>,
19    pending_approvals: HashMap<String, PendingApproval>,
20}
21
22#[derive(Clone)]
23pub struct ActiveTurn {
24    pub turn: AppTurn,
25    pub handle: RunHandle,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct PendingApproval {
30    pub turn_id: String,
31    pub request_id: String,
32}
33
34impl ThreadStateManager {
35    pub async fn subscribe(&self, thread_id: impl Into<String>, connection_id: ConnectionId) {
36        self.inner
37            .lock()
38            .await
39            .subscribers
40            .entry(thread_id.into())
41            .or_default()
42            .insert(connection_id);
43    }
44
45    pub async fn subscribers(&self, thread_id: &str) -> Vec<ConnectionId> {
46        self.inner
47            .lock()
48            .await
49            .subscribers
50            .get(thread_id)
51            .map(|subscribers| subscribers.iter().copied().collect())
52            .unwrap_or_default()
53    }
54
55    pub async fn set_active_turn(&self, thread_id: impl Into<String>, active_turn: ActiveTurn) {
56        self.inner
57            .lock()
58            .await
59            .active_turns
60            .insert(thread_id.into(), active_turn);
61    }
62
63    pub async fn active_turn(&self, thread_id: &str) -> Option<ActiveTurn> {
64        self.inner.lock().await.active_turns.get(thread_id).cloned()
65    }
66
67    pub async fn clear_active_turn(&self, thread_id: &str, turn_id: &str) {
68        let mut inner = self.inner.lock().await;
69        if inner
70            .active_turns
71            .get(thread_id)
72            .is_some_and(|active| active.turn.id == turn_id)
73        {
74            inner.active_turns.remove(thread_id);
75        }
76    }
77
78    pub async fn cancel_turn(&self, thread_id: &str, turn_id: &str) -> bool {
79        let active = self.active_turn(thread_id).await;
80        let Some(active) = active else {
81            return false;
82        };
83        if active.turn.id != turn_id {
84            return false;
85        }
86        active.handle.cancel();
87        true
88    }
89
90    pub async fn set_pending_approval(
91        &self,
92        thread_id: impl Into<String>,
93        turn_id: impl Into<String>,
94        request_id: impl Into<String>,
95    ) {
96        self.inner.lock().await.pending_approvals.insert(
97            thread_id.into(),
98            PendingApproval {
99                turn_id: turn_id.into(),
100                request_id: request_id.into(),
101            },
102        );
103    }
104
105    pub async fn pending_approval(&self, thread_id: &str) -> Option<PendingApproval> {
106        self.inner
107            .lock()
108            .await
109            .pending_approvals
110            .get(thread_id)
111            .cloned()
112    }
113
114    pub async fn clear_pending_approval(&self, thread_id: &str, request_id: &str) {
115        let mut inner = self.inner.lock().await;
116        if inner
117            .pending_approvals
118            .get(thread_id)
119            .is_some_and(|pending| pending.request_id == request_id)
120        {
121            inner.pending_approvals.remove(thread_id);
122        }
123    }
124}
vv_agent/app_server/thread_state.rs

vv_agent/app_server/
thread_state.rs
