1use crate::{
3 async_trait,
4 wire::{self, Envelope, Initialized, ServiceDescriptor},
5 Value,
6};
7use std::{collections::HashMap, io, sync::Arc, time::Duration};
8use tokio::{
9 io::{AsyncRead, AsyncWrite},
10 sync::{mpsc, watch},
11 task::JoinSet,
12};
13
14const MAX_CALLS: usize = 32;
15const MAX_ID: u64 = 9_007_199_254_740_991;
16
17#[derive(Debug, Clone)]
19pub struct CallError {
20 pub code: String,
21 pub message: String,
22}
23impl CallError {
24 pub fn new(code: impl Into<String>, message: impl Into<String>) -> Self {
25 Self {
26 code: code.into(),
27 message: message.into(),
28 }
29 }
30}
31
32#[derive(Clone)]
34pub struct RequestContext {
35 pub id: u64,
36 canceled: watch::Receiver<bool>,
37}
38impl RequestContext {
39 pub fn is_canceled(&self) -> bool {
40 *self.canceled.borrow() || self.canceled.has_changed().is_err()
41 }
42 pub async fn canceled(&self) {
43 let mut canceled = self.canceled.clone();
44 if !*canceled.borrow() {
45 let _ = canceled.changed().await;
46 }
47 }
48}
49
50#[async_trait]
51pub trait Adapter: Send + Sync + 'static {
52 async fn initialize(
54 &self,
55 configuration: Value,
56 context: RequestContext,
57 ) -> Result<Vec<ServiceDescriptor>, CallError>;
58 async fn call(
61 &self,
62 method: &str,
63 params: Value,
64 context: RequestContext,
65 ) -> Result<Value, CallError>;
66}
67
68pub fn run(adapter: impl Adapter) -> io::Result<()> {
69 let runtime = tokio::runtime::Builder::new_multi_thread()
70 .enable_all()
71 .build()?;
72 let result = runtime.block_on(serve(adapter, tokio::io::stdin(), tokio::io::stdout()));
73 runtime.shutdown_timeout(Duration::from_millis(100));
76 result
77}
78
79fn failed(id: u64, error: CallError) -> Envelope {
80 let valid = matches!(
81 error.code.as_str(),
82 "invalid"
83 | "closed"
84 | "aborted"
85 | "denied"
86 | "unavailable"
87 | "busy"
88 | "failed"
89 | "deadline"
90 ) && error.message.len() <= 4096;
91 Envelope::Error {
92 v: 1,
93 id,
94 code: if valid { error.code } else { "failed".into() },
95 message: if valid {
96 error.message
97 } else {
98 "Adapter returned an invalid error".into()
99 },
100 }
101}
102fn invalid(message: &str) -> io::Error {
103 io::Error::other(message)
104}
105
106pub async fn serve(
110 adapter: impl Adapter,
111 mut input: impl AsyncRead + Unpin + Send + 'static,
112 mut output: impl AsyncWrite + Unpin + Send + 'static,
113) -> io::Result<()> {
114 let adapter = Arc::new(adapter);
115 let (incoming, mut receive) = mpsc::channel(4);
116 let (send, mut outgoing) = mpsc::channel::<Vec<u8>>(MAX_CALLS * 2);
117 let mut pipes = JoinSet::new();
120 pipes.spawn(async move {
121 loop {
122 let frame = wire::read_frame_or_eof(&mut input).await;
123 let ended = !matches!(frame, Ok(Some(_)));
124 if incoming.send(frame).await.is_err() || ended {
125 break;
126 }
127 }
128 std::future::pending::<io::Result<()>>().await
129 });
130 pipes.spawn(async move {
131 while let Some(bytes) = outgoing.recv().await {
132 tokio::time::timeout(
133 Duration::from_secs(10),
134 wire::write_frame(&mut output, &bytes),
135 )
136 .await
137 .map_err(|_| invalid("Adapter output pipe stalled"))??;
138 }
139 Ok(())
140 });
141 let mut calls = JoinSet::new();
142 let mut cancellation: HashMap<u64, watch::Sender<bool>> = HashMap::new();
143 let mut last_id = 0;
144 let mut catalog: Option<Vec<ServiceDescriptor>> = None;
145 let result = loop {
146 tokio::select! {
147 ended = pipes.join_next() => {
148 break match ended {
149 Some(Ok(result)) => result,
150 _ => Err(invalid("Adapter protocol pipe failed")),
151 };
152 }
153 completed = calls.join_next(), if !calls.is_empty() => {
154 let Some(Ok((id, reply, initialized))) = completed else {
155 break Err(invalid("Adapter request task failed"));
156 };
157 cancellation.remove(&id);
158 if let Some(services) = initialized { catalog = Some(services); }
159 let bytes = match wire::encode(&reply) {
160 Ok(bytes) => bytes,
161 Err(error) => break Err(error),
162 };
163 if send.try_send(bytes).is_err() {
164 break Err(invalid("Adapter output queue is full or closed"));
165 }
166 }
167 message = receive.recv() => {
168 let message = match message {
169 Some(Ok(Some(message))) => message,
170 Some(Ok(None)) => break Ok(()),
171 Some(Err(error)) => break Err(error),
172 None => break Err(invalid("Adapter input stopped")),
173 };
174 let (id, method, params) = match message {
175 Envelope::Cancel {v: 1, id} if id > 0 && id <= last_id => {
176 if let Some(cancel) = cancellation.get(&id) { cancel.send_replace(true); }
177 continue;
178 }
179 Envelope::Request {v: 1, id, method, params} if id > last_id && id <= MAX_ID => (id, method, params),
180 _ => break Err(invalid("Invalid adapter request envelope or request identity")),
181 };
182 let initialize = last_id == 0 && method == "system.adapter.initialize";
183 last_id = id;
184 if !initialize && catalog.is_none() {
185 break Err(invalid("Initialize the adapter before sending service requests"));
186 }
187 let rejection = if calls.len() >= MAX_CALLS {
188 Some(CallError::new("busy", "Adapter has too many active requests"))
189 } else if !initialize && !catalog.as_ref().unwrap().iter().any(|service| service.methods.contains(&method)) {
190 Some(CallError::new("unavailable", "The adapter does not advertise this method"))
191 } else { None };
192 if let Some(error) = rejection {
193 let bytes = match wire::encode(&failed(id, error)) {
194 Ok(bytes) => bytes,
195 Err(error) => break Err(error),
196 };
197 if send.try_send(bytes).is_err() { break Err(invalid("Adapter output queue is full or closed")); }
198 continue;
199 }
200 let (cancel, canceled) = watch::channel(false);
201 cancellation.insert(id, cancel);
202 let context = RequestContext {id, canceled};
203 let adapter = adapter.clone();
204 calls.spawn(async move {
205 let mut services = None;
206 let value = if initialize {
207 #[derive(serde::Deserialize)]
208 #[serde(deny_unknown_fields)]
209 struct Setup {protocol: u8, configuration: Value}
210 match serde_json::from_value::<Setup>(params) {
211 Ok(setup) if setup.protocol == 1 => match adapter.initialize(setup.configuration, context).await {
212 Ok(catalog) => {
213 let initialized = Initialized {protocol: 1, services: catalog};
214 if wire::validate(&initialized) {
215 let value = serde_json::to_value(&initialized).expect("Protocol descriptor");
216 services = Some(initialized.services);
217 Ok(value)
218 } else { Err(CallError::new("invalid", "Adapter advertised an invalid service catalog")) }
219 }
220 Err(error) => Err(error),
221 },
222 _ => Err(CallError::new("invalid", "Unsupported adapter initialization")),
223 }
224 } else { adapter.call(&method, params, context).await };
225 let reply = match value {
226 Ok(value) => Envelope::Result {v: 1, id, value},
227 Err(error) => failed(id, error),
228 };
229 (id, reply, services)
230 });
231 }
232 }
233 };
234 for cancel in cancellation.values() {
235 cancel.send_replace(true);
236 }
237 let _ = tokio::time::timeout(Duration::from_secs(5), async {
238 while calls.join_next().await.is_some() {}
239 })
240 .await;
241 result
242}