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 {
317 retry_after_seconds,
318 ..
319 } => SourceError::RateLimited {
320 retry_after_seconds,
321 message: Some(message),
322 },
323 SourceError::Config { .. } => SourceError::Config { message },
324 SourceError::Auth { .. } => SourceError::Auth { message },
325 SourceError::Refused { .. } => SourceError::Refused { message },
326 SourceError::Malformed { .. } => SourceError::Malformed { message },
327 SourceError::Unavailable { .. } => SourceError::Unavailable { message },
328 }
329}
330
331#[async_trait]
332impl TaskSource for SubprocessSource {
333 fn kind(&self) -> &'static str {
334 self.kind
335 }
336
337 fn capabilities(&self) -> Capabilities {
338 self.capabilities.clone()
339 }
340
341 async fn health(&self) -> Result<Health, SourceError> {
342 self.ask("health", json!({})).await
343 }
344
345 async fn get_task(&self, id: &NativeId) -> Result<Option<Task>, SourceError> {
346 let result: TaskResult = self
347 .ask("get_task", params(&IdParams { id: id.clone() }))
348 .await?;
349 Ok(result.task)
350 }
351
352 async fn get_project(&self, id: &NativeId) -> Result<Option<Project>, SourceError> {
353 let result: ProjectResult = self
354 .ask("get_project", params(&IdParams { id: id.clone() }))
355 .await?;
356 Ok(result.project)
357 }
358
359 async fn query_tasks(
360 &self,
361 query: &TaskQuery,
362 page: &PageRequest,
363 ) -> Result<Page<Task>, SourceError> {
364 self.ask(
365 "query_tasks",
366 params(&TaskQueryParams {
367 query: query.clone(),
368 page: page.clone(),
369 }),
370 )
371 .await
372 }
373
374 async fn query_projects(
375 &self,
376 query: &ProjectQuery,
377 page: &PageRequest,
378 ) -> Result<Page<Project>, SourceError> {
379 self.ask(
380 "query_projects",
381 params(&ProjectQueryParams {
382 query: query.clone(),
383 page: page.clone(),
384 }),
385 )
386 .await
387 }
388
389 async fn labels(&self, page: &PageRequest) -> Result<Page<Label>, SourceError> {
390 self.ask("labels", params(&LabelParams { page: page.clone() }))
391 .await
392 }
393
394 async fn task_dependencies(
395 &self,
396 id: &NativeId,
397 direction: Direction,
398 page: &PageRequest,
399 ) -> Result<Page<DependencyEdge>, SourceError> {
400 self.ask(
401 "task_dependencies",
402 params(&DependencyParams {
403 id: id.clone(),
404 direction,
405 page: page.clone(),
406 }),
407 )
408 .await
409 }
410
411 async fn project_dependencies(
412 &self,
413 id: &NativeId,
414 direction: Direction,
415 page: &PageRequest,
416 ) -> Result<Page<DependencyEdge>, SourceError> {
417 self.ask(
418 "project_dependencies",
419 params(&DependencyParams {
420 id: id.clone(),
421 direction,
422 page: page.clone(),
423 }),
424 )
425 .await
426 }
427
428 fn writes(&self) -> WriteSupport {
429 self.writes
430 }
431
432 async fn write_task(&self, write: &ItemWrite<Task>) -> Result<NativeId, SourceError> {
433 let result: WriteResult = self
434 .ask(
435 "write_task",
436 params(&TaskWriteParams {
437 write: write.clone(),
438 }),
439 )
440 .await?;
441 Ok(result.id)
442 }
443
444 async fn write_project(&self, write: &ItemWrite<Project>) -> Result<NativeId, SourceError> {
445 let result: WriteResult = self
446 .ask(
447 "write_project",
448 params(&ProjectWriteParams {
449 write: write.clone(),
450 }),
451 )
452 .await?;
453 Ok(result.id)
454 }
455
456 async fn delete_task(&self, id: &NativeId) -> Result<(), SourceError> {
457 let _: IgnoredResult = self
458 .ask("delete_task", params(&DeleteParams { id: id.clone() }))
459 .await?;
460 Ok(())
461 }
462
463 async fn delete_project(&self, id: &NativeId) -> Result<(), SourceError> {
464 let _: IgnoredResult = self
465 .ask("delete_project", params(&DeleteParams { id: id.clone() }))
466 .await?;
467 Ok(())
468 }
469
470 async fn get_document(&self, id: &NativeId) -> Result<Option<Document>, SourceError> {
471 let result: DocumentResult = self
472 .ask("get_document", params(&IdParams { id: id.clone() }))
473 .await?;
474 Ok(result.document)
475 }
476
477 async fn query_documents(
478 &self,
479 query: &DocumentQuery,
480 page: &PageRequest,
481 ) -> Result<Page<Document>, SourceError> {
482 self.ask(
483 "query_documents",
484 params(&DocumentQueryParams {
485 query: query.clone(),
486 page: page.clone(),
487 }),
488 )
489 .await
490 }
491
492 async fn write_document(&self, write: &ItemWrite<Document>) -> Result<NativeId, SourceError> {
493 let result: WriteResult = self
494 .ask(
495 "write_document",
496 params(&DocumentWriteParams {
497 write: write.clone(),
498 }),
499 )
500 .await?;
501 Ok(result.id)
502 }
503
504 async fn delete_document(&self, id: &NativeId) -> Result<(), SourceError> {
505 let _: IgnoredResult = self
506 .ask("delete_document", params(&DeleteParams { id: id.clone() }))
507 .await?;
508 Ok(())
509 }
510}
511
512#[derive(serde::Deserialize)]
521struct IgnoredResult {}
522
523fn params<T: serde::Serialize>(value: &T) -> Value {
528 serde_json::to_value(value).expect("method parameters are plain data")
529}