1use std::collections::VecDeque;
2use std::fmt;
3use std::time::Duration;
4
5use serde::de::DeserializeOwned;
6use serde_json::{json, Map, 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, ApprovalResolveParams, InitializeParams,
13 InitializeResponse, JsonRpcErrorBody, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest,
14 JsonRpcResponse, ModelListParams, ModelListResponse, RequestId, SchemaExportResponse,
15 ThreadArchiveParams, ThreadArchiveResponse, ThreadListParams, ThreadListResponse,
16 ThreadReadParams, ThreadReadResponse, ThreadResumeParams, ThreadResumeResponse,
17 ThreadStartParams, ThreadStartResponse, ThreadUnsubscribeParams, ThreadUnsubscribeResponse,
18 TurnFollowUpParams, TurnFollowUpResponse, TurnInterruptParams, TurnInterruptResponse,
19 TurnResumeParams, TurnResumeResponse, TurnStartParams, TurnStartResponse, TurnSteerParams,
20 TurnSteerResponse,
21};
22use crate::app_server::transport::ConnectionId;
23
24const RESPONSE_TIMEOUT: Duration = Duration::from_secs(3);
25
26pub struct AppServerClient {
27 processor: MessageProcessor,
28 outgoing: mpsc::Receiver<OutgoingEnvelope>,
29 connection_id: ConnectionId,
30 next_request_id: i64,
31 backlog: VecDeque<JsonRpcMessage>,
32 closed: bool,
33}
34
35impl AppServerClient {
36 pub fn new_for_processor(
37 processor: MessageProcessor,
38 outgoing: mpsc::Receiver<OutgoingEnvelope>,
39 connection_id: ConnectionId,
40 ) -> Self {
41 Self {
42 processor,
43 outgoing,
44 connection_id,
45 next_request_id: 1,
46 backlog: VecDeque::new(),
47 closed: false,
48 }
49 }
50
51 pub async fn initialize(
52 &mut self,
53 info: AppClientInfo,
54 ) -> Result<InitializeResponse, AppServerClientError> {
55 let response = self
56 .send_request(
57 "initialize",
58 serde_json::to_value(InitializeParams {
59 client_info: info,
60 capabilities: AppClientCapabilities::default(),
61 })?,
62 )
63 .await?;
64 self.send_notification("initialized", None).await?;
65 Ok(response)
66 }
67
68 pub async fn start_thread(
69 &mut self,
70 params: ThreadStartParams,
71 ) -> Result<ThreadStartResponse, AppServerClientError> {
72 self.send_request("thread/start", serde_json::to_value(params)?)
73 .await
74 }
75
76 pub async fn resume_thread(
77 &mut self,
78 params: ThreadResumeParams,
79 ) -> Result<ThreadResumeResponse, AppServerClientError> {
80 self.send_request("thread/resume", serde_json::to_value(params)?)
81 .await
82 }
83
84 pub async fn read_thread(
85 &mut self,
86 params: ThreadReadParams,
87 ) -> Result<ThreadReadResponse, AppServerClientError> {
88 self.send_request("thread/read", serde_json::to_value(params)?)
89 .await
90 }
91
92 pub async fn list_threads(
93 &mut self,
94 params: ThreadListParams,
95 ) -> Result<ThreadListResponse, AppServerClientError> {
96 self.send_request("thread/list", serde_json::to_value(params)?)
97 .await
98 }
99
100 pub async fn archive_thread(
101 &mut self,
102 params: ThreadArchiveParams,
103 ) -> Result<ThreadArchiveResponse, AppServerClientError> {
104 self.send_request("thread/archive", serde_json::to_value(params)?)
105 .await
106 }
107
108 pub async fn unsubscribe_thread(
109 &mut self,
110 params: ThreadUnsubscribeParams,
111 ) -> Result<ThreadUnsubscribeResponse, AppServerClientError> {
112 self.send_request("thread/unsubscribe", serde_json::to_value(params)?)
113 .await
114 }
115
116 pub async fn start_turn(
117 &mut self,
118 params: TurnStartParams,
119 ) -> Result<TurnStartResponse, AppServerClientError> {
120 self.send_request("turn/start", serde_json::to_value(params)?)
121 .await
122 }
123
124 pub async fn resume_turn(
125 &mut self,
126 params: TurnResumeParams,
127 ) -> Result<TurnResumeResponse, AppServerClientError> {
128 self.send_request("turn/resume", serde_json::to_value(params)?)
129 .await
130 }
131
132 pub async fn interrupt_turn(
133 &mut self,
134 params: TurnInterruptParams,
135 ) -> Result<TurnInterruptResponse, AppServerClientError> {
136 self.send_request("turn/interrupt", serde_json::to_value(params)?)
137 .await
138 }
139
140 pub async fn steer_turn(
141 &mut self,
142 params: TurnSteerParams,
143 ) -> Result<TurnSteerResponse, AppServerClientError> {
144 self.send_request("turn/steer", serde_json::to_value(params)?)
145 .await
146 }
147
148 pub async fn follow_up_turn(
149 &mut self,
150 params: TurnFollowUpParams,
151 ) -> Result<TurnFollowUpResponse, AppServerClientError> {
152 self.send_request("turn/followUp", serde_json::to_value(params)?)
153 .await
154 }
155
156 pub async fn resolve_approval(
157 &mut self,
158 params: ApprovalResolveParams,
159 ) -> Result<(), AppServerClientError> {
160 self.send_response(
161 RequestId::String(params.request_id),
162 json!({
163 "decision": params.decision.as_wire(),
164 "reason": params.reason,
165 "metadata": params.metadata,
166 }),
167 )
168 .await
169 }
170
171 pub async fn resolve_approval_request(
172 &mut self,
173 params: ApprovalResolveParams,
174 ) -> Result<(), AppServerClientError> {
175 let _: Map<String, Value> = self
176 .send_request("approval/resolve", serde_json::to_value(params)?)
177 .await?;
178 Ok(())
179 }
180
181 pub async fn send_response(
182 &mut self,
183 request_id: RequestId,
184 result: Value,
185 ) -> Result<(), AppServerClientError> {
186 self.ensure_open()?;
187 self.processor
188 .process_message(
189 self.connection_id,
190 JsonRpcMessage::Response(JsonRpcResponse {
191 id: request_id,
192 result,
193 }),
194 )
195 .await;
196 Ok(())
197 }
198
199 pub async fn list_models(
200 &mut self,
201 params: ModelListParams,
202 ) -> Result<ModelListResponse, AppServerClientError> {
203 self.send_request("model/list", serde_json::to_value(params)?)
204 .await
205 }
206
207 pub async fn export_schema(&mut self) -> Result<SchemaExportResponse, AppServerClientError> {
208 self.send_request("schema/export", json!({})).await
209 }
210
211 pub async fn next_message(&mut self) -> Option<JsonRpcMessage> {
212 self.try_next_message().await.ok()
213 }
214
215 pub async fn try_next_message(&mut self) -> Result<JsonRpcMessage, AppServerClientError> {
216 self.ensure_open()?;
217 if let Some(message) = self.backlog.pop_front() {
218 return Ok(message);
219 }
220 self.next_outgoing_message("timed out waiting for App Server message")
221 .await
222 }
223
224 pub async fn close(&mut self) -> bool {
225 if self.closed {
226 return false;
227 }
228 self.closed = true;
229 self.backlog.clear();
230 self.processor
231 .disconnect_connection(self.connection_id)
232 .await;
233 true
234 }
235
236 async fn send_request<T: DeserializeOwned>(
237 &mut self,
238 method: &str,
239 params: Value,
240 ) -> Result<T, AppServerClientError> {
241 self.ensure_open()?;
242 let id = RequestId::Integer(self.next_request_id);
243 self.next_request_id += 1;
244 self.processor
245 .process_message(
246 self.connection_id,
247 JsonRpcMessage::Request(JsonRpcRequest {
248 id: id.clone(),
249 method: method.to_string(),
250 params: Some(params),
251 }),
252 )
253 .await;
254 self.wait_response(id, method).await
255 }
256
257 async fn send_notification(
258 &mut self,
259 method: &str,
260 params: Option<Value>,
261 ) -> Result<(), AppServerClientError> {
262 self.ensure_open()?;
263 self.processor
264 .process_message(
265 self.connection_id,
266 JsonRpcMessage::Notification(JsonRpcNotification {
267 method: method.to_string(),
268 params,
269 }),
270 )
271 .await;
272 Ok(())
273 }
274
275 async fn wait_response<T: DeserializeOwned>(
276 &mut self,
277 id: RequestId,
278 method: &str,
279 ) -> Result<T, AppServerClientError> {
280 for _ in 0..self.backlog.len() {
281 let message = self
282 .backlog
283 .pop_front()
284 .expect("backlog length was captured before scanning");
285 match message {
286 JsonRpcMessage::Response(response) if response.id == id => {
287 return Self::decode_response(response.result, method);
288 }
289 JsonRpcMessage::Error(error) if error.id == id => {
290 return Err(AppServerClientError::server(error.error));
291 }
292 message => self.backlog.push_back(message),
293 }
294 }
295
296 loop {
297 let timeout_message = format!("timed out waiting for App Server {method} response");
298 match self.next_outgoing_message(&timeout_message).await? {
299 JsonRpcMessage::Response(response) if response.id == id => {
300 return Self::decode_response(response.result, method);
301 }
302 JsonRpcMessage::Error(error) if error.id == id => {
303 return Err(AppServerClientError::server(error.error));
304 }
305 message => self.backlog.push_back(message),
306 }
307 }
308 }
309
310 fn decode_response<T: DeserializeOwned>(
311 result: Value,
312 method: &str,
313 ) -> Result<T, AppServerClientError> {
314 serde_json::from_value(result).map_err(|error| {
315 AppServerClientError::client(format!(
316 "failed decoding App Server {method} response: {error}"
317 ))
318 })
319 }
320
321 async fn next_outgoing_message(
322 &mut self,
323 timeout_message: &str,
324 ) -> Result<JsonRpcMessage, AppServerClientError> {
325 loop {
326 let envelope = tokio::time::timeout(RESPONSE_TIMEOUT, self.outgoing.recv())
327 .await
328 .map_err(|_| AppServerClientError::timeout(timeout_message))?
329 .ok_or_else(|| AppServerClientError::transport("outgoing channel closed"))?;
330 if envelope.connection_id == self.connection_id {
331 return Ok(envelope.message);
332 }
333 }
334 }
335
336 fn ensure_open(&self) -> Result<(), AppServerClientError> {
337 if self.closed {
338 Err(AppServerClientError::closed())
339 } else {
340 Ok(())
341 }
342 }
343}
344
345#[derive(Debug, Clone, PartialEq, Eq)]
346pub struct AppServerClientError {
347 message: String,
348 code: Option<i64>,
349 data: Option<Value>,
350}
351
352impl AppServerClientError {
353 fn server(error: JsonRpcErrorBody) -> Self {
354 Self {
355 message: error.message,
356 code: Some(error.code),
357 data: error.data,
358 }
359 }
360
361 fn timeout(message: impl Into<String>) -> Self {
362 Self::client(message)
363 }
364
365 fn transport(message: impl Into<String>) -> Self {
366 Self::client(message)
367 }
368
369 fn closed() -> Self {
370 Self::client("App Server client is closed")
371 }
372
373 fn client(message: impl Into<String>) -> Self {
374 Self {
375 message: message.into(),
376 code: None,
377 data: None,
378 }
379 }
380
381 pub fn message(&self) -> &str {
382 &self.message
383 }
384
385 pub fn code(&self) -> Option<i64> {
386 self.code
387 }
388
389 pub fn data(&self) -> Option<&Value> {
390 self.data.as_ref()
391 }
392}
393
394impl fmt::Display for AppServerClientError {
395 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
396 formatter.write_str(&self.message)
397 }
398}
399
400impl std::error::Error for AppServerClientError {}
401
402impl From<serde_json::Error> for AppServerClientError {
403 fn from(error: serde_json::Error) -> Self {
404 Self::client(error.to_string())
405 }
406}