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(Debug)]
30struct PendingServerRequest {
31    connection_id: ConnectionId,
32    method: String,
33    thread_id: Option<String>,
34    turn_id: Option<String>,
35    callback: oneshot::Sender<ServerRequestResult>,
36}
37
38#[derive(Clone)]
39pub struct OutgoingMessageSender {
40    tx: mpsc::Sender<OutgoingEnvelope>,
41    connections: Arc<Mutex<HashMap<ConnectionId, ConnectionOutgoingState>>>,
42    pending_server_requests: Arc<Mutex<HashMap<RequestId, PendingServerRequest>>>,
43    next_server_request_id: Arc<AtomicI64>,
44}
45
46impl OutgoingMessageSender {
47    pub fn channel(capacity: usize) -> (Self, mpsc::Receiver<OutgoingEnvelope>) {
48        let (tx, rx) = mpsc::channel(capacity);
49        (
50            Self {
51                tx,
52                connections: Arc::new(Mutex::new(HashMap::new())),
53                pending_server_requests: Arc::new(Mutex::new(HashMap::new())),
54                next_server_request_id: Arc::new(AtomicI64::new(1)),
55            },
56            rx,
57        )
58    }
59
60    pub async fn register_connection(&self, connection_id: ConnectionId) {
61        self.connections
62            .lock()
63            .await
64            .entry(connection_id)
65            .or_default();
66    }
67
68    pub async fn is_connection_registered(&self, connection_id: ConnectionId) -> bool {
69        self.connections.lock().await.contains_key(&connection_id)
70    }
71
72    pub async fn unregister_connection(&self, connection_id: ConnectionId) {
73        self.connections.lock().await.remove(&connection_id);
74        let pending = {
75            let mut requests = self.pending_server_requests.lock().await;
76            let request_ids = requests
77                .iter()
78                .filter_map(|(request_id, pending)| {
79                    (pending.connection_id == connection_id).then_some(request_id.clone())
80                })
81                .collect::<Vec<_>>();
82            request_ids
83                .into_iter()
84                .filter_map(|request_id| requests.remove(&request_id))
85                .collect::<Vec<_>>()
86        };
87        let error = JsonRpcErrorBody {
88            code: AppServerErrorCode::InternalError.code(),
89            message: "client_disconnected".to_string(),
90            data: None,
91        };
92        for pending in pending {
93            let _ = pending.callback.send(Err(error.clone()));
94        }
95    }
96
97    pub async fn configure_connection(
98        &self,
99        connection_id: ConnectionId,
100        opt_out_notification_methods: HashSet<String>,
101    ) {
102        let mut connections = self.connections.lock().await;
103        let state = connections.entry(connection_id).or_default();
104        state.opt_out_notification_methods = opt_out_notification_methods;
105    }
106
107    pub async fn mark_ready_for_notifications(&self, connection_id: ConnectionId) {
108        self.connections
109            .lock()
110            .await
111            .entry(connection_id)
112            .or_default()
113            .ready_for_notifications = true;
114    }
115
116    pub async fn send_response(
117        &self,
118        connection_id: ConnectionId,
119        id: RequestId,
120        result: Value,
121    ) -> Result<(), AppServerError> {
122        self.send_message(
123            connection_id,
124            JsonRpcMessage::Response(JsonRpcResponse { id, result }),
125        )
126        .await
127    }
128
129    pub async fn send_error(
130        &self,
131        connection_id: ConnectionId,
132        id: RequestId,
133        error: AppServerError,
134    ) -> Result<(), AppServerError> {
135        self.send_message(
136            connection_id,
137            JsonRpcMessage::Error(error.into_json_rpc_error(id)),
138        )
139        .await
140    }
141
142    pub async fn send_notification(
143        &self,
144        connection_id: ConnectionId,
145        notification: ServerNotification,
146    ) -> Result<(), AppServerError> {
147        if !self
148            .should_send_notification(connection_id, &notification)
149            .await
150        {
151            return Ok(());
152        }
153        let message = server_notification_message(notification)?;
154        self.send_message(connection_id, message).await
155    }
156
157    pub async fn broadcast_initialized_client_notification(
158        &self,
159        notification: ServerNotification,
160    ) -> Result<(), AppServerError> {
161        let connection_ids: Vec<ConnectionId> = {
162            self.connections
163                .lock()
164                .await
165                .iter()
166                .filter_map(|(connection_id, state)| {
167                    if state.ready_for_notifications
168                        && !state
169                            .opt_out_notification_methods
170                            .contains(notification_method(&notification))
171                    {
172                        Some(*connection_id)
173                    } else {
174                        None
175                    }
176                })
177                .collect()
178        };
179
180        for connection_id in connection_ids {
181            self.send_notification(connection_id, notification.clone())
182                .await?;
183        }
184        Ok(())
185    }
186
187    pub async fn send_server_request(
188        &self,
189        connection_id: ConnectionId,
190        request: ServerRequest,
191    ) -> Result<(RequestId, oneshot::Receiver<ServerRequestResult>), AppServerError> {
192        let id = RequestId::String(format!(
193            "srvreq_{}",
194            self.next_server_request_id.fetch_add(1, Ordering::Relaxed)
195        ));
196        let callback_rx = self
197            .send_server_request_with_id(connection_id, id.clone(), request)
198            .await?;
199        Ok((id, callback_rx))
200    }
201
202    pub async fn send_server_request_with_timeout(
203        &self,
204        connection_id: ConnectionId,
205        request: ServerRequest,
206        timeout: std::time::Duration,
207    ) -> Result<Value, AppServerError> {
208        let (id, callback_rx) = self.send_server_request(connection_id, request).await?;
209        match tokio::time::timeout(timeout, callback_rx).await {
210            Ok(Ok(Ok(value))) => Ok(value),
211            Ok(Ok(Err(error))) => Err(AppServerError::new(
212                AppServerErrorCode::InternalError,
213                error.message,
214            )
215            .with_data(error.data.unwrap_or(Value::Null))),
216            Ok(Err(_)) => Err(AppServerError::internal("server request callback dropped")),
217            Err(_) => {
218                self.pending_server_requests.lock().await.remove(&id);
219                Err(AppServerError::internal("server request timed out"))
220            }
221        }
222    }
223
224    pub async fn pending_server_request_count(&self) -> usize {
225        self.pending_server_requests.lock().await.len()
226    }
227
228    pub async fn send_server_request_with_id(
229        &self,
230        connection_id: ConnectionId,
231        id: RequestId,
232        request: ServerRequest,
233    ) -> Result<oneshot::Receiver<ServerRequestResult>, AppServerError> {
234        if matches!(id, RequestId::Null) {
235            return Err(AppServerError::invalid_params(
236                "Server request id cannot be null",
237            ));
238        }
239        let (callback_tx, callback_rx) = oneshot::channel();
240        let (method, params) = tagged_method_params(&request)?;
241        let (thread_id, turn_id) = server_request_scope(&request);
242        {
243            let mut requests = self.pending_server_requests.lock().await;
244            if requests.contains_key(&id) {
245                return Err(AppServerError::invalid_params(
246                    "Duplicate server request id",
247                ));
248            }
249            requests.insert(
250                id.clone(),
251                PendingServerRequest {
252                    connection_id,
253                    method: method.clone(),
254                    thread_id,
255                    turn_id,
256                    callback: callback_tx,
257                },
258            );
259        }
260        if let Err(error) = self
261            .send_message(
262                connection_id,
263                JsonRpcMessage::Request(crate::app_server::protocol::JsonRpcRequest {
264                    id: id.clone(),
265                    method,
266                    params,
267                }),
268            )
269            .await
270        {
271            self.pending_server_requests.lock().await.remove(&id);
272            return Err(error);
273        }
274        Ok(callback_rx)
275    }
276
277    pub async fn send_server_request_with_id_and_timeout(
278        &self,
279        connection_id: ConnectionId,
280        id: RequestId,
281        request: ServerRequest,
282        timeout: std::time::Duration,
283    ) -> Result<Value, AppServerError> {
284        let callback_rx = self
285            .send_server_request_with_id(connection_id, id.clone(), request)
286            .await?;
287        match tokio::time::timeout(timeout, callback_rx).await {
288            Ok(Ok(Ok(value))) => Ok(value),
289            Ok(Ok(Err(error))) => Err(AppServerError::new(
290                AppServerErrorCode::InternalError,
291                error.message,
292            )
293            .with_data(error.data.unwrap_or(Value::Null))),
294            Ok(Err(_)) => Err(AppServerError::internal("server request callback dropped")),
295            Err(_) => {
296                self.pending_server_requests.lock().await.remove(&id);
297                Err(AppServerError::internal("server request timed out"))
298            }
299        }
300    }
301
302    pub async fn send_json_rpc_request(
303        &self,
304        connection_id: ConnectionId,
305        id: RequestId,
306        method: String,
307        params: Option<Value>,
308    ) -> Result<(), AppServerError> {
309        self.send_message(
310            connection_id,
311            JsonRpcMessage::Request(crate::app_server::protocol::JsonRpcRequest {
312                id,
313                method,
314                params,
315            }),
316        )
317        .await
318    }
319
320    pub async fn resolve_server_response(
321        &self,
322        connection_id: ConnectionId,
323        response: JsonRpcResponse,
324    ) -> bool {
325        self.resolve_server_response_bound(connection_id, None, None, None, response)
326            .await
327    }
328
329    pub async fn resolve_server_response_bound(
330        &self,
331        connection_id: ConnectionId,
332        method: Option<&str>,
333        thread_id: Option<&str>,
334        turn_id: Option<&str>,
335        response: JsonRpcResponse,
336    ) -> bool {
337        let pending = {
338            let mut requests = self.pending_server_requests.lock().await;
339            let matches = requests.get(&response.id).is_some_and(|pending| {
340                pending.connection_id == connection_id
341                    && method.is_none_or(|method| pending.method == method)
342                    && thread_id
343                        .is_none_or(|thread_id| pending.thread_id.as_deref() == Some(thread_id))
344                    && turn_id.is_none_or(|turn_id| pending.turn_id.as_deref() == Some(turn_id))
345            });
346            matches.then(|| requests.remove(&response.id)).flatten()
347        };
348        pending.is_some_and(|pending| pending.callback.send(Ok(response.result)).is_ok())
349    }
350
351    pub async fn resolve_server_error(
352        &self,
353        connection_id: ConnectionId,
354        error: JsonRpcError,
355    ) -> bool {
356        let pending = {
357            let mut requests = self.pending_server_requests.lock().await;
358            let matches = requests
359                .get(&error.id)
360                .is_some_and(|pending| pending.connection_id == connection_id);
361            matches.then(|| requests.remove(&error.id)).flatten()
362        };
363        pending.is_some_and(|pending| pending.callback.send(Err(error.error)).is_ok())
364    }
365
366    async fn send_message(
367        &self,
368        connection_id: ConnectionId,
369        message: JsonRpcMessage,
370    ) -> Result<(), AppServerError> {
371        if !self.is_connection_registered(connection_id).await {
372            return Err(AppServerError::internal("client_disconnected"));
373        }
374        self.tx
375            .send(OutgoingEnvelope {
376                connection_id,
377                message,
378            })
379            .await
380            .map_err(|_| AppServerError::internal("outgoing channel closed"))
381    }
382
383    async fn should_send_notification(
384        &self,
385        connection_id: ConnectionId,
386        notification: &ServerNotification,
387    ) -> bool {
388        self.connections
389            .lock()
390            .await
391            .get(&connection_id)
392            .is_some_and(|state| {
393                state.ready_for_notifications
394                    && !state
395                        .opt_out_notification_methods
396                        .contains(notification_method(notification))
397            })
398    }
399}
400
401fn server_request_scope(request: &ServerRequest) -> (Option<String>, Option<String>) {
402    match request {
403        ServerRequest::ApprovalRequest(params) => {
404            (Some(params.thread_id.clone()), Some(params.turn_id.clone()))
405        }
406    }
407}
408
409fn server_notification_message(
410    notification: ServerNotification,
411) -> Result<JsonRpcMessage, AppServerError> {
412    let (method, params) = tagged_method_params(&notification)?;
413    Ok(JsonRpcMessage::Notification(JsonRpcNotification {
414        method,
415        params,
416    }))
417}
418
419fn tagged_method_params<T: Serialize>(
420    value: &T,
421) -> Result<(String, Option<Value>), AppServerError> {
422    let mut value =
423        serde_json::to_value(value).map_err(|error| AppServerError::internal(error.to_string()))?;
424    let object = value
425        .as_object_mut()
426        .ok_or_else(|| AppServerError::internal("tagged protocol value is not an object"))?;
427    let method = object
428        .remove("method")
429        .and_then(|method| method.as_str().map(str::to_string))
430        .ok_or_else(|| AppServerError::internal("tagged protocol value is missing method"))?;
431    let params = object.remove("params");
432    Ok((method, params))
433}
434
435fn notification_method(notification: &ServerNotification) -> &'static str {
436    match notification {
437        ServerNotification::ThreadStarted(_) => "thread/started",
438        ServerNotification::ThreadArchived(_) => "thread/archived",
439        ServerNotification::ThreadClosed(_) => "thread/closed",
440        ServerNotification::ThreadStatusChanged(_) => "thread/status/changed",
441        ServerNotification::TurnStarted(_) => "turn/started",
442        ServerNotification::TurnCompleted(_) => "turn/completed",
443        ServerNotification::ItemStarted(_) => "item/started",
444        ServerNotification::AgentMessageDelta(_) => "item/agentMessage/delta",
445        ServerNotification::ToolCallDelta(_) => "item/toolCall/delta",
446        ServerNotification::ItemCompleted(_) => "item/completed",
447        ServerNotification::ApprovalRequested(_) => "approval/requested",
448        ServerNotification::ApprovalResolved(_) => "approval/resolved",
449        ServerNotification::ErrorWarning(_) => "error/warning",
450    }
451}
vv_agent/app_server/outgoing.rs

vv_agent/app_server/
outgoing.rs
