1use std::{sync::Arc, time::Duration};
2
3use eoka_protocol::sidecar::{
4 validate_params, validate_safe_numbers, ErrorCode, ErrorData, RecordingStartParams, Request,
5 Response, ServerError, MAX_IN_FLIGHT,
6};
7use serde::Deserialize;
8use serde_json::{json, Value};
9use tokio::{
10 io::{AsyncBufRead, AsyncWrite, AsyncWriteExt},
11 sync::{mpsc, watch, OwnedSemaphorePermit, Semaphore},
12 task::JoinSet,
13};
14
15use crate::{
16 dispatch,
17 operations::sidecar,
18 state::AppState,
19 transport::{self, Frame, FrameReader},
20};
21
22struct Job {
23 request: Request,
24 operation_id: Option<String>,
25 permit: OwnedSemaphorePermit,
26}
27
28struct Outgoing {
29 id: Value,
30 operation_id: Option<String>,
31 response: Response,
32 _permit: Option<OwnedSemaphorePermit>,
33}
34
35fn operation_id(request: &Request) -> Option<String> {
36 request
37 .params
38 .get("operation_id")
39 .or_else(|| {
40 request
41 .method
42 .starts_with("operations.")
43 .then(|| request.params.get("id"))
44 .flatten()
45 })
46 .and_then(Value::as_str)
47 .map(str::to_owned)
48}
49
50fn with_operation(mut error: ServerError, operation_id: Option<&str>) -> ServerError {
51 if let Some(id) = operation_id {
52 error
53 .data
54 .get_or_insert_with(Box::<ErrorData>::default)
55 .operation_id = Some(id.to_owned());
56 }
57 error
58}
59
60impl Outgoing {
61 fn result(job: Job, result: Result<Value, ServerError>) -> Self {
62 Self::new(job.request.id, result, job.operation_id, Some(job.permit))
63 }
64
65 fn new(
66 id: Value,
67 result: Result<Value, ServerError>,
68 operation_id: Option<String>,
69 permit: Option<OwnedSemaphorePermit>,
70 ) -> Self {
71 let result = result.and_then(|result| {
72 validate_safe_numbers(&result).map_err(|error| {
73 ServerError::internal(format!("unsafe integer in response: {}", error.message))
74 })?;
75 Ok(result)
76 });
77 let response = match result {
78 Ok(result) => Response::ok(id.clone(), result),
79 Err(error) => Response::err(id.clone(), with_operation(error, operation_id.as_deref())),
80 };
81 Self {
82 id,
83 operation_id,
84 response,
85 _permit: permit,
86 }
87 }
88}
89
90fn emit(output: &mpsc::Sender<Outgoing>, response: Outgoing) -> anyhow::Result<()> {
91 output
92 .try_send(response)
93 .map_err(|_| anyhow::anyhow!("response queue is full or closed"))
94}
95
96async fn cancelled(shutdown: &mut watch::Receiver<bool>) {
97 if *shutdown.borrow() {
98 return;
99 }
100 let _ = shutdown.changed().await;
101}
102
103async fn worker(
104 mut state: AppState,
105 mut jobs: mpsc::Receiver<Job>,
106 output: mpsc::Sender<Outgoing>,
107 mut shutdown: watch::Receiver<bool>,
108 stop: watch::Sender<bool>,
109) -> Result<(), ServerError> {
110 loop {
111 let mut job = tokio::select! {
112 biased;
113 _ = cancelled(&mut shutdown) => break,
114 job = jobs.recv() => match job { Some(job) => job, None => break },
115 };
116 let protected = matches!(
117 job.request.method.as_str(),
118 "browser.launch" | "browser.connect"
119 ) || matches!(
120 job.request.method.as_str(),
121 "actions.perform" | "agent.step"
122 ) || job.request.params.get("operation_id").is_some();
123 let result = {
124 let params = std::mem::take(&mut job.request.params);
125 let dispatch = dispatch::dispatch(&mut state, &job.request.method, params);
126 tokio::pin!(dispatch);
127 tokio::select! {
128 biased;
129 _ = cancelled(&mut shutdown) => {
130 if protected {
131 match tokio::time::timeout(Duration::from_secs(35), &mut dispatch).await {
132 Ok(result) => result,
133 Err(_) => Err(ServerError::new(ErrorCode::Cancelled, "request exceeded the shutdown launch/admission grace deadline; inspect the operation ID before retrying")),
134 }
135 } else {
136 Err(ServerError::new(ErrorCode::Cancelled, "request interrupted by server shutdown"))
137 }
138 },
139 result = &mut dispatch => result,
140 }
141 };
142 if emit(&output, Outgoing::result(job, result)).is_err() {
143 break;
144 }
145 if state.should_shutdown() {
146 stop.send_replace(true);
147 break;
148 }
149 }
150 stop.send_replace(true);
151 jobs.close();
152 while let Some(job) = jobs.recv().await {
153 let _ = emit(
154 &output,
155 Outgoing::result(
156 job,
157 Err(ServerError::new(
158 ErrorCode::ShuttingDown,
159 "request was not dispatched because the server is shutting down",
160 )),
161 ),
162 );
163 }
164 state.shutdown().await
165}
166
167async fn writer<W: AsyncWrite + Unpin>(
168 mut output: W,
169 mut responses: mpsc::Receiver<Outgoing>,
170) -> anyhow::Result<()> {
171 while let Some(response) = responses.recv().await {
172 let bytes = transport::encode(
173 &response.response,
174 &response.id,
175 response.operation_id.as_deref(),
176 )?;
177 tokio::time::timeout(Duration::from_secs(2), async {
178 output.write_all(&bytes).await?;
179 output.flush().await
180 })
181 .await
182 .map_err(|_| {
183 anyhow::anyhow!("stdout backpressure exceeded the 2 second write deadline")
184 })??;
185 }
186 Ok(())
187}
188
189pub async fn serve<R, W>(input: R, output: W, state: AppState) -> anyhow::Result<()>
190where
191 R: AsyncBufRead + Unpin,
192 W: AsyncWrite + Unpin + Send + 'static,
193{
194 let mut input = FrameReader::new(input);
195 let operations = state.operations.clone();
196 let recordings = state.recordings.clone();
197 let pending = Arc::new(Semaphore::new(MAX_IN_FLIGHT));
198 let (jobs, job_receiver) = mpsc::channel(MAX_IN_FLIGHT);
199 let (responses, response_receiver) = mpsc::channel(MAX_IN_FLIGHT + 1);
200 let (stop, shutdown) = watch::channel(false);
201 let mut intake_shutdown = shutdown.clone();
202 let mut worker = tokio::spawn(worker(
203 state,
204 job_receiver,
205 responses.clone(),
206 shutdown.clone(),
207 stop.clone(),
208 ));
209 let mut writer = tokio::spawn(writer(output, response_receiver));
210 let mut waiters = JoinSet::new();
211 let mut shutdown_id = None;
212 let mut worker_result = None;
213 let mut writer_result = None;
214 let mut intake_error = None;
215 let mut input_closed = false;
216 loop {
217 tokio::select! {
218 biased;
219 _ = cancelled(&mut intake_shutdown) => break,
220 result = &mut writer => {
221 writer_result = Some(result);
222 break;
223 },
224 result = &mut worker => {
225 worker_result = Some(result);
226 break;
227 },
228 Some(result) = waiters.join_next(), if !waiters.is_empty() => {
229 match result {
230 Err(error) => {
231 intake_error = Some(anyhow::anyhow!("operation waiter failed: {error}"));
232 break;
233 },
234 Ok(Err(error)) => {
235 intake_error = Some(error);
236 break;
237 },
238 Ok(Ok(())) => {},
239 }
240 },
241 frame = input.next() => {
242 let frame = match frame {
243 Ok(Some(frame)) => frame,
244 Ok(None) => { input_closed = true; break; },
245 Err(error) => { intake_error = Some(error.into()); break; },
246 };
247 let request = match frame {
248 Frame::Empty => continue,
249 Frame::Error { id, error } => {
250 if let Err(error) = emit(&responses, Outgoing::new(id, Err(error), None, None)) {
251 intake_error = Some(error);
252 break;
253 }
254 continue;
255 },
256 Frame::Request(request) => request,
257 };
258 if request.method == "server.shutdown" {
259 match validate_params(&request.method, &request.params) {
260 Ok(()) => { shutdown_id = Some(request.id); break; },
261 Err(error) => {
262 if let Err(error) = emit(&responses, Outgoing::new(request.id, Err(error), None, None)) {
263 intake_error = Some(error);
264 break;
265 }
266 continue;
267 },
268 }
269 }
270 let permit = match pending.clone().try_acquire_owned() {
271 Ok(permit) => permit,
272 Err(_) => {
273 let id = operation_id(&request);
274 if let Err(error) = emit(&responses, Outgoing::new(request.id, Err(ServerError::new(ErrorCode::ServerBusy, "maxInFlight pending requests exceeded")), id, None)) {
275 intake_error = Some(error);
276 break;
277 }
278 continue;
279 },
280 };
281 let mut job = Job { operation_id: operation_id(&request), request, permit };
282 if job.request.method == "recording.start" {
283 let admission = RecordingStartParams::deserialize(&job.request.params)
284 .map_err(|error| ServerError::invalid_params(error.to_string()))
285 .and_then(|params| recordings.admit(params));
286 if let Err(error) = admission {
287 if let Err(error) = emit(&responses, Outgoing::result(job, Err(error))) {
288 intake_error = Some(error);
289 break;
290 }
291 continue;
292 }
293 }
294 if job.request.method == "server.hello" {
295 let result = dispatch::hello(std::mem::take(&mut job.request.params));
296 if let Err(error) = emit(&responses, Outgoing::result(job, result)) {
297 intake_error = Some(error);
298 break;
299 }
300 } else if crate::recordings::supports_control(&job.request.method)
301 || sidecar::supports_control(&job.request.method)
302 {
303 let operations = operations.clone();
304 let recordings = recordings.clone();
305 let output = responses.clone();
306 let jobs = jobs.clone();
307 let mut shutdown = shutdown.clone();
308 waiters.spawn(async move {
309 let result = tokio::select! {
310 biased;
311 _ = cancelled(&mut shutdown) => Some(Err(ServerError::new(ErrorCode::ShuttingDown, "wait ended because the server is shutting down; cancelling a waiter does not cancel its operation or recording"))),
312 result = async {
313 if matches!(job.request.method.as_str(), "operations.get" | "operations.cancel") {
314 sidecar::existing_control(&operations, &job.request.method, &job.request.params).await
315 } else {
316 let params = std::mem::take(&mut job.request.params);
317 Some(if crate::recordings::supports_control(&job.request.method) {
318 crate::recordings::control(&recordings, &job.request.method, params).await
319 } else {
320 sidecar::control(&operations, &job.request.method, params).await
321 })
322 }
323 } => result,
324 };
325 match result {
326 Some(result) => emit(&output, Outgoing::result(job, result)),
327 None => match jobs.try_send(job) {
328 Ok(()) => Ok(()),
329 Err(error) => emit(
330 &output,
331 Outgoing::result(
332 error.into_inner(),
333 Err(ServerError::new(ErrorCode::ShuttingDown, "dispatch worker is shutting down")),
334 ),
335 ),
336 },
337 }
338 });
339 } else if let Err(error) = jobs.try_send(job) {
340 let job = error.into_inner();
341 if let Err(error) = emit(&responses, Outgoing::result(job, Err(ServerError::new(ErrorCode::ShuttingDown, "dispatch worker is shutting down")))) {
342 intake_error = Some(error);
343 }
344 break;
345 }
346 },
347 }
348 }
349 stop.send_replace(true);
350 drop(jobs);
351 let cleanup = match worker_result {
352 Some(result) => result,
353 None => worker.await,
354 }
355 .map_err(|error| anyhow::anyhow!("dispatch worker failed: {error}"))?;
356 while let Some(result) = waiters.join_next().await {
357 if let Err(error) = result {
358 intake_error.get_or_insert_with(|| anyhow::anyhow!("operation waiter failed: {error}"));
359 } else if let Ok(Err(error)) = result {
360 intake_error.get_or_insert(error);
361 }
362 }
363 if let Some(id) = shutdown_id {
364 let result = match &cleanup {
365 Ok(()) => Ok(json!({})),
366 Err(error) => Err(error.clone()),
367 };
368 if let Err(error) = emit(&responses, Outgoing::new(id, result, None, None)) {
369 intake_error.get_or_insert(error);
370 }
371 }
372 drop(responses);
373 let written = match writer_result {
374 Some(result) => result.map_err(anyhow::Error::from)?,
375 None => match tokio::time::timeout(Duration::from_secs(2), &mut writer).await {
376 Ok(result) => result.map_err(anyhow::Error::from)?,
377 Err(_) => {
378 writer.abort();
379 let _ = writer.await;
380 Err(anyhow::anyhow!(
381 "response drain exceeded the 2 second deadline"
382 ))
383 }
384 },
385 };
386 cleanup.map_err(|error| anyhow::anyhow!(error.message))?;
387 let parent_lost = written.as_ref().err().is_some_and(|error| {
388 error
389 .downcast_ref::<std::io::Error>()
390 .is_some_and(|error| error.kind() == std::io::ErrorKind::BrokenPipe)
391 });
392 if input_closed || parent_lost {
393 return Ok(());
394 }
395 written?;
396 match intake_error {
397 Some(error) => Err(error),
398 None => Ok(()),
399 }
400}