Skip to main content

vv_agent/app_server/
client.rs

1use std::collections::VecDeque;
2use std::fmt;
3use std::time::Duration;
4
5use serde::de::DeserializeOwned;
6use serde_json::{json, Value};
7use tokio::sync::mpsc;
8
9use crate::app_server::outgoing::OutgoingEnvelope;
10use crate::app_server::processor::MessageProcessor;
11use crate::app_server::protocol::{
12    AppClientCapabilities, AppClientInfo, ApprovalDecision, ApprovalResolveParams,
13    InitializeParams, InitializeResponse, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest,
14    JsonRpcResponse, RequestId, ThreadReadParams, ThreadReadResponse, ThreadStartParams,
15    ThreadStartResponse, TurnStartParams, TurnStartResponse,
16};
17use crate::app_server::transport::ConnectionId;
18
19pub struct AppServerClient {
20    processor: MessageProcessor,
21    outgoing: mpsc::Receiver<OutgoingEnvelope>,
22    connection_id: ConnectionId,
23    next_request_id: i64,
24    backlog: VecDeque<JsonRpcMessage>,
25}
26
27impl AppServerClient {
28    pub fn new_for_processor(
29        processor: MessageProcessor,
30        outgoing: mpsc::Receiver<OutgoingEnvelope>,
31        connection_id: ConnectionId,
32    ) -> Self {
33        Self {
34            processor,
35            outgoing,
36            connection_id,
37            next_request_id: 1,
38            backlog: VecDeque::new(),
39        }
40    }
41
42    pub async fn initialize(
43        &mut self,
44        info: AppClientInfo,
45    ) -> Result<InitializeResponse, AppServerClientError> {
46        let response = self
47            .send_request(
48                "initialize",
49                serde_json::to_value(InitializeParams {
50                    client_info: info,
51                    capabilities: AppClientCapabilities::default(),
52                })?,
53            )
54            .await?;
55        self.processor
56            .process_message(
57                self.connection_id,
58                JsonRpcMessage::Notification(JsonRpcNotification {
59                    method: "initialized".to_string(),
60                    params: None,
61                }),
62            )
63            .await;
64        Ok(response)
65    }
66
67    pub async fn start_thread(
68        &mut self,
69        params: ThreadStartParams,
70    ) -> Result<ThreadStartResponse, AppServerClientError> {
71        self.send_request("thread/start", serde_json::to_value(params)?)
72            .await
73    }
74
75    pub async fn start_turn(
76        &mut self,
77        params: TurnStartParams,
78    ) -> Result<TurnStartResponse, AppServerClientError> {
79        self.send_request("turn/start", serde_json::to_value(params)?)
80            .await
81    }
82
83    pub async fn resolve_approval(
84        &mut self,
85        params: ApprovalResolveParams,
86    ) -> Result<(), AppServerClientError> {
87        let decision = match params.decision {
88            ApprovalDecision::Allow => "allow",
89            ApprovalDecision::Deny => "deny",
90        };
91        self.processor
92            .process_message(
93                self.connection_id,
94                JsonRpcMessage::Response(JsonRpcResponse {
95                    id: RequestId::String(params.request_id),
96                    result: json!({ "decision": decision }),
97                }),
98            )
99            .await;
100        Ok(())
101    }
102
103    pub async fn read_thread(
104        &mut self,
105        params: ThreadReadParams,
106    ) -> Result<ThreadReadResponse, AppServerClientError> {
107        self.send_request("thread/read", serde_json::to_value(params)?)
108            .await
109    }
110
111    pub async fn next_message(&mut self) -> Option<JsonRpcMessage> {
112        if let Some(message) = self.backlog.pop_front() {
113            return Some(message);
114        }
115        self.next_outgoing_message().await.ok()
116    }
117
118    async fn send_request<T: DeserializeOwned>(
119        &mut self,
120        method: &str,
121        params: Value,
122    ) -> Result<T, AppServerClientError> {
123        let id = RequestId::Integer(self.next_request_id);
124        self.next_request_id += 1;
125        self.processor
126            .process_message(
127                self.connection_id,
128                JsonRpcMessage::Request(JsonRpcRequest {
129                    id: id.clone(),
130                    method: method.to_string(),
131                    params: Some(params),
132                }),
133            )
134            .await;
135        self.wait_response(id).await
136    }
137
138    async fn wait_response<T: DeserializeOwned>(
139        &mut self,
140        id: RequestId,
141    ) -> Result<T, AppServerClientError> {
142        loop {
143            match self.next_outgoing_message().await? {
144                JsonRpcMessage::Response(response) if response.id == id => {
145                    return Ok(serde_json::from_value(response.result)?);
146                }
147                JsonRpcMessage::Error(error) if error.id == id => {
148                    return Err(AppServerClientError::server(error.error.message));
149                }
150                message => self.backlog.push_back(message),
151            }
152        }
153    }
154
155    async fn next_outgoing_message(&mut self) -> Result<JsonRpcMessage, AppServerClientError> {
156        loop {
157            let envelope = tokio::time::timeout(Duration::from_secs(3), self.outgoing.recv())
158                .await
159                .map_err(|_| AppServerClientError::timeout("timed out waiting for message"))?
160                .ok_or_else(|| AppServerClientError::transport("outgoing channel closed"))?;
161            if envelope.connection_id == self.connection_id {
162                return Ok(envelope.message);
163            }
164        }
165    }
166}
167
168#[derive(Debug, Clone, PartialEq, Eq)]
169pub struct AppServerClientError {
170    message: String,
171}
172
173impl AppServerClientError {
174    fn server(message: impl Into<String>) -> Self {
175        Self {
176            message: message.into(),
177        }
178    }
179
180    fn timeout(message: impl Into<String>) -> Self {
181        Self {
182            message: message.into(),
183        }
184    }
185
186    fn transport(message: impl Into<String>) -> Self {
187        Self {
188            message: message.into(),
189        }
190    }
191}
192
193impl fmt::Display for AppServerClientError {
194    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
195        formatter.write_str(&self.message)
196    }
197}
198
199impl std::error::Error for AppServerClientError {}
200
201impl From<serde_json::Error> for AppServerClientError {
202    fn from(error: serde_json::Error) -> Self {
203        Self {
204            message: error.to_string(),
205        }
206    }
207}