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}