1use std::collections::BTreeMap;
12use std::num::NonZeroU64;
13use std::time::Duration;
14
15use async_trait::async_trait;
16use onetaskgraph_plugin_api::{
17 Capabilities, DependencyEdge, Direction, Document, DocumentQuery, Health, ItemWrite, Label,
18 NativeId, Page, PageRequest, Project, ProjectQuery, SourceError, SourceName, Task, TaskQuery,
19 TaskSource, WriteSupport,
20};
21use serde::Deserialize;
22use serde_json::{Value, json};
23
24use super::connection::{Connection, Peer};
25use super::wire::{
26 DeleteParams, DependencyParams, DocumentQueryParams, DocumentResult, DocumentWriteParams,
27 EngineIdentity, IdParams, InitializeParams, InitializeResult, LabelParams, PROTOCOL_VERSION,
28 ProjectQueryParams, ProjectResult, ProjectWriteParams, Request, TaskQueryParams, TaskResult,
29 TaskWriteParams, WriteResult,
30};
31
32const HANDSHAKE_ID: &str = "0";
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub struct RequestDeadline(NonZeroU64);
40
41impl RequestDeadline {
42 pub const DEFAULT: Self = Self(NonZeroU64::new(30_000).expect("non-zero default"));
44
45 #[must_use]
47 pub const fn from_millis(milliseconds: NonZeroU64) -> Self {
48 Self(milliseconds)
49 }
50
51 #[must_use]
53 pub const fn milliseconds(self) -> NonZeroU64 {
54 self.0
55 }
56
57 fn duration(self) -> Duration {
58 Duration::from_millis(self.0.get())
59 }
60}
61
62pub struct SubprocessSource {
64 kind: &'static str,
72 capabilities: Capabilities,
74 writes: WriteSupport,
80 connection: Connection,
82}
83
84impl std::fmt::Debug for SubprocessSource {
85 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88 f.debug_struct("SubprocessSource")
89 .field("kind", &self.kind)
90 .finish_non_exhaustive()
91 }
92}
93
94impl SubprocessSource {
95 pub fn connect(
104 program: &str,
105 args: &[String],
106 name: &SourceName,
107 config: &Value,
108 secrets: BTreeMap<String, String>,
109 ) -> Result<Self, SourceError> {
110 Self::connect_with_deadline(
111 program,
112 args,
113 name,
114 config,
115 secrets,
116 RequestDeadline::DEFAULT,
117 )
118 }
119
120 pub fn connect_with_deadline(
122 program: &str,
123 args: &[String],
124 name: &SourceName,
125 config: &Value,
126 secrets: BTreeMap<String, String>,
127 deadline: RequestDeadline,
128 ) -> Result<Self, SourceError> {
129 Self::adopt(
130 Peer::spawn(program, args, deadline.duration())?,
131 name,
132 config,
133 secrets,
134 )
135 }
136
137 pub fn over(
152 to_plugin: impl std::io::Write + Send + 'static,
153 from_plugin: impl std::io::Read + Send + 'static,
154 name: &SourceName,
155 config: &Value,
156 secrets: BTreeMap<String, String>,
157 ) -> Result<Self, SourceError> {
158 Self::over_with_request_deadline(
159 to_plugin,
160 from_plugin,
161 name,
162 config,
163 secrets,
164 RequestDeadline::DEFAULT,
165 )
166 }
167
168 pub fn over_with_request_deadline(
174 to_plugin: impl std::io::Write + Send + 'static,
175 from_plugin: impl std::io::Read + Send + 'static,
176 name: &SourceName,
177 config: &Value,
178 secrets: BTreeMap<String, String>,
179 deadline: RequestDeadline,
180 ) -> Result<Self, SourceError> {
181 Self::adopt(
182 Peer::over(to_plugin, from_plugin, deadline.duration()),
183 name,
184 config,
185 secrets,
186 )
187 }
188
189 fn adopt(
191 mut peer: Peer,
192 name: &SourceName,
193 config: &Value,
194 secrets: BTreeMap<String, String>,
195 ) -> Result<Self, SourceError> {
196 let result = Self::handshake(&mut peer, name, config, secrets);
197 let InitializeResult {
198 protocol_version,
199 kind,
200 capabilities,
201 writes,
202 } = match result {
203 Ok(result) => result,
204 Err(error) => return Err(with_diagnostics(error, &mut peer)),
205 };
206 let kind = kind.into_string();
207 if protocol_version != Some(PROTOCOL_VERSION) {
208 return Err(SourceError::Config {
209 message: match protocol_version {
210 Some(spoken) => format!(
211 "the {kind:?} plugin was asked for protocol version \
212 {PROTOCOL_VERSION} and answered in version {spoken}; the two are \
213 incompatible and this engine does not guess between them"
214 ),
215 None => format!(
216 "the {kind:?} plugin did not say which protocol version it \
217 answered in; this engine speaks version {PROTOCOL_VERSION} and \
218 does not guess"
219 ),
220 },
221 });
222 }
223 Ok(Self {
224 kind: String::leak(kind),
225 capabilities,
226 writes: writes.unwrap_or(WriteSupport::Unsupported),
227 connection: Connection::adopt(peer),
228 })
229 }
230
231 fn handshake(
233 peer: &mut Peer,
234 name: &SourceName,
235 config: &Value,
236 secrets: BTreeMap<String, String>,
237 ) -> Result<InitializeResult, SourceError> {
238 let params = InitializeParams {
239 protocol_version: PROTOCOL_VERSION,
240 engine: EngineIdentity {
241 name: "onetaskgraph".to_owned(),
242 version: env!("CARGO_PKG_VERSION").to_owned(),
243 },
244 source_name: name.as_str().to_owned(),
245 config: config.clone(),
246 secrets,
247 };
248 let request = Request {
249 id: HANDSHAKE_ID.to_owned(),
250 method: "initialize".to_owned(),
251 params: serde_json::to_value(¶ms).expect("a handshake is plain data"),
254 };
255 let line = peer.exchange(
256 &serde_json::to_string(&request).expect("a handshake request is plain data"),
257 )?;
258 let response: super::wire::Response =
259 serde_json::from_str(&line).map_err(|error| SourceError::Malformed {
260 message: format!(
261 "the plugin's handshake answer is not a response envelope: {error}"
262 ),
263 })?;
264 if response.id != HANDSHAKE_ID {
269 return Err(SourceError::Malformed {
270 message: format!(
271 "the plugin answered the handshake with an envelope addressed to {:?} \
272 rather than to {HANDSHAKE_ID:?}",
273 response.id
274 ),
275 });
276 }
277 let outcome = response.outcome().ok_or_else(|| SourceError::Malformed {
278 message: "the plugin's handshake answer carried both a result and an error, or \
279 neither"
280 .to_owned(),
281 })?;
282 let result = outcome?;
283 serde_json::from_value(result).map_err(|error| SourceError::Malformed {
284 message: format!("the plugin's handshake answer is not an initialize result: {error}"),
285 })
286 }
287
288 async fn ask<T: for<'de> Deserialize<'de>>(
290 &self,
291 method: &str,
292 params: Value,
293 ) -> Result<T, SourceError> {
294 let result = self.connection.call(method, params).await?;
295 serde_json::from_value(result).map_err(|error| SourceError::Malformed {
296 message: format!(
297 "the plugin's answer to {method} is not the shape it promises: {error}"
298 ),
299 })
300 }
301}
302
303fn with_diagnostics(error: SourceError, peer: &mut Peer) -> SourceError {
308 let said = peer.said();
309 if said.is_empty() {
310 return error;
311 }
312 let message = format!("{error}; the plugin wrote: {said}");
313 match error {
314 SourceError::RateLimited { .. } => error,
315 SourceError::Config { .. } => SourceError::Config { message },
316 SourceError::Auth { .. } => SourceError::Auth { message },
317 SourceError::Refused { .. } => SourceError::Refused { message },
318 SourceError::Malformed { .. } => SourceError::Malformed { message },
319 SourceError::Unavailable { .. } => SourceError::Unavailable { message },
320 }
321}
322
323#[async_trait]
324impl TaskSource for SubprocessSource {
325 fn kind(&self) -> &'static str {
326 self.kind
327 }
328
329 fn capabilities(&self) -> Capabilities {
330 self.capabilities.clone()
331 }
332
333 async fn health(&self) -> Result<Health, SourceError> {
334 self.ask("health", json!({})).await
335 }
336
337 async fn get_task(&self, id: &NativeId) -> Result<Option<Task>, SourceError> {
338 let result: TaskResult = self
339 .ask("get_task", params(&IdParams { id: id.clone() }))
340 .await?;
341 Ok(result.task)
342 }
343
344 async fn get_project(&self, id: &NativeId) -> Result<Option<Project>, SourceError> {
345 let result: ProjectResult = self
346 .ask("get_project", params(&IdParams { id: id.clone() }))
347 .await?;
348 Ok(result.project)
349 }
350
351 async fn query_tasks(
352 &self,
353 query: &TaskQuery,
354 page: &PageRequest,
355 ) -> Result<Page<Task>, SourceError> {
356 self.ask(
357 "query_tasks",
358 params(&TaskQueryParams {
359 query: query.clone(),
360 page: page.clone(),
361 }),
362 )
363 .await
364 }
365
366 async fn query_projects(
367 &self,
368 query: &ProjectQuery,
369 page: &PageRequest,
370 ) -> Result<Page<Project>, SourceError> {
371 self.ask(
372 "query_projects",
373 params(&ProjectQueryParams {
374 query: query.clone(),
375 page: page.clone(),
376 }),
377 )
378 .await
379 }
380
381 async fn labels(&self, page: &PageRequest) -> Result<Page<Label>, SourceError> {
382 self.ask("labels", params(&LabelParams { page: page.clone() }))
383 .await
384 }
385
386 async fn task_dependencies(
387 &self,
388 id: &NativeId,
389 direction: Direction,
390 page: &PageRequest,
391 ) -> Result<Page<DependencyEdge>, SourceError> {
392 self.ask(
393 "task_dependencies",
394 params(&DependencyParams {
395 id: id.clone(),
396 direction,
397 page: page.clone(),
398 }),
399 )
400 .await
401 }
402
403 async fn project_dependencies(
404 &self,
405 id: &NativeId,
406 direction: Direction,
407 page: &PageRequest,
408 ) -> Result<Page<DependencyEdge>, SourceError> {
409 self.ask(
410 "project_dependencies",
411 params(&DependencyParams {
412 id: id.clone(),
413 direction,
414 page: page.clone(),
415 }),
416 )
417 .await
418 }
419
420 fn writes(&self) -> WriteSupport {
421 self.writes
422 }
423
424 async fn write_task(&self, write: &ItemWrite<Task>) -> Result<NativeId, SourceError> {
425 let result: WriteResult = self
426 .ask(
427 "write_task",
428 params(&TaskWriteParams {
429 write: write.clone(),
430 }),
431 )
432 .await?;
433 Ok(result.id)
434 }
435
436 async fn write_project(&self, write: &ItemWrite<Project>) -> Result<NativeId, SourceError> {
437 let result: WriteResult = self
438 .ask(
439 "write_project",
440 params(&ProjectWriteParams {
441 write: write.clone(),
442 }),
443 )
444 .await?;
445 Ok(result.id)
446 }
447
448 async fn delete_task(&self, id: &NativeId) -> Result<(), SourceError> {
449 let _: IgnoredResult = self
450 .ask("delete_task", params(&DeleteParams { id: id.clone() }))
451 .await?;
452 Ok(())
453 }
454
455 async fn delete_project(&self, id: &NativeId) -> Result<(), SourceError> {
456 let _: IgnoredResult = self
457 .ask("delete_project", params(&DeleteParams { id: id.clone() }))
458 .await?;
459 Ok(())
460 }
461
462 async fn get_document(&self, id: &NativeId) -> Result<Option<Document>, SourceError> {
463 let result: DocumentResult = self
464 .ask("get_document", params(&IdParams { id: id.clone() }))
465 .await?;
466 Ok(result.document)
467 }
468
469 async fn query_documents(
470 &self,
471 query: &DocumentQuery,
472 page: &PageRequest,
473 ) -> Result<Page<Document>, SourceError> {
474 self.ask(
475 "query_documents",
476 params(&DocumentQueryParams {
477 query: query.clone(),
478 page: page.clone(),
479 }),
480 )
481 .await
482 }
483
484 async fn write_document(&self, write: &ItemWrite<Document>) -> Result<NativeId, SourceError> {
485 let result: WriteResult = self
486 .ask(
487 "write_document",
488 params(&DocumentWriteParams {
489 write: write.clone(),
490 }),
491 )
492 .await?;
493 Ok(result.id)
494 }
495
496 async fn delete_document(&self, id: &NativeId) -> Result<(), SourceError> {
497 let _: IgnoredResult = self
498 .ask("delete_document", params(&DeleteParams { id: id.clone() }))
499 .await?;
500 Ok(())
501 }
502}
503
504#[derive(serde::Deserialize)]
513struct IgnoredResult {}
514
515fn params<T: serde::Serialize>(value: &T) -> Value {
520 serde_json::to_value(value).expect("method parameters are plain data")
521}