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" | "thread/unsubscribe" | "turn/start" | "turn/resume"
67            | "turn/interrupt" | "turn/steer" | "turn/followUp" | "approval/resolve" => {
68                thread_id(params).map(Self::exclusive_thread)
69            }
70            "thread/list" => Some(Self::shared_global("thread/list")),
71            "model/list" => Some(Self::shared_global("model")),
72            "schema/export" => Some(Self::shared_global("schema")),
73            _ => None,
74        }
75    }
76
77    pub fn key(&self) -> &RequestSerializationQueueKey {
78        &self.key
79    }
80
81    pub fn access(&self) -> RequestSerializationAccess {
82        self.access
83    }
84}
85
86#[derive(Clone, Default)]
87pub struct RequestSerializationQueue {
88    inner: Arc<Mutex<HashMap<RequestSerializationQueueKey, QueueState>>>,
89}
90
91impl RequestSerializationQueue {
92    pub async fn run<F, T>(&self, scope: RequestSerializationScope, future: F) -> T
93    where
94        F: Future<Output = T>,
95    {
96        self.acquire(scope.clone()).await;
97        let result = future.await;
98        self.release(scope).await;
99        result
100    }
101
102    async fn acquire(&self, scope: RequestSerializationScope) {
103        let notify = Arc::new(Notify::new());
104        let mut queues = self.inner.lock().await;
105        let state = queues.entry(scope.key.clone()).or_default();
106
107        if state.can_start_now(scope.access) {
108            state.start(scope.access);
109            return;
110        }
111
112        state.waiters.push_back(QueuedRequest {
113            access: scope.access,
114            notify: notify.clone(),
115        });
116        drop(queues);
117        notify.notified().await;
118    }
119
120    async fn release(&self, scope: RequestSerializationScope) {
121        let mut queues = self.inner.lock().await;
122        let Some(state) = queues.get_mut(&scope.key) else {
123            return;
124        };
125        state.finish(scope.access);
126        state.wake_next();
127        if state.is_empty() {
128            queues.remove(&scope.key);
129        }
130    }
131}
132
133#[derive(Default)]
134struct QueueState {
135    active_shared: usize,
136    active_exclusive: bool,
137    waiters: VecDeque<QueuedRequest>,
138}
139
140impl QueueState {
141    fn can_start_now(&self, access: RequestSerializationAccess) -> bool {
142        self.waiters.is_empty()
143            && match access {
144                RequestSerializationAccess::Shared => !self.active_exclusive,
145                RequestSerializationAccess::Exclusive => {
146                    !self.active_exclusive && self.active_shared == 0
147                }
148            }
149    }
150
151    fn start(&mut self, access: RequestSerializationAccess) {
152        match access {
153            RequestSerializationAccess::Shared => self.active_shared += 1,
154            RequestSerializationAccess::Exclusive => self.active_exclusive = true,
155        }
156    }
157
158    fn finish(&mut self, access: RequestSerializationAccess) {
159        match access {
160            RequestSerializationAccess::Shared => {
161                self.active_shared = self.active_shared.saturating_sub(1);
162            }
163            RequestSerializationAccess::Exclusive => {
164                self.active_exclusive = false;
165            }
166        }
167    }
168
169    fn wake_next(&mut self) {
170        if self.active_exclusive || self.active_shared > 0 {
171            return;
172        }
173
174        match self.waiters.front().map(|waiter| waiter.access) {
175            Some(RequestSerializationAccess::Exclusive) => {
176                if let Some(waiter) = self.waiters.pop_front() {
177                    self.active_exclusive = true;
178                    waiter.notify.notify_one();
179                }
180            }
181            Some(RequestSerializationAccess::Shared) => {
182                while self
183                    .waiters
184                    .front()
185                    .is_some_and(|waiter| waiter.access == RequestSerializationAccess::Shared)
186                {
187                    if let Some(waiter) = self.waiters.pop_front() {
188                        self.active_shared += 1;
189                        waiter.notify.notify_one();
190                    }
191                }
192            }
193            None => {}
194        }
195    }
196
197    fn is_empty(&self) -> bool {
198        self.active_shared == 0 && !self.active_exclusive && self.waiters.is_empty()
199    }
200}
201
202struct QueuedRequest {
203    access: RequestSerializationAccess,
204    notify: Arc<Notify>,
205}
206
207fn thread_id(params: Option<&Value>) -> Option<String> {
208    params
209        .and_then(|params| params.get("threadId"))
210        .and_then(Value::as_str)
211        .map(str::to_string)
212}
vv_agent/app_server/request_serialization.rs

vv_agent/app_server/
request_serialization.rs
