vv_agent/app_server/
outgoing.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::atomic::{AtomicI64, Ordering};
3use std::sync::Arc;
4
5use serde::Serialize;
6use serde_json::Value;
7use tokio::sync::{mpsc, oneshot, Mutex};
8
9use crate::app_server::protocol::{
10    AppServerError, AppServerErrorCode, JsonRpcError, JsonRpcErrorBody, JsonRpcMessage,
11    JsonRpcNotification, JsonRpcResponse, RequestId, ServerNotification, ServerRequest,
12};
13use crate::app_server::transport::ConnectionId;
14
15pub type ServerRequestResult = Result<Value, JsonRpcErrorBody>;
16
17#[derive(Debug, Clone, PartialEq)]
18pub struct OutgoingEnvelope {
19    pub connection_id: ConnectionId,
20    pub message: JsonRpcMessage,
21}
22
23#[derive(Debug, Clone, Default)]
24struct ConnectionOutgoingState {
25    ready_for_notifications: bool,
26    opt_out_notification_methods: HashSet<String>,
27}
28
29#[derive(Clone)]
30pub struct OutgoingMessageSender {
31    tx: mpsc::Sender<OutgoingEnvelope>,
32    connections: Arc<Mutex<HashMap<ConnectionId, ConnectionOutgoingState>>>,
33    pending_server_requests: Arc<Mutex<HashMap<RequestId, oneshot::Sender<ServerRequestResult>>>>,
34    next_server_request_id: Arc<AtomicI64>,
35}
36
37impl OutgoingMessageSender {
38    pub fn channel(capacity: usize) -> (Self, mpsc::Receiver<OutgoingEnvelope>) {
39        let (tx, rx) = mpsc::channel(capacity);
40        (
41            Self {
42                tx,
43                connections: Arc::new(Mutex::new(HashMap::new())),
44                pending_server_requests: Arc::new(Mutex::new(HashMap::new())),
45                next_server_request_id: Arc::new(AtomicI64::new(1)),
46            },
47            rx,
48        )
49    }
50
51    pub async fn register_connection(&self, connection_id: ConnectionId) {
52        self.connections
53            .lock()
54            .await
55            .entry(connection_id)
56            .or_default();
57    }
58
59    pub async fn configure_connection(
60        &self,
61        connection_id: ConnectionId,
62        opt_out_notification_methods: HashSet<String>,
63    ) {
64        let mut connections = self.connections.lock().await;
65        let state = connections.entry(connection_id).or_default();
66        state.opt_out_notification_methods = opt_out_notification_methods;
67    }
68
69    pub async fn mark_ready_for_notifications(&self, connection_id: ConnectionId) {
70        self.connections
71            .lock()
72            .await
73            .entry(connection_id)
74            .or_default()
75            .ready_for_notifications = true;
76    }
77
78    pub async fn send_response(
79        &self,
80        connection_id: ConnectionId,
81        id: RequestId,
82        result: Value,
83    ) -> Result<(), AppServerError> {
84        self.send_message(
85            connection_id,
86            JsonRpcMessage::Response(JsonRpcResponse { id, result }),
87        )
88        .await
89    }
90
91    pub async fn send_error(
92        &self,
93        connection_id: ConnectionId,
94        id: RequestId,
95        error: AppServerError,
96    ) -> Result<(), AppServerError> {
97        self.send_message(
98            connection_id,
99            JsonRpcMessage::Error(error.into_json_rpc_error(id)),
100        )
101        .await
102    }
103
104    pub async fn send_notification(
105        &self,
106        connection_id: ConnectionId,
107        notification: ServerNotification,
108    ) -> Result<(), AppServerError> {
109        if !self
110            .should_send_notification(connection_id, &notification)
111            .await
112        {
113            return Ok(());
114        }
115        let message = server_notification_message(notification)?;
116        self.send_message(connection_id, message).await
117    }
118
119    pub async fn broadcast_initialized_client_notification(
120        &self,
121        notification: ServerNotification,
122    ) -> Result<(), AppServerError> {
123        let connection_ids: Vec<ConnectionId> = {
124            self.connections
125                .lock()
126                .await
127                .iter()
128                .filter_map(|(connection_id, state)| {
129                    if state.ready_for_notifications
130                        && !state
131                            .opt_out_notification_methods
132                            .contains(notification_method(&notification))
133                    {
134                        Some(*connection_id)
135                    } else {
136                        None
137                    }
138                })
139                .collect()
140        };
141
142        for connection_id in connection_ids {
143            self.send_notification(connection_id, notification.clone())
144                .await?;
145        }
146        Ok(())
147    }
148
149    pub async fn send_server_request(
150        &self,
151        connection_id: ConnectionId,
152        request: ServerRequest,
153    ) -> Result<(RequestId, oneshot::Receiver<ServerRequestResult>), AppServerError> {
154        let id = RequestId::String(format!(
155            "srvreq_{}",
156            self.next_server_request_id.fetch_add(1, Ordering::Relaxed)
157        ));
158        let (callback_tx, callback_rx) = oneshot::channel();
159        self.pending_server_requests
160            .lock()
161            .await
162            .insert(id.clone(), callback_tx);
163        let (method, params) = tagged_method_params(&request)?;
164        if let Err(error) = self
165            .send_message(
166                connection_id,
167                JsonRpcMessage::Request(crate::app_server::protocol::JsonRpcRequest {
168                    id: id.clone(),
169                    method,
170                    params,
171                }),
172            )
173            .await
174        {
175            self.pending_server_requests.lock().await.remove(&id);
176            return Err(error);
177        }
178        Ok((id, callback_rx))
179    }
180
181    pub async fn send_server_request_with_timeout(
182        &self,
183        connection_id: ConnectionId,
184        request: ServerRequest,
185        timeout: std::time::Duration,
186    ) -> Result<Value, AppServerError> {
187        let (id, callback_rx) = self.send_server_request(connection_id, request).await?;
188        match tokio::time::timeout(timeout, callback_rx).await {
189            Ok(Ok(Ok(value))) => Ok(value),
190            Ok(Ok(Err(error))) => Err(AppServerError::new(
191                AppServerErrorCode::InternalError,
192                error.message,
193            )
194            .with_data(error.data.unwrap_or(Value::Null))),
195            Ok(Err(_)) => Err(AppServerError::internal("server request callback dropped")),
196            Err(_) => {
197                self.pending_server_requests.lock().await.remove(&id);
198                Err(AppServerError::internal("server request timed out"))
199            }
200        }
201    }
202
203    pub async fn pending_server_request_count(&self) -> usize {
204        self.pending_server_requests.lock().await.len()
205    }
206
207    pub async fn send_server_request_with_id(
208        &self,
209        connection_id: ConnectionId,
210        id: RequestId,
211        request: ServerRequest,
212    ) -> Result<oneshot::Receiver<ServerRequestResult>, AppServerError> {
213        let (callback_tx, callback_rx) = oneshot::channel();
214        self.pending_server_requests
215            .lock()
216            .await
217            .insert(id.clone(), callback_tx);
218        let (method, params) = tagged_method_params(&request)?;
219        if let Err(error) = self
220            .send_message(
221                connection_id,
222                JsonRpcMessage::Request(crate::app_server::protocol::JsonRpcRequest {
223                    id: id.clone(),
224                    method,
225                    params,
226                }),
227            )
228            .await
229        {
230            self.pending_server_requests.lock().await.remove(&id);
231            return Err(error);
232        }
233        Ok(callback_rx)
234    }
235
236    pub async fn send_server_request_with_id_and_timeout(
237        &self,
238        connection_id: ConnectionId,
239        id: RequestId,
240        request: ServerRequest,
241        timeout: std::time::Duration,
242    ) -> Result<Value, AppServerError> {
243        let callback_rx = self
244            .send_server_request_with_id(connection_id, id.clone(), request)
245            .await?;
246        match tokio::time::timeout(timeout, callback_rx).await {
247            Ok(Ok(Ok(value))) => Ok(value),
248            Ok(Ok(Err(error))) => Err(AppServerError::new(
249                AppServerErrorCode::InternalError,
250                error.message,
251            )
252            .with_data(error.data.unwrap_or(Value::Null))),
253            Ok(Err(_)) => Err(AppServerError::internal("server request callback dropped")),
254            Err(_) => {
255                self.pending_server_requests.lock().await.remove(&id);
256                Err(AppServerError::internal("server request timed out"))
257            }
258        }
259    }
260
261    pub async fn send_json_rpc_request(
262        &self,
263        connection_id: ConnectionId,
264        id: RequestId,
265        method: String,
266        params: Option<Value>,
267    ) -> Result<(), AppServerError> {
268        self.send_message(
269            connection_id,
270            JsonRpcMessage::Request(crate::app_server::protocol::JsonRpcRequest {
271                id,
272                method,
273                params,
274            }),
275        )
276        .await
277    }
278
279    pub async fn resolve_server_response(&self, response: JsonRpcResponse) -> bool {
280        self.pending_server_requests
281            .lock()
282            .await
283            .remove(&response.id)
284            .is_some_and(|callback| callback.send(Ok(response.result)).is_ok())
285    }
286
287    pub async fn resolve_server_error(&self, error: JsonRpcError) -> bool {
288        self.pending_server_requests
289            .lock()
290            .await
291            .remove(&error.id)
292            .is_some_and(|callback| callback.send(Err(error.error)).is_ok())
293    }
294
295    async fn send_message(
296        &self,
297        connection_id: ConnectionId,
298        message: JsonRpcMessage,
299    ) -> Result<(), AppServerError> {
300        self.tx
301            .send(OutgoingEnvelope {
302                connection_id,
303                message,
304            })
305            .await
306            .map_err(|_| AppServerError::internal("outgoing channel closed"))
307    }
308
309    async fn should_send_notification(
310        &self,
311        connection_id: ConnectionId,
312        notification: &ServerNotification,
313    ) -> bool {
314        self.connections
315            .lock()
316            .await
317            .get(&connection_id)
318            .is_some_and(|state| {
319                state.ready_for_notifications
320                    && !state
321                        .opt_out_notification_methods
322                        .contains(notification_method(notification))
323            })
324    }
325}
326
327fn server_notification_message(
328    notification: ServerNotification,
329) -> Result<JsonRpcMessage, AppServerError> {
330    let (method, params) = tagged_method_params(&notification)?;
331    Ok(JsonRpcMessage::Notification(JsonRpcNotification {
332        method,
333        params,
334    }))
335}
336
337fn tagged_method_params<T: Serialize>(
338    value: &T,
339) -> Result<(String, Option<Value>), AppServerError> {
340    let mut value =
341        serde_json::to_value(value).map_err(|error| AppServerError::internal(error.to_string()))?;
342    let object = value
343        .as_object_mut()
344        .ok_or_else(|| AppServerError::internal("tagged protocol value is not an object"))?;
345    let method = object
346        .remove("method")
347        .and_then(|method| method.as_str().map(str::to_string))
348        .ok_or_else(|| AppServerError::internal("tagged protocol value is missing method"))?;
349    let params = object.remove("params");
350    Ok((method, params))
351}
352
353fn notification_method(notification: &ServerNotification) -> &'static str {
354    match notification {
355        ServerNotification::ThreadStarted(_) => "thread/started",
356        ServerNotification::ThreadArchived(_) => "thread/archived",
357        ServerNotification::TurnStarted(_) => "turn/started",
358        ServerNotification::TurnCompleted(_) => "turn/completed",
359        ServerNotification::ItemStarted(_) => "item/started",
360        ServerNotification::AgentMessageDelta(_) => "item/agentMessage/delta",
361        ServerNotification::ToolCallDelta(_) => "item/toolCall/delta",
362        ServerNotification::ItemCompleted(_) => "item/completed",
363        ServerNotification::ApprovalRequested(_) => "approval/requested",
364        ServerNotification::ApprovalResolved(_) => "approval/resolved",
365        ServerNotification::ErrorWarning(_) => "error/warning",
366    }
367}
vv_agent/app_server/outgoing.rs

vv_agent/app_server/
outgoing.rs
