vv_agent/app_server/
request_serialization.rs

1use std::collections::{HashMap, VecDeque};
2use std::future::Future;
3use std::sync::Arc;
4
5use serde_json::Value;
6use tokio::sync::{Mutex, Notify};
7
8#[derive(Debug, Clone, PartialEq, Eq, Hash)]
9pub struct RequestSerializationQueueKey(String);
10
11impl RequestSerializationQueueKey {
12    pub fn thread(thread_id: impl Into<String>) -> Self {
13        Self(format!("thread:{}", thread_id.into()))
14    }
15
16    pub fn global(name: impl Into<String>) -> Self {
17        Self(format!("global:{}", name.into()))
18    }
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum RequestSerializationAccess {
23    Shared,
24    Exclusive,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct RequestSerializationScope {
29    key: RequestSerializationQueueKey,
30    access: RequestSerializationAccess,
31}
32
33impl RequestSerializationScope {
34    pub fn shared_thread(thread_id: impl Into<String>) -> Self {
35        Self {
36            key: RequestSerializationQueueKey::thread(thread_id),
37            access: RequestSerializationAccess::Shared,
38        }
39    }
40
41    pub fn exclusive_thread(thread_id: impl Into<String>) -> Self {
42        Self {
43            key: RequestSerializationQueueKey::thread(thread_id),
44            access: RequestSerializationAccess::Exclusive,
45        }
46    }
47
48    pub fn shared_global(name: impl Into<String>) -> Self {
49        Self {
50            key: RequestSerializationQueueKey::global(name),
51            access: RequestSerializationAccess::Shared,
52        }
53    }
54
55    pub fn exclusive_global(name: impl Into<String>) -> Self {
56        Self {
57            key: RequestSerializationQueueKey::global(name),
58            access: RequestSerializationAccess::Exclusive,
59        }
60    }
61
62    pub fn for_method(method: &str, params: Option<&Value>) -> Option<Self> {
63        match method {
64            "thread/start" => Some(Self::exclusive_global("thread")),
65            "thread/resume" | "thread/read" => thread_id(params).map(Self::shared_thread),
66            "thread/archive" | "turn/start" | "turn/interrupt" | "turn/steer"
67            | "approval/resolve" => thread_id(params).map(Self::exclusive_thread),
68            "thread/list" => Some(Self::shared_global("thread/list")),
69            "model/list" => Some(Self::shared_global("model")),
70            "schema/export" => Some(Self::shared_global("schema")),
71            _ => None,
72        }
73    }
74
75    pub fn key(&self) -> &RequestSerializationQueueKey {
76        &self.key
77    }
78
79    pub fn access(&self) -> RequestSerializationAccess {
80        self.access
81    }
82}
83
84#[derive(Clone, Default)]
85pub struct RequestSerializationQueue {
86    inner: Arc<Mutex<HashMap<RequestSerializationQueueKey, QueueState>>>,
87}
88
89impl RequestSerializationQueue {
90    pub async fn run<F, T>(&self, scope: RequestSerializationScope, future: F) -> T
91    where
92        F: Future<Output = T>,
93    {
94        self.acquire(scope.clone()).await;
95        let result = future.await;
96        self.release(scope).await;
97        result
98    }
99
100    async fn acquire(&self, scope: RequestSerializationScope) {
101        let notify = Arc::new(Notify::new());
102        let mut queues = self.inner.lock().await;
103        let state = queues.entry(scope.key.clone()).or_default();
104
105        if state.can_start_now(scope.access) {
106            state.start(scope.access);
107            return;
108        }
109
110        state.waiters.push_back(QueuedRequest {
111            access: scope.access,
112            notify: notify.clone(),
113        });
114        drop(queues);
115        notify.notified().await;
116    }
117
118    async fn release(&self, scope: RequestSerializationScope) {
119        let mut queues = self.inner.lock().await;
120        let Some(state) = queues.get_mut(&scope.key) else {
121            return;
122        };
123        state.finish(scope.access);
124        state.wake_next();
125        if state.is_empty() {
126            queues.remove(&scope.key);
127        }
128    }
129}
130
131#[derive(Default)]
132struct QueueState {
133    active_shared: usize,
134    active_exclusive: bool,
135    waiters: VecDeque<QueuedRequest>,
136}
137
138impl QueueState {
139    fn can_start_now(&self, access: RequestSerializationAccess) -> bool {
140        self.waiters.is_empty()
141            && match access {
142                RequestSerializationAccess::Shared => !self.active_exclusive,
143                RequestSerializationAccess::Exclusive => {
144                    !self.active_exclusive && self.active_shared == 0
145                }
146            }
147    }
148
149    fn start(&mut self, access: RequestSerializationAccess) {
150        match access {
151            RequestSerializationAccess::Shared => self.active_shared += 1,
152            RequestSerializationAccess::Exclusive => self.active_exclusive = true,
153        }
154    }
155
156    fn finish(&mut self, access: RequestSerializationAccess) {
157        match access {
158            RequestSerializationAccess::Shared => {
159                self.active_shared = self.active_shared.saturating_sub(1);
160            }
161            RequestSerializationAccess::Exclusive => {
162                self.active_exclusive = false;
163            }
164        }
165    }
166
167    fn wake_next(&mut self) {
168        if self.active_exclusive || self.active_shared > 0 {
169            return;
170        }
171
172        match self.waiters.front().map(|waiter| waiter.access) {
173            Some(RequestSerializationAccess::Exclusive) => {
174                if let Some(waiter) = self.waiters.pop_front() {
175                    self.active_exclusive = true;
176                    waiter.notify.notify_one();
177                }
178            }
179            Some(RequestSerializationAccess::Shared) => {
180                while self
181                    .waiters
182                    .front()
183                    .is_some_and(|waiter| waiter.access == RequestSerializationAccess::Shared)
184                {
185                    if let Some(waiter) = self.waiters.pop_front() {
186                        self.active_shared += 1;
187                        waiter.notify.notify_one();
188                    }
189                }
190            }
191            None => {}
192        }
193    }
194
195    fn is_empty(&self) -> bool {
196        self.active_shared == 0 && !self.active_exclusive && self.waiters.is_empty()
197    }
198}
199
200struct QueuedRequest {
201    access: RequestSerializationAccess,
202    notify: Arc<Notify>,
203}
204
205fn thread_id(params: Option<&Value>) -> Option<String> {
206    params
207        .and_then(|params| params.get("threadId"))
208        .and_then(Value::as_str)
209        .map(str::to_string)
210}
vv_agent/app_server/request_serialization.rs

vv_agent/app_server/
request_serialization.rs
